|
|
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.DEBUG) |
|
|
|
|
|
DEEPINFRA_API_URL = "https://api.deepinfra.com/v1/openai/chat/completions" |
|
|
API_KEY = os.environ.get("API_KEY") |
|
|
|
|
|
def authenticate(): |
|
|
auth_header = request.headers.get("Authorization") |
|
|
if not auth_header or not auth_header.startswith("Bearer "): |
|
|
return False |
|
|
token = auth_header.split(" ")[1] |
|
|
return token == API_KEY |
|
|
|
|
|
@app.route('/hf/v1/chat/completions', methods=['POST']) |
|
|
def chat_completions(): |
|
|
logging.debug(f"Received headers: {request.headers}") |
|
|
logging.debug(f"Received body: {request.get_data(as_text=True)}") |
|
|
|
|
|
if not authenticate(): |
|
|
logging.warning("Unauthorized access attempt") |
|
|
return jsonify({"error": "Unauthorized"}), 401 |
|
|
|
|
|
try: |
|
|
openai_request = request.json |
|
|
except json.JSONDecodeError: |
|
|
logging.error("Invalid JSON in request body") |
|
|
return jsonify({"error": "Invalid JSON in request body"}), 400 |
|
|
|
|
|
logging.info(f"Received request: {openai_request}") |
|
|
|
|
|
deepinfra_request = { |
|
|
"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), |
|
|
"messages": openai_request.get("messages", []) |
|
|
} |
|
|
|
|
|
headers = { |
|
|
"Content-Type": "application/json", |
|
|
"Accept": "text/event-stream" if deepinfra_request["stream"] else "application/json" |
|
|
} |
|
|
|
|
|
try: |
|
|
response = requests.post(DEEPINFRA_API_URL, json=deepinfra_request, headers=headers, stream=deepinfra_request["stream"]) |
|
|
response.raise_for_status() |
|
|
logging.debug(f"DeepInfra API response status: {response.status_code}") |
|
|
logging.debug(f"DeepInfra API response headers: {response.headers}") |
|
|
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 not line: |
|
|
logging.warning("Received empty line from DeepInfra API") |
|
|
continue |
|
|
try: |
|
|
line_text = line.decode('utf-8') |
|
|
if line_text.startswith('data: '): |
|
|
data_text = line_text.split('data: ', 1)[1] |
|
|
if data_text == "[DONE]": |
|
|
yield f"data: [DONE]\n\n" |
|
|
break |
|
|
data = json.loads(data_text) |
|
|
|
|
|
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" |
|
|
except json.JSONDecodeError as e: |
|
|
logging.error(f"JSON decode error: {e}. Raw line: {line}") |
|
|
continue |
|
|
except Exception as e: |
|
|
logging.error(f"Error processing line: {e}. Raw line: {line}") |
|
|
continue |
|
|
|
|
|
|
|
|
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: |
|
|
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 |
|
|
|
|
|
@app.route('/') |
|
|
def home(): |
|
|
return "Welcome to the API proxy server. Please use the /hf/v1/chat/completions endpoint for chat completions." |
|
|
|
|
|
if __name__ == '__main__': |
|
|
app.run(host='0.0.0.0', port=7860) |
