Update app.py

app.py

@@ -89,7 +89,7 @@ if st.session_state.df is not None and st.session_state.show_preview:
 
 
 
-def ask_gpt4o_for_visualization(query, df, llm):
+"""def ask_gpt4o_for_visualization(query, df, llm):
     columns = ', '.join(df.columns)
     prompt = f"""
     Analyze the query and suggest one or more relevant visualizations.
@@ -110,7 +110,85 @@ def ask_gpt4o_for_visualization(query, df, llm):
         return json.loads(response)
     except json.JSONDecodeError:
         st.error("⚠️ GPT-4o failed to generate a valid suggestion.")
-        return None
+        return None"""
+
+def ask_gpt4o_for_visualization(query, df, llm, retries=2):
+    import json
+
+    # Identify numeric and categorical columns
+    numeric_columns = df.select_dtypes(include='number').columns.tolist()
+    categorical_columns = df.select_dtypes(exclude='number').columns.tolist()
+
+    # Enhanced Prompt with Clear Instructions
+    prompt = f"""
+    Analyze the following query and suggest the most suitable visualization(s) using the dataset.
+
+    **Query:** "{query}"
+
+    **Numeric Columns (for Y-axis):** {', '.join(numeric_columns) if numeric_columns else 'None'}
+    **Categorical Columns (for X-axis or grouping):** {', '.join(categorical_columns) if categorical_columns else 'None'}
+
+    Suggest visualizations in this exact JSON format:
+    [
+      {{
+        "chart_type": "bar/box/line/scatter/pie/heatmap",
+        "x_axis": "categorical_or_time_column",
+        "y_axis": "numeric_column",
+        "group_by": "optional_column_for_grouping",
+        "title": "Title of the chart",
+        "description": "Why this chart is suitable"
+      }}
+    ]
+
+    **Examples:**  
+    - For salary distribution:  
+      {{
+        "chart_type": "box",
+        "x_axis": "job_title",
+        "y_axis": "salary_in_usd",
+        "group_by": "experience_level",
+        "title": "Salary Distribution by Job Title and Experience",
+        "description": "A box plot showing salary ranges across job titles and experience levels."
+      }}
+
+    - For trend analysis:  
+      {{
+        "chart_type": "line",
+        "x_axis": "year",
+        "y_axis": "revenue",
+        "group_by": null,
+        "title": "Revenue Growth Over Years",
+        "description": "A line chart showing the trend of revenue over the years."
+      }}
+
+    Only suggest visualizations that make sense for the data and the query.
+    """
+
+    for attempt in range(retries + 1):
+        try:
+            # Generate response from the model
+            response = llm.generate(prompt)
+
+            # Load JSON response
+            suggestions = json.loads(response)
+
+            # Validate response structure
+            if isinstance(suggestions, list):
+                valid_suggestions = [
+                    s for s in suggestions if all(k in s for k in ["chart_type", "x_axis", "y_axis"])
+                ]
+                if valid_suggestions:
+                    return valid_suggestions
+                else:
+                    st.warning("⚠️ GPT-4o did not suggest valid visualizations.")
+                    return None
+
+            elif isinstance(suggestions, dict):
+                if all(k in suggestions for k in ["chart_type", "x_axis", "y_axis"]):
+                    return [suggestions]
+                else:
+                    st.warning("⚠️ GPT-4o's suggestion is
+
 
 def add_stats_to_figure(fig, df, y_axis, chart_type):
     """