mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-09-06 08:06:37 +08:00
优化Gemini响应处理逻辑，改进错误处理和代码结构
This commit is contained in:
yinpeng
@@ -65,17 +65,12 @@ class ChatService:
|
|||||||
) -> Optional[Dict[str, Any]]:
|
) -> Optional[Dict[str, Any]]:
|
||||||
"""Convert Gemini response to OpenAI format"""
|
"""Convert Gemini response to OpenAI format"""
|
||||||
if stream:
|
if stream:
|
||||||
if not response.get("candidates"):
|
|
||||||
return None
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
if response.get("candidates"):
|
||||||
candidate = response["candidates"][0]
|
candidate = response["candidates"][0]
|
||||||
content = candidate.get("content", {})
|
content = candidate.get("content", {})
|
||||||
parts = content.get("parts", [])
|
parts = content.get("parts", [])
|
||||||
|
|
||||||
if not parts:
|
|
||||||
return None
|
|
||||||
|
|
||||||
if "text" in parts[0]:
|
if "text" in parts[0]:
|
||||||
text = parts[0].get("text")
|
text = parts[0].get("text")
|
||||||
elif "executableCode" in parts[0]:
|
elif "executableCode" in parts[0]:
|
||||||
@@ -88,6 +83,8 @@ class ChatService:
|
|||||||
text = self.format_execution_result(parts[0]["codeExecutionResult"])
|
text = self.format_execution_result(parts[0]["codeExecutionResult"])
|
||||||
else:
|
else:
|
||||||
text = ""
|
text = ""
|
||||||
|
else:
|
||||||
|
text = ""
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||||
@@ -180,7 +177,6 @@ class ChatService:
|
|||||||
}
|
}
|
||||||
|
|
||||||
if stream:
|
if stream:
|
||||||
|
|
||||||
async def generate():
|
async def generate():
|
||||||
retries = 0
|
retries = 0
|
||||||
MAX_RETRIES = 3
|
MAX_RETRIES = 3
|
||||||
@@ -190,22 +186,12 @@ class ChatService:
|
|||||||
try:
|
try:
|
||||||
async with httpx.AsyncClient() as client:
|
async with httpx.AsyncClient() as client:
|
||||||
stream_url = f"https://generativelanguage.googleapis.com/v1beta/models/{gemini_model}:streamGenerateContent?alt=sse&key={current_api_key}"
|
stream_url = f"https://generativelanguage.googleapis.com/v1beta/models/{gemini_model}:streamGenerateContent?alt=sse&key={current_api_key}"
|
||||||
async with client.stream(
|
async with client.stream("POST", stream_url, json=payload) as response:
|
||||||
"POST", stream_url, json=payload
|
|
||||||
) as response:
|
|
||||||
if response.status_code != 200:
|
if response.status_code != 200:
|
||||||
|
error_msg = await response.text()
|
||||||
|
logger.error(f"API error: {response.status_code}, {error_msg}")
|
||||||
if retries < MAX_RETRIES - 1:
|
if retries < MAX_RETRIES - 1:
|
||||||
logger.warning(
|
current_api_key = await self.key_manager.handle_api_failure(current_api_key)
|
||||||
f"API error: {response.status_code}, attempting retry {retries + 1}"
|
|
||||||
)
|
|
||||||
current_api_key = (
|
|
||||||
await self.key_manager.handle_api_failure(
|
|
||||||
current_api_key
|
|
||||||
)
|
|
||||||
)
|
|
||||||
logger.info(
|
|
||||||
f"Switched to new API key: {current_api_key}"
|
|
||||||
)
|
|
||||||
retries += 1
|
retries += 1
|
||||||
continue
|
continue
|
||||||
else:
|
else:
|
||||||
@@ -219,27 +205,21 @@ class ChatService:
|
|||||||
if line.startswith("data: "):
|
if line.startswith("data: "):
|
||||||
try:
|
try:
|
||||||
chunk = json.loads(line[6:])
|
chunk = json.loads(line[6:])
|
||||||
openai_chunk = (
|
openai_chunk = self.convert_gemini_response_to_openai(
|
||||||
self.convert_gemini_response_to_openai(
|
|
||||||
chunk, model, stream=True, finish_reason=None
|
chunk, model, stream=True, finish_reason=None
|
||||||
)
|
)
|
||||||
)
|
|
||||||
if openai_chunk:
|
if openai_chunk:
|
||||||
yield f"data: {json.dumps(openai_chunk)}\n\n"
|
yield f"data: {json.dumps(openai_chunk)}\n\n"
|
||||||
except json.JSONDecodeError:
|
except json.JSONDecodeError:
|
||||||
continue
|
continue
|
||||||
yield f"data: {json.dumps({'finish_reason': 'stop'})}\n\n"
|
yield f"data: {json.dumps(self.convert_gemini_response_to_openai({}, model,stream=True, finish_reason='stop'))}\n\n"
|
||||||
yield "data: [DONE]\n\n"
|
yield "data: [DONE]\n\n"
|
||||||
return # 成功完成,退出重试循环
|
return
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
logger.warning(f"Stream error: {str(e)}, attempting retry {retries + 1}")
|
||||||
if retries < MAX_RETRIES - 1:
|
if retries < MAX_RETRIES - 1:
|
||||||
logger.warning(
|
current_api_key = await self.key_manager.handle_api_failure(current_api_key)
|
||||||
f"Stream error: {str(e)}, attempting retry {retries + 1}"
|
|
||||||
)
|
|
||||||
current_api_key = await self.key_manager.handle_api_failure(
|
|
||||||
current_api_key
|
|
||||||
)
|
|
||||||
retries += 1
|
retries += 1
|
||||||
continue
|
continue
|
||||||
else:
|
else:
|
||||||
@@ -249,11 +229,17 @@ class ChatService:
|
|||||||
|
|
||||||
return generate()
|
return generate()
|
||||||
else:
|
else:
|
||||||
|
try:
|
||||||
async with httpx.AsyncClient() as client:
|
async with httpx.AsyncClient() as client:
|
||||||
url = f"https://generativelanguage.googleapis.com/v1beta/models/{gemini_model}:generateContent?key={api_key}"
|
url = f"https://generativelanguage.googleapis.com/v1beta/models/{gemini_model}:generateContent?key={api_key}"
|
||||||
response = await client.post(url, json=payload)
|
response = await client.post(url, json=payload)
|
||||||
|
if response.status_code != 200:
|
||||||
|
raise Exception(f"API error: {response.status_code}")
|
||||||
gemini_response = response.json()
|
gemini_response = response.json()
|
||||||
return self.convert_gemini_response_to_openai(gemini_response, model, finish_reason="stop")
|
return self.convert_gemini_response_to_openai(gemini_response, model, finish_reason="stop")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error in non-stream completion: {str(e)}")
|
||||||
|
raise
|
||||||
|
|
||||||
async def _openai_chat_completion(
|
async def _openai_chat_completion(
|
||||||
self,
|
self,
|
||||||
|
|||||||
Reference in New Issue
Block a user