Update app.py

app.py

@@ -5,18 +5,13 @@ import io
 import re
 import gc
 import os
-from typing import Tuple, Optional
+from typing import Tuple, Optional, List
 
-# ---------- Helper functions ----------
+# ---------- Helpers for uploaded file reading ----------
 def read_uploaded_file(file):
-    """
-    Try multiple ways to get bytes from Gradio upload objects.
-    Returns (bytes_content, filename or None, error_message)
-    """
+    """Try multiple ways to get bytes from Gradio upload objects."""
     if file is None:
         return None, None, "No file uploaded."
-
-    # 1) If object has read() (file-like), use it
     try:
         if hasattr(file, "read"):
             content = file.read()
@@ -24,8 +19,6 @@ def read_uploaded_file(file):
             return content, name, None
     except Exception:
         pass
-
-    # 2) If it's a path-like string, open it
     try:
         if isinstance(file, (str, os.PathLike)):
             path = str(file)
@@ -35,25 +28,18 @@ def read_uploaded_file(file):
                 return content, os.path.basename(path), None
     except Exception:
         pass
-
-    # 3) If it's a dict-like object (some envs), try common keys
     try:
         if isinstance(file, dict):
-            # Many times keys could be 'name' and 'data' or 'content'
             name = file.get("name") or file.get("filename")
             data = file.get("data") or file.get("content") or file.get("bytes")
             if isinstance(data, (bytes, bytearray)):
                 return data, name, None
-            if isinstance(data, str):
-                # maybe base64 or path — try open
-                if os.path.exists(data):
-                    with open(data, "rb") as f:
-                        content = f.read()
-                    return content, name or os.path.basename(data), None
+            if isinstance(data, str) and os.path.exists(data):
+                with open(data, "rb") as f:
+                    content = f.read()
+                return content, name or os.path.basename(data), None
     except Exception:
         pass
-
-    # 4) Fallback: try to get .name and open that path
     try:
         name = getattr(file, "name", None)
         if name and isinstance(name, str) and os.path.exists(name):
@@ -62,22 +48,16 @@ def read_uploaded_file(file):
             return content, os.path.basename(name), None
     except Exception:
         pass
-
     return None, None, "Uploaded file format not supported by this server environment."
 
 def load_file_bytes_to_df(file) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
-    """
-    Read uploaded file bytes into a pandas DataFrame WITHOUT saving to disk.
-    Returns (df, error_message).
-    """
+    """Read bytes and convert to DataFrame (no disk writes)."""
     content, name, err = read_uploaded_file(file)
     if err:
         return None, f"Error reading file: {err}"
     if content is None:
         return None, "No content read from uploaded file."
-
     try:
-        # Basic heuristic for CSV vs Excel
         fname = (name or "").lower()
         if fname.endswith(".csv") or (isinstance(content, (bytes, bytearray)) and b"," in content[:200]):
             df = pd.read_csv(io.BytesIO(content))
@@ -86,7 +66,6 @@ def load_file_bytes_to_df(file) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
     except Exception as e:
         return None, f"Error reading file: {e}"
     finally:
-        # remove raw bytes quickly
         try:
             del content
         except Exception:
@@ -94,18 +73,89 @@ def load_file_bytes_to_df(file) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
         gc.collect()
     return df, None
 
-# ---------- Natural language to action ----------
+# ---------- Column matching in queries ----------
+def find_columns_in_query(columns: List[str], query: str, max_matches: int = 3) -> List[str]:
+    """Return a list of best matching column names from the DataFrame for words in the query."""
+    q = query.lower()
+    found = []
+    # exact word matches first
+    for col in columns:
+        cl = col.lower()
+        # exact full word present
+        if re.search(r"\b" + re.escape(cl) + r"\b", q):
+            found.append(col)
+            if len(found) >= max_matches:
+                return found
+    # partial matches (any token)
+    q_tokens = set(re.findall(r"[a-z0-9_]+", q))
+    for col in columns:
+        if col in found:
+            continue
+        cl = col.lower()
+        col_tokens = set(re.findall(r"[a-z0-9_]+", cl))
+        if q_tokens & col_tokens:
+            found.append(col)
+            if len(found) >= max_matches:
+                return found
+    # fallback: if query contains "department" but no exact column, look for column names containing department
+    for col in columns:
+        if "department" in col.lower() and col not in found:
+            found.append(col)
+            if len(found) >= max_matches:
+                return found
+    return found
+
+# ---------- Aggregation helpers ----------
+def group_count(df: pd.DataFrame, group_col: str, top_n: Optional[int] = None):
+    res = df.groupby(group_col).size().reset_index(name="count").sort_values("count", ascending=False).reset_index(drop=True)
+    if top_n:
+        return res.head(top_n)
+    return res
+
+def group_agg(df: pd.DataFrame, group_col: str, value_col: str, agg: str):
+    if agg in ("mean", "avg", "average"):
+        res = df.groupby(group_col)[value_col].mean().reset_index().rename(columns={value_col: "average"})
+    elif agg in ("sum",):
+        res = df.groupby(group_col)[value_col].sum().reset_index().rename(columns={value_col: "sum"})
+    elif agg in ("max",):
+        res = df.groupby(group_col)[value_col].max().reset_index().rename(columns={value_col: "max"})
+    elif agg in ("min",):
+        res = df.groupby(group_col)[value_col].min().reset_index().rename(columns={value_col: "min"})
+    else:
+        res = df.groupby(group_col)[value_col].agg(agg).reset_index().rename(columns={value_col: agg})
+    return res.sort_values(res.columns[-1], ascending=False).reset_index(drop=True)
+
+def compute_percentage_counts(df: pd.DataFrame, group_col: str):
+    counts = group_count(df, group_col)
+    total = counts["count"].sum()
+    counts["percentage"] = (counts["count"] / total * 100).round(2)
+    return counts
+
+def compute_percentage_of_value(df: pd.DataFrame, group_col: str, value_col: str):
+    # percent share of value_col per group
+    sums = df.groupby(group_col)[value_col].sum().reset_index().rename(columns={value_col: "sum"})
+    total = sums["sum"].sum()
+    sums["percentage"] = (sums["sum"] / total * 100).round(2)
+    return sums.sort_values("sum", ascending=False).reset_index(drop=True)
+
+# ---------- Natural language parser & action ----------
 def simple_nl_to_action(df: pd.DataFrame, query: str):
-    """NL parser that also supports 'wise head count' groupby counting."""
     q = (query or "").strip().lower()
     if q == "":
-        return None, "Please type a question like: 'show columns', 'show first 5 rows', 'describe', 'department wise head count', or 'filter where Region = India'."
+        return None, "Please type a question like: 'department wise head count', 'percentage of employees by department', 'average salary by department', or 'show columns'."
+
+    cols = list(df.columns)
+    matched = find_columns_in_query(cols, q, max_matches=3)  # up to 3 column matches
 
-    # 1) show columns
+    # direct commands
     if "columns" in q or "show columns" in q or "list columns" in q:
-        return pd.DataFrame({"columns": df.columns}), None
+        return pd.DataFrame({"columns": cols}), None
+
+    # overall totals
+    if re.search(r"\b(total|how many|count of rows|number of rows|total employees|total employee)\b", q):
+        return pd.DataFrame({"total_rows": [len(df)]}), None
 
-    # 2) head / first N rows
+    # show first N rows
     m = re.search(r"(first|head)\s*(\d+)?", q)
     if "head" in q or "first" in q:
         n = 5
@@ -113,88 +163,90 @@ def simple_nl_to_action(df: pd.DataFrame, query: str):
             n = int(m.group(2))
         return df.head(n), None
 
-    # 3) describe / summary stats
+    # describe / summary
     if "describe" in q or "summary" in q or "statistics" in q:
         return df.describe(include='all').reset_index(), None
 
-    # 4) show column X
-    m = re.search(r"show (?:column )?([a-z0-9_ ]+)", q)
-    if m:
-        col = m.group(1).strip()
-        matches = [c for c in df.columns if c.lower() == col.lower()]
-        if matches:
-            return df[[matches[0]]].head(100), None
+    # HEADCOUNT / COUNT requests (department wise head count etc.)
+    if any(w in q for w in ["headcount", "head count", "head-count", "headcounts", "head count", "number of employees", "how many", "count by", "count of", "count"]):
+        # If a grouping column is mentioned, use it
+        if matched:
+            group_col = matched[0]
+            # if user mentions percentage as well
+            if "%" in q or "percentage" in q or "percent" in q or "share" in q:
+                return compute_percentage_counts(df, group_col), None
+            # If they asked which has maximum
+            if any(w in q for w in ["most", "maximum", "max", "highest", "where max", "to which"]):
+                counts = group_count(df, group_col)
+                top = counts.head(1)
+                # also show full counts for context
+                summary = counts
+                # build a small output that includes top and summary (we'll return summary; top is first row)
+                return summary, f"Top: {top.iloc[0,0]} with {top.iloc[0,1]} (rows)."
+            # just return counts
+            return group_count(df, group_col), None
         else:
-            return None, f"Column '{col}' not found. Try 'show columns' to see exact names."
+            # no group column mentioned: return total rows
+            return pd.DataFrame({"total_rows": [len(df)]}), None
 
-    # 5) filter pattern
-    m = re.search(r"filter where ([a-z0-9_ ]+?)\s*(=|>|<|>=|<=)\s*'?(?P<val>[^']+?)'?$", q)
-    if m:
-        col_text = m.group(1).strip()
-        op = m.group(2)
-        val = m.group('val').strip()
-        matches = [c for c in df.columns if c.lower() == col_text.lower()]
-        if not matches:
-            return None, f"Column '{col_text}' not found. Use 'show columns' to check names."
-        colname = matches[0]
-        try:
-            if pd.api.types.is_numeric_dtype(df[colname]):
-                val_num = float(val)
-                if op == "=":
-                    res = df[df[colname] == val_num]
-                elif op == ">":
-                    res = df[df[colname] > val_num]
-                elif op == "<":
-                    res = df[df[colname] < val_num]
-                elif op == ">=":
-                    res = df[df[colname] >= val_num]
-                elif op == "<=":
-                    res = df[df[colname] <= val_num]
-                return res.head(200), None
-            else:
-                if op == "=":
-                    res = df[df[colname].astype(str).str.lower() == val.lower()]
-                    return res.head(200), None
-                else:
-                    return None, f"Operator {op} not supported for non-numeric column '{colname}'."
-        except Exception as e:
-            return None, f"Error applying filter: {e}"
-
-    # 6) department-wise / <column> wise head count / count by column
-    # Examples caught: "department wise head count", "count by Department", "headcount by department"
-    if re.search(r"(head[\s-]?count|headcount|count)", q) and (("wise" in q) or (" by " in q) or ("by " in q)):
-        # try "X wise head count" pattern
-        m = re.search(r"([a-z0-9_ ]+?)\s*(?:wise|by)\s*(?:head[\s-]?count|headcount|count)", q)
-        col_candidate = None
-        if m:
-            col_candidate = m.group(1).strip()
+    # AGGREGATION requests (average, mean, sum, max/min of a numeric column grouped by another)
+    if any(w in q for w in ["average", "mean", "avg", "sum", "total", "maximum", "minimum", "max", "min"]):
+        # try to detect grouping and value column
+        if len(matched) >= 2:
+            group_col = matched[0]
+            value_col = matched[1]
+        elif len(matched) == 1:
+            # ambiguous: user mentioned one column. If that's numeric, perhaps they want overall average
+            cand = matched[0]
+            if pd.api.types.is_numeric_dtype(df[cand]):
+                # overall stat
+                if any(w in q for w in ["average", "mean", "avg"]):
+                    return pd.DataFrame({f"overall_{cand}_average": [df[cand].mean()]}), None
+                if "sum" in q or "total" in q:
+                    return pd.DataFrame({f"overall_{cand}_sum": [df[cand].sum()]}), None
+            # else ask for more clarity
+            return None, "I found one column but couldn't tell grouping vs value column. Please ask like 'average Salary by Department' or 'sum Sales by Region'."
         else:
-            # try "count by X"
-            m2 = re.search(r"(?:head[\s-]?count|headcount|count)\s*(?:by\s*)([a-z0-9_ ]+)", q)
-            if m2:
-                col_candidate = m2.group(1).strip()
+            return None, "Please mention columns. Example: 'average Salary by Department' or 'sum Sales by Region'."
+        # determine aggregation type
+        if any(w in q for w in ["average", "mean", "avg"]):
+            return group_agg(df, group_col, value_col, "mean"), None
+        if any(w in q for w in ["sum", "total"]):
+            return group_agg(df, group_col, value_col, "sum"), None
+        if any(w in q for w in ["max", "maximum", "highest"]):
+            return group_agg(df, group_col, value_col, "max"), None
+        if any(w in q for w in ["min", "minimum", "lowest"]):
+            return group_agg(df, group_col, value_col, "min"), None
 
-        if col_candidate:
-            # match to actual column name
-            matches = [c for c in df.columns if c.lower() == col_candidate.lower()]
-            if not matches:
-                # try partial match (contains)
-                partials = [c for c in df.columns if col_candidate.lower() in c.lower()]
-                if partials:
-                    colname = partials[0]
-                else:
-                    return None, f"Column '{col_candidate}' not found. Use 'show columns' to see exact names."
+    # PERCENTAGE requests for a numeric column per group
+    if any(w in q for w in ["percentage", "%", "percent", "share"]):
+        # if two columns mentioned, assume first is group, second is numeric value
+        if len(matched) >= 2:
+            group_col = matched[0]
+            value_col = matched[1]
+            if pd.api.types.is_numeric_dtype(df[value_col]):
+                return compute_percentage_of_value(df, group_col, value_col), None
             else:
-                colname = matches[0]
+                return None, f"Column '{value_col}' is not numeric; cannot compute percentage of values."
+        elif len(matched) == 1:
+            group_col = matched[0]
+            # percent of counts
+            return compute_percentage_counts(df, group_col), None
+        else:
+            return None, "Please mention the group column (and optionally a numeric column). Example: 'percentage of Salary by Department' or 'percentage of employees by Department'."
 
-            try:
-                counts = df.groupby(colname).size().reset_index(name="count").sort_values("count", ascending=False)
-                return counts.reset_index(drop=True), None
-            except Exception as e:
-                return None, f"Error computing counts: {e}"
+    # SHOW specific columns (e.g., 'show Department and Salary')
+    m = re.search(r"show (.+)", q)
+    if m:
+        # try to extract column names from matched list
+        if matched:
+            # if user asked show with two columns, return them
+            return df[matched].head(200), None
+        else:
+            return None, "Couldn't identify columns to show. Use 'show columns' to view exact names."
 
-    # fallback: first 10 rows w/ message
-    return df.head(10), "Couldn't parse exact request — showing first 10 rows. Try: 'show columns', 'describe', 'department wise head count', or 'filter where Column = Value'."
+    # fallback: return first 10 rows with suggestion
+    return df.head(10), "Couldn't parse exact request — showing first 10 rows. Try: 'show columns', 'department wise head count', 'percentage of employees by department', or 'average Salary by Department'."
 
 # ---------- Processing wrapper ----------
 def process(file, query):
@@ -229,11 +281,8 @@ def process(file, query):
     else:
         return None, (msg or "No result")
 
-# ---------- Clear / reset function ----------
+# ---------- Clear / reset ----------
 def clear_all():
-    """
-    Returns Gradio update objects that clear inputs and outputs.
-    """
     return (
         gr.File.update(value=None),
         gr.Textbox.update(value=""),
@@ -243,10 +292,10 @@ def clear_all():
 
 # ---------- Gradio UI ----------
 with gr.Blocks() as demo:
-    gr.Markdown("# Chat-with-CSV — private ephemeral uploads (cleared on Reset)")
+    gr.Markdown("# Chat-with-CSV — enhanced analysis (ephemeral uploads)")
     with gr.Row():
         file_input = gr.File(label="Upload CSV or XLSX (will not be saved)", file_count="single")
-        query_input = gr.Textbox(label="Ask a question (example: 'department wise head count' or 'filter where Country = India')", placeholder="Type your question here")
+        query_input = gr.Textbox(label="Ask a question (examples: 'department wise head count', 'percentage of Salary by Department', 'average Salary by Department')", placeholder="Type your question here")
     with gr.Row():
         submit = gr.Button("Run query")
         clear_btn = gr.Button("Clear / Reset (remove uploaded file & results)")