docs/dab_integration.md

Dataset switching + DataAgentBench integration

How the agent moves off the bundled PKDD dataset, and how DataAgentBench (DAB) plugs into the benchmark harness.

---

1. How dataset switching works

Every dataset the agent can query is a DatasetHandle (one DuckDB file behind it). Switching datasets = building a new handle and assigning it to ctx.dataset. After that, inspect_data / run_sql / the modelling tools all operate against the new dataset — because they go through ctx.duck(), which opens ctx.dataset.duckdb_path.

Kind	Loader	Substrate
bundled	load_pkdd_handle()	shipped artifacts/financial.duckdb
upload	load_upload_handle(UploadSpec)	per-session DuckDB, one table per CSV/Parquet
connector	load_s3_connector_handle(S3ConnectorSpec)	per-session DuckDB with read_parquet(s3://…) views
attached	load_attached_handle([AttachSpec…])	per-session DuckDB with several sqlite/duckdb DBs materialized into it

Runtime switch (the agent does it itself): the connect_datalake tool builds an S3 handle and sets ctx.dataset = handle mid-run — the only tool with that side-effect. The DAB loader reuses the exact same contract: build a handle, assign it. No new switching machinery.

Why attached materializes instead of ATTACH-ing live: ATTACH is connection-scoped, but each tool opens its own ctx.duck() connection. The S3 path survives reopen because read_parquet(s3://…) views are self-contained. sqlite/duckdb sources have no equivalent self-contained view (a duckdb file can only be read via ATTACH), so load_attached_handle copies each source table into the session file as <alias>_<table>. The session file is then self-contained — fresh connections just work. Trade-off: a one-time copy at load (≈5 s for a 660k-row sqlite table); after that, single-DB-speed joins.

---

2. DataAgentBench at a glance

DAB tests data agents on multi-database enterprise tasks: 17 datasets across sqlite / duckdb / postgres / mongo. Each task is:

query_<dataset>/
  db_config.yaml            # which DBs, what type, file path
  db_description.txt        # schema prose for the agent
  query<N>/
    query.json              # the question (a JSON-quoted string)
    ground_truth.csv        # the reference answer
    validate.py             # def validate(llm_output) -> (bool, reason)

Scoring is Pass@1: run the agent, take its final answer string, call the task's own validate(answer) — which checks the ground truth is present in the answer (substring / name-near-value match). We reuse DAB's validate.py verbatim rather than reimplementing it.

---

3. What we can run today

load_attached_handle folds all four DB systems DAB uses into one session DuckDB (needs the connectors extra for the last two):

Source	How it's materialized
sqlite / duckdb	ATTACH the file → CREATE TABLE AS SELECT
postgres	restore the pg_dump .sql into an embedded pgserver (no Docker) → postgres scanner ATTACH → copy
mongo	read the mongodump .bson server-lessly via bson → flatten (nested → JSON text) → copy

So 17 of 17 datasets are loadable (was 5). All db_types reduce to the same <alias>_<table> shape in one DuckDB, so cross-system joins are ordinary SQL — verified on cve (sqlite+duckdb+postgres+mongo, 8 tables, ~2M rows, ~4s).

Note: DAB ships most DB files in-repo and a few via download.sh / git-lfs (PATENTS is a 5 GB Drive download). build_handle raises FileNotFoundError naming the missing file if a DB isn't present — only PATENTS is missing locally today. Install the deps with uv sync --extra connectors.

---

4. Running it

# 1. Clone DataAgentBench somewhere + pull its DB files
git clone https://github.com/ucbepic/DataAgentBench /tmp/DataAgentBench

# 2. Run the agent against one dataset (needs the Lexsi text env for a
#    real planner; --llm stub just exercises the wiring)
export SDK_ACCESS_TOKEN=… LEXSI_WORKSPACE_NAME=… LEXSI_TEXT_PROJECT_NAME=…
python -m benchmark.dab.runner \
    --dab-root /tmp/DataAgentBench \
    --dataset DEPS_DEV_V1 \
    --llm lexsi \
    --iterations 1

Output:

DAB DEPS_DEV_V1: bound 3 table(s): ['package_database_packageinfo',
  'project_database_project_info', 'project_database_project_packageversion']
DAB DEPS_DEV_V1/query1: pass@1=1/1  (All name-version pairs validated…)

=== DAB results ===
  [PASS] DEPS_DEV_V1/query1  steps=6  All name-version pairs validated…
  total: 1/1 passed

Omit --query to run all tasks in the dataset; bump --iterations for a real Pass@1 over N runs.

---

5. Code map

File	Role
lexsi_ds/agent/datasource.py → load_attached_handle	multi-DB fold-in; dispatches sqlite/duckdb/postgres/mongo. Live-DB loaders load_postgres_connector_handle / load_mongo_connector_handle reuse the same copy/flatten helpers
benchmark/dab/adapter.py	parse db_config.yaml (db_path/sql_file/dump_folder) → AttachSpecs → DatasetHandle
benchmark/dab/runner.py	per-task: run agent → validate(final_answer) → Pass@1
lexsi_ds/agent/table_index.py + tools/find_relevant_tables.py	schema linking for big many-table DBs (text_to_sql auto-caps too)
tests/test_dab_adapter.py, tests/test_pg_mongo_materialize.py	self-contained (synthetic sqlite/duckdb/pg/mongo dumps), no external checkout

---

6. Extending coverage

Postgres + Mongo are done (embedded pgserver / server-less bson; the same helpers back the live connect_datalake connectors). Remaining:

• PATENTS data — its Postgres dump is a 5 GB Drive download not pulled locally; everything else runs. build_handle reports it cleanly.
• UI dataset picker — the Gradio app hardcodes load_pkdd_handle() in _build_ui_ctx. A dropdown swapping the handle would expose switching (and now connect_datalake postgres/mongo) to users.
• DAB as a bench tier — wire the runner's Pass@1 into the scoreboard alongside PKDD-Curated (see v1_benchmarking.md §4.2).