From bf2cb1fcabfdcee166ce01b7ccf330c1fbbfc89c Mon Sep 17 00:00:00 2001 From: Ankit Khullar Date: Fri, 26 Jun 2026 22:30:41 +0530 Subject: [PATCH 1/9] Add uv setup and Codex-backed BEAVER eval method (myagent) Set up the repo with uv (pyproject.toml, .python-version pinned to 3.10, uv.lock) covering eval + retrieve + data dependencies. Add eval/myagent/: a drop-in BEAVER text-to-SQL method whose agent delegates generation to the local Codex CLI (codex exec). Includes the full pipeline (execute -> unify) mirroring the existing baselines, plus two gold-blind, execution-guided enhancements in agent.py: - self-fix (CODEX_SQL_FIX): run own SQL read-only; repair execution errors from DB error feedback only. - explore/verify (CODEX_SQL_EXPLORE): run read-only queries against the real tables, inspect the rows its own queries return, self-check, then finalize. Neither ever sees gold; all DB access is mediated read-only by the agent process. RESULTS.md records the study on dw (100q): one-shot 23% -> setting 2 + explore/verify + fix 34%, with lever analysis (hints > explore > fix/effort). Raw per-run outputs and unified-output/ (gold SQL from the gated dataset) are gitignored; results are captured as aggregate metrics in RESULTS.md. Co-Authored-By: Claude Opus 4.8 (1M context) --- .python-version | 1 + eval/myagent/.gitignore | 4 + eval/myagent/README.md | 80 ++ eval/myagent/RESULTS.md | 64 + eval/myagent/agent.py | 313 +++++ eval/myagent/execute.py | 91 ++ eval/myagent/prompt.py | 180 +++ eval/myagent/run.sh | 97 ++ eval/myagent/unify.py | 84 ++ eval/myagent/utils.py | 78 ++ pyproject.toml | 36 + uv.lock | 2663 +++++++++++++++++++++++++++++++++++++++ 12 files changed, 3691 insertions(+) create mode 100644 .python-version create mode 100644 eval/myagent/.gitignore create mode 100644 eval/myagent/README.md create mode 100644 eval/myagent/RESULTS.md create mode 100644 eval/myagent/agent.py create mode 100644 eval/myagent/execute.py create mode 100644 eval/myagent/prompt.py create mode 100755 eval/myagent/run.sh create mode 100644 eval/myagent/unify.py create mode 100644 eval/myagent/utils.py create mode 100644 pyproject.toml create mode 100644 uv.lock diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..c8cfe39 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.10 diff --git a/eval/myagent/.gitignore b/eval/myagent/.gitignore new file mode 100644 index 0000000..1f9d07d --- /dev/null +++ b/eval/myagent/.gitignore @@ -0,0 +1,4 @@ +# Raw per-run generation artifacts (regenerable; can be large). The meaningful +# results are captured in RESULTS.md; full SQL outputs live under +# eval/unified-output/ (gitignored: contains gold SQL from the gated dataset). +output/ diff --git a/eval/myagent/README.md b/eval/myagent/README.md new file mode 100644 index 0000000..bab7562 --- /dev/null +++ b/eval/myagent/README.md @@ -0,0 +1,80 @@ +# myagent — Codex-backed text-to-SQL method + +A drop-in BEAVER method folder (mirrors `fewshot/`) whose agent delegates SQL +generation to the local **Codex CLI** via `codex exec` (headless, one-shot). + +## Backend +`agent.py::run_agent` builds a single prompt from the question + schema (+ hints +for the active `--setting`), runs: + +``` +codex exec --sandbox read-only --skip-git-repo-check -C \ + -c model_reasoning_effort= -o "" +``` + +in a throwaway working directory (Codex never touches the repo), and reads the +final agent message back via `--output-last-message`. `…` / ```` ```sql ```` +wrappers are stripped by `clean_sql`. + +Uses your existing Codex login. **With a ChatGPT login, only Codex-supported +models work** (e.g. `gpt-5-codex`, `gpt-5.5`); arbitrary names like `gpt-5-mini` +are rejected — so BEAVER's `--model` is a label only and is *not* forwarded to +Codex. Pick the Codex model with `CODEX_MODEL`. + +### Config (env vars, optional) +| var | default | meaning | +|-----|---------|---------| +| `CODEX_MODEL` | account default | Codex model (`codex -m`) | +| `CODEX_REASONING_EFFORT` | `low` | minimal / low / medium / high / xhigh | +| `CODEX_SANDBOX` | `read-only` | sandbox policy | +| `CODEX_TIMEOUT` | `300` | per codex-call seconds | +| `CODEX_BIN` | `codex` | path to the codex binary | + +### Execution-guided self-fix (optional) +With `CODEX_SQL_FIX=1`, after generating, the agent runs its own SQL against the +live MySQL database; if it **fails to execute**, it feeds *only the database error* +back to Codex and asks for a corrected query, looping up to `CODEX_FIX_ATTEMPTS` +times. It is **gold-blind** — it never sees the correct/expected rows, only whether +its own query errors — and stops as soon as the query executes (rows may still be +wrong). This fixes syntax/execution errors, not semantic correctness. Execution is +read-only guarded (refuses non-SELECT statements) and time-capped. + +| var | default | meaning | +|-----|---------|---------| +| `CODEX_SQL_FIX` | `0` | `1` to enable the execute-and-fix loop | +| `CODEX_FIX_ATTEMPTS` | `2` | max correction rounds on execution error | +| `CODEX_FIX_TIMEOUT_MS` | `15000` | SELECT execution cap during the fix check | + +```bash +CODEX_SQL_FIX=1 CODEX_REASONING_EFFORT=high ./run.sh --dataset dw --setting 1 +``` +Needs the dataset's MySQL DB loaded + `MYSQL_*` creds (read from env or nearest `.env`). + +## Files +| file | role | edit? | +|------|------|-------| +| `agent.py` | Codex-backed `run_agent` (the seam) | swap backend here | +| `prompt.py` | builds per-question context + hints per `--setting` | rarely | +| `execute.py` | runs the agent in parallel (default 4 workers), resume support | no | +| `unify.py` | reshapes output into `unified-output/myagent//{generated,gold}` | no | +| `run.sh` | CLI wrapper: execute → unify | no | + +## Run +```bash +cd eval/myagent +./run.sh --dataset dw --setting 1 # setting 1/2 need no retrieval +# ./run.sh --dataset dw --setting 0 # setting 0 needs retrieve/retrieve.py first +# CODEX_MODEL=gpt-5-codex ./run.sh --dataset dw --setting 1 +``` + +## Score (from `eval/`) +```bash +RUN= +uv run python evaluate_ex_acc.py --dataset dw --input_dir unified-output/myagent/$RUN +uv run python evaluate_subtasks.py --dataset dw --model gpt-5-mini --input_dir unified-output/myagent/$RUN +``` +`evaluate_ex_acc.py` needs MySQL loaded + creds in `.env`. `evaluate_subtasks.py` +needs an LLM key (its `--model` is the *grader*, unrelated to the Codex backend). + +> Each `codex exec` runs ~30 s at `low` effort; keep `--num_workers` modest to +> avoid Codex rate limits. The full 100-question `dev_sampled` run is sized accordingly. diff --git a/eval/myagent/RESULTS.md b/eval/myagent/RESULTS.md new file mode 100644 index 0000000..e9c2fa7 --- /dev/null +++ b/eval/myagent/RESULTS.md @@ -0,0 +1,64 @@ +# myagent — Codex on BEAVER `dw`: results + +Evaluation of the Codex-backed agent (`agent.py`) on the BEAVER `dw` benchmark. +All numbers are **execution accuracy** (generated SQL's result set equals gold's) +on the standard 100-question `dev_sampled` set, scored against the live MySQL +`dw` database. Generation uses the local Codex CLI (model `gpt-5.5` via ChatGPT +login). Every enhancement below is **gold-blind** — the agent never sees the +expected/correct answer. + +## Scoreboard (`dw`, 100 questions, `high` reasoning effort) + +| Config | exec acc | correct | SQL errors | empty | mismatch | +|--------|:--------:|:-------:|:----------:|:-----:|:--------:| +| setting 1, one-shot (baseline) | 23% | 23 | 11 | 2 | 64 | +| setting 1 + self-fix | 23% | 23 | 0 | 5 | 72 | +| setting 1 + explore/verify | 26% | 26 | 5 | 3 | 66 | +| setting 2 + self-fix | 30% | 30 | 1 | 4 | 65 | +| **setting 2 + explore/verify + fix** | **34%** | **34** | **0** | 2 | 64 | + +5-question sanity samples (setting 1): low 20%, high 40%, xhigh 40%. +Full-100, setting 1: high 23%, xhigh 25% (within run-to-run noise). + +## Settings +- **setting 1** — hints: gold tables, column mapping, join keys. +- **setting 2** — setting 1 + domain knowledge + query decomposition. +- (setting 0 = retrieved tables only; not run here — needs `retrieve/`.) + +## Enhancements (all in `agent.py`, all gold-blind) +- **self-fix** (`CODEX_SQL_FIX=1`) — runs its own SQL read-only; on an execution + error, feeds back *only the DB error* and asks Codex to fix, looping up to + `CODEX_FIX_ATTEMPTS`. Stops once the query executes. → guarantees executable + output (SQL errors → 0); does not change correctness on its own. +- **explore/verify** (`CODEX_SQL_EXPLORE=1`) — the agent runs read-only queries + against the real tables (sample rows, counts, its candidate query) and inspects + the rows *its own* queries return, then self-checks and revises before + finalizing. Never shown gold rows. → +3–4 pts. + +## Findings +- **Lever ranking:** richer hints (+7) > explore/verify (+3–4) > self-fix + (robustness, ~0 acc) ≈ reasoning effort (within noise, high↔xhigh). +- Levers **stack independently**: setting 2 (30%) + explore/verify + fix → **34%**. +- **explore and fix are complementary**: explore improves correctness; the final + fix pass drives SQL errors to 0 (explore alone left 5). +- The dominant failure mode is **semantic** (~64 "values mismatch") at every + config — queries execute but return the wrong rows. Even with all 5 oracle + hints (setting 2), ~66% still miss, i.e. SQL *construction* on these enterprise + schemas is hard even when schema-linking is handed over. +- **Headroom:** the union of correct sets across runs (~39) exceeds any single + run, indicating self-consistency / majority-vote ensembling is the next lever. + +## Reproduce +```bash +# best config +cd eval/myagent +CODEX_REASONING_EFFORT=high CODEX_SQL_EXPLORE=1 CODEX_SQL_FIX=1 \ + ./run.sh --dataset dw --setting 2 + +# score (from eval/) +cd .. && uv run python evaluate_ex_acc.py --dataset dw \ + --input_dir unified-output/myagent/ +``` +Requires the `dw` MySQL DB loaded + `MYSQL_*` creds in `.env`. Per-run SQL outputs +and `summary_ex_acc.json` are written under `eval/unified-output/myagent/` +(gitignored — contains gold SQL from the gated dataset). diff --git a/eval/myagent/agent.py b/eval/myagent/agent.py new file mode 100644 index 0000000..3729b8a --- /dev/null +++ b/eval/myagent/agent.py @@ -0,0 +1,313 @@ +""" +============================================================================ + Codex-backed agent for BEAVER text-to-SQL. + Three generation modes (pick via env), all gold-blind: + * plain : one `codex exec` call -> SQL (default) + * self-fix : CODEX_SQL_FIX=1 -> run own SQL, fix execution errors + * explore+verify : CODEX_SQL_EXPLORE=1 -> run read-only queries against the + real tables, see the rows returned, self-check, finalize +============================================================================ + +All DB access is mediated by THIS process: the model emits queries through a +text protocol and we execute them read-only and feed back the rows. The model +never gets credentials, network, or filesystem access to the gold files. It can +see the INPUT data (tables) and the rows ITS OWN queries return — never the +correct/expected answer (gold SQL/results live in the JSON files, not in MySQL). + + instance fields (built in prompt.py::build_instances): + id, question, db, tables, prompt (chat messages), hints, record + -> NOTE: do not read record['sql']; that is the gold answer. + +Config (env vars, all optional): + CODEX_MODEL Codex model (`codex -m`); unset -> account default. + CODEX_REASONING_EFFORT minimal|low|medium|high|xhigh (default: low) + CODEX_SANDBOX read-only|workspace-write|danger-full-access (default: read-only) + CODEX_TIMEOUT per codex-call seconds (default: 300) + CODEX_BIN path to codex (default: codex) + + CODEX_SQL_FIX 1 -> fix execution errors via error feedback (default 0) + CODEX_FIX_ATTEMPTS max fix rounds (default 2) + + CODEX_SQL_EXPLORE 1 -> explore/verify loop (overrides SQL_FIX) (default 0) + CODEX_EXPLORE_STEPS max exploratory query rounds (default 4) + CODEX_EXPLORE_ROWS max rows returned per exploratory query (default 20) + CODEX_FIX_TIMEOUT_MS SELECT execution cap, ms (default 15000) + MYSQL_HOST/USER/PASSWORD DB creds (env or nearest .env) +""" +import os +import re +import shutil +import tempfile +import subprocess + +CODEX_BIN = os.getenv("CODEX_BIN", "codex") +CODEX_MODEL = os.getenv("CODEX_MODEL") +CODEX_REASONING_EFFORT = os.getenv("CODEX_REASONING_EFFORT", "low") +CODEX_SANDBOX = os.getenv("CODEX_SANDBOX", "read-only") +CODEX_TIMEOUT = int(os.getenv("CODEX_TIMEOUT", "300")) + +CODEX_SQL_FIX = os.getenv("CODEX_SQL_FIX", "0") not in ("0", "", "false", "False") +CODEX_FIX_ATTEMPTS = int(os.getenv("CODEX_FIX_ATTEMPTS", "2")) + +CODEX_SQL_EXPLORE = os.getenv("CODEX_SQL_EXPLORE", "0") not in ("0", "", "false", "False") +CODEX_EXPLORE_STEPS = int(os.getenv("CODEX_EXPLORE_STEPS", "4")) +CODEX_EXPLORE_ROWS = int(os.getenv("CODEX_EXPLORE_ROWS", "20")) +CODEX_FIX_TIMEOUT_MS = int(os.getenv("CODEX_FIX_TIMEOUT_MS", "15000")) + +_READ_STMTS = ("select", "with", "show", "describe", "desc", "explain") +_MAX_FEEDBACK_CHARS = 4000 +_MAX_CELL_CHARS = 80 + + +def clean_sql(text: str) -> str: + """Strip tags and ```sql fences from a model response.""" + if not text: + return "" + if "" in text: + text = text.split("")[-1] + if "" in text: + text = text.split("")[0] + if "```sql" in text: + text = text.split("```sql")[-1] + if "```" in text: + text = text.split("```")[0] + return text.strip() + + +def render_prompt(instance: dict) -> str: + """Flatten the chat-style prompt into the single string `codex exec` expects.""" + blocks = [] + user_seen = 0 + for msg in instance["prompt"]: + role, content = msg["role"], msg["content"] + if role == "system": + blocks.append(content) + elif role == "assistant": + blocks.append(f"### Example answer\n{content}") + elif role == "user": + user_seen += 1 + header = "### Example input" if user_seen == 1 else "### Now answer this" + blocks.append(f"{header}\n{content}") + return "\n\n".join(blocks) + + +def _codex_call(prompt: str, model: str) -> str: + """One headless `codex exec` call -> raw final message text.""" + workdir = tempfile.mkdtemp(prefix="codex_beaver_") + last_msg = os.path.join(workdir, "_last_message.txt") + cmd = [ + CODEX_BIN, "exec", "--sandbox", CODEX_SANDBOX, "--skip-git-repo-check", + "-C", workdir, "-c", f"model_reasoning_effort={CODEX_REASONING_EFFORT}", "-o", last_msg, + ] + if CODEX_MODEL: + cmd += ["-m", CODEX_MODEL] + cmd.append(prompt) + try: + proc = subprocess.run( + cmd, stdin=subprocess.DEVNULL, capture_output=True, text=True, timeout=CODEX_TIMEOUT, + ) + raw = "" + if os.path.exists(last_msg): + with open(last_msg) as f: + raw = f.read() + if not raw.strip() and proc.returncode != 0: + tail = (proc.stdout or "") + (proc.stderr or "") + raise RuntimeError(f"codex exec failed (rc={proc.returncode}): {tail[-800:].strip()}") + return raw + except subprocess.TimeoutExpired: + raise RuntimeError(f"codex exec timed out after {CODEX_TIMEOUT}s") + finally: + shutil.rmtree(workdir, ignore_errors=True) + + +def _codex_generate(prompt: str, model: str) -> str: + return clean_sql(_codex_call(prompt, model)) + + +# ----------------------- read-only DB access (gold-blind) ----------------------- + +def _load_db_creds(): + if not os.getenv("MYSQL_HOST"): + try: + from dotenv import load_dotenv, find_dotenv + load_dotenv(find_dotenv(usecwd=True)) + except Exception: + pass + return os.getenv("MYSQL_HOST", "localhost"), os.getenv("MYSQL_USER", "root"), os.getenv("MYSQL_PASSWORD", "") + + +def _is_read_only(sql: str) -> bool: + first = (sql.lstrip().split(None, 1)[0].lower() if sql.strip() else "") + return first in _READ_STMTS + + +def _connect(db): + import mysql.connector + return mysql.connector.connect( + host=_load_db_creds()[0], user=_load_db_creds()[1], password=_load_db_creds()[2], + database=db, connection_timeout=10, + ) + + +def _execute_sql(sql: str, db: str): + """Run read-only with a time cap. Returns (ok, error_str). Never inspects rows.""" + if not _is_read_only(sql): + return False, f"refusing to execute non-read statement (starts with {sql.split(None,1)[:1]})" + conn = None + try: + conn = _connect(db) + cur = conn.cursor() + try: + cur.execute(f"SET SESSION max_execution_time={CODEX_FIX_TIMEOUT_MS}") + except Exception: + pass + cur.execute(sql) + cur.fetchmany(1) + return True, "" + except Exception as e: + return False, str(e) + finally: + try: + conn and conn.close() + except Exception: + pass + + +def _query_preview(sql: str, db: str, max_rows: int): + """Run read-only and return a compact text preview of up to max_rows rows. + Returns feedback string for the model (cols + rows, or the error).""" + if not _is_read_only(sql): + return "ERROR: only read-only queries (SELECT/WITH/SHOW/DESCRIBE/EXPLAIN) are allowed here." + conn = None + try: + conn = _connect(db) + cur = conn.cursor() + try: + cur.execute(f"SET SESSION max_execution_time={CODEX_FIX_TIMEOUT_MS}") + except Exception: + pass + cur.execute(sql) + rows = cur.fetchmany(max_rows) + cols = [d[0] for d in cur.description] if cur.description else [] + extra = "" + if len(rows) == max_rows: + extra = f"\n... (truncated at {max_rows} rows)" + def fmt(v): + s = "NULL" if v is None else str(v) + return s if len(s) <= _MAX_CELL_CHARS else s[:_MAX_CELL_CHARS] + "…" + lines = [" | ".join(cols)] + for r in rows: + lines.append(" | ".join(fmt(v) for v in r)) + body = "\n".join(lines) + extra + if len(body) > _MAX_FEEDBACK_CHARS: + body = body[:_MAX_FEEDBACK_CHARS] + "\n… (output truncated)" + return f"{len(rows)} row(s) returned:\n{body}" + except Exception as e: + return f"ERROR: {e}" + finally: + try: + conn and conn.close() + except Exception: + pass + + +# ------------------------------- modes ------------------------------- + +def _fix_prompt(base, db, bad_sql, error): + return ( + base + f"\n\n### Execution feedback\nThe query below was run against the `{db}` MySQL " + f"database and FAILED TO EXECUTE. Fix it so it runs without error, keeping the intended " + f"logic.\n\nFailed SQL:\n{bad_sql}\n\nDatabase error:\n{error}\n\n" + f"Return only the corrected MySQL query wrapped in ." + ) + + +def _fix_loop(base, model, db, sql): + """Given a candidate SQL, repair execution errors (error feedback only).""" + if not sql: + return sql + for _ in range(CODEX_FIX_ATTEMPTS): + ok, err = _execute_sql(sql, db) + if ok: + break + fixed = _codex_generate(_fix_prompt(base, db, sql, err), model) + if not fixed or fixed == sql: + break + sql = fixed + return sql + + +_EXPLORE_PROTOCOL = """\ + +### Database access (read-only) — verification is REQUIRED +You have live read-only access to the `{db}` MySQL database. You will be shown the rows YOUR queries return — you will NOT be shown the expected/correct answer. + +You MUST run at least one query to CHECK your candidate answer before finalizing: inspect the real tables (sample rows, distinct values, counts, ranges) to confirm your assumptions, then run your candidate query and verify the returned rows make sense for the question (right columns, plausible row count, filters/joins working, not empty when it shouldn't be). Revise if the results look wrong. + +Respond with EXACTLY ONE of the following each turn (nothing else): + +1) To run a read-only query (SELECT/WITH/SHOW/DESCRIBE/EXPLAIN), output: +RUN_SQL: + + +2) Only after you have verified, output your final answer as: +YOUR FINAL MYSQL QUERY + +You have at most {steps} queries. Do not give until you have run at least one verification query.\ +""" + + +def _parse_action(text: str): + """Return ('answer', sql) or ('run', query) from a model turn.""" + if "" in text: + return "answer", clean_sql(text) + m = re.search(r"RUN_SQL:\s*", text) + if m: + q = text[m.end():].strip() + # strip code fences if present + if "```" in q: + q = q.split("```")[1] if q.count("```") >= 2 else q.replace("```sql", "").replace("```", "") + q = q.replace("sql\n", "", 1).strip() if q.lower().startswith("sql") else q.strip() + # cut at a trailing if model appended one + q = q.split("")[0].strip() + return "run", q.strip() + # no protocol token -> treat as a (possibly fenced) final SQL + return "answer", clean_sql(text) + + +def _run_explore(base, model, db): + transcript = base + _EXPLORE_PROTOCOL.format(db=db, steps=CODEX_EXPLORE_STEPS) + last_sql = "" + for step in range(CODEX_EXPLORE_STEPS): + resp = _codex_call(transcript, model) + action, payload = _parse_action(resp) + if action == "answer" and payload: + return payload + if action == "run" and payload: + last_sql = payload + result = _query_preview(payload, db, CODEX_EXPLORE_ROWS) + transcript += ( + f"\n\n### Your query (step {step + 1})\n{payload}\n\n### Result\n{result}\n\n" + f"Run another RUN_SQL query, or output your final ...." + ) + else: + break + # out of steps (or no parseable action): force a final answer + final = _codex_call( + transcript + "\n\nYou must now output ONLY your final answer as YOUR MYSQL QUERY.", + model, + ) + return clean_sql(final) or last_sql + + +def run_agent(instance: dict, model: str) -> str: + base = render_prompt(instance) + db = instance.get("db") or "dw" + if CODEX_SQL_EXPLORE: + sql = _run_explore(base, model, db) + if CODEX_SQL_FIX: # final error-repair pass on the explored answer + sql = _fix_loop(base, model, db, sql) + return sql + if CODEX_SQL_FIX: + return _fix_loop(base, model, db, _codex_generate(base, model)) + return _codex_generate(base, model) diff --git a/eval/myagent/execute.py b/eval/myagent/execute.py new file mode 100644 index 0000000..ea9c2fc --- /dev/null +++ b/eval/myagent/execute.py @@ -0,0 +1,91 @@ +"""Generation driver for a custom agent. + +Builds one instance per question, runs your agent (agent.run_agent) over them in +parallel with resume support, and writes each prediction to + //predicted_0.sql (+ generation.log) +which unify.py then reshapes into the unified generated/ + gold/ layout. + +You normally do NOT edit this file — put your agent in agent.py. +""" +import os +import argparse +import traceback +from concurrent.futures import ThreadPoolExecutor + +from tqdm import tqdm + +from prompt import build_instances +from utils import EvalConfig +from agent import run_agent + + +def _output_paths(output_dir, instance_id): + instance_dir = os.path.join(output_dir, instance_id) + return instance_dir, os.path.join(instance_dir, "predicted_0.sql"), os.path.join(instance_dir, "generation.log") + + +def _process(instance, model, output_dir): + instance_id = instance["id"] + instance_dir, sql_file, log_file = _output_paths(output_dir, instance_id) + try: + sql = run_agent(instance, model) or "" + log = f"Question:\n{instance['question']}\n\nGenerated SQL:\n{sql}" + except Exception as e: + sql = "" + log = f"Error: {e}\n{traceback.format_exc()}" + print(f"Error on {instance_id}: {e}") + + os.makedirs(instance_dir, exist_ok=True) + with open(sql_file, "w") as f: + f.write(sql) + with open(log_file, "w") as f: + f.write(log) + return instance_id, bool(sql) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=str, default="codex") + parser.add_argument("--dataset", type=str, default="dw") + parser.add_argument("--q_fn", type=str, default="dev_sampled") + parser.add_argument("--output_dir", type=str, required=True) + parser.add_argument("--data_dir", type=str, default="../../data") + parser.add_argument("--num_workers", type=int, default=4) + # --setting flags (set by run.sh) + parser.add_argument("--gold_tables", action="store_true") + parser.add_argument("--join_keys", action="store_true") + parser.add_argument("--mapping", action="store_true") + parser.add_argument("--knowledge", action="store_true") + parser.add_argument("--decomp", action="store_true") + args = parser.parse_args() + + eval_config = EvalConfig( + gold_tables=args.gold_tables, join_keys=args.join_keys, + mapping=args.mapping, knowledge=args.knowledge, decomp=args.decomp, + ) + print(eval_config) + + instances = build_instances(args.dataset, args.q_fn, eval_config, args.data_dir) + os.makedirs(args.output_dir, exist_ok=True) + + # Resume: skip instances that already have a non-empty prediction. + todo = [] + for inst in instances: + _, sql_file, _ = _output_paths(args.output_dir, inst["id"]) + if not os.path.exists(sql_file) or os.path.getsize(sql_file) == 0: + todo.append(inst) + + print(f"Output directory: {args.output_dir}") + if not todo: + print("All instances already processed.") + return + print(f"Running agent on {len(todo)} / {len(instances)} instances (model={args.model})...") + + with ThreadPoolExecutor(max_workers=args.num_workers) as ex: + futures = [ex.submit(_process, inst, args.model, args.output_dir) for inst in todo] + for fut in tqdm(futures, total=len(futures)): + fut.result() + + +if __name__ == "__main__": + main() diff --git a/eval/myagent/prompt.py b/eval/myagent/prompt.py new file mode 100644 index 0000000..c930239 --- /dev/null +++ b/eval/myagent/prompt.py @@ -0,0 +1,180 @@ +"""Build per-question instances for the agent. + +This mirrors the context/hint assembly used by the `fewshot` baseline +(`eval/fewshot/prompt.py`) so the --setting flag behaves identically, but it +returns a list of *aligned* instance dicts instead of two parallel lists. Each +instance carries both: + * a ready-to-send OpenAI-style chat `prompt` (system + few-shot example + user) + for agents that just want to call a chat model, and + * the structured fields (question, db, tables, and the active hint strings) for + agents that do their own multi-step reasoning / tool use. +""" +from pathlib import Path + +from utils import ( + read_json, format_tables, EvalConfig, format_join, system, user, assistant, +) + + +def get_mapping(mapping): + desc = [] + for sq in mapping: + cols_desc = [] + for col in mapping[sq]: + table_name, col = col.split(".") + cols_desc.append(f"column {col} in table {table_name.lower()}") + desc.append(f'"{sq}" in the user question refers to {", ".join(cols_desc)}') + return "\n".join(desc) + + +def get_join_keys(join_keys): + return "\n".join(format_join(jk) for jk in join_keys) + + +def get_knowledge(domain_knowledge): + return "\n".join(domain_knowledge) + + +def get_decomp(decomps): + return "\n".join(f"Subquery {i + 1}: {sq}" for i, sq in enumerate(decomps)) + + +def get_user_prompt(q, tables, corpus_tables, eval_config: EvalConfig): + desc = [ + format_tables(tables, corpus_tables, eval_config.instances), + f"User question: {q['question']}", + ] + if eval_config.mapping: + desc.append(f"Mapping:\n{get_mapping(q['column_mapping'])}") + if eval_config.join_keys: + desc.append(f"Join keys:\n{get_join_keys(q['join_keys'])}") + if eval_config.knowledge and get_knowledge(q['domain_knowledge']) != '': + desc.append(f"Domain knowledge:\n{get_knowledge(q['domain_knowledge'])}") + if eval_config.decomp and get_decomp(q["sub_questions"]) != "": + desc.append(f"Query decomposition:\n{get_decomp(q['sub_questions'])}") + return user("\n\n".join(desc)) + + +def get_retrieved_tables(dataset: str, data_dir="../../data"): + retrieved_fn = Path(f"{data_dir}/{dataset}/retrieval/retrieved_tables.json") + reranked_fn = Path(f"{data_dir}/{dataset}/retrieval/reranked_tables.json") + if not retrieved_fn.exists(): + raise FileNotFoundError( + f"No retrieved tables at {retrieved_fn}. At --setting 0 the agent gets the " + f"retrieved candidate tables; run `python retrieve/retrieve.py --dataset {dataset} ...` " + f"first, or use --setting 1/2 which inject gold tables and need no retrieval." + ) + if reranked_fn.exists(): + print(f"Loading reranked tables from {reranked_fn}") + return read_json(reranked_fn) + print(f"Loading retrieved tables from {retrieved_fn}") + return read_json(retrieved_fn) + + +def _build_instruction(eval_config, q_knowledge, q_decomp, q, structures): + db_type = "MySQL" + instruction = ["You are given a list of tables", "a user question"] + if eval_config.join_keys: + instruction.append("join keys among the provided tables") + if eval_config.mapping: + instruction.append("a mapping from information mentioned in the user question to columns in the provided tables") + if q_knowledge: + instruction.append("domain knowledge") + if q_decomp: + instruction.append("decomposition of the user question") + instruction[-1] = f"and {instruction[-1]}" + instruction = ", ".join(instruction) + ", " + + instruction += ( + f"your task is output a {db_type} SQL statement that can be used to answer the user " + f"question based on the provided information. You need to ensure that syntax and functions " + f"used in your SQL statement are appropriate for {db_type} database. If you are unable to " + f"determine the SQL statement, output None. " + ) + if eval_config.mapping: + instruction += "You should use the provided mapping to determine which columns and tables should be used in the SQL statement. " + if eval_config.join_keys: + instruction += "You should use the provided join keys to determine how to connect the tables in the SQL statement. " + if q_knowledge: + instruction += "You should use the provided domain knowledge to determine which tables, columns, and literals should be used in the SQL statement. " + if q_decomp: + instruction += ( + "You must answer each subquery individually and then combine them to form the complete " + "SQL statement. Each subquery you generate must be explicitly used in the final SQL " + "statement, without being simplified. " + ) + instruction += ( + "Below is the structure of the SQL statement with subqueries denoted. Each provided " + "subquery is used in the final SQL statement in such a structure." + ) + structure_name = q.get("detailed_category") + if structure_name and structure_name != 'real' and structure_name in structures: + structure = structures[structure_name] + instruction += f"\n\n{structure['structure']}" + instruction += f"\n\n{structure['subquery_decomposition']} " + + instruction += "The SQL statement need to be wrapped in tags." + return instruction + + +def build_instances(dataset: str, q_fn: str, eval_config: EvalConfig, data_dir: str = "../../data"): + """Return a list of aligned instance dicts, one per question that has table context. + + Each instance: + id : question id + question : natural-language question + db : target database name + tables : list of candidate table names provided to the agent + prompt : OpenAI-style chat messages (system + few-shot example + user) + hints : {mapping, join_keys, domain_knowledge, decomposition} active strings + record : full question/gold dict from .json (avoid peeking at gold `sql`) + """ + structures = read_json(f"{data_dir}/template_structure.json") + qs = read_json(f"{data_dir}/{dataset}/{q_fn}.json") + example = read_json(f"{data_dir}/{dataset}/example.json") + dev_tables = read_json(f"{data_dir}/{dataset}/dev_tables.json") + + retrieved_tables = None if eval_config.gold_tables else get_retrieved_tables(dataset, data_dir) + + example_prompt = [ + get_user_prompt(example, example["tables"], dev_tables, eval_config), + assistant(f"SQL: {example['sql']}"), + ] + + instances = [] + skipped = 0 + for q in qs: + q_id = q['id'] + q_knowledge = eval_config.knowledge and get_knowledge(q['domain_knowledge']) != '' + q_decomp = eval_config.decomp and get_decomp(q["sub_questions"]) != '' + + if eval_config.gold_tables: + tables = q["tables"] + else: + if q_id not in retrieved_tables: + skipped += 1 + continue + tables = retrieved_tables[q_id] + + instruction = _build_instruction(eval_config, q_knowledge, q_decomp, q, structures) + prompt = [system(instruction)] + example_prompt + [get_user_prompt(q, tables, dev_tables, eval_config)] + + instances.append({ + "id": q_id, + "question": q["question"], + "db": q.get("db", dataset), + "tables": tables, + "prompt": prompt, + "hints": { + "mapping": get_mapping(q['column_mapping']) if eval_config.mapping else None, + "join_keys": get_join_keys(q['join_keys']) if eval_config.join_keys else None, + "domain_knowledge": get_knowledge(q['domain_knowledge']) if q_knowledge else None, + "decomposition": get_decomp(q['sub_questions']) if q_decomp else None, + }, + "record": q, + }) + + if skipped: + print(f"Skipped {skipped} question(s) with no retrieved tables.") + print(f"#instances: {len(instances)}") + return instances diff --git a/eval/myagent/run.sh b/eval/myagent/run.sh new file mode 100755 index 0000000..9a44371 --- /dev/null +++ b/eval/myagent/run.sh @@ -0,0 +1,97 @@ +#!/bin/bash +set -e + +# ============================================================================ +# Custom-agent generation pipeline for BEAVER +# Usage: ./run.sh --model gpt-5-mini --dataset dw --setting 0 +# +# Step 1 (execute.py): run your agent (agent.py) over every question +# Step 2 (unify.py): reshape outputs into unified-output/myagent// +# +# Then score with (from eval/): +# uv run python evaluate_ex_acc.py --dataset --input_dir unified-output/myagent/ +# uv run python evaluate_subtasks.py --dataset --model gpt-5-mini --input_dir unified-output/myagent/ +# ============================================================================ + +MODEL="codex" +DATASET="dw" +SETTING=0 +NUM_WORKERS=4 + +while [[ $# -gt 0 ]]; do + case $1 in + --model) MODEL="$2"; shift 2 ;; + --dataset) DATASET="$2"; shift 2 ;; + --setting) SETTING="$2"; shift 2 ;; + --num_workers) NUM_WORKERS="$2"; shift 2 ;; + --help|-h) + echo "Usage: ./run.sh --model