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

feat: 将 Gemini API 调用迁移至 ChatService 并支持 API Key 验证

Browse Source
This commit is contained in:
yinpeng
2024-12-18 21:35:49 +08:00
parent 914922624e
commit 8dfe617468
5 changed files with 64 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
app/services/chat_service.py
View File

@@ -312,21 +312,45 @@ class ChatService:
api_key: str
) -> AsyncGenerator:
"""调用Gemini API流式生成内容"""
url = f"{self.base_url}/models/{model_name}:streamGenerateContent?alt=sse&key={api_key}"
timeout = httpx.Timeout(60.0, read=60.0) # 连接超时60秒，读取超时60秒
async with httpx.AsyncClient(timeout=timeout) as client:
retries = 0
MAX_RETRIES = 3
current_api_key = api_key
while retries < MAX_RETRIES:
try:
async with client.stream('POST', url, json=request.model_dump()) as response:
if response.status_code != 200:
error_text = response.text
logger.error(f"Error: {response.status_code}")
logger.error(error_text)
raise Exception(f"API request failed with status {response.status_code}: {error_text}")
async for line in response.aiter_lines():
print(line)
yield line + "\n\n"
url = f"{self.base_url}/models/{model_name}:streamGenerateContent?alt=sse&key={current_api_key}"
timeout = httpx.Timeout(60.0, read=60.0)
async with httpx.AsyncClient(timeout=timeout) as client:
async with client.stream('POST', url, json=request.model_dump()) as response:
if response.status_code != 200:
error_text = await response.text()
logger.error(f"Error: {response.status_code}: {error_text}")
if retries < MAX_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
continue
raise Exception(f"API request failed with status {response.status_code}: {error_text}")
async for line in response.aiter_lines():
yield line + "\n\n"
return
except httpx.ReadTimeout:
logger.warning(f"Read timeout occurred, attempting retry {retries + 1}")
if retries < MAX_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
continue
raise
except Exception as e:
logger.error(f"Streaming request failed: {str(e)}")
if retries < MAX_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
continue
raise