Update main.py

main.py

@@ -32,17 +32,28 @@ app = Flask(__name__)
 cors = CORS(app)
 
 class FlaskResponse(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'])
-               return
-        def format_other(self, result):
-               st.write(result['value'])
-               return
+  def __init__(self, context) -> None:
+      super().__init__(context)
+
+  def format_dataframe(self, result):
+      return result['value'].to_html()
+
+  def format_plot(self, result):
+        # Save the plot using savefig
+      try:
+          
+          img_path = result['value']
+      
+        
+      except ValueError:
+          img_path = str(result['value'])
+          print("value error!", img_path)
+        
+      print("response_class_path:", img_path)
+      return img_path
+
+  def format_other(self, result):
+      return str(result['value'])
 
 gemini_api_key = os.environ['Gemini']
 
@@ -94,11 +105,9 @@ def bot():
 
     lake = SmartDatalake([inventory_df, sales_df], config={"llm":llm, "response_parser":FlaskResponse, "enable_cache": False, "save_logs":False})
     response = lake.chat(user_question)
-    resp = str(response)
     print(user_question)
     
-    print(resp)
-    return jsonify(resp)
+    return jsonify(response)
 
 
 if __name__ == "__main__":