diff --git a/examples/grpo_sqlite_sql/README.md b/examples/grpo_sqlite_sql/README.md new file mode 100644 index 0000000000..f899f174f2 --- /dev/null +++ b/examples/grpo_sqlite_sql/README.md @@ -0,0 +1,70 @@ +# Example: Multi-turn GRPO for Text-to-SQL on SQLite (BIRD) + +This example trains an agent that explores a real SQLite database with read-only queries and then +submits one SQL query, using the [BIRD](https://bird-bench.github.io/) benchmark. It complements the +[InterCode-SQL example](../grpo_intercode_sql/README.md): it needs **no Docker or MySQL server**, +because each rollout opens the `.sqlite` file directly. + +## How the workflow works + +`sqlite_sql_workflow` (`trinity/common/workflows/envs/sqlite_sql/`) gives the model the database +schema, the question and BIRD's external knowledge. At each turn the model replies with exactly one +of: + +- `SQL`: run a read-only query; the first `max_obs_rows` rows (or the error) come back + as the next user message +- `SQL`: submit the final query and end the episode + +The reward is BIRD / Spider execution accuracy: 1.0 if the submitted query returns the same set of +rows as the reference query, else 0.0. Episodes that run out of turns score 0 unless +`auto_submit: true`, which grades the last query that ran successfully. + +Because the reward comes from executing model-written SQL, every query runs in a sandbox: the +database is opened read-only behind an SQLite authorizer that only allows reads (no writes, `ATTACH`, +`PRAGMA` or extensions), one statement is executed per query, and a query is interrupted after +`timeout` seconds. Concurrent rollouts therefore cannot change the data other rollouts are scored +on. Experience `info` also records `gold_error` (the reference query fails) and `gold_empty` (the +reference returns no rows, so any query returning nothing would match), so these tasks can be +filtered. + +## Prepare data + +1. Download the BIRD train set and databases from , and BIRD + Mini-Dev (500 SQLite questions, disjoint from train) from + . +2. Create the taskset: + +```bash +python examples/grpo_sqlite_sql/get_bird_data.py \ + --train_json train/train.json \ + --test_json minidev/MINIDEV/mini_dev_sqlite.json +``` + +Each line of `bird_data/train.jsonl` and `bird_data/test.jsonl` has `db_id`, `question`, `evidence` +and `SQL`. Any dataset in this format with `//.sqlite` databases (e.g. Spider) +works the same way. + +## Run + +```bash +export BIRD_TRAIN_DB_ROOT=/path/to/train/train_databases +export BIRD_DEV_DB_ROOT=/path/to/minidev/MINIDEV/dev_databases +trinity run --config examples/grpo_sqlite_sql/sqlite_sql.yaml +``` + +Workflow arguments (`workflow_args`): + +| Argument | Default | Description | +|---|---|---| +| `db_root` | required | Directory with `/.sqlite` | +| `max_env_steps` | 5 | Maximum number of model turns | +| `max_obs_rows` | 10 | Rows of a query result shown to the model | +| `timeout` | 30 | Per-query timeout in seconds (the official BIRD evaluator uses 30) | +| `auto_submit` | false | Grade the last successful query if the model never submits | + +Notes on BIRD Mini-Dev as a reward source: + +- 16 of the 500 reference queries return the single value `1`, so a constant `SELECT 1` earns the + full reward on 3.2% of the tasks. Watch rollouts for constant answers on hard tasks. +- Two reference queries (question ids 518 and 701) take longer than 30 s on a typical machine, + so they score 0 for every model at the default timeout, as they do in the official evaluator. diff --git a/examples/grpo_sqlite_sql/get_bird_data.py b/examples/grpo_sqlite_sql/get_bird_data.py new file mode 100644 index 0000000000..e40a276b06 --- /dev/null +++ b/examples/grpo_sqlite_sql/get_bird_data.py @@ -0,0 +1,53 @@ +""" +Create the Trinity taskset for `sqlite_sql_workflow` from BIRD. + +Download the BIRD train set (https://bird-bench.github.io/) and BIRD Mini-Dev +(https://bird-bench.oss-cn-beijing.aliyuncs.com/minidev.zip), then run: + + python examples/grpo_sqlite_sql/get_bird_data.py \ + --train_json train/train.json \ + --test_json minidev/MINIDEV/mini_dev_sqlite.json + +This writes `examples/grpo_sqlite_sql/bird_data/train.jsonl` and `test.jsonl` with the fields +the workflow reads: `db_id`, `question`, `evidence` and `SQL`. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +FIELDS = ("db_id", "question", "evidence", "SQL") + + +def load(path: Path) -> list[dict]: + text = path.read_text(encoding="utf-8") + if path.suffix == ".jsonl": + return [json.loads(line) for line in text.splitlines() if line.strip()] + return json.loads(text) + + +def write(records: list[dict], path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + for record in records: + row = {key: record.get(key) or "" for key in FIELDS} + f.write(json.dumps(row, ensure_ascii=False) + "\n") + print(f"Wrote {len(records)} tasks to {path}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--train_json", type=Path, required=True) + parser.add_argument("--test_json", type=Path, required=True) + parser.add_argument( + "--output_dir", type=Path, default=Path("examples/grpo_sqlite_sql/bird_data") + ) + args = parser.parse_args() + write(load(args.train_json), args.output_dir / "train.jsonl") + write(load(args.test_json), args.output_dir / "test.jsonl") + + +if __name__ == "__main__": + main() diff --git a/examples/grpo_sqlite_sql/sqlite_sql.yaml b/examples/grpo_sqlite_sql/sqlite_sql.yaml new file mode 100644 index 0000000000..a98b8a649a --- /dev/null +++ b/examples/grpo_sqlite_sql/sqlite_sql.yaml @@ -0,0 +1,81 @@ +project: "SQLITE_SQL" +name: "BIRD_SQLITE_SQL_RFT" +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +algorithm: + algorithm_type: grpo + repeat_times: 8 + optimizer: + lr: 1e-6 +model: + model_path: ${oc.env:TRINITY_MODEL_PATH,Qwen/Qwen2.5-7B-Instruct} + max_prompt_tokens: 12288 + max_response_tokens: 1024 + max_model_len: 13312 +cluster: + node_num: 1 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 16 + explorer_input: + taskset: + name: bird_sqlite_sql + storage_type: file + path: 'examples/grpo_sqlite_sql/bird_data' + split: train + format: + prompt_key: 'question' + rollout_args: + temperature: 1.0 + logprobs: 0 + workflow_args: + db_root: ${oc.env:BIRD_TRAIN_DB_ROOT,train/train_databases} + max_env_steps: 5 + max_obs_rows: 10 + timeout: 30 + eval_tasksets: + - name: bird_minidev_eval + storage_type: file + path: 'examples/grpo_sqlite_sql/bird_data' + split: test + format: + prompt_key: 'question' + rollout_args: + temperature: 0.0 + logprobs: 0 + workflow_args: + db_root: ${oc.env:BIRD_DEV_DB_ROOT,minidev/MINIDEV/dev_databases} + max_env_steps: 5 + max_obs_rows: 10 + timeout: 30 + repeat_times: 1 + default_workflow_type: 'sqlite_sql_workflow' + trainer_input: + experience_buffer: + name: sqlite_sql_buffer + storage_type: queue + path: 'sqlite:///sqlite_sql.db' +explorer: + eval_on_startup: true + eval_interval: 20 + runner_per_model: 16 + max_timeout: 3600 + rollout_model: + engine_num: 2 + tensor_parallel_size: 2 + enable_prefix_caching: true + enforce_eager: true + dtype: bfloat16 + seed: 42 + gpu_memory_utilization: 0.7 + enable_chunked_prefill: true +synchronizer: + sync_method: 'nccl' + sync_interval: 8 + sync_timeout: 1200 +trainer: + save_interval: 100 + grad_clip: 1.0 + use_dynamic_bsz: true + max_token_len_per_gpu: 13312 + ulysses_sequence_parallel_size: 1 diff --git a/tests/common/sqlite_sql_env_test.py b/tests/common/sqlite_sql_env_test.py new file mode 100644 index 0000000000..ca3abfda2f --- /dev/null +++ b/tests/common/sqlite_sql_env_test.py @@ -0,0 +1,109 @@ +import sqlite3 + +import pytest + +from trinity.common.workflows.envs.sqlite_sql.sqlite_sql_env import ( + SQLiteSQLEnv, + db_file, + execute_sql, + parse_response, + schema_text, +) + +GOLD = "SELECT name FROM customers WHERE city = 'Hangzhou'" + + +@pytest.fixture() +def db_path(tmp_path): + path = db_file(str(tmp_path), "shop") + path.parent.mkdir() + conn = sqlite3.connect(path) + conn.executescript( + """ + CREATE TABLE customers(id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT, city TEXT); + INSERT INTO customers(name, city) VALUES ('Ann','Hangzhou'),('Bob','Shanghai'),('Cat','Hangzhou'); + """ + ) + conn.commit() + conn.close() + return path + + +def test_parse_response(): + assert parse_response("xSELECT 1;") == ("action", "SELECT 1") + assert parse_response("\n```sql\nSELECT 2\n```\n") == ("answer", "SELECT 2") + assert parse_response("SELECT 1 then SELECT 3") == ( + "answer", + "SELECT 3", + ) + # tags inside the thinking section are ignored + assert parse_response("SELECT 9") == (None, "") + assert parse_response("SELECT 1") == (None, "") + + +def test_schema_text_skips_internal_tables(db_path): + schema = schema_text(db_path) + assert "CREATE TABLE customers" in schema + assert "sqlite_sequence" not in schema + + +def test_explore_then_answer(db_path): + env = SQLiteSQLEnv(db_path, GOLD, max_obs_rows=2) + + obs = env.step("SELECT name, city FROM customers ORDER BY id") + assert not obs.done and obs.error is None + assert obs.observation.splitlines()[0] == "name | city" + assert "Ann | Hangzhou" in obs.observation and "Cat" not in obs.observation + assert "(3 rows in total)" in obs.observation + + result = env.step("SELECT name FROM customers WHERE city LIKE 'Hang%'") + assert result.done and result.reward == 1.0 + assert env.submitted and env.num_queries == 1 + + +def test_wrong_answer(db_path): + result = SQLiteSQLEnv(db_path, GOLD).step("SELECT name FROM customers") + assert result.done and result.reward == 0.0 + + +@pytest.mark.parametrize( + "sql", + [ + "DELETE FROM customers", + "UPDATE customers SET city = 'Hangzhou'", + "DROP TABLE customers", + "ATTACH DATABASE ':memory:' AS other", + "PRAGMA table_info(customers)", + "SELECT 1; DELETE FROM customers", + ], +) +def test_writes_are_rejected_and_do_not_reach_the_database(db_path, sql): + env = SQLiteSQLEnv(db_path, GOLD) + obs = env.step(f"{sql}") + assert obs.observation.startswith("Error:") and not obs.done + assert env.num_errors == 1 + rows, _, _ = execute_sql(db_path, "SELECT COUNT(*) FROM customers") + assert rows == [(3,)] + + +def test_runaway_query_is_interrupted(db_path): + env = SQLiteSQLEnv(db_path, GOLD, timeout=0.5) + obs = env.step( + "WITH RECURSIVE r(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM r) " + "SELECT COUNT(*) FROM r" + ) + assert obs.observation == "Error: timeout" + + +def test_invalid_format_keeps_episode_open(db_path): + obs = SQLiteSQLEnv(db_path, GOLD).step("I think the answer is Ann.") + assert not obs.done and obs.error == "no action" + + +def test_broken_or_empty_reference(db_path): + broken = SQLiteSQLEnv(db_path, "SELECT nope FROM customers") + assert broken.gold_error is not None + assert broken.step(f"{GOLD}").reward == 0.0 + + empty = SQLiteSQLEnv(db_path, "SELECT name FROM customers WHERE city = 'Beijing'") + assert empty.gold_empty diff --git a/trinity/common/workflows/__init__.py b/trinity/common/workflows/__init__.py index d862218397..2db48c7bd2 100644 --- a/trinity/common/workflows/__init__.py +++ b/trinity/common/workflows/__init__.py @@ -36,6 +36,7 @@ "sciworld_workflow": "trinity.common.workflows.envs.sciworld.sciworld_workflow.SciWorldWorkflow", "webshop_workflow": "trinity.common.workflows.envs.webshop.webshop_workflow.WebShopWorkflow", "intercode_sql_workflow": "trinity.common.workflows.envs.intercode_sql.intercode_sql_workflow.InterCodeSQLWorkflow", + "sqlite_sql_workflow": "trinity.common.workflows.envs.sqlite_sql.sqlite_sql_workflow.SQLiteSQLWorkflow", # general multi-turn "step_wise_alfworld_workflow": "trinity.common.workflows.envs.alfworld.alfworld_workflow.StepWiseAlfworldWorkflow", # ruler/llm_as_a_judge diff --git a/trinity/common/workflows/envs/sqlite_sql/__init__.py b/trinity/common/workflows/envs/sqlite_sql/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/trinity/common/workflows/envs/sqlite_sql/sqlite_sql_env.py b/trinity/common/workflows/envs/sqlite_sql/sqlite_sql_env.py new file mode 100644 index 0000000000..e6bb156d54 --- /dev/null +++ b/trinity/common/workflows/envs/sqlite_sql/sqlite_sql_env.py @@ -0,0 +1,162 @@ +# -*- coding: utf-8 -*- +"""A Docker-free, read-only SQLite environment for multi-turn text-to-SQL (BIRD / Spider layout). + +The agent may run exploratory queries and then submits one final query. The reward is BIRD / Spider +execution accuracy: 1.0 when the submitted query returns the same set of rows as the reference. + +Every rollout opens its own read-only connection behind an SQLite authorizer that only allows reads, +so concurrent rollouts on the same database file cannot write to it, ATTACH other files, run +PRAGMAs or load extensions, and a runaway query is interrupted after ``timeout`` seconds. +""" + +import re +import sqlite3 +import time +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional, Tuple + +_ALLOWED_ACTIONS = { + sqlite3.SQLITE_SELECT, + sqlite3.SQLITE_READ, + sqlite3.SQLITE_FUNCTION, + sqlite3.SQLITE_RECURSIVE, +} +_DENIED_FUNCTIONS = {"load_extension", "readfile", "writefile", "edit", "fts3_tokenizer"} +_TAG_RE = re.compile(r"<(action|answer)>(.*?)", re.IGNORECASE | re.DOTALL) +_FENCE_RE = re.compile(r"^```[ \t]*(?:sql|sqlite)?[ \t]*\n?(.*?)```$", re.IGNORECASE | re.DOTALL) +_THINK_RE = re.compile(r".*?", re.DOTALL) + + +def _authorizer(action, arg1, arg2, db_name, trigger): + if action not in _ALLOWED_ACTIONS: + return sqlite3.SQLITE_DENY + if action == sqlite3.SQLITE_FUNCTION and arg2 and arg2.lower() in _DENIED_FUNCTIONS: + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + +def db_file(db_root: str, db_id: str) -> Path: + """``//.sqlite``, the BIRD and Spider database layout.""" + return Path(db_root) / db_id / f"{db_id}.sqlite" + + +def execute_sql( + path: Path, sql: str, timeout: float = 30.0, max_rows: int = 100_000 +) -> Tuple[Optional[List[Tuple]], Optional[List[str]], Optional[str]]: + """Run one read-only statement. Returns ``(rows, column_names, None)`` or ``(None, None, error)``.""" + sql = (sql or "").strip().rstrip(";").strip() + if not sql: + return None, None, "empty query" + if not Path(path).is_file(): + return None, None, f"database not found: {path}" + deadline = time.monotonic() + timeout + conn = None + try: + conn = sqlite3.connect( + f"{Path(path).resolve().as_uri()}?mode=ro", uri=True, check_same_thread=False + ) + conn.text_factory = lambda b: b.decode("utf-8", errors="replace") + conn.set_authorizer(_authorizer) + conn.set_progress_handler(lambda: 1 if time.monotonic() > deadline else 0, 1000) + cursor = conn.execute(sql) + rows = cursor.fetchmany(max_rows + 1) + if len(rows) > max_rows: + return None, None, f"result exceeds {max_rows} rows" + columns = [d[0] for d in cursor.description or []] + return [tuple(r) for r in rows], columns, None + except sqlite3.OperationalError as e: + return None, None, "timeout" if "interrupted" in str(e) else f"OperationalError: {e}" + except Exception as e: # sqlite3.Warning (several statements), DatabaseError, ... + return None, None, f"{type(e).__name__}: {e}" + finally: + if conn is not None: + conn.close() + + +def schema_text(path: Path) -> str: + """The CREATE TABLE statements of every user table.""" + rows, _, error = execute_sql(path, "SELECT name, sql FROM sqlite_master WHERE type = 'table'") + if error: + raise ValueError(f"cannot read schema of {path}: {error}") + return "\n\n".join(sql for name, sql in rows if sql and name != "sqlite_sequence") + + +def parse_response(text: str) -> Tuple[Optional[str], str]: + """Return ``(kind, sql)`` for the last ```` or ```` tag outside ````.""" + matches = _TAG_RE.findall(_THINK_RE.sub("", text or "")) + if not matches: + return None, "" + kind, body = matches[-1] + body = body.strip() + fenced = _FENCE_RE.match(body) + if fenced: + body = fenced.group(1) + return kind.lower(), body.strip().rstrip(";").strip() + + +def format_result(rows: List[Tuple], columns: List[str], max_rows: int) -> str: + shown = rows[:max_rows] + lines = [" | ".join(columns)] + [" | ".join(str(v) for v in row) for row in shown] + more = f"\n... ({len(rows)} rows in total)" if len(rows) > max_rows else f"\n({len(rows)} rows)" + return "\n".join(lines) + more + + +@dataclass +class StepResult: + observation: str + done: bool + reward: float + error: Optional[str] = None + + +class SQLiteSQLEnv: + """One episode on one database. ``step`` takes a model response and returns a StepResult.""" + + def __init__( + self, + db_path: Path, + gold_sql: str, + timeout: float = 30.0, + max_obs_rows: int = 10, + max_rows: int = 100_000, + ): + self.db_path = Path(db_path) + self.timeout = timeout + self.max_obs_rows = max_obs_rows + self.max_rows = max_rows + self.gold_rows, _, self.gold_error = execute_sql(self.db_path, gold_sql, timeout, max_rows) + self.num_queries = 0 + self.num_errors = 0 + self.submitted = False + + @property + def gold_empty(self) -> bool: + return self.gold_error is None and not self.gold_rows + + def score(self, sql: str) -> Tuple[float, Optional[str]]: + rows, _, error = execute_sql(self.db_path, sql, self.timeout, self.max_rows) + if error or self.gold_error: + return 0.0, error + return float(set(rows) == set(self.gold_rows)), None + + def step(self, response_text: str) -> StepResult: + kind, sql = parse_response(response_text) + if kind is None: + return StepResult( + "Invalid format. Reply with one SQL to run a query, " + "or SQL to submit your final query.", + False, + 0.0, + "no action", + ) + if kind == "answer": + self.submitted = True + reward, error = self.score(sql) + return StepResult("Submitted.", True, reward, error) + self.num_queries += 1 + rows, columns, error = execute_sql(self.db_path, sql, self.timeout, self.max_rows) + if error: + self.num_errors += 1 + return StepResult(f"Error: {error}", False, 0.0, error) + return StepResult(format_result(rows, columns, self.max_obs_rows), False, 0.0) diff --git a/trinity/common/workflows/envs/sqlite_sql/sqlite_sql_workflow.py b/trinity/common/workflows/envs/sqlite_sql/sqlite_sql_workflow.py new file mode 100644 index 0000000000..0154b9f43a --- /dev/null +++ b/trinity/common/workflows/envs/sqlite_sql/sqlite_sql_workflow.py @@ -0,0 +1,113 @@ +# -*- coding: utf-8 -*- +"""Multi-turn text-to-SQL workflow on local SQLite databases (BIRD / Spider), no Docker needed.""" + +from dataclasses import asdict +from typing import List, Optional + +from trinity.common.experience import Experience +from trinity.common.models.model import ModelWrapper +from trinity.common.workflows.envs.sqlite_sql.sqlite_sql_env import ( + SQLiteSQLEnv, + db_file, + parse_response, + schema_text, +) +from trinity.common.workflows.workflow import MultiTurnWorkflow, Task + +SQLITE_SQL_SYSTEM_PROMPT = """ +You are a data analyst answering questions about a SQLite database. + +At each turn, think briefly and then reply with exactly one of: +- SQL to run a read-only SQLite query and see up to {max_obs_rows} rows of its result +- SQL to submit the final query that answers the question + +You have at most {max_env_steps} turns. Only the submitted query is graded: it is correct when it +returns the same rows as the reference answer. +""" + +USER_PROMPT = """Database schema: +{schema} + +Question: {question} +{evidence}""" + + +class SQLiteSQLWorkflow(MultiTurnWorkflow): + """Agent explores a SQLite database with read-only queries, then submits one query. + + Task fields (BIRD format): ``db_id``, ``question``, ``SQL`` (reference query), optional + ``evidence``. ``workflow_args``: ``db_root`` (required, ``//.sqlite``), + ``max_env_steps`` (default 5), ``max_obs_rows`` (default 10), ``timeout`` (seconds, default 30), + ``auto_submit`` (default False: grade the last successful query if the agent never submits). + """ + + is_async: bool = True + + def __init__( + self, + model: ModelWrapper, + task: Task, + auxiliary_models: Optional[List] = None, + ): + super().__init__(model=model, task=task, auxiliary_models=auxiliary_models) + args = task.workflow_args + self.raw_task = task.raw_task or {} + if not args.get("db_root"): + raise ValueError("SQLiteSQLWorkflow requires `workflow_args.db_root`.") + self.db_root = args["db_root"] + self.max_env_steps = int(args.get("max_env_steps", 5)) + self.max_obs_rows = int(args.get("max_obs_rows", 10)) + self.timeout = float(args.get("timeout", 30)) + self.auto_submit = bool(args.get("auto_submit", False)) + + def build_messages(self, db_path) -> List[dict]: + evidence = self.raw_task.get("evidence") + system = SQLITE_SQL_SYSTEM_PROMPT.format( + max_obs_rows=self.max_obs_rows, max_env_steps=self.max_env_steps + ).strip() + user = USER_PROMPT.format( + schema=schema_text(db_path), + question=self.raw_task["question"], + evidence=f"External knowledge: {evidence}" if evidence else "", + ).strip() + return [{"role": "system", "content": system}, {"role": "user", "content": user}] + + async def run_async(self) -> List[Experience]: + db_path = db_file(self.db_root, self.raw_task["db_id"]) + env = SQLiteSQLEnv( + db_path, self.raw_task["SQL"], timeout=self.timeout, max_obs_rows=self.max_obs_rows + ) + if env.gold_error: + self.logger.warning(f"Reference SQL fails on {db_path}: {env.gold_error}") + memory = self.build_messages(db_path) + reward, rounds, last_ok_sql = 0.0, 0, None + # Honour the task's sampling settings (e.g. eval temperature); one response per turn. + rollout_args = asdict(self.task.rollout_args) + rollout_args["n"] = 1 + + for rounds in range(1, self.max_env_steps + 1): + responses = await self.model.chat_async(memory, **rollout_args) + text = responses[0].response_text + memory.append({"role": "assistant", "content": text}) + result = env.step(text) + if result.done: + reward = result.reward + break + if result.error is None: + last_ok_sql = parse_response(text)[1] + memory.append({"role": "user", "content": result.observation}) + else: + if self.auto_submit and last_ok_sql: + reward, _ = env.score(last_ok_sql) + + info = { + "env_rounds": rounds, + "submitted": int(env.submitted), + "num_queries": env.num_queries, + "num_query_errors": env.num_errors, + "gold_error": int(env.gold_error is not None), + "gold_empty": int(env.gold_empty), + "final_reward": reward, + } + experience = await self.process_messages_to_experience_async(memory, reward, info) + return [experience]