Upload retry_middleware.py

retry_middleware.py

@@ -4,15 +4,39 @@ from fastapi.responses import JSONResponse
 import asyncio
 import json
 import logging
+import random
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class RetryMiddleware(BaseHTTPMiddleware):
-    def __init__(self, app, max_retries: int = 3, delay: float = 1.0):
+    def __init__(self, app, max_retries: int = 5, initial_delay: float = 0.5):
         super().__init__(app)
         self.max_retries = max_retries
-        self.delay = delay
+        self.initial_delay = initial_delay
+        
+    async def should_retry_response(self, response_data):
+        """检查响应是否需要重试"""
+        if isinstance(response_data, dict):
+            # 检查错误信息
+            error = response_data.get('error', '')
+            if isinstance(error, str) and 'content is not safe' in error.lower():
+                return True
+                
+            # 检查响应内容
+            choices = response_data.get('choices', [])
+            if choices:
+                content = choices[0].get('message', {}).get('content', '')
+                if 'content is not safe' in content.lower():
+                    return True
+                    
+                # 检查是否有部分内容包含错误信息
+                if isinstance(content, str):
+                    lines = content.split('\n')
+                    for line in lines:
+                        if 'content is not safe' in line.lower():
+                            return True
+        return False
 
     async def dispatch(self, request: Request, call_next):
         # 只处理 /api/v1/chat/completions 路径的请求
@@ -21,53 +45,75 @@ class RetryMiddleware(BaseHTTPMiddleware):
 
         # 读取原始请求体
         body = await request.body()
+        original_response = None
+        best_response = None
+        retry_count = 0
         
-        for attempt in range(self.max_retries):
+        while retry_count < self.max_retries:
             try:
-                # 构造新的请求
+                # 构造请求
                 async def receive():
                     return {
                         "type": "http.request",
                         "body": body,
                         "more_body": False,
                     }
-                
-                # ���送请求并获取响应
+
+                # 发送请求并获取响应
                 response = await call_next(Request(request.scope, receive))
-                
-                # 读取响应内容
                 response_body = b""
                 async for chunk in response.body_iterator:
                     response_body += chunk
-                
+
                 try:
                     response_data = json.loads(response_body)
-                    if isinstance(response_data, dict):
-                        error = response_data.get('error', '')
-                        if isinstance(error, str) and 'content is not safe' in error.lower():
-                            if attempt < self.max_retries - 1:
-                                logger.info(f"检测到内容安全问题，等待 {self.delay} 秒后进行第 {attempt + 2} 次重试...")
-                                await asyncio.sleep(self.delay)
-                                continue
+                    
+                    # 第一次响应，保存作为原始响应
+                    if original_response is None:
+                        original_response = response_data
+                    
+                    # 检查响应是否需要重试
+                    if await self.should_retry_response(response_data):
+                        retry_count += 1
+                        if retry_count < self.max_retries:
+                            # 使用指数退避和随机抖动
+                            delay = self.initial_delay * (2 ** retry_count) * (0.5 + random.random())
+                            logger.info(f"检测到内容安全问题，等待 {delay:.2f} 秒后进行第 {retry_count + 1} 次重试...")
+                            await asyncio.sleep(delay)
+                            continue
+                    else:
+                        # 如果响应正常，保存为最佳响应
+                        best_response = response_data
+                        break
+                        
                 except json.JSONDecodeError:
-                    pass
-                
-                # 如果没有错误或是最后一次尝试，返回响应
-                return Response(
-                    content=response_body,
-                    status_code=response.status_code,
-                    headers=dict(response.headers),
-                    media_type=response.media_type
-                )
-                
+                    # 如果响应不是JSON格式，直接返回
+                    return Response(
+                        content=response_body,
+                        status_code=response.status_code,
+                        headers=dict(response.headers),
+                        media_type=response.media_type
+                    )
+                    
             except Exception as e:
                 logger.error(f"重试过程中发生错误: {str(e)}")
-                if attempt == self.max_retries - 1:
+                retry_count += 1
+                if retry_count < self.max_retries:
+                    delay = self.initial_delay * (2 ** retry_count) * (0.5 + random.random())
+                    await asyncio.sleep(delay)
+                    continue
+                else:
                     return JSONResponse(
                         status_code=500,
                         content={"error": f"在 {self.max_retries} 次尝试后仍然失败: {str(e)}"}
                     )
-                await asyncio.sleep(self.delay)
+
+        # 返回最佳响应，如果没有最佳响应则返回原始响应
+        final_response = best_response or original_response
+        return JSONResponse(
+            content=final_response,
+            status_code=200
+        )
 
 # 使用方法：
 """