Update app.py

app.py

@@ -16,76 +16,26 @@ from PIL import Image
 import io
 import base64
 
-
-
 class StreamLitResponse(ResponseParser):
     def __init__(self, context):
         super().__init__(context)
+        self.results = []
 
     def format_dataframe(self, result):
-        st.dataframe(result['value'])
+        self.results.append({"type": "dataframe", "value": 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}")
+        self.results.append({"type": "plot", "value": result['value']})
+        return
 
     def format_other(self, result):
-        st.write(result['value'])
+        self.results.append({"type": "text", "value": result['value']})
         return
 
+    def get_results(self):
+        return self.results
+
 
 load_dotenv()  # Load environment variables at the beginning
 GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY') #Use .get to handle if the variable is not present