Update openai_ondemand_adapter.py

openai_ondemand_adapter.py

@@ -238,13 +238,7 @@ def create_session(apikey, external_user_id=None, plugin_ids=None):
 def format_openai_sse_delta(chunk_data_dict):
     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):
-    """
-    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 = {
         "query": query_str,
@@ -259,7 +253,7 @@ def _execute_one_stream_attempt(apikey_for_attempt, session_id_for_attempt, quer
     }
 
     accumulated_text_parts = []
-    api_error_yielded = False # Flag to indicate if an API error was processed and yielded as SSE
+    api_error_yielded = False 
 
     logging.info(f"【流式请求子尝试 {current_attempt_num_logging}】发送到 OnDemand: Session={session_id_for_attempt}, Endpoint={endpoint_id}, Key={keymgr.display_key(apikey_for_attempt)}")
 
@@ -278,7 +272,7 @@ def _execute_one_stream_attempt(apikey_for_attempt, session_id_for_attempt, quer
                 }
                 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
+                return "".join(accumulated_text_parts).strip(), api_error_yielded 
 
             first_chunk_sent = False
             last_line_str = "" 
@@ -295,7 +289,6 @@ 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}")
                         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]:"):
@@ -308,7 +301,7 @@ def _execute_one_stream_attempt(apikey_for_attempt, session_id_for_attempt, quer
                             error_obj = {"message": error_json_str, "type": "on_demand_stream_error_format"}
                         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
+                        return "".join(accumulated_text_parts).strip(), api_error_yielded 
                     
                     else:
                         try:
@@ -337,28 +330,23 @@ 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}]
                                 }
-                                yield format_openai_sse_delta(openai_chunk) # Yield immediately
+                                yield format_openai_sse_delta(openai_chunk) 
                         except json.JSONDecodeError:
                             logging.warning(f"【OnDemand流】无法解析JSON (子尝试 {current_attempt_num_logging}): {data_part[:100]}... Session: {session_id_for_attempt}")
                             continue
             
-            # 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
+            if not api_error_yielded: 
+                if not last_line_str.startswith("data: [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
+        raise 
     
     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
+        api_error_yielded = True 
         error_payload = {
             "error": {"message": f"Unknown error during streaming (Attempt {current_attempt_num_logging}): {str(e)}", "type": "unknown_streaming_error_in_attempt"}
         }
@@ -418,7 +406,6 @@ def chat_completions():
     
     def attempt_ondemand_request_wrapper(current_apikey_from_wrapper, current_session_id_from_wrapper):
         if is_stream_request:
-            # 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'
@@ -446,9 +433,17 @@ def chat_completions():
                 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}")
+                status_code_from_exc = None
+                if hasattr(http_err_outer, 'response') and http_err_outer.response is not None:
+                    status_code_from_exc = http_err_outer.response.status_code
+                
+                logging.warning(f"【请求处理 - Key轮换尝试 {key_retry_count}】HTTP/请求错误。Status: {status_code_from_exc}, 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: 
-                    keymgr.mark_bad(selected_apikey_for_outer_retry)
+                    if status_code_from_exc == 524:
+                        logging.info(f"【KeyManager】Key {keymgr.display_key(selected_apikey_for_outer_retry)} not marked bad due to 524 error.")
+                    else: 
+                        keymgr.mark_bad(selected_apikey_for_outer_retry)
                 
                 if key_retry_count >= max_key_retries:
                     logging.error(f"【请求处理】所有Key轮换尝试均失败。最后错误: {last_exception_for_key_retry}")
@@ -468,28 +463,18 @@ def chat_completions():
         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)
-        # 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 format_openai_sse_delta({"error": {"message": error_message, "type": "proxy_setup_error", "code": "proxy_error_503"}}) # Changed code
                 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 jsonify({"error": error_message, "code": "proxy_error_503"}), 503 # Added code here
 
     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):
-    """
-    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
     
@@ -510,9 +495,16 @@ def handle_stream_request(initial_apikey, initial_session_id, query_str, endpoin
             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): 
-                    keymgr.mark_bad(current_apikey_for_attempt)
+                    status_code_from_exc = None
+                    if hasattr(e_key_session, 'response') and e_key_session.response is not None:
+                        status_code_from_exc = e_key_session.response.status_code
+                    if status_code_from_exc == 524:
+                        logging.info(f"【KeyManager】Key {keymgr.display_key(current_apikey_for_attempt)} not marked bad for 524 error during key/session acquisition for retry.")
+                    else:
+                        keymgr.mark_bad(current_apikey_for_attempt)
+
                 if empty_retry_attempt_num >= max_empty_response_retries:
-                    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 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", "code": "proxy_retry_setup_failed"}})
                     yield "data: [DONE]\n\n"
                     return 
                 time.sleep(1)
@@ -522,8 +514,6 @@ def handle_stream_request(initial_apikey, initial_session_id, query_str, endpoin
         log_attempt_str = f"初始尝试" if empty_retry_attempt_num == 1 else f"空回复重试 {empty_retry_attempt_num-1}"
 
         try:
-            # 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, 
@@ -536,50 +526,52 @@ def handle_stream_request(initial_apikey, initial_session_id, query_str, endpoin
             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}")
-            if current_apikey_for_attempt: # Ensure key is marked bad if one was used
-                 keymgr.mark_bad(current_apikey_for_attempt)
+            log_key_display = keymgr.display_key(current_apikey_for_attempt) if current_apikey_for_attempt else "N/A"
+            status_code_from_exc_stream = None
+            if hasattr(e_req, 'response') and e_req.response is not None:
+                status_code_from_exc_stream = e_req.response.status_code
+            
+            logging.warning(f"【流式请求】({log_attempt_str} using key {log_key_display}) 发生请求级错误: {e_req}, Status: {status_code_from_exc_stream}")
+            
+            if current_apikey_for_attempt:
+                if status_code_from_exc_stream == 524:
+                    logging.info(f"【KeyManager】Key {log_key_display} not marked bad for 524 error during stream attempt.")
+                else:
+                    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"【流式请求】在最后一次空回复重试时发生请求错误。")
-                yield format_openai_sse_delta({"error": {"message": f"Request failed on final empty stream retry attempt: {str(e_req)}", "type": "internal_proxy_error"}})
+                yield format_openai_sse_delta({"error": {"message": f"Request failed on final empty stream retry attempt: {str(e_req)}", "type": "internal_proxy_error", "code": "proxy_final_retry_failed"}})
                 yield "data: [DONE]\n\n"
                 return 
             time.sleep(1)
-            continue # To the next iteration of the empty_retry_attempt_num loop (will try new key/session)
+            continue 
             
-        # 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"}
+                "error": {"message": "返回空回复，请重试", "type": "empty_response_after_retries", "code": "empty_response_please_retry"} # Changed message and code
             })
             yield "data: [DONE]\n\n"
             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 format_openai_sse_delta({"error": {"message": "Unexpected error in stream handling.", "type": "internal_proxy_error", "code":"unexpected_stream_exit"}})
     yield "data: [DONE]\n\n"
 
 
@@ -605,9 +597,15 @@ def handle_non_stream_request(initial_apikey, initial_session_id, query_str, end
             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): 
-                    keymgr.mark_bad(current_apikey_for_attempt)
+                    status_code_from_exc = None
+                    if hasattr(e_key_session, 'response') and e_key_session.response is not None:
+                        status_code_from_exc = e_key_session.response.status_code
+                    if status_code_from_exc == 524:
+                        logging.info(f"【KeyManager】Key {keymgr.display_key(current_apikey_for_attempt)} not marked bad for 524 error during key/session acquisition for non-stream retry.")
+                    else:
+                        keymgr.mark_bad(current_apikey_for_attempt)
                 if empty_retry_attempt_num >= max_empty_response_retries:
-                    return jsonify({"error": f"Failed to get new key/session for final empty response retry. Error: {str(e_key_session)}"}), 503
+                    return jsonify({"error": f"Failed to get new key/session for final empty response retry. Error: {str(e_key_session)}", "code":"proxy_retry_setup_failed"}), 503
                 time.sleep(1)
                 current_apikey_for_attempt = None 
                 continue 
@@ -644,32 +642,43 @@ def handle_non_stream_request(initial_apikey, initial_session_id, query_str, end
                 if empty_retry_attempt_num >= max_empty_response_retries:
                     logging.error(f"【同步请求】达到最大空回复重试次数 ({max_empty_response_retries})。")
                     return jsonify({
-                        "error": f"Model returned an empty response after {max_empty_response_retries} retries.",
+                        "error": "返回空回复，请重试", # Changed message
                         "id": "chatcmpl-" + str(uuid.uuid4())[:12], "object": "chat.completion", "created": int(time.time()),
                         "model": openai_model_name_for_response,
-                        "choices": [{"index": 0, "message": {"role": "assistant", "content": ""}, "finish_reason": "length"}],
-                        "usage": {}
-                    }), 500 
+                        "choices": [{"index": 0, "message": {"role": "assistant", "content": ""}, "finish_reason": "length"}], # or "error" finish_reason
+                        "usage": {},
+                        "code": "empty_response_please_retry" # Added code
+                    }), 503 # Changed status code to 503 for "please retry"
                 logging.info(f"【同步请求】空回复，将在1秒后重试下一个Key。当前总尝试 {empty_retry_attempt_num}/{max_empty_response_retries}")
                 time.sleep(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}")
-            if current_apikey_for_attempt: # Ensure key is marked bad
-                keymgr.mark_bad(current_apikey_for_attempt)
+            log_key_display_sync = keymgr.display_key(current_apikey_for_attempt) if current_apikey_for_attempt else "N/A"
+            status_code_from_exc_sync = None
+            if hasattr(e_req, 'response') and e_req.response is not None:
+                status_code_from_exc_sync = e_req.response.status_code
+            
+            logging.warning(f"【同步请求】({log_attempt_str} using key {log_key_display_sync}) 发生请求级错误: {e_req}, Status: {status_code_from_exc_sync}")
+            
+            if current_apikey_for_attempt:
+                if status_code_from_exc_sync == 524:
+                    logging.info(f"【KeyManager】Key {log_key_display_sync} not marked bad for 524 error during sync attempt.")
+                else:
+                    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:
                 logging.error(f"【同步请求】在最后一次空回复重试时发生请求错误。")
-                return jsonify({"error": f"Request failed on final empty response retry attempt. Last error: {str(e_req)}"}), 503
+                return jsonify({"error": f"Request failed on final empty response retry attempt. Last error: {str(e_req)}", "code":"proxy_final_retry_failed"}), 503
             time.sleep(1)
             continue 
         except (ValueError, KeyError, json.JSONDecodeError) as e_parse: 
             logging.error(f"【同步请求】({log_attempt_str}) 处理响应时出错: {e_parse}", exc_info=True)
-            return jsonify({"error": f"Error processing OnDemand sync response: {str(e_parse)}"}), 502 
+            return jsonify({"error": f"Error processing OnDemand sync response: {str(e_parse)}", "code": "on_demand_parse_error"}), 502 
     
     logging.error(f"【同步请求】意外退出空回复重试循环。") 
-    return jsonify({"error": "Unexpected error in non-stream handling after empty response retries."}), 500
+    return jsonify({"error": "返回空回复，请重试", "code": "unexpected_empty_retry_exit_please_retry"}), 503 # Changed message and code, and status
 
 
 @app.route("/v1/models", methods=["GET"])