Update app.py

app.py

@@ -16,6 +16,8 @@ from PIL import Image
 import io
 import base64
 
+
+
 class StreamLitResponse(ResponseParser):
     def __init__(self, context):
         super().__init__(context)
@@ -128,8 +130,7 @@ def get_vectorstore(text_chunks):
     vectorstore = FAISS.from_texts(texts=text_chunks, embedding=embeddings)
     return vectorstore
 
-
-# Handling user questions
+#handle user input
 def handle_userinput(question, pdf_vectorstore, dfs):
     if pdf_vectorstore and st.session_state.conversation:
         response = st.session_state.conversation({"question": question})
@@ -137,18 +138,16 @@ def handle_userinput(question, pdf_vectorstore, dfs):
 
         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_type = assistant_response.get('type')
             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
+                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": assistant_response})
 
@@ -157,25 +156,12 @@ def handle_userinput(question, pdf_vectorstore, dfs):
 
         st.rerun()
 
-    elif dfs:  # PandasAI case
-        answer = generateResponse(question, dfs)
-        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']
+    elif dfs:
+        results = generateResponse(question, dfs)
 
-            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.session_state.chat_history.append({"role": "user", "content": question})
+        for result in results:
+            st.session_state.chat_history.append({"role": "assistant", **result})
 
         st.rerun()
 
@@ -192,43 +178,50 @@ def get_conversation_chain(vectorstore):
     )
     return conversation_chain
 
-
 def main():
-    
+    st.set_page_config(page_title="Chat with PDFs and Data", page_icon=":books:")
+
     if "conversation" not in st.session_state:
         st.session_state.conversation = None
     if "chat_history" not in st.session_state:
         st.session_state.chat_history = []
-    if "vectorstore" not in st.session_state: #Store vectorstore in session state
+    if "vectorstore" not in st.session_state:
         st.session_state.vectorstore = None
-    
-
-    st.header("Chat with PDFs and Data :books: :bar_chart:")
+    if "dfs" not in st.session_state:
+        st.session_state.dfs = None
 
-    user_question = st.chat_input("Ask a question about your documents or data:")
+    st.title("Chat with PDFs and Data :books: :bar_chart:")
 
     # Chat display
     for message in st.session_state.chat_history:
         with st.chat_message(message["role"]):
             if "dataframe" in message:
-                st.dataframe(message["dataframe"]) 
+                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"])
+                elif isinstance(message["plot"], bytes):
+                    try:
+                        image = Image.open(io.BytesIO(message["plot"]))
+                        st.image(image)
+                    except Exception as e:
+                        st.error(f"Error displaying image: {e}")
                 else:
                     st.write("Unsupported plot format")
             else:
                 st.write(message["content"])
 
+    user_question = st.chat_input("Ask a question about your documents or data:")
+
     if user_question:
         handle_userinput(user_question, st.session_state.vectorstore, st.session_state.dfs)
 
     with st.sidebar:
         st.subheader("Your files")
         uploaded_files = st.file_uploader(
-            "Upload PDFs, CSVs, or Excel files (up to 3)", accept_multiple_files=True, key="file_uploader" # Add a key
+            "Upload PDFs, CSVs, or Excel files (up to 3)", accept_multiple_files=True, key="file_uploader"
         )
 
         if st.button("Process"):
@@ -247,7 +240,6 @@ def main():
                                 st.session_state.dfs = None
                             data_uploaded = False
                             st.warning("Switching to PDF mode. Data files removed.")
-                             # Rerun to clear file uploader
                         pdf_docs.append(uploaded_file)
                         pdf_uploaded = True
                     elif file_extension in ["csv", "xlsx", "xls"]:
@@ -257,7 +249,6 @@ def main():
                                 st.session_state.conversation = None
                             pdf_uploaded = False
                             st.warning("Switching to Data mode. PDF files removed.")
-                             # Rerun to clear file uploader
                         try:
                             if file_extension == 'csv':
                                 df = pd.read_csv(uploaded_file)
@@ -288,4 +279,4 @@ def main():
             st.rerun()
 
 if __name__ == "__main__":
-    main()
+    main()