Skip to content

About

RAG text-to-SQL: ChromaDB retrieval + Gemini generation behind a FastAPI service, wired into a PostgreSQL quiz app with sandboxed execution and answer-key validation.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Repository files navigation

RAG Text-to-SQL for the IDS706 SQL Quiz

CI Live API Live app

A retrieval-augmented text-to-SQL service for the SQL quiz built in Week 7. A student types a question in plain English on the quiz page, the service retrieves the relevant schema notes, business rules, and example queries from a vector database, asks Gemini to write PostgreSQL, and the quiz page runs the result in its read-only sandbox and checks it against the answer key.

Try it: open https://luel-du.github.io/ids706-sql-quiz/, sign in (any email, password of six or more characters), and use Ask the SQL assistant under the editor.

API: https://text2sql-rag.onrender.com (free tier; the first request after idle can take up to a minute).

Contents

Architecture

flowchart LR
  P[Quiz page<br/>GitHub Pages] -- "question + student token" --> A[FastAPI service<br/>Render]
  A -- "verify token" --> S[Supabase Auth]
  A -- "embed question" --> G1[Gemini<br/>gemini-embedding-001]
  A -- "nearest chunks" --> C[(ChromaDB<br/>schema, rules, examples)]
  A -- "context + question" --> G2[Gemini<br/>gemini-3.8-flash]
  A -- "SQL + context" --> P
  P -- "run_quiz_sql(sql)" --> D[(Supabase Postgres<br/>read-only sandbox)]
Loading
Component Technology Role
Vector database ChromaDB 1.5, persistent, cosine distance Stores 18 knowledge chunks with 768-dimension Gemini embeddings
Embeddings Gemini gemini-embedding-001 Embeds chunks as documents and questions as queries
Generation Gemini gemini-3.8-flash, temperature 0, JSON output Writes the SQL, an explanation, and an answerable flag
API FastAPI on Render Verifies the student's session, runs retrieval and generation, rate-limits per user
Execution The quiz database's run_quiz_sql() function Runs the SQL under a read-only role, caps rows, logs to the audit table, reports whether the result matches the answer key
Front end The existing quiz page One input, one button; shows SQL, explanation, retrieved context, rows, and the match check

How a question is answered

  1. The page sends POST /generate with the question and the student's Supabase access token.
  2. The service asks Supabase Auth who the token belongs to. No token, no answer.
  3. The question is embedded with gemini-embedding-001 using the RETRIEVAL_QUERY task type.
  4. ChromaDB returns the nearest chunks per kind: 2 tables, 2 column notes, 2 business rules, 2 examples.
  5. Schema linking: for each retrieved example, the tables its SQL reads from are added to the context if the table search missed them. The page shows these as "via example-qN".
  6. If the best distance is above 0.45 the request is flagged low confidence and the prompt tells the model to say so if the question is not about this database.
  7. A prompt is assembled from a fixed system instruction, the retrieved chunks, and the question. Gemini returns JSON: answerable, sql, explanation.
  8. The page places the SQL in the editor and triggers Run, which calls run_quiz_sql in the quiz database with the student's token. Rows come back with matched: true/false against the stored answer key.

The knowledge base

text2sql/knowledge.py holds the chunks as plain Python. Each chunk is one self-contained piece of text.

Kind Count Example
table 4 enrollments(...): what it stores, which joins it supports, which questions it answers, synonyms students use
column 5 grade_points is on a 0.00 to 4.00 scale; "average grade" means avg(grade_points)
rule 4 "more than the average student" means a per-student CTE compared with a scalar subquery
example 5 the five quiz questions with their answer-key SQL

For examples only the question text is embedded; the SQL travels as metadata so it does not blur the vector. Rules that apply to every query (single SELECT, plain table names, ORDER BY with a tiebreaker) are not retrieved; they live in the system instruction in text2sql/generate.py.

The ChromaDB collection stores a fingerprint of the knowledge, the embedding model, and the index layout. Any change to those triggers a rebuild on the next request, so the index can never silently serve stale vectors.

Retrieval design

Why not paste the whole schema into the prompt? For four tables that would work. The design here is the one that still works at four hundred: retrieve by kind so a table definition is always present, use past queries as a schema-linking signal, keep always-on rules out of the retrieval pool, and flag weak matches instead of answering anyway. Each of these was added after watching the naive version fail on this database:

  • With a flat top-3, a generic rule mentioning "courses" crowded out the table definition for question 1.
  • "How many students take each class?" retrieved the students and instructors tables, not enrollments, because the literal word "students" outranked any description. Better table descriptions moved courses into the top two; the example link is what brought enrollments in.
  • "What is the capital of France?" still returns three nearest chunks, because a vector search always has a nearest neighbour. The distance threshold and the instruction together let the model decline.

ask.py --context prints the retrieved chunks and distances for any question, which is how these cases were found.

API

Method Path Auth Body Returns
GET /health none status, model names, chunk count
POST /generate Bearer Supabase token {"question": "..."} answerable, sql, explanation, context[], best_distance, low_confidence, model
POST /run Bearer Supabase token {"sql": "...", "question_id": 3} the sandbox's response: ok, rows, row_count, truncated, matched

Interactive docs are served at /docs. Requests are limited to 20 per minute per user.

The quiz page uses /generate and then the quiz's own Run button, so /run is there for other clients, for example a command-line tool or a notebook.

Security model

Concern Mechanism
Who can call the service Only holders of a valid Supabase session from the quiz; the token is checked against Supabase Auth on every request
What the service can do to the database Nothing directly. It holds no database password. Execution goes through run_quiz_sql as the signed-in student, under a NOLOGIN role with SELECT on four tables, a 200-row cap, and an append-only audit log
Secrets GEMINI_API_KEY is the only secret, set as an environment variable on the host. SUPABASE_URL and the anon key are public values that already ship in the quiz page
Prompt content Schema descriptions, business rules, examples, and the question. Never database rows: generation happens before execution
Abuse of the Gemini quota Per-user rate limit; CORS restricted to the quiz's origin

Running locally

git clone https://github.com/luel-du/text2sql-rag.git
cd text2sql-rag
python3 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
cp .env.example .env     # set GEMINI_API_KEY, SUPABASE_URL, SUPABASE_ANON_KEY

Command line, no server needed:

python ask.py --context "How many students take each class?"   # retrieved chunks + SQL
python ask.py --prompt "Who teaches Linear Algebra?"            # the exact prompt sent to Gemini
python ask.py                                                   # built-in demo questions

The service:

uvicorn api.main:app --reload --port 8080
# open http://localhost:8080/docs and authorize with a Supabase access token from the quiz page

The first run builds chroma_db/ from knowledge.py; later runs reuse it until the knowledge changes.

Tests

pytest -q

Twenty-one tests run in about a second with no API key and no network. They cover the knowledge base's shape, table detection in example SQL (including the alias case that once pulled in the wrong table), the stale-index fingerprint, prompt assembly, and the API's behaviour with Supabase and Gemini replaced by fakes: token required, bad token refused, prompt never returned to the client, the student's token forwarded to the sandbox, and the per-user rate limit. GitHub Actions runs them on every push.

Retrieval quality is deliberately not tested in CI, since that needs live embedding calls; use ask.py --context for that by hand.

Deploying

render.yaml describes the service for Render's Blueprint flow. Create a Blueprint from this repository, enter GEMINI_API_KEY, SUPABASE_URL, and SUPABASE_ANON_KEY when prompted, and deploy. The free instance has an ephemeral disk, so CHROMA_PATH points at /tmp and the index is rebuilt at startup, which takes about a second. ALLOWED_ORIGINS must list the quiz page's origin.

The quiz page reads the service URL from docs/app.js (RAG_API), overridable through config.js.

Screenshots

A question and its generated SQL, with the retrieved context expanded.

Question, generated SQL, and retrieved context

Executed SQL, result table, and validation: the generated query ran in the sandbox and matched the answer key.

Executed result with the match check

Project structure

.
|-- README.md
|-- pyproject.toml              pytest configuration
|-- .env.example                variable names with placeholders
|-- .gitignore                  .env, chroma_db/, .venv/
|-- requirements.txt            chromadb, google-genai, fastapi, uvicorn, httpx, python-dotenv, pytest (pinned)
|-- render.yaml                 Render Blueprint for the API
|-- ask.py                      command-line front end: question -> context -> SQL
|-- .github/workflows/ci.yml    runs the tests on every push
|-- tests/                      offline tests: knowledge, store logic, prompt, API
|-- api/
|   `-- main.py                 FastAPI service: auth, rate limit, /generate, /run, /health
|-- text2sql/
|   |-- knowledge.py            the chunks: tables, column semantics, rules, examples
|   |-- store.py                ChromaDB build/retrieve, per-kind retrieval, example-to-table linking, fingerprint
|   `-- generate.py             system instruction, prompt assembly, Gemini call with JSON output
`-- assets/                     README screenshots

Built for IDS706 Data Engineering as an extension of the Week 7 SQL quiz. The database is a fictional university registrar.

About

RAG text-to-SQL: ChromaDB retrieval + Gemini generation behind a FastAPI service, wired into a PostgreSQL quiz app with sandboxed execution and answer-key validation.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages