Update app.py

app.py

@@ -2,8 +2,10 @@ import os
 import requests
 from flask import Flask, request, Response, stream_with_context, jsonify
 import json
+import logging
 
 app = Flask(__name__)
+logging.basicConfig(level=logging.INFO)
 
 DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions"
 API_KEY = os.environ.get("API_KEY")
@@ -20,12 +22,11 @@ def chat_completions():
     if not authenticate():
         return jsonify({"error": "Unauthorized"}), 401
 
-    # 获取OpenAI格式的请求
     openai_request = request.json
+    logging.info(f"Received request: {openai_request}")
 
-    # 转换为DeepInfra格式
     deepinfra_request = {
-        "model": openai_request.get("model", "meta-llama/Meta-Llama-3.1-70B-Instruct"),
+        "model": openai_request.get("model", "meta-llama/Meta-Llama-3.1-405B-Instruct"),
         "temperature": openai_request.get("temperature", 0.7),
         "max_tokens": openai_request.get("max_tokens", 1000),
         "stream": openai_request.get("stream", False),
@@ -37,37 +38,88 @@ def chat_completions():
         "Accept": "text/event-stream" if deepinfra_request["stream"] else "application/json"
     }
 
-    # 发送请求到DeepInfra API
-    response = requests.post(DEEPINFRA_API_URL, json=deepinfra_request, headers=headers, stream=True)
+    try:
+        response = requests.post(DEEPINFRA_API_URL, json=deepinfra_request, headers=headers, stream=deepinfra_request["stream"])
+        response.raise_for_status()
+    except requests.RequestException as e:
+        logging.error(f"Error calling DeepInfra API: {str(e)}")
+        return jsonify({"error": "Failed to call DeepInfra API"}), 500
 
     if deepinfra_request["stream"]:
-        # 流式响应
         def generate():
+            full_content = ""
             for line in response.iter_lines():
                 if line:
-                    yield f"data: {line.decode('utf-8')}\n\n"
+                    data = json.loads(line.decode('utf-8').split('data: ')[1])
+                    if data == "[DONE]":
+                        yield f"data: [DONE]\n\n"
+                        break
+                    
+                    delta_content = data['choices'][0]['delta'].get('content', '')
+                    full_content += delta_content
+                    
+                    openai_format = {
+                        "id": data['id'],
+                        "object": "chat.completion.chunk",
+                        "created": data['created'],
+                        "model": data['model'],
+                        "choices": [
+                            {
+                                "index": 0,
+                                "delta": {
+                                    "content": delta_content
+                                },
+                                "finish_reason": data['choices'][0].get('finish_reason')
+                            }
+                        ]
+                    }
+                    yield f"data: {json.dumps(openai_format)}\n\n"
+            
+            # Send the final usage information
+            if 'usage' in data:
+                final_chunk = {
+                    "id": data['id'],
+                    "object": "chat.completion.chunk",
+                    "created": data['created'],
+                    "model": data['model'],
+                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+                    "usage": data['usage']
+                }
+                yield f"data: {json.dumps(final_chunk)}\n\n"
+
         return Response(stream_with_context(generate()), content_type='text/event-stream')
     else:
-        # 非流式响应
-        deepinfra_response = response.json()
-        openai_response = {
-            "id": deepinfra_response.get("id", ""),
-            "object": "chat.completion",
-            "created": deepinfra_response.get("created", 0),
-            "model": deepinfra_response.get("model", ""),
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": deepinfra_response["choices"][0]["message"]["content"]
-                    },
-                    "finish_reason": deepinfra_response["choices"][0].get("finish_reason", "stop")
-                }
-            ],
-            "usage": deepinfra_response.get("usage", {})
-        }
-        return json.dumps(openai_response), 200, {'Content-Type': 'application/json'}
+        try:
+            deepinfra_response = response.json()
+            logging.info(f"Received response from DeepInfra: {deepinfra_response}")
+
+            if 'error' in deepinfra_response:
+                return jsonify({"error": deepinfra_response['error']}), 400
+
+            if 'choices' not in deepinfra_response or not deepinfra_response['choices']:
+                return jsonify({"error": "Unexpected response format from DeepInfra"}), 500
+
+            openai_response = {
+                "id": deepinfra_response.get("id", ""),
+                "object": "chat.completion",
+                "created": deepinfra_response.get("created", 0),
+                "model": deepinfra_response.get("model", ""),
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": deepinfra_response["choices"][0]["message"]["content"]
+                        },
+                        "finish_reason": deepinfra_response["choices"][0].get("finish_reason", "stop")
+                    }
+                ],
+                "usage": deepinfra_response.get("usage", {})
+            }
+            return json.dumps(openai_response), 200, {'Content-Type': 'application/json'}
+        except Exception as e:
+            logging.error(f"Error processing DeepInfra response: {str(e)}")
+            return jsonify({"error": "Failed to process DeepInfra response"}), 500
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)