Update app.py

app.py

@@ -1,23 +1,18 @@
-from flask import Flask,jsonify,request,Response, stream_with_context
+from flask import Flask, jsonify, request, Response, stream_with_context
 import requests
 import os
 import logging
 from dotenv import load_dotenv
 
-# 加载.env文件中的环境变量
 load_dotenv()
 
+logging.basicConfig(level=logging.INFO) # 统一使用 logging 模块
 
-# 配置日志输出（可调整级别为DEBUG以获得更详细日志）
-logging.basicConfig(level=logging.INFO)
-
-API_BASE_URL = os.getenv("API_BASE_URL","")
-
-AUTH_KEY = os.getenv("AUTH_KEY","callme")
-
+API_BASE_URL = os.getenv("API_BASE_URL", "")
+AUTH_KEY = os.getenv("AUTH_KEY", "callme")
 API_KEYS = os.getenv("API_KEYS")
 
-API_KEY_LIST = API_KEYS.split(",")
+API_KEY_LIST = API_KEYS.split(",") if API_KEYS else [] # 确保 API_KEYS 为空时不会出错
 
 key_index = 0
 
@@ -25,95 +20,92 @@ app = Flask(__name__)
 
 def getAPI_KEY():
     global key_index
-    
-    #先获取当前的 key
+    if not API_KEY_LIST:
+        logging.warning("API_KEYS is not configured.")
+        return "" # 或者抛出异常
     key = API_KEY_LIST[key_index]
-    
-    #再计算下一个 index，使用取模运算来自动循环
     key_index = (key_index + 1) % len(API_KEY_LIST)
-    
     return key
 
 @app.before_request
 def check_api_key():
     key = request.headers.get("Authorization")
-    if key != "Bearer "+AUTH_KEY:
-        return jsonify({"success":False,"message": "Unauthorized: Invalid API key"}), 403
+    if key != "Bearer " + AUTH_KEY:
+        return jsonify({"success": False, "message": "Unauthorized: Invalid API key"}), 403
 
-@app.route("/v1/models",methods=['GET'])
+@app.route("/v1/models", methods=['GET'])
 def getModels():
     api_key = getAPI_KEY()
-    print(api_key)
+    if not api_key:
+        return jsonify({"success": False, "message": "API Key not available"}), 500
+
     headers = {
-        "Authorization":"Bearer "+api_key,
-        "Content-Type":"application/json"
+        "Authorization": "Bearer " + api_key,
+        "Content-Type": "application/json"
     }
     try:
-        response = requests.get(API_BASE_URL+"/models",headers=headers)
+        response = requests.get(API_BASE_URL + "/models", headers=headers, timeout=30) # 增加超时
         response.raise_for_status()
         response_headers = {'Content-Type': response.headers.get('content-type', 'application/json')}
+        # 直接返回内容和头部，这里不需要流式处理
         return (response.content, response.status_code, response_headers)
+    except requests.exceptions.RequestException as e: # 捕获requests特有异常
+        logging.error("Get models error. %s", e)
+        return jsonify({"success": False, "message": f"Failed to fetch models: {e}"}), 500
     except Exception as e:
-        logging("get models error. %s" ,e)
+        logging.error("An unexpected error occurred in getModels: %s", e)
+        return jsonify({"success": False, "message": str(e)}), 500
 
-@app.route("/v1/chat/completions",methods=['POST'])
+@app.route("/v1/chat/completions", methods=['POST'])
 def chat():
+    api_key = getAPI_KEY()
+    if not api_key:
+        return jsonify({"success": False, "message": "API Key not available"}), 500
+
     headers = {
-        "Authorization":"Bearer "+getAPI_KEY(),
-        "Content-Type":"application/json"
+        "Authorization": "Bearer " + api_key,
+        "Content-Type": "application/json"
     }
     data = request.get_json()
-    stream_flag = data.get('stream', False)
-
-    def generate():
-        try:
-            with requests.post(API_BASE_URL+"/chat/completions", headers=headers, json=data, stream=stream_flag) as response:
-                response.raise_for_status() # 检查上游请求是否成功
-                for chunk in response.iter_content(chunk_size=1024):
-                    yield chunk
-        except requests.exceptions.RequestException as e:
-            logging.error("Request to upstream API failed: %s", e)
-            # 在这里处理上游请求失败，例如可以 yield 一个错误消息或者抛出异常
-            # 但请注意，一旦开始 yield 数据，就不能改变 HTTP 状态码和头部了
-            yield b'{"error": "Upstream API request failed"}' # 作为 JSON 错误返回
-        except Exception as e:
-            logging.error("Unexpected error during streaming: %s", e)
-            yield b'{"error": "Internal server error during streaming"}'
 
+    # 重点在这里：根据客户端请求中的 'stream' 字段决定是否进行流式转发
+    # 如果客户端没有提供 'stream' 字段，我们假设它需要非流式响应（或者默认值取决于上游API的约定）
+    # 但为了明确支持非流式，这里我们设为 False
+    client_wants_stream = data.get('stream', False) # 客户端请求中 stream 的值
 
     try:
-        # 如果不是流式请求，可以考虑不使用生成器，或者根据 stream_flag 来判断
-        if not stream_flag:
-            # 对于非流式请求，直接返回完整响应
-            response = requests.post(API_BASE_URL+"/chat/completions", headers=headers, json=data, stream=False)
-            response.raise_for_status()
-            return Response(response.content,
-                            status=response.status_code,
-                            content_type=response.headers.get('content-type'))
+        # 使用 requests.post 发送请求到上游 API
+        # 上游 API 的 'stream' 参数应该与客户端请求的 'stream' 字段保持一致
+        upstream_response = requests.post(
+            API_BASE_URL + "/chat/completions",
+            headers=headers,
+            json=data, # 客户端请求的 payload，包括 stream 字段
+            stream=client_wants_stream, # 控制 requests 是否以流式接收上游响应
+            timeout= (600 if client_wants_stream else 60) # 流式请求可以有更长的超时
+        )
+        upstream_response.raise_for_status() # 检查上游 API 响应的 HTTP 状态码
+
+        # 根据客户端是否想要流式响应来处理
+        if client_wants_stream:
+            # 流式响应：使用 stream_with_context 逐块发送
+            return Response(stream_with_context(upstream_response.iter_content(chunk_size=1024)),
+                            status=upstream_response.status_code,
+                            content_type=upstream_response.headers.get('content-type', 'application/json'))
         else:
-            # 对于流式请求，使用生成器
-            # 注意：在生成器中处理异常时，如果已经开始发送数据，状态码和头部就不能更改了。
-            # 所以最好是在生成器开始之前捕获requests.post的异常。
-            initial_response = requests.post(API_BASE_URL+"/chat/completions", headers=headers, json=data, stream=True)
-            initial_response.raise_for_status() # 检查初始请求是否成功
-
-            return Response(generate_from_response(initial_response),
-                            status=initial_response.status_code,
-                            content_type=initial_response.headers.get('content-type'))
+            # 非流式响应：直接返回完整的响应内容
+            # 这里 upstream_response.content 会等待所有数据接收完毕
+            response_headers = {'Content-Type': upstream_response.headers.get('content-type', 'application/json')}
+            return (upstream_response.content, upstream_response.status_code, response_headers)
 
     except requests.exceptions.RequestException as e:
-        logging.error("Initial upstream API request failed: %s", e)
+        logging.error("Chat completion request error to upstream API: %s", e)
         return jsonify({"success": False, "message": f"Upstream API request failed: {e}"}), 500
     except Exception as e:
-        logging.error("Error setting up chat completion: %s", e)
+        logging.error("An unexpected error occurred in chat completion: %s", e)
         return jsonify({"success": False, "message": str(e)}), 500
 
-def generate_from_response(upstream_response):
-    # 这是一个辅助函数，用于将上游响应的迭代器包装成一个生成器
-    for chunk in upstream_response.iter_content(chunk_size=1024):
-        yield chunk
-
 
 if __name__ == '__main__':
-    print("test")
+    print("Starting Flask app...")
+    # 在生产环境，不建议使用 debug=True，且应通过 Gunicorn 等 WSGI 服务器运行
     app.run(host='0.0.0.0', port=7860, debug=True)