# 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](README.md) and the [v0 design doc](docs/v0_design.md).

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

---

## TL;DR

```bash
# 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)

```bash
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

```bash
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

```python
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

```bash
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)
```