Update openai_ondemand_adapter.py

openai_ondemand_adapter.py

@@ -101,15 +101,13 @@ Claude always responds to the person in the language they use or request. If the
 Claude is now being connected with a person."""
 
 # ====== 读取 Huggingface Secret 配置的私有key =======
-# 用于保护此代理服务本身，防止未授权访问
-PRIVATE_KEY = os.environ.get("PRIVATE_KEY", "") # 如果在Huggingface Spaces运行，这里会读取Secrets
-SAFE_HEADERS = ["Authorization", "X-API-KEY"] # 允许传递私有key的请求头
+PRIVATE_KEY = os.environ.get("PRIVATE_KEY", "") 
+SAFE_HEADERS = ["Authorization", "X-API-KEY"] 
 
 # 全局接口访问权限检查
 def check_private_key():
-    # 根路径和favicon通常不需要鉴权
     if request.path in ["/", "/favicon.ico"]:
-        return None # 显式返回 None 表示通过
+        return None 
 
     key_from_header = None
     for header_name in SAFE_HEADERS:
@@ -119,28 +117,22 @@ def check_private_key():
                 key_from_header = key_from_header[len("Bearer "):].strip()
             break
     
-    if not PRIVATE_KEY: # 如果没有设置 PRIVATE_KEY，则不进行鉴权 (方便本地测试)
+    if not PRIVATE_KEY: 
         logging.warning("PRIVATE_KEY 未设置，服务将不进行鉴权！")
         return None
 
     if not key_from_header or key_from_header != PRIVATE_KEY:
         logging.warning(f"未授权访问尝试: Path={request.path}, IP={request.remote_addr}, Key Provided='{key_from_header[:10]}...'")
         return jsonify({"error": "Unauthorized. Correct 'Authorization: Bearer <PRIVATE_KEY>' or 'X-API-KEY: <PRIVATE_KEY>' header is required."}), 401
-    return None # 鉴权通过
+    return None 
 
-# 应用所有API鉴权
 app = Flask(__name__)
 app.before_request(check_private_key)
 
-# ========== OnDemand API KEY池（从环境变量读取，每行一个KEY，用逗号分隔）==========
 ONDEMAND_APIKEYS_STR = os.environ.get("ONDEMAND_APIKEYS", "")
 ONDEMAND_APIKEYS = [key.strip() for key in ONDEMAND_APIKEYS_STR.split(',') if key.strip()]
+BAD_KEY_RETRY_INTERVAL = 600  
 
-BAD_KEY_RETRY_INTERVAL = 600  # 标记为坏的KEY的重试间隔（秒），例如10分钟
-# SESSION_TIMEOUT 已移除，因为我们现在每次都用新会话
-
-# ========== OnDemand模型映射 ==========
-# 将 OpenAI 风格的模型名称映射到 OnDemand 服务的 endpointId
 MODEL_MAP = {
     "gpto3-mini": "predefined-openai-gpto3-mini",
     "gpt-4o": "predefined-openai-gpt4o",
@@ -154,106 +146,85 @@ MODEL_MAP = {
     "gemini-2.0-flash": "predefined-gemini-2.0-flash",
 }
 DEFAULT_ONDEMAND_MODEL = "predefined-openai-gpt4o"
-# ==========================================
 
 class KeyManager:
-    """管理 OnDemand API 密钥池"""
     def __init__(self, key_list):
-        self.key_list = list(key_list) # 存储可用的API密钥
-        self.lock = threading.Lock()   # 线程锁，用于同步访问密钥状态
-        # 存储每个密钥的状态：是否被标记为“坏的”以及标记的时间戳
+        self.key_list = list(key_list) 
+        self.lock = threading.Lock()   
         self.key_status = {key: {"bad": False, "bad_ts": None} for key in self.key_list}
-        self.idx = 0 # 用于轮询密钥的索引
+        self.idx = 0 
 
     def display_key(self, key):
-        """返回部分隐藏的密钥，用于日志输出"""
         if not key or len(key) < 10:
             return "INVALID_KEY_FORMAT"
         return f"{key[:6]}...{key[-4:]}"
 
     def get(self):
-        """获取一个可用的API密钥"""
         with self.lock:
-            if not self.key_list: # 如果密钥池为空
+            if not self.key_list: 
                 logging.error("【KeyManager】API密钥池为空！无法提供密钥。")
                 raise ValueError("API key pool is empty.")
 
             now = time.time()
             num_keys = len(self.key_list)
 
-            for i in range(num_keys): # 尝试遍历所有密钥最多一次
+            for i in range(num_keys): 
                 current_key_candidate = self.key_list[self.idx]
-                self.idx = (self.idx + 1) % num_keys # 移动到下一个密钥，循环使用
-
+                self.idx = (self.idx + 1) % num_keys 
                 status = self.key_status[current_key_candidate]
 
-                if not status["bad"]: # 如果密钥状态良好
+                if not status["bad"]: 
                     logging.info(f"【KeyManager】选择API KEY: {self.display_key(current_key_candidate)} [状态：正常]")
                     return current_key_candidate
                 
-                # 如果密钥被标记为坏的，检查是否已达到重试时间
                 if status["bad_ts"] and (now - status["bad_ts"] >= BAD_KEY_RETRY_INTERVAL):
                     logging.info(f"【KeyManager】API KEY: {self.display_key(current_key_candidate)} 达到重试周期，恢复为正常。")
                     status["bad"] = False
                     status["bad_ts"] = None
                     return current_key_candidate
             
-            # 如果所有密钥都被标记为坏的，并且都未达到重试时间
-            # 强制重置所有密钥状态并返回第一个，这是一种降级策略
             logging.warning("【KeyManager】所有API KEY均被标记为不良且未到重试时间。将强制重置所有KEY状态并尝试第一个。")
             for key_to_reset in self.key_list:
                 self.key_status[key_to_reset]["bad"] = False
                 self.key_status[key_to_reset]["bad_ts"] = None
             self.idx = 0
-            if self.key_list: # 再次检查，以防在极小概率下key_list变空
+            if self.key_list: 
                  selected_key = self.key_list[0]
                  logging.info(f"【KeyManager】强制选择API KEY: {self.display_key(selected_key)} [状态：强制重试]")
                  return selected_key
-            else: # 理论上不应该到这里，因为前面有检查
+            else: 
                 logging.error("【KeyManager】在强制重试逻辑中发现密钥池为空！")
                 raise ValueError("API key pool became empty during forced retry logic.")
 
-
     def mark_bad(self, key):
-        """将指定的API密钥标记为“坏的”"""
         with self.lock:
             if key in self.key_status and not self.key_status[key]["bad"]:
                 logging.warning(f"【KeyManager】禁用API KEY: {self.display_key(key)}。将在 {BAD_KEY_RETRY_INTERVAL // 60} 分钟后自动重试。")
                 self.key_status[key]["bad"] = True
                 self.key_status[key]["bad_ts"] = time.time()
 
-# 初始化 KeyManager
 if not ONDEMAND_APIKEYS:
     logging.warning("【启动警告】ONDEMAND_APIKEYS 环境变量未设置或为空。服务可能无法正常工作。")
 keymgr = KeyManager(ONDEMAND_APIKEYS)
 
-ONDEMAND_API_BASE = "https://api.on-demand.io/chat/v1" # OnDemand API 的基础URL
+ONDEMAND_API_BASE = "https://api.on-demand.io/chat/v1" 
 
 def get_endpoint_id(openai_model_name):
-    """根据用户提供的OpenAI模型名称，从MODEL_MAP中查找对应的OnDemand endpointId"""
     normalized_model_name = str(openai_model_name or "").lower().replace(" ", "")
     return MODEL_MAP.get(normalized_model_name, DEFAULT_ONDEMAND_MODEL)
 
 def create_session(apikey, external_user_id=None, plugin_ids=None):
-    """
-    向 OnDemand API 创建一个新的会话。
-    :param apikey: OnDemand API 密钥。
-    :param external_user_id: 可选，外部用户ID。
-    :param plugin_ids: 可选，插件ID列表。
-    :return: 新创建的会话ID。
-    :raises: requests.HTTPError 如果API调用失败。
-    """
     url = f"{ONDEMAND_API_BASE}/sessions"
-    payload = {"externalUserId": external_user_id or str(uuid.uuid4())} # 如果未提供，则生成UUID
-    if plugin_ids is not None: # 通常聊天场景可能不需要插件
+    payload = {"externalUserId": external_user_id or str(uuid.uuid4())} 
+    if plugin_ids is not None: 
         payload["pluginIds"] = plugin_ids
     
     headers = {"apikey": apikey, "Content-Type": "application/json"}
     
     logging.info(f"【OnDemand】尝试创建新会话... URL: {url}, Key: {keymgr.display_key(apikey)}")
     try:
-        resp = requests.post(url, json=payload, headers=headers, timeout=20) # 设置超时
-        resp.raise_for_status() # 如果状态码不是2xx，则抛出HTTPError
+        resp = requests.post(url, json=payload, headers=headers, timeout=20) 
+        resp.raise_for_status() 
         session_id = resp.json()["data"]["id"]
         logging.info(f"【OnDemand】新会话创建成功: {session_id}, Key: {keymgr.display_key(apikey)}")
         return session_id
@@ -265,17 +236,14 @@ def create_session(apikey, external_user_id=None, plugin_ids=None):
         raise
 
 def format_openai_sse_delta(chunk_data_dict):
-    """将数据块格式化为 OpenAI SSE (Server-Sent Events) 流格式"""
     return f"data: {json.dumps(chunk_data_dict, ensure_ascii=False)}\n\n"
 
+# Modified: This function is now a generator and returns accumulated text + error flag at the end.
 def _execute_one_stream_attempt(apikey_for_attempt, session_id_for_attempt, query_str, endpoint_id, openai_model_name_for_response, current_attempt_num_logging):
     """
-    执行一次流式请求尝试。
-    返回: (generated_sse_strings, accumulated_text_content, api_error_occurred)
-    generated_sse_strings: 此尝试生成的所有SSE事件字符串列表。
-    accumulated_text_content: 从流中累积的纯文本内容。
-    api_error_occurred: 布尔值，指示此尝试是否遇到可直接转换为SSE错误事件的API错误。
-                       requests.RequestException (如超时) 会被直接抛出。
+    Executes one streaming request attempt.
+    Yields SSE event strings.
+    Returns a tuple: (accumulated_text_content, api_error_yielded_flag)
     """
     url = f"{ONDEMAND_API_BASE}/sessions/{session_id_for_attempt}/query"
     payload = {
@@ -290,16 +258,15 @@ def _execute_one_stream_attempt(apikey_for_attempt, session_id_for_attempt, quer
         "Accept": "text/event-stream"
     }
 
-    generated_sse_strings = []
     accumulated_text_parts = []
-    api_error_handled_as_sse = False
+    api_error_yielded = False # Flag to indicate if an API error was processed and yielded as SSE
 
     logging.info(f"【流式请求子尝试 {current_attempt_num_logging}】发送到 OnDemand: Session={session_id_for_attempt}, Endpoint={endpoint_id}, Key={keymgr.display_key(apikey_for_attempt)}")
 
     try:
         with requests.post(url, json=payload, headers=headers, stream=True, timeout=180) as resp:
             if resp.status_code != 200:
-                api_error_handled_as_sse = True
+                api_error_yielded = True
                 error_text = resp.text
                 logging.error(f"【OnDemand流错误】请求失败 (子尝试 {current_attempt_num_logging})。状态码: {resp.status_code}, Session: {session_id_for_attempt}, 响应: {error_text[:500]}")
                 error_payload = {
@@ -309,9 +276,9 @@ def _execute_one_stream_attempt(apikey_for_attempt, session_id_for_attempt, quer
                         "code": resp.status_code
                     }
                 }
-                generated_sse_strings.append(format_openai_sse_delta(error_payload))
-                generated_sse_strings.append("data: [DONE]\n\n")
-                return generated_sse_strings, "".join(accumulated_text_parts), api_error_handled_as_sse
+                yield format_openai_sse_delta(error_payload)
+                yield "data: [DONE]\n\n"
+                return "".join(accumulated_text_parts).strip(), api_error_yielded # Return after yielding error
 
             first_chunk_sent = False
             last_line_str = "" 
@@ -327,19 +294,22 @@ def _execute_one_stream_attempt(apikey_for_attempt, session_id_for_attempt, quer
                     
                     if data_part == "[DONE]":
                         logging.info(f"【OnDemand流】接收到 [DONE] 信号 (子尝试 {current_attempt_num_logging})。Session: {session_id_for_attempt}")
-                        generated_sse_strings.append("data: [DONE]\n\n")
-                        break 
+                        yield "data: [DONE]\n\n"
+                        # No break here, let the function return after the loop
+                        return "".join(accumulated_text_parts).strip(), api_error_yielded
+                    
                     elif data_part.startswith("[ERROR]:"):
-                        api_error_handled_as_sse = True 
+                        api_error_yielded = True 
                         error_json_str = data_part[len("[ERROR]:"):].strip()
                         logging.warning(f"【OnDemand流】接收到错误事件 (子尝试 {current_attempt_num_logging}): {error_json_str}。Session: {session_id_for_attempt}")
                         try:
                             error_obj = json.loads(error_json_str)
                         except json.JSONDecodeError:
                             error_obj = {"message": error_json_str, "type": "on_demand_stream_error_format"}
-                        generated_sse_strings.append(format_openai_sse_delta({"error": error_obj}))
-                        generated_sse_strings.append("data: [DONE]\n\n") 
-                        break
+                        yield format_openai_sse_delta({"error": error_obj})
+                        yield "data: [DONE]\n\n" 
+                        return "".join(accumulated_text_parts).strip(), api_error_yielded # Return after yielding error
+                    
                     else:
                         try:
                             event_data = json.loads(data_part)
@@ -367,31 +337,39 @@ def _execute_one_stream_attempt(apikey_for_attempt, session_id_for_attempt, quer
                                     "model": openai_model_name_for_response,
                                     "choices": [{"delta": choice_delta, "index": 0, "finish_reason": None}]
                                 }
-                                generated_sse_strings.append(format_openai_sse_delta(openai_chunk))
+                                yield format_openai_sse_delta(openai_chunk) # Yield immediately
                         except json.JSONDecodeError:
                             logging.warning(f"【OnDemand流】无法解析JSON (子尝试 {current_attempt_num_logging}): {data_part[:100]}... Session: {session_id_for_attempt}")
                             continue
             
-            if not last_line_str.startswith("data: [DONE]") and not api_error_handled_as_sse:
-                 logging.info(f"【OnDemand流】(子尝试 {current_attempt_num_logging}) 流迭代完成，补充发送 [DONE]。Session: {session_id_for_attempt}")
-                 generated_sse_strings.append("data: [DONE]\n\n")
-
-    except requests.exceptions.RequestException: # Let specific RequestExceptions be handled by the caller
-        raise 
-    except Exception as e: # Catch other unexpected errors during stream processing
-        api_error_handled_as_sse = True 
+            # This part is reached if the loop finishes without an explicit [DONE] or [ERROR] from the stream data itself.
+            # This might happen if the stream just ends.
+            if not api_error_yielded: # If no error was yielded, and no [DONE] was in data, yield a [DONE]
+                if not last_line_str.startswith("data: [DONE]"): # Check if last processed line was not already DONE
+                     logging.info(f"【OnDemand流】(子尝试 {current_attempt_num_logging}) 流迭代完成，补充发送 [DONE]。Session: {session_id_for_attempt}")
+                     yield "data: [DONE]\n\n"
+        
+    except requests.exceptions.RequestException: 
+        # Network/request level errors before or during streaming
+        # These should be caught by the caller (handle_stream_request) to decide on retries
+        logging.error(f"【OnDemand流】请求过程中发生网络或请求异常 (子尝试 {current_attempt_num_logging}): Session: {session_id_for_attempt}", exc_info=False)
+        raise # Re-raise for the caller to handle
+    
+    except Exception as e: 
+        # Unexpected Python errors during stream processing
         logging.error(f"【OnDemand流】处理流时发生未知错误 (子尝试 {current_attempt_num_logging}): {e}, Session: {session_id_for_attempt}", exc_info=True)
+        api_error_yielded = True # Mark that an error occurred and we are yielding an SSE for it
         error_payload = {
             "error": {"message": f"Unknown error during streaming (Attempt {current_attempt_num_logging}): {str(e)}", "type": "unknown_streaming_error_in_attempt"}
         }
-        generated_sse_strings.append(format_openai_sse_delta(error_payload))
-        generated_sse_strings.append("data: [DONE]\n\n")
+        yield format_openai_sse_delta(error_payload)
+        yield "data: [DONE]\n\n"
     
-    return generated_sse_strings, "".join(accumulated_text_parts).strip(), api_error_handled_as_sse
+    return "".join(accumulated_text_parts).strip(), api_error_yielded
+
 
 @app.route("/v1/chat/completions", methods=["POST"])
 def chat_completions():
-    """处理聊天补全请求，模拟 OpenAI /v1/chat/completions 接口"""
     try:
         request_data = request.json
     except Exception as e:
@@ -435,20 +413,20 @@ def chat_completions():
     if not formatted_query_parts:
         return jsonify({"error": "No valid content found in 'messages'."}), 400
 
-    # Use the globally defined CLAUDE_SYSTEM_PROMPT
     start_prompt = CLAUDE_SYSTEM_PROMPT + "\n\n" + """下面是对话历史. 你是Assitant角色，请遵从User指令，并用中文尽可能详细的回复。注意，请直接回复! 请不要在开头提出"根据上下文及历史记录"相关的话语。\n"""
     final_query_to_ondemand = start_prompt + "\n".join(formatted_query_parts)
     
     def attempt_ondemand_request_wrapper(current_apikey_from_wrapper, current_session_id_from_wrapper):
-        # This inner function is what with_valid_key_and_session calls.
-        # It receives the *initial* apikey and session_id.
         if is_stream_request:
-            return handle_stream_request(current_apikey_from_wrapper, current_session_id_from_wrapper, final_query_to_ondemand, target_endpoint_id, openai_model_name)
+            # Pass the generator directly to Response
+            return Response(
+                handle_stream_request(current_apikey_from_wrapper, current_session_id_from_wrapper, final_query_to_ondemand, target_endpoint_id, openai_model_name),
+                content_type='text/event-stream'
+            )
         else:
             return handle_non_stream_request(current_apikey_from_wrapper, current_session_id_from_wrapper, final_query_to_ondemand, target_endpoint_id, openai_model_name)
 
     def with_valid_key_and_session(action_func_to_wrap):
-        # This is the outer retry loop for API key/session issues for the *first* attempt of action_func_to_wrap
         max_key_retries = len(ONDEMAND_APIKEYS) * 2 if ONDEMAND_APIKEYS else 1 
         key_retry_count = 0
         last_exception_for_key_retry = None
@@ -460,60 +438,70 @@ def chat_completions():
                 selected_apikey_for_outer_retry = keymgr.get() 
                 logging.info(f"【请求处理 - Key轮换尝试 {key_retry_count}/{max_key_retries}】使用 API Key: {keymgr.display_key(selected_apikey_for_outer_retry)}，准备创建新会话...")
                 ondemand_session_id_for_outer_retry = create_session(selected_apikey_for_outer_retry)
-                
-                # Call the action_func_to_wrap (which is attempt_ondemand_request_wrapper)
-                # This action_func_to_wrap will then call handle_stream_request or handle_non_stream_request
-                # If handle_..._request fails its *first* attempt due to RequestException, it re-raises it here.
                 return action_func_to_wrap(selected_apikey_for_outer_retry, ondemand_session_id_for_outer_retry)
 
-            except ValueError as ve: # keymgr.get() failed
+            except ValueError as ve: 
                 logging.critical(f"【请求处理 - Key轮换尝试 {key_retry_count}】KeyManager 错误: {ve}")
                 last_exception_for_key_retry = ve
-                break # Cannot get any key, fatal for this request.
-            except requests.exceptions.RequestException as http_err_outer: # Covers create_session failure or re-raised error from action_func's first attempt
+                break 
+            except requests.exceptions.RequestException as http_err_outer: 
                 last_exception_for_key_retry = http_err_outer
                 logging.warning(f"【请求处理 - Key轮换尝试 {key_retry_count}】HTTP/请求错误。Key: {keymgr.display_key(selected_apikey_for_outer_retry) if selected_apikey_for_outer_retry else 'N/A'}, Error: {http_err_outer}")
-                if selected_apikey_for_outer_retry: # If a key was involved in this failure
+                if selected_apikey_for_outer_retry: 
                     keymgr.mark_bad(selected_apikey_for_outer_retry)
                 
                 if key_retry_count >= max_key_retries:
                     logging.error(f"【请求处理】所有Key轮换尝试均失败。最后错误: {last_exception_for_key_retry}")
-                    break # Exhausted key retries
+                    break 
                 
                 logging.info(f"【请求处理】Key轮换尝试 {key_retry_count} 失败，等待后重试下一个Key...")
                 time.sleep(1) 
-                continue # To the next iteration of the key_retry_count loop
-            except Exception as e_outer: # Other unexpected errors during the initial setup/call
+                continue 
+            except Exception as e_outer: 
                 last_exception_for_key_retry = e_outer
                 logging.error(f"【请求处理 - Key轮换尝试 {key_retry_count}】发生意外严重错误: {e_outer}", exc_info=True)
                 if selected_apikey_for_outer_retry:
                     keymgr.mark_bad(selected_apikey_for_outer_retry) 
-                # For truly unexpected errors, might be better to fail fast
-                break # Break outer retry loop
+                break 
         
         error_message = "All attempts to process the request failed after multiple key/session retries."
         if last_exception_for_key_retry:
             error_message += f" Last known error during key/session phase: {str(last_exception_for_key_retry)}"
         logging.error(error_message)
-        return jsonify({"error": error_message}), 503
+        # For stream requests, if with_valid_key_and_session fails, we can't return jsonify directly
+        # This part might need adjustment if the action_func_to_wrap for stream is expected to return a Response object
+        # However, if action_func_to_wrap (attempt_ondemand_request_wrapper) for stream returns a Response,
+        # then this jsonify will only be hit if create_session or keymgr.get fails repeatedly.
+        if is_stream_request:
+             # Construct a generator that yields an error SSE
+            def error_stream_gen():
+                yield format_openai_sse_delta({"error": {"message": error_message, "type": "proxy_setup_error", "code": 503}})
+                yield "data: [DONE]\n\n"
+            return Response(error_stream_gen(), content_type='text/event-stream', status=503)
+        else:
+            return jsonify({"error": error_message}), 503
 
     return with_valid_key_and_session(attempt_ondemand_request_wrapper)
 
-
+# Modified: This function is now a generator that uses `yield from`
 def handle_stream_request(initial_apikey, initial_session_id, query_str, endpoint_id, openai_model_name_for_response):
-    """处理流式聊天补全请求，包含空回复重试逻辑 (每次空回复重试使用新Key/Session)"""
+    """
+    Handles streaming chat completion requests with empty response retries.
+    Each empty response retry uses a new API key and session.
+    Yields SSE event strings directly.
+    """
     max_empty_response_retries = 5
     empty_retry_attempt_num = 0
     
     current_apikey_for_attempt = initial_apikey
     current_session_id_for_attempt = initial_session_id
     
-    final_sse_strings_to_yield = []
-
     while empty_retry_attempt_num < max_empty_response_retries:
         empty_retry_attempt_num += 1
+        accumulated_text_this_attempt = ""
+        api_error_in_attempt = False
         
-        if empty_retry_attempt_num > 1: # This is an empty-response retry, get new key/session
+        if empty_retry_attempt_num > 1: 
             logging.info(f"【流式请求-空回复重试 {empty_retry_attempt_num-1}】获取新Key/Session...")
             try:
                 current_apikey_for_attempt = keymgr.get()
@@ -521,95 +509,98 @@ def handle_stream_request(initial_apikey, initial_session_id, query_str, endpoin
                 logging.info(f"【流式请求-空回复重试 {empty_retry_attempt_num-1}】新Key/Session获取成功: Key={keymgr.display_key(current_apikey_for_attempt)}, Session={current_session_id_for_attempt}")
             except (ValueError, requests.exceptions.RequestException) as e_key_session:
                 logging.warning(f"【流式请求-空回复重��� {empty_retry_attempt_num-1}】获取新Key/Session失败: {e_key_session}")
-                if current_apikey_for_attempt and not isinstance(e_key_session, ValueError): # If key was obtained but create_session failed
+                if current_apikey_for_attempt and not isinstance(e_key_session, ValueError): 
                     keymgr.mark_bad(current_apikey_for_attempt)
                 if empty_retry_attempt_num >= max_empty_response_retries:
-                    final_sse_strings_to_yield = [
-                        format_openai_sse_delta({"error": {"message": f"Failed to get new key/session for final empty stream retry. Error: {str(e_key_session)}", "type": "internal_proxy_error"}}),
-                        "data: [DONE]\n\n"
-                    ]
-                    break # Break empty retry loop
+                    yield format_openai_sse_delta({"error": {"message": f"Failed to get new key/session for final empty stream retry. Error: {str(e_key_session)}", "type": "internal_proxy_error"}})
+                    yield "data: [DONE]\n\n"
+                    return 
                 time.sleep(1)
-                current_apikey_for_attempt = None # Ensure it's reset if keymgr.get() failed
-                continue # Next iteration of empty_retry_attempt_num loop
+                current_apikey_for_attempt = None 
+                continue 
 
-        # Log which attempt this is (1st overall, or nth empty-response retry)
         log_attempt_str = f"初始尝试" if empty_retry_attempt_num == 1 else f"空回复重试 {empty_retry_attempt_num-1}"
 
         try:
-            sse_strings_this_attempt, accumulated_text_this_attempt, api_error_in_attempt = \
-                _execute_one_stream_attempt(current_apikey_for_attempt, current_session_id_for_attempt, query_str, endpoint_id, openai_model_name_for_response, f"{log_attempt_str} (Overall attempt {empty_retry_attempt_num})")
-
-            final_sse_strings_to_yield = sse_strings_this_attempt 
-
-            if api_error_in_attempt: 
-                logging.warning(f"【流式请求】({log_attempt_str}) OnDemand 服务返回错误或处理内部错误，将返回此错误信息。")
-                break 
-
-            if accumulated_text_this_attempt: 
-                logging.info(f"【流式请求】({log_attempt_str}) 成功获取非空内容。")
-                break 
-            
-            logging.warning(f"【流式请求】({log_attempt_str}) 返回空内容。")
-            if empty_retry_attempt_num >= max_empty_response_retries:
-                logging.error(f"【流式请求】达到最大空回复重试次数 ({max_empty_response_retries})。")
-                empty_error_payload = {
-                    "error": {"message": f"Model returned an empty stream after {max_empty_response_retries} retries.", "type": "empty_stream_error_after_retries", "code": "empty_response"}
-                }
-                final_sse_strings_to_yield = [format_openai_sse_delta(empty_error_payload), "data: [DONE]\n\n"]
-                break 
-            
-            logging.info(f"【流式请求】空回复，将在1秒后重试下一个Key。当前总尝试 {empty_retry_attempt_num}/{max_empty_response_retries}")
-            time.sleep(1)
+            # Yield from the sub-generator; result_tuple will be (accumulated_text, api_error_yielded_flag)
+            # This is where the true streaming to the client happens chunk by chunk.
+            result_tuple = yield from _execute_one_stream_attempt(
+                current_apikey_for_attempt, 
+                current_session_id_for_attempt, 
+                query_str, 
+                endpoint_id, 
+                openai_model_name_for_response, 
+                f"{log_attempt_str} (Overall attempt {empty_retry_attempt_num})"
+            )
+            accumulated_text_this_attempt = result_tuple[0]
+            api_error_in_attempt = result_tuple[1]
 
         except requests.exceptions.RequestException as e_req:
             logging.warning(f"【流式请求】({log_attempt_str} using key {keymgr.display_key(current_apikey_for_attempt)}) 发生请求级错误: {e_req}")
-            keymgr.mark_bad(current_apikey_for_attempt)
-            if empty_retry_attempt_num == 1: # This was the initial_apikey provided by with_valid_key_and_session
-                raise e_req # Re-raise for the outer key retry mechanism
+            if current_apikey_for_attempt: # Ensure key is marked bad if one was used
+                 keymgr.mark_bad(current_apikey_for_attempt)
+
+            if empty_retry_attempt_num == 1: 
+                # This was the initial_apikey. Re-raise for the outer key retry mechanism in with_valid_key_and_session.
+                # The with_valid_key_and_session will then try a new key for the *entire* operation.
+                raise e_req 
             
             # If it's an empty-response retry (attempt_num > 1) that failed with RequestException
             if empty_retry_attempt_num >= max_empty_response_retries:
                 logging.error(f"【流式请求】在���后一次空回复重试时发生请求错误。")
-                final_sse_strings_to_yield = [
-                    format_openai_sse_delta({"error": {"message": f"Request failed on final empty stream retry attempt: {str(e_req)}", "type": "internal_proxy_error"}}),
-                    "data: [DONE]\n\n"
-                ]
-                break # Break empty retry loop
+                yield format_openai_sse_delta({"error": {"message": f"Request failed on final empty stream retry attempt: {str(e_req)}", "type": "internal_proxy_error"}})
+                yield "data: [DONE]\n\n"
+                return 
             time.sleep(1)
-            # Loop continues, will try to get another new key for the next empty-response retry
-            continue
+            continue # To the next iteration of the empty_retry_attempt_num loop (will try new key/session)
             
-    def final_generator_for_response():
-        if not final_sse_strings_to_yield: 
-            logging.error("【流式请求】final_sse_strings_to_yield 为空，返回通用错误。")
-            yield format_openai_sse_delta({"error": {"message": "Unexpected empty result in streaming.", "type": "internal_proxy_error"}})
+        # Check results after _execute_one_stream_attempt has finished for this attempt
+        if api_error_in_attempt: 
+            logging.warning(f"【流式请求】({log_attempt_str}) OnDemand 服务返回错误或处理内部错误，已将错误信息流式传输。")
+            # Error already yielded by _execute_one_stream_attempt, so we just stop.
+            return 
+
+        if accumulated_text_this_attempt: 
+            logging.info(f"【流式请求】({log_attempt_str}) 成功获取非空内容。")
+            # Content already yielded by _execute_one_stream_attempt. We are done.
+            return 
+        
+        # If we reach here, content was empty and no API error was yielded by _execute_one_stream_attempt
+        logging.warning(f"【流式请求】({log_attempt_str}) 返回空内容。")
+        if empty_retry_attempt_num >= max_empty_response_retries:
+            logging.error(f"【流式请求】达到最大空回复重试次数 ({max_empty_response_retries})。")
+            yield format_openai_sse_delta({
+                "error": {"message": f"Model returned an empty stream after {max_empty_response_retries} retries.", "type": "empty_stream_error_after_retries", "code": "empty_response"}
+            })
             yield "data: [DONE]\n\n"
-        else:
-            for sse_str in final_sse_strings_to_yield:
-                yield sse_str
-            
-    return Response(final_generator_for_response(), content_type='text/event-stream')
+            return 
+        
+        logging.info(f"【流式请求】空回复，将在1秒后重试下一个Key。当前总尝试 {empty_retry_attempt_num}/{max_empty_response_retries}")
+        time.sleep(1)
+    # Fallback if loop finishes unexpectedly (shouldn't happen with current logic)
+    logging.error("【流式请求】意外退出空回复重试循环。")
+    yield format_openai_sse_delta({"error": {"message": "Unexpected error in stream handling.", "type": "internal_proxy_error"}})
+    yield "data: [DONE]\n\n"
+
 
 def handle_non_stream_request(initial_apikey, initial_session_id, query_str, endpoint_id, openai_model_name_for_response):
-    """处理非流式聊天补全请求，包含空回复重试逻辑 (每次空回复重试使用新Key/Session)"""
     max_empty_response_retries = 5
-    empty_retry_attempt_num = 0 # Counts total attempts including initial one
+    empty_retry_attempt_num = 0 
 
     current_apikey_for_attempt = initial_apikey
     current_session_id_for_attempt = initial_session_id
     
-    url = f"{ONDEMAND_API_BASE}/sessions/{current_session_id_for_attempt}/query" # URL will change if session_id changes
+    url = f"{ONDEMAND_API_BASE}/sessions/{current_session_id_for_attempt}/query" 
 
     while empty_retry_attempt_num < max_empty_response_retries:
         empty_retry_attempt_num += 1
         
-        if empty_retry_attempt_num > 1: # This is an empty-response retry, get new key/session
+        if empty_retry_attempt_num > 1: 
             logging.info(f"【同步请求-空回复重试 {empty_retry_attempt_num-1}】获取新Key/Session...")
             try:
                 current_apikey_for_attempt = keymgr.get()
                 current_session_id_for_attempt = create_session(current_apikey_for_attempt)
-                url = f"{ONDEMAND_API_BASE}/sessions/{current_session_id_for_attempt}/query" # Update URL with new session
+                url = f"{ONDEMAND_API_BASE}/sessions/{current_session_id_for_attempt}/query" 
                 logging.info(f"【同步请求-空回复重试 {empty_retry_attempt_num-1}】新Key/Session获取成功: Key={keymgr.display_key(current_apikey_for_attempt)}, Session={current_session_id_for_attempt}")
             except (ValueError, requests.exceptions.RequestException) as e_key_session:
                 logging.warning(f"【同步请求-空回复重试 {empty_retry_attempt_num-1}】获取新Key/Session失败: {e_key_session}")
@@ -664,7 +655,8 @@ def handle_non_stream_request(initial_apikey, initial_session_id, query_str, end
         
         except requests.exceptions.RequestException as e_req:
             logging.warning(f"【同步请求】({log_attempt_str} using key {keymgr.display_key(current_apikey_for_attempt)}) 发生请求级错误: {e_req}")
-            keymgr.mark_bad(current_apikey_for_attempt)
+            if current_apikey_for_attempt: # Ensure key is marked bad
+                keymgr.mark_bad(current_apikey_for_attempt)
             if empty_retry_attempt_num == 1: 
                 raise e_req 
             if empty_retry_attempt_num >= max_empty_response_retries:
@@ -676,13 +668,12 @@ def handle_non_stream_request(initial_apikey, initial_session_id, query_str, end
             logging.error(f"【同步请求】({log_attempt_str}) 处理响应时出错: {e_parse}", exc_info=True)
             return jsonify({"error": f"Error processing OnDemand sync response: {str(e_parse)}"}), 502 
     
-    logging.error(f"【同步请求】意外退出空回复重试循环。") # Should be unreachable
+    logging.error(f"【同步请求】意外退出空回复重试循环。") 
     return jsonify({"error": "Unexpected error in non-stream handling after empty response retries."}), 500
 
 
 @app.route("/v1/models", methods=["GET"])
 def list_models():
-    """返回此代理支持的模型列表，模拟 OpenAI /v1/models 接口"""
     model_objects = []
     for model_key_alias in MODEL_MAP.keys():
         model_objects.append({
@@ -698,7 +689,6 @@ def list_models():
 
 @app.route("/", methods=["GET"])
 def health_check():
-    """简单的健康检查端点或首页"""
     num_keys = len(ONDEMAND_APIKEYS)
     key_status_summary = {keymgr.display_key(k): ("OK" if not v["bad"] else f"BAD (since {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(v['bad_ts'])) if v['bad_ts'] else 'N/A'})") for k, v in keymgr.key_status.items()}
     
@@ -716,7 +706,6 @@ def health_check():
 
 if __name__ == "__main__":
     log_format = '[%(asctime)s] %(levelname)s in %(module)s (%(funcName)s): %(message)s'
-    # Use LOG_LEVEL from env if set, otherwise default to INFO
     logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO").upper(), format=log_format)
 
     if not PRIVATE_KEY: