|
from asgiref.wsgi import WsgiToAsgi |
|
from flask import Flask, request, Response |
|
import requests |
|
import json |
|
import random |
|
import os |
|
import jwt |
|
import uuid |
|
from datetime import datetime, timedelta |
|
import string |
|
|
|
app = Flask(__name__) |
|
|
|
def get_github_username_zed_userid_list(): |
|
""" |
|
Get GitHub username and Zed user ID from the USER environment variable. |
|
|
|
Returns: |
|
list: A list containing a single tuple with the GitHub username (str) and Zed user ID (str). |
|
""" |
|
user_env = os.environ.get('USER', 'default_user,123456') |
|
try: |
|
username, user_id = user_env.split(',') |
|
return [(username.strip(), user_id.strip())] |
|
except ValueError: |
|
print("Warning: Invalid format in USER environment variable. Using default values.") |
|
return [("default_user", "123456")] |
|
|
|
def create_jwt(github_user_login: str, user_id: int) -> str: |
|
""" |
|
Create a JSON Web Token (JWT) for a given GitHub user. |
|
|
|
Args: |
|
github_user_login (str): The GitHub username of the user. |
|
user_id (int): The user's ID. |
|
|
|
Returns: |
|
str: A JWT encoded string containing user information and authentication details. |
|
|
|
Note: |
|
The token has a lifetime of 1 hour and includes the following claims: |
|
- iat: Issued at time |
|
- exp: Expiration time |
|
- jti: Unique token identifier |
|
- userId: User's ID |
|
- githubUserLogin: GitHub username |
|
- isStaff: Boolean indicating staff status (default: False) |
|
- hasLlmClosedBetaFeatureFlag: Boolean for LLM closed beta feature (default: False) |
|
- plan: User's plan (default: "Free") |
|
""" |
|
LLM_TOKEN_LIFETIME = timedelta(hours=1) |
|
now = datetime.utcnow() |
|
|
|
payload = { |
|
"iat": int(now.timestamp()), |
|
"exp": int((now + LLM_TOKEN_LIFETIME).timestamp()), |
|
"jti": str(uuid.uuid4()), |
|
"userId": user_id, |
|
"githubUserLogin": github_user_login, |
|
"isStaff": False, |
|
"hasLlmClosedBetaFeatureFlag": False, |
|
"plan": "Free" |
|
} |
|
|
|
return jwt.encode(payload, 'llm-secret', algorithm='HS256') |
|
|
|
@app.route('/chat/completions', methods=['POST']) |
|
async def chat(): |
|
""" |
|
Handle chat completion requests. |
|
|
|
This function processes incoming POST requests to the '/chat/completions' endpoint. |
|
It prepares the payload for the LLM API, generates a JWT for authentication, |
|
and streams the response from the LLM API back to the client. |
|
|
|
Returns: |
|
Response: A streaming response containing the LLM API's output. |
|
|
|
Note: |
|
- The function uses environment variables for proxy configuration and user data. |
|
- The LLM model defaults to "claude-3-5-sonnet-20240620" if not specified. |
|
""" |
|
|
|
payload = request.json |
|
|
|
|
|
model = payload.get('model', 'claude-3-5-sonnet-20240620') |
|
|
|
|
|
url = "https://llm.zed.dev/completion?" |
|
|
|
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": "" |
|
} |
|
} |
|
|
|
github_username, zed_user_id = get_github_username_zed_userid_list()[0] |
|
jwt = create_jwt(github_username, zed_user_id) |
|
|
|
headers = { |
|
'Host': 'llm.zed.dev', |
|
'accept': '*/*', |
|
'content-type': 'application/json', |
|
'authorization': f'Bearer {jwt}', |
|
'user-agent': 'Zed/0.149.3 (macos; aarch64)' |
|
} |
|
|
|
|
|
proxy = os.environ.get('HTTP_PROXY', None) |
|
proxies = {'http': proxy, 'https': proxy} if proxy else None |
|
|
|
async def generate(): |
|
with requests.post(url, headers=headers, json=llm_payload, stream=True, proxies=proxies) as response: |
|
for chunk in response.iter_content(chunk_size=1024): |
|
if chunk: |
|
yield chunk |
|
|
|
return Response(generate(), content_type='application/octet-stream') |
|
|
|
|
|
asgi_app = WsgiToAsgi(app) |
|
|
|
if __name__ == '__main__': |
|
import uvicorn |
|
uvicorn.run(asgi_app, host="0.0.0.0", port=8000) |
|
|