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. 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) 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:

{"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:

@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:

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

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

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

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

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

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

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 (largest tool — ~900 LoC; absorbs most of the Lexsi-SDK quirks documented in docs/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

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.

3.9 explain_prediction

Source: 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

The final user-facing narrative. Reads cached SQL result, predictions, and XAI; asks the LLM (SUMMARIZER_SYSTEM prompt) 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 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.