Update main.py

main.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from hugchat import hugchat
 from hugchat.login import Login
@@ -53,34 +54,34 @@ async def startup_event():
 @app.post("/generate")
 async def generate_response(request: QuestionRequest):
     """
-    Generates a response from the AI based on the provided question.
+    Generates a streaming response from the AI based on the provided question.
     
     Args:
         request (QuestionRequest): JSON body containing the question.
     
     Returns:
-        dict: A dictionary containing the AI's response or an error message.
+        StreamingResponse: A streaming response with the AI's reply.
     """
     global chatbot
     if chatbot is None:
         raise HTTPException(status_code=500, detail="Chatbot not initialized. Please try again later.")
     
-    try:
-        # Generate response (non-streaming for simplicity)
-        response_data = chatbot.chat(request.question, stream=False)
-        
-        # Extract the actual response text
-        # The response may be a dictionary; check for 'gen' or other keys
-        if isinstance(response_data, dict):
-            # Assuming 'gen' contains the response text (list of strings)
-            response_text = "".join(response_data.get("gen", [])) if response_data.get("gen") else "Here's what we can do: Let's discuss your vision!"
-            if not response_text:
-                # Fallback to a default premium response if 'gen' is empty
-                response_text = f"Welcome, valued client! How can Abdullah Ali and our premium team bring your vision to life with a custom website or AI chatbot?"
-        else:
-            response_text = response_data  # Direct string response (if hugchat returns string)
-        
-        return {"response": response_text}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to generate response: {str(e)}")
+    async def generate():
+        try:
+            # Use streaming to match client expectations
+            for token in chatbot.chat(request.question, stream=True):
+                # Handle hugchat response (token may be dict or string)
+                if isinstance(token, dict):
+                    # Extract token from dictionary (e.g., {"type": "stream", "token": "text"})
+                    token_text = token.get("token", "")
+                else:
+                    token_text = str(token)
+                
+                # Encode and yield the token
+                yield token_text.encode('utf-8')
+        except Exception as e:
+            error_message = f"Error: Failed to generate response: {str(e)}"
+            yield error_message.encode('utf-8')
+    
+    return StreamingResponse(generate(), media_type="text/plain")