Update app.py

app.py

@@ -1,8 +1,9 @@
-from flask import Flask, request, Response, jsonify
-import requests
+from flask import Flask, request, Response, stream_with_context, jsonify
+from openai import OpenAI
+import json
 
 app = Flask(__name__)
-
+        
 @app.route('/')
 def index():
     return "Hello, this is the root page of your Flask application!"
@@ -10,57 +11,73 @@ def index():
 @app.route('/hf/v1/chat/completions', methods=['POST'])
 def chat():
     try:
-        # 从请求头中获取 API 密钥和基础 URL
+        # 验证请求头中的API密钥
         auth_header = request.headers.get('Authorization')
         if not auth_header or not auth_header.startswith('Bearer '):
-            return Response("Unauthorized", status=401)
+            return jsonify({"error": "Unauthorized"}), 401
         
-        api_key, base_url = auth_header.split(" ")[1:]
-
-        # 构建目标 URL
-        target_url = f"{base_url}/v1/chat/completions"
+        api_key = auth_header.split(" ")[1]
+        base_url= auth_header.split(" ")[2]
 
-        # 准备请求头
-        headers = {
-            'Authorization': f'Bearer {api_key}',
-            'Content-Type': 'application/json'
-        }
+        data = request.json
+        #print("Received data:", data)  # 打印请求体以进行调试
+        
+        # 验证请求格式
+        if not data or 'messages' not in data or 'model' not in data:
+            return jsonify({"error": "Missing 'messages' or 'model' in request body"}), 400
 
-        # 检查是否为流式请求
-        is_stream = request.json.get('stream', False)
+        model = data['model']
+        messages = data['messages']
+        temperature = data.get('temperature', 0.7)  # 默认值0.7
+        #max_tokens = calculate_max_tokens(model, messages, requested_max_tokens)
+        top_p = data.get('top_p', 1.0)              # 默认值1.0
+        n = data.get('n', 1)                        # 默认值1
+        stream = data.get('stream', False)          # 默认值False
+        functions = data.get('functions', None)     # Functions for function calling
+        function_call = data.get('function_call', None)  # Specific function call request
 
-        # 转发请求
-        response = requests.post(
-            target_url,
-            headers=headers,
-            json=request.json,
-            stream=is_stream
+        # 创建每个请求的 OpenAI 客户端实例
+        client = OpenAI(
+            api_key=api_key,
+            base_url=base_url,
         )
-        print(target_url,api_key)
 
-        if is_stream:
-            # 流式响应
+        # 处理模型响应
+        if stream:
+            # 处理流式响应
             def generate():
-                for chunk in response.iter_content(chunk_size=1024):
-                    yield chunk
+                    response = client.chat.completions.create(
+                        model=model,
+                        messages=messages,
+                        temperature=temperature,
+                        #max_tokens=max_tokens,
+                        top_p=top_p,
+                        n=n,
+                        stream=True,
+                        functions=functions,
+                        function_call=function_call
+                    )
+                    for chunk in response:
+                        yield f"data: {json.dumps(chunk.to_dict())}\n\n"
 
-            return Response(
-                generate(),
-                content_type=response.headers['Content-Type'],
-                status=response.status_code
-            )
+            return Response(stream_with_context(generate()), content_type='text/event-stream')
         else:
-            # 非流式响应
-            return Response(
-                response.content,
-                content_type=response.headers['Content-Type'],
-                status=response.status_code
+            # 处理非流式响应
+            response = client.chat.completions.create(
+                model=model,
+                messages=messages,
+                temperature=temperature,
+                #max_tokens=max_tokens,
+                top_p=top_p,
+                n=n,
+                functions=functions,
+                function_call=function_call,
             )
+            return jsonify(response.to_dict())
 
     except Exception as e:
         print("Exception:", e)
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=7860, threaded=True)
-
+    app.run(host='0.0.0.0', port=7860, threaded=True)