Update routes.py

src/api/routes.py

@@ -3,10 +3,10 @@ import json
 
 from fastapi import APIRouter, Request
 from fastapi.responses import JSONResponse, StreamingResponse
-from src.utils.helpers import get_clean_text
 
 from src.core.config import settings
 from src.core.engine import engine
+from src.utils.helpers import get_clean_text
 
 router = APIRouter()
 
@@ -25,23 +25,79 @@ async def chat_completions(request: Request):
     stream = data.get("stream", True)
 
     async def stream_generator():
+        # Ensure sequential processing: acquire the engine lock for the whole generation
         async with engine.lock:
-            output = engine.generate(
-                messages,
-                data.get("max_tokens", settings.DEFAULT_MAX_TOKENS),
-                data.get("temperature", settings.DEFAULT_TEMP),
-            )
-            for chunk in output:
-                yield f"data: {json.dumps(chunk)}\n\n"
-            yield "data: [DONE]\n\n"
+            import threading
+
+            # Use an asyncio.Queue to safely transfer chunks from a blocking worker thread to the async generator
+            q: asyncio.Queue = asyncio.Queue()
+            stop_event = threading.Event()
+            loop = asyncio.get_running_loop()
+
+            def worker():
+                try:
+                    for chunk in engine.llm.create_chat_completion(
+                        messages=messages,
+                        max_tokens=int(
+                            data.get("max_tokens", settings.DEFAULT_MAX_TOKENS)
+                        ),
+                        temperature=float(
+                            data.get("temperature", settings.DEFAULT_TEMP)
+                        ),
+                        stream=True,
+                    ):
+                        # stop early if requested (e.g. client disconnected)
+                        if stop_event.is_set():
+                            break
+                        loop.call_soon_threadsafe(q.put_nowait, chunk)
+                except Exception as e:
+                    # Pass exception to the async side so we can surface an error or terminate cleanly
+                    loop.call_soon_threadsafe(q.put_nowait, {"__error": str(e)})
+                finally:
+                    # Sentinel to mark completion
+                    loop.call_soon_threadsafe(q.put_nowait, None)
+
+            # Run the blocking model iteration in a thread so it doesn't block the event loop
+            worker_future = loop.run_in_executor(None, worker)
+
+            try:
+                while True:
+                    item = await q.get()
+                    if item is None:
+                        # worker finished normally
+                        break
+                    # If worker reported an error, stream it and break
+                    if isinstance(item, dict) and item.get("__error"):
+                        yield f"data: {json.dumps({'error': item['__error']})}\n\n"
+                        break
+                    yield f"data: {json.dumps(item)}\n\n"
+                yield "data: [DONE]\n\n"
+            except asyncio.CancelledError:
+                # Client disconnected: signal the worker to stop and wait for it to finish, then re-raise to terminate streaming
+                stop_event.set()
+                try:
+                    await worker_future
+                except Exception:
+                    pass
+                raise
+            finally:
+                # Ensure worker is signalled to stop and awaited (idempotent)
+                stop_event.set()
+                try:
+                    await worker_future
+                except Exception:
+                    pass
 
     if stream:
         return StreamingResponse(stream_generator(), media_type="text/event-stream")
 
+    # For non-streaming responses: keep sequential processing but run blocking work in a thread
     async with engine.lock:
-        return engine.generate(
+        result = await asyncio.to_thread(
+            engine.generate,
             messages,
             data.get("max_tokens", settings.DEFAULT_MAX_TOKENS),
             data.get("temperature", settings.DEFAULT_TEMP),
             stream=False,
         )
+        return result