Update app.py

app.py

@@ -164,11 +164,11 @@ You are an advanced AI specializing in retail analytics. You help users:
 1. Find similar stores based on category and sales performance.
 2. Generate insights from sales data.
 3. Create visualizations like bar charts, line charts, and trends.
-4. If a visualization request is made, generate Python code to create the chart.
+4. If a visualization request is made like a heatmap or chart, give proper and full Python code to create the chart.
 """
 
 agent = create_pandas_dataframe_agent(
-    ChatOpenAI(model="gpt-3.5-turbo", api_key=openai_api_key),
+    ChatOpenAI(model="gpt-4", api_key=openai_api_key),
     df_kpi,
     verbose=True,
     agent_type="openai-tools",
@@ -187,17 +187,18 @@ if user_query:
         if "import matplotlib.pyplot" in response:
             st.write("🛠 AI is generating a chart based on your request...")
 
-            # ✅ Extract and Execute Python Code Safely
+            # ✅ Create a dictionary to store execution variables
+            exec_globals = {"df": df_kpi}
+
+            # ✅ Execute the AI-generated code
             try:
-                exec_globals = {}
-                exec(response, {}, exec_globals)  # Execute AI-generated code
+                exec(response, {}, exec_globals)
 
                 # ✅ Capture Matplotlib Figure
                 fig = plt.gcf()
                 buf = BytesIO()
                 fig.savefig(buf, format="png")
                 buf.seek(0)
-                image_base64 = base64.b64encode(buf.read()).decode()
 
                 # ✅ Display Image in Streamlit
                 st.image(buf, caption="📊 AI-Generated Chart", use_column_width=True)
@@ -212,6 +213,4 @@ if user_query:
         st.error(f"⚠️ Error: {e}")
 
 st.markdown("---")
-st.markdown("🚀 **Powered by WaysAhead Global AI Engine**")
-
-
+st.markdown("🚀 **Powered by WaysAhead Global AI Engine**")