Update app.py

app.py

@@ -160,13 +160,15 @@ if selected_mall:
 st.subheader("💬 Chat with WaysAhead AI Engine")
 
 system_prompt = """
-You are an advanced AI engine that understands the retail industry in the UAE. 
-You help users find the best store recommendations, analyze competitors, and generate visualizations.
-If the user requests a chart, generate it and return the image properly formatted for display.
+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.
 """
 
 agent = create_pandas_dataframe_agent(
-    ChatOpenAI(model="gpt-4", api_key=openai_api_key),
+    ChatOpenAI(model="gpt-3.5-turbo", api_key=openai_api_key),
     df_kpi,
     verbose=True,
     agent_type="openai-tools",
@@ -175,21 +177,34 @@ agent = create_pandas_dataframe_agent(
     system_prompt=system_prompt
 )
 
+# ✅ User Input
 user_query = st.text_input("🔍 Ask AI for similar stores or analytics:")
 if user_query:
     try:
         response = agent.run(user_query)
 
-        # ✅ Detect if the response is an image
-        if "<figure>" in response and "<img" in response and "base64" in response:
-            # Extract Base64 Image Data
-            base64_img = response.split('base64,')[1].split('"')[0]
+        # ✅ Detect if AI Generated Python Code for a Chart
+        if "import matplotlib.pyplot" in response:
+            st.write("🛠 AI is generating a chart based on your request...")
 
-            # Decode and Display Image
-            image_bytes = BytesIO(base64.b64decode(base64_img))
-            image = plt.imread(image_bytes, format="png")
+            # ✅ Extract and Execute Python Code Safely
+            try:
+                exec_globals = {}
+                exec(response, {}, exec_globals)  # Execute AI-generated code
+
+                # ✅ 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)
+
+            except Exception as e:
+                st.error(f"⚠️ Error generating chart: {e}")
 
-            st.image(image, caption="Generated Visualization", use_column_width=True)
         else:
             st.markdown(f"🤖 **AI Response:** {response}")