Update app.py

app.py

@@ -12,6 +12,8 @@ import ctypes
 from wasmtime import Store, Module, Linker
 import re
 import transformers
+import queue
+import threading
 
 # -------------------------- 初始化 tokenizer --------------------------
 chat_tokenizer_dir = "THUDM/chatglm2-6b"  # 使用现成的模型tokenizer
@@ -114,14 +116,13 @@ DEEPSEEK_COMPLETION_URL = f"https://{DEEPSEEK_HOST}/api/v0/chat/completion"
 
 BASE_HEADERS = {
     'Host': "chat.deepseek.com",
-    'User-Agent': "DeepSeek/1.0.7 Android/34",
+    'User-Agent': "DeepSeek/1.0.13 Android/35",
     'Accept': "application/json",
     'Accept-Encoding': "gzip",
     'Content-Type': "application/json",
     'x-client-platform': "android",
-    'x-client-version': "1.0.7",
+    'x-client-version': "1.0.13",
     'x-client-locale': "zh_CN",
-    'x-rangers-id': "7883327620434123524",
     'accept-charset': "UTF-8",
 }
 
@@ -184,21 +185,36 @@ def login_deepseek_via_account(account):
     return new_token
 
 # ----------------------------------------------------------------------
-# (4) 从 accounts 中随机选择一个未忙且未尝试过的账号
-# ----------------------------------------------------------------------
-def choose_new_account(exclude_ids):
-    accounts = CONFIG.get("accounts", [])
-    available = [
-        acc for acc in accounts
-        if get_account_identifier(acc) not in exclude_ids and get_account_identifier(acc) not in active_accounts
-    ]
-    if available:
-        chosen = random.choice(available)
-        app.logger.info(f"[choose_new_account] 新选择账号: {get_account_identifier(chosen)}")
-        return chosen
-    app.logger.warning("[choose_new_account] 没有可用的账号")
+# -------------------------- 全局账号队列 --------------------------
+account_queue = []  # 维护所有可用账号
+
+def init_account_queue():
+    """初始化时从配置加载账号"""
+    global account_queue
+    account_queue = CONFIG.get("accounts", [])[:]  # 深拷贝
+    random.shuffle(account_queue)  # 初始随机排序
+
+init_account_queue()
+
+def choose_new_account():
+    """选择策略：
+    1. 遍历队列，找到第一个未被 exclude_ids 包含的账号
+    2. 从队列中移除该账号
+    3. 返回该账号（由后续逻辑保证最终会重新入队）
+    """
+    for i in range(len(account_queue)):
+        acc = account_queue[i]
+        acc_id = get_account_identifier(acc)
+        if acc_id:
+            # 从队列中移除并返回
+            return account_queue.pop(i)
+    app.logger.warning("[choose_new_account] 没有可用的账号或所有账号都在使用中")
     return None
 
+def release_account(account):
+    """将账号重新加入队列末尾"""
+    account_queue.append(account)
+
 # ----------------------------------------------------------------------
 # (5) 判断调用模式：配置模式 vs 用户自带 token
 # ----------------------------------------------------------------------
@@ -218,10 +234,10 @@ def determine_mode_and_token():
     if caller_key in config_keys:
         g.use_config_token = True
         g.tried_accounts = []  # 初始化已尝试账号
-        selected_account = choose_new_account(g.tried_accounts)
+        selected_account = choose_new_account()
         if not selected_account:
-            return Response(json.dumps({"error": "No accounts configured."}),
-                            status=500, mimetype="application/json")
+            return Response(json.dumps({"error": "No accounts configured or all accounts are busy."}),
+                            status=429, mimetype="application/json")
         if not selected_account.get("token", "").strip():
             try:
                 login_deepseek_via_account(selected_account)
@@ -229,15 +245,11 @@ def determine_mode_and_token():
                 app.logger.error(f"[determine_mode_and_token] 账号 {get_account_identifier(selected_account)} 登录失败：{e}")
                 return Response(json.dumps({"error": "Account login failed."}),
                                 status=500, mimetype="application/json")
-        else:
-            app.logger.info(f"[determine_mode_and_token] 账号 {get_account_identifier(selected_account)} 已有 token，无需重新登录")
         g.deepseek_token = selected_account.get("token")
         g.account = selected_account
-        app.logger.info(f"[determine_mode_and_token] 配置模式：使用账号 {get_account_identifier(selected_account)} 的 token")
     else:
         g.use_config_token = False
         g.deepseek_token = caller_key
-        app.logger.info("[determine_mode_and_token] 使用用户自带 DeepSeek token")
     return None
 
 def get_auth_headers():
@@ -299,7 +311,7 @@ def create_session(max_attempts=3):
                     g.tried_accounts = []
                 if current_id not in g.tried_accounts:
                     g.tried_accounts.append(current_id)
-                new_account = choose_new_account(g.tried_accounts)
+                new_account = choose_new_account()
                 if new_account is None:
                     break
                 try:
@@ -464,7 +476,7 @@ def get_pow_response(max_attempts=3):
                     g.tried_accounts = []
                 if current_id not in g.tried_accounts:
                     g.tried_accounts.append(current_id)
-                new_account = choose_new_account(g.tried_accounts)
+                new_account = choose_new_account()
                 if new_account is None:
                     break
                 try:
@@ -564,8 +576,8 @@ def messages_prepare(messages: list) -> str:
         else:
             parts.append(text)
     final_prompt = "".join(parts)
-    # 移除 markdown 图片格式：
-    final_prompt = re.sub(r"!", "", final_prompt)
+    # 仅移除 markdown 图片格式(不全部移除 !）
+    final_prompt = re.sub(r"!\[(.*?)\]\((.*?)\)", r"[\1](\2)", final_prompt)
     return final_prompt
 
 # ----------------------------------------------------------------------
@@ -577,22 +589,6 @@ def chat_completions():
     if mode_resp:
         return mode_resp
 
-    # 如果使用配置模式，检查账号是否正忙；如果忙则尝试切换账号
-    if g.use_config_token:
-        account_id = get_account_identifier(g.account)
-        if account_id in active_accounts:
-            g.tried_accounts.append(account_id)
-            new_account = choose_new_account(g.tried_accounts)
-            if new_account is None:
-                return jsonify({"error": "All accounts are busy."}), 503
-            try:
-                login_deepseek_via_account(new_account)
-            except Exception as e:
-                return jsonify({"error": "Account login failed."}), 500
-            g.account = new_account
-            g.deepseek_token = new_account.get("token")
-            account_id = get_account_identifier(new_account)
-        active_accounts.add(account_id)
     try:
         req_data = request.json or {}
         app.logger.info(f"[chat_completions] 收到请求: {req_data}")
@@ -661,29 +657,81 @@ def chat_completions():
                                 status=deepseek_resp.status_code,
                                 mimetype="application/json")
 
+            # 添加保活超时配置（5秒）
+            KEEP_ALIVE_TIMEOUT = 5
+
             def sse_stream():
                 try:
                     final_text = ""
                     final_thinking = ""
                     first_chunk_sent = False
+                    result_queue = queue.Queue()
+                    last_send_time = time.time()
                     citation_map = {}  # 用于存储引用链接的字典
-                    for raw_line in deepseek_resp.iter_lines(chunk_size=512):
+
+                    def process_data():
                         try:
-                            line = raw_line.decode("utf-8")
+                            for raw_line in deepseek_resp.iter_lines():
+                                try:
+                                    line = raw_line.decode("utf-8")
+                                except Exception as e:
+                                    app.logger.warning(f"[sse_stream] 解码失败: {e}")
+                                    busy_content_str = '{"choices":[{"index":0,"delta":{"content":"服务器繁忙，请稍候再试","type":"text"}}],"model":"","chunk_token_usage":1,"created":0,"message_id":-1,"parent_id":-1}'
+                                    busy_content = json.loads(busy_content_str)
+                                    result_queue.put(busy_content)
+                                    result_queue.put(None)
+                                    break
+                                if not line:
+                                    continue
+                                if line.startswith("data:"):
+                                    data_str = line[5:].strip()
+                                    if data_str == "[DONE]":
+                                        result_queue.put(None)  # 结束信号
+                                        break
+                                    try:
+                                        chunk = json.loads(data_str)
+                                        # 处理搜索索引数据
+                                        if chunk.get("choices", [{}])[0].get("delta", {}).get("type") == "search_index":
+                                            search_indexes = chunk["choices"][0]["delta"].get("search_indexes", [])
+                                            for idx in search_indexes:
+                                                citation_map[str(idx.get("cite_index"))] = idx.get("url", "")
+                                            continue
+                                        result_queue.put(chunk)  # 将数据放入队列
+                                    except Exception as e:
+                                        app.logger.warning(f"[sse_stream] 无法解析: {data_str}, 错误: {e}")
+                                        busy_content_str = '{"choices":[{"index":0,"delta":{"content":"服务器繁忙，请稍候再试","type":"text"}}],"model":"","chunk_token_usage":1,"created":0,"message_id":-1,"parent_id":-1}'
+                                        busy_content = json.loads(busy_content_str)
+                                        result_queue.put(busy_content)
+                                        result_queue.put(None)
+                                        break
                         except Exception as e:
-                            app.logger.warning(f"[sse_stream] 解码失败: {e}")
+                            app.logger.warning(f"[sse_stream] 错误: {e}")
+                            busy_content_str = '{"choices":[{"index":0,"delta":{"content":"服务器繁忙，请稍候再试","type":"text"}}],"model":"","chunk_token_usage":1,"created":0,"message_id":-1,"parent_id":-1}'
+                            busy_content = json.loads(busy_content_str)
+                            result_queue.put(busy_content)
+                            result_queue.put(None)
+                        finally:
+                            deepseek_resp.close()
+
+                    process_thread = threading.Thread(target=process_data)
+                    process_thread.start()
+
+                    while True:
+                        current_time = time.time()
+                        if current_time - last_send_time >= KEEP_ALIVE_TIMEOUT:
+                            yield ": keep-alive\n\n"
+                            last_send_time = current_time
                             continue
-                        if not line:
-                            continue
-                        if line.startswith("data:"):
-                            data_str = line[5:].strip()
-                            if data_str == "[DONE]":
+                        try:
+                            chunk = result_queue.get(timeout=0.1)
+                            if chunk is None:
+                                # 发送最终统计信息
                                 prompt_tokens = len(tokenizer.encode(final_prompt))
                                 completion_tokens = len(tokenizer.encode(final_text))
                                 usage = {
                                     "prompt_tokens": prompt_tokens,
                                     "completion_tokens": completion_tokens,
-                                    "total_tokens": prompt_tokens + completion_tokens
+                                    "total_tokens": prompt_tokens + completion_tokens,
                                 }
                                 finish_chunk = {
                                     "id": completion_id,
@@ -691,30 +739,25 @@ def chat_completions():
                                     "created": created_time,
                                     "model": model,
                                     "choices": [
-                                        {"delta": {}, "index": 0, "finish_reason": "stop"}
+                                        {
+                                            "delta": {},
+                                            "index": 0,
+                                            "finish_reason": "stop",
+                                        }
                                     ],
-                                    "usage": usage
+                                    "usage": usage,
                                 }
                                 yield f"data: {json.dumps(finish_chunk, ensure_ascii=False)}\n\n"
                                 yield "data: [DONE]\n\n"
+                                last_send_time = current_time
                                 break
-                            try:
-                                chunk = json.loads(data_str)
-                                app.logger.debug(f"[sse_stream] 解析到 chunk: {chunk}")
-                                # 处理搜索索引数据
-                                if chunk.get("choices", [{}])[0].get("delta", {}).get("type") == "search_index":
-                                    search_indexes = chunk["choices"][0]["delta"].get("search_indexes", [])
-                                    for idx in search_indexes:
-                                        citation_map[str(idx.get("cite_index"))] = idx.get("url", "")
-                                    continue
-                            except Exception as e:
-                                app.logger.warning(f"[sse_stream] 无法解析: {data_str}, 错误: {e}")
-                                continue
                             new_choices = []
                             for choice in chunk.get("choices", []):
                                 delta = choice.get("delta", {})
                                 ctype = delta.get("type")
                                 ctext = delta.get("content", "")
+                                if choice.get("finish_reason") == "backend_busy":
+                                    ctext = '服务器繁忙，请稍候再试'
                                 if search_enabled and ctext.startswith("[citation:"):
                                     ctext = ""
                                 if ctype == "thinking":
@@ -732,96 +775,141 @@ def chat_completions():
                                 elif ctype == "text":
                                     delta_obj["content"] = ctext
                                 if delta_obj:
-                                    new_choices.append({"delta": delta_obj, "index": choice.get("index", 0)})
+                                    new_choices.append(
+                                        {
+                                            "delta": delta_obj,
+                                            "index": choice.get("index", 0),
+                                        }
+                                    )
                             if new_choices:
                                 out_chunk = {
                                     "id": completion_id,
                                     "object": "chat.completion.chunk",
                                     "created": created_time,
                                     "model": model,
-                                    "choices": new_choices
+                                    "choices": new_choices,
                                 }
                                 yield f"data: {json.dumps(out_chunk, ensure_ascii=False)}\n\n"
+                                last_send_time = current_time
+                        except queue.Empty:
+                            continue
                 except Exception as e:
                     app.logger.error(f"[sse_stream] 异常: {e}")
                 finally:
                     deepseek_resp.close()
                     if g.use_config_token:
-                        active_accounts.discard(get_account_identifier(g.account))
+                        release_account(g.account)
             return Response(stream_with_context(sse_stream()), content_type="text/event-stream")
         else:
             # 非流式响应处理
             think_list = []
             text_list = []
+            result = None
             citation_map = {}  # 用于存储引用链接的字典
-            try:
-                for raw_line in deepseek_resp.iter_lines(chunk_size=512):
-                    try:
-                        line = raw_line.decode("utf-8")
-                    except Exception as e:
-                        app.logger.warning(f"[chat_completions] 解码失败: {e}")
-                        continue
-                    if not line:
-                        continue
-                    if line.startswith("data:"):
-                        data_str = line[5:].strip()
-                        if data_str == "[DONE]":
-                            break
+
+            data_queue = queue.Queue()
+
+            def collect_data():
+                nonlocal result
+                try:
+                    for raw_line in deepseek_resp.iter_lines():
                         try:
-                            chunk = json.loads(data_str)
-                            app.logger.debug(f"[chat_completions] 非流式 chunk: {chunk}")
-                            # 处理搜索索引数据
-                            if chunk.get("choices", [{}])[0].get("delta", {}).get("type") == "search_index":
-                                search_indexes = chunk["choices"][0]["delta"].get("search_indexes", [])
-                                for idx in search_indexes:
-                                    citation_map[str(idx.get("cite_index"))] = idx.get("url", "")
-                                continue
+                            line = raw_line.decode("utf-8")
                         except Exception as e:
-                            app.logger.warning(f"[chat_completions] 无法解析: {data_str}, 错误: {e}")
+                            app.logger.warning(f"[chat_completions] 解码失败: {e}")
+                            ctext = '服务器繁忙，请稍候再试'
+                            text_list.append(ctext)
+                            data_queue.put(None)
+                            break
+                        if not line:
                             continue
-                        for choice in chunk.get("choices", []):
-                            delta = choice.get("delta", {})
-                            ctype = delta.get("type")
-                            ctext = delta.get("content", "")
-                            if search_enabled and ctext.startswith("[citation:"):
-                                ctext = ""
-                            if ctype == "thinking" and thinking_enabled:
-                                think_list.append(ctext)
-                            elif ctype == "text":
+                        if line.startswith("data:"):
+                            data_str = line[5:].strip()
+                            if data_str == "[DONE]":
+                                data_queue.put(None)
+                                break
+                            try:
+                                chunk = json.loads(data_str)
+                                if chunk.get("choices", [{}])[0].get("delta", {}).get("type") == "search_index":
+                                    search_indexes = chunk["choices"][0]["delta"].get("search_indexes", [])
+                                    for idx in search_indexes:
+                                        citation_map[str(idx.get("cite_index"))] = idx.get("url", "")
+                                    continue
+                                for choice in chunk.get("choices", []):
+                                    delta = choice.get("delta", {})
+                                    ctype = delta.get("type")
+                                    ctext = delta.get("content", "")
+                                    if choice.get("finish_reason") == "backend_busy":
+                                        ctext = '服务器繁忙，请稍候再试'
+                                    if search_enabled and ctext.startswith("[citation:"):
+                                        ctext = ""
+                                    if ctype == "thinking" and thinking_enabled:
+                                        think_list.append(ctext)
+                                    elif ctype == "text":
+                                        text_list.append(ctext)
+                            except Exception as e:
+                                app.logger.warning(f"[collect_data] 无法解析: {data_str}, 错误: {e}")
+                                ctext = '服务器繁忙，请稍候再试'
                                 text_list.append(ctext)
-            finally:
-                deepseek_resp.close()
-            final_reasoning = "".join(think_list)
-            final_content = "".join(text_list)
-            prompt_tokens = len(tokenizer.encode(final_prompt))
-            completion_tokens = len(tokenizer.encode(final_content))
-            total_tokens = prompt_tokens + completion_tokens
-            result = {
-                "id": completion_id,
-                "object": "chat.completion",
-                "created": created_time,
-                "model": model,
-                "choices": [
-                    {
-                        "index": 0,
-                        "message": {
-                            "role": "assistant",
-                            "content": final_content,
-                            "reasoning_content": final_reasoning
+                                data_queue.put(None)
+                                break
+                except Exception as e:
+                    app.logger.warning(f"[collect_data] 错误: {e}")
+                    ctext = '服务器繁忙，请稍候再试'
+                    text_list.append(ctext)
+                    data_queue.put(None)
+                finally:
+                    deepseek_resp.close()
+                    final_reasoning = "".join(think_list)
+                    final_content = "".join(text_list)
+                    prompt_tokens = len(tokenizer.encode(final_prompt))
+                    completion_tokens = len(tokenizer.encode(final_content))
+                    result = {
+                        "id": completion_id,
+                        "object": "chat.completion",
+                        "created": created_time,
+                        "model": model,
+                        "choices": [
+                            {
+                                "index": 0,
+                                "message": {
+                                    "role": "assistant",
+                                    "content": final_content,
+                                    "reasoning_content": final_reasoning,
+                                },
+                                "finish_reason": "stop",
+                            }
+                        ],
+                        "usage": {
+                            "prompt_tokens": prompt_tokens,
+                            "completion_tokens": completion_tokens,
+                            "total_tokens": prompt_tokens + completion_tokens,
                         },
-                        "finish_reason": "stop"
                     }
-                ],
-                "usage": {
-                    "prompt_tokens": prompt_tokens,
-                    "completion_tokens": completion_tokens,
-                    "total_tokens": total_tokens
-                }
-            }
-            return jsonify(result), 200
+                    data_queue.put("DONE")
+
+            collect_thread = threading.Thread(target=collect_data)
+            collect_thread.start()
+
+            def generate():
+                last_send_time = time.time()
+                while True:
+                    current_time = time.time()
+                    if current_time - last_send_time >= KEEP_ALIVE_TIMEOUT:
+                        yield ""
+                        last_send_time = current_time
+                    if not collect_thread.is_alive() and result is not None:
+                        yield json.dumps(result)
+                        break
+                    time.sleep(0.1)
+
+            return Response(generate(), mimetype="application/json")
+    except Exception as e:
+        app.logger.error(f"[chat_completions] 未知异常: {e}")
+        return jsonify({"error": "Internal Server Error"}), 500
     finally:
         if g.use_config_token:
-            active_accounts.discard(get_account_identifier(g.account))
+            release_account(g.account)
 
 # ----------------------------------------------------------------------
 # (11) 路由：/