mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-06-02 06:09:39 +08:00
feat: 增强流式响应处理,支持使用元数据
本次提交对流式响应处理进行了增强,主要变更包括: - **参数更新**: - 在 `_handle_openai_stream_response` 方法中新增 `usage_metadata` 参数,以支持传递使用情况的元数据。 - **数据结构调整**: - 在返回的响应中,若提供了 `usage_metadata`,则将其包含在返回的 JSON 结构中,确保更全面的响应信息。 - **伪流式逻辑更新**: - 在 `OpenAIChatService` 中的多个方法中,更新了对流式响应的调用,确保在处理响应时也能传递和使用元数据。 这些更改旨在提升流式响应的灵活性和信息丰富性,改善用户体验。
This commit is contained in:
@@ -37,7 +37,7 @@ class GeminiResponseHandler(ResponseHandler):
|
||||
|
||||
|
||||
def _handle_openai_stream_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]:
|
||||
text, tool_calls = _extract_result(
|
||||
response, model, stream=True, gemini_format=False
|
||||
@@ -48,14 +48,16 @@ def _handle_openai_stream_response(
|
||||
delta = {"content": text, "role": "assistant"}
|
||||
if tool_calls:
|
||||
delta["tool_calls"] = tool_calls
|
||||
|
||||
return {
|
||||
template_chunk = {
|
||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
||||
}
|
||||
if usage_metadata:
|
||||
template_chunk["usage"] = {"prompt_tokens": usage_metadata.get("promptTokenCount", 0), "completion_tokens": usage_metadata.get("candidatesTokenCount",0), "total_tokens": usage_metadata.get("totalTokenCount", 0)}
|
||||
return template_chunk
|
||||
|
||||
|
||||
def _handle_openai_normal_response(
|
||||
@@ -101,7 +103,7 @@ class OpenAIResponseHandler(ResponseHandler):
|
||||
usage_metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
if stream:
|
||||
return _handle_openai_stream_response(response, model, finish_reason)
|
||||
return _handle_openai_stream_response(response, model, finish_reason, usage_metadata)
|
||||
return _handle_openai_normal_response(response, model, finish_reason, usage_metadata)
|
||||
|
||||
def handle_image_chat_response(
|
||||
|
||||
Reference in New Issue
Block a user