Update app.py

app.py

@@ -129,11 +129,11 @@ def messages():
                     "id": msg_id,
                     "type": "message",
                     "role": "assistant",
-                    "content": [],
                     "model": model,
+                    "content": [],
                     "stop_reason": None,
                     "stop_sequence": None,
-                    "usage": {"input_tokens": input_tokens, "output_tokens": 0},
+                    "usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
                 },
             })
             yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
@@ -158,7 +158,7 @@ def messages():
                 sio.emit('perplexity_ask', (previous_messages, emit_data))
 
             def on_query_progress(data):
-                nonlocal response_text, total_output_tokens
+                nonlocal total_output_tokens
                 if 'text' in data:
                     text = json.loads(data['text'])
                     chunk = text['chunks'][-1] if text['chunks'] else None
@@ -173,21 +173,9 @@ def messages():
                             "delta": {"type": "text_delta", "text": chunk},
                         })
 
-                        yield create_event("message_delta", {
-                            "type": "message_delta",
-                            "delta": {},
-                            "usage": {
-                                "input_tokens": input_tokens,
-                                "output_tokens": total_output_tokens
-                            },
-                        })
-
                 if data.get('final', False):
                     response_event.set()
 
-            def on_query_complete(data):
-                response_event.set()
-
             def on_disconnect():
                 logging.info("Disconnected from Perplexity AI")
                 response_event.set()
@@ -199,7 +187,6 @@ def messages():
 
             sio.on('connect', on_connect)
             sio.on('query_progress', on_query_progress)
-            sio.on('query_complete', on_query_complete)
             sio.on('disconnect', on_disconnect)
             sio.on('connect_error', on_connect_error)
 
@@ -224,7 +211,7 @@ def messages():
             yield create_event("message_delta", {
                 "type": "message_delta",
                 "delta": {"stop_reason": "end_turn", "stop_sequence": None},
-                "usage": {"input_tokens": input_tokens, "output_tokens": total_output_tokens},
+                "usage": {"output_tokens": total_output_tokens},
             })
             yield create_event("message_stop", {"type": "message_stop"})