Update app.py

app.py

@@ -4,31 +4,89 @@ import pandas as pd
 import io
 import re
 import gc
+import os
 from typing import Tuple, Optional
 
 # ---------- Helper functions ----------
+def read_uploaded_file(file):
+    """
+    Try multiple ways to get bytes from Gradio upload objects.
+    Returns (bytes_content, filename or None, error_message)
+    """
+    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()
+            name = getattr(file, "name", None)
+            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)
+            if os.path.exists(path):
+                with open(path, "rb") as f:
+                    content = f.read()
+                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
+    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):
+            with open(name, "rb") as f:
+                content = f.read()
+            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 DataFrame WITHOUT saving to disk.
+    Read uploaded file bytes into a pandas DataFrame WITHOUT saving to disk.
     Returns (df, error_message).
     """
-    if file is None:
-        return None, "No file uploaded."
+    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:
-        # file is a TemporaryFile-like object in Gradio; read bytes once
-        content = file.read()
-        # use BytesIO to load into pandas
-        # decide by filename or by sniffing bytes
-        name = getattr(file, "name", "") or ""
         # Basic heuristic for CSV vs Excel
-        if name.lower().endswith(".csv") or (isinstance(content, (bytes, bytearray)) and b"," in content[:200]):
+        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))
         else:
             df = pd.read_excel(io.BytesIO(content))
     except Exception as e:
         return None, f"Error reading file: {e}"
     finally:
-        # Immediately try to remove raw bytes if present to minimize memory lifetime
+        # remove raw bytes quickly
         try:
             del content
         except Exception:
@@ -36,15 +94,18 @@ def load_file_bytes_to_df(file) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
         gc.collect()
     return df, None
 
+# ---------- Natural language to action ----------
 def simple_nl_to_action(df: pd.DataFrame, query: str):
-    """Very small NL parser that recognizes a few common requests and returns a DataFrame or (None, message)."""
+    """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 column sales', or 'filter where Region = India and Year >= 2021'."
+        return None, "Please type a question like: 'show columns', 'show first 5 rows', 'describe', 'department wise head count', or 'filter where Region = India'."
 
+    # 1) show columns
     if "columns" in q or "show columns" in q or "list columns" in q:
         return pd.DataFrame({"columns": df.columns}), None
 
+    # 2) head / first N rows
     m = re.search(r"(first|head)\s*(\d+)?", q)
     if "head" in q or "first" in q:
         n = 5
@@ -52,18 +113,21 @@ def simple_nl_to_action(df: pd.DataFrame, query: str):
             n = int(m.group(2))
         return df.head(n), None
 
+    # 3) describe / summary stats
     if "describe" in q or "summary" in q or "statistics" in q:
         return df.describe(include='all').reset_index(), None
 
-    m = re.search(r"show (column )?([a-z0-9_ ]+)", q)
+    # 4) show column X
+    m = re.search(r"show (?:column )?([a-z0-9_ ]+)", q)
     if m:
-        col = m.group(2).strip()
+        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
         else:
             return None, f"Column '{col}' not found. Try 'show columns' to see exact names."
 
+    # 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()
@@ -96,11 +160,44 @@ def simple_nl_to_action(df: pd.DataFrame, query: str):
         except Exception as e:
             return None, f"Error applying filter: {e}"
 
-    return df.head(10), "Couldn't parse exact request — showing first 10 rows. Try: 'show columns', 'show first 5 rows', 'describe', or 'filter where Column = Value'."
+    # 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()
+        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()
+
+        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."
+            else:
+                colname = matches[0]
+
+            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}"
+
+    # 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'."
 
 # ---------- Processing wrapper ----------
 def process(file, query):
-    # Load into memory-only DataFrame
     df, err = load_file_bytes_to_df(file)
     if err:
         try:
@@ -110,7 +207,6 @@ def process(file, query):
         gc.collect()
         return None, err
 
-    # Run the NLP-to-action
     try:
         res, msg = simple_nl_to_action(df, query)
         if isinstance(res, pd.DataFrame):
@@ -121,7 +217,6 @@ def process(file, query):
         out_df = None
         msg = f"Error while processing: {e}"
 
-    # Remove references to large objects immediately
     try:
         del df
         del file
@@ -151,11 +246,10 @@ with gr.Blocks() as demo:
     gr.Markdown("# Chat-with-CSV — private ephemeral uploads (cleared on Reset)")
     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: 'show columns' or 'filter where Country = India')", placeholder="Type your question here")
+        query_input = gr.Textbox(label="Ask a question (example: 'department wise head count' or 'filter where Country = India')", placeholder="Type your question here")
     with gr.Row():
         submit = gr.Button("Run query")
         clear_btn = gr.Button("Clear / Reset (remove uploaded file & results)")
-    # Use headers=None to be compatible with the Gradio version in Spaces
     output_table = gr.Dataframe(headers=None, label="Result table")
     status = gr.Textbox(label="Status / Messages", interactive=False)