Changed Generate stream to async

main/api.py

@@ -2,7 +2,8 @@ import os
 from pathlib import Path
 from threading import Thread
 import torch
-from typing import Optional, Iterator, List
+from typing import Optional, List, AsyncIterator
+import asyncio
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 from .utils.logging import setup_logger
 
@@ -248,12 +249,12 @@ class LLMApi:
             self.logger.error(f"Error generating response: {str(e)}")
             raise
 
-    def generate_stream(
+    async def generate_stream(
             self,
             prompt: str,
             system_message: Optional[str] = None,
             max_new_tokens: Optional[int] = None
-    ) -> Iterator[str]:
+    ) -> AsyncIterator[str]:
         """
         Generate a streaming response for the given prompt.
         """
@@ -287,10 +288,12 @@ class LLMApi:
             thread = Thread(target=self.generation_model.generate, kwargs=generation_kwargs)
             thread.start()
 
-            # Yield the generated text in chunks
+            # Use async generator to yield chunks
             for new_text in streamer:
                 self.logger.debug(f"Generated chunk: {new_text[:50]}...")
                 yield new_text
+                # Add a small delay to allow other tasks to run
+                await asyncio.sleep(0)
 
         except Exception as e:
             self.logger.error(f"Error in streaming generation: {str(e)}")