Browse Source

feat: 增加appName、输入、输出、总的token字段,后续可通过appName+timestamp统计总量

zhangwl 1 week ago
parent
commit
23f8258395
2 changed files with 18 additions and 0 deletions
  1. 8 0
      app/db/mongo.py
  2. 10 0
      app/routers/chat.py

+ 8 - 0
app/db/mongo.py

@@ -72,6 +72,10 @@ def save_chat_history(
     thinking: str = None,
     thinking: str = None,
     searching: str = None,
     searching: str = None,
     attachments: list = None,
     attachments: list = None,
+    input_tokens: int = None,
+    output_tokens: int = None,
+    total_tokens: int = None,
+    app_name:str=None,
 ):
 ):
     try:
     try:
         _ensure_index()
         _ensure_index()
@@ -83,7 +87,11 @@ def save_chat_history(
             "thinking": thinking,
             "thinking": thinking,
             "searching": searching,
             "searching": searching,
             "response_id": response_id,
             "response_id": response_id,
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "total_tokens": total_tokens,
             "timestamp": timestamp,
             "timestamp": timestamp,
+            "app_name": app_name,
         }
         }
         if attachments:
         if attachments:
             doc["attachments"] = attachments
             doc["attachments"] = attachments

+ 10 - 0
app/routers/chat.py

@@ -320,6 +320,7 @@ async def generate_stream_response(request: ChatRequest, user_context: UserConte
         accumulated_thinking = ""
         accumulated_thinking = ""
         accumulated_searching = ""
         accumulated_searching = ""
         response_id = None
         response_id = None
+        input_tokens = output_tokens = total_tokens = None
 
 
         # 将同步阻塞的 stream 迭代放入子线程,通过 Queue 传递给异步生成器
         # 将同步阻塞的 stream 迭代放入子线程,通过 Queue 传递给异步生成器
         # 避免阻塞事件循环,保证每个 chunk 到达时立即 yield 推送给前端
         # 避免阻塞事件循环,保证每个 chunk 到达时立即 yield 推送给前端
@@ -389,6 +390,11 @@ async def generate_stream_response(request: ChatRequest, user_context: UserConte
                 response_obj = getattr(chunk, 'response', None)
                 response_obj = getattr(chunk, 'response', None)
                 if response_obj and hasattr(response_obj, 'id'):
                 if response_obj and hasattr(response_obj, 'id'):
                     response_id = response_obj.id
                     response_id = response_obj.id
+                usage_obj = getattr(response_obj, 'usage', None) if response_obj else None
+                if usage_obj:
+                    input_tokens = getattr(usage_obj, 'input_tokens', None)
+                    output_tokens = getattr(usage_obj, 'output_tokens', None)
+                    total_tokens = getattr(usage_obj, 'total_tokens', None)
                 save_chat_log(
                 save_chat_log(
                     user_id=user_context.user_id,
                     user_id=user_context.user_id,
                     question=latest_user_msg.content,
                     question=latest_user_msg.content,
@@ -410,6 +416,10 @@ async def generate_stream_response(request: ChatRequest, user_context: UserConte
                 response_id=response_id,
                 response_id=response_id,
                 thinking=accumulated_thinking or None,
                 thinking=accumulated_thinking or None,
                 searching=accumulated_searching or None,
                 searching=accumulated_searching or None,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                total_tokens=total_tokens,
+                app_name=user_context.app_name,
             )
             )
             yield f"data: {StreamResponse(content='', finished=True, model=config.MODEL_NAME, timestamp=datetime.now()).model_dump_json()}\n\n"
             yield f"data: {StreamResponse(content='', finished=True, model=config.MODEL_NAME, timestamp=datetime.now()).model_dump_json()}\n\n"