Update main.py

main.py

@@ -11,11 +11,9 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = FastAPI()
-MODEL_NAME = 'krishna_choudhary/AI_Assistant_Chatbot' # Your specific model name
+MODEL_NAME = 'krishna_choudhary/AI_Assistant_Chatbot'
 
 def get_llm():
-    # Note: StreamingStdOutCallbackHandler is for console output, not for API streaming.
-    # For actual API streaming, you'd integrate a custom callback that yields chunks.
     callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
     return Ollama(model=MODEL_NAME, callback_manager=callback_manager)
 
@@ -30,8 +28,6 @@ def read_root():
 async def ask_question(question: Question):
     try:
         llm = get_llm()
-        # For simplicity, we're using invoke() which returns the full response at once.
-        # If you need true streaming to the client, you'd use llm.stream() and a custom StreamingResponse.
         response = llm.invoke(question.text)
         return {"response": response}
     except Exception as e:
@@ -45,4 +41,4 @@ async def startup_event():
 
 @app.on_event("shutdown")
 async def shutdown_event():
-    logger.info("Shutting down")
+    logger.info("Shutting down")