Issue fix

app.py

@@ -150,16 +150,35 @@ if prompt := st.chat_input("Ask a follow-up question..."):
             #
             # st.markdown(assistant_response)
             # st.session_state.messages.append({"role": "assistant", "content": assistant_response})
+        # with st.spinner("Thinking..."):
+        #     if len(st.session_state.messages) > 1:
+        #         response = llm.invoke([{"role": m["role"], "content": m["content"]} for m in st.session_state.messages])
+        #         response = response.content
+        #     else:
+        #         response = rag_chain.invoke(prompt)
+        #         response = response['result']
+        #
+        #     st.markdown(response)
+        #     st.session_state.messages.append({"role": "assistant", "content": response})
+
         with st.spinner("Thinking..."):
             if len(st.session_state.messages) > 1:
-                response = llm.invoke([{"role": m["role"], "content": m["content"]} for m in st.session_state.messages])
-                response = response['result']
+                conversation_context = "\n".join(
+                    f"{m['role']}: {m['content']}"
+                    for m in st.session_state.messages[:-1]  # Exclude current prompt
+                )
+                augmented_prompt = (
+                    f"Conversation history:\n{conversation_context}\n\n"
+                    f"Current question: {prompt}"
+                )
+                response = rag_chain.invoke({"query": augmented_prompt})
+                assistant_response = response['result']
             else:
-                response = rag_chain.invoke(prompt)
-                response = response['result']
+                response = rag_chain.invoke({"query": prompt})
+                assistant_response = response['result']
 
-            st.markdown(response)
-            st.session_state.messages.append({"role": "assistant", "content": response})
+            st.markdown(assistant_response)
+            st.session_state.messages.append({"role": "assistant", "content": assistant_response})
 
 if st.session_state.messages and st.button("📄 Download Chat as PDF"):
     pdf_file = export_chat_to_pdf(st.session_state.messages)