Update main.py

main.py

@@ -24,20 +24,25 @@ cors = CORS(app)
 
 
 class FlaskResponse(ResponseParser):
-    def __init__(self, context) -> None:
+    def __init__(self, context):
         super().__init__(context)
 
     def format_dataframe(self, result):
         return result['value'].to_html()
 
     def format_plot(self, result):
+        # Here we assume that result['value'] is a matplotlib Figure.
         try:
-            img_path = result['value']
-        except ValueError:
-            img_path = str(result['value'])
-            print("ValueError:", img_path)
-        print("response_class_path:", img_path)
-        return img_path
+            fig = result['value']
+            buf = io.BytesIO()
+            fig.savefig(buf, format="png")
+            buf.seek(0)
+            image_base64 = base64.b64encode(buf.read()).decode("utf-8")
+            # Return a data URL that can be rendered in an <img> tag.
+            return f"data:image/png;base64,{image_base64}"
+        except Exception as e:
+            print("Error processing plot:", e)
+            return str(result['value'])
 
     def format_other(self, result):
         return str(result['value'])
@@ -59,40 +64,65 @@ model = genai.GenerativeModel(
   model_name="gemini-2.0-flash-thinking-exp",
   generation_config=generation_config,
 )
+
+
 # Endpoint for chat
 @app.route("/chat", methods=["POST"])
 @cross_origin()
 def bot():
-
+    # Retrieve parameters from the request
     json_table = request.json.get("json_table")
     user_question = request.json.get("user_question")
-    #data = request.get_json(force=True)TRye
-    #print(req_body)
-    #data = eval(req_body)
-    #json_table = data["json_table"]
-    #user_question = data["user_question"]
-    #print(json_table)
-    print(user_question)
+    print("User question:", user_question)
+
+    # Convert the table data into a dataframe
     data = eval(str(json_table))
     df = pd.DataFrame(data)
-    print(list(df))
-    pandas_agent = SmartDataframe(df,config={"llm":llm, "response_parser":FlaskResponse, "custom_whitelisted_dependencies": [
-        "os",
-        "io",
-        "sys",
-        "chr",
-        "glob",
-        "b64decoder",
-        "collections",
-        "geopy",
-        "geopandas",
-        "wordcloud",
-        "builtins"
-    ], "security":"none"})
+    print("Columns in dataframe:", list(df.columns))
+
+    # Create a SmartDataframe instance using your configuration.
+    pandas_agent = SmartDataframe(
+        df,
+        config={
+            "llm": llm,
+            "response_parser": FlaskResponse,
+            "custom_whitelisted_dependencies": [
+                "os",
+                "io",
+                "sys",
+                "chr",
+                "glob",
+                "b64decoder",
+                "collections",
+                "geopy",
+                "geopandas",
+                "wordcloud",
+                "builtins"
+            ],
+            "security": "none"
+        }
+    )
+
+    # Get the answer from the agent
     answer = pandas_agent.chat(user_question)
-    return jsonify(answer)
-    return answer
-    #df = df.rename(co
+
+    # Process the answer based on its type
+    formatted_answer = None
+    if isinstance(answer, pd.DataFrame):
+        formatted_answer = answer.to_html()
+    elif isinstance(answer, plt.Figure):
+        buf = io.BytesIO()
+        answer.savefig(buf, format="png")
+        buf.seek(0)
+        image_base64 = base64.b64encode(buf.read()).decode("utf-8")
+        formatted_answer = f"data:image/png;base64,{image_base64}"
+    elif isinstance(answer, (int, float)):
+        formatted_answer = str(answer)
+    else:
+        formatted_answer = str(answer)
+
+    # Return the formatted answer as JSON.
+    return jsonify({"answer": formatted_answer})
 
 # Reports endpoint
 @app.route("/report", methods=["POST"])