Update app.py

app.py

@@ -1,60 +1,37 @@
+from asgiref.wsgi import WsgiToAsgi
 from flask import Flask, request, Response
 import requests
 import json
-import os
-import sys
-import logging
-from helper import create_jwt
+import random
+from helper import create_jwt, get_github_username_zed_userid_list
 
 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('/chat/completions', methods=['POST'])
-def chat():
+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.
+        - It generates random GitHub username and Zed user ID for each request.
+        - The LLM model defaults to "claude-3-5-sonnet-20240620" if not specified.
     """
-    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: {json.dumps(dict(request.headers), indent=2)}")
-    
     # Get the payload from the request
     payload = request.json
-    logger.info(f"Request payload: {json.dumps(payload, indent=2)}")
 
     # Get the model from the payload, defaulting to "claude-3-5-sonnet-20240620"
     model = payload.get('model', 'claude-3-5-sonnet-20240620')
-    logger.info(f"Using model: {model}")
-
-    # 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 '):
-        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"Extracted GitHub username: {github_username}")
-        logger.info(f"Extracted Zed user ID: {zed_user_id}")
-    except ValueError:
-        logger.error("Invalid Authorization header format")
-        return Response('Invalid Authorization header format', status=401)
 
     # Prepare the request for the LLM API
-    url = "https://llm.zed.dev/completion"
-    logger.info(f"LLM API URL: {url}")
+    url = "https://llm.zed.dev/completion?"
     
     llm_payload = {
         "provider": "anthropic",
@@ -68,10 +45,9 @@ def chat():
             "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}")
+    github_username, zed_user_id = get_github_username_zed_userid_list()
+    jwt = create_jwt(github_username, zed_user_id)
 
     headers = {
         'Host': 'llm.zed.dev',
@@ -80,40 +56,22 @@ def chat():
         'authorization': f'Bearer {jwt}',
         'user-agent': 'Zed/0.149.3 (macos; aarch64)'
     }
-    logger.info(f"Request headers for LLM API: {json.dumps(headers, indent=2)}")
 
     # Get proxy from environment variable
     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: {json.dumps(dict(response.headers), indent=2)}")
-                
-                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()
+    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
 
-    logger.info("Returning streaming response")
     return Response(generate(), content_type='application/octet-stream')
 
-@app.route('/', methods=['GET'])
-def home():
-    logger.info("Received request to home page")
-    return "Welcome to the Chat Completion API", 200
+# Convert the Flask app to an ASGI app
+asgi_app = WsgiToAsgi(app)
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=8000)
+    import uvicorn
+    uvicorn.run(asgi_app, host="0.0.0.0", port=8000)