Update app.py

app.py

@@ -130,15 +130,17 @@ def get_vectorstore(text_chunks):
 
 
 # Handling user questions
-def handle_userinput(question, pdf_vectorstore, dfs):  # Add dfs argument
-    if pdf_vectorstore and st.session_state.conversation:  # Check both
+def handle_userinput(question, pdf_vectorstore, dfs):
+    if pdf_vectorstore and st.session_state.conversation:
         response = st.session_state.conversation({"question": question})
-        st.session_state.chat_history = response['chat_history']
-        answer = response['answer']
-        st.write(answer)
-    elif dfs:  # Handle pandasai if dataframes are present
+        st.session_state.chat_history.append({"role": "user", "content": question}) # Add to chat history
+        st.session_state.chat_history.append({"role": "assistant", "content": response['answer']}) # Add to chat history
+        st.rerun()  # Rerun to update the chat display
+    elif dfs:
         answer = generateResponse(question, dfs)
-        st.write(answer)
+        st.session_state.chat_history.append({"role": "user", "content": question})  # Add to chat history
+        st.session_state.chat_history.append({"role": "assistant", "content": answer}) # Add to chat history
+        st.rerun() # Rerun to update the chat display
     else:
         st.write("Please upload and process your documents/data first.")
 
@@ -161,10 +163,16 @@ def main():
         st.session_state.chat_history = None
     if "vectorstore" not in st.session_state: #Store vectorstore in session state
         st.session_state.vectorstore = None
+    
 
     st.header("Chat with PDFs and Data :books: :bar_chart:")
 
-    user_question = st.text_input("Ask a question about your documents or data:")
+    user_question = st.chat_input("Ask a question about your documents or data:")
+
+    # Chat display
+    for message in st.session_state.chat_history:
+        with st.chat_message(message["role"]):
+            st.write(message["content"])
 
     if user_question:
         handle_userinput(user_question, st.session_state.vectorstore, st.session_state.dfs)
@@ -227,5 +235,9 @@ def main():
                 else:
                     st.session_state.dfs = None
 
+        if st.button("Clear Chat"):
+            st.session_state.chat_history = []
+            st.rerun()
+
 if __name__ == "__main__":
     main()