APIForwarder / app.py
tianlong12's picture
Update app.py
4965971 verified
raw
history blame
3.1 kB
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!"
@app.route('/hf/v1/chat/completions', methods=['POST'])
def chat():
try:
# 验证请求头中的API密钥
auth_header = request.headers.get('Authorization')
if not auth_header or not auth_header.startswith('Bearer '):
return jsonify({"error": "Unauthorized"}), 401
api_key = auth_header.split(" ")[1]
base_url= auth_header.split(" ")[2]
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
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
# 创建每个请求的 OpenAI 客户端实例
client = OpenAI(
api_key=api_key,
base_url=base_url,
)
# 处理模型响应
if stream:
# 处理流式响应
def generate():
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(stream_with_context(generate()), content_type='text/event-stream')
else:
# 处理非流式响应
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)