Update app.py

app.py

@@ -114,6 +114,7 @@ def messages():
         msg_id = str(uuid.uuid4())
         response_event = Event()
         response_text = []
+        total_output_tokens = 0
 
         if not stream:
             return handle_non_stream(previous_messages, msg_id, model, input_tokens)
@@ -121,7 +122,7 @@ def messages():
         log_request(request.remote_addr, request.path, 200)
 
         def generate():
-            total_output_tokens = 0
+            nonlocal total_output_tokens
 
             yield create_event("message_start", {
                 "type": "message_start",
@@ -139,6 +140,16 @@ def messages():
             yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
             yield create_event("ping", {"type": "ping"})
 
+            def on_query_progress(data):
+                nonlocal 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
+
             def on_connect():
                 logging.info("Connected to Perplexity AI")
                 emit_data = {
@@ -157,25 +168,6 @@ def messages():
                 }
                 sio.emit('perplexity_ask', (previous_messages, emit_data))
 
-            def on_query_progress(data):
-                nonlocal 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
-
-                        yield create_event("content_block_delta", {
-                            "type": "content_block_delta",
-                            "index": 0,
-                            "delta": {"type": "text_delta", "text": chunk},
-                        })
-
-                if data.get('final', False):
-                    response_event.set()
-
             def on_disconnect():
                 logging.info("Disconnected from Perplexity AI")
                 response_event.set()
@@ -195,6 +187,13 @@ def messages():
                 
                 while not response_event.is_set():
                     sio.sleep(0.1)
+                    while response_text:
+                        chunk = response_text.pop(0)
+                        yield create_event("content_block_delta", {
+                            "type": "content_block_delta",
+                            "index": 0,
+                            "delta": {"type": "text_delta", "text": chunk},
+                        })
                 
             except Exception as e:
                 logging.error(f"Error during socket connection: {str(e)}")