Update app.py

app.py

@@ -192,21 +192,45 @@ COLUMN_SYNONYMS = {
 
 
 # Fuzzy matcher for mapping query terms to dataset columns
-def fuzzy_match_columns(query, n=2):
+def fuzzy_match_columns(query):
     query = query.lower()
     all_synonyms = {synonym: col for col, synonyms in COLUMN_SYNONYMS.items() for synonym in synonyms}
-    
     words = query.replace("and", "").replace("vs", "").replace("by", "").split()
-    
+
     matched_columns = []
     for word in words:
-        matches = get_close_matches(word, all_synonyms.keys(), n=n, cutoff=0.6)
-        for match in matches:
-            matched_columns.append(all_synonyms[match])
+        matches = get_close_matches(word, all_synonyms.keys(), n=1, cutoff=0.6)
+        matched_columns.extend([all_synonyms[match] for match in matches])
 
     return list(dict.fromkeys(matched_columns))
 
-# Statistical annotations for plots
+# Ask LLM to suggest relevant columns if fuzzy matching fails
+def ask_llm_for_columns(query, llm, df):
+    columns = ', '.join(df.columns)
+    prompt = f"""
+    Analyze this user query and suggest the most relevant dataset columns for visualization.
+
+    Query: "{query}"
+
+    Available Columns: {columns}
+
+    Respond in this JSON format:
+    {{
+      "x_axis": "column_name",
+      "y_axis": "column_name",
+      "group_by": "optional_column_name"
+    }}
+    """
+
+    response = llm.generate(prompt)
+    try:
+        suggestion = json.loads(response)
+        return suggestion
+    except json.JSONDecodeError:
+        st.error("⚠️ Failed to interpret AI response. Please refine your query.")
+        return None
+
+# Add min, max, and average salary annotations to the chart
 def add_stats_to_figure(fig, df, y_axis):
     min_salary = df[y_axis].min()
     max_salary = df[y_axis].max()
@@ -222,18 +246,25 @@ def add_stats_to_figure(fig, df, y_axis):
     )
     return fig
 
-# Visualization generator
-def generate_visual_from_query(query, df):
+# Unified visualization function with LLM fallback
+def generate_visual_from_query(query, df, llm=None):
     try:
         matched_columns = fuzzy_match_columns(query)
 
-        # Detect and handle multiple grouping columns
+        # Fallback to LLM if fuzzy matching fails
+        if not matched_columns and llm:
+            st.info("🤖 No match found. Asking AI for suggestions...")
+            suggestion = ask_llm_for_columns(query, llm, df)
+            if suggestion:
+                matched_columns = [suggestion.get("x_axis"), suggestion.get("group_by")]
+
+        # Handle cases when we have columns to plot
         if len(matched_columns) >= 2:
             x_axis, group_by = matched_columns[0], matched_columns[1]
         elif len(matched_columns) == 1:
             x_axis, group_by = matched_columns[0], None
         else:
-            st.warning("❓ No matching columns found. Try rephrasing your query.")
+            st.warning("❓ No matching columns found. Please refine your query.")
             return None
 
         # Handle distribution queries
@@ -258,16 +289,15 @@ def generate_visual_from_query(query, df):
                           title=f"Salary Trend Over Years by {x_axis.replace('_', ' ').title()}")
             return add_stats_to_figure(fig, df, "salary_in_usd")
 
-        # Handle remote work queries
+        # Handle remote work impact
         elif "remote" in query:
             grouped_df = df.groupby(["remote_ratio"] + ([group_by] if group_by else []))["salary_in_usd"].mean().reset_index()
             fig = px.bar(grouped_df, x="remote_ratio", y="salary_in_usd", color=group_by,
                          title="Remote Work Impact on Salary")
             return add_stats_to_figure(fig, df, "salary_in_usd")
 
-        # Default behavior if query doesn't match anything specific
         else:
-            st.warning("❓ No suitable visualization generated. Try refining your query.")
+            st.warning("⚠️ No suitable visualization generated. Please refine your query.")
             return None
 
     except Exception as e:
@@ -275,6 +305,7 @@ def generate_visual_from_query(query, df):
         return None
 
 
+
 # SQL-RAG Analysis
 if st.session_state.df is not None:
     temp_dir = tempfile.TemporaryDirectory()