🧠 MSGE-SQL — Multi-Strategy Generation & Evaluation for SQL

A Text-to-SQL agent that turns natural-language questions into validated, high-accuracy SQL queries — combining retrieval-augmented context, multi-strategy SQL generation, and an LLM-as-a-Judge selection stage.

---

📑 Table of Contents

• Overview
• How the Text2SQL Agent Works
  • 1. Offline Preprocessing
  • 2. Context Awareness
  • 3. SQL Generation
  • 4. SQL Validation
  • 5. SQL Selection
• Tech Stack
• Repository Structure
• Runtime Notes & Token Limits
• Getting Started
• Contributing
• License & Credits

---

🔍 Overview

MSGE-SQL is a multi-agent Text-to-SQL pipeline built to answer natural-language questions over a relational database with high accuracy and minimal hallucination. Instead of relying on a single LLM call to translate a question directly into SQL, the framework runs an offline preprocessing stage followed by a four-step online pipeline — context retrieval, multi-strategy generation, iterative validation, and judge-based selection — so that only the most accurate, executable query reaches the user.

Key components:

• Retrieval layer — vector database lookups for schema, evidence, categorical values, and few-shot examples.
• Generation layer — three parallel SQL-generation strategies to maximize candidate diversity.
• Validation layer — a syntax/semantic self-correction loop (sql_validation_agent).
• Selection layer — an LLM-as-a-Judge node that ranks and picks the best candidate.

---

⚙️ How the Text2SQL Agent Works

1. Offline Preprocessing

Before the agent can answer any question, a one-time offline phase prepares the database for retrieval:

• The schema is represented in two formats: a Markdown schema (table/column names, descriptions, foreign keys, relation types) for general-purpose LLMs, and a DDL schema for code-specialized LLMs.
• Textual cell values, schema representations, question–SQL example pairs, and expert/business knowledge are all indexed into a vector database, so the online pipeline retrieves only what's relevant instead of overloading the LLM's context window.

2. Context Awareness

At query time, the agent gathers the context needed to generate an accurate query:

• Dual schema representation — DDL for code-oriented models, Markdown for general-purpose models, since each model family interprets one format better than the other.
• Masked few-shot retrieval — the user's question is masked (entities/values stripped) before being used to retrieve similar past question–SQL pairs, so retrieval matches on intent and structure rather than surface wording.
• Categorical value grounding — real database values are surfaced so generated filters match actual stored data (e.g., using "Casa" instead of "Casablanca" if that's how the value is stored), avoiding silent empty-result bugs.
• Expert knowledge injection — domain-specific business rules and guidelines are added so generated SQL respects operational constraints, not just schema correctness.

3. SQL Generation

Three complementary strategies run to maximize the odds that at least one high-quality candidate is produced:

Strategy	Approach
In-Context Learning (ICL)	Uses retrieved few-shot examples with Chain-of-Thought prompting to leverage large-LLM pattern recognition for rapid candidate generation.
Divide-and-Prompt (Clause-by-Clause)	Decomposes the task: schema linking → execution planning → clause-by-clause SQL generation, in the order FROM → filters/aggregation → SELECT, using a code-specialized LLM on the DDL schema. Best for complex, multi-join queries.
Skeleton-Based Generation	A reasoning model first drafts the query's structural skeleton (filters, aggregation, grouping, joins) without schema context, then fills in details once schema and categorical values are provided — reducing premature anchoring on surface details.

4. SQL Validation

Each candidate passes through a two-stage, iterative self-correction loop:

1. Syntax Checker — an LLM-based component detects and repairs invalid SQL syntax, re-checking until the query is valid or a retry limit is reached.
2. Semantic Checker — a reasoning-oriented LLM verifies the query actually reflects user intent, catching logical errors the syntax checker can't see.

If a semantic fix introduces a new syntax issue, the query is routed back to the Syntax Checker — the two run in a feedback loop until the candidate is both syntactically valid and semantically aligned.

5. SQL Selection

Finally, an LLM-as-a-Judge evaluates the surviving candidates on logical structure, reasoning accuracy, and alignment with the user's goal:

• Candidates that return identical execution results are grouped as semantically equivalent, and only one representative per group is sent to the judge — reducing evaluation cost.
• The highest-scoring candidate is returned as the final answer.

---

🧰 Tech Stack

Category	Technologies
Orchestration	LangGraph (multi-agent workflow graph), LangChain
API / Backend	FastAPI
LLMs	Qwen, GPT-OSS (e.g. openai/gpt-oss-120b), reasoning + code-specialized models
Vector Store	ChromaDB (persistent local store, chroma_db_store/)
Structured Output	Pydantic
Language	Python 3.10+

---

📂 Repository Structure

.
├── models.py                  # LLM and vector DB client configuration
├── prompts.py                 # System and user prompt templates
├── states.py                  # Typed state definitions for agents and merge rules
├── text2sql.py                # Main workflow graph: retrieval → generation → validation
├── sql_validation_agent.py    # Subgraph running syntax_checker / semantic_checker loops
├── main.py                    # Small runner to invoke the workflow
└── chroma_db_store/           # Persistent Chroma DB storage (grounding collections)

---

⚠️ Runtime Notes & Token Limits

• The semantic_checker node uses the llm defined in models.py (defaults to openai/gpt-oss-120b in this workspace). This model and your API plan have a tokens-per-minute (TPM) limit — running multiple generation branches in parallel with large prompts / long full_context can trigger 429 rate-limit errors.
• Short-term mitigations:
  • Reduce parallel fan-out in text2sql.py so only one generator runs per request.
  • Use a smaller/cheaper model for semantic checks in models.py.
  • Trim full_context passed into semantic prompts (prompts.py) to fewer lines.

---

🚀 Getting Started

# Run the workflow locally
python main.py

This will execute the graph end-to-end and save the workflow diagram for inspection.

---

🤝 Contributing

• Run the workflow locally via python main.py and inspect the saved diagram.
• Please open issues for edge cases (e.g., inconsistent state types) and include the full traceback and input question.

---

📄 License & Credits

Project created for a Master-level Text2SQL research exercise. Reuse and modification are allowed for research and educational purposes.