# Lexsi DS Agent — Developer / KT Guide A code-level companion to [`DSAGENT_OVERVIEW.md`](DSAGENT_OVERVIEW.md) (which is the what/why). This is the **how**: file map, control flow, the contracts you'll touch, configuration, and the full benchmark-running mechanics. Read the overview first, then this. --- ## 1. Mental model in one paragraph One **question** → one `AgentLoop.run()` → a **ReAct loop**: build a system prompt (framing + active-dataset summary + tool catalog), ask the LLM for a JSON decision (`{"thought","tool","args"}` or `{"thought","final_answer"}`), validate args with the tool's pydantic schema, run the tool against a shared `AgentContext`, append the observation to history, repeat until a final answer or the step budget. Tools read and write a shared `ctx.cache` scratchpad. The data substrate is always **DuckDB**; tabular ML runs on **remote Lexsi pods** via the SDK. The loop is generic — it has no opinion on which tools to call. --- ## 2. Repo layout (code-oriented) ``` lexsi_ds/ agent/ loop.py AgentLoop.run() — the ReAct orchestrator (start here) context.py AgentContext, DatasetHandle, TableInfo, ColumnInfo prompts.py PLANNER_SYSTEM, render_tool_catalog, tool-selector (intent routing) datasource.py DuckDB handle loaders + multi-DB materialization + DB cache tools/ base.py Tool / ToolSpec / ToolResult contract (the interface) __init__.py REGISTRY: dict[str, Tool] (30 tools) .py one tool per file, each exports `TOOL: Tool` _tabular_common.py shared helpers for the model-lifecycle tools (resolve_tab_project, active_model_name, df_records, …) kg_index.py embedding/lexical index over the PKDD knowledge graph table_index.py embedding index over table schemas (schema linking, big DBs) session.py SessionContext + multi-turn merge turn_classifier.py classify a turn: new / refine / abandon reference_resolver.py resolve "it/that/those" against prior turns samples.py curated demo questions (by_id, SAMPLES) — UI chips + tests llm/client.py LLMClient protocol, StubLLMClient, LexsiTextClient, cache, timeout schema/ raw.py PKDD DDL introspection (introspect()) context_graph.py ContextGraph — the banking knowledge graph data/loader.py PKDD TSVs -> artifacts/financial.duckdb paths.py DUCKDB_PATH, ARTIFACTS_DIR, etc. app/gradio_app.py the UI (build_ui(); demo.launch on GRADIO_PORT) benchmark/ dab/ adapter.py parse a DAB dataset -> DatasetHandle + tasks (+ DB cache key) runner.py run agent per DAB task, score Pass@1 via the task's validate.py DataAgentBench/ the (gitignored) upstream checkout + DB files pkdd_runner.py run agent over bench/pkdd/questions.yaml, auto-score datagen/ template pipeline that GENERATES bench/pkdd/*.yaml bench/pkdd/ questions.yaml 40 single-turn bench items (analytic/predictive/adversarial) sessions.yaml multi-turn sessions p2_feature_engineering.yaml 10 feature-eng items (separate `assertions` schema) benchmark/fixtures/pkdd_seed.duckdb pinned PKDD snapshot the bench scores against scripts/ repro_tabicl/ TabICL inference + per-case XAI failure repro repro_model_lifecycle/ model-lifecycle SDK probe harness (self-contained): data/ the PKDD CSVs (local copy) repro_model_lifecycle.py TabICL-first (clean) repro_model_lifecycle_xgboost_first.py XGBoost-first (#05-013 cascade repro) tests/ pytest suite (test_b_tools, test_v1_multiturn, test_dab_*, …) docs/ design + results + this guide ``` --- ## 3. The run lifecycle ([loop.py](../lexsi_ds/agent/loop.py)) `AgentLoop(ctx, registry=REGISTRY, max_steps=18).run(question) -> AgentRun` ``` run(question): 1. (multi-turn) classify_turn() -> new/refine/abandon; resolve_reference() - "abandon" short-circuits with a polite answer. 2. ctx.cache["question"] = question 3. system = _render_system(...) # built ONCE per run, re-sent every step - _select_tools(question): LLM intent classifier -> render only relevant tool group(s) - render_dataset_block(ctx) + render_tool_catalog(tools) + PLANNER_SYSTEM 4. history = ["USER QUESTION: ..."] 5. while n < max_steps: user_msg = "\n\n".join(history) + "You have {max_steps-n} tool calls left..." decision = parse(ctx.llm.complete(system, user_msg)) if "final_answer": finalize + return tool = registry.get(decision["tool"]) result = _call_tool(tool, decision["args"]) # pydantic-validate then tool.run() history.append("TOOL CALL ...\nOBSERVATION: " + result.summary[:1500]) # anti-loop: same tool fails 2x in a row -> inject "[LOOP GUARD] re-plan" directive if result.meta.get("pause_for_user"): return (clarify pause) # interactive only 6. budget exhausted -> _force_final(): ask once for a final answer ``` Key methods: `run`, `_render_system`, `_select_tools`, `_call_tool`, `_force_final`. Dispatch **always uses the full `self.registry`**; the tool selector only narrows what the planner *sees* in the prompt, never what it can call. Data classes: - `AgentRun` — `run_id, question, dataset_id, steps[], final_answer, ok, error, pending_clarification, system_prompt`. - `Step` — `n, thought, tool, args, result(ToolResult), raw_llm_text, latency_s, prompt`. The full per-turn user message lives on each Step (so trajectories are reproducible). --- ## 4. Core data structures ([context.py](../lexsi_ds/agent/context.py)) **`AgentContext`** — built once per run, passed to every tool: - `dataset: DatasetHandle`, `run_id` - `org / text_project / tab_project / tab_projects` — Lexsi SDK handles (may be None offline) - `llm` — an `LLMClient` - `cache: dict` — run-scoped scratchpad (see §6) - `session` — optional `SessionContext` (multi-turn) - `interactive: bool` — True in the UI (clarify pauses); False in bench/headless (clarify falls back to its `default`) - `.duck(read_only=None)` — open a DuckDB connection on `dataset.duckdb_path` (read-only for bundled, writable otherwise) **`DatasetHandle`** — `id, kind, duckdb_path, tables: list[TableInfo], connector_id, kg`. `kind ∈ {bundled, upload, connector, attached}`. Everything funnels through `ctx.duck()` → `duckdb_path`, so "switch dataset" = build a new handle and assign `ctx.dataset`. **`TableInfo`** = `name, columns: list[ColumnInfo], n_rows, comment`. **`ColumnInfo`** = `name, dtype, comment, sample` (the `sample` value is the per-column representative value the planner sees in the schema — populated at describe time). --- ## 5. The tool contract — how to add a tool ([tools/base.py](../lexsi_ds/agent/tools/base.py)) A tool is any object with `.spec: ToolSpec` and `.run(args, ctx) -> ToolResult`. ```python # lexsi_ds/agent/tools/my_tool.py from pydantic import BaseModel, Field from lexsi_ds.agent.context import AgentContext from lexsi_ds.agent.tools.base import Tool, ToolResult, ToolSpec class MyToolArgs(BaseModel): table_or_label: str = Field(..., description="...") # match sibling arg names! def _run(args: MyToolArgs, ctx: AgentContext) -> ToolResult: con = ctx.duck() ... return ToolResult(ok=True, summary="short text the LLM sees (<~200 tok)", payload={"dataframe": df}) # rich data; LLM never sees it class _MyTool: spec = ToolSpec(name="my_tool", description="...(LLM sees verbatim)...", args_schema=MyToolArgs, returns="what the observation looks like") def run(self, args, ctx): return _run(args, ctx) TOOL: Tool = _MyTool() ``` Then register it: add `from .my_tool import TOOL as my_tool_tool` and an entry in `REGISTRY` in [tools/__init__.py](../lexsi_ds/agent/tools/__init__.py). If it fits an intent group, add it to `_TOOL_GROUPS` in [prompts.py](../lexsi_ds/agent/prompts.py) so the tool selector surfaces it (see §8). Contract notes: - `summary` is the **only** thing the LLM sees; `payload` is for the UI / downstream tools. - Return `ok=False` with an **actionable** `summary` on expected failures — the planner reads it and re-plans. Unexpected exceptions propagate; the loop wraps them. - **Arg-name convention:** sibling tools that take a source use `table_or_label`. Match it (or alias via pydantic `AliasChoices`) — mismatches show up as `validation_error` and waste planner steps (this bit us; see `transform_column`). - Args are validated in `_call_tool` via `tool.spec.args_schema.model_validate(raw_args)`. --- ## 6. The cache contract (`ctx.cache`) Run-scoped (session-scoped in multi-turn). Producers → consumers, by key: | Key | Producer | Consumer | |---|---|---| | `question` | loop | tool selector | | `tool_intents::` | `_select_tools` | (memoize the intent classification) | | `sql_result: