text-to-sql-agent-langgraph

CLAUDE.md — Project Context for Claude Code Sessions

This file orients any future Claude Code session working in this repo. Read this before making changes.

What this project is

A Text-to-SQL dashboard connected to one or more real, user-configured databases. A user types a natural-language question in Streamlit, a LangGraph agent turns it into SQL against the configured database (schema retrieved via ChromaDB, embedded from live schema introspection — not a hardcoded sample — so only relevant tables are shown to the LLM), the SQL is validated (SELECT-only allowlist), executed read-only, and the result is rendered as a table + auto-picked Plotly chart. The LLM runs locally via Ollama — no network calls for the LLM, no API keys required for that part. Database connectivity is fully config-driven via .env; there is no hardcoded connection string, host, or schema anywhere in the codebase.

Multiple databases: DB_CONNECTIONS in .env can list more than one named connection (config.settings.DatabaseConnectionConfig, one full DB_<NAME>_* field set per name) instead of the single legacy DB_* block. When more than one is configured, retrieve_schema_node auto-routes each question to whichever database’s schema looks most relevant (embeddings.retriever.select_database) before generating SQL — there is no manual database picker anywhere in the UI/API. A plain single-database .env (the common case, and everything this file describes elsewhere unless multi-database is called out explicitly) still works unchanged: it’s internally treated as one connection named "default". See “Multi-database auto-routing” under Key design decisions below.

History note: the project originally shipped with a bundled sample DuckDB e-commerce database for demo purposes. That was fully removed (by explicit user decision, no fallback/demo mode kept) in favor of connecting only to a real database via SQLAlchemy. If you see references to DuckDB, db/schema.sql, or scripts/seed_db.py anywhere (docs, old branches, stale comments), they’re leftover from that phase and should be treated as wrong, not as a parallel supported mode.

Multi-source (optional, off by default): ENABLE_MULTI_SOURCE_ROUTER=true switches ui/app.py/api/main.py from calling agent.graph.run_agent directly to calling agent.orchestrator.graph.run_orchestrated, which adds a router in front of the SQL pipeline and can fan a question out to up to three more sources: an “documents” and a separate, more access-sensitive “policies” PDF collection (agentic RAG, native SQL Server VECTOR storage), and live web search (Tavily). The SQL pipeline itself is never modified by any of this — see “Multi-source orchestration” under Key design decisions below, and docs/MULTI_SOURCE_GUIDE.md for how to configure each source.

Tech stack

Concern Choice
LLM runtime Ollama, default model llama3.1:8b (swap via .env / config/settings.py, e.g. sqlcoder, duckdb-nsql)
Orchestration LangGraph — explicit state machine, not a black-box agent. Two graphs: agent/graph.py (the SQL pipeline, always present) and, when multi-source is enabled, agent/orchestrator/graph.py (router + fan-out, sitting in front of it)
Schema retrieval ChromaDB (persisted locally) — embeds table DDL synthesized from live introspection, retrieves top-k relevant tables per question
Database User’s own — PostgreSQL, MySQL, SQL Server, or Oracle, via SQLAlchemy. Config-driven (DB_TYPE + connection params in .env), pluggable per db.connection.SUPPORTED_DB_TYPES. One or more named connections (DB_CONNECTIONS in .env); the agent auto-routes each question to the right one when more than one is configured
SQL parsing/validation sqlglot — parses generated SQL and checks statement type against an allowlist, in the dialect matching DB_TYPE
Document/policy RAG SQL Server 2025+/Azure SQL native VECTOR column type (rag/store.py) — a dedicated connection (RAG_STORE_CONNECTION_STRING), separate from DB_CONNECTIONS. Optional, off by default (ENABLE_DOCUMENT_RAG/ENABLE_POLICY_RAG)
Web search Configurable provider (search/web_search.py), Tavily implemented today. Optional, off by default (ENABLE_WEB_SEARCH + WEB_SEARCH_API_KEY)
UI Streamlit + Plotly. ui/app.py (chat) + ui/pages/1_Knowledge_Sources.py (PDF upload/management, Streamlit’s native multipage convention)
Python 3.11 is the target per project spec. This machine only has 3.14 installed (no 3.11 on PATH via py -0p) — the venv was created against 3.14. If a future session hits a wheel-availability issue for a pinned dependency, that’s why (see “Python 3.14 gotchas” below for two real ones already hit and fixed). Re-run py -0p to check if 3.11 has since been installed and consider recreating .venv against it if so.

Python 3.14 gotchas already hit (fixed, but worth knowing about)

Folder conventions

Key design decisions

Self-correcting retry loop (LangGraph)

The full graph (agent/graph.py) is eight nodes, not four: sanitize_input → classify_followup → retrieve_schema → generate_sql → validate_sql → estimate_cost → execute_sql → generate_insight. On a validation, cost-estimate, or execution failure, a conditional edge routes back to generate_sql (or, for a “missing reference” execution error, back to retrieve_schema) with the error message appended to the state’s history, so the LLM sees what went wrong and can correct itself. Capped at MAX_RETRIES = 3 (config/settings.py) — after that, the graph ends in a terminal failed state and the UI surfaces the last error rather than looping forever. This is the interview-relevant piece: it’s a small explicit state machine, not a ReAct-style free-form agent, specifically so the retry/error-feedback path is inspectable and boundable. See docs/ARCHITECTURE.md for the full per-node walkthrough and the complete retry-routing table.

Schema scoping (why ChromaDB at all)

For a large real schema, dumping every table’s DDL into the prompt burns context and increases hallucinated joins on irrelevant tables. Each table’s synthesized DDL (from live introspection) is embedded as one chunk; at query time we retrieve the top-k (SCHEMA_TOP_K, default 4) most relevant tables and only inject those into the generation prompt. This matters a lot more now than it did with the old bundled 5-table sample schema — a real production database can easily have hundreds of tables, which is exactly the case this code path is written for.

Multi-database auto-routing

DB_CONNECTIONS in .env can name more than one database (config.settings.DatabaseConnectionConfig, collected into Settings.databases). Each configured database gets its own Chroma collection (embeddings.schema_indexer.get_collection’s db_name param) — never a shared one, because FK-bridge/keyword-match expansion in embeddings/retriever.py only makes sense within one database’s own foreign-key graph, and a shared collection would also risk table-name collisions between two databases that happen to share a table name.

Routing itself (embeddings.retriever.select_database) is a cheap, separate pre-step: with one configured database it short-circuits immediately (no Chroma query, no behavior/latency change for a plain single-database setup — the overwhelmingly common case); with several, it queries every database’s collection for its own single best-matching table (n_results=1) and picks the database that wins. The existing, unmodified top-k/FK-bridge/keyword-fallback retrieval logic then runs exactly as before, scoped to that one winning database’s collection.

retrieve_schema_node calls select_database only on the first pass through a question and stores the result in AgentState["selected_database"]. The one retry path that re-enters retrieve_schema (execute_sql’s missing_reference retry — see the self-correcting retry loop above) reuses that stored value rather than re-routing: a retry must keep targeting the same database attempt 1 already generated/executed SQL against. Every downstream dialect/engine resolution (validate_sql_node, estimate_query_cost_node, execute_sql_node, and ui/app.py’s “Confirm and Run”) reads db.connection.get_connection(settings, state["selected_database"]) rather than a single global Settings.db_type.

Two things deliberately left alone by this design (documented, not silently ignored): the eval benchmark’s per-case database: label (see the eval/ folder note above) and config/table_descriptions.yaml/ config/sensitive_columns.yaml, which are keyed by bare table name, not (database, table) — a note/classification for one configured database’s table could in principle also apply to a same-named table in another. Both are real, narrow limitations worth knowing about if you’re extending this further, not oversights to silently work around.

Multi-source orchestration (router + subgraphs)

ENABLE_MULTI_SOURCE_ROUTER (default false) puts a router in front of the SQL pipeline, in agent/orchestrator/. Router + subgraphs, not a single tool-calling agent, was a deliberate choice, for the same reason the SQL pipeline itself is an explicit LangGraph state machine and not a ReAct-style agent: every source’s safety boundary (the SQL validator, the policy-sensitivity gate, the “web content is untrusted” framing) stays separately testable and inspectable rather than folded into one model’s implicit tool-selection reasoning.

agent.orchestrator.graph.run_orchestrated is the one entry point ui/app.py/api/main.py call, and it is deliberately a two-path function, not a graph with one trivial branch:

See docs/ARCHITECTURE.md’s “Multi-source orchestration” section for the full diagram and per-node walkthrough, and docs/MULTI_SOURCE_GUIDE.md for how to configure and use each source (including where the Tavily API key goes and how to upload a policy PDF).

Document/policy agentic RAG (rag/)

Same code serves “documents” (general uploads) and “policies” (more access-sensitive), parameterized by collection name (rag.graph.build_rag_subgraph(collection)), not two near-duplicate modules — they’re structurally identical and only differ in how generate_node treats a sensitive chunk. The subgraph is retrieve → grade → (rewrite → retry, bounded by RAG_MAX_RETRIES) → generate-with-citations, or an insufficient-information fallback after retries are exhausted — the same bounded self-correction philosophy as the SQL pipeline’s retry loop, a separate knob because a bad chunk retrieval and a bad SQL parse aren’t the same kind of budget.

Storage (rag/store.py) is SQL Server 2025+/Azure SQL’s native VECTOR column type — confirmed against this project’s actual target instance before being built, not assumed, including one real bug found and fixed along the way: a long (~7000+ character) embedding JSON string gets bound by pyodbc as ntext rather than nvarchar, and SQL Server’s VECTOR cast rejects ntext as a source type (“Explicit conversion from data type ntext to vector is not allowed”) — fixed by casting through NVARCHAR(MAX) first (rag.store._VECTOR_CAST). Storage is a dedicated connection (RAG_STORE_CONNECTION_STRING), never one of DB_CONNECTIONS — chunk/ embedding storage isn’t business data and shouldn’t share a schema or connection pool with a configured database.

Policy sensitivity (compensation/disciplinary/legal, set per-document at upload time — rag/store.py’s SensitivityCategory) is enforced by rag/graph.py’s generate_node refusing to summarize a sensitive chunk into an answer at all, since this app has no per-user authorization system to check who’s allowed to see it — the same fail-closed philosophy as agent/sql_validator.py’s SAFETY_VIOLATION_TYPES, applied to a different data shape (a document/chunk tag instead of a (table, column) pair).

Retrieved chunk text is framed as untrusted data, never instructions in rag/graph.py’s _GENERATE_SYSTEM_PROMPT — the same “SQL is untrusted output” principle below, applied to what a poisoned/malicious uploaded PDF could contain, since that’s this feature’s realistic injection vector.

search/web_search.py mirrors db/connection.py’s SUPPORTED_DB_TYPES pattern exactly: SUPPORTED_SEARCH_PROVIDERS maps a provider name to its call function, so WEB_SEARCH_PROVIDER is a .env change, not a code change. Only tavily is implemented today. Results are wrapped in a fixed WebResult shape before they ever reach a prompt and are explicitly framed as external/live/untrusted data in the answer-generation prompt (agent.orchestrator.nodes.web_search_node) — the answer text itself always opens with “According to a live web search:”, so it’s never presented as if it came from the company’s own systems, and the same “data, not instructions” principle as ingested PDF content applies here too, since a search result’s content is exactly as attacker-influenceable as a stored database value or an uploaded document.

SQL is untrusted output, always

The LLM’s SQL is never trusted at face value. agent/sql_validator.py parses it with sqlglot (in the dialect matching DB_TYPE) and rejects anything that isn’t a single SELECT/UNION/EXCEPT/INTERSECT statement (explicit allowlist of the parsed statement type, not a regex blocklist). Execution happens on a read-only-by-convention SQLAlchemy engine (db.connection.get_read_only_engine()), with a row cap (MAX_RESULT_ROWS, enforced both via LIMIT in the SQL text and independently via fetchmany() at the cursor level, so a malformed/mistranslated query can’t bypass it just by lacking a working LIMIT) and a query timeout enforced at the driver level where a cheap session-level SET exists (Postgres, MySQL) and via forced connection-abort otherwise (SQL Server, Oracle — see db/execution.py::_execute_with_timeout). This validation step runs every time SQL is about to be displayed or executed for the user, including after the user hand-edits the SQL box in the UI — an edit is exactly as untrusted as an LLM generation.

One nuance worth knowing if you’re reading ui/app.py: the LangGraph agent’s own internal retry loop does execute candidate SQL automatically (that’s how it detects and self-corrects runtime errors like an unknown column) — those internal executions are safe (read-only, validated, row-capped, timed out) but are never shown to the user. Nothing is rendered until the user clicks Confirm and Run, and that button always re-validates and re-executes the current SQL text fresh, rather than trusting whatever the agent’s last internal attempt produced.

True read-only enforcement is layered, not just code-level

get_read_only_engine() does not itself strip write privileges — there’s no generic, cross-database way to do that purely at the SQLAlchemy layer. The real guarantee is two layers: (1) the SQL validator, described above, and (2) the .env DB_USER should be a genuinely read-only database role/account (documented in README’s Security section, not silently assumed). If you’re asked to “harden” this further, that’s the layer to push on — a DB-level read-only user, not more code-level checks, since the validator is already an AST-based allowlist rather than a blocklist.

Caching

How to run

See README.md for full setup. Short version:

.venv\Scripts\Activate.ps1
ollama pull llama3.1:8b
# fill in .env with your real DB connection details first
python scripts\test_db_connection.py
python scripts\build_embeddings.py
streamlit run ui\app.py

How to run tests / lint

.\tasks.ps1 test     # pytest
.\tasks.ps1 lint      # ruff check + black --check + mypy
.\tasks.ps1 format    # black + ruff --fix

Equivalent make test, make lint, make format targets exist in the Makefile for anyone on WSL/macOS/Linux. All pytest tests are fully mocked (no real DB, no Ollama) — scripts/integration_test.py is the separate, manual, real-DB-required script; it is never run by pytest or CI.

Common commands

Task PowerShell Make
Create venv + install deps .\tasks.ps1 setup make setup
Verify DB connection python scripts\test_db_connection.py python scripts/test_db_connection.py
Build/refresh embeddings python scripts\build_embeddings.py python scripts/build_embeddings.py
Run app .\tasks.ps1 run make run
Run tests .\tasks.ps1 test make test
Lint .\tasks.ps1 lint make lint
Manual real-DB integration check python scripts\integration_test.py python scripts/integration_test.py

Windows / Visual Studio-specific notes

Coding standards

Known gaps / follow-ups (multi-source RAG)

Named explicitly rather than silently left for a future session to rediscover:

This site is open source. Improve this page.