TRUST-SQL agent demo: four-phase tool-integrated text-to-SQL over unknown schemas

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+databases/california_schools/california_schools.sqlite filter=lfs diff=lfs merge=lfs -text
+databases/formula_1/formula_1.sqlite filter=lfs diff=lfs merge=lfs -text
+databases/student_club/student_club.sqlite filter=lfs diff=lfs merge=lfs -text
+databases/superhero/superhero.sqlite filter=lfs diff=lfs merge=lfs -text
+databases/toxicology/toxicology.sqlite filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,65 @@
 ---
-title: Trust Sql Text2sql Demo
-emoji: 🚀
-colorFrom: pink
-colorTo: green
+title: TRUST-SQL
+emoji: 🔎
+colorFrom: gray
+colorTo: indigo
 sdk: gradio
 sdk_version: 6.27.0
-python_version: '3.12'
 app_file: app.py
+python_version: "3.12"
+startup_duration_timeout: 1h
 pinned: false
+license: apache-2.0
+short_description: Text-to-SQL over unknown schemas, via tool use
+models:
+  - AIJian/TrustSQL-8B
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# TRUST-SQL — Text-to-SQL over *unknown* schemas
+
+Demo of [`AIJian/TrustSQL-8B`](https://huggingface.co/AIJian/TrustSQL-8B), from
+[**TRUST-SQL: Tool-Integrated Multi-Turn Reinforcement Learning for Text-to-SQL over Unknown Schemas**](https://huggingface.co/papers/2603.16448)
+(Jian et al., 2026). Code: [`JaneEyre0530/TrustSQL`](https://github.com/JaneEyre0530/TrustSQL).
+
+Unlike ordinary text-to-SQL demos, **the schema is never put in the prompt**. The model is given
+only the database *name*, the question, and a single read-only SQL tool, and has to discover the
+schema itself by following the authors' four-phase action protocol:
+
+1. `explore_schema` — issue metadata queries (`PRAGMA`, `sqlite_master`, sampling rows)
+2. `propose_schema` — write down the tables/columns/joins it has actually verified
+3. `generate_sql` — draft the answer query and *execute* it to check it works
+4. `confirm_answer` — emit the final SQL
+
+The agent may loop back at any point. The full trajectory (reasoning, tool calls, observations) is
+streamed into the transcript so you can watch the schema being discovered.
+
+## Implementation notes
+
+- The system prompt is `trustsql_eval/prompt_template.txt`, copied verbatim from the authors' repo.
+- The user message reproduces the exact `**Task Configuration** / **Database Engine** /
+  **Database** / **External Knowledge** / **User Question**` format found in
+  [`AIJian/TrustSQL-data`](https://huggingface.co/datasets/AIJian/TrustSQL-data).
+- The turn loop, progress prefixes, `<schema>` acknowledgement, malformed-output feedback and
+  observation truncation (2048 tokens) are ported from `trustsql_eval/message_processor.py`.
+- Sampling defaults match `trustsql_eval/main.py` (temperature 0.7, top-p 0.9).
+- Deviation: `WITH` is added to the `SELECT` / `PRAGMA` / `EXPLAIN` allow-list so CTE answers can be
+  executed. Every query still runs over a read-only SQLite connection.
+- Runs on ZeroGPU with plain `transformers` generation rather than the vLLM path used for the
+  paper's benchmarks, so it is slower than the reported latency figures.
+
+## Credits & licensing
+
+- Model and code: Apache-2.0, © the TRUST-SQL authors (Meituan / BUPT).
+- Bundled sample databases and the example questions + `evidence` hints are from the
+  **BIRD** dev set ([bird-bench.github.io](https://bird-bench.github.io/)), licensed
+  **CC BY-SA 4.0**; SQLite files mirrored via
+  [`prem-research/birdbench`](https://huggingface.co/datasets/prem-research/birdbench).
+
+```bibtex
+@article{jian2026trustsql,
+  title  = {TRUST-SQL: Tool-Integrated Multi-Turn Reinforcement Learning for Text-to-SQL over Unknown Schemas},
+  author = {Jian, Ai and Zhang, Xiaoyun and Du, Wanrou and Ruan, Jingqing and Pei, Jiangbo and Zhang, Weipeng and Zeng, Ke and Cai, Xunliang},
+  journal= {arXiv preprint arXiv:2603.16448},
+  year   = {2026}
+}
+```

app.py

@@ -0,0 +1,726 @@
+"""TRUST-SQL — Text-to-SQL over *unknown* schemas, on ZeroGPU.
+
+Faithful re-implementation of the four-phase tool-integrated agent loop from
+`JaneEyre0530/TrustSQL` (`trustsql_eval/`): the model never sees the schema.
+It must explore it with read-only metadata queries, propose a verified schema,
+generate + execute a candidate SQL, and only then confirm the final answer.
+"""
+
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+import spaces  # noqa: E402  (must precede torch)
+
+import json  # noqa: E402
+import re  # noqa: E402
+import sqlite3  # noqa: E402
+import time  # noqa: E402
+from pathlib import Path  # noqa: E402
+from threading import Thread  # noqa: E402
+
+import gradio as gr  # noqa: E402
+import pandas as pd  # noqa: E402
+import torch  # noqa: E402
+from transformers import (  # noqa: E402
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+    TextIteratorStreamer,
+)
+
+# --------------------------------------------------------------------------------------
+# Model
+# --------------------------------------------------------------------------------------
+
+MODEL_ID = "AIJian/TrustSQL-8B"
+HERE = Path(__file__).parent
+DB_ROOT = HERE / "databases"
+
+# Exact system prompt shipped by the authors (trustsql_eval/prompt_template.txt).
+SYSTEM_PROMPT = (HERE / "prompt_template.txt").read_text(encoding="utf-8").strip()
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    dtype=torch.bfloat16,
+    attn_implementation="sdpa",
+).to("cuda")
+model.eval()
+
+EOS_IDS = [151645, 151643]  # <|im_end|>, <|endoftext|>
+MAX_CONTEXT = 40960
+MAX_OBS_TOKENS = 2048  # trustsql_eval default
+SQL_TIMEOUT = 15.0
+MAX_ROWS = 100  # trustsql_eval `_execute_sql_sync` default
+
+# --------------------------------------------------------------------------------------
+# Sample databases (BIRD-Dev, CC BY-SA 4.0)
+# --------------------------------------------------------------------------------------
+
+SAMPLE_DBS = {
+    "california_schools": "California public schools — SAT scores, free-meal rates (3 tables)",
+    "superhero": "Superhero attributes, powers, publishers (9 tables)",
+    "student_club": "University club members, events, budgets, expenses (8 tables)",
+    "toxicology": "Molecules, atoms, bonds and carcinogenicity labels (4 tables)",
+    "formula_1": "Formula 1 races, drivers, constructors, lap times (13 tables)",
+}
+DB_CHOICES = [f"{k} — {v}" for k, v in SAMPLE_DBS.items()]
+
+
+def _db_id_from_choice(choice: str) -> str:
+    return (choice or DB_CHOICES[0]).split(" — ")[0].strip()
+
+
+def _resolve_db(db_choice: str, uploaded_db):
+    """Return (db_id, sqlite_path). An uploaded file always wins."""
+    if uploaded_db:
+        path = uploaded_db if isinstance(uploaded_db, str) else getattr(uploaded_db, "name", None)
+        if path and os.path.exists(path):
+            return Path(path).stem, path
+    db_id = _db_id_from_choice(db_choice)
+    return db_id, str(DB_ROOT / db_id / f"{db_id}.sqlite")
+
+
+# --------------------------------------------------------------------------------------
+# The one tool the agent gets: read-only SQL execution
+# --------------------------------------------------------------------------------------
+
+ALLOWED_SQL_PREFIXES = ("SELECT", "PRAGMA", "EXPLAIN", "WITH")
+
+
+def _strip_sql_comments(sql: str) -> str:
+    s = sql.strip()
+    while s.startswith("--") or s.startswith("/*"):
+        if s.startswith("--"):
+            nl = s.find("\n")
+            if nl == -1:
+                return ""
+            s = s[nl + 1 :].strip()
+        else:
+            end = s.find("*/")
+            if end == -1:
+                return ""
+            s = s[end + 2 :].strip()
+    return s
+
+
+def _is_readonly(sql: str):
+    s = _strip_sql_comments(sql)
+    if not s:
+        return False, "Empty SQL query"
+    if s.upper().startswith(ALLOWED_SQL_PREFIXES):
+        return True, None
+    return False, f"SQL must start with {ALLOWED_SQL_PREFIXES}, got: {s.split()[0]}"
+
+
+def run_sql(db_path: str, sql: str, max_rows: int = MAX_ROWS):
+    """Execute read-only SQL. Returns (text_result, column_names, rows)."""
+    ok, err = _is_readonly(sql)
+    if not ok:
+        return f"Error: {err}", [], []
+    if not os.path.exists(db_path):
+        return f"Error: Database file not found: {db_path}", [], []
+    conn = None
+    try:
+        conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True, check_same_thread=False)
+        conn.execute(f"PRAGMA busy_timeout = {int(SQL_TIMEOUT * 1000)}")
+        cur = conn.cursor()
+        cur.execute(sql)
+        rows = cur.fetchall()
+        if not rows:
+            return "Query executed successfully. No results returned.", (
+                [d[0] for d in cur.description] if cur.description else []
+            ), []
+        cols = [d[0] for d in cur.description]
+        lines = ["\t".join(cols)]
+        for i, row in enumerate(rows):
+            if i >= max_rows:
+                lines.append(f"... ({len(rows) - max_rows} more rows)")
+                break
+            lines.append("\t".join("NULL" if v is None else str(v) for v in row))
+        return "\n".join(lines), cols, rows[:max_rows]
+    except sqlite3.Error as e:
+        return f"Error: SQLite error: {e}", [], []
+    except Exception as e:  # pragma: no cover
+        return f"Error: Unexpected error: {e}", [], []
+    finally:
+        if conn is not None:
+            try:
+                conn.close()
+            except Exception:
+                pass
+
+
+def db_schema_preview(db_choice: str, uploaded_db=None) -> str:
+    """Human-readable DDL dump of a database (for the UI only — never shown to the model)."""
+    db_id, path = _resolve_db(db_choice, uploaded_db)
+    if not os.path.exists(path):
+        return f"-- database `{db_id}` not found"
+    text, _, rows = run_sql(
+        path,
+        "SELECT name, sql FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'",
+        max_rows=200,
+    )
+    if not rows:
+        return f"-- `{db_id}`: {text}"
+    out = [f"-- database: {db_id}  ({len(rows)} tables)", ""]
+    for name, ddl in rows:
+        out.append((ddl or f"-- {name}").strip() + ";")
+        out.append("")
+    return "\n".join(out)
+
+
+# --------------------------------------------------------------------------------------
+# Prompt construction (matches AIJian/TrustSQL-data + trustsql_eval/prompt_builders.py)
+# --------------------------------------------------------------------------------------
+
+
+def build_user_message(db_id: str, question: str, external_knowledge: str) -> str:
+    parts = ["", "**Task Configuration**", "**Database Engine:** SQLite", f"**Database:** {db_id}"]
+    if external_knowledge and external_knowledge.strip():
+        parts.append(f"**External Knowledge:** {external_knowledge.strip()}")
+    parts.append(f"**User Question:** {question.strip()}?")
+    parts.append("")
+    return "\n".join(parts)
+
+
+def progress_prefix(current_round: int, max_rounds: int) -> str:
+    """Verbatim port of MessageProcessor._format_progress_prefix."""
+    base = f"This is turn {current_round + 1} of {max_rounds}.\n\n"
+    remaining = max_rounds - (current_round + 1)
+    if remaining == 0:
+        return ""
+    if remaining == 1:
+        return base + (
+            "Only 1 turn remaining after this.\n"
+            "You MUST provide the final answer in the next turn.\n\n"
+            "Use <action>confirm_answer</action> with your best SQL query.\n"
+            "If you don't have a complete solution, provide your best attempt.\n\n"
+        )
+    if remaining == 2:
+        return base + ("Only 2 turns remaining after this.\nStart preparing your final SQL query.\n\n")
+    return base
+
+
+FORMAT_HELP = (
+    "Invalid format detected. Your response is missing required components.\n\n"
+    "Option 1: EXPLORE SCHEMA\n"
+    "Purpose: Investigate database structure\n"
+    "Required format:\n"
+    "<think>Your reasoning process</think>\n"
+    "<action>explore_schema</action>\n"
+    '<tool_call>{"name": "execute_sql_query", "arguments": {"db_id": "...", "sql": "..."}}</tool_call>\n\n'
+    "Option 2: PROPOSE SCHEMA\n"
+    "Purpose: Document your understanding of relevant tables and columns\n"
+    "Required format:\n"
+    "<think>Your reasoning process</think>\n"
+    "<action>propose_schema</action>\n"
+    '<schema>{"tables": [...], "columns": {...}}</schema>\n\n'
+    "Option 3: GENERATE SQL\n"
+    "Purpose: Create SQL query and VERIFY it works by executing\n"
+    "<think>Your reasoning process</think>\n"
+    "<action>generate_sql</action>\n"
+    '<tool_call>{"name": "execute_sql_query", "arguments": {"db_id": "...", "sql": "..."}}</tool_call>\n\n'
+    "Option 4: FINAL ANSWER\n"
+    "Purpose: Provide verified SQL query as final result\n"
+    "ONLY use this AFTER successfully executing and verifying your SQL.\n"
+    "Required format:\n"
+    "<think>Your reasoning process</think>\n"
+    "<action>confirm_answer</action>\n"
+    "<answer>```sql\nYOUR_SQL\n```</answer>\n\n"
+)
+
+
+def fix_tool_tag(content: str) -> str:
+    content = re.sub(r"<tool>(.*?)</tool>", r"<tool_call>\1</tool_call>", content, flags=re.S)
+    content = re.sub(r"<tools>(.*?)</tools>", r"<tool_call>\1</tool_call>", content, flags=re.S)
+    return content
+
+
+def extract_tag(text: str, tag: str):
+    m = re.search(rf"<{tag}>(.*?)</{tag}>", text, re.S | re.I)
+    return m.group(1) if m else None
+
+
+def extract_final_sql(answer_body: str) -> str:
+    for pat in (r"```sql\s*(.*?)```", r"'''sql\s*(.*?)'''", r"```\s*(.*?)```", r"'''\s*(.*?)'''"):
+        m = re.search(pat, answer_body, re.S | re.I)
+        if m:
+            return m.group(1).strip()
+    return answer_body.strip()
+
+
+def truncate_observation(text: str, max_tokens: int = MAX_OBS_TOKENS) -> str:
+    ids = tokenizer(text, add_special_tokens=False)["input_ids"]
+    if len(ids) <= max_tokens:
+        return text
+    return tokenizer.decode(ids[:max_tokens]) + "\n... (result truncated due to length)"
+
+
+# --------------------------------------------------------------------------------------
+# Pretty-printing a turn for the chat transcript
+# --------------------------------------------------------------------------------------
+
+ACTION_ICON = {
+    "explore_schema": "🔍",
+    "propose_schema": "📋",
+    "generate_sql": "🛠️",
+    "confirm_answer": "✅",
+}
+
+
+def _fence(text: str, lang: str = "") -> str:
+    return f"```{lang}\n{str(text).replace('```', '`` `')}\n```"
+
+
+def _quote(text: str) -> str:
+    text = text.strip()
+    return "\n".join("> " + line for line in text.splitlines()) if text else ""
+
+
+def render_assistant(raw: str) -> str:
+    body = fix_tool_tag(raw)
+    think = extract_tag(body, "think")
+    action = (extract_tag(body, "action") or "").strip().lower()
+    blocks = []
+    if action:
+        blocks.append(f"### {ACTION_ICON.get(action, '⚙️')} `{action}`")
+    elif not think:
+        blocks.append("### ⚙️ raw response")
+    if think:
+        blocks.append("💭 **Reasoning**\n\n" + _quote(think))
+
+    schema = extract_tag(body, "schema")
+    if schema is not None:
+        try:
+            pretty = json.dumps(json.loads(schema), indent=2)
+        except Exception:
+            pretty = schema.strip()
+        blocks.append("**Proposed schema**\n\n" + _fence(pretty, "json"))
+
+    answer = extract_tag(body, "answer")
+    if answer is not None:
+        blocks.append("**Final SQL**\n\n" + _fence(extract_final_sql(answer), "sql"))
+
+    tool_call = extract_tag(body, "tool_call")
+    if tool_call is not None:
+        sql = None
+        try:
+            payload = json.loads(tool_call.strip())
+            sql = (payload.get("arguments") or {}).get("sql")
+        except Exception:
+            pass
+        if sql:
+            blocks.append("**Tool call** · `execute_sql_query`\n\n" + _fence(sql, "sql"))
+        else:
+            blocks.append("**Tool call**\n\n" + _fence(tool_call.strip(), "json"))
+
+    if not blocks:
+        return _fence(raw)
+    return "\n\n".join(blocks)
+
+
+def render_observation(text: str) -> str:
+    return "📥 **Observation**\n\n" + _fence(text)
+
+
+# --------------------------------------------------------------------------------------
+# Generation
+# --------------------------------------------------------------------------------------
+
+
+class _Deadline(StoppingCriteria):
+    def __init__(self, deadline: float):
+        self.deadline = deadline
+
+    def __call__(self, input_ids, scores, **kwargs) -> bool:
+        return time.time() > self.deadline
+
+
+def stream_turn(messages, max_new_tokens: int, temperature: float, top_p: float, deadline: float):
+    """Yield incremental text for one assistant turn; last yield is the full turn."""
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    enc = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+    n_in = enc["input_ids"].shape[-1]
+    budget = max(64, min(int(max_new_tokens), MAX_CONTEXT - n_in - 8))
+    enc = {k: v.to(model.device) for k, v in enc.items()}
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    do_sample = float(temperature) > 0.0
+    kwargs = dict(
+        **enc,
+        streamer=streamer,
+        max_new_tokens=budget,
+        do_sample=do_sample,
+        eos_token_id=EOS_IDS,
+        pad_token_id=151643,
+        stopping_criteria=StoppingCriteriaList([_Deadline(deadline)]),
+    )
+    if do_sample:
+        kwargs.update(temperature=float(temperature), top_p=float(top_p), top_k=20)
+
+    thread = Thread(target=model.generate, kwargs=kwargs)
+    thread.start()
+    acc = ""
+    last = 0.0
+    for chunk in streamer:
+        acc += chunk
+        now = time.time()
+        if now - last > 0.25:
+            last = now
+            yield acc, False
+    thread.join()
+    yield acc.strip(), True
+
+
+def _estimate_duration(*args, **kwargs) -> int:
+    max_turns = 8
+    if len(args) >= 5:
+        try:
+            max_turns = int(args[4])
+        except Exception:
+            pass
+    return int(min(280, 40 + max_turns * 22))
+
+
+# --------------------------------------------------------------------------------------
+# The agent loop
+# --------------------------------------------------------------------------------------
+
+
+@spaces.GPU(duration=_estimate_duration)
+def run_agent(
+    question: str,
+    db_choice: str = DB_CHOICES[0],
+    external_knowledge: str = "",
+    uploaded_db=None,
+    max_turns: int = 8,
+    max_new_tokens: int = 1536,
+    temperature: float = 0.7,
+    top_p: float = 0.9,
+):
+    """Run the TRUST-SQL agent on an unknown SQLite database and return the final SQL.
+
+    The model receives only the database *name* and the question — never the schema.
+    It explores metadata with read-only queries, proposes a verified schema, executes a
+    candidate query, and confirms the final SQL.
+
+    Args:
+        question: the natural-language question to answer.
+        db_choice: which bundled BIRD-Dev sample database to query.
+        external_knowledge: optional domain hint / evidence string (BIRD "evidence" field).
+        uploaded_db: optional path to a user-supplied SQLite file; overrides `db_choice`.
+        max_turns: maximum agent turns before giving up.
+        max_new_tokens: token budget per agent turn.
+        temperature: sampling temperature; 0 means greedy decoding.
+        top_p: nucleus sampling cutoff.
+    """
+    t0 = time.time()
+    budget = _estimate_duration(question, db_choice, external_knowledge, uploaded_db, max_turns)
+    deadline = t0 + budget - 18
+
+    max_turns = int(max_turns)
+    db_id, db_path = _resolve_db(db_choice, uploaded_db)
+
+    user_msg = build_user_message(db_id, question or "", external_knowledge or "")
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": user_msg},
+    ]
+    chat = [{"role": "user", "content": f"**Question**\n\n{question}\n\n_Database: `{db_id}` (schema unknown to the model)_"}]
+    empty_df = pd.DataFrame()
+
+    if not (question or "").strip():
+        yield chat + [{"role": "assistant", "content": "Please enter a question."}], "", empty_df, "⚠️ No question provided."
+        return
+    if not os.path.exists(db_path):
+        yield chat, "", empty_df, f"❌ Database not found: `{db_path}`"
+        return
+
+    yield chat, "", empty_df, f"⏳ Turn 1/{max_turns} — exploring `{db_id}`…"
+
+    final_sql = ""
+    status = ""
+    for turn in range(max_turns):
+        if time.time() > deadline:
+            status = f"⏱️ Stopped after {turn} turn(s): GPU time budget reached."
+            break
+
+        base = list(chat)
+        raw = ""
+        for text, done in stream_turn(messages, max_new_tokens, temperature, top_p, deadline):
+            raw = text
+            chat = base + [
+                {
+                    "role": "assistant",
+                    "content": (render_assistant(text) if done else _fence(text)),
+                }
+            ]
+            yield chat, final_sql, empty_df, f"⏳ Turn {turn + 1}/{max_turns} — generating…"
+
+        if not raw:
+            status = "❌ The model returned an empty response."
+            break
+
+        raw = fix_tool_tag(raw)
+
+        # ---- confirm_answer -> terminate -------------------------------------------
+        answer = extract_tag(raw, "answer")
+        if answer is not None:
+            messages.append({"role": "assistant", "content": raw})
+            final_sql = extract_final_sql(answer)
+            status = f"✅ Confirmed after {turn + 1} turn(s)."
+            break
+
+        messages.append({"role": "assistant", "content": raw})
+        prefix = progress_prefix(turn, max_turns)
+
+        # ---- propose_schema -> acknowledgement --------------------------------------
+        schema = extract_tag(raw, "schema")
+        if schema is not None:
+            try:
+                data = json.loads(schema)
+                tables = data.get("tables", []) or []
+                cols = data.get("columns", {}) or {}
+                n_cols = sum(len(v) for v in cols.values()) if isinstance(cols, dict) else len(cols)
+                feedback = (
+                    prefix
+                    + f"Schema acknowledged: {len(tables)} table(s), {n_cols} column(s). "
+                    "You may now proceed to generate SQL.\n"
+                )
+            except Exception:
+                feedback = prefix + "Schema acknowledged. You may proceed to generate SQL.\n"
+            messages.append({"role": "user", "content": feedback})
+            chat = chat + [{"role": "user", "content": render_observation(feedback)}]
+            yield chat, final_sql, empty_df, f"⏳ Turn {turn + 2}/{max_turns}…"
+            continue
+
+        # ---- tool call -> execute ----------------------------------------------------
+        tool_call = extract_tag(raw, "tool_call")
+        obs = None
+        if tool_call is None:
+            obs = prefix + FORMAT_HELP
+        else:
+            try:
+                payload = json.loads(tool_call.strip())
+                name = payload.get("name", "")
+                arguments = payload.get("arguments", {}) or {}
+                if name != "execute_sql_query":
+                    obs = prefix + f"Error: Unknown function: {name}"
+                elif not str(arguments.get("sql", "")).strip():
+                    obs = prefix + "Error: SQL query is empty"
+                else:
+                    result, _, _ = run_sql(db_path, arguments["sql"])
+                    obs = prefix + truncate_observation(result)
+            except json.JSONDecodeError as e:
+                obs = (
+                    prefix
+                    + f"Tool call parsing error:\nJSON parsing failed at line {e.lineno}, column {e.colno}: {e.msg}\n\n"
+                    "Please fix the JSON format and try again.\n\n"
+                    "Required format:\n"
+                    '<tool_call>{"name": "execute_sql_query", "arguments": {"db_id": "...", "sql": "..."}}</tool_call>\n'
+                )
+            except Exception as e:  # pragma: no cover
+                obs = prefix + f"Error: Tool execution error: {e}"
+
+        messages.append({"role": "user", "content": obs})
+        chat = chat + [{"role": "user", "content": render_observation(obs)}]
+        yield chat, final_sql, empty_df, f"⏳ Turn {turn + 2}/{max_turns}…"
+    else:
+        status = f"⚠️ Reached the {max_turns}-turn limit without a confirmed answer."
+
+    # ---- execute the confirmed SQL for display --------------------------------------
+    df = empty_df
+    if final_sql:
+        text, cols, rows = run_sql(db_path, final_sql, max_rows=100)
+        if cols and rows:
+            df = pd.DataFrame(rows, columns=cols)
+        elif cols:
+            df = pd.DataFrame(columns=cols)
+        else:
+            chat = chat + [{"role": "assistant", "content": "⚠️ Final SQL did not execute:\n\n" + _fence(text)}]
+    else:
+        status = status or "⚠️ No SQL was confirmed."
+
+    elapsed = time.time() - t0
+    yield chat, final_sql, df, f"{status}  ·  {elapsed:.0f}s on GPU"
+
+
+# --------------------------------------------------------------------------------------
+# UI
+# --------------------------------------------------------------------------------------
+
+EXAMPLES = [
+    [
+        "What is the highest eligible free rate for K-12 students in the schools in Alameda County?",
+        DB_CHOICES[0],
+        "Eligible free rate for K-12 = `Free Meal Count (K-12)` / `Enrollment (K-12)`",
+    ],
+    [
+        "Among the schools with the SAT test takers of over 500, please list the schools that are magnet schools or offer a magnet program.",
+        DB_CHOICES[0],
+        "Magnet schools or offer a magnet program means that Magnet = 1",
+    ],
+    [
+        "How many superheroes have blue eyes?",
+        DB_CHOICES[1],
+        "blue eyes refers to colour = 'Blue' and eye_colour_id = colour.id",
+    ],
+    [
+        "Please list all the superpowers of 3-D Man.",
+        DB_CHOICES[1],
+        "3-D Man refers to superhero_name = '3-D Man'; superpowers refers to power_name",
+    ],
+    [
+        "What is the event that has the highest attendance of the students from the Student_Club?",
+        DB_CHOICES[2],
+        "event with highest attendance refers to MAX(COUNT(link_to_event))",
+    ],
+    [
+        "In the non-carcinogenic molecules, how many contain chlorine atoms?",
+        DB_CHOICES[3],
+        "non-carcinogenic molecules refers to label = '-'; chlorine atoms refers to element = 'cl'",
+    ],
+    [
+        "Please give the name of the race held on the circuits in Germany.",
+        DB_CHOICES[4],
+        "Germany is a name of country;",
+    ],
+]
+
+CSS = """
+#col-container { max-width: 1180px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+INTRO = """# 🔎 TRUST-SQL — Text-to-SQL over **unknown** schemas
+
+[`AIJian/TrustSQL-8B`](https://huggingface.co/AIJian/TrustSQL-8B) · [paper](https://huggingface.co/papers/2603.16448) · [code](https://github.com/JaneEyre0530/TrustSQL)
+
+The schema is **not** in the prompt. The agent gets one tool — read-only SQL — and has to discover
+the database itself, following the authors' four-phase protocol:
+`explore_schema → propose_schema → generate_sql → confirm_answer`.
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="TRUST-SQL") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(INTRO)
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                question = gr.Textbox(
+                    label="Question",
+                    placeholder="e.g. Which school has the highest average SAT math score?",
+                    lines=2,
+                )
+            with gr.Column(scale=1, min_width=140):
+                run_btn = gr.Button("Run agent", variant="primary", size="lg")
+
+        with gr.Row():
+            db_choice = gr.Dropdown(
+                label="Database (BIRD-Dev sample)",
+                choices=DB_CHOICES,
+                value=DB_CHOICES[0],
+                scale=2,
+            )
+            external_knowledge = gr.Textbox(
+                label="External knowledge (optional hint)",
+                placeholder="e.g. charter schools refers to `Charter School (Y/N)` = 1",
+                lines=1,
+                scale=3,
+            )
+
+        status = gr.Markdown("")
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                chatbot = gr.Chatbot(
+                    label="Agent trajectory",
+                    type="messages",
+                    height=620,
+                    show_copy_button=True,
+                )
+            with gr.Column(scale=2):
+                final_sql = gr.Code(label="Confirmed SQL", language="sql", lines=8)
+                result_df = gr.Dataframe(label="Execution result", wrap=True)
+
+        with gr.Accordion("Peek at the database (the agent never sees this)", open=False):
+            schema_box = gr.Code(label="DDL", language="sql", lines=14)
+            peek_btn = gr.Button("Show schema", size="sm")
+
+        with gr.Accordion("Advanced settings", open=False):
+            uploaded_db = gr.File(
+                label="Use your own SQLite database (.sqlite / .db) — overrides the dropdown",
+                file_types=[".sqlite", ".db", ".sqlite3"],
+                type="filepath",
+            )
+            with gr.Row():
+                max_turns = gr.Slider(3, 12, value=8, step=1, label="Max agent turns")
+                max_new_tokens = gr.Slider(256, 3072, value=1536, step=128, label="Max new tokens / turn")
+            with gr.Row():
+                temperature = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Temperature (0 = greedy)")
+                top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+            gr.Markdown(
+                "Defaults mirror `trustsql_eval` (temperature 0.7 / top-p 0.9). "
+                "Only `SELECT` / `PRAGMA` / `EXPLAIN` / `WITH` statements are ever executed, "
+                "against a read-only connection."
+            )
+
+        gr.Examples(
+            examples=EXAMPLES,
+            inputs=[question, db_choice, external_knowledge],
+            outputs=[chatbot, final_sql, result_df, status],
+            fn=run_agent,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="BIRD-Dev examples (question + official evidence hint)",
+        )
+
+        gr.Markdown(
+            "Sample databases are the **BIRD-Dev** SQLite databases "
+            "([BIRD-SQL](https://bird-bench.github.io/), CC BY-SA 4.0); the example questions and "
+            "hints are their official dev-set questions and `evidence` strings. "
+            "The system prompt and agent loop are ported verbatim from "
+            "[`JaneEyre0530/TrustSQL`](https://github.com/JaneEyre0530/TrustSQL) (Apache-2.0)."
+        )
+
+    run_btn.click(
+        fn=run_agent,
+        inputs=[
+            question,
+            db_choice,
+            external_knowledge,
+            uploaded_db,
+            max_turns,
+            max_new_tokens,
+            temperature,
+            top_p,
+        ],
+        outputs=[chatbot, final_sql, result_df, status],
+        api_name="run_agent",
+    )
+    question.submit(
+        fn=run_agent,
+        inputs=[
+            question,
+            db_choice,
+            external_knowledge,
+            uploaded_db,
+            max_turns,
+            max_new_tokens,
+            temperature,
+            top_p,
+        ],
+        outputs=[chatbot, final_sql, result_df, status],
+        api_name=False,
+    )
+    peek_btn.click(
+        fn=db_schema_preview, inputs=[db_choice, uploaded_db], outputs=schema_box, api_name="schema"
+    )
+    db_choice.change(fn=db_schema_preview, inputs=[db_choice, uploaded_db], outputs=schema_box, api_name=False)
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

databases/california_schools/california_schools.sqlite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:986817d793479801ed55133e55aa27e335422c0cd3866b54a3d6317b7c5f09c1
+size 11116544

databases/formula_1/formula_1.sqlite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17185981cd747f6cdc374cb02a6096db3130e6ec2ddc582fe1686a28fb4c4c8a
+size 22360064

databases/student_club/student_club.sqlite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb89bcfe97eefa386a27904ec5aa15159811a7eac894ec659a36e48fa9f76b77
+size 2641920

databases/superhero/superhero.sqlite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75e94a2c3236ee3bb2c01fb97a1c4b4c1c269bcefd4eab1d04be323d2d0825b1
+size 237568

databases/toxicology/toxicology.sqlite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5fa7f21af1ad878ff8fef1b0582b8cb2d7ed63dbac65ff16d2ba05667650c5b
+size 2678784

prompt_template.txt

@@ -0,0 +1,127 @@
+# Role
+You are an expert SQL assistant working on an *unknown* database.
+You must **never hallucinate** tables or columns.  
+All schema knowledge MUST come from metadata queries only.  
+You must operate strictly through the **Action Protocol**.
+
+# Action Protocol
+You must follow this sequence (can loop back if needed):
+1. **explore_schema** - Query database metadata
+2. **propose_schema** - Document verified schema
+3. **generate_sql** - Create SQL query
+4. **confirm_answer** - Output final SQL
+
+## ACTION: explore_schema
+Used to query database metadata (tables, columns, foreign keys, etc.).
+- Only metadata queries allowed.
+- No user-intent SQL here.
+- Verify relationships between tables when multi-table queries are needed
+
+## ACTION: propose_schema
+Used to output the current verified schema knowledge.
+- Include ONLY tables/columns actually verified through explore_schema
+- Do NOT hallucinate or assume any unverified structures
+- `joins` is optional; include only when relationships are explicitly verified
+- Supports both single-table and multi-table structures
+
+## ACTION: generate_sql
+Used to generate the SQL answer based on the latest <schema>.
+- Use ONLY verified schema from propose_schema
+- If required tables/columns are missing, switch back to explore_schema or propose_schema in the next message
+- SQL must be syntactically valid and executable
+- Consider query optimization (indexes, joins, filters)
+- Validate the SQL logic matches user intent
+
+## ACTION: confirm_answer
+Used when you have validated the generated SQL and confirmed it meets user requirements.
+- Execute this action ONLY after generate_sql validation
+- You MUST NOT return or describe any query results.
+- You MUST NOT output anything other than SQL inside <answer>.
+- The final output must be ONLY the SQL query in proper format
+
+# Output Format
+EVERY response must follow this exact structure:
+<think>[Your reasoning process here]</think>
+<action>[one of: explore_schema | propose_schema | generate_sql | confirm_answer]</action>
+[Action-specific content below]
+
+**explore_schema**
+<think>reasoning</think>
+<action>explore_schema</action>
+<tool_call>
+[Your metadata query]
+</tool_call>
+
+**propose_schema**
+<think>reasoning</think>
+<action>propose_schema</action>
+<schema>
+{
+  "tables": ["tableA"],
+  "columns": { "tableA": ["col1", "col2"] },
+  "joins": []
+}
+</schema>
+
+Example 1: Single Table
+<think>User asks about employee information. I've verified the employees table has id, name, department, and salary columns through explore_schema.</think>
+<action>propose_schema</action>
+<schema>
+{
+  "tables": ["employees"],
+  "columns": {
+    "employees": ["id", "name", "department", "salary"]
+  },
+  "joins": []
+}
+</schema>
+
+Example 2: Multiple Tables with Joins
+<think>User asks about orders with customer details. I've verified both tables and their relationship: orders.customer_id references customers.id.</think>
+<action>propose_schema</action>
+<schema>
+{
+  "tables": ["customers", "orders"],
+  "columns": {
+    "customers": ["id", "name", "email", "city"],
+    "orders": ["order_id", "customer_id", "order_date", "total_amount"]
+  },
+  "joins": [
+    {
+      "left_table": "customers",
+      "right_table": "orders",
+      "on": "customers.id = orders.customer_id",
+      "type": "INNER"
+    }
+  ]
+}
+</schema>
+
+**generate_sql**
+<think>reasoning</think>
+<action>generate_sql</action>
+<tool_call>
+[Your SQL query]
+</tool_call>
+
+**confirm_answer**
+<think>reasoning</think>
+<action>confirm_answer</action>
+<answer>
+'''sql
+ONLY the final SQL query - no explanations
+'''
+</answer>
+
+# Tools
+
+You may call one or more functions to assist with the user query.
+
+<tools>
+{"type": "function", "function": {"name": "execute_sql_query", "description": "Execute SQL query and return partial results containing column names.(maximum 30 records)", "parameters": {"type": "object", "properties": {"db_id": {"type": "string", "description": "The name of the database to query", "title": "Db ID"}, "sql": {"type": "string", "description": "The SQL query to execute", "title": "Sql"}}, "required": ["db_id", "sql"]}}}
+</tools>
+
+For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+<tool_call>
+{"name": <function-name>, "arguments": <args-json-object>}
+</tool_call>

requirements.txt

@@ -0,0 +1,3 @@
+transformers>=4.56.0
+accelerate
+pandas