mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-09-08 17:36:37 +08:00
feat: 添加重试机制和消息转换器，并支持Gemini v1beta API
This commit is contained in:
yinpeng
@@ -0,0 +1,49 @@
|
||||
# app/services/chat/api_client.py
|
||||
|
||||
from typing import Dict, Any, AsyncGenerator
|
||||
import httpx
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
class ApiClient(ABC):
|
||||
"""API客户端基类"""
|
||||
|
||||
@abstractmethod
|
||||
async def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def stream_generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> AsyncGenerator[str, None]:
|
||||
pass
|
||||
|
||||
class GeminiApiClient(ApiClient):
|
||||
"""Gemini API客户端"""
|
||||
|
||||
def __init__(self, base_url: str, timeout: int = 300):
|
||||
self.base_url = base_url
|
||||
self.timeout = timeout
|
||||
|
||||
def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
|
||||
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
||||
if model.endswith("-search"):
|
||||
model = model[:-7]
|
||||
with httpx.Client(timeout=timeout) as client:
|
||||
url = f"{self.base_url}/models/{model}:generateContent?key={api_key}"
|
||||
response = client.post(url, json=payload)
|
||||
if response.status_code != 200:
|
||||
error_content = response.text
|
||||
raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
|
||||
return response.json()
|
||||
|
||||
async def stream_generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> AsyncGenerator[str, None]:
|
||||
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
||||
if model.endswith("-search"):
|
||||
model = model[:-7]
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
url = f"{self.base_url}/models/{model}:streamGenerateContent?alt=sse&key={api_key}"
|
||||
async with client.stream("POST", url, json=payload) as response:
|
||||
if response.status_code != 200:
|
||||
error_content = await response.aread()
|
||||
error_msg = error_content.decode("utf-8")
|
||||
raise Exception(f"API call failed with status code {response.status_code}, {error_msg}")
|
||||
async for line in response.aiter_lines():
|
||||
yield line
|
||||
@@ -0,0 +1,50 @@
|
||||
# app/services/chat/message_converter.py
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Dict, Any
|
||||
|
||||
class MessageConverter(ABC):
|
||||
"""消息转换器基类"""
|
||||
|
||||
@abstractmethod
|
||||
def convert(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
pass
|
||||
|
||||
class OpenAIMessageConverter(MessageConverter):
|
||||
"""OpenAI消息格式转换器"""
|
||||
|
||||
def convert(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
converted_messages = []
|
||||
for msg in messages:
|
||||
role = "user" if msg["role"] == "user" else "model"
|
||||
parts = []
|
||||
|
||||
if isinstance(msg["content"], str):
|
||||
parts.append({"text": msg["content"]})
|
||||
elif isinstance(msg["content"], list):
|
||||
for content in msg["content"]:
|
||||
if isinstance(content, str):
|
||||
parts.append({"text": content})
|
||||
elif isinstance(content, dict):
|
||||
if content["type"] == "text":
|
||||
parts.append({"text": content["text"]})
|
||||
elif content["type"] == "image_url":
|
||||
parts.append(self._convert_image(content["image_url"]["url"]))
|
||||
|
||||
converted_messages.append({"role": role, "parts": parts})
|
||||
|
||||
return converted_messages
|
||||
|
||||
def _convert_image(self, image_url: str) -> Dict[str, Any]:
|
||||
if image_url.startswith("data:image"):
|
||||
return {
|
||||
"inline_data": {
|
||||
"mime_type": "image/jpeg",
|
||||
"data": image_url.split(",")[1]
|
||||
}
|
||||
}
|
||||
return {
|
||||
"image_url": {
|
||||
"url": image_url
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,322 @@
|
||||
# app/services/chat/response_handler.py
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Any, Optional
|
||||
import time
|
||||
import uuid
|
||||
from app.core.config import settings
|
||||
|
||||
class ResponseHandler(ABC):
|
||||
"""响应处理器基类"""
|
||||
|
||||
@abstractmethod
|
||||
def handle_response(self, response: Dict[str, Any], model: str, stream: bool = False) -> Dict[str, Any]:
|
||||
pass
|
||||
|
||||
class GeminiResponseHandler(ResponseHandler):
|
||||
"""Gemini响应处理器"""
|
||||
|
||||
def __init__(self):
|
||||
self.thinking_first = True
|
||||
self.thinking_status = False
|
||||
|
||||
def handle_response(self, response: Dict[str, Any], model: str, stream: bool = False) -> Dict[str, Any]:
|
||||
if stream:
|
||||
return self._handle_stream_response(response, model, stream)
|
||||
return self._handle_normal_response(response, model, stream)
|
||||
|
||||
def _handle_stream_response(self, response: Dict[str, Any], model: str, stream: bool) -> Dict[str, Any]:
|
||||
text = self._extract_text(response, model, stream=stream)
|
||||
content = {"parts": [{"text": text}],"role": "model"}
|
||||
response["candidates"][0]["content"] = content
|
||||
return response
|
||||
|
||||
def _handle_normal_response(self, response: Dict[str, Any], model: str, stream: bool) -> Dict[str, Any]:
|
||||
text = self._extract_text(response, model, stream=stream)
|
||||
content = {"parts": [{"text": text}],"role": "model"}
|
||||
response["candidates"][0]["content"] = content
|
||||
return response
|
||||
|
||||
def _extract_text(self, response: Dict[str, Any], model: str, stream: bool = False) -> str:
|
||||
text = ""
|
||||
if stream:
|
||||
if response.get("candidates"):
|
||||
candidate = response["candidates"][0]
|
||||
content = candidate.get("content", {})
|
||||
parts = content.get("parts", [])
|
||||
if "thinking" in model:
|
||||
if settings.SHOW_THINKING_PROCESS:
|
||||
if len(parts) == 1:
|
||||
if self.thinking_first:
|
||||
self.thinking_first = False
|
||||
self.thinking_status = True
|
||||
text = "> thinking\n\n" + parts[0].get("text")
|
||||
else:
|
||||
text = parts[0].get("text")
|
||||
|
||||
if len(parts) == 2:
|
||||
self.thinking_status = False
|
||||
if self.thinking_first:
|
||||
self.thinking_first = False
|
||||
text = (
|
||||
"> thinking\n\n"
|
||||
+ parts[0].get("text")
|
||||
+ "\n\n---\n> output\n\n"
|
||||
+ parts[1].get("text")
|
||||
)
|
||||
else:
|
||||
text = (
|
||||
parts[0].get("text")
|
||||
+ "\n\n---\n> output\n\n"
|
||||
+ parts[1].get("text")
|
||||
)
|
||||
else:
|
||||
if len(parts) == 1:
|
||||
if self.thinking_first:
|
||||
self.thinking_first = False
|
||||
self.thinking_status = True
|
||||
text = ""
|
||||
elif self.thinking_status:
|
||||
text = ""
|
||||
else:
|
||||
text = parts[0].get("text")
|
||||
|
||||
if len(parts) == 2:
|
||||
self.thinking_status = False
|
||||
if self.thinking_first:
|
||||
self.thinking_first = False
|
||||
text = parts[1].get("text")
|
||||
else:
|
||||
text = parts[1].get("text")
|
||||
else:
|
||||
if "text" in parts[0]:
|
||||
text = parts[0].get("text")
|
||||
elif "executableCode" in parts[0]:
|
||||
text = _format_code_block(parts[0]["executableCode"])
|
||||
elif "codeExecution" in parts[0]:
|
||||
text = _format_code_block(parts[0]["codeExecution"])
|
||||
elif "executableCodeResult" in parts[0]:
|
||||
text = _format_execution_result(
|
||||
parts[0]["executableCodeResult"]
|
||||
)
|
||||
elif "codeExecutionResult" in parts[0]:
|
||||
text = _format_execution_result(
|
||||
parts[0]["codeExecutionResult"]
|
||||
)
|
||||
else:
|
||||
text = ""
|
||||
text = _add_search_link_text(model, candidate, text)
|
||||
else:
|
||||
if response.get("candidates"):
|
||||
candidate = response["candidates"][0]
|
||||
if "thinking" in model:
|
||||
if settings.SHOW_THINKING_PROCESS:
|
||||
if len(candidate["content"]["parts"]) == 2:
|
||||
text = (
|
||||
"> thinking\n\n"
|
||||
+ candidate["content"]["parts"][0]["text"]
|
||||
+ "\n\n---\n> output\n\n"
|
||||
+ candidate["content"]["parts"][1]["text"]
|
||||
)
|
||||
else:
|
||||
text = candidate["content"]["parts"][0]["text"]
|
||||
else:
|
||||
if len(candidate["content"]["parts"]) == 2:
|
||||
text = candidate["content"]["parts"][1]["text"]
|
||||
else:
|
||||
text = candidate["content"]["parts"][0]["text"]
|
||||
else:
|
||||
text = candidate["content"]["parts"][0]["text"]
|
||||
text = _add_search_link_text(model, candidate, text)
|
||||
else:
|
||||
text = "暂无返回"
|
||||
return text
|
||||
|
||||
|
||||
|
||||
class OpenAIResponseHandler(ResponseHandler):
|
||||
"""OpenAI响应处理器"""
|
||||
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self.thinking_first = True
|
||||
self.thinking_status = False
|
||||
|
||||
def handle_response(
|
||||
self,
|
||||
response: Dict[str, Any],
|
||||
model: str,
|
||||
stream: bool = False,
|
||||
finish_reason: str = None
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
if stream:
|
||||
return self._handle_stream_response(response, model, finish_reason)
|
||||
return self._handle_normal_response(response, model, finish_reason)
|
||||
|
||||
def _handle_stream_response(self, response: Dict[str, Any], model: str, finish_reason: str) -> Dict[str, Any]:
|
||||
text = self._extract_text(response, model, stream=True)
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {"content": text} if text else {},
|
||||
"finish_reason": finish_reason
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
def _handle_normal_response(self, response: Dict[str, Any], model: str, finish_reason: str) -> Dict[str, Any]:
|
||||
text = self._extract_text(response, model, stream=False)
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": text
|
||||
},
|
||||
"finish_reason": finish_reason
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"total_tokens": 0
|
||||
}
|
||||
}
|
||||
|
||||
def _extract_text(self, response: Dict[str, Any], model: str, stream: bool = False) -> str:
|
||||
text = ""
|
||||
if stream:
|
||||
if response.get("candidates"):
|
||||
candidate = response["candidates"][0]
|
||||
content = candidate.get("content", {})
|
||||
parts = content.get("parts", [])
|
||||
if "thinking" in model:
|
||||
if settings.SHOW_THINKING_PROCESS:
|
||||
if len(parts) == 1:
|
||||
if self.thinking_first:
|
||||
self.thinking_first = False
|
||||
self.thinking_status = True
|
||||
text = "> thinking\n\n" + parts[0].get("text")
|
||||
else:
|
||||
text = parts[0].get("text")
|
||||
|
||||
if len(parts) == 2:
|
||||
self.thinking_status = False
|
||||
if self.thinking_first:
|
||||
self.thinking_first = False
|
||||
text = (
|
||||
"> thinking\n\n"
|
||||
+ parts[0].get("text")
|
||||
+ "\n\n---\n> output\n\n"
|
||||
+ parts[1].get("text")
|
||||
)
|
||||
else:
|
||||
text = (
|
||||
parts[0].get("text")
|
||||
+ "\n\n---\n> output\n\n"
|
||||
+ parts[1].get("text")
|
||||
)
|
||||
else:
|
||||
if len(parts) == 1:
|
||||
if self.thinking_first:
|
||||
self.thinking_first = False
|
||||
self.thinking_status = True
|
||||
text = ""
|
||||
elif self.thinking_status:
|
||||
text = ""
|
||||
else:
|
||||
text = parts[0].get("text")
|
||||
|
||||
if len(parts) == 2:
|
||||
self.thinking_status = False
|
||||
if self.thinking_first:
|
||||
self.thinking_first = False
|
||||
text = parts[1].get("text")
|
||||
else:
|
||||
text = parts[1].get("text")
|
||||
else:
|
||||
if "text" in parts[0]:
|
||||
text = parts[0].get("text")
|
||||
elif "executableCode" in parts[0]:
|
||||
text = _format_code_block(parts[0]["executableCode"])
|
||||
elif "codeExecution" in parts[0]:
|
||||
text = _format_code_block(parts[0]["codeExecution"])
|
||||
elif "executableCodeResult" in parts[0]:
|
||||
text = _format_execution_result(
|
||||
parts[0]["executableCodeResult"]
|
||||
)
|
||||
elif "codeExecutionResult" in parts[0]:
|
||||
text = _format_execution_result(
|
||||
parts[0]["codeExecutionResult"]
|
||||
)
|
||||
else:
|
||||
text = ""
|
||||
text = _add_search_link_text(model, candidate, text)
|
||||
else:
|
||||
if response.get("candidates"):
|
||||
candidate = response["candidates"][0]
|
||||
if "thinking" in model:
|
||||
if settings.SHOW_THINKING_PROCESS:
|
||||
if len(candidate["content"]["parts"]) == 2:
|
||||
text = (
|
||||
"> thinking\n\n"
|
||||
+ candidate["content"]["parts"][0]["text"]
|
||||
+ "\n\n---\n> output\n\n"
|
||||
+ candidate["content"]["parts"][1]["text"]
|
||||
)
|
||||
else:
|
||||
text = candidate["content"]["parts"][0]["text"]
|
||||
else:
|
||||
if len(candidate["content"]["parts"]) == 2:
|
||||
text = candidate["content"]["parts"][1]["text"]
|
||||
else:
|
||||
text = candidate["content"]["parts"][0]["text"]
|
||||
else:
|
||||
text = candidate["content"]["parts"][0]["text"]
|
||||
text = _add_search_link_text(model, candidate, text)
|
||||
else:
|
||||
text = "暂无返回"
|
||||
return text
|
||||
|
||||
|
||||
def _format_code_block(code_data: dict) -> str:
|
||||
"""格式化代码块输出"""
|
||||
language = code_data.get("language", "").lower()
|
||||
code = code_data.get("code", "").strip()
|
||||
return f"""\n\n---\n\n【代码执行】\n```{language}\n{code}\n```\n"""
|
||||
|
||||
|
||||
def _add_search_link_text(model:str, candidate:dict, text:str) -> str:
|
||||
if (
|
||||
settings.SHOW_SEARCH_LINK
|
||||
and model.endswith("-search")
|
||||
and "groundingMetadata" in candidate
|
||||
and "groundingChunks" in candidate["groundingMetadata"]
|
||||
):
|
||||
grounding_chunks = candidate["groundingMetadata"]["groundingChunks"]
|
||||
text += "\n\n---\n\n"
|
||||
text += "**【引用来源】**\n\n"
|
||||
for _, grounding_chunk in enumerate(grounding_chunks, 1):
|
||||
if "web" in grounding_chunk:
|
||||
text += _create_search_link(grounding_chunk["web"])
|
||||
return text
|
||||
else:
|
||||
return text
|
||||
|
||||
|
||||
def _create_search_link(grounding_chunk: dict) -> str:
|
||||
return f'\n- [{grounding_chunk["title"]}]({grounding_chunk["uri"]})'
|
||||
|
||||
|
||||
def _format_execution_result(result_data: dict) -> str:
|
||||
"""格式化执行结果输出"""
|
||||
outcome = result_data.get("outcome", "")
|
||||
output = result_data.get("output", "").strip()
|
||||
return f"""\n【执行结果】\n> outcome: {outcome}\n\n【输出结果】\n```plaintext\n{output}\n```\n\n---\n\n"""
|
||||
@@ -0,0 +1,40 @@
|
||||
# app/services/chat/retry_handler.py
|
||||
|
||||
from typing import TypeVar, Callable
|
||||
from functools import wraps
|
||||
from app.core.logger import get_retry_logger
|
||||
from app.services.key_manager import KeyManager
|
||||
|
||||
T = TypeVar('T')
|
||||
logger = get_retry_logger()
|
||||
|
||||
class RetryHandler:
|
||||
"""重试处理装饰器"""
|
||||
|
||||
def __init__(self, max_retries: int = 3, key_manager: KeyManager = None, key_arg: str = "api_key"):
|
||||
self.max_retries = max_retries
|
||||
self.key_manager = key_manager
|
||||
self.key_arg = key_arg
|
||||
|
||||
def __call__(self, func: Callable[..., T]) -> Callable[..., T]:
|
||||
@wraps(func)
|
||||
async def wrapper(*args, **kwargs) -> T:
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(self.max_retries):
|
||||
try:
|
||||
return await func(*args, **kwargs)
|
||||
except Exception as e:
|
||||
last_exception = e
|
||||
logger.warning(f"API call failed with error: {str(e)}. Attempt {attempt + 1} of {self.max_retries}")
|
||||
|
||||
if self.key_manager:
|
||||
old_key = kwargs.get(self.key_arg)
|
||||
new_key = await self.key_manager.handle_api_failure(old_key)
|
||||
kwargs[self.key_arg] = new_key
|
||||
logger.info(f"Switched to new API key: {new_key}")
|
||||
|
||||
logger.error(f"All retry attempts failed, raising final exception: {str(last_exception)}")
|
||||
raise last_exception
|
||||
|
||||
return wrapper
|
||||
Reference in New Issue
Block a user