changes

app.py

@@ -7,6 +7,8 @@ from langchain_community.llms import Ollama
 from langchain_core.messages import HumanMessage
 import logging
 from functools import lru_cache
+from langchain.callbacks.manager import CallbackManager
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -16,7 +18,8 @@ MODEL_NAME = 'gemma2:2b'
 
 @lru_cache()
 def get_llm():
-    return Ollama(model=MODEL_NAME)
+    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
+    return Ollama(model=MODEL_NAME, callback_manager=callback_manager)
 
 class Question(BaseModel):
     text: str
@@ -45,20 +48,9 @@ async def ask_question_stream(question: Question):
 
         async def generate():
             full_response = ""
-            buffer = ""
             async for chunk in llm.astream(question.text):
-                buffer += chunk
-                words = re.findall(r'\S+|\s+', buffer)
-                
-                for word in words[:-1]:
-                    full_response += word
-                    yield word
-                
-                buffer = words[-1] if words else ""
-            
-            if buffer:
-                full_response += buffer
-                yield buffer
+                full_response += chunk
+                yield chunk
             
             # Log the full response after streaming is complete
             logger.info(f"Full streamed response: {full_response}")
@@ -67,7 +59,7 @@ async def ask_question_stream(question: Question):
     except Exception as e:
         logger.error(f"Error in /ask_stream endpoint: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
-
+    
 @app.on_event("startup")
 async def startup_event():
     logger.info(f"Starting up with model: {MODEL_NAME}")