Update app.py

app.py

@@ -1,4 +1,5 @@
 import json
+import re
 import time
 import asyncio
 import uvicorn
@@ -74,7 +75,7 @@ def get_headers():
     # 负载均衡，轮询选择token
     if len(DEEPSIDER_TOKEN) > 0:
         current_token = DEEPSIDER_TOKEN[TOKEN_INDEX % len(DEEPSIDER_TOKEN)]
-        TOKEN_INDEX = (TOKEN_INDEX + 1) % 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"]:
@@ -149,6 +150,7 @@ async def initialize_token_status():
 
             active = False
             quota_info = {}
+            count = 0
 
             if response.status_code == 200:
                 data = response.json()
@@ -159,7 +161,7 @@ async def initialize_token_status():
                     for item in quota_list:
                         item_type = item.get('type', '')
                         available = item.get('available', 0)
-
+                        count += available
                         if available > 0:
                             active = True
 
@@ -173,7 +175,8 @@ async def initialize_token_status():
                 "active": active,
                 "quota": quota_info,
                 "last_checked": datetime.now(),
-                "failed_count": 0
+                "failed_count": 0,
+                "count": count
             }
 
             logger.info(f"Token {token[:8]}... 状态：{'活跃' if active else '无效'}")
@@ -213,7 +216,10 @@ def format_messages_for_deepsider(messages: List[ChatMessage]) -> str:
         elif role == "user":
             prompt += f"Human: {msg.content}\n\n"
         elif role == "assistant":
+            content = msg.content
+            re.sub(r'\[clId:(.*)]','' ,content)
             prompt += f"Assistant: {msg.content}\n\n"
+
         else:
             # 其他角色按用户处理
             prompt += f"Human ({role}): {msg.content}\n\n"
@@ -251,8 +257,13 @@ def update_token_status(token: str, success: bool, error_message: str = None):
             logger.warning(f"Token {token[:8]}... 连续失败{token_status[token]['failed_count']}次，已标记为不活跃")
     else:
         # 成功则重置失败计数
+        token_status[token]["count"] -= 1
         token_status[token]["failed_count"] = 0
 
+    if token_status[token]["count"] <= 0:
+        token_status[token]["active"] = False
+        logger.warning(f"Token {token[:8]}... 余额不足，已标记为不活跃")
+
 
 async def generate_openai_response(full_response: str, request_id: str, model: str) -> Dict:
     """生成符合OpenAI API响应格式的完整响应"""
@@ -284,6 +295,7 @@ async def stream_openai_response(response, request_id: str, model: str, token: s
     """流式返回OpenAI API格式的响应"""
     timestamp = int(time.time())
     full_response = ""
+    codeFlag = False
 
     try:
         # 将DeepSider响应流转换为OpenAI流格式
@@ -319,9 +331,31 @@ async def stream_openai_response(response, request_id: str, model: str, token: s
                                     }
                                 ]
                             }
+
+                            if '验证码提示' in content:
+                                codeFlag = True
                             yield f"data: {json.dumps(chunk)}\n\n"
 
                     elif data.get('code') == 203:
+                        # 尾巴
+                        if codeFlag:
+                            chunk = {
+                                "id": f"chatcmpl-{request_id}",
+                                "object": "chat.completion.chunk",
+                                "created": timestamp,
+                                "model": model,
+                                "choices": [
+                                    {
+                                        "index": 0,
+                                        "delta": {
+                                            "content": f"\n[clId:{data.get('data', {}).get('clId')}]"
+                                        },
+                                        "finish_reason": None
+                                    }
+                                ]
+                            }
+                            yield f"data: {json.dumps(chunk)}\n\n"
+
                         # 生成完成信号
                         chunk = {
                             "id": f"chatcmpl-{request_id}",
@@ -410,8 +444,21 @@ async def create_chat_completion(
     # 映射模型
     deepsider_model = map_openai_to_deepsider_model(chat_request.model)
 
+    # 验证验证码
+    isCode = False
+    clId = None
+
+    if len(chat_request.messages) > 1:
+        msg = chat_request.messages[-2]
+        if msg.role == 'assistant' and '验证码提示' in msg.content and 'clId' in msg.content:
+            isCode = True
+            clId = re.search(r'\[clId:(.*)]', msg.content).group(1)
+
     # 准备DeepSider API所需的提示
-    prompt = format_messages_for_deepsider(chat_request.messages)
+    if isCode:
+        prompt = chat_request.messages[-1].content
+    else:
+        prompt = format_messages_for_deepsider(chat_request.messages)
 
     # 准备请求体
     payload = {
@@ -421,6 +468,9 @@ async def create_chat_completion(
         "timezone": "Asia/Shanghai"
     }
 
+    if isCode:
+        payload["clId"] = clId
+
     # 获取当前token
     headers = get_headers()
     current_token = headers["authorization"].replace("Bearer ", "")