Semantic column-type inference + override UI [prototype]

aidrin/agentic/retriever.py

@@ -304,7 +304,7 @@ def _compress_one(item: dict[str, Any]) -> str:
             context_texts = [item.get("full_text", "") for item in retrieved]
 
         prompt_context = "\n\n".join(
-            f"[Source: {item.get('source','unknown')}]\n{self._sanitize(text)}"
+            f"[Source: {item.get('source', 'unknown')}]\n{self._sanitize(text)}"
             for item, text in zip(retrieved, context_texts)
         )
 

tests/unit/test_column_roles.py

@@ -0,0 +1,88 @@
+"""Tests for semantic column-role inference (continuous / categorical / identifier)."""
+
+import numpy as np
+import pandas as pd
+
+from web.routes.utils import infer_column_roles
+
+
+def _df(n=1000):
+    rng = np.random.default_rng(42)
+    return pd.DataFrame(
+        {
+            # genuine continuous float measures
+            "pt": rng.normal(50, 10, n).astype("float32"),
+            "eta": rng.normal(0, 2, n).astype("float32"),
+            # low-cardinality integer code -> categorical
+            "flavor": rng.choice([-1, 1, 2, 5, 21], n).astype("int32"),
+            # near-unique integer -> identifier (by cardinality)
+            "jetIndex": np.arange(n, dtype="uint64"),
+            # moderate-cardinality integer with id-like name -> identifier (by name)
+            "eventIndex": rng.integers(0, n // 5, n).astype("uint64"),
+            # genuine string column -> categorical
+            "label": rng.choice(["a", "b", "c"], n),
+        }
+    )
+
+
+def test_continuous_floats():
+    roles = infer_column_roles(_df())
+    assert roles["pt"] == "continuous"
+    assert roles["eta"] == "continuous"
+
+
+def test_low_cardinality_int_is_categorical():
+    roles = infer_column_roles(_df())
+    assert roles["flavor"] == "categorical"
+
+
+def test_near_unique_int_is_identifier():
+    roles = infer_column_roles(_df())
+    assert roles["jetIndex"] == "identifier"
+
+
+def test_idlike_name_with_moderate_cardinality_is_identifier():
+    roles = infer_column_roles(_df())
+    assert roles["eventIndex"] == "identifier"
+
+
+def test_string_column_is_categorical():
+    roles = infer_column_roles(_df())
+    assert roles["label"] == "categorical"
+
+
+def test_user_override_wins():
+    roles = infer_column_roles(_df(), overrides={"flavor": "continuous", "pt": "categorical"})
+    assert roles["flavor"] == "continuous"
+    assert roles["pt"] == "categorical"
+
+
+def test_invalid_override_ignored():
+    roles = infer_column_roles(_df(), overrides={"flavor": "bogus", "missing": "categorical"})
+    assert roles["flavor"] == "categorical"  # heuristic retained
+    assert "missing" not in roles
+
+
+def test_bool_column_is_categorical():
+    df = pd.DataFrame({"flag": [True, False, True, False] * 10})
+    assert infer_column_roles(df)["flag"] == "categorical"
+
+
+def test_empty_dataframe():
+    assert infer_column_roles(pd.DataFrame()) == {}
+
+
+def test_small_dataset_all_distinct_int_is_continuous():
+    # In a tiny dataset every value is trivially distinct — a numeric column
+    # must not be mistaken for an identifier or (via low nunique) a category.
+    df = pd.DataFrame({"age": [25, 30, 35], "active": [True, False, True]})
+    roles = infer_column_roles(df)
+    assert roles["age"] == "continuous"
+    assert roles["active"] == "categorical"
+
+
+def test_identifier_requires_enough_distinct_values():
+    # An id-like name with only a handful of distinct values in a small dataset
+    # should not be forced to "identifier".
+    df = pd.DataFrame({"user_id": [1, 2, 3, 4, 5]})
+    assert infer_column_roles(df)["user_id"] != "identifier"

web/routes/core.py

@@ -18,9 +18,11 @@
 from werkzeug.utils import secure_filename
 from aidrin.file_handling.file_parser import SUPPORTED_FILE_TYPES, READER_MAP
 from web.routes.utils import (
+    categorical_bars,
     clear_all_user_cache,
     ensure_json_serializable,
     get_current_user_id,
+    infer_column_roles,
     load_dataframe,
     summary_histograms,
 )
@@ -56,6 +58,8 @@ def inspector():
             session["uploaded_file_name"] = display_name
             session["uploaded_file_path"] = file_path
             session["uploaded_file_type"] = request.form.get("fileTypeSelector")
+            # Column-role overrides are per-dataset; drop any from a prior file.
+            session.pop("column_roles", None)
 
             # Track files this session created so /clear removes only these,
             # never files belonging to other concurrent sessions.
@@ -349,43 +353,40 @@ def summary_statistics():
         if load_error:
             return jsonify({"success": False, "message": load_error}), 200
 
-        # Restrict to numeric columns: describe() would otherwise fall back to
-        # object-column stats (top/freq are strings) and the numeric formatter
-        # below would fail on them. Datasets with no numerical features simply
-        # get an empty numerical summary (issue #125).
-        numeric_df = df.select_dtypes(include="number")
+        # Classify columns by *semantic role* rather than raw dtype, so numeric
+        # codes (e.g. a class label) and identifiers (e.g. a near-unique index)
+        # are not summarized as if they were continuous measurements. User
+        # overrides stored in the session take precedence over the heuristic.
+        role_overrides = session.get("column_roles", {})
+        column_roles = infer_column_roles(df, role_overrides)
+        continuous_columns = [c for c in df.columns if column_roles[c] == "continuous"]
+        categorical_columns = [c for c in df.columns if column_roles[c] == "categorical"]
+        identifier_columns = [c for c in df.columns if column_roles[c] == "identifier"]
+
+        # describe() runs only on continuous columns. Datasets with no continuous
+        # features simply get an empty numerical summary (issue #125).
+        numeric_df = df[continuous_columns]
         if numeric_df.shape[1] == 0:
             summary_statistics = {}
         else:
             summary_statistics = numeric_df.describe().map(
                 lambda x: round(x, 2) if x == 0 or abs(x) >= 0.001 else f"{x:.2e}"
             ).to_dict()
 
-        histograms = summary_histograms(df)
-
-        # Booleans are treated as categorical (they're excluded from describe()
-        # and select_dtypes("number") above, so they belong with the categorical
-        # summary, not the numerical one).
-        numerical_columns = [
-            col for col, dtype in df.dtypes.items()
-            if pd.api.types.is_numeric_dtype(dtype) and not pd.api.types.is_bool_dtype(dtype)
-        ]
-        categorical_columns = [
-            col for col, dtype in df.dtypes.items()
-            if pd.api.types.is_string_dtype(dtype) or pd.api.types.is_bool_dtype(dtype)
-        ]
-        all_features = numerical_columns + categorical_columns
+        # KDE curves for continuous columns; value-count bars for categorical
+        # ones (a KDE of a discrete code is meaningless). Identifiers get neither.
+        histograms = summary_histograms(df, columns=continuous_columns)
+        categorical_histograms = categorical_bars(df, categorical_columns)
 
         for v in summary_statistics.values():
             for old_key in list(v.keys()):
                 if old_key in ["25%", "50%", "75%"]:
                     new_key = old_key.replace("%", "th percentile")
                     v[new_key] = v.pop(old_key)
 
-        # Per-column summary for categorical features (describe() above only
-        # covers numerical columns). Reports non-null count, distinct values,
-        # the most frequent value, its frequency, and that frequency as a
-        # percentage of the non-null values.
+        # Per-column summary for categorical features (numeric codes included).
+        # Reports non-null count, distinct values, the most frequent value, its
+        # frequency, and that frequency as a percentage of the non-null values.
         categorical_summary = {}
         for col in categorical_columns:
             counts = df[col].value_counts(dropna=True)
@@ -399,24 +400,65 @@ def summary_statistics():
                 "freq_pct": round(freq / count * 100, 1) if count else 0.0,
             }
 
+        # Compact per-identifier note (excluded from stats/plots).
+        identifier_summary = {
+            str(col): {
+                "count": int(df[col].notna().sum()),
+                "unique": int(df[col].nunique(dropna=True)),
+            }
+            for col in identifier_columns
+        }
+
+        all_features = continuous_columns + categorical_columns
+
         response_data = ensure_json_serializable({
             "success": True,
             "message": "File uploaded successfully",
             "records_count": len(df),
             "features_count": len(df.columns),
             "categorical_features": list(categorical_columns),
-            "numerical_features": list(numerical_columns),
+            "numerical_features": list(continuous_columns),
+            "identifier_features": list(identifier_columns),
             "all_features": all_features,
+            "column_roles": {str(c): r for c, r in column_roles.items()},
             "summary_statistics": summary_statistics,
             "categorical_summary": categorical_summary,
+            "identifier_summary": identifier_summary,
             "histograms": histograms,
+            "categorical_histograms": categorical_histograms,
         })
         return jsonify(response_data)
     except Exception as e:
         file_upload_time_log.error("Error computing summary statistics: %s", e, exc_info=True)
         return jsonify({"success": False, "message": "An internal error occurred"})
 
 
+@core_bp.route("/column-roles", methods=["POST"])
+def set_column_roles():
+    """Persist user overrides of inferred column roles for the current session.
+
+    Accepts a JSON body ``{"column_roles": {col: role}}`` (or a bare
+    ``{col: role}`` map). The full override set replaces any previous one, so
+    sending ``{}`` resets to the inferred defaults. The Data Overview re-fetches
+    ``/summary-statistics`` afterwards to recompute with the new roles.
+    """
+    try:
+        data = request.get_json(silent=True) or {}
+        overrides = data.get("column_roles", data)
+        if not isinstance(overrides, dict):
+            return jsonify({"success": False, "message": "Invalid payload"}), 200
+        valid_roles = {"continuous", "categorical", "identifier"}
+        cleaned = {str(c): r for c, r in overrides.items() if r in valid_roles}
+        session["column_roles"] = cleaned
+        # Overrides change which columns each metric treats as numeric/categorical,
+        # so any cached metric results are now stale.
+        clear_all_user_cache()
+        return jsonify({"success": True, "column_roles": cleaned})
+    except Exception as e:
+        file_upload_time_log.error("Error setting column roles: %s", e, exc_info=True)
+        return jsonify({"success": False, "message": "An internal error occurred"}), 200
+
+
 @core_bp.route("/feature-set", methods=["POST"])
 def extract_features():
     try: