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

优化Gemini响应处理逻辑，改进错误处理和代码结构

Browse Source
This commit is contained in:
yinpeng
2024-12-15 15:23:25 +08:00
parent 17d1750329
commit a5adcdae48
1 changed files with 37 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/services/chat_service.py
+37 -51
View File
@@ -65,27 +65,24 @@ 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:
candidate = response["candidates"][0] if response.get("candidates"):
content = candidate.get("content", {}) candidate = response["candidates"][0]
parts = content.get("parts", []) content = candidate.get("content", {})
parts = content.get("parts", [])
if not parts: if "text" in parts[0]:
return None text = parts[0].get("text")
elif "executableCode" in parts[0]:
if "text" in parts[0]: text = self.format_code_block(parts[0]["executableCode"])
text = parts[0].get("text") elif "codeExecution" in parts[0]:
elif "executableCode" in parts[0]: text = self.format_code_block(parts[0]["codeExecution"])
text = self.format_code_block(parts[0]["executableCode"]) elif "executableCodeResult" in parts[0]:
elif "codeExecution" in parts[0]: text = self.format_execution_result(parts[0]["executableCodeResult"])
text = self.format_code_block(parts[0]["codeExecution"]) elif "codeExecutionResult" in parts[0]:
elif "executableCodeResult" in parts[0]: text = self.format_execution_result(parts[0]["codeExecutionResult"])
text = self.format_execution_result(parts[0]["executableCodeResult"]) else:
elif "codeExecutionResult" in parts[0]: text = ""
text = self.format_execution_result(parts[0]["codeExecutionResult"])
else: else:
text = "" text = ""
@@ -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:
async with httpx.AsyncClient() as client: try:
url = f"https://generativelanguage.googleapis.com/v1beta/models/{gemini_model}:generateContent?key={api_key}" async with httpx.AsyncClient() as client:
response = await client.post(url, json=payload) url = f"https://generativelanguage.googleapis.com/v1beta/models/{gemini_model}:generateContent?key={api_key}"
gemini_response = response.json() response = await client.post(url, json=payload)
return self.convert_gemini_response_to_openai(gemini_response, model, finish_reason="stop") if response.status_code != 200:
raise Exception(f"API error: {response.status_code}")
gemini_response = response.json()
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,