Update utils.py

utils.py

@@ -2,7 +2,29 @@ from huggingface_hub import InferenceClient
 import os
 import pandas as pd
 
+# ------------------- Helper Functions -------------------
+
+def split_multi_value_columns(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Automatically splits any column that contains multiple comma-separated values
+    into separate columns.
+    """
+    new_df = df.copy()
+    for col in df.columns:
+        # Check if the first non-null row contains a comma
+        sample = df[col].dropna().iloc[0] if not df[col].dropna().empty else ""
+        if isinstance(sample, str) and "," in sample:
+            # Split the column into multiple columns
+            split_cols = df[col].str.split(",", expand=True)
+            split_cols = split_cols.rename(columns=lambda i: f"{col}_{i+1}")
+            new_df = new_df.drop(columns=[col]).join(split_cols)
+    return new_df
+
+
 def summarize_dataframe(df: pd.DataFrame, max_rows: int = 30) -> str:
+    """
+    Returns a text summary of the dataframe for LLM prompts.
+    """
     summary = f"Columns: {', '.join(df.columns)}\n\n"
     if len(df) > max_rows:
         sample = df.sample(max_rows, random_state=42)
@@ -13,23 +35,32 @@ def summarize_dataframe(df: pd.DataFrame, max_rows: int = 30) -> str:
     summary += sample.to_string(index=False)
     return summary
 
+# ------------------- Main Query Function -------------------
 
 def query_agent(df: pd.DataFrame, query: str) -> str:
+    """
+    Analyzes a dataframe to answer queries. Supports:
+    - Direct analysis of most common values (single or multiple columns)
+    - Fallback to LLM using google/gemma-2b-it
+    """
+    # Automatically split multi-value columns
+    df = split_multi_value_columns(df)
+
     query_lower = query.lower()
 
-    # ----------------- Direct Analysis for Most Common -----------------
+    # -------- Direct Analysis for Most Common Values --------
     try:
         if "most common" in query_lower or "most frequent" in query_lower:
-            # Look for multiple columns in query
+            # Find all columns mentioned in the query
             cols_in_query = [col for col in df.columns if col.lower() in query_lower]
-            
+
             if len(cols_in_query) == 1:
                 col = cols_in_query[0]
                 value = df[col].mode()[0]
                 return f"The most common value in column '{col}' is '{value}'."
-            
+
             elif len(cols_in_query) > 1:
-                # Compute most common combination of values across the columns
+                # Most common combination across multiple columns
                 combo_series = df[cols_in_query].apply(lambda row: tuple(row), axis=1)
                 most_common_combo = combo_series.mode()[0]
                 combo_str = ", ".join(f"{col}={val}" for col, val in zip(cols_in_query, most_common_combo))
@@ -38,7 +69,7 @@ def query_agent(df: pd.DataFrame, query: str) -> str:
     except Exception as e:
         print("Direct analysis failed:", e)
 
-    # ----------------- Use LLM if direct analysis fails -----------------
+    # -------- Use LLM Fallback if Direct Analysis Fails --------
     data_text = summarize_dataframe(df)
     prompt = f"""
 You are a data analysis assistant with expertise in statistics and data interpretation.
@@ -54,7 +85,7 @@ Question:
 Answer (with explanation):
 """
 
-    # Initialize client with explicit provider
+    # Initialize Hugging Face InferenceClient with explicit provider
     client = InferenceClient(
         model="google/gemma-2b-it",
         provider="hf-inference",