Update app.py

app.py

@@ -3,7 +3,7 @@ from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 import requests
 import json
-import asyncio
+from typing import AsyncIterator
 
 app = FastAPI()
 
@@ -35,40 +35,39 @@ async def chat(request: ChatRequest):
         "max_tokens": request.max_tokens
     }
 
-    # Function to stream the response
-    async def event_stream():
-        try:
-            with requests.post(url, headers=headers, data=json.dumps(payload), stream=True) as response:
-                # Check if the request was successful
-                if response.status_code == 200:
-                    # Stream the response
-                    for line in response.iter_lines():
-                        if line:
-                            # Decode the line
-                            decoded_line = line.decode('utf-8')
+    # Make the POST request with streaming
+    try:
+        response = requests.post(url, headers=headers, data=json.dumps(payload), stream=True)
 
-                            # Check if the line starts with "data: "
-                            if decoded_line.startswith("data: "):
-                                # Extract the JSON part
-                                json_data = decoded_line[6:]  # Remove "data: " prefix
+        # Check if the request was successful
+        if response.status_code == 200:
+            async def event_stream() -> AsyncIterator[str]:
+                # Stream the response line by line
+                for line in response.iter_lines():
+                    if line:
+                        # Decode the line
+                        decoded_line = line.decode('utf-8')
+                        # Check if the line starts with "data: "
+                        if decoded_line.startswith("data: "):
+                            # Extract the JSON part
+                            json_data = decoded_line[6:]  # Remove "data: " prefix
+                            # Parse the JSON
+                            try:
+                                parsed_data = json.loads(json_data)
+                                # Check if 'choices' and 'delta' exist and yield the content
+                                if 'choices' in parsed_data and len(parsed_data['choices']) > 0:
+                                    content = parsed_data['choices'][0]['delta'].get('content', '')
+                                    if content:
+                                        yield content
+                            except json.JSONDecodeError:
+                                continue
 
-                                # Parse the JSON
-                                try:
-                                    parsed_data = json.loads(json_data)
-                                    # Check if 'choices' and 'delta' exist and yield the content
-                                    if 'choices' in parsed_data and len(parsed_data['choices']) > 0:
-                                        content = parsed_data['choices'][0]['delta'].get('content', '')
-                                        if content:
-                                            yield content
-                                except json.JSONDecodeError:
-                                    continue
-                else:
-                    raise HTTPException(status_code=response.status_code, detail=response.text)
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-
-    return StreamingResponse(event_stream(), media_type="text/event-stream")
+            return StreamingResponse(event_stream(), media_type="text/event-stream")
+        else:
+            raise HTTPException(status_code=response.status_code, detail=response.text)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8083)
+    uvicorn.run(app, host="0.0.0.0", port=8083)