Upload 2 files

Dockerfile

@@ -1,5 +1,8 @@
 FROM hpyp/bbapi:latest
 
+# 安装必要的依赖
+RUN pip install starlette
+
 # 复制重试中间件文件到容器中
 COPY retry_middleware.py /app/retry_middleware.py
 

retry_middleware.py

@@ -1,112 +1,73 @@
+from starlette.middleware.base import BaseHTTPMiddleware
 from fastapi import Request, Response
 from fastapi.responses import JSONResponse
 import asyncio
 import json
-from typing import Callable
 import logging
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-class RetryMiddleware:
-    def __init__(
-        self,
-        app,
-        max_retries: int = 3,
-        delay: float = 1.0
-    ):
-        self.app = app
+class RetryMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app, max_retries: int = 3, delay: float = 1.0):
+        super().__init__(app)
         self.max_retries = max_retries
         self.delay = delay
 
-    async def __call__(self, scope, receive, send):
-        if scope["type"] != "http":
-            return await self.app(scope, receive, send)
+    async def dispatch(self, request: Request, call_next):
+        # 只处理 /api/v1/chat/completions 路径的请求
+        if not request.url.path.endswith('/api/v1/chat/completions'):
+            return await call_next(request)
 
-        async def wrapped_send(message):
-            if message["type"] == "http.response.body":
-                body = message.get("body", b"")
-                if body:
-                    try:
-                        response_data = json.loads(body.decode())
-                        if isinstance(response_data, dict) and response_data.get("error"):
-                            error_msg = response_data["error"].lower()
-                            if "content is not safe" in error_msg:
-                                # 需要���试
-                                return await self.handle_retry(scope, receive, send)
-                    except json.JSONDecodeError:
-                        pass
-            await send(message)
-
-        await self.app(scope, receive, wrapped_send)
-
-    async def handle_retry(self, scope, receive, send):
-        original_receive = receive
+        # 读取原始请求体
+        body = await request.body()
+        
         for attempt in range(self.max_retries):
             try:
-                logger.info(f"正在进行第 {attempt + 1} 次尝试...")
+                # 构造新的请求
+                async def receive():
+                    return {
+                        "type": "http.request",
+                        "body": body,
+                        "more_body": False,
+                    }
                 
-                # 重新构造请求体
-                request_body = b""
-                more_body = True
-                while more_body:
-                    message = await original_receive()
-                    if message["type"] == "http.request":
-                        request_body += message.get("body", b"")
-                        more_body = message.get("more_body", False)
-
-                async def modified_receive():
-                    if not hasattr(modified_receive, 'called'):
-                        modified_receive.called = True
-                        return {
-                            "type": "http.request",
-                            "body": request_body,
-                            "more_body": False
-                        }
-                    return {"type": "http.disconnect"}
-
-                response_sent = False
-                async def modified_send(message):
-                    nonlocal response_sent
-                    if message["type"] == "http.response.start":
-                        await send(message)
-                    elif message["type"] == "http.response.body":
-                        body = message.get("body", b"")
-                        if body:
-                            try:
-                                response_data = json.loads(body.decode())
-                                if isinstance(response_data, dict) and response_data.get("error"):
-                                    error_msg = response_data["error"].lower()
-                                    if "content is not safe" in error_msg:
-                                        if attempt < self.max_retries - 1:
-                                            logger.info(f"检测到内容安全问题，等待 {self.delay} 秒后重试...")
-                                            await asyncio.sleep(self.delay)
-                                            return
-                            except json.JSONDecodeError:
-                                pass
-                        response_sent = True
-                        await send(message)
-
-                await self.app(scope, modified_receive, modified_send)
-                if response_sent:
-                    break
+                # ���送请求并获取响应
+                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
+                except json.JSONDecodeError:
+                    pass
+                
+                # 如果没有错误或是最后一次尝试，返回响应
+                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:
-                    # 如果是最后一次尝试，发送错误响应
-                    error_response = {
-                        "error": f"在 {self.max_retries} 次尝试后仍然失败: {str(e)}"
-                    }
-                    await send({
-                        "type": "http.response.start",
-                        "status": 500,
-                        "headers": [(b"content-type", b"application/json")]
-                    })
-                    await send({
-                        "type": "http.response.body",
-                        "body": json.dumps(error_response).encode()
-                    })
+                    return JSONResponse(
+                        status_code=500,
+                        content={"error": f"在 {self.max_retries} 次尝试后仍然失败: {str(e)}"}
+                    )
+                await asyncio.sleep(self.delay)
 
 # 使用方法：
 """