docs/tool_reference.md

# Lexsi Data Science Agent — Tool Reference (v0)

What's in the registry right now, how each tool works, and how they fit
together in a run. Generated from the current code, not from the
design doc — when those disagree, this file wins.


---

## 1. Agent loop in one minute

A run = one user question. The loop:

1. Build a system prompt: framing + active dataset summary + tool catalogue.
2. Ask the LLM for a JSON decision: call a tool **OR** emit a final answer.
3. If tool call: validate args via the tool's pydantic schema, run, append
   the observation to the conversation, loop.
4. If final answer: stop and return.

Implementation: [`lexsi_ds/agent/loop.py`](../lexsi_ds/agent/loop.py).
Hard rail: `max_steps=12` (v0 default). On exhaustion the loop asks
once for a final answer with `force_final=True`.

The loop has **no opinion about which tools to call** that's the
agent's job. The planner sees:

- A short system framing ([`PLANNER_SYSTEM`](../lexsi_ds/agent/prompts.py))
  with 11 numbered behavior rules (use `text_to_sql` for non-trivial
  SQL, retry on empty SQL results, call `summarize_result` after
  `predict`, etc.).
- The active dataset summary (id, kind, table list, KG availability).
- The tool catalogue — every registered tool's `name` + `description` +
  JSON schema + `returns` hint rendered into the prompt verbatim.

The planner responds with a SINGLE JSON object, one of:

```json
{"thought": "...", "tool": "<tool_name>", "args": {...}}
{"thought": "...", "final_answer": "<answer for the user>"}
```
---

## 2. The tool contract

Every tool implements [`lexsi_ds/agent/tools/base.py`](../lexsi_ds/agent/tools/base.py):

```python
@dataclass(frozen=True)
class ToolSpec:
    name: str
    description: str           # full paragraph; LLM sees verbatim
    args_schema: type[BaseModel]
    returns: str               # one-line observation shape

@dataclass
class ToolResult:
    ok: bool
    summary: str               # short observation text the LLM sees
    payload: Any = None        # rich Python (DataFrames, dicts) — LLM never sees
    error: str | None = None
    meta: dict[str, Any] = {}  # timing, provider trace ids, etc.

class Tool(Protocol):
    spec: ToolSpec
    def run(self, args: BaseModel, ctx: AgentContext) -> ToolResult: ...
```

**Two-channel design** is the most important pattern: `summary` is for
the planner (keep under ~200 tokens), `payload` is for the UI and
downstream tools (rich Python, never serialised into the prompt). The
distinction is what keeps prompts small as runs grow.

### `AgentContext` — what tools share

From [`lexsi_ds/agent/context.py`](../lexsi_ds/agent/context.py):

| Field | Purpose |
|---|---|
| `dataset: DatasetHandle` | Active dataset (PKDD / upload / S3 connector) |
| `run_id: str` | UUID-derived per-run id; tools use it to mint unique resource names |
| `org, text_project, tab_project, tab_projects` | Lexsi SDK handles |
| `llm: LLMClient` | Planner + summarizer LLM |
| `cache: dict[str, Any]` | Run-scoped scratchpad shared between tools (see §4) |
| `duck(read_only=…)` | Open a fresh DuckDB connection against the active dataset |

### Registry pattern

[`lexsi_ds/agent/tools/__init__.py`](../lexsi_ds/agent/tools/__init__.py)
maps name → `Tool` instance. Adding a tool = drop a file with a
`TOOL: Tool` module-level constant + add a line to `REGISTRY`. The
loop never imports a tool directly.

---

## 3. Tool catalogue

Ten tools across four groups. Order in the table matches the order in
the system prompt (related tools grouped so the LLM sees them in a
sensible neighborhood).

| Group | Tool | Side effects | Lexsi SDK? |
|---|---|---|---|
| Discovery | `connect_datalake` | binds new dataset | `create_data_connectors` |
| Discovery | `inspect_data` | none | no |
| Discovery | `query_kg` | none | no |
| Retrieval | `text_to_sql` | LLM call only | optional |
| Retrieval | `run_sql` | writes `cache["sql_result:*"]` | no |
| Retrieval | `sample_values` | none | no |
| Modeling | `train_tabular_model` | trains a Lexsi model, mints a project | `TabularProject.upload_data + train_model` |
| Modeling | `predict` | uploads predict tag, runs inference | `TabularProject.model_inference` |
| Modeling | `explain_prediction` | per-case SHAP | `case_predict` + `xai_*` |
| Narration | `summarize_result` | LLM call only | `TextProject.chat_completion` |

### 3.1 `connect_datalake`

**Source:** [`tools/connect_datalake.py`](../lexsi_ds/agent/tools/connect_datalake.py)

Register an S3 datalake via Lexsi `create_data_connectors` and bind one
or more parquet files as queryable DuckDB views. After a successful
call the connector becomes the active dataset; subsequent
`inspect_data` / `run_sql` calls operate against the bound views.

| Arg | Type | Default | Notes |
|---|---|---|---|
| `kind` | `Literal["s3"]` | `"s3"` | GCS / SFTP / GDrive are fast-follow |
| `bucket` | `str` | — | required |
| `prefix` | `str` | `""` | key prefix under the bucket |
| `region` | `str` | `"us-east-1"` | |
| `aws_access_key_id` | `str \| None` | reads `AWS_ACCESS_KEY_ID` from env | not logged, not surfaced |
| `aws_secret_access_key` | `str \| None` | reads `AWS_SECRET_ACCESS_KEY` | same |
| `files` | `list[str]` | — | each file → a view named after its stem |

**Observation:** list of bound view names + the Lexsi connector id.
**Use when:** user wants their own data instead of the bundled PKDD.

### 3.2 `inspect_data`

**Source:** [`tools/inspect_data.py`](../lexsi_ds/agent/tools/inspect_data.py)

Return the schema (tables, columns, types, comments) for the active
dataset as Markdown. Cheap and side-effect-free.

| Arg | Type | Default | Notes |
|---|---|---|---|
| `table_filter` | `list[str] \| None` | `None` | If omitted, all tables |

**Observation:** Markdown schema with table names, column types, and
KG-attached pointers (when the dataset has a KG, the observation
suggests `query_kg` for disambiguation).

**Cap:** `_SCHEMA_BUDGET_CHARS = 10000` — truncates between tables,
never mid-table, so the LLM always sees coherent column lists.

**Use when:** unfamiliar dataset, OR mid-run to remind the planner
which columns exist.

### 3.3 `query_kg`

**Source:** [`tools/query_kg.py`](../lexsi_ds/agent/tools/query_kg.py)

Explore the business knowledge graph attached to the active dataset.
The KG knows which column to pick when wording is ambiguous
("customer" → `fin_client`, not `fin_account`), which join paths are
correct, and what business metrics exist.

| Arg | Type | Default | Notes |
|---|---|---|---|
| `section` | `Literal[…] \| None` | `None` | concepts / disambigs / metrics / join_paths / tables |
| `search` | `str \| None` | `None` | Free-text — pass the user's question **verbatim**, not extracted keywords |
| `node` | `str \| None` | `None` | Fetch one node by label or full id |
| `top_k` | `int` | `8` | Max results for `search=` |

Three modes:
- No args → outline of all sections.
- `search=` / `section=` → top-K matches or full section listing.
- `node=` → one node's full record.

**Why pass the full question:** the KG indexes on trigger phrases like
"lives in", "owned by", "customer of". Single nouns lose
disambiguation signal. If initial scores are weak, the tool
auto-broadens using the active run's question.

**Use when:** the question contains potentially-ambiguous wording.
**Critical:** the planner is instructed to call this *before* writing
SQL on any ambiguous noun.

### 3.4 `text_to_sql`

**Source:** [`tools/text_to_sql.py`](../lexsi_ds/agent/tools/text_to_sql.py)

Convert a natural-language question into DuckDB SQL against the active
dataset. The default SQL-writer — `run_sql` inline is reserved for
trivial single-table queries.

| Arg | Type | Default | Notes |
|---|---|---|---|
| `question` | `str` | — | required |
| `mode` | `Literal["analytic","predictive","auto"]` | `"auto"` | see below |

Modes:

| Mode | Output |
|---|---|
| `analytic` | One SQL string that answers the question directly. |
| `predictive` | Two SQLs (`context_sql` for label-known rows, `predict_sql` for label-unknown) + a `task` spec (`task_type`, `target_column`, `entity_column`). Wired straight into `train_tabular_model` + `predict` next. |
| `auto` | Picks based on keywords (`predict`, `forecast`, `likely`, `will`, `propensity`, …). |

**Schema rendering:** For PKDD, uses the introspected `RawSchema` with
DDL + comments. For uploaded/connector datasets, builds a plain
DESCRIBE-style listing from `ctx.dataset.tables`.

**Why the default:** the planner is hard-rule-instructed (see
[`PLANNER_SYSTEM` rule 6](../lexsi_ds/agent/prompts.py)) to call
`text_to_sql` for any SQL involving JOINs, conditional aggregates, or
string-value filters. Inline `run_sql(sql=…)` allowed only for trivial
single-table queries. The tool sees the full KG and is far less likely
to pick the wrong join path.

### 3.5 `run_sql`

**Source:** [`tools/run_sql.py`](../lexsi_ds/agent/tools/run_sql.py)

Execute DuckDB-compatible SQL against the active dataset and return the
result as a DataFrame.

| Arg | Type | Default | Notes |
|---|---|---|---|
| `sql` | `str` | — | required |
| `label` | `str \| None` | `None` | Human label later tools can reference |

**Observation:** row count × col count + head preview (max 20 rows).
**Payload:** the full `pd.DataFrame`.
**Cache writes:**

- `ctx.cache["last_sql_result"]` = the DataFrame (always).
- `ctx.cache["sql_result:<label>"]` = the DataFrame (when `label` is passed).
- `ctx.cache["last_sql"]` = the executed SQL.

**Critical pattern:** when the planner runs `text_to_sql(mode="predictive")` →
`run_sql(sql=context_sql, label="context_df")` →
`run_sql(sql=predict_sql, label="predict_df")`, downstream
`train_tabular_model(df_label="context_df")` reads the DataFrame by
label out of the cache. **`context_df` is a cache key, not a DuckDB
table.** Planner rule 11 spells this out.

### 3.6 `sample_values`

**Source:** [`tools/sample_values.py`](../lexsi_ds/agent/tools/sample_values.py)

Preview the distinct values of a column. Use **before** writing
`WHERE col = 'value'` on any string column whose values aren't already
known.

| Arg | Type | Default | Notes |
|---|---|---|---|
| `table` | `str` | — | must exist in the active dataset |
| `column` | `str` | — | |
| `k` | `int` | `20` | hard cap `100` |
| `order_by_freq` | `bool` | `True` | most common first — surfaces canonical spelling |

Real-world cases this catches: `'Prague'` vs `'Hl.m. Praha'`, `'gold'`
vs `'G'`, loan-status codes, transaction-type codes.

**Safety:** identifiers are quoted defensively (`_quote_ident` rejects
anything with whitespace, parens, or quotes). The planner's args are
already pydantic-validated, but the values come from an LLM, so we
double-check at the identifier boundary.

**Planner rule 8** (system prompt): "String filters: sample first."

### 3.7 `train_tabular_model`

**Source:** [`tools/train_tabular_model.py`](../lexsi_ds/agent/tools/train_tabular_model.py)
(largest tool — ~900 LoC; absorbs most of the Lexsi-SDK quirks
documented in [`docs/sdk_issues.md`](sdk_issues.md))

Train a Lexsi tabular model on a previously-fetched DataFrame.

| Arg | Type | Default | Notes |
|---|---|---|---|
| `df_label` | `str` | — | references `ctx.cache["sql_result:<label>"]` |
| `target_column` | `str` | — | |
| `entity_column` | `str \| None` | auto-detect | first unique int/string column |
| `task_type` | `Literal["classification","regression"]` | `"classification"` | |
| `model_name` | `str \| None` | `None` | Informational; not the Lexsi-side model name |
| `model_type` | `str \| None` | `None` | `"XGBoost"`, `"TabPFN"`, `"TabICL"`, … |
| `compute_type` | `str \| None` | `None` | required for foundation models |
| `xai_method` | `list[str]` | `[]` | empty by default (see SDK issue #4) |
| `sample_percentage` | `float \| None` | `None` | |

**What it actually does:**

1. Reads the DataFrame from cache by label.
2. Resolves a per-run Lexsi project via `_get_or_create_run_project` —
   mints `agent<task[:5]><runid[:8]>` and calls
   `workspace.create_project(modality="tabular", project_type=task_type)`.
   Each agent run gets a fresh project (workaround for the SDK's
   "Config already exists" guard — no `delete_config` on the SDK).
3. Builds a `ProjectConfig` dict (12 keys, all populated; see
   `_build_project_config`). `pred_label="Prediction"` unconditionally —
   the server-side pipeline KeyErrors on `pred_label=None`.
4. Adds a placeholder `"Prediction"` column to the DataFrame mirroring
   the target — workaround for `#03-013` "Target must have ≥2 classes"
   when `pred_label` is configured but absent.
5. If the project has no saved config (fresh-project path),
   `upload_data(data, tag, config=…, compute_type=…)` does configure +
   upload + train + XAI in one shot.
6. If it does (returning-project path),
   `upload_data(data, tag)` then `train_model(model_type, compute_type, data_config=…)`.
7. Polls for completion (default `LEXSI_TRAIN_TIMEOUT_S=900s`).

**Cache writes:**

- `ctx.cache["last_model_id"]` = Lexsi-minted model name (e.g. `XGBoost_v1`)
- `ctx.cache["last_model_task"]` = `{task_type, target_column, entity_column}`
- `ctx.cache["last_train_tag"]` = upload tag
- `ctx.cache["last_train_df_label"]` = the input cache label
- `ctx.cache[f"agent_run_tab_project:{task_type}"]` = the per-run `TabularProject` handle

**Failure recovery:** when `upload_data` raises with an empty message
(`Exception("")`), the tool fetches `project.recent_events()` and
attaches the latest failure detail to the `ToolResult.summary` so the
planner sees the real error.

### 3.8 `predict`

**Source:** [`tools/predict.py`](../lexsi_ds/agent/tools/predict.py)

Run a trained Lexsi model against label-unknown rows.

| Arg | Type | Default | Notes |
|---|---|---|---|
| `df_label` | `str` | — | references `ctx.cache["sql_result:<label>"]` |
| `model_id` | `str \| None` | `cache["last_model_id"]` | |
| `top_k` | `int \| None` | `10` | observation shows top-K only; full df cached |
| `pod` | `str \| None` | `"small"` | workaround for SDK `UnboundLocalError` (#1) |

**What it actually does:**

1. Loads DataFrame from cache; strips an all-null target column if
   present (Lexsi rejects predict-tag uploads that include the
   configured `true_label`).
2. Checks the per-run predict cache (`ctx.cache["predict_cache:<hash>"]`)
   — same `(model_id, df_hash)` is a no-op. Replays cached side effects.
3. Resolves the right `TabularProject` from `ctx.cache[f"agent_run_tab_project:{task_type}"]`.
4. Mints a predict tag (`agentpredict<runid><label>`, alphanumeric only),
   `upload_data(predict_df, tag=…)` without config.
5. `model_inference(tag=…, model_name=…, pod="small")`.
6. `_normalize_predictions` finds the prediction + probability columns by
   candidate lookup — current candidates include `"Predicted_value_AutoML"`,
   `"pred_proba_AutoML"`, generic `"Prediction"`/`"Probability"`, etc.
   Resilient to SDK shape changes.

**Cache writes:**

- `ctx.cache["last_predictions"]` = full predictions DataFrame
- `ctx.cache["last_predict_tag"]` = the Lexsi tag
- `ctx.cache["last_predict_df_label"]` = input df_label
- `ctx.cache["last_prediction_columns"]` = `{pred, prob, entity}` column names

**Known SDK bug worked around:** `pod=None` triggers
`UnboundLocalError: custom_batch_servers` in
`lexsi_sdk.core.tabular.model_inference`. We always pass `pod="small"`
unless the planner overrides. Upstream PR open at
[Lexsi-Labs/Lexsi-sdk#67](https://github.com/Lexsi-Labs/Lexsi-sdk/pull/67).

### 3.9 `explain_prediction`

**Source:** [`tools/explain_prediction.py`](../lexsi_ds/agent/tools/explain_prediction.py)

Per-case SHAP + plain-language summary for selected predictions.

| Arg | Type | Default | Notes |
|---|---|---|---|
| `case_ids` | `list[str\|int] \| None` | top-K of `last_predictions` | |
| `top_k` | `int` | `10` | default when `case_ids` omitted |
| `model_id` | `str \| None` | `cache["last_model_id"]` | |
| `include_similar` | `bool` | `False` | one extra API call per case |
| `include_summary` | `bool` | `True` | calls `case.xai_summary()` per case |

**What it actually does:** for each case_id, calls
`TabularProject.case_predict(tag, model_name, entity_value, xai_method=["shap"])`.
For each returned case object, optionally calls `case.xai_summary()`
(plain-language) and `case.xai_similar_cases()` (k-NN over training set).

**Cache writes:**

- `ctx.cache["last_xai"]` = list of `{case_id, shap_values, summary?, similar?}`
- `ctx.cache["last_xai_failures"]` = per-case failure detail (so the
  summarizer can mention how many cases failed XAI without scanning logs)

**Observation:** compact table of (case_id, top-3 SHAP drivers,
plain-language summary). Full XAI lives on `payload["xai"]`.

### 3.10 `summarize_result`

**Source:** [`tools/summarize_result.py`](../lexsi_ds/agent/tools/summarize_result.py)

The final user-facing narrative. Reads cached SQL result, predictions,
and XAI; asks the LLM ([`SUMMARIZER_SYSTEM` prompt](../lexsi_ds/agent/prompts.py))
to write a concise answer that cites the numbers.

| Arg | Type | Default | Notes |
|---|---|---|---|
| `question` | `str` | — | the original user question verbatim |
| `sql_result_label` | `str \| None` | most recent | which `run_sql` result to cite |
| `include_xai` | `bool` | `True` | include `last_xai` if cached |
| `notes` | `str \| None` | `None` | planner-side caveat / what was skipped |

**What it actually does:** builds a short prompt with the question, a
~15-row preview of the SQL result, the first 10 predictions if any,
the XAI summary if any, and the planner's optional notes. Sends to
the LLM via `ctx.llm.complete(SUMMARIZER_SYSTEM, user_msg)`.

**Observation = the narrative itself.** The planner usually emits this
as `final_answer` immediately after — planner rule 11 makes this
explicit: "After `predict` succeeds, your next tool call MUST be
`summarize_result`."

---

## 4. Shared cache contract

The `ctx.cache` dict is how tools talk to each other without
re-querying the SDK or re-running SQL. Keys grouped by producer:

| Key | Producer | Consumer(s) | Type |
|---|---|---|---|
| `sql_result:<label>` | `run_sql` | `train_tabular_model`, `predict`, `summarize_result` | `pd.DataFrame` |
| `last_sql_result` | `run_sql` | `summarize_result` | `pd.DataFrame` |
| `last_sql` | `run_sql` | UI / debug | `str` |
| `last_model_id` | `train_tabular_model` | `predict`, `explain_prediction`, `summarize_result` | `str` |
| `last_model_task` | `train_tabular_model` | `predict`, `explain_prediction` | `dict` |
| `last_train_tag` | `train_tabular_model` | (UI) | `str` |
| `last_train_df_label` | `train_tabular_model` | (UI) | `str` |
| `agent_run_tab_project:<task>` | `train_tabular_model` | `predict`, `explain_prediction` | `TabularProject` |
| `agent_workspace` | `train_tabular_model` | reused on subsequent `_get_or_create_run_project` | `Workspace` |
| `last_predictions` | `predict` | `explain_prediction`, `summarize_result` | `pd.DataFrame` |
| `last_predict_tag` | `predict` | `explain_prediction` | `str` |
| `last_predict_df_label` | `predict` | (UI) | `str` |
| `last_prediction_columns` | `predict` | `summarize_result` | `dict` |
| `predict_cache:<hash>` | `predict` | `predict` (replay) | `ToolResult` |
| `last_xai` | `explain_prediction` | `summarize_result` | `list[dict]` |
| `last_xai_failures` | `explain_prediction` | `summarize_result` | `list[dict]` |
| `question` | `AgentLoop.run` | `query_kg` (for question-aware broadening) | `str` |

All keys are run-scoped and wiped at the next `AgentLoop.run` call.

---

## 5. North-star trace

The agent's predictive demo question, instrumented step by step:

> *"For loans currently being repaid, which are likely to default?
> Show me the top 10 highest-risk loans and tell me why."*

```
1. inspect_data()
   → schema for 8 PKDD tables + KG pointer

2. query_kg(search="<full question>")
   → disambig: 'default' → loan_status in ('B','D')

3. text_to_sql(question=…, mode="predictive")
   → context_sql (label-known: status A/B)
   + predict_sql  (label-unknown: status C/D)
   + task = {task_type: classification, target_column: y_default,
             entity_column: loan_id}

4. run_sql(sql=context_sql, label="context_df")
   → cache["sql_result:context_df"] = 682-row DataFrame

5. run_sql(sql=predict_sql, label="predict_df")
   → cache["sql_result:predict_df"] = 403-row DataFrame

6. train_tabular_model(df_label="context_df", target_column="y_default",
                       entity_column="loan_id", task_type="classification",
                       model_type="XGBoost", compute_type="T4.small")
   → cache["last_model_id"] = "XGBoost_v1"

7. predict(df_label="predict_df", top_k=10)
   → cache["last_predictions"] = 403-row predictions DataFrame
   → observation: top-10 rows by probability

8. explain_prediction(top_k=10)
   → cache["last_xai"] = per-case SHAP for the 10 highest-risk loans

9. summarize_result(question=<verbatim>, include_xai=True)
   → narrative: top-10 loan ids + probabilities + top SHAP drivers

10. final_answer = narrative
```

A simpler question ("how many accounts?") collapses to
`text_to_sql` + `run_sql` + `summarize_result` — three steps. An
adversarial one ("what's the meaning of life?") should end at step 0
with `final_answer` declining. The planner picks the subset; the
catalogue makes the subset available.

---

## 6. How tools FAIL — patterns to know

Three failure modes the catalogue is shaped around:

### 6.1 Empty `ToolResult.summary`

Lexsi SDK occasionally raises `Exception("")` — empty message —
when the server returns `{success: false, details: ""}`. The training
tool catches this and fetches `project.recent_events()` so the planner
sees the actual server error instead of a bare `Exception`. The
predict tool does the equivalent for `model_inference`.

### 6.2 Bare-name cache misuse

The agent sometimes tries `SELECT * FROM predict_df` thinking
`predict_df` is a DuckDB table. It's not — it's a key in
`ctx.cache["sql_result:predict_df"]`. Rule 11 in
[`PLANNER_SYSTEM`](../lexsi_ds/agent/prompts.py) calls this out:

> "Do not try `run_sql` against the predict `df_label` (it's a cache
> key, not a DuckDB table)."

### 6.3 Skipped `summarize_result`

After `predict` succeeds the planner sometimes emits `final_answer`
directly with prose like "the predictions are cached, here are next
steps". Rule 11 (same): "After `predict` succeeds, call
`summarize_result` next." The cached predictions DataFrame is what
the summarizer reads — even when `pred_col` / `prob_col` couldn't be
auto-detected.

---

## 7. Adding a new tool

```
1. Drop `lexsi_ds/agent/tools/<name>.py` with:
     - a pydantic `*Args(BaseModel)` class
     - a `_run(args, ctx) -> ToolResult` function
     - a module-level `TOOL: Tool` instance with a `ToolSpec`

2. Add to REGISTRY in `lexsi_ds/agent/tools/__init__.py`
   (keep related tools grouped — order is what the LLM sees).

3. If the tool writes to `ctx.cache`, document the key contract in
   §4 of this doc.

4. If the tool needs Lexsi SDK access, check `ctx.org`/`ctx.tab_project`
   for None first — the loop should still run offline (stub LLM,
   bundled dataset) for tests.

5. Add a question to bench/pkdd/questions.yaml that exercises the
   tool's happy path AND at least one error path.
```

That's it. No inheritance, no framework registration, no plugin
loader. The Tool Protocol from §2 is the only contract.

---

## 8. What's coming in v1


- `ToolRouter` retrieves a small relevant subset per step instead of
  rendering the full catalogue every prompt.
- `ToolSpec` extended with `capabilities`, `cost_class`, `requires`,
  `produces`, `parallel_safe`.
- `SessionContext` layered over `AgentContext` so cached state
  survives across turns.
- `ResultNode` graph replaces the string-only `final_answer`.