docs/MODELING_TOOLS.md

Modeling Tools - added by Aditya

This document describes the modeling-tools layer of the Lexsi DS Agent v0: the three production tools owned by Aditya per the design doc §4 (train_tabular_model, predict, explain_prediction) plus the supporting test infrastructure. It assumes you've read the project README and the v0 design doc.

For the rest of the agent (data connectors, NL→SQL, KG retrieval, summarization), see Bhavish's tools under lexsi_ds/agent/tools/.

---

TL;DR

# Run the full modeling flow offline — no Lexsi auth required
python scripts/demo_modeling_offline.py

# Run the test suite (47 hermetic modeling tests + 32 PKDD-gated tests)
python -m pytest tests/ -q

# Bundle a zip for review
bash scripts/package.sh

---

What was delivered

1. Three production tools (live SDK)

Each replaces a not_implemented stub with a full implementation against the lexsi-sdk==0.1.46 surface. Argument schemas match the design doc §4 contract; downstream observation shapes match what summarize_result expects in cache.

File	What it does	Key SDK call
lexsi_ds/agent/tools/train_tabular_model.py	Uploads a context DataFrame + trains a model in one shot. Defaults to XGBoost; supports TabPFN / OrionMSP / TabICL / TabDPT / OrionBix / Mitra / ContextTab (requires compute_type).	TabularProject.upload_data(data, tag, config=ProjectConfig, compute_type, tunning_strategy)
lexsi_ds/agent/tools/predict.py	Uploads the predict DataFrame under a sibling tag, runs batch inference, surfaces top-K by probability.	TabularProject.model_inference(tag, model_name, pod)
lexsi_ds/agent/tools/explain_prediction.py	Per-case XAI for the top-K cases (or explicit case_ids). Reads shap_feature_importance directly, calls xai_summary and xai_similar_cases opt-in. Handles both flat and multiclass-nested SHAP dicts.	TabularProject.case_predict(unique_identifier, tag, model_name, xai)

2. Offline test infrastructure

The three tools talk to Lexsi over the network. To run them in tests and demos without an SDK token, this delivery includes a faithful test double that mirrors the SDK contract in executable form.

File	Purpose
lexsi_ds/agent/testing/__init__.py	Re-exports FakeTabularProject, FakeCaseTabular, make_offline_ctx.
lexsi_ds/agent/testing/fakes.py	FakeTabularProject mirrors upload_data's two-branch behavior (first-call sets project config + trains; subsequent calls just store tags), foundation-model compute_type check, model_inference, case_predict. FakeCaseTabular mirrors the live class field-for-field for the subset explain_prediction reads. Same role StubLLMClient plays for text. If the SDK ever changes a kwarg name, this file is the canary — tests break before production does.

3. Production bootstrap

File	Purpose
lexsi_ds/agent/lexsi_bootstrap.py	bootstrap_lexsi_handles() — env-driven login + project resolution. Mirrors the LexsiTextClient.from_env pattern. Returns LexsiHandles{org, workspace, tab_project, text_project}. Asserts isinstance(tab_project, TabularProject) so wrong-modality misconfiguration fails at boot, not at first tool call.

4. Test suite

File	Tests	What it covers
tests/test_train_tabular_model.py	15	XGBoost / TabPFN / sample_percentage happy paths; missing_df, bad_target, empty_target, missing_compute_type, no_tab_project error branches; entity-column resolution, tag minting, ProjectConfig 12-key shape, active-model resolution.
tests/test_predict.py	13	Happy path, explicit model_id, top_k bounding, null-target stripping; missing_df, empty_predict_df, no_model_id, no_tab_project; column-name normalization, top-K helpers.
tests/test_explain_prediction.py	15	Auto / explicit case_ids, include_similar, summary caching; no_model_id, no_tab_project, no_predict_tag, no_cases, all_cases_failed; SHAP top-k for flat / nested-multiclass / empty / garbage dicts.
tests/test_modeling_e2e.py	4	Full 3-step flow with cache handoff; predict-without-train fails clean; explain-without-predict fails clean; registry membership.
tests/conftest.py	(fixtures)	Adds offline fixtures: loan_context_df, loan_predict_df, fake_tab_project, ctx_with_fake, ctx_offline, seed_sql_result.

5. Demo + packaging

File	Purpose
scripts/demo_modeling_offline.py	Runnable end-to-end demo: train → predict → explain against FakeTabularProject. Flags: --rows, --predict-rows, --top-k, --model-type, --compute-type, --seed, --lexsi (swap fake for real SDK), -v.
scripts/package.sh	Produces dist/lexsi_ds_agent-modeling.zip excluding caches / sessions / runs / PKDD source. OUT= env overrides the output path.

---

Architectural decisions

1. predict uses batch model_inference; explain_prediction uses per-row case_predict. The per-case XAI path is needed only for the K rows we actually want to explain. Running case_predict on the full predict_df would cost K × the inference time and the bias-monitor / policy hooks would fire K times for no extra information.

2. No external polling. The live upload_data already calls poll_events internally and blocks until training is done. The agent loop layers wall-clock guards via daemon threads so a stuck backend cannot hang the loop:

   Env var	Default	Affects
   LEXSI_TRAIN_TIMEOUT_S	900 (15 min)	upload_data train call
   LEXSI_PREDICT_TIMEOUT_S	600 (10 min)	upload_data predict upload, model_inference
   LEXSI_CASE_TIMEOUT_S	60	per-case case_predict
   LEXSI_XAI_BUDGET_S	300 (5 min)	total across all explained cases

3. Offline-runnable via FakeTabularProject mirroring the live SDK surface, exactly the way StubLLMClient mirrors LexsiTextClient. Pins the SDK contract in executable form so contract drift breaks tests, not production.

---

Cache contract

This is what links the three tools together and what summarize_result reads in the final step.

After train_tabular_model → predict → explain_prediction:

Cache key	Type	Set by
last_model_id	str (the live Lexsi model_name)	train
last_train_tag	str	train
last_train_df_label	str	train
last_model_task	dict{target_column, entity_column, task_type, model_type}	train
last_predict_tag	str	predict
last_predict_df_label	str	predict
last_predictions	pd.DataFrame (input cols + Prediction + Probability)	predict
last_prediction_columns	dict{pred, prob, entity}	predict
last_xai	list[dict] (one entry per case: case_id, pred_value, pred_category, shap_top, shap_full, summary, optionally similar_preview)	explain
last_xai_failures	list[dict]	explain

---

How to run

Offline (no auth, no compute, no network)

python scripts/demo_modeling_offline.py --rows 100 --predict-rows 25 --top-k 5

Runs all three tools against FakeTabularProject and prints the cache state. The SHAP drivers in the demo are deterministic per seed (not per case) because the fake uses a single seeded RNG — the live Lexsi backend produces real per-case SHAP.

Against real Lexsi

export SDK_ACCESS_TOKEN=...                  # from app.lexsi.ai/sdk
export LEXSI_ORG_NAME=personal               # or your org
export LEXSI_WORKSPACE_NAME=your_ws
export LEXSI_TABULAR_PROJECT_NAME=your_proj

python scripts/demo_modeling_offline.py --lexsi --model-type XGBoost

# foundation model:
python scripts/demo_modeling_offline.py --lexsi --model-type TabPFN --compute-type small

Wiring into the live agent loop

from lexsi_ds.agent.context import AgentContext
from lexsi_ds.agent.lexsi_bootstrap import bootstrap_lexsi_handles
from lexsi_ds.llm.client import factory as llm_factory

handles = bootstrap_lexsi_handles()
ctx = AgentContext(
    dataset=...,                        # from Bhavish's connect_datalake / load_pkdd_handle()
    run_id=...,
    org=handles.org,
    text_project=handles.text_project,
    tab_project=handles.tab_project,
    llm=llm_factory("lexsi"),
)
# Three modeling tools are already in the registry; the planner can call them by name.

---

How to test

python -m pytest tests/                       # all 47 modeling tests + 32 PKDD-gated
python -m pytest tests/test_modeling_e2e.py   # end-to-end only
python -m pytest tests/ -k "happy_path"       # happy paths only
python -m pytest tests/ -k "error or no_"     # error branches only

All modeling tests are hermetic — no network, no DuckDB, no Lexsi auth. Bhavish's PKDD-gated tests skip cleanly when the DuckDB isn't bootstrapped.

---

SDK cross-verification

Every SDK call site was cross-verified against lexsi-sdk==0.1.46 via introspection of installed signatures + reading source bodies for xai_summary, xai_similar_cases, and Workspace.project.

Verified call sites

Call site	Live signature in lexsi-sdk==0.1.46	Result
TabularProject.upload_data(data, tag, config, compute_type, tunning_strategy)	matches exactly (preserving the SDK's tunning_strategy typo, not tuning_strategy)	✅
TabularProject.model_inference(tag, model_name, pod)	matches	✅
TabularProject.case_predict(unique_identifier, tag, model_name, xai)	matches	✅
TabularProject.active_model() / .models()	returns pd.DataFrame with model_name column — matches my reads	✅
CaseTabular field reads (shap_feature_importance, pred_value, pred_category, summary, similar_cases_data, model_name, data_id, unique_identifier)	all present	✅
CaseTabular.xai_summary() / .xai_similar_cases()	source confirmed — xai_summary returns string (caches on self.summary); xai_similar_cases returns pd.DataFrame | str; my code handles both	✅
Login chain: xai.login → xai.organization → Organization.workspace → Workspace.project	matches; Workspace.project() is modality-dispatching	✅
ProjectConfig TypedDict shape	exact 12-key match	✅

---

Known limitations / deliberately deferred

• No bias-monitor / risk-policy hookup. The SDK exposes these on case_predict(..., risk_policies=True); not wired in v0 because design doc scope is risk scoring and explanation, not policy enforcement. Single arg passthrough when needed.
• No PEFT / fine-tune path. Foundation models default to tunning_strategy="inference" (zero-shot) to keep the demo path fast. PEFT / base-ft requires peft_config / tunning_config — the SDK supports them; the tool args don't expose them yet. Easy follow-on if needed.
• Single active model assumed. _resolve_active_model_name falls back to the most recent row in .models() if .active_model() is empty, but doesn't model multiple concurrent active models. The live SDK enforces a single-active-model invariant, so this should be fine.
• Fake SHAP is deterministic per project, not per case. In the offline demo, top-K cases show the same SHAP drivers because the fake uses a single seeded RNG. The live Lexsi backend produces per-case SHAP — this is a cosmetic limitation of the test double, not the production tools.

---

File index — exactly what was added or changed

NEW    lexsi_ds/agent/lexsi_bootstrap.py
NEW    lexsi_ds/agent/testing/__init__.py
NEW    lexsi_ds/agent/testing/fakes.py
REPL   lexsi_ds/agent/tools/train_tabular_model.py    (stub → production)
REPL   lexsi_ds/agent/tools/predict.py                (stub → production)
REPL   lexsi_ds/agent/tools/explain_prediction.py     (stub → production)
NEW    tests/__init__.py
MERGE  tests/conftest.py                              (added offline fixtures alongside existing PKDD fixtures)
NEW    tests/test_train_tabular_model.py              (15 tests)
NEW    tests/test_predict.py                          (13 tests)
NEW    tests/test_explain_prediction.py               (15 tests)
NEW    tests/test_modeling_e2e.py                     (4 tests)
NEW    scripts/demo_modeling_offline.py
NEW    scripts/package.sh
NEW    HANDOFF.md                                     (delivery handoff note)
NEW    MODELING_TOOLS.md                              (this file)