Predict Module

The predict/ module is Eunice’s supervised/unsupervised ML surface: dataset profiling, LLM-driven model advisement, training with SHAP explanations, per-row scoring, and a small REST surface for the sidebar page. It is exercised primarily through the chat, train_model, get_model_info, explain_prediction, list_trained_models, find_model, but the same plumbing backs the /api/v1/predict/* endpoints.

For how training fits into the broader tool-dispatch loop, see Engine and the tool system overview. For the analyze_data namespace that train_model reads from, see Tools.

Module Files

File	Responsibility
predict/trainer.py	ModelTrainer, supervised/unsupervised training, evaluation, SHAP, scoring, on-disk persistence
predict/profiler.py	DataProfiler, per-column statistical profile for the advisor
predict/advisor.py	ModelAdvisor, LLM (fast role) recommends target/algorithm/features from a profile
predict/router.py	/api/v1/predict/* endpoints (upload, recommend, train, score, explain, datasets, models) plus tabular/spatial upload handling
tools/prediction.py	Tool handlers _exec_train_model, _exec_get_model_info, _exec_explain_prediction, _exec_list_trained_models, _exec_find_model

The trained-model registry is a shared mutable dict from the DI container (container.predict_models()), keyed by "{tenant_id}:{model_id}" for isolation. ModelTrainer writes entries during training, tools/prediction.py and predict/router.py read them; on restart ModelTrainer.load_persisted_models() re-hydrates from models/{tenant_id}/*.meta.json.

trainer.py, ModelTrainer

ModelTrainer is wired as a DI singleton and takes a CostTracker so training operations show up in usage_log (predict/trainer.py:50-53). The pickle root defaults to models/ but is configurable.

Supported Algorithms

SUPPORTED_ALGORITHMS (predict/trainer.py:28-33) maps tool-facing keys to estimator classes:

Key	Class	Task
xgboost_classifier	xgboost.XGBClassifier	Supervised classification
xgboost_regressor	xgboost.XGBRegressor	Supervised regression
random_forest	sklearn.ensemble.RandomForestClassifier	Supervised classification
isolation_forest	sklearn.ensemble.IsolationForest	Unsupervised anomaly detection

train(df, target_column, algorithm, tenant_id, ...) -> dict

ModelTrainer.train is one function from predict/trainer.py:55-273. Order matters; the steps below mirror the source:

1. Feature selection. Use feature_columns if supplied, otherwise everything except target_column.
2. Auto-encode non-numeric features. Anything outside np.number | bool is passed through a LabelEncoder with _missing_ as the missing sentinel. Fitted encoders are kept in fitted_encoders so scoring uses the same mapping (predict/trainer.py:84-97). The Arrow-backed string dtypes from the catalog are caught here too, the check is “not already numeric/bool”, not a denylist of dtype names.
3. Drop the unencodable. Anything still non-numeric after encoding is dropped with an info log. If nothing usable remains, the trainer raises a ValueError that names the columns and points the user at analyze_data for preparation (predict/trainer.py:107-114).
4. Auto-encode non-numeric targets for classifiers (e.g. Y/N → 0/1) via a second LabelEncoder, logging the class list (predict/trainer.py:121-129).
5. Median-impute missing values in X (predict/trainer.py).
6. Train/test split. 80/20 with random_state=42; isolation_forest uses all data because it’s unsupervised (predict/trainer.py:137-145).
7. Build the model with defaults: XGBoost, n_estimators=100, max_depth=6, learning_rate=0.1, random_state=42; Random Forest, n_estimators=100, random_state=42; Isolation Forest, contamination=0.1. The classifier eval_metric is picked dynamically: mlogloss for 3+ classes, logloss for binary, because XGBoost auto-selects multi:softprob for >2 classes and logloss against an n_samples × n_classes prediction crashes (predict/trainer.py:148-169).
8. Fit, with progress. For XGBoost the trainer attaches a TrainingCallback that calls progress_callback every 10 rounds with {ds}-{metric}: value; the callbacks are explicitly cleared off the model before evaluation so they don’t leak into cross-validation or the pickle (predict/trainer.py:186-213).
9. Evaluate. Classifiers: accuracy, precision/recall/F1 (weighted, zero_division=0), AUC-ROC for binary, confusion matrix. Regressors: RMSE, MAE, R². Isolation Forest: anomaly count + percentage (predict/trainer.py:275-311).
10. Cross-validate (5-fold, accuracy or r2) when len(X) >= 10 and not isolation-forest.
11. SHAP. shap.TreeExplainer over the first 200 test rows; for list-valued outputs (binary) take class 1; rank features by mean absolute SHAP. Isolation Forest skips SHAP with an explanatory note (predict/trainer.py:313-347).
12. Score every row. _generate_scores produces up to 50 entries sorted by risk_score descending. Classifiers with predict_proba use class-1 probability and bucket HIGH ≥ 0.7 / MEDIUM ≥ 0.4 / LOW < 0.4; regressors use the raw prediction with a 75th-percentile cutoff. Isolation Forest normalizes decision_function to 0-1 and tags ANOMALY/NORMAL. Score entries carry up to two identifier columns picked by _pick_identifier_columns (name contains id, identifier, name, code, or is id/*_id; falls back to the first object-typed column) so the model output is human-readable without bloating every row (predict/trainer.py:349-442).
13. Persist a model bundle. _save_model writes models/{tenant_id}/{model_id}.pkl containing {"model": estimator, "preprocessing": fitted_encoders} (predict/trainer.py:444-461). Older pickles were bare estimators, load sites check isinstance(bundle, dict).
14. Record cost. CostRecord(module="predict", operation="model_training", cost_usd=0.0, metadata={"elapsed_seconds", "rows", "features"}). Local compute is free, but the row gives FinOps a time series.

The returned dict carries model_id, algorithm, target_column, feature_columns, metrics, shap_summary, scores, training_time_seconds, rows_trained, rows_tested, model_path, dataset_id. The tool-handler layer in tools/prediction.py adds _dataset_id, _tenant_id, _profile, user_goal, source_dataset_name, created_at before stashing into the shared registry (see “Trained-model registry” below).

progress_callback

Optional Callable[[str], None] invoked at every stage label and every 10 XGBoost rounds. The chat handler wires this to the display-handle status() method so the thinking stream renders training progress live.

explain_entity(model_path, df, entity_index, feature_columns) -> dict

Loads the pickle, runs SHAP on a single row, returns explanations sorted by abs(shap_value) descending with feature, value, shap_value, and direction (increases_risk/decreases_risk), plus the explainer’s base_value (predict/trainer.py:505-540).

explain_entity unpickles with pickle.load(f) and unwraps the {"model": ..., "preprocessing": ...} bundle dict shape via bundle["model"] if isinstance(bundle, dict) else bundle (predict/trainer.py:524), this used to be a gap (scoring unwrapped the bundle, explain didn’t, so SHAP got handed a dict and raised AttributeError: 'dict' object has no attribute 'predict') but the fix landed; the isinstance check also keeps old bare-model pickles loadable.

load_persisted_models() -> dict

Called on startup. Walks models/{tenant_id}/*.meta.json, restores each into a tenant-prefixed registry key "{tenant_id}:{model_id}", tags _tenant_id on the result. Skips malformed JSON silently (predict/trainer.py:477-503).

save_model_metadata(result)

Writes a .meta.json next to the pickle, stripping non-serializable df and capping scores to 50 entries (predict/trainer.py:463-475). This is what load_persisted_models reads back.

profiler.py, DataProfiler

DataProfiler.profile(df, filename) returns the profile that drives ModelAdvisor. Every column gets a logical type plus type-specific stats; all floats are run through sanitize_for_json() so NaN/Inf can be safely serialized (predict/profiler.py:17-31).

Column Type Detection

_profile_column (predict/profiler.py:33-99):

Type	Detection	Extra fields
boolean	pd.api.types.is_bool_dtype	true_pct
datetime	is_datetime64_any_dtype	min, max
datetime (epoch)	Numeric in (1_000_000_000, 4_102_444_800)	min, max, detected_as="unix_epoch"
numeric	Remaining numeric	min, max, mean, median, std, skew, outlier_pct OR distribution="skewed" if >15% beyond IQR fences
categorical	Everything else	top_values (top 5 by frequency)

Note that the in-code logic does NOT split string columns by cardinality into categorical vs text, anything non-numeric, non-bool, non-datetime falls into categorical with a top-5 value-counts dict. The “skewed” branch deliberately suppresses an outlier_pct reading when >15% of values look like outliers, that’s a distribution shape, not a quality flag.

Every column also gets up to 3 sample_values (deterministic, random_state=42) coerced through _safe_value to convert numpy scalars and turn NaN/Inf into None.

Profile Output

{
    "filename": str,
    "rows": int,
    "columns": int,
    "column_profiles": [{
        "name": str,
        "type": "boolean" | "datetime" | "numeric" | "categorical",
        "dtype": str,
        "null_count": int, "null_pct": float,
        "unique_count": int, "cardinality_pct": float,
        # type-specific keys above
        "sample_values": [Any, ...],
    }],
    "memory_mb": float,
}

advisor.py, ModelAdvisor

ModelAdvisor.recommend(profile, tenant_id, user_goal) is an async method that:

1. Anonymizes the profile, strips sample_values, keeps schema-level stats only (predict/advisor.py:82-103). The advisor never sees raw data values.
2. Loads the fast role from model_config via get_model_for_role("fast") and the MODEL_ADVISOR system prompt from shared/prompts.py (predict/advisor.py:20-31). max_tokens is min(settings.max_tokens_advisor, fast_config["max_output_tokens"]), never a hardcoded literal.
3. Calls the LLM through create_llm_client_for_role("fast") so provider routing is centralized.
4. Records cost via CostTracker.record(CostRecord(module="predict", operation="model_recommendation", ...)) with tokens_in/out, cache-read/write tokens, and cost_usd computed by calculate_cost() (predict/advisor.py:46-63).
5. Parses JSON with extract_json_from_response(). On failure, returns a sentinel {target_column: None, algorithm: "xgboost_classifier", reasoning: <raw answer>} so the caller can still log the LLM’s prose explanation (predict/advisor.py:66-75).

Recommendation shape returned to the caller:

{
    "target_column": str | None,
    "feature_columns": list[str],
    "algorithm": str,
    "reasoning": str,
    "prep_steps": list[str],
    "cost": float, "tokens_in": int, "tokens_out": int,
}

Tool Handlers, tools/prediction.py

The chat-facing surface. Five handlers, all returning ToolResult. _container is the DI container, injected by the app on startup via set_container().

_exec_train_model, auto-routing data

The auto-data-routing logic that makes “train a risk model on this dataset” Just Work lives in tools/prediction.py:67-322. Source-of-truth order:

1. Explicit dataset_id in context → catalog.get_entry_by_id.
2. Otherwise catalog.default_dataset() (most-recent non-internal entry; never a _dream_* frame).
3. If dataset_name was passed as an input:
   • First try catalog.exists(name) → catalog.load(name).
   • If not in the catalog, check the analyze_data namespace via tools.analysis.get_namespace(conversation_id) (tools/analysis.py:379-381). If the name resolves to a pd.DataFrame in scratch, use it (_source_name is tagged "… (analyze_data namespace)").
   • If neither resolves, build a structured error message that lists scratch DataFrames AND catalog datasets so the LLM can pick a real name instead of guessing (tools/prediction.py:121-165). The catalog and scratch partitions are kept separate because they share the same dict, without the split, scratch frames get hidden by catalog entries that sort first by insertion order.

This is what “auto-use of cleaned data” means in practice, anything df_cleaned = df[df['valid'] == True] produced in analyze_data is reachable by name; nothing about the underlying training pipeline cares whether the DataFrame came from a catalog arrow_table or a namespace variable.

After data resolution:

• The pre-exclusion DataFrame is held as df_full so the eventual scored_{target} dataset can include lat/lon, IDs, and every original column even when the model only trained on a numeric subset.
• If target_column is explicit and present, the advisor LLM call is skipped (cost=0) and a sensible default algorithm is chosen by target cardinality (xgboost_regressor if nunique() > 20, else xgboost_classifier).
• Progress is wired through the display handle if present (display.progress(msg, display_id="train_progress")) otherwise the legacy progress_callback.
• After trainer.train(), the handler unpickles the saved bundle, applies the SAME LabelEncoder instances to map unknown categorical values to _missing_ (sentinel constant _MISSING_SENTINEL lives at tools/prediction.py:16), produces a fully-scored DataFrame, and saves it as scored_{target} in the catalog so the model can interact with all rows via analyze_data instead of being capped at the top 50 in the card.
• The returned ToolResult.content is a plain-English summary from _summarize_results, accuracy band, confusion-matrix breakdown, top-3 SHAP features, and a pointer at scored_{target} (tools/prediction.py:27-64).
• The training_results card carries the top 20 scores. The datasets channel hands the full scored_{target} DataFrame back to the orchestrator for registration.

Async-in-sync: _exec_train_model is a sync handler that needs an await advisor.recommend(...) inside. It detects whether an event loop is running and either calls asyncio.run directly or submits to a ThreadPoolExecutor (tools/prediction.py:316-322).

_exec_get_model_info

Supports info_type ∈ {"scores", "importance", "metrics", "all"} and respects an optional model_id input (falls back to the most-recently inserted registry key). All branches attach a training_results card so the artifact panel always shows the model context.

• scores (default), buckets HIGH/MEDIUM/LOW (or ANOMALY count for isolation), top-N renderer with discovered identifier columns, plus a sentence that tells the model where the full scored_{target} dataset lives so it doesn’t retry hoping for more rows. top_n is caller-configurable (default 10).
• importance, top 10 SHAP features as a bar chart of █ blocks.
• metrics, flat key/value listing, skipping the confusion matrix.

_exec_explain_prediction

Per-entity SHAP. The dataset-resolution cascade (tools/prediction.py:436-462) tries, in order:

1. scored_{target_column} (most likely to have every original column).
2. The training _dataset_id from the model registry.
3. Any catalog dataset whose columns are a superset of feature_columns.
4. catalog.default_dataset() as a last resort.

If none of those have an arrow_table, return a “re-upload the dataset” error message.

Entity reference parsing (tools/prediction.py:465-476) accepts row N, index N, entry N, #N, or FAC-N (1-based). Otherwise it picks the highest-risk row from the registry’s scores list. The renderer shows the top 3 explanations as “feature = value (increases/decreases risk)” and emits a shap_explanation card.

_exec_list_trained_models and _exec_find_model

Layer-3 discovery tools. [SESSION_STATE] caps inline trained-model listings at 3, so without these the LLM can’t reach the 4th-most-recent model.

• list_trained_models, tenant-scoped enumeration with optional target, algorithm, dataset_name, limit filters. Sorts newest-first by created_at. Each entry is rendered by _format_model_for_listing (tools/prediction.py:504-539): model_id, original user_goal quote, target/algorithm, metrics (first one from a quality-ordered list auc_roc | auc | accuracy | f1 | rmse | mae | r2), source dataset, feature preview, trained timestamp.
• find_model, weighted token search over user_goal (×3), target_column (×2), source_dataset_name (×2), and feature_columns (×0.5), with a small quality nudge from metrics[auc_roc|auc|accuracy|f1] × 0.1 to break ties. Returns up to limit (default 5, max 25) matches. No vector embeddings here, this is keyword tokenization, not the vector_searcher path used for tools/primitives/catalog. If the trained-model registry grows past a few hundred entries the ranking gets noisy; a future migration onto the system-corpus embedding path is the obvious move (cf. the Knowledge-Is-Routing rule).

Trained-Model Registry

container.predict_models() is a providers.Object({}), a plain dict passed by reference to every consumer. Three writers (ModelTrainer.train indirectly, tools/prediction.py:_exec_train_model, predict/router.py:train_model), one reloader (ModelTrainer.load_persisted_models at startup), and several readers (_exec_get_model_info, _exec_explain_prediction, _exec_list_trained_models, _exec_find_model, the router endpoints).

Keys are "{tenant_id}:{model_id}". Each value carries the full training result dict plus the disambiguation fields written by the tool handler:

Field	Source	Purpose
_dataset_id	training input	Re-resolve the original dataset for SHAP
_tenant_id	context	Tenant isolation cross-check
_profile	catalog._build_profile	What the advisor saw
user_goal	tool input	”What did the user ask this model to do?”, primary signal for find_model
source_dataset_name	catalog or namespace var	Human-readable origin (with (analyze_data namespace) suffix when applicable)
created_at	UTC ISO at training time	Sort key for list_trained_models

predict/router.py:train_model writes the same disambiguation fields with user_goal="" because the REST TrainRequest schema doesn’t carry it (predict/router.py:762-771).

router.py, Predict Endpoints

/api/v1/predict/*, all @requires("predict").

Method	Path	Body / Query	Purpose
POST	/upload	UploadFile	Tabular/spatial upload, dedupe, profile, catalog-register. Handles CSV/TSV/XLSX/XLS/CSV.gz, GeoJSON/KML/KMZ/GPKG/SHP zips, mixed ZIPs (peeks inside to decide spatial vs tabular), .xlsb → .xlsx via headless LibreOffice. ZipSlip-safe, capped by settings.max_upload_bytes and settings.max_archive_extracted_bytes.
POST	/recommend	RecommendRequest	Profile-driven ModelAdvisor.recommend for a registered dataset.
POST	/train	TrainRequest	Train a model on a catalog dataset (REST path; the chat tool uses _exec_train_model).
GET	/datasets	,	Tenant-scoped dataset summaries (id, name, rows, columns).
GET	/models	,	Trained models for the current tenant (id, algorithm, target, metrics, training time).
POST	/score	ScoreRequest	Per-entity scores from a previously trained model.
POST	/explain	ExplainRequest	SHAP for a single entity by index.

Trailing whitespace is stripped from object/string columns on every tabular upload (predict/router.py:653-656), EPA/government CSVs often have padded values like 'Y ' that silently break joins and value mapping downstream. XLSX/XLS files are preserved on disk under data/uploads/{content_hash[:12]}.xlsx so manipulate_excel can still see formulas, formatting, and other sheets; CSV/TSV files are ephemeral.

The KML reader (_read_kml_all_layers) bypasses GDAL for multi-folder KMLs by parsing XML directly with lxml, which makes the national FRS KML (~19k folders, 163k features) usable in seconds instead of hours. Spatial ZIPs (and KMZ) are extracted via shared.safe_zip.safe_extract*, which enforces ZipSlip safety and an extraction-byte budget.

Tool Boundary and the Uniform Contract

Every handler in tools/prediction.py still returns a ToolResult, which the orchestrator reads for content, cards, datasets, and cost, that source shape hasn’t changed.

What sits on top of it has: the Uniform Tool Contract is live in the running tree (engine/tool_registry.py:_apply_contract wraps every registered handler at registration time), not a pending merge. tests/contract_deferred_baseline.json, the “still on the legacy path” snapshot the migration ratchet checks against, is an empty list as of 2026-09, meaning every tool in SYSTEM_TOOLS now speaks the contract at the boundary, one way or another. Every tool boundary speaks one of three named result variants instead of a free-form envelope:

• Ok, measured schema: list[ColumnSpec], rows, data, optional advisory warnings, and a typed Effects(cards, datasets, exports, cost) side-effect bag.
• Hint, recoverable boundary message (hint_kind: "schema" | "fetch" | "ambiguous") with ranked alternatives. Surfaced when the model asked for a thing the platform couldn’t bind (unknown column, missing filter, ambiguous catalog entry). Carries NO effects.
• Err, actual infra/bug failure. Carries NO effects.

A single @uniform_contract decorator wraps every registered handler at the dispatch boundary (ToolRegistry._apply_contract), catches HintException BEFORE generic Exception, validates the envelope, fires contract_envelope telemetry, and adapts back to ToolResult via envelope_to_tool_result so the orchestrator path is unchanged. Bulk migration uses migrate_tool(handler, name, is_tabular=…) which does not change the handler’s source, the wrapper is installed at registration time and translates ToolResult / legacy dict shapes into envelopes. Handlers that need recoverable behavior (e.g. “dataset name doesn’t resolve in catalog or namespace”) will eventually raise HintException(kind="schema", message=…, alternatives=[…]) directly.

For predict specifically:

• The five predict tools are spec-only migrated via migrate_tool. Source returns ToolResult; the decorator coerces.
• _exec_train_model’s current “Dataset ‘X’ not found, here’s what IS available” diagnostic (tools/prediction.py:121-165) is exactly the shape that becomes a Hint(hint_kind="schema", alternatives=[…]) once the boundary speaks the contract natively. The behavior is unchanged; the wire format is structured.
• _exec_get_model_info and _exec_explain_prediction produce tabular data (scores, SHAP table), so they register with is_tabular=True; the validator will warn if an Ok comes back with an empty schema.
• Cost handling: prediction does NOT roll up child costs the way delegate does, so is_cost_rollup stays False. The advisor LLM call’s cost flows through CostTracker.record exactly as it does today.

If you’re touching predict code and need to move a tool from spec-only to source-decorated, see tool-system.md for the contract definition and engine.md for the dispatch path.

Data Flow

User: "train a model on the latest upload"
  → engine dispatch → _exec_train_model (tools/prediction.py)
  → resolve df: explicit dataset_id → catalog.default_dataset()
                → dataset_name → catalog.exists/load OR analyze_data namespace
                → structured "what IS available" error on miss
  → DataProfiler.profile(df) via catalog._build_profile
  → ModelAdvisor.recommend(profile)   [skipped if explicit target_column]
  → ModelTrainer.train(df, target, algorithm)
       encode → split → fit → evaluate → CV → SHAP → score → pickle bundle
  → score full df_full via saved LabelEncoders → catalog.save("scored_<target>")
  → ToolResult(content=summary, cards=[training_results], datasets={"scored_<target>": df})
  → container.predict_models()[ "<tenant>:<model_id>" ] = result (+ disambiguation fields)
  → ModelTrainer.save_model_metadata(result) → models/<tenant>/<model_id>.meta.json

Cross-references: dataset resolution and the namespace bridge are owned by tools/analysis.py; the catalog primitive (catalog.save, catalog.default_dataset, _build_profile) is owned by Fabric; the LLM-role lookup and cost calculator are owned by the Shared layer.