Update main.py

main.py

@@ -25,63 +25,40 @@ cors = CORS(app)
 
 
 
-class StreamLitResponse(ResponseParser):
+class FlaskResponse(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):
-        """
-        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}
+        return result["value"].to_html()
 
     def format_plot(self, result):
-        img_path = result["value"]
-        # If img_path is a matplotlib figure, handle it accordingly.
-        if hasattr(img_path, "savefig"):
+        val = result["value"]
+        # If val is a matplotlib figure, handle it accordingly.
+        if hasattr(val, "savefig"):
             try:
                 buf = io.BytesIO()
-                img_path.savefig(buf, format="png")
+                val.savefig(buf, format="png")
                 buf.seek(0)
                 image_base64 = base64.b64encode(buf.read()).decode("utf-8")
-                return {'type': 'plot', 'value': image_base64}
+                return f"data:image/png;base64,{image_base64}"
             except Exception as e:
                 print("Error processing figure:", e)
-                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:
+                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:
                 data = file.read()
             base64_data = base64.b64encode(data).decode("utf-8")
-            return {'type': 'plot', 'value': base64_data}
+            return f"data:image/png;base64,{base64_data}"
         # Fallback: return as a string.
-        return {'type': 'text', 'value': str(img_path)}
+        return str(val)
 
     def format_other(self, result):
-        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'])}
-
+        # For non-image responses, simply return the value as a string.
+        return str(result["value"])
 
 
 gemini_api_key = os.getenv('Gemini')