Update app.py

app.py

@@ -2,27 +2,17 @@ from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse
 import httpx
 import json
-import random
 import os
-
 from helper import create_jwt, generate_random_tuple
 
 app = FastAPI()
 
 @app.post('/ai/v1/chat/completions')
 async def chat(request: Request):
-    """
-    Handle chat completion requests.
-    This function processes incoming POST requests to the '/ai/v1/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:
-        StreamingResponse: 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.
-    """
+    # Generate JWT token
+    github_username, user_id = generate_random_tuple()
+    jwt_token = create_jwt(github_username, user_id)
+
     # Get the payload from the request
     payload = await request.json()
 
@@ -45,14 +35,11 @@ async def chat(request: Request):
         }
     }
 
-    github_username, zed_user_id = generate_random_tuple()
-    jwt = create_jwt(github_username, zed_user_id)
-
     headers = {
         'Host': 'llm.zed.dev',
         'accept': '*/*',
         'content-type': 'application/json',
-        'authorization': f'Bearer {jwt}',
+        'authorization': f'Bearer {jwt_token}',  # Use the generated JWT token
         'user-agent': 'Zed/0.149.3 (macos; aarch64)'
     }
 
@@ -68,10 +55,6 @@ async def chat(request: Request):
 
     return StreamingResponse(generate(), media_type='application/octet-stream')
 
-@app.get("/")
-async def root():
-    return {"message": "Welcome to the AI Chat Completions API"}
-
 if __name__ == '__main__':
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)