# Dataset switching + DataAgentBench integration

How the agent moves off the bundled PKDD dataset, and how
[DataAgentBench](https://github.com/ucbepic/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

```bash
# 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`) → `AttachSpec`s → `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`](v1_benchmarking.md) §4.2).