Update app.py
Browse files
app.py
CHANGED
@@ -75,6 +75,11 @@ def calculate_tokens(text):
|
|
75 |
tokens = text.split()
|
76 |
return len(tokens)
|
77 |
|
|
|
|
|
|
|
|
|
|
|
78 |
@app.route('/')
|
79 |
def root():
|
80 |
log_request(request.remote_addr, request.path, 200)
|
@@ -217,6 +222,20 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
|
|
217 |
try:
|
218 |
response_event = Event()
|
219 |
response_text = []
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
220 |
|
221 |
def on_connect():
|
222 |
logging.info("Connected to Perplexity AI")
|
@@ -236,37 +255,13 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
|
|
236 |
}
|
237 |
sio.emit('perplexity_ask', (previous_messages, emit_data))
|
238 |
|
239 |
-
def on_query_progress(data):
|
240 |
-
nonlocal response_text
|
241 |
-
if 'text' in data:
|
242 |
-
text = json.loads(data['text'])
|
243 |
-
chunk = text['chunks'][-1] if text['chunks'] else None
|
244 |
-
if chunk:
|
245 |
-
response_text.append(chunk)
|
246 |
-
|
247 |
-
if data.get('final', False):
|
248 |
-
response_event.set()
|
249 |
-
|
250 |
-
def on_disconnect():
|
251 |
-
logging.info("Disconnected from Perplexity AI")
|
252 |
-
response_event.set()
|
253 |
-
|
254 |
-
def on_connect_error(data):
|
255 |
-
logging.error(f"Connection error: {data}")
|
256 |
-
response_text.append(f"Error connecting to Perplexity AI: {data}")
|
257 |
-
response_event.set()
|
258 |
-
|
259 |
sio.on('connect', on_connect)
|
260 |
sio.on('query_progress', on_query_progress)
|
261 |
-
sio.on('disconnect', on_disconnect)
|
262 |
-
sio.on('connect_error', on_connect_error)
|
263 |
|
264 |
sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
|
265 |
|
266 |
response_event.wait(timeout=30)
|
267 |
|
268 |
-
output_tokens = calculate_tokens(''.join(response_text))
|
269 |
-
|
270 |
full_response = {
|
271 |
"content": [{"text": ''.join(response_text), "type": "text"}],
|
272 |
"id": msg_id,
|
@@ -277,7 +272,7 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
|
|
277 |
"type": "message",
|
278 |
"usage": {
|
279 |
"input_tokens": input_tokens,
|
280 |
-
"output_tokens":
|
281 |
},
|
282 |
}
|
283 |
return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
|
@@ -300,14 +295,9 @@ def server_error(error):
|
|
300 |
log_request(request.remote_addr, request.path, 500)
|
301 |
return "Something broke!", 500
|
302 |
|
303 |
-
def create_event(event, data):
|
304 |
-
if isinstance(data, dict):
|
305 |
-
data = json.dumps(data, ensure_ascii=False)
|
306 |
-
return f"event: {event}\ndata: {data}\n\n"
|
307 |
-
|
308 |
if __name__ == '__main__':
|
309 |
port = int(os.environ.get('PORT', 8081))
|
310 |
logging.info(f"Perplexity proxy listening on port {port}")
|
311 |
if not API_KEY:
|
312 |
logging.warning("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
|
313 |
-
app.run(host='0.0.0.0', port=port)
|
|
|
75 |
tokens = text.split()
|
76 |
return len(tokens)
|
77 |
|
78 |
+
def create_event(event, data):
|
79 |
+
if isinstance(data, dict):
|
80 |
+
data = json.dumps(data, ensure_ascii=False)
|
81 |
+
return f"event: {event}\ndata: {data}\n\n"
|
82 |
+
|
83 |
@app.route('/')
|
84 |
def root():
|
85 |
log_request(request.remote_addr, request.path, 200)
|
|
|
222 |
try:
|
223 |
response_event = Event()
|
224 |
response_text = []
|
225 |
+
total_output_tokens = 0
|
226 |
+
|
227 |
+
def on_query_progress(data):
|
228 |
+
nonlocal response_text, total_output_tokens
|
229 |
+
if 'text' in data:
|
230 |
+
text = json.loads(data['text'])
|
231 |
+
chunk = text['chunks'][-1] if text['chunks'] else None
|
232 |
+
if chunk:
|
233 |
+
response_text.append(chunk)
|
234 |
+
chunk_tokens = calculate_tokens(chunk)
|
235 |
+
total_output_tokens += chunk_tokens
|
236 |
+
|
237 |
+
if data.get('final', False):
|
238 |
+
response_event.set()
|
239 |
|
240 |
def on_connect():
|
241 |
logging.info("Connected to Perplexity AI")
|
|
|
255 |
}
|
256 |
sio.emit('perplexity_ask', (previous_messages, emit_data))
|
257 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
258 |
sio.on('connect', on_connect)
|
259 |
sio.on('query_progress', on_query_progress)
|
|
|
|
|
260 |
|
261 |
sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
|
262 |
|
263 |
response_event.wait(timeout=30)
|
264 |
|
|
|
|
|
265 |
full_response = {
|
266 |
"content": [{"text": ''.join(response_text), "type": "text"}],
|
267 |
"id": msg_id,
|
|
|
272 |
"type": "message",
|
273 |
"usage": {
|
274 |
"input_tokens": input_tokens,
|
275 |
+
"output_tokens": total_output_tokens,
|
276 |
},
|
277 |
}
|
278 |
return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
|
|
|
295 |
log_request(request.remote_addr, request.path, 500)
|
296 |
return "Something broke!", 500
|
297 |
|
|
|
|
|
|
|
|
|
|
|
298 |
if __name__ == '__main__':
|
299 |
port = int(os.environ.get('PORT', 8081))
|
300 |
logging.info(f"Perplexity proxy listening on port {port}")
|
301 |
if not API_KEY:
|
302 |
logging.warning("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
|
303 |
+
app.run(host='0.0.0.0', port=port)
|