Upload 2 files

Dockerfile

@@ -11,7 +11,7 @@ RUN echo 'from fastapi import FastAPI\n\
 from retry_middleware import RetryMiddleware\n\
 from main import app\n\
 \n\
-app.add_middleware(RetryMiddleware, max_retries=5, initial_delay=0.5)' > /app/wrapper.py
+app.add_middleware(RetryMiddleware)' > /app/wrapper.py
 
 # ENV APP_SECRET=
 

retry_middleware.py

@@ -1,20 +1,16 @@
 from starlette.middleware.base import BaseHTTPMiddleware
 from fastapi import Request, Response
 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 = 5, initial_delay: float = 0.5):
+    def __init__(self, app):
         super().__init__(app)
-        self.max_retries = max_retries
-        self.initial_delay = initial_delay
-        
+
     async def should_retry_response(self, response_data):
         """检查响应是否需要重试"""
         if isinstance(response_data, dict):
@@ -45,11 +41,9 @@ class RetryMiddleware(BaseHTTPMiddleware):
 
         # 读取原始请求体
         body = await request.body()
-        original_response = None
-        best_response = None
         retry_count = 0
         
-        while retry_count < self.max_retries:
+        while True:  # 无限循环，直到获得正确响应
             try:
                 # 构造请求
                 async def receive():
@@ -68,23 +62,19 @@ class RetryMiddleware(BaseHTTPMiddleware):
                 try:
                     response_data = json.loads(response_body)
                     
-                    # 第一次响应，保存作为原始响应
-                    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
+                        logger.info(f"检测到内容安全问题，立即进行第 {retry_count + 1} 次重试...")
+                        continue
                     else:
-                        # 如果响应正常，保存为最佳响应
-                        best_response = response_data
-                        break
+                        # 如果响应正常，直接返回
+                        return Response(
+                            content=response_body,
+                            status_code=response.status_code,
+                            headers=dict(response.headers),
+                            media_type=response.media_type
+                        )
                         
                 except json.JSONDecodeError:
                     # 如果响应不是JSON格式，直接返回
@@ -98,22 +88,7 @@ class RetryMiddleware(BaseHTTPMiddleware):
             except Exception as e:
                 logger.error(f"重试过程中发生错误: {str(e)}")
                 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)}"}
-                    )
-
-        # 返回最佳响应，如果没有最佳响应则返回原始响应
-        final_response = best_response or original_response
-        return JSONResponse(
-            content=final_response,
-            status_code=200
-        )
+                continue
 
 # 使用方法：
 """