Upload 2 files

Dockerfile

@@ -3,9 +3,17 @@ FROM hpyp/bbapi:latest
 # 复制重试中间件文件到容器中
 COPY retry_middleware.py /app/retry_middleware.py
 
+# 创建一个新的入口点脚本
+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=3, delay=1.0)' > /app/wrapper.py
+
 # ENV APP_SECRET=
 
 EXPOSE 8001
 
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8001"]
+# 使用新的入口点脚本
+CMD ["uvicorn", "wrapper:app", "--host", "0.0.0.0", "--port", "8001"]
 

retry_middleware.py

@@ -1,48 +1,119 @@
-import time
-from functools import wraps
-
-def retry_on_unsafe_content(max_retries=3, delay=1):
-    """
-    装饰器函数，用于在收到 'content is not safe' 响应时自动重试
-    
-    Args:
-        max_retries (int): 最大重试次数
-        delay (int): 重试之间的延迟时间（秒）
-    """
-    def decorator(func):
-        @wraps(func)
-        async def wrapper(*args, **kwargs):
-            retries = 0
-            while retries < max_retries:
-                try:
-                    response = await func(*args, **kwargs)
-                    
-                    # 检查响应中是否包含 "content is not safe" 错误
-                    if isinstance(response, dict) and response.get('error') and 'content is not safe' in response['error'].lower():
-                        retries += 1
-                        if retries < max_retries:
-                            print(f"检测到内容安全问题，正在进行第 {retries} 次重试...")
-                            time.sleep(delay)
-                            continue
-                    return response
-                except Exception as e:
-                    if 'content is not safe' in str(e).lower():
-                        retries += 1
-                        if retries < max_retries:
-                            print(f"检测到内容安全问题，正在进行第 {retries} 次重试...")
-                            time.sleep(delay)
-                            continue
-                    raise e
-            
-            # 如果所有重试都失败，返回最后一次的响应
-            return response
-        return wrapper
-    return decorator
-
-# 使用示例:
+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
+        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 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
+        for attempt in range(self.max_retries):
+            try:
+                logger.info(f"正在进行第 {attempt + 1} 次尝试...")
+                
+                # 重新构造请求体
+                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
+                
+            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()
+                    })
+
+# 使用方法：
 """
-@retry_on_unsafe_content(max_retries=3, delay=1)
-async def your_ai_chat_function(message):
-    # 原有的AI聊天处理逻辑
-    pass
+在 main.py 中添加:
+from fastapi import FastAPI
+from retry_middleware import RetryMiddleware
+
+app = FastAPI()
+app.add_middleware(RetryMiddleware, max_retries=3, delay=1.0)
 """