Update api/utils.py

api/utils.py

@@ -2,6 +2,7 @@ from datetime import datetime
 import json
 from typing import Any, Dict, Optional
 import uuid
+import asyncio
 
 import httpx
 from api import validate
@@ -71,6 +72,9 @@ def message_to_dict(message):
 
 
 async def process_streaming_response(request: ChatRequest):
+    max_retries = 3
+    retry_delay = 1
+    
     json_data = {
         "messages": [message_to_dict(msg) for msg in request.messages],
         "previewToken": None,
@@ -94,47 +98,82 @@ async def process_streaming_response(request: ChatRequest):
         "validated": validate.getHid()
     }
 
-    async with httpx.AsyncClient() as client:
+    for attempt in range(max_retries):
         try:
-            async with client.stream(
-                "POST",
-                f"{BASE_URL}/api/chat",
-                headers=headers,
-                json=json_data,
-                timeout=100,
-            ) as response:
-                try:
+            async with httpx.AsyncClient(timeout=120.0) as client:
+                async with client.stream(
+                    "POST",
+                    f"{BASE_URL}/api/chat",
+                    headers=headers,
+                    json=json_data,
+                ) as response:
+                    if response.status_code == 500:
+                        # 如果是500错误，尝试刷新hid
+                        validate.getHid(True)
+                        # 更新json_data中的validated值
+                        json_data["validated"] = validate.getHid()
+                        if attempt < max_retries - 1:
+                            await asyncio.sleep(retry_delay)
+                            continue
+                        
                     response.raise_for_status()
                     async for line in response.aiter_lines():
                         timestamp = int(datetime.now().timestamp())
                         if line:
-                            content = line + "\n"
-                            if "https://www.blackbox.ai" in content:
-                                validate.getHid(True)
-                                content = "hid已刷新,重新对话即可\n"
-                                yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
-                                break
-                            if content.startswith("$@$v=undefined-rv1$@$"):
-                                yield f"data: {json.dumps(create_chat_completion_data(content[21:], request.model, timestamp))}\n\n"
-                            else:
+                            try:
+                                content = line + "\n"
+                                if "https://www.blackbox.ai" in content:
+                                    validate.getHid(True)
+                                    content = "正在刷新会话，请重试\n"
+                                    yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
+                                    yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
+                                    yield "data: [DONE]\n\n"
+                                    return
+
+                                if content.startswith("$@$v=undefined-rv1$@$"):
+                                    content = content[21:]
+
                                 yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n"
+                            except Exception as e:
+                                logger.error(f"Error processing line: {e}")
+                                continue
 
                     yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
                     yield "data: [DONE]\n\n"
-                except httpx.HTTPStatusError as e:
-                    logger.error(f"HTTP error occurred: {e}")
-                    error_message = "抱歉，服务器暂时无法处理您的请求。请稍后重试。"
-                    timestamp = int(datetime.now().timestamp())
-                    yield f"data: {json.dumps(create_chat_completion_data(error_message, request.model, timestamp))}\n\n"
-                    yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
-                    yield "data: [DONE]\n\n"
+                    return
+
+        except httpx.HTTPStatusError as e:
+            logger.error(f"HTTP error occurred: {e}")
+            if attempt == max_retries - 1:
+                error_message = "服务暂时不可用，请稍后重试"
+                timestamp = int(datetime.now().timestamp())
+                yield f"data: {json.dumps(create_chat_completion_data(error_message, request.model, timestamp))}\n\n"
+                yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
+                yield "data: [DONE]\n\n"
+                return
+            await asyncio.sleep(retry_delay)
+
         except httpx.RequestError as e:
-            logger.error(f"Error occurred during request: {e}")
-            error_message = "抱歉，连接服务器时出现错误。请检查网络连接并重试。"
-            timestamp = int(datetime.now().timestamp())
-            yield f"data: {json.dumps(create_chat_completion_data(error_message, request.model, timestamp))}\n\n"
-            yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
-            yield "data: [DONE]\n\n"
+            logger.error(f"Request error occurred: {e}")
+            if attempt == max_retries - 1:
+                error_message = "网络连接错误，请检查网络后重试"
+                timestamp = int(datetime.now().timestamp())
+                yield f"data: {json.dumps(create_chat_completion_data(error_message, request.model, timestamp))}\n\n"
+                yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
+                yield "data: [DONE]\n\n"
+                return
+            await asyncio.sleep(retry_delay)
+
+        except Exception as e:
+            logger.error(f"Unexpected error: {e}")
+            if attempt == max_retries - 1:
+                error_message = "发生未知错误，请重试"
+                timestamp = int(datetime.now().timestamp())
+                yield f"data: {json.dumps(create_chat_completion_data(error_message, request.model, timestamp))}\n\n"
+                yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
+                yield "data: [DONE]\n\n"
+                return
+            await asyncio.sleep(retry_delay)
 
 
 async def process_non_streaming_response(request: ChatRequest):