Update app.py

app.py

@@ -133,14 +133,52 @@ def get_vectorstore(text_chunks):
 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.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:
+        st.session_state.chat_history.append({"role": "user", "content": question})
+
+        assistant_response = response['answer']
+
+        # Check if the response contains image or dataframe information (adapt as needed)
+
+        if isinstance(assistant_response, dict) and 'value' in assistant_response:
+            content_type = assistant_response.get('type') # Assuming your StreamLitResponse sets this
+            content_value = assistant_response['value']
+
+            if content_type == "dataframe":
+                st.session_state.chat_history.append({"role": "assistant", "content": "DataFrame"}) # Store a placeholder
+                st.session_state.chat_history.append({"role": "assistant", "dataframe": content_value})  # Store the DataFrame separately
+            elif content_type == "plot":  # Or whatever you name the plot type
+                st.session_state.chat_history.append({"role": "assistant", "content": "Plot"})  # Placeholder
+                st.session_state.chat_history.append({"role": "assistant", "plot": content_value}) # Store the image/plot data
+            else:
+                st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
+
+        else:
+            st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
+
+        st.rerun()
+
+    elif dfs:  # PandasAI case
         answer = generateResponse(question, dfs)
-        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
+        st.session_state.chat_history.append({"role": "user", "content": question})
+
+        if isinstance(answer, dict) and 'value' in answer:
+            content_type = answer.get('type')
+            content_value = answer['value']
+
+            if content_type == "dataframe":
+                st.session_state.chat_history.append({"role": "assistant", "content": "DataFrame"})
+                st.session_state.chat_history.append({"role": "assistant", "dataframe": content_value})
+            elif content_type == "plot":
+                st.session_state.chat_history.append({"role": "assistant", "content": "Plot"})
+                st.session_state.chat_history.append({"role": "assistant", "plot": content_value})
+            else:
+                st.session_state.chat_history.append({"role": "assistant", "content": answer})
+
+        else:
+            st.session_state.chat_history.append({"role": "assistant", "content": answer})
+
+        st.rerun()
+
     else:
         st.write("Please upload and process your documents/data first.")
 
@@ -172,7 +210,17 @@ def main():
     # Chat display
     for message in st.session_state.chat_history:
         with st.chat_message(message["role"]):
-            st.write(message["content"])
+            if "dataframe" in message:
+                st.dataframe(message["dataframe"]) 
+            elif "plot" in message:
+                if isinstance(message["plot"], Image.Image):
+                    st.image(message["plot"])
+                elif isinstance(message["plot"], go.Figure):
+                    st.plotly_chart(message["plot"])
+                else:
+                    st.write("Unsupported plot format")
+            else:
+                st.write(message["content"])
 
     if user_question:
         handle_userinput(user_question, st.session_state.vectorstore, st.session_state.dfs)