Update main.py

main.py

@@ -25,40 +25,63 @@ cors = CORS(app)
 
 
 
-class FlaskResponse(ResponseParser):
+class StreamLitResponse(ResponseParser):
     def __init__(self, context):
         super().__init__(context)
+        # Ensure the export directory exists
+        os.makedirs("/home/user/app/exports/charts", exist_ok=True)
 
     def format_dataframe(self, result):
-        return result["value"].to_html()
+        """
+        Convert a DataFrame to an image using dataframe_image,
+        and return a dict with type 'plot' to match the expected output.
+        """
+        try:
+            df = result['value']
+            # Apply styling if desired
+            styled_df = df.style
+            img_path = f"/home/user/app/exports/charts/{uuid.uuid4().hex}.png"
+            dfi.export(styled_df, img_path)
+        except Exception as e:
+            print("Error in format_dataframe:", e)
+            # Fallback to a string representation if needed
+            img_path = str(result['value'])
+        print("response_class_path (dataframe):", img_path)
+        # Return as a dict with type 'plot'
+        return {'type': 'plot', 'value': img_path}
 
     def format_plot(self, result):
-        val = result["value"]
-        # If val is a matplotlib figure, handle it accordingly.
-        if hasattr(val, "savefig"):
+        img_path = result["value"]
+        # If img_path is a matplotlib figure, handle it accordingly.
+        if hasattr(img_path, "savefig"):
             try:
                 buf = io.BytesIO()
-                val.savefig(buf, format="png")
+                img_path.savefig(buf, format="png")
                 buf.seek(0)
                 image_base64 = base64.b64encode(buf.read()).decode("utf-8")
-                return f"data:image/png;base64,{image_base64}"
+                return {'type': 'plot', 'value': image_base64}
             except Exception as e:
                 print("Error processing figure:", e)
-                return str(val)
-        # If val is a string and is a valid file path, read and encode it.
-        if isinstance(val, str) and os.path.isfile(os.path.join(val)):
-            image_path = os.path.join(val)
-            print("My image path:", image_path)
-            with open(image_path, "rb") as file:
+                return {'type': 'text', 'value': str(img_path)}
+        # If img_path is a string and is a valid file path, read and encode it.
+        if isinstance(img_path, str) and os.path.isfile(img_path):
+            print("My image path:", img_path)
+            with open(img_path, "rb") as file:
                 data = file.read()
             base64_data = base64.b64encode(data).decode("utf-8")
-            return f"data:image/png;base64,{base64_data}"
+            return {'type': 'plot', 'value': base64_data}
         # Fallback: return as a string.
-        return str(val)
+        return {'type': 'text', 'value': str(img_path)}
 
     def format_other(self, result):
-        # For non-image responses, simply return the value as a string.
-        return str(result["value"])
+        if isinstance(result.get('value'), str) and (
+            result.get('value').startswith('iVBOR') or 
+            result.get('value').startswith('data:image')
+        ):
+            return self.format_plot({'type': 'plot', 'value': result['value']})
+        
+        return {'type': 'text', 'value': str(result['value'])}
+
 
 
 gemini_api_key = os.getenv('Gemini')