Update app.py

app.py

@@ -1,7 +1,9 @@
 from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 import requests
 import json
+import asyncio
 
 app = FastAPI()
 
@@ -33,38 +35,39 @@ async def chat(request: ChatRequest):
         "max_tokens": request.max_tokens
     }
 
-    # Make the POST request with streaming
-    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:
-                result = ""
-                # Stream the response
-                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 append the content
-                                if 'choices' in parsed_data and len(parsed_data['choices']) > 0:
-                                    content = parsed_data['choices'][0]['delta'].get('content', '')
-                                    if content:
-                                        result += content
-                            except json.JSONDecodeError:
-                                continue
-                return {"response": result}
-            else:
-                raise HTTPException(status_code=response.status_code, detail=response.text)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    # 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')
+
+                            # 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
+                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")
 
 if __name__ == "__main__":
     import uvicorn