fix: 修复 OpenAI 和 Gemini API 调用重试逻辑及日志记录

.gitignore

@@ -36,6 +36,7 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+.idea/
 
 # PyInstaller
 #  Usually these files are written by a python script from a template

app/api/gemini_routes.py

@@ -1,14 +1,13 @@
-from email.header import Header
 from fastapi import APIRouter, Depends
 from fastapi.responses import StreamingResponse
 
+from app.core.config import settings
+from app.core.logger import get_gemini_logger
 from app.core.security import SecurityService
+from app.schemas.gemini_models import GeminiRequest
 from app.services.chat_service import ChatService
 from app.services.key_manager import KeyManager
 from app.services.model_service import ModelService
-from app.schemas.gemini_models import GeminiRequest
-from app.core.config import settings
-from app.core.logger import get_gemini_logger
 
 router = APIRouter(prefix="/gemini/v1beta")
 logger = get_gemini_logger()
@@ -19,6 +18,7 @@ key_manager = KeyManager(settings.API_KEYS)
 model_service = ModelService(settings.MODEL_SEARCH)
 chat_service = ChatService(base_url=settings.BASE_URL, key_manager=key_manager)
 
+
 @router.get("/models")
 async def list_models(
         key: str = None,
@@ -31,6 +31,7 @@ async def list_models(
     logger.info(f"Using API key: {api_key}")
     return model_service.get_gemini_models(api_key)
 
+
 @router.post("/models/{model_name}:generateContent")
 async def generate_content(
         model_name: str,
@@ -66,6 +67,7 @@ async def generate_content(
             if retries >= MAX_RETRIES:
                 logger.error(f"Max retries ({MAX_RETRIES}) reached. Raising error")
 
+
 @router.post("/models/{model_name}:streamGenerateContent")
 async def stream_generate_content(
         model_name: str,

app/api/openai_routes.py

@@ -48,9 +48,9 @@ async def chat_completion(
     api_key = await key_manager.get_next_working_key()
     logger.info(f"Using API key: {api_key}")
     retries = 0
-    MAX_RETRIES = 3
+    max_retries = 3
 
-    while retries < MAX_RETRIES:
+    while retries < max_retries:
         try:
             response = await chat_service.create_chat_completion(
                 request=request,
@@ -64,13 +64,13 @@ async def chat_completion(
 
         except Exception as e:
             logger.warning(
-                f"API call failed with error: {str(e)}. Attempt {retries + 1} of {MAX_RETRIES}"
+                f"API call failed with error: {str(e)}. Attempt {retries + 1} of {max_retries}"
             )
             api_key = await key_manager.handle_api_failure(api_key)
             logger.info(f"Switched to new API key: {api_key}")
             retries += 1
-            if retries >= MAX_RETRIES:
-                logger.error(f"Max retries ({MAX_RETRIES}) reached. Raising error")
+            if retries >= max_retries:
+                logger.error(f"Max retries ({max_retries}) reached. Raising error")
                 raise
 
 

app/core/config.py

@@ -9,6 +9,7 @@ class Settings(BaseSettings):
     MODEL_SEARCH: List[str] = ["gemini-2.0-flash-exp"]
     TOOLS_CODE_EXECUTION_ENABLED: bool = False
     SHOW_SEARCH_LINK: bool = True
+
     class Config:
         env_file = ".env"
 

app/core/logger.py

@@ -15,6 +15,7 @@ COLORS = {
 # Windows系统启用ANSI支持
 if platform.system() == 'Windows':
     import ctypes
+
     kernel32 = ctypes.windll.kernel32
     kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), 7)
 
@@ -23,6 +24,7 @@ class ColoredFormatter(logging.Formatter):
     """
     自定义的日志格式化器,添加颜色支持
     """
+
     def format(self, record):
         # 获取对应级别的颜色代码
         color = COLORS.get(record.levelname, '')
@@ -30,6 +32,7 @@ class ColoredFormatter(logging.Formatter):
         record.levelname = f"{color}{record.levelname}\033[0m"
         return super().format(record)
 
+
 # 日志格式
 FORMATTER = ColoredFormatter(
     "%(asctime)s - %(name)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s"
@@ -44,7 +47,11 @@ LOG_LEVELS = {
     "critical": logging.CRITICAL,
 }
 
+
 class Logger:
+    def __init__(self):
+        pass
+
     _loggers: Dict[str, logging.Logger] = {}
 
     @staticmethod
@@ -82,30 +89,39 @@ class Logger:
         """
         return Logger._loggers.get(name)
 
+
 # 预定义的loggers
 def get_openai_logger():
     return Logger.setup_logger("openai")
 
+
 def get_gemini_logger():
     return Logger.setup_logger("gemini")
 
+
 def get_chat_logger():
     return Logger.setup_logger("chat")
 
+
 def get_model_logger():
     return Logger.setup_logger("model")
 
+
 def get_security_logger():
     return Logger.setup_logger("security")
 
+
 def get_key_manager_logger():
     return Logger.setup_logger("key_manager")
 
+
 def get_main_logger():
     return Logger.setup_logger("main")
 
+
 def get_embeddings_logger():
     return Logger.setup_logger("embeddings")
 
+
 def get_request_logger():
     return Logger.setup_logger("request")

app/middleware/request_logging_middleware.py

@@ -3,9 +3,9 @@ from starlette.middleware.base import BaseHTTPMiddleware
 import json
 from app.core.logger import get_request_logger
 
-
 logger = get_request_logger()
 
+
 # 添加中间件类
 class RequestLoggingMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):

app/schemas/gemini_models.py

@@ -1,9 +1,12 @@
 from typing import List, Optional, Dict, Any, Literal
 from pydantic import BaseModel
 
+
 class SafetySetting(BaseModel):
-    category: Optional[Literal["HARM_CATEGORY_HATE_SPEECH", "HARM_CATEGORY_DANGEROUS_CONTENT", "HARM_CATEGORY_HARASSMENT", "HARM_CATEGORY_SEXUALLY_EXPLICIT"]] = None
-    threshold: Optional[Literal["HARM_BLOCK_THRESHOLD_UNSPECIFIED", "BLOCK_LOW_AND_ABOVE", "BLOCK_MEDIUM_AND_ABOVE","BLOCK_ONLY_HIGH","BLOCK_NONE","OFF"]] = None
+    category: Optional[Literal[
+        "HARM_CATEGORY_HATE_SPEECH", "HARM_CATEGORY_DANGEROUS_CONTENT", "HARM_CATEGORY_HARASSMENT", "HARM_CATEGORY_SEXUALLY_EXPLICIT"]] = None
+    threshold: Optional[Literal[
+        "HARM_BLOCK_THRESHOLD_UNSPECIFIED", "BLOCK_LOW_AND_ABOVE", "BLOCK_MEDIUM_AND_ABOVE", "BLOCK_ONLY_HIGH", "BLOCK_NONE", "OFF"]] = None
 
 
 class GenerationConfig(BaseModel):

app/services/chat_service.py

@@ -11,12 +11,7 @@ from app.schemas.openai_models import ChatRequest
 logger = get_chat_logger()
 
 
-class ChatService:
-    def __init__(self, base_url: str, key_manager=None):
-        self.base_url = base_url
-        self.key_manager = key_manager
-
-    def convert_messages_to_gemini_format(self, messages: list) -> list:
+def convert_messages_to_gemini_format(messages: list) -> list:
     """Convert OpenAI message format to Gemini format"""
     converted_messages = []
     for msg in messages:
@@ -60,6 +55,23 @@ class ChatService:
 
     return converted_messages
 
+
+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"""
+
+
+def create_search_link(web):
+    return f'\n- [{web["title"]}]({web["uri"]})'
+
+
+class ChatService:
+    def __init__(self, base_url: str, key_manager=None):
+        self.base_url = base_url
+        self.key_manager = key_manager
+
     def convert_gemini_response_to_openai(
             self,
             response: Dict[str, Any],
@@ -82,28 +94,17 @@ class ChatService:
                     elif "codeExecution" in parts[0]:
                         text = self.format_code_block(parts[0]["codeExecution"])
                     elif "executableCodeResult" in parts[0]:
-                        text = self.format_execution_result(
+                        text = format_execution_result(
                             parts[0]["executableCodeResult"]
                         )
                     elif "codeExecutionResult" in parts[0]:
-                        text = self.format_execution_result(
+                        text = format_execution_result(
                             parts[0]["codeExecutionResult"]
                         )
                     else:
                         text = ""
 
-                    if (
-                        settings.SHOW_SEARCH_LINK
-                        and model.endswith("-search")
-                        and "groundingMetadata" in candidate
-                        and "groundingChunks" in candidate["groundingMetadata"]
-                    ):
-                        groundingChunks = candidate["groundingMetadata"]["groundingChunks"]
-                        text += "\n\n---\n\n"
-                        text += f"**【引用来源】**\n\n"
-                        for _, groundingChunk in enumerate(groundingChunks, 1):
-                            if "web" in groundingChunk:
-                                text += self.create_search_link(groundingChunk["web"])
+                    text = self.add_search_link_text(model, candidate, text)
                 else:
                     text = ""
 
@@ -150,18 +151,7 @@ class ChatService:
                 if response.get("candidates"):
                     text = response["candidates"][0]["content"]["parts"][0]["text"]
                     candidate = response["candidates"][0]
-                    if (
-                        settings.SHOW_SEARCH_LINK
-                        and model.endswith("-search")
-                        and "groundingMetadata" in candidate
-                        and "groundingChunks" in candidate["groundingMetadata"]
-                    ):
-                        groundingChunks = candidate["groundingMetadata"]["groundingChunks"]
-                        text += "\n\n---\n\n"
-                        text += f"**【引用来源】**\n\n"
-                        for _, groundingChunk in enumerate(groundingChunks, 1):
-                            if "web" in groundingChunk:
-                                text += self.create_search_link(groundingChunk["web"])
+                    text = self.add_search_link_text(model, candidate, text)
                     res["choices"][0]["message"]["content"] = text
                     return res
                 else:
@@ -173,6 +163,23 @@ class ChatService:
                 res["choices"][0]["message"]["content"] = f"Error converting Gemini response: {str(e)}"
                 return res
 
+    def add_search_link_text(self, model, candidate, text):
+        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 += f"**【引用来源】**\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
+        
     async def create_chat_completion(
             self,
             request: ChatRequest,
@@ -210,7 +217,7 @@ class ChatService:
             gemini_model = model[:-7]  # Remove -search suffix
         else:
             gemini_model = model
-        gemini_messages = self.convert_messages_to_gemini_format(messages)
+        gemini_messages = convert_messages_to_gemini_format(messages)
 
         if not stream:
             # 非流式模式下，移除代码执行工具
@@ -247,26 +254,26 @@ class ChatService:
         if stream:
             async def generate():
                 retries = 0
-                MAX_RETRIES = 3
+                max_retries = 3
                 current_api_key = api_key
 
-                while retries < MAX_RETRIES:
+                while retries < max_retries:
                     try:
                         timeout = httpx.Timeout(
                             60.0, read=60.0
                         )  # 连接超时60秒，读取超时60秒
-                        async with httpx.AsyncClient(timeout=timeout) as client:
+                        async with httpx.AsyncClient(timeout=timeout) as async_client:
                             stream_url = f"https://generativelanguage.googleapis.com/v1beta/models/{gemini_model}:streamGenerateContent?alt=sse&key={current_api_key}"
-                            async with client.stream(
+                            async with async_client.stream(
                                     "POST", stream_url, json=payload
-                            ) as response:
-                                if response.status_code != 200:
-                                    error_content = await response.read()
+                            ) as async_response:
+                                if async_response.status_code != 200:
+                                    error_content = await async_response.read()
                                     error_msg = error_content.decode("utf-8")
                                     logger.error(
-                                        f"API error: {response.status_code}, {error_msg}"
+                                        f"API error: {async_response.status_code}, {error_msg}"
                                     )
-                                    if retries < MAX_RETRIES - 1:
+                                    if retries < max_retries - 1:
                                         current_api_key = (
                                             await self.key_manager.handle_api_failure(
                                                 current_api_key
@@ -276,12 +283,12 @@ class ChatService:
                                         continue
                                     else:
                                         logger.error(
-                                            f"Max retries reached. Final error: {response.status_code}, {error_msg}"
+                                            f"Max retries reached. Final error: {async_response.status_code}, {error_msg}"
                                         )
-                                        yield f"data: {json.dumps({'error': f'API error: {response.status_code}, {error_msg}'})}\n\n"
+                                        yield f"data: {json.dumps({'error': f'API error: {async_response.status_code}, {error_msg}'})}\n\n"
                                         return
 
-                                async for line in response.aiter_lines():
+                                async for line in async_response.aiter_lines():
                                     if line.startswith("data: "):
                                         try:
                                             chunk = json.loads(line[6:])
@@ -305,7 +312,7 @@ class ChatService:
                         logger.warning(
                             f"Read timeout occurred, attempting retry {retries + 1}"
                         )
-                        if retries < MAX_RETRIES - 1:
+                        if retries < max_retries - 1:
                             current_api_key = await self.key_manager.handle_api_failure(
                                 current_api_key
                             )
@@ -322,7 +329,7 @@ class ChatService:
                         logger.exception(
                             f"Stream error: {str(e)}, attempting retry {retries + 1}"
                         )
-                        if retries < MAX_RETRIES - 1:
+                        if retries < max_retries - 1:
                             current_api_key = await self.key_manager.handle_api_failure(
                                 current_api_key
                             )
@@ -362,12 +369,6 @@ class ChatService:
 
         return f"""\n【代码执行】\n```{language}\n{code}\n```\n"""
 
-    def format_execution_result(self, 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"""
-
     async def generate_content(
             self, model_name: str, request: GeminiRequest, api_key: str
     ) -> dict:
@@ -451,6 +452,3 @@ class ChatService:
                     retries += 1
                     continue
                 raise
-
-    def create_search_link(self, web):
-        return f'\n- [{web["title"]}]({web["uri"]})'

app/services/embedding_service.py

@@ -1,5 +1,8 @@
+from typing import Union, List
+
 import openai
-from typing import Union, List, Dict, Any
+from openai.types import CreateEmbeddingResponse
+
 from app.core.logger import get_embeddings_logger
 
 logger = get_embeddings_logger()
@@ -11,7 +14,7 @@ class EmbeddingService:
 
     async def create_embedding(
         self, input_text: Union[str, List[str]], model: str, api_key: str
-    ) -> Dict[str, Any]:
+    ) -> CreateEmbeddingResponse:
         """Create embeddings using OpenAI API"""
         try:
             client = openai.OpenAI(api_key=api_key, base_url=self.base_url)