Spaces:
Sleeping
Sleeping
| # 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) | |
| ``` | |