GithubBackup/gemini-balance
1
0
Fork 0
mirror of https://github.com/snailyp/gemini-balance.git synced 2026-09-05 07:26:39 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 在 OpenAI 聊天响应中集成 usage_metadata 以跟踪 token 使用情况

Browse Source 
此更改将 `usage_metadata` 参数添加到了 `app/handler/response_handler.py` 和 `app/service/chat/openai_chat_service.py` 中的相关函数。

`usage_metadata`（通常包含 token 计数：prompt_tokens, completion_tokens, total_tokens）现在会从 OpenAI API 响应中提取，并用于填充标准化响应格式中的 `usage` 字段。

这样可以更准确地跟踪 OpenAI 聊天完成接口的 token 消耗。
This commit is contained in:
snaily
2025-05-06 18:32:47 +08:00
parent d4a3ed3a57
commit f13a4fba5f
2 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/handler/response_handler.py
+5 -4
View File
@@ -29,7 +29,7 @@ class GeminiResponseHandler(ResponseHandler):
self.thinking_status = False self.thinking_status = False
def handle_response( def handle_response(
self, response: Dict[str, Any], model: str, stream: bool = False self, response: Dict[str, Any], model: str, stream: bool = False, usage_metadata: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]: ) -> Dict[str, Any]:
if stream: if stream:
return _handle_gemini_stream_response(response, model, stream) return _handle_gemini_stream_response(response, model, stream)
@@ -59,7 +59,7 @@ def _handle_openai_stream_response(
def _handle_openai_normal_response( def _handle_openai_normal_response(
response: Dict[str, Any], model: str, finish_reason: str response: Dict[str, Any], model: str, finish_reason: str, usage_metadata: Optional[Dict[str, Any]]
) -> Dict[str, Any]: ) -> Dict[str, Any]:
text, tool_calls = _extract_result( text, tool_calls = _extract_result(
response, model, stream=False, gemini_format=False response, model, stream=False, gemini_format=False
@@ -80,7 +80,7 @@ def _handle_openai_normal_response(
"finish_reason": finish_reason, "finish_reason": finish_reason,
} }
], ],
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, "usage": {"prompt_tokens": usage_metadata.get("promptTokenCount", 0), "completion_tokens": usage_metadata.get("candidatesTokenCount",0), "total_tokens": usage_metadata.get("totalTokenCount", 0)},
} }
@@ -98,10 +98,11 @@ class OpenAIResponseHandler(ResponseHandler):
model: str, model: str,
stream: bool = False, stream: bool = False,
finish_reason: str = None, finish_reason: str = None,
usage_metadata: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, Any]]: ) -> Optional[Dict[str, Any]]:
if stream: if stream:
return _handle_openai_stream_response(response, model, finish_reason) return _handle_openai_stream_response(response, model, finish_reason)
return _handle_openai_normal_response(response, model, finish_reason) return _handle_openai_normal_response(response, model, finish_reason, usage_metadata)
def handle_image_chat_response( def handle_image_chat_response(
self, image_str: str, model: str, stream=False, finish_reason="stop" self, image_str: str, model: str, stream=False, finish_reason="stop"
app/service/chat/openai_chat_service.py
+2 -1
View File
@@ -204,10 +204,11 @@ class OpenAIChatService:
response = None response = None
try: try:
response = await self.api_client.generate_content(payload, model, api_key) response = await self.api_client.generate_content(payload, model, api_key)
usage_metadata = response.get("usageMetadata", {})
is_success = True is_success = True
status_code = 200 status_code = 200
return self.response_handler.handle_response( return self.response_handler.handle_response(
response, model, stream=False, finish_reason="stop" response, model, stream=False, finish_reason="stop", usage_metadata=usage_metadata
) )
except Exception as e: except Exception as e:
is_success = False is_success = False