Update app.py

app.py

@@ -1,33 +1,30 @@
-from asgiref.wsgi import WsgiToAsgi
-from flask import Flask, request, Response
-import requests
+from fastapi import FastAPI, Request, Response
+from fastapi.responses import StreamingResponse
+import httpx
 import json
 import random
 import os
 
 from helper import create_jwt, generate_random_tuple
 
-app = Flask(__name__)
+app = FastAPI()
 
-@app.route('/ai/v1/chat/completions', methods=['POST'])
-async def chat():
+@app.post('/ai/v1/chat/completions')
+async def chat(request: Request):
     """
     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.
-
+        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.
     """
     # Get the payload from the request
-    payload = request.json
+    payload = await request.json()
 
     # Get the model from the payload, defaulting to "claude-3-5-sonnet-20240620"
     model = payload.get('model', 'claude-3-5-sonnet-20240620')
@@ -48,7 +45,6 @@ async def chat():
         }
     }
 
-    # github_username, zed_user_id = random.choice(github_username_zed_userid_list)
     github_username, zed_user_id = generate_random_tuple()
     jwt = create_jwt(github_username, zed_user_id)
 
@@ -65,16 +61,13 @@ async def chat():
     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:
+        async with httpx.AsyncClient(proxies=proxies) as client:
+            async with client.stream('POST', url, headers=headers, json=llm_payload) as response:
+                async for chunk in response.aiter_bytes():
                     yield chunk
 
-    return Response(generate(), content_type='application/octet-stream')
-
-# Convert the Flask app to an ASGI app
-asgi_app = WsgiToAsgi(app)
+    return StreamingResponse(generate(), media_type='application/octet-stream')
 
 if __name__ == '__main__':
     import uvicorn
-    uvicorn.run(asgi_app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)