Update app.py

app.py

@@ -141,7 +141,7 @@ def messages():
             yield create_event("ping", {"type": "ping"})
 
             def on_query_progress(data):
-                nonlocal total_output_tokens
+                nonlocal total_output_tokens, response_text
                 if 'text' in data:
                     text = json.loads(data['text'])
                     chunk = text['chunks'][-1] if text['chunks'] else None
@@ -150,6 +150,9 @@ def messages():
                         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")
                 emit_data = {
@@ -168,19 +171,8 @@ def messages():
                 }
                 sio.emit('perplexity_ask', (previous_messages, emit_data))
 
-            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)
 
             try:
                 sio.connect('wss://www.perplexity.ai/', **connect_opts, headers=sio_opts['extraHeaders'])