Natural language to SQL query agent built with LangGraph and Google Gemini.

🌐 Now with Streamlit Web UI! - Beautiful, interactive web interface for easy access.

# 1. Setup
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

# 2. Add API key to .env
echo "GOOGLE_API_KEY=your_key_here" > .env

# 3. Launch web app
streamlit run app.py

Then open http://localhost:8501 and start asking questions!

• 🌐 Streamlit Web Interface - Modern UI with live results and schema viewer
• 🗣️ Natural Language Interface - Ask questions in plain English
• 🔒 Safety Checks - Prevents destructive SQL operations (DROP, DELETE, UPDATE)
• 🔄 Refinement Loop - Automatically retries unsafe queries with feedback
• 🤖 LLM-Powered - Uses Google Gemini 1.5 Flash for fast query generation
• 📊 Schema-Aware - Automatically reads database structure
• ✨ Human-Readable Answers - Converts SQL results to natural language

Built with LangGraph state machine:

START → Generate SQL → Safety Check → Execute Query → Summarize → END
                             ↓
                        (if unsafe + retries left)
                             ↓
                        Refine & Retry ──┐
                             ↑           │
                             └───────────┘
                             
                        (if max retries)
                             ↓
                        Error Message → END

• Query Generator: Converts natural language → SQL using schema context
• Safety Check: Validates query is read-only (regex pattern matching)
• Refine: Loops back to generator with feedback when unsafe query detected
• Execution: Runs approved SQL against SQLite database
• Summary: Converts results → natural language answer

When an unsafe query is detected, the agent:

1. Detects the safety violation (e.g., DELETE, DROP, UPDATE)
2. Provides feedback to the LLM about what went wrong
3. Retries query generation with the safety context
4. Limits attempts to max_retries (default: 2) to prevent infinite loops

This makes the agent resilient to edge cases where the LLM might initially generate unsafe SQL.

1. Install dependencies:

python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

2. Add your Gemini API key to .env:

GOOGLE_API_KEY=your_key_here

Get your key from: https://aistudio.google.com/apikey

3. Run the agent:

Streamlit Web UI (Recommended) 🌐

streamlit run app.py
# Opens at http://localhost:8501

CLI Interface

python main.py

4. View the graph visualization:

python visualize_graph.py
# Opens graph.png with the visual workflow diagram

Modern web interface with:

• 📝 Natural language input
• 📊 Visual results display
• 📋 Live database schema viewer
• 💡 Example question templates
• ⚙️ Configurable retry settings
• 🔍 Detailed query inspector

Terminal-based for automation and scripting.

sql-query-agent/
├── src/
│   ├── state.py       # State schema (TypedDict)
│   ├── nodes.py       # Node functions (generate, check, execute, summarize)
│   ├── tools.py       # Database tools (schema extractor, query executor)
│   └── graph.py       # LangGraph construction & compilation
├── data/
│   ├── company.db     # Sample SQLite database
│   └── create_sample_db.py
├── tests/             # Test suite
│   ├── test_nodes.py
│   ├── test_full_flow.py
│   ├── test_graph.py
│   └── test_refinement*.py
├── scripts/
│   └── visualize_graph.py
├── docs/
│   ├── program-req.md
│   └── streamlit-guide.md
├── app.py             # 🌐 Streamlit web interface
├── main.py            # 💻 CLI application
└── requirements.txt

❓ Your question: How many employees work in Engineering?

📊 Generated SQL:
   SELECT COUNT(*) FROM employees e 
   JOIN departments d ON e.department_id = d.id 
   WHERE d.name = 'Engineering'

🔒 Safety Check: SAFE
   Rows returned: 1

✨ Answer:
   There are 4 employees working in the Engineering department.

The company.db contains:

• departments: id, name, budget (4 departments)
• employees: id, name, age, department_id, salary (8 employees)

Blocks queries containing:

• DROP, DELETE, UPDATE, INSERT
• ALTER, TRUNCATE, CREATE, REPLACE
• EXEC, EXECUTE

Only SELECT queries are allowed.

# Test individual nodes
python test_nodes.py

# Test full pipeline
python test_full_flow.py

# Test the graph
python test_graph.py

• LangGraph - State machine framework
• LangChain - LLM orchestration
• Google Gemini - LLM for SQL generation
• SQLAlchemy - Database interaction
• SQLite - Sample database