Update stories.py

stories.py

@@ -191,51 +191,55 @@ def transcribe_audio(audio_file):
 # PandasAI Response for DataFrame (using SmartDataframe and ChatSambaNovaCloud)
 # -----------------------
 def generateResponse(prompt, df):
-    
-    """Generate response using PandasAI with SmartDataframe and the ChatSambaNovaCloud LLM."""
-    
+    """
+    Return either a base64-encoded PNG string like 'data:image/png;base64,...'
+    if the answer is a chart, or a fallback string if the answer is something else.
+    """
     pandas_agent = SmartDataframe(
         df,
         config={
             "llm": llm,
-            "response_parser": FlaskResponse,
+            "response_parser": FlaskResponse,  # You can still use it for internal logic
             "custom_whitelisted_dependencies": [
-                "os",
-                "io",
-                "sys",
-                "chr",
-                "glob",
-                "b64decoder",
-                "collections",
-                "geopy",
-                "geopandas",
-                "wordcloud",
-                "builtins"
+                "os", "io", "sys", "chr", "glob", "b64decoder",
+                "collections", "geopy", "geopandas", "wordcloud", "builtins"
             ],
-            "security": "none", "save_charts_path": user_defined_path, "save_charts": False, "enable_cache": False,
+            "security": "none",
+            "save_charts_path": user_defined_path,
+            "save_charts": False,
+            "enable_cache": False,
         }
     )
 
-    # Get the answer from the agent
-    answer = pandas_agent.chat(user_question)
+    answer = pandas_agent.chat(prompt)
 
-    # Process the answer based on its type
-    formatted_answer = None
+    # Convert 'answer' into a base64 string or fallback
     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 answer.to_html()
 
-    # Return the formatted answer as JSON.
-    return jsonify({"answer": formatted_answer})
+    elif hasattr(answer, "savefig"):  # e.g. a Matplotlib figure
+        try:
+            buf = io.BytesIO()
+            answer.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 None
+
+    elif isinstance(answer, str):
+        # Could be a file path or just a textual answer
+        if os.path.isfile(answer):
+            with open(answer, "rb") as f:
+                data = f.read()
+            b64 = base64.b64encode(data).decode("utf-8")
+            return f"data:image/png;base64,{b64}"
+        else:
+            return answer
+    else:
+        # fallback
+        return str(answer)
 
 # -----------------------
 # DataFrame-Based Story Generation (for CSV/Excel files)