Text-to-SQL with Schema Linking

Ask a database a question in plain English and get back SQL that actually runs. The distinctive part isn't the LLM call — it's schema linking: the relevant tables are retrieved by embedding similarity before prompting, so the model sees a small, focused schema instead of the whole database. And quality is measured by execution accuracy (does the generated SQL return the correct rows?), not by string-matching SQL.

This is the fourth Python project in a series of embedding + vector-database systems, and like the others it pairs a clean, testable core with a measurable benchmark.

What this project demonstrates

• Schema linking as retrieval — each table is embedded as a short document (name + columns + relationships) and indexed in a vector store. A question retrieves only the top-k relevant tables, which keeps the prompt small and on-topic on wide schemas where dumping every table would blow the context window and invite mistakes.
• Foreign-key expansion — retrieving tables by similarity alone misses the bridge tables a join needs (e.g. an order_items table between orders and products that the question never names). Retrieved tables pull in their FK neighbours so join paths stay intact.
• Execution accuracy, done right — two different SQL strings can be equally correct, so the benchmark runs both the predicted and gold queries and compares their result sets, with order sensitivity inferred from whether the gold query has an ORDER BY.
• Metrics that localise failures — alongside execution accuracy, schema-linking recall isolates retrieval errors from generation errors (high recall but low accuracy ⇒ the LLM is at fault, not the linker), and valid-SQL rate separates syntax failures from wrong answers.
• A built-in ablation — every question is run twice, with the full schema and with the linked subset, so the report directly shows what linking buys.
• Runs fully offline — an in-memory vector store + a deterministic hashing embedder mean the entire linking path (and the whole test suite) runs with no Qdrant, no model download, and no LLM. SQLite is the database, so there's no server to run at all for the demo.

Architecture

                    ┌────────────────────────────┐
                    │  SQLite database (.sqlite)  │
                    └──────────────┬──────────────┘
                       introspect  │ (PRAGMA)
                                  ▼
                    Schema (tables, columns, FKs)
                                  │  one document per table
                                  ▼
                    Embedder (SentenceTransformer | Hashing)
                                  │
                                  ▼
                    Vector store (Qdrant | in-memory)   ◄── schema index
                                  ▲
   question ─► embed ────────────┘  top-k tables
                                  │  + FK-neighbour expansion
                                  ▼
                         LinkedSchema (subset)
                                  │  pseudo-DDL in prompt
                                  ▼
                    LLM  ─►  SQL  ─►  extract_sql()
                                  │
                                  ▼
        read-only execute() on SQLite  ─►  rows
                                  │
                                  ▼
   Evaluation: predicted rows vs gold rows  ─►  execution accuracy
               + schema-linking recall + valid-SQL rate
               (full-schema vs linked-schema ablation)

Quickstart

Requirements: Python 3.10+. Qdrant and an LLM key are optional — the demo runs offline without them.

# 1. Install
pip install -r requirements.txt && pip install -e .

# 2. Unit tests (offline: no Qdrant, no model, no LLM)
pytest -q

# 3. Build the demo database + gold benchmark
python scripts/build_demo_db.py

Offline run (no API key, no Qdrant)

Uses the in-memory store + hashing embedder for linking. SQL generation needs an LLM, so the offline path is best explored via the benchmark's gold generator and the test suite; for live generation, add a key as below.

export VECTOR_BACKEND=memory EMBEDDER=hashing
python scripts/index_schema.py

Full run (real embeddings + LLM)

docker compose up -d                      # Qdrant, if VECTOR_BACKEND=qdrant
export ANTHROPIC_API_KEY=sk-...           # or OPENAI_API_KEY + LLM_PROVIDER=openai
export VECTOR_BACKEND=qdrant EMBEDDER=sentence-transformer

python scripts/index_schema.py
python scripts/ask.py "Which customer has placed the most orders?"
python scripts/run_benchmark.py

Benchmark output

run_benchmark.py prints the schema-linking ablation:

mode              exec_acc  valid_sql  link_recall  avg_tables
--------------------------------------------------------------
full_schema           ...        ...          ...         9.0
linked_schema         ...        ...          ...         ~

avg_tables is the average number of tables placed in the prompt — a direct proxy for prompt size and cost. On the small bundled schema the win is modest; the wider the real database, the larger the gap between full_schema and linked_schema while linking recall stays high. Per-example details (including the exact predicted SQL for every question) are written to results/per_example.json for error analysis.

Project structure

text-to-sql/
├── src/text2sql/
│   ├── config.py                 # env-based settings
│   ├── models.py                 # Schema, Table, Column, LinkedSchema, results (dataclasses)
│   ├── pipeline.py               # factories + Text2SqlPipeline facade
│   ├── schema/
│   │   ├── introspect.py         # read schema from SQLite (PRAGMA)
│   │   └── documents.py          # table -> embedding document
│   ├── linking/
│   │   ├── embedder.py           # SentenceTransformer | Hashing (offline)
│   │   ├── vector_store.py       # InMemory | Qdrant (same interface)
│   │   └── linker.py             # retrieve tables + FK expansion (the core)
│   ├── generation/
│   │   ├── llm.py                # Anthropic | OpenAI client
│   │   ├── prompt.py             # prompt building + defensive SQL extraction
│   │   └── generator.py          # SqlGenerator interface + LLM impl
│   ├── execution/
│   │   ├── executor.py           # read-only SQLite execution
│   │   └── comparison.py         # result-set equality (execution accuracy)
│   ├── evaluation/
│   │   ├── metrics.py            # linking recall, reference-table extraction
│   │   └── runner.py             # full-vs-linked ablation benchmark
│   └── simulation/demo_db.py     # deterministic demo DB + gold benchmark
├── scripts/                      # build_demo_db / index_schema / ask / run_benchmark
├── tests/                        # offline unit tests (real SQLite + fakes)
├── data/README.md                # demo vs bring-your-own database
├── docker-compose.yml            # local Qdrant
└── docs/DESIGN_DECISIONS.md

Design decisions

See docs/DESIGN_DECISIONS.md — why execution accuracy over string match, why FK expansion, the offline in-memory store, and the Protocol-based components that keep the core testable without a model, a database server, or an LLM.

Limitations & possible extensions

• Table-level linking — one vector per table. Column-level linking would help on very wide tables; the document builder is the only thing that changes.
• Single-database benchmark — the runner already generalises to any .sqlite + benchmark.json; wiring in a Spider subset would give a headline, comparable number.
• No self-correction loop — feeding an execution error back to the LLM for one repair attempt is a well-known, high-impact addition.
• Set-based comparison — column permutations aren't normalised (a documented, standard limitation of execution accuracy).

License

MIT