fix: Streaming LLM and API outputs

nilai-api/src/nilai_api/auth.py

@@ -9,7 +9,6 @@
 
 def get_user(credentials: HTTPAuthorizationCredentials = Security(bearer_scheme)):
     token = credentials.credentials
-    print(token)
     user = UserManager.check_api_key(token)
     if user:
         return user

nilai-api/src/nilai_api/routers/private.py

@@ -1,6 +1,7 @@
 # Fast API and serving
 import logging
 import os
+import asyncio
 from base64 import b64encode
 from typing import AsyncGenerator, Union
 
@@ -168,6 +169,9 @@ async def stream_response() -> AsyncGenerator[str, None]:
                         async for chunk in response.aiter_lines():
                             if chunk:  # Skip empty lines
                                 yield f"{chunk}\n"
+                                await asyncio.sleep(
+                                    0
+                                )  # Add an await to return inmediately
             except httpx.HTTPStatusError as e:
                 raise HTTPException(
                     status_code=e.response.status_code,

nilai-models/src/nilai_models/models/llama_model.py

@@ -1,6 +1,7 @@
+import asyncio
 import json
 import logging
-from typing import Any, Generator
+from typing import AsyncGenerator
 
 from fastapi import HTTPException
 from fastapi.responses import StreamingResponse
@@ -71,15 +72,20 @@ async def chat_completion(
         # Streaming response logic
         if req.stream:
 
-            def generate() -> Generator[str, Any, None]:
+            async def generate() -> AsyncGenerator[str, None]:
                 try:
                     # Create a generator for the streamed output
-                    for output in self.model.create_chat_completion(
-                        prompt,  # type: ignore
-                        stream=True,
-                        temperature=req.temperature if req.temperature else 0.2,
-                        max_tokens=req.max_tokens,
-                    ):
+                    loop = asyncio.get_event_loop()
+                    output_generator = await loop.run_in_executor(
+                        None,
+                        lambda: self.model.create_chat_completion(
+                            prompt,  # type: ignore
+                            stream=True,
+                            temperature=req.temperature if req.temperature else 0.2,
+                            max_tokens=req.max_tokens,
+                        ),
+                    )
+                    for output in output_generator:
                         # Extract delta content from output
                         choices = output.get("choices", [])  # type: ignore
                         if not choices or "delta" not in choices[0]:
@@ -92,6 +98,7 @@ def generate() -> Generator[str, Any, None]:
                         )  # Create a ChoiceChunk
                         completion_chunk = ChatCompletionChunk(choices=[chunk])
                         yield f"data: {completion_chunk.model_dump_json()}\n\n"  # Stream the chunk
+                        await asyncio.sleep(0)  # Add an await to return inmediately
 
                     yield "data: [DONE]\n\n"
                 except Exception as e: