Update app.py

app.py

@@ -142,6 +142,7 @@ def messages():
         log_request(request.remote_addr, request.path, 200)
 
         def generate():
+            nonlocal output_tokens  # 使用 nonlocal 声明 output_tokens
             yield create_event("message_start", {
                 "type": "message_start",
                 "message": {
@@ -152,7 +153,7 @@ def messages():
                     "model": model,  # 动态模型
                     "stop_reason": None,
                     "stop_sequence": None,
-                    "usage": {"input_tokens": input_tokens, "output_tokens": 0},  # 动态 input_tokens，初始化 output_tokens
+                    "usage": {"input_tokens": input_tokens, "output_tokens": 1},  # 动态 input_tokens
                 },
             })
             yield create_event("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}})
@@ -177,33 +178,24 @@ def messages():
                 sio.emit('perplexity_ask', (previous_messages, emit_data))
 
             def on_query_progress(data):
-                nonlocal response_text, output_tokens
+                nonlocal response_text, output_tokens  # 使用 nonlocal 声明 response_text 和 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)
-                        # 动态计算当前块的 token 数量并累积到 output_tokens
-                        output_tokens += calculate_tokens(chunk)
+                        output_tokens = calculate_tokens(''.join(response_text))  # 更新 output_tokens
 
-                        # 发送 message_delta 事件，包含当前累积的 output_tokens
+                        # 在 message_delta 事件中包含 output_tokens
                         yield create_event("message_delta", {
                             "type": "message_delta",
                             "delta": {"stop_reason": None, "stop_sequence": None},
-                            "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},  # 累积 output_tokens
+                            "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},
                         })
 
                 # 检查是否是最终响应
                 if data.get('final', False):
-                    # 在响应结束时发送结束事件
-                    yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
-                    yield create_event("message_delta", {
-                        "type": "message_delta",
-                        "delta": {"stop_reason": "end_turn", "stop_sequence": None},
-                        "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},  # 最终 output_tokens
-                    })
-                    yield create_event("message_stop", {"type": "message_stop"})  # 确保发送 message_stop 事件
-                    response_event.set()  # 触发响应结束
+                    response_event.set()
 
             def on_query_complete(data):
                 response_event.set()
@@ -228,13 +220,6 @@ 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)}")
@@ -247,6 +232,14 @@ def messages():
                 if sio.connected:
                     sio.disconnect()
 
+            yield create_event("content_block_stop", {"type": "content_block_stop", "index": 0})
+            yield create_event("message_delta", {
+                "type": "message_delta",
+                "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+                "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},  # 动态 output_tokens
+            })
+            yield create_event("message_stop", {"type": "message_stop"})  # 确保发送 message_stop 事件
+
         return Response(generate(), content_type='text/event-stream')
 
     except Exception as e: