Update app.py

app.py

@@ -191,15 +191,99 @@ COLUMN_SYNONYMS = {
 }
 
 
-# Helper function to map user query terms to dataset columns
-#def map_query_to_column(query):
-#    for col, synonyms in COLUMN_SYNONYMS.items():
-#        for term in synonyms:
-#            if term in query:
-#                return col
-#    return None
-
-def map_query_to_column(query):
+# Fuzzy match to map query terms to dataset columns
+def fuzzy_match_columns(query, n=2):
+    query = query.lower()
+    all_synonyms = {synonym: col for col, synonyms in COLUMN_SYNONYMS.items() for synonym in synonyms}
+    words = query.replace("and", "").replace("vs", "").split()  # Remove "and"/"vs" for better matching
+    
+    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])
+
+    # Remove duplicates while preserving order
+    matched_columns = list(dict.fromkeys(matched_columns))
+    return matched_columns
+
+# Visualization generator with dynamic groupby handling
+def generate_visual_from_query(query, df):
+    try:
+        # Step 1: Fuzzy match columns mentioned in the query
+        matched_columns = fuzzy_match_columns(query)
+
+        # Step 2: Detect groupby intent (handling "and", "vs", "by")
+        if "and" in query or "vs" in query or "by" in query or len(matched_columns) > 1:
+            if len(matched_columns) >= 2:
+                x_axis = matched_columns[0]
+                group_by = matched_columns[1]
+            else:
+                x_axis, group_by = matched_columns[0], None
+        else:
+            x_axis = matched_columns[0] if matched_columns else None
+            group_by = None
+
+        # Step 3: Visualization logic
+        if "distribution" in query and x_axis:
+            fig = px.box(df, x=x_axis, y="salary_in_usd", color=group_by,
+                         title=f"Salary Distribution by {x_axis.replace('_', ' ').title()}"
+                               + (f" and {group_by.replace('_', ' ').title()}" if group_by else ""))
+            return fig
+
+        elif "average" in query or "mean" in query:
+            grouped_df = df.groupby([x_axis] + ([group_by] if group_by else []))["salary_in_usd"].mean().reset_index()
+            fig = px.bar(grouped_df, x=x_axis, y="salary_in_usd", color=group_by,
+                         barmode="group",
+                         title=f"Average Salary by {x_axis.replace('_', ' ').title()}"
+                               + (f" and {group_by.replace('_', ' ').title()}" if group_by else ""))
+            return fig
+
+        elif "trend" in query and "work_year" in df.columns and x_axis:
+            grouped_df = df.groupby(["work_year", x_axis])["salary_in_usd"].mean().reset_index()
+            fig = px.line(grouped_df, x="work_year", y="salary_in_usd", color=x_axis,
+                          title=f"Salary Trend over Years by {x_axis.replace('_', ' ').title()}")
+            return fig
+
+        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,
+                         barmode="group", title="Remote Work Impact on Salary")
+            return fig
+
+        elif "company size" in query:
+            grouped_df = df.groupby(["company_size"] + ([group_by] if group_by else []))["salary_in_usd"].mean().reset_index()
+            fig = px.bar(grouped_df, x="company_size", y="salary_in_usd", color=group_by,
+                         title=f"Salary by Company Size"
+                               + (f" and {group_by.replace('_', ' ').title()}" if group_by else ""))
+            return fig
+
+        elif "country" in query or "location" in query:
+            grouped_df = df.groupby(["employee_residence"] + ([group_by] if group_by else []))["salary_in_usd"].mean().reset_index()
+            fig = px.bar(grouped_df, x="employee_residence", y="salary_in_usd", color=group_by,
+                         title=f"Salary by Employee Residence"
+                               + (f" and {group_by.replace('_', ' ').title()}" if group_by else ""))
+            return fig
+
+        else:
+            st.warning("❓ No suitable visualization detected. Please refine your query.")
+            return None
+
+    except Exception as e:
+        st.error(f"Error generating visualization: {e}")
+        return None
+
+
+
+
+
+
+
+
+
+
+
+"""def map_query_to_column(query):
     query = query.lower()
     all_synonyms = {synonym: col for col, synonyms in COLUMN_SYNONYMS.items() for synonym in synonyms}
     matches = get_close_matches(query, all_synonyms.keys(), n=1, cutoff=0.6)
@@ -210,10 +294,10 @@ def map_query_to_column(query):
         for col, synonyms in COLUMN_SYNONYMS.items():
             if any(term in query for term in synonyms):
                 return col
-    return None
+    return None"""
 
 
-# Visualization generator with synonym handling
+"""# Visualization generator with synonym handling
 def generate_visual_from_query(query, df):
     try:
         query = query.lower()
@@ -253,7 +337,7 @@ def generate_visual_from_query(query, df):
 
     except Exception as e:
         st.error(f"Error generating visualization: {e}")
-        return None
+        return None"""
 
 # SQL-RAG Analysis
 if st.session_state.df is not None: