Update app.py

app.py

@@ -1,49 +1,48 @@
-import logging
-from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import StreamingResponse
-import httpx
+from asgiref.wsgi import WsgiToAsgi
+from flask import Flask, request, Response
+import requests
 import json
 import os
-
 from helper import create_jwt
 
-# 设置日志
-logging.basicConfig(level=logging.DEBUG)
-logger = logging.getLogger(__name__)
+app = Flask(__name__)
 
-app = FastAPI()
+@app.route('/ai/v1/chat/completions', methods=['POST'])
+async def chat():
+    """
+    Handle chat completion requests.
 
-LLM_API_URL = "https://llm.zed.dev/completion?"
+    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.
 
-@app.post('/ai/v1/chat/completions')
-async def chat(request: Request):
-    logger.debug("Received request")
-    
-    # 获取客户端的 authorization 头
-    auth_header = request.headers.get('authorization')
-    if not auth_header or not auth_header.startswith('Bearer '):
-        raise HTTPException(status_code=401, detail="Invalid authorization header")
-    
-    # 提取 github_username 和 user_id
-    try:
-        _, auth_data = auth_header.split('Bearer ', 1)
-        github_username, user_id_str = auth_data.split(',')
-        user_id = int(user_id_str)
-    except ValueError:
-        raise HTTPException(status_code=401, detail="Invalid authorization format")
-    
-    # 生成 JWT token
-    jwt_token = create_jwt(github_username, user_id)
-    logger.debug(f"Generated JWT token: {jwt_token}")
+    Returns:
+        Response: A streaming response containing the LLM API's output.
 
-    # 获取请求 payload
-    payload = await request.json()
-    logger.debug(f"Received payload: {payload}")
+    Note:
+        - The function uses environment variables for proxy configuration.
+        - It extracts GitHub username and Zed user ID from the Authorization header.
+        - The LLM model defaults to "claude-3-5-sonnet-20240620" if not specified.
+    """
+    # Get the payload from the request
+    payload = request.json
 
-    # 获取模型，默认为 "claude-3-5-sonnet-20240620"
+    # Get the model from the payload, defaulting to "claude-3-5-sonnet-20240620"
     model = payload.get('model', 'claude-3-5-sonnet-20240620')
 
-    # 准备 LLM API 请求
+    # Extract GitHub username and Zed user ID from Authorization header
+    auth_header = request.headers.get('Authorization')
+    if not auth_header or not auth_header.startswith('Bearer '):
+        return Response('Invalid Authorization header', status=401)
+
+    try:
+        github_username, zed_user_id = auth_header[7:].split(',')
+    except ValueError:
+        return Response('Invalid Authorization header format', status=401)
+
+    # Prepare the request for the LLM API
+    url = "https://llm.zed.dev/completion?"
+    
     llm_payload = {
         "provider": "anthropic",
         "model": model,
@@ -56,40 +55,32 @@ async def chat(request: Request):
             "system": ""
         }
     }
-    logger.debug(f"LLM payload: {llm_payload}")
+
+    jwt = create_jwt(github_username, int(zed_user_id))
 
     headers = {
+        'Host': 'llm.zed.dev',
         'accept': '*/*',
         'content-type': 'application/json',
-        'authorization': f'Bearer {jwt_token}',
+        'authorization': f'Bearer {jwt}',
         'user-agent': 'Zed/0.149.3 (macos; aarch64)'
     }
-    logger.debug(f"Request headers: {headers}")
 
-    # 获取代理设置
+    # Get proxy from environment variable
     proxy = os.environ.get('HTTP_PROXY', None)
     proxies = {'http': proxy, 'https': proxy} if proxy else None
-    logger.debug(f"Using proxies: {proxies}")
 
     async def generate():
-        async with httpx.AsyncClient(proxies=proxies) as client:
-            try:
-                async with client.stream('POST', LLM_API_URL, headers=headers, json=llm_payload) as response:
-                    logger.debug(f"LLM API response status: {response.status_code}")
-                    logger.debug(f"LLM API response headers: {response.headers}")
-                    if response.status_code != 200:
-                        error_content = await response.aread()
-                        logger.error(f"LLM API error response: {error_content}")
-                        yield f"Error: {response.status_code} - {error_content.decode()}"
-                    else:
-                        async for chunk in response.aiter_bytes():
-                            yield chunk
-            except Exception as e:
-                logger.error(f"Error during LLM API request: {str(e)}")
-                yield f"Error: {str(e)}"
-
-    return StreamingResponse(generate(), media_type='application/octet-stream')
+        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')
+
+# Convert the Flask app to an ASGI app
+asgi_app = WsgiToAsgi(app)
 
 if __name__ == '__main__':
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(asgi_app, host="0.0.0.0", port=8000)