Update app.py

app.py

@@ -15,7 +15,7 @@ import os
 from dotenv import load_dotenv
 
 from apscheduler.schedulers.background import BackgroundScheduler
-  
+
 # 加载环境变量
 load_dotenv()
 
@@ -50,22 +50,22 @@ TOKEN_INDEX = 0
 # 模型映射表
 MODEL_MAPPING = {
     "gpt-4o-mini": "openai/gpt-4o-mini",
-      "gpt-4o": "openai/gpt-4o",
-      "o1": "openai/o1",
-      "o3-mini": "oopenai/o3-mini",
-      "claude-3.5-sonnet": "anthropic/claude-3.5-sonnet",
-      "claude-3.7-sonnet": "anthropic/claude-3.7-sonnet",
-      "grok-3": "x-ai/grok-3",
-      "grok-3-reasoner":"x-ai/grok-3-reasoner",
-      "deepseek-v3":"deepseek/deepseek-chat",
-      "deepseek-r1":"deepseek/deepseek-r1",
-      "gemini-2.0-flash":"google/gemini-2.0-flash",
-      "gemini-2.0-pro-exp":"google/gemini-2.0-pro-exp-02-05",
-      "gemini-2.0-flash-thinking-exp":"google/gemini-2.0-flash-thinking-exp-1219",
-      "qwq-32b":"qwen/qwq-32b",
-      "qwen-max":"qwen/qwen-max",
-      "DeepSeek-V3-0324":"deepseek/deepseek-chat-v3-0324",
-       "gpt-4o-image": "openai/gpt-4o-image"
+    "gpt-4o": "openai/gpt-4o",
+    "o1": "openai/o1",
+    "o3-mini": "openai/o3-mini",
+    "claude-3.5-sonnet": "anthropic/claude-3.5-sonnet",
+    "claude-3.7-sonnet": "anthropic/claude-3.7-sonnet",
+    "grok-3": "x-ai/grok-3",
+    "grok-3-reasoner": "x-ai/grok-3-reasoner",
+    "deepseek-v3": "deepseek/deepseek-chat",
+    "deepseek-r1": "deepseek/deepseek-r1",
+    "gemini-2.0-flash": "google/gemini-2.0-flash",
+    "gemini-2.0-pro-exp": "google/gemini-2.0-pro-exp-02-05",
+    "gemini-2.0-flash-thinking-exp": "google/gemini-2.0-flash-thinking-exp-1219",
+    "qwq-32b": "qwen/qwq-32b",
+    "qwen-max": "qwen/qwen-max",
+    "DeepSeek-V3-0324": "deepseek/deepseek-chat-v3-0324",
+    "gpt-4o-image": "openai/gpt-4o-image"
 }
 
 # Token负载均衡状态
@@ -131,15 +131,17 @@ class ChatCompletionRequest(BaseModel):
     frequency_penalty: Optional[float] = 0
     user: Optional[str] = None
 
+
 def reset_task():
     try:
-        for key,value in token_status.items():
+        for key, value in token_status.items():
             token_status[key]["total"] = value["count"]
             token_status[key]["active"] = True
         print(f"执行重置任务... 当前时间: {datetime.now()}")
     except Exception as e:
         print(f"任务执行出错: {e}")
 
+
 # 初始化token状态
 async def initialize_token_status():
     """初始化检查所有token的状态和余额"""
@@ -228,7 +230,8 @@ def format_messages_for_deepsider(messages: List[ChatMessage]) -> str:
             next_code = False
             continue
 
-        if msg.role == 'assistant' and ('验证码提示' in msg.content or '验证码已发送至您的邮箱' in msg.content) and 'clId' in msg.content:
+        if msg.role == 'assistant' and (
+                '验证码提示' in msg.content or '验证码已发送至您的邮箱' in msg.content) and 'clId' in msg.content:
             next_code = True
             continue
 
@@ -241,7 +244,7 @@ def format_messages_for_deepsider(messages: List[ChatMessage]) -> str:
             prompt += f"Human: {msg.content}\n\n"
         elif role == "assistant":
             content = msg.content
-            re.sub(r'\[clId:(.*)]','' ,content)
+            re.sub(r'\[clId:(.*)]', '', content)
             prompt += f"Assistant: {msg.content}\n\n"
 
         else:
@@ -363,7 +366,7 @@ async def stream_openai_response(response, request_id: str, model: str, token: s
 
                     elif data.get('code') == 203:
                         # 尾巴
-                        if codeFlag:
+                        if codeFlag or model == 'gpt-4o-image':
                             chunk = {
                                 "id": f"chatcmpl-{request_id}",
                                 "object": "chat.completion.chunk",
@@ -471,19 +474,24 @@ async def create_chat_completion(
 
     # 映射模型
     deepsider_model = map_openai_to_deepsider_model(chat_request.model)
-
+    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 or '验证码已发送至您的邮箱' in msg.content) and 'clId' in msg.content:
+        if msg.role == 'assistant' and (
+                '验证码提示' in msg.content or '验证码已发送至您的邮箱' in msg.content) and 'clId' in msg.content:
             isCode = True
-            clId = re.search(r'\[clId:(.*)]', msg.content).group(1)
-
-    # 准备DeepSider API所需的提示
-    if isCode:
+        
+        pattern = r'\[clId:(.*)]'
+        match = re.search(pattern, msg.content)
+        if match:
+            clId = match.group(1)
+
+        # 准备DeepSider API所需的提示
+    if isCode or model == 'gpt-4o-image':
         prompt = chat_request.messages[-1].content
     else:
         prompt = format_messages_for_deepsider(chat_request.messages)
@@ -496,7 +504,7 @@ async def create_chat_completion(
         "timezone": "Asia/Shanghai"
     }
 
-    if isCode:
+    if isCode or model == 'gpt-4o-image':
         payload["clId"] = clId
 
     # 获取当前token
@@ -632,6 +640,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__":
     scheduler = BackgroundScheduler()