Update app.py

app.py

@@ -75,6 +75,11 @@ def calculate_tokens(text):
         tokens = text.split()
         return len(tokens)
 
+def create_event(event, data):
+    if isinstance(data, dict):
+        data = json.dumps(data, ensure_ascii=False)
+    return f"event: {event}\ndata: {data}\n\n"
+
 @app.route('/')
 def root():
     log_request(request.remote_addr, request.path, 200)
@@ -217,6 +222,20 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
     try:
         response_event = Event()
         response_text = []
+        total_output_tokens = 0
+
+        def on_query_progress(data):
+            nonlocal response_text, total_output_tokens
+            if 'text' in data:
+                text = json.loads(data['text'])
+                chunk = text['chunks'][-1] if text['chunks'] else None
+                if chunk:
+                    response_text.append(chunk)
+                    chunk_tokens = calculate_tokens(chunk)
+                    total_output_tokens += chunk_tokens
+
+            if data.get('final', False):
+                response_event.set()
 
         def on_connect():
             logging.info("Connected to Perplexity AI")
@@ -236,37 +255,13 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
             }
             sio.emit('perplexity_ask', (previous_messages, emit_data))
 
-        def on_query_progress(data):
-            nonlocal response_text
-            if 'text' in data:
-                text = json.loads(data['text'])
-                chunk = text['chunks'][-1] if text['chunks'] else None
-                if chunk:
-                    response_text.append(chunk)
-
-            if data.get('final', False):
-                response_event.set()
-
-        def on_disconnect():
-            logging.info("Disconnected from Perplexity AI")
-            response_event.set()
-
-        def on_connect_error(data):
-            logging.error(f"Connection error: {data}")
-            response_text.append(f"Error connecting to Perplexity AI: {data}")
-            response_event.set()
-
         sio.on('connect', on_connect)
         sio.on('query_progress', on_query_progress)
-        sio.on('disconnect', on_disconnect)
-        sio.on('connect_error', on_connect_error)
 
         sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])
         
         response_event.wait(timeout=30)
 
-        output_tokens = calculate_tokens(''.join(response_text))
-
         full_response = {
             "content": [{"text": ''.join(response_text), "type": "text"}],
             "id": msg_id,
@@ -277,7 +272,7 @@ def handle_non_stream(previous_messages, msg_id, model, input_tokens):
             "type": "message",
             "usage": {
                 "input_tokens": input_tokens,
-                "output_tokens": output_tokens,
+                "output_tokens": total_output_tokens,
             },
         }
         return Response(json.dumps(full_response, ensure_ascii=False), content_type='application/json')
@@ -300,14 +295,9 @@ def server_error(error):
     log_request(request.remote_addr, request.path, 500)
     return "Something broke!", 500
 
-def create_event(event, data):
-    if isinstance(data, dict):
-        data = json.dumps(data, ensure_ascii=False)
-    return f"event: {event}\ndata: {data}\n\n"
-
 if __name__ == '__main__':
     port = int(os.environ.get('PORT', 8081))
     logging.info(f"Perplexity proxy listening on port {port}")
     if not API_KEY:
         logging.warning("Warning: PPLX_KEY environment variable is not set. API key validation will fail.")
-    app.run(host='0.0.0.0', port=port)
+    app.run(host='0.0.0.0', port=port)