Update app.py

app.py

@@ -11,19 +11,78 @@ import pandas as pd
 from pandasai import SmartDataframe, SmartDatalake
 from pandasai.responses.response_parser import  ResponseParser
 from pandasai.llm import GoogleGemini
+import plotly.graph_objects as go
+from PIL import Image
+import io
+import base64
 
 class StreamLitResponse(ResponseParser):
-        def __init__(self,context) -> None:
-              super().__init__(context)
-        def format_dataframe(self,result):
-               st.dataframe(result['value'])
-               return
-        def format_plot(self,result):
-               st.image(result['value'], use_container_width=True)
-               return
-        def format_other(self, result):
-               st.write(result['value'])
-               return
+    def __init__(self, context):
+        super().__init__(context)
+
+    def format_dataframe(self, result):
+        st.dataframe(result['value'])
+        return
+
+    def format_plot(self, result):
+        try:
+            image = result['value']
+
+            if isinstance(image, Image.Image):  # PIL Image
+                img_bytes = io.BytesIO()
+                image.save(img_bytes, format="PNG")
+                img_bytes = img_bytes.getvalue()
+                encoded = base64.b64encode(img_bytes).decode("ascii")
+
+                st.image(image)  # Display the image
+                fig = go.Figure(data=[go.Image(source=f'data:image/png;base64,{encoded}')])
+                fig.update_layout(
+                    margin=dict(l=0, r=0, b=0, t=0),
+                    xaxis=dict(visible=False),
+                    yaxis=dict(visible=False),
+                )
+                st.plotly_chart(fig)
+
+            elif isinstance(image, bytes):  # Bytes
+                encoded = base64.b64encode(image).decode("ascii")
+                fig = go.Figure(data=[go.Image(source=f'data:image/png;base64,{encoded}')])
+                fig.update_layout(
+                    margin=dict(l=0, r=0, b=0, t=0),
+                    xaxis=dict(visible=False),
+                    yaxis=dict(visible=False),
+                )
+                st.plotly_chart(fig)
+
+            elif isinstance(image, str) and os.path.exists(image):  # File Path
+                with open(image, "rb") as f:
+                    encoded = base64.b64encode(f.read()).decode("ascii")
+                fig = go.Figure(data=[go.Image(source=f'data:image/png;base64,{encoded}')])
+                fig.update_layout(
+                    margin=dict(l=0, r=0, b=0, t=0),
+                    xaxis=dict(visible=False),
+                    yaxis=dict(visible=False),
+                )
+                st.plotly_chart(fig)
+
+            elif isinstance(image, str):  # Base64 encoded string
+                fig = go.Figure(data=[go.Image(source=f'data:image/png;base64,{image}')])
+                fig.update_layout(
+                    margin=dict(l=0, r=0, b=0, t=0),
+                    xaxis=dict(visible=False),
+                    yaxis=dict(visible=False),
+                )
+                st.plotly_chart(fig)
+
+            else:
+                st.write("Unsupported Image format")
+
+        except Exception as e:
+            st.image(image)
+            st.write(f"Error displaying image: {e}")
+
+    def format_other(self, result):
+        st.write(result['value'])
+        return
 
 
 load_dotenv()  # Load environment variables at the beginning
@@ -113,7 +172,7 @@ def main():
     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
+            "Upload PDFs, CSVs, or Excel files (up to 3)", accept_multiple_files=True, key="file_uploader" # Add a key
         )
 
         if st.button("Process"):
@@ -127,20 +186,22 @@ def main():
                     file_extension = uploaded_file.name.split(".")[-1].lower()
 
                     if file_extension == "pdf":
-                        if data_uploaded:  # Check if data was already uploaded
+                        if data_uploaded:
                             if st.session_state.dfs:
                                 st.session_state.dfs = None
                             data_uploaded = False
-                            st.warning("Switching to PDF mode. Data files removed.")  # Inform user
+                            st.warning("Switching to PDF mode. Data files removed.")
+                            st.experimental_rerun() # Rerun to clear file uploader
                         pdf_docs.append(uploaded_file)
                         pdf_uploaded = True
                     elif file_extension in ["csv", "xlsx", "xls"]:
-                        if pdf_uploaded:  # Check if PDF was already uploaded
+                        if pdf_uploaded:
                             if st.session_state.vectorstore:
                                 st.session_state.vectorstore = None
                                 st.session_state.conversation = None
                             pdf_uploaded = False
-                            st.warning("Switching to Data mode. PDF files removed.")  # Inform user
+                            st.warning("Switching to Data mode. PDF files removed.")
+                            st.rerun() # Rerun to clear file uploader
                         try:
                             if file_extension == 'csv':
                                 df = pd.read_csv(uploaded_file)
@@ -166,6 +227,15 @@ def main():
                 else:
                     st.session_state.dfs = None
 
+# Example of creating a bar chart with Plotly Express (replace with your actual data)
+if st.session_state.dfs:
+    for df_ in st.session_state.dfs: #Iterate over dataframes
+        try:
+            fig = px.bar(df_, x=df_.columns[0], y=df_.columns[1], title="Example Bar Chart") #Assumes first two columns are x and y
+            fig.update_layout(xaxis_tickangle=-45) # Rotate x-axis labels for better readability
+            st.plotly_chart(fig)
+        except Exception as e:
+            st.write(f"Error plotting chart for dataframe: {e}") #Handle exceptions when plotting
+
 if __name__ == "__main__":
     main()
-