Update app.py

app.py

@@ -1,5 +1,6 @@
 from fastapi import HTTPException, FastAPI, File, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
+from asyncio import Queue, create_task
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
 import aiofiles
@@ -100,6 +101,8 @@ async def upload_file(file: UploadFile = File(...)):
         if tmp_path and os.path.exists(tmp_path):
             os.unlink(tmp_path)
 
+from asyncio import Queue, create_task
+
 @app.post("/chat")
 async def chat(chat_req: ChatRequest):
     hf_token = os.environ.get("HF_TOKEN")
@@ -121,8 +124,7 @@ async def chat(chat_req: ChatRequest):
     
     if not context_chunks:
         async def empty_stream():
-            ymsg = "I couldn't find relevant information in the document."
-            yield f"data: {json.dumps({'token': ymsg})}\n\n"
+            yield f"data: {json.dumps({'token': 'I couldn't find relevant information in the document.'})}\n\n"
             yield "data: [DONE]\n\n"
         return StreamingResponse(empty_stream(), media_type="text/event-stream")
     
@@ -132,47 +134,51 @@ async def chat(chat_req: ChatRequest):
     messages.extend(chat_req.history)
     messages.append({"role": "user", "content": f"Context:\n{context_text}\n\nQuestion: {chat_req.message}"})
 
-    def generate_tokens():
-        """Synchronous generator that yields SSE formatted tokens"""
-        for model in MODELS:
-            success = False
-            try:
-                client = InferenceClient(model, token=hf_token)
-                for chunk in client.chat_completion(messages=messages, max_tokens=512, stream=True):
-                    text = chunk.choices[0].delta.content
-                    if text:
-                        yield f"data: {json.dumps({'token': text})}\n\n"
-                        success = True
-               
-                yield "data: [DONE]\n\n"
-                return
-            except Exception as e:
-                if success:
-                    yield "data: [DONE]\n\n"
+
+    queue: Queue = Queue()
+
+    async def consumer():
+        try:
+            for model in MODELS:
+                success = False
+                try:
+                    client = InferenceClient(model, token=hf_token)
+                    for chunk in client.chat_completion(messages=messages, max_tokens=512, stream=True):
+                        text = chunk.choices[0].delta.content
+                        if text:
+                            await queue.put(f"data: {json.dumps({'token': text})}\n\n")
+                            success = True
+                    await queue.put("data: [DONE]\n\n")
                     return
-                continue
-        # All models failed
-        yield f"data: {json.dumps({'token': 'Sorry, all models are currently unavailable. Try again later.'})}\n\n"
-        yield "data: [DONE]\n\n"
+                except Exception:
+                    if success:
+                        await queue.put("data: [DONE]\n\n")
+                        return
+                    continue
+            await queue.put(f"data: {json.dumps({'token': 'Sorry, all models are currently unavailable.'})}\n\n")
+            await queue.put("data: [DONE]\n\n")
+        except Exception as e:
+            await queue.put(f"data: {json.dumps({'token': 'Error generating response.'})}\n\n")
+            await queue.put("data: [DONE]\n\n")
+        finally:
+            await queue.put(None)  
+
+
+    task = create_task(consumer())
 
     async def stream_wrapper():
-        """Safe async wrapper for the generator"""
         try:
-            # Use a safer way to iterate over the generator in a thread
-            gen = generate_tokens()
             while True:
-                try:
-                    chunk = await asyncio.to_thread(next, gen)
-                    yield chunk
-                except StopIteration:
-                    break
-                except Exception as e:
-                    yield f"data: {json.dumps({'token': 'Error occurred during streaming.'})}\n\n"
-                    yield "data: [DONE]\n\n"
+                item = await queue.get()
+                if item is None:
                     break
+                yield item
+                queue.task_done()
         except Exception:
-            yield f"data: {json.dumps({'token': 'Unexpected error in response generation.'})}\n\n"
+            yield f"data: {json.dumps({'token': 'Streaming error occurred.'})}\n\n"
             yield "data: [DONE]\n\n"
+        finally:
+            task.cancel()
 
     return StreamingResponse(stream_wrapper(), media_type="text/event-stream")