Update app.py

app.py

@@ -121,18 +121,26 @@ async def handle_message(message: cl.Message):
             await cl.Message(content=f"⚠️ Error generating image: {str(e)}").send()
     
     else:
-        # Handle regular QA queries with streaming
-        chain = cl.user_session.get("qa_chain")
-        if chain:
-            try:
-                # Create a message placeholder
-                msg = cl.Message(content="")
-                await msg.send()
-                
-                # Stream the response
-                async for chunk in chain.astream({"query": message.content}):
-                    if "result" in chunk:
-                        await msg.stream_token(chunk["result"])
-                        
-            except Exception as e:
-                await cl.Message(content=f"⚠️ Error: {str(e)}").send()
+        retriever = cl.user_session.get("qa_chain").retriever
+        prompt_template = rag_prompt
+        llm = cl.user_session.get("qa_chain").llm
+
+        try:
+            # Create a message placeholder
+            msg = cl.Message(content="")
+            await msg.send()
+
+            # Step 1: Retrieve relevant documents
+            docs = retriever.get_relevant_documents(message.content)
+            context = "\n\n".join([doc.page_content for doc in docs])
+
+            # Step 2: Format prompt
+            formatted_prompt = prompt_template.format(context=context, question=message.content)
+
+            # Step 3: Stream the response directly from the LLM
+            async for chunk in llm.astream(formatted_prompt):
+                if hasattr(chunk, "content") and chunk.content:
+                    await msg.stream_token(chunk.content)
+
+        except Exception as e:
+            await cl.Message(content=f"⚠️ Error: {str(e)}").send()