Update main.py

main.py

@@ -14,6 +14,7 @@ from langchain.chains import LLMChain
 from datetime import datetime
 import matplotlib.pyplot as plt
 import google.generativeai as genai
+import uuid
 
 
 load_dotenv()
@@ -28,24 +29,36 @@ class FlaskResponse(ResponseParser):
         super().__init__(context)
 
     def format_dataframe(self, result):
-        return result['value'].to_html()
+        return result["value"].to_html()
 
     def format_plot(self, result):
-        # Here we assume that result['value'] is a matplotlib Figure.
-        try:
-            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'])
+        val = result["value"]
+        # If val is a matplotlib figure, handle it accordingly.
+        if hasattr(val, "savefig"):
+            try:
+                buf = io.BytesIO()
+                val.savefig(buf, format="png")
+                buf.seek(0)
+                image_base64 = base64.b64encode(buf.read()).decode("utf-8")
+                return f"data:image/png;base64,{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:
+                data = file.read()
+            base64_data = base64.b64encode(data).decode("utf-8")
+            return f"data:image/png;base64,{base64_data}"
+        # Fallback: return as a string.
+        return str(val)
 
     def format_other(self, result):
-        return str(result['value'])
+        # For non-image responses, simply return the value as a string.
+        return str(result["value"])
+
 
 gemini_api_key = os.getenv('Gemini')
 llm = ChatGoogleGenerativeAI(api_key=gemini_api_key, model='gemini-2.0-flash-thinking-exp', temperature=0.1)
@@ -66,6 +79,11 @@ model = genai.GenerativeModel(
 )
 
 
+guid = uuid.uuid4()
+new_filename = f"{guid}"
+user_defined_path = os.path.join("/exports/charts", new_filename)
+
+
 # Endpoint for chat
 @app.route("/chat", methods=["POST"])
 @cross_origin()
@@ -99,7 +117,7 @@ def bot():
                 "wordcloud",
                 "builtins"
             ],
-            "security": "none"
+            "security": "none", "save_charts_path": user_defined_path, "save_charts": False, , "enable_cache": False, "conversational":True
         }
     )