A working Python reproduction of LinkedIn's enterprise Text-to-SQL system (Enterprise Text-to-SQL, Chen et al., KDD 2025 AAE Workshop, arXiv:2507.14372) — knowledge graph, ICA table clustering, Query Writer / Data Finder / Query Fixer / Researcher / Q&A agents — running on the BIRD dev set (SQLite) with the Claude Agent SDK as the multi-agent framework and a Gradio chat UI.
📖 Paper digest (my write-ups of the original paper): English · 繁體中文
┌────────────────────────── Gradio Chat UI ──────────────────────────┐
│ Intent Classifier ──► Query Writer │ Data Finder │ Fixer │ Q&A │
└──────────────────────────────┬─────────────────────────────────────┘
Query Writer pipeline: │ agentic agents get MCP tools:
① Retrieve (EBR, K_ret=20) │ search_tables / get_table_schema /
② Rank (LLM, K_rnk=7) │ get_table_metadata / validate_sql /
③ Write (JSON contract) │ run_sql / search_examples / knowledge
④ Validate & Fix (≤2) ◄── Researcher Agent (hallucination resolution)
┌──────────────────────────────┴─────────────────────────────────────┐
│ Knowledge Graph: table/column · usage · ICA clusters · examples · │
│ domain knowledge (instant refresh) · jargon map │
└────────────────────────────────────────────────────────────────────┘
Requires the Claude Code CLI installed and authenticated (the Claude Agent SDK drives it), plus uv.
# 1. Install dependencies
uv venv
uv pip install claude-agent-sdk gradio scikit-learn numpy sqlglot
# 2. Get the BIRD dev set (~330 MB zip)
curl -L -o data/dev.zip https://bird-bench.oss-cn-beijing.aliyuncs.com/dev.zip
cd data; unzip dev.zip; unzip dev_20240627/dev_databases.zip -d dev_20240627/; cd ..
# 3. Build the Knowledge Graph (schemas + usage mining + ICA clustering, ~1 min)
uv run python scripts/build_kg.py
# 4a. CLI smoke test — one BIRD question end-to-end
uv run python scripts/demo.py 0
# 4b. Chat UI
uv run python -m sqlbot.app # http://127.0.0.1:7860| Paper component | Paper source | This repo |
|---|---|---|
| Table/Column index | DataHub + Glean AI descriptions | SQLite PRAGMA + BIRD database_description CSVs; template AI descriptions (sqlbot/kg/builder.py) |
| Usage index | Trino EXPLAIN logs | table/column popularity + common joins mined from dev.json gold SQL via sqlglot |
| Table cluster index | FastICA on 3-month user-table access matrix (N=200, T_c=20) | FastICA on question-table matrix (N=16, T_c=8) — sqlbot/kg/clustering.py, soft clusters, Algorithm 4 candidate assembly |
| Example query index | wiki/code-repo queries | dev.json (question, SQL) pairs |
| Domain knowledge + jargon | crowdsourced, instant refresh | BIRD evidence strings; "X refers to Y" jargon mining; user submissions via UI refresh instantly |
| EBR (E5-large-v2 / ada-002) | embedding search | TF-IDF cosine (offline, swappable) — sqlbot/retrieval.py |
| Query Writer (4 steps, K_ret=20, K_rnk=7) | langchain + gpt-4o | Claude Agent SDK — sqlbot/agents/query_writer.py (ranker sees no schemas, per the paper's finding) |
Trino EXPLAIN VALIDATE |
syntax check | SQLite EXPLAIN QUERY PLAN — sqlbot/validation.py |
| Hallucination validator | custom, all-at-once | sqlglot identifier check vs. KG — validation.py |
| Researcher Agent | self-reflecting, gpt-4o-mini search | agentic SDK session + MCP tools, Haiku tier — agents/researcher.py |
| Multi-Agent UI + intent router | 4 agents, quick replies, knowledge contribution | Gradio Blocks — sqlbot/app.py |
| Validate & fix loop ≤ 2 | §3.4 | MAX_FIX_ROUNDS = 2 in config.py |
The default build (build_kg.py) mines the usage/example/cluster indexes from
dev.json — the same questions you'd evaluate on — so dev.json numbers from
that build are inflated by data leakage and aren't a fair benchmark.
scripts/build_kg_synth.py builds an alternative KG whose every index is
derived from sources independent of the eval questions, so it can be frozen
and scored against all of BIRD dev in a single pass — no LOO/K-fold masking
needed:
| KG index | Leaky build (kg.json) |
Leakage-free build (kg_synth.json) |
|---|---|---|
| table/column | schema | schema (same) |
| common joins | co-occurrence in dev gold SQL | FK graph from dev_tables.json |
| usage popularity | dev gold SQL counts | synthetic-query counts + FK centrality |
| ICA clusters | access matrix from dev gold SQL | access matrix from synthetic personas |
| example queries | dev.json (question, SQL) | deterministically generated, execute-filtered |
| domain knowledge / jargon | dev.json evidence |
column description CSVs (schema_doc) |
Synthetic queries come from templates over the schema + FK edges (COUNT,
GROUP BY on dimensions, aggregations on metrics, FK joins), each executed
against the real database and kept only if it runs. A runtime tripwire
(dev_json_tripwire) aborts the build if any code opens dev.json, and the
audit shows 0 question / SQL / evidence overlap with the eval set.
uv run python scripts/build_kg_synth.py # -> data/kg/kg_synth.json
$env:SQLBOT_KG="synth"; uv run python scripts/demo.py 0 # score against it
$env:SQLBOT_KG="synth"; uv run python -m sqlbot.app # UI on the fair KG
# Execution-accuracy (EX) over BIRD dev, leakage-free:
uv run python scripts/eval.py --kg synth --sample 100 --concurrency 6
uv run python scripts/eval.py --kg synth --difficulty simple --out ex.jsonscripts/eval.py runs the Query Writer over sampled dev questions and computes
BIRD execution accuracy (predicted vs. gold result-set equality), with a
per-difficulty breakdown, valid-SQL rate, and average fix rounds. Questions
whose gold SQL fails to execute are excluded from the denominator. --kg leaky
runs the same eval against the leaky KG as an inflated upper-reference.
SQLBOT_KG=synth switches load_kg() (and thus the demo, UI, and any eval
harness) to the leakage-free build. The generator never sees dev.json, so a
frozen kg_synth.json gives a fair execution-accuracy number over the full dev
set — the number then partly reflects synthetic-history coverage, so report it
alongside the schema-only cold-start config as a floor.
- BIRD has no employee query logs, so the ICA access matrix uses dev.json questions as pseudo user-sessions and their gold SQL as accessed tables. The gold SQL also feeds the example-query and usage indexes — fine for a system demo, but it means dev.json questions are not a fair benchmark for this build (the answer's neighbors are in the index).
- We can execute queries (BIRD is public data), so the UI shows a result preview — LinkedIn could not, which is why their self-consistency attempts failed (paper §Negative Results).
- Embeddings are TF-IDF instead of dense vectors; swap
TfidfRetrieverfor a sentence-transformers encoder for better recall. - Models are configurable via env vars (
SQLBOT_MODEL_WRITER,SQLBOT_MODEL_RESEARCHER, …); the researcher/intent/Q&A default to the Haiku tier, mirroring the paper's gpt-4o-mini speed split.
sqlbot/
config.py hyperparameters (K_ret=20, K_rnk=7, ICA dims, models)
kg/graph.py KG data model (5 indexes)
kg/builder.py BIRD -> KG (schemas, usage mining, jargon, clustering)
kg/clustering.py FastICA soft clustering (paper Algorithm 1)
retrieval.py EBR (TF-IDF) over tables / examples / knowledge
validation.py EXPLAIN + all-at-once hallucination validator + executor
llm.py Claude Agent SDK wrapper (ask / ask_json)
tools.py in-process SDK MCP server (7 KG tools)
agents/
router.py intent classifier
query_writer.py ① retrieve ② rank ③ write ④ validate & fix
researcher.py self-reflecting hallucination resolver
data_finder.py retrieve + rank only
query_fixer.py agentic debugging with validate/run tools
qa.py long-tail Q&A (difficulty-gated tool use)
app.py Gradio multi-agent chat UI
scripts/
build_kg.py build the knowledge graph (leaky — mines dev.json)
build_kg_synth.py build the leakage-free KG (synthetic query history)
demo.py CLI end-to-end smoke test
eval.py BIRD execution-accuracy (EX) over dev, per-difficulty