Update app.py

app.py

@@ -60,6 +60,7 @@ MODEL_MAPPING = {
 # Token负载均衡状态
 token_status = {}
 
+
 # 请求头
 def get_headers():
     global TOKEN_INDEX
@@ -67,7 +68,7 @@ def get_headers():
     if len(DEEPSIDER_TOKEN) > 0:
         current_token = DEEPSIDER_TOKEN[TOKEN_INDEX % len(DEEPSIDER_TOKEN)]
         TOKEN_INDEX = (TOKEN_INDEX + 1) % len(DEEPSIDER_TOKEN)
-        
+
         # 检查token状态
         if current_token in token_status and not token_status[current_token]["active"]:
             # 如果token不可用，尝试下一个
@@ -79,7 +80,7 @@ def get_headers():
                     break
     else:
         current_token = ""
-    
+
     return {
         "accept": "*/*",
         "accept-encoding": "gzip, deflate, br, zstd",
@@ -98,12 +99,14 @@ def get_headers():
         "authorization": f"Bearer {current_token}"
     }
 
+
 # OpenAI API请求模型
 class ChatMessage(BaseModel):
     role: str
     content: str
     name: Optional[str] = None
 
+
 class ChatCompletionRequest(BaseModel):
     model: str
     messages: List[ChatMessage]
@@ -116,57 +119,58 @@ class ChatCompletionRequest(BaseModel):
     presence_penalty: Optional[float] = 0
     frequency_penalty: Optional[float] = 0
     user: Optional[str] = None
-    
+
+
 # 初始化token状态
 async def initialize_token_status():
     """初始化检查所有token的状态和余额"""
     global token_status
-    
+
     for token in DEEPSIDER_TOKEN:
         headers = {
             "accept": "*/*",
             "content-type": "application/json",
             "authorization": f"Bearer {token}"
         }
-        
+
         try:
             # 获取账户余额信息
             response = requests.get(
                 f"{DEEPSIDER_API_BASE.replace('/v2', '')}/quota/retrieve",
                 headers=headers
             )
-            
+
             active = False
             quota_info = {}
-            
+
             if response.status_code == 200:
                 data = response.json()
                 if data.get('code') == 0:
                     quota_list = data.get('data', {}).get('list', [])
-                    
+
                     # 解析余额信息
                     for item in quota_list:
                         item_type = item.get('type', '')
                         available = item.get('available', 0)
-                        
+
                         if available > 0:
                             active = True
-                        
+
                         quota_info[item_type] = {
                             "total": item.get('total', 0),
                             "available": available,
                             "title": item.get('title', '')
                         }
-            
+
             token_status[token] = {
                 "active": active,
                 "quota": quota_info,
                 "last_checked": datetime.now(),
                 "failed_count": 0
             }
-            
+
             logger.info(f"Token {token[:8]}... 状态：{'活跃' if active else '无效'}")
-            
+
         except Exception as e:
             logger.warning(f"检查Token {token[:8]}... 出错：{str(e)}")
             token_status[token] = {
@@ -176,6 +180,7 @@ async def initialize_token_status():
                 "failed_count": 0
             }
 
+
 # 工具函数
 def verify_api_key(api_key: str = Header(..., alias="Authorization")):
     """验证API密钥"""
@@ -183,10 +188,12 @@ def verify_api_key(api_key: str = Header(..., alias="Authorization")):
         raise HTTPException(status_code=401, detail="Invalid API key format")
     return api_key.replace("Bearer ", "")
 
+
 def map_openai_to_deepsider_model(model: str) -> str:
     """将OpenAI模型名称映射到DeepSider模型名称"""
     return MODEL_MAPPING.get(model, "anthropic/claude-3.7-sonnet")
 
+
 def format_messages_for_deepsider(messages: List[ChatMessage]) -> str:
     """格式化消息列表为DeepSider API所需的提示格式"""
     prompt = ""
@@ -203,17 +210,18 @@ def format_messages_for_deepsider(messages: List[ChatMessage]) -> str:
         else:
             # 其他角色按用户处理
             prompt += f"Human ({role}): {msg.content}\n\n"
-    
+
     # 如果最后一个消息不是用户的 添加一个Human前缀引导模型回答
     if messages and messages[-1].role != "user":
         prompt += "Human: "
-    
+
     return prompt.strip()
 
+
 def update_token_status(token: str, success: bool, error_message: str = None):
     """更新token的状态"""
     global token_status
-    
+
     if token not in token_status:
         token_status[token] = {
             "active": True,
@@ -221,15 +229,15 @@ def update_token_status(token: str, success: bool, error_message: str = None):
             "last_checked": datetime.now(),
             "failed_count": 0
         }
-    
+
     if not success:
         token_status[token]["failed_count"] += 1
-        
+
         # 如果失败消息包含余额不足，标记为不活跃
         if error_message and ("配额不足" in error_message or "quota" in error_message.lower()):
             token_status[token]["active"] = False
             logger.warning(f"Token {token[:8]}... 余额不足，已标记为不活跃")
-            
+
         # 连续失败5次，也标记为不活跃
         if token_status[token]["failed_count"] >= 5:
             token_status[token]["active"] = False
@@ -238,6 +246,7 @@ def update_token_status(token: str, success: bool, error_message: str = None):
         # 成功则重置失败计数
         token_status[token]["failed_count"] = 0
 
+
 async def generate_openai_response(full_response: str, request_id: str, model: str) -> Dict:
     """生成符合OpenAI API响应格式的完整响应"""
     timestamp = int(time.time())
@@ -263,27 +272,30 @@ async def generate_openai_response(full_response: str, request_id: str, model: s
         }
     }
 
+
 async def stream_openai_response(response, request_id: str, model: str, token: str):
     """流式返回OpenAI API格式的响应"""
     timestamp = int(time.time())
     full_response = ""
-    
+
     try:
         # 将DeepSider响应流转换为OpenAI流格式
         for line in response.iter_lines():
             if not line:
                 continue
-                
+
             if line.startswith(b'data: '):
                 try:
                     data = json.loads(line[6:].decode('utf-8'))
-                    
+                    if data.get('code') == 1005:
+                        raise Exception(data.get("message"))
+
                     if data.get('code') == 202 and data.get('data', {}).get('type') == "chat":
                         # 获取正文内容
                         content = data.get('data', {}).get('content', '')
                         if content:
                             full_response += content
-                            
+
                             # 生成OpenAI格式的流式响应
                             chunk = {
                                 "id": f"chatcmpl-{request_id}",
@@ -301,7 +313,7 @@ async def stream_openai_response(response, request_id: str, model: str, token: s
                                 ]
                             }
                             yield f"data: {json.dumps(chunk)}\n\n"
-                    
+
                     elif data.get('code') == 203:
                         # 生成完成信号
                         chunk = {
@@ -319,18 +331,18 @@ async def stream_openai_response(response, request_id: str, model: str, token: s
                         }
                         yield f"data: {json.dumps(chunk)}\n\n"
                         yield "data: [DONE]\n\n"
-                        
+
                 except json.JSONDecodeError:
                     logger.warning(f"无法解析响应: {line}")
-                    
+
         # 更新token状态（成功）
         update_token_status(token, True)
-                    
+
     except Exception as e:
         logger.error(f"流式响应处理出错: {str(e)}")
         # 更新token状态（失败）
         update_token_status(token, False, str(e))
-        
+
         # 返回错误信息
         error_chunk = {
             "id": f"chatcmpl-{request_id}",
@@ -350,11 +362,13 @@ async def stream_openai_response(response, request_id: str, model: str, token: s
         yield f"data: {json.dumps(error_chunk)}\n\n"
         yield "data: [DONE]\n\n"
 
+
 # 路由定义
 @app.get("/")
 async def root():
     return {"message": "OpenAI API Proxy服务已启动 连接至DeepSider API"}
 
+
 @app.get("/v1/models")
 async def list_models(api_key: str = Depends(verify_api_key)):
     """列出可用的模型"""
@@ -366,31 +380,32 @@ async def list_models(api_key: str = Depends(verify_api_key)):
             "created": int(time.time()),
             "owned_by": "openai-proxy"
         })
-    
+
     return {
         "object": "list",
         "data": models
     }
 
+
 @app.post("/v1/chat/completions")
 async def create_chat_completion(
-    request: Request,
-    api_key: str = Depends(verify_api_key)
+        request: Request,
+        api_key: str = Depends(verify_api_key)
 ):
     """创建聊天完成API - 支持普通请求和流式请求"""
     # 解析请求体
     body = await request.json()
     chat_request = ChatCompletionRequest(**body)
-    
+
     # 生成唯一请求ID
     request_id = datetime.now().strftime("%Y%m%d%H%M%S") + str(time.time_ns())[-6:]
-    
+
     # 映射模型
     deepsider_model = map_openai_to_deepsider_model(chat_request.model)
-    
+
     # 准备DeepSider API所需的提示
     prompt = format_messages_for_deepsider(chat_request.messages)
-    
+
     # 准备请求体
     payload = {
         "model": deepsider_model,
@@ -398,11 +413,11 @@ async def create_chat_completion(
         "webAccess": "close",  # 默认关闭网络访问
         "timezone": "Asia/Shanghai"
     }
-    
+
     # 获取当前token
     headers = get_headers()
     current_token = headers["authorization"].replace("Bearer ", "")
-    
+
     try:
         # 发送请求到DeepSider API
         response = requests.post(
@@ -411,7 +426,7 @@ async def create_chat_completion(
             json=payload,
             stream=True
         )
-        
+
         # 检查响应状态
         if response.status_code != 200:
             error_msg = f"DeepSider API请求失败: {response.status_code}"
@@ -420,14 +435,14 @@ async def create_chat_completion(
                 error_msg += f" - {error_data.get('message', '')}"
             except:
                 error_msg += f" - {response.text}"
-                
+
             logger.error(error_msg)
-            
+
             # 更新token状态
             update_token_status(current_token, False, error_msg)
-            
+
             raise HTTPException(status_code=response.status_code, detail="API请求失败")
-        
+
         # 处理流式或非流式响应
         if chat_request.stream:
             # 返回流式响应
@@ -441,25 +456,25 @@ async def create_chat_completion(
             for line in response.iter_lines():
                 if not line:
                     continue
-                    
+
                 if line.startswith(b'data: '):
                     try:
                         data = json.loads(line[6:].decode('utf-8'))
-                        
+
                         if data.get('code') == 202 and data.get('data', {}).get('type') == "chat":
                             content = data.get('data', {}).get('content', '')
                             if content:
                                 full_response += content
-                                
+
                     except json.JSONDecodeError:
                         pass
-            
+
             # 更新token状态（成功）
             update_token_status(current_token, True)
-            
+
             # 返回OpenAI格式的完整响应
             return await generate_openai_response(full_response, request_id, chat_request.model)
-            
+
     except HTTPException:
         raise
     except Exception as e:
@@ -468,6 +483,7 @@ async def create_chat_completion(
         update_token_status(current_token, False, str(e))
         raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(e)}")
 
+
 # 查看token状态的端点
 @app.get("/admin/tokens")
 async def get_token_status(admin_key: str = Header(None, alias="X-Admin-Key")):
@@ -476,15 +492,16 @@ async def get_token_status(admin_key: str = Header(None, alias="X-Admin-Key")):
     expected_admin_key = os.getenv("ADMIN_KEY", "admin")
     if not admin_key or admin_key != expected_admin_key:
         raise HTTPException(status_code=403, detail="Unauthorized")
-    
+
     # 脱敏token，只显示前8位
     safe_status = {}
     for token, status in token_status.items():
         token_display = token[:8] + "..." if len(token) > 8 else token
         safe_status[token_display] = status
-    
+
     return {"tokens": safe_status, "active_tokens": sum(1 for s in token_status.values() if s["active"])}
 
+
 # 手动刷新token状态
 @app.post("/admin/refresh-tokens")
 async def refresh_token_status(admin_key: str = Header(None, alias="X-Admin-Key")):
@@ -493,10 +510,11 @@ async def refresh_token_status(admin_key: str = Header(None, alias="X-Admin-Key"
     expected_admin_key = os.getenv("ADMIN_KEY", "admin")
     if not admin_key or admin_key != expected_admin_key:
         raise HTTPException(status_code=403, detail="Unauthorized")
-    
+
     await initialize_token_status()
     return {"message": "所有token状态已刷新", "active_tokens": sum(1 for s in token_status.values() if s["active"])}
 
+
 # 模拟模型的路由
 @app.get("/v1/engines")
 @app.get("/v1/engines/{engine_id}")
@@ -504,6 +522,7 @@ async def engines_handler():
     """兼容旧的引擎API"""
     raise HTTPException(status_code=404, detail="引擎API已被弃用 请使用模型API")
 
+
 # 错误处理器
 @app.exception_handler(404)
 async def not_found_handler(request, exc):
@@ -515,6 +534,7 @@ async def not_found_handler(request, exc):
         }
     }, 404
 
+
 # 启动事件
 @app.on_event("startup")
 async def startup_event():
@@ -527,6 +547,7 @@ async def startup_event():
         active_tokens = sum(1 for s in token_status.values() if s["active"])
         logger.info(f"初始化完成 活跃token: {active_tokens}/{len(DEEPSIDER_TOKEN)}")
 
+
 # 主程序
 if __name__ == "__main__":
     # 启动服务器