|
from flask import Flask, request, Response |
|
import requests |
|
import json |
|
import os |
|
import sys |
|
import logging |
|
from helper import create_jwt |
|
|
|
app = Flask(__name__) |
|
|
|
|
|
logging.basicConfig( |
|
level=logging.INFO, |
|
format='%(asctime)s [%(levelname)s] %(message)s', |
|
handlers=[ |
|
logging.StreamHandler(sys.stdout) |
|
] |
|
) |
|
logger = logging.getLogger(__name__) |
|
|
|
@app.route('/ai/v1/chat/completions', methods=['POST']) |
|
async def chat(): |
|
""" |
|
Handle chat completion requests. |
|
""" |
|
logger.info("Received chat completion request") |
|
|
|
|
|
logger.info(f"Request method: {request.method}") |
|
logger.info(f"Request URL: {request.url}") |
|
logger.info(f"Request headers: {dict(request.headers)}") |
|
|
|
|
|
payload = request.json |
|
logger.info(f"Request payload: {json.dumps(payload, indent=2)}") |
|
|
|
|
|
model = payload.get('model', 'claude-3-5-sonnet-20240620') |
|
logger.info(f"Using model: {model}") |
|
|
|
|
|
auth_header = request.headers.get('Authorization') |
|
if not auth_header or not auth_header.startswith('Bearer '): |
|
logger.error("Invalid Authorization header") |
|
return Response('Invalid Authorization header', status=401) |
|
|
|
try: |
|
github_username, zed_user_id = auth_header[7:].split(',') |
|
logger.info(f"GitHub username: {github_username}, Zed user ID: {zed_user_id}") |
|
except ValueError: |
|
logger.error("Invalid Authorization header format") |
|
return Response('Invalid Authorization header format', status=401) |
|
|
|
|
|
url = "https://llm.zed.dev/completion" |
|
logger.info(f"LLM API URL: {url}") |
|
|
|
llm_payload = { |
|
"provider": "anthropic", |
|
"model": model, |
|
"provider_request": { |
|
"model": model, |
|
"max_tokens": payload.get('max_tokens', 8192), |
|
"temperature": payload.get('temperature', 0), |
|
"top_p": payload.get('top_p', 0.7), |
|
"messages": payload['messages'], |
|
"system": "" |
|
} |
|
} |
|
logger.info(f"LLM API payload: {json.dumps(llm_payload, indent=2)}") |
|
|
|
jwt = create_jwt(github_username, int(zed_user_id)) |
|
logger.info(f"Generated JWT token: {jwt}") |
|
|
|
headers = { |
|
'Host': 'llm.zed.dev', |
|
'accept': '*/*', |
|
'content-type': 'application/json', |
|
'authorization': f'Bearer {jwt}', |
|
'user-agent': 'Zed/0.149.3 (macos; aarch64)' |
|
} |
|
logger.info(f"Request headers: {json.dumps(headers, indent=2)}") |
|
|
|
|
|
proxy = os.environ.get('HTTP_PROXY', None) |
|
proxies = {'http': proxy, 'https': proxy} if proxy else None |
|
logger.info(f"Using proxy: {proxy}") |
|
|
|
def generate(): |
|
try: |
|
logger.info("Sending request to LLM API") |
|
with requests.post(url, headers=headers, json=llm_payload, stream=True, proxies=proxies, allow_redirects=True) as response: |
|
logger.info(f"LLM API response status: {response.status_code}") |
|
logger.info(f"LLM API response headers: {dict(response.headers)}") |
|
|
|
if response.status_code == 301: |
|
new_location = response.headers.get('Location') |
|
logger.warning(f"Received 301 redirect. New location: {new_location}") |
|
|
|
|
|
for chunk in response.iter_content(chunk_size=1024): |
|
if chunk: |
|
logger.debug(f"Received chunk of size: {len(chunk)} bytes") |
|
yield chunk |
|
except Exception as e: |
|
logger.error(f"Error during API request: {str(e)}") |
|
yield str(e).encode() |
|
|
|
logger.info("Returning streaming response") |
|
return Response(generate(), content_type='application/octet-stream') |
|
|
|
@app.route('/', methods=['GET']) |
|
def home(): |
|
return "Welcome to the Chat Completion API", 200 |
|
|
|
if __name__ == '__main__': |
|
logger.info("Starting the application") |
|
app.run(host="0.0.0.0", port=8000) |
|
|