SQL AI Agent

An LLM-based AI agent capable of answering questions about company employees and their work on projects. The agent automatically generates and executes SQL queries against a PostgreSQL database.

Features

• ✅ Automatic SQL query generation from natural language (any language)
• ✅ Responds in the same language as the question
• ✅ SQL query error correction (up to 3 attempts)
• ✅ Complex question decomposition into multiple simpler queries
• ✅ Query optimization analysis and suggestions
• ✅ Protection against expensive unbounded queries
• ✅ Memory of previous query results
• ✅ Support for different LLM models via OpenRouter
• ✅ Interactive and single-shot modes

Project Structure

sql-ai-agent/
├── agent.py                    # Main CLI script
├── sql_agent.py               # Agent core with agentic loop
├── llm_client.py              # LLM client
├── db_connector.py            # PostgreSQL connection
├── query_executor.py          # SQL execution with optimization analysis
├── setup_database.py          # Database setup script
├── import_vacation_data.py    # Vacation data importer
├── requirements.txt           # Python dependencies
├── setup.sh                   # Initial setup script
├── .env.example              # Configuration template
├── .env                      # Your configuration (not in git)
└── knowledge/
    ├── database.sql          # PostgreSQL dump
    ├── vacation_requests.json # Vacation data from ClickUp
    └── vacation_requests.py   # Pydantic schemas

Installation

1. Clone and Install Dependencies

# Navigate to project directory
cd sql-ai-agent

# Run setup script
./setup.sh

This script will:

• Create virtual environment
• Install all dependencies
• Create .env file from template

2. Configure .env File

Edit the .env file and provide your credentials:

# OpenRouter API Key
OPENROUTER_API_KEY=your-api-key-here

# PostgreSQL Configuration
DB_HOST=localhost
DB_PORT=5432
DB_NAME=sql_agent_db
DB_USER=postgres
DB_PASSWORD=your-password-here

3. Setup Database

Make sure PostgreSQL is running, then execute:

# Activate virtual environment
source venv/bin/activate

# Run database setup
python3 setup_database.py

This script will:

• Create the database
• Import data from knowledge/database.sql
• Import vacation data from knowledge/vacation_requests.json

Usage

Interactive Mode

python3 agent.py

Example:

Question: How many teams are in the database?
Question: What projects does team Alpha have?
Question: exit

Single-Shot Mode

python3 agent.py "Your question"

Examples:

# English
python3 agent.py "How many vacation days did each member of team Alpha take?"

# Ukrainian
python3 agent.py "Скільки днів відпустки взяв кожен член команди Alpha?"

Command Line Parameters

# Use a different model
python3 agent.py --model "anthropic/claude-3.5-sonnet" "Question"

# Disable verbose logs
python3 agent.py --no-verbose "Question"

# Change maximum retry attempts
python3 agent.py --max-retries 5 "Question"

# Change warning threshold for large result sets
python3 agent.py --max-rows-warning 500 "Question"

Test Questions

1. Project Activity

   Which is the most active project in team Fusion by hours logged last week?
   

2. Team Composition

   Describe the composition of each team by developer specialization
   

3. Vacations

   How many vacation days did each member of team Alpha take starting from the new year? List all vacation periods for each employee.
   

4. Estimate vs Actual Time

   By how much on average does the task estimate for team Alpha exceed the actual time spent?
   

Architecture

How It Works

User Question → SQL Agent → LLM → SQL Query → PostgreSQL
                    ↑                              ↓
                    ←── Results/Errors ←──────────┘

The SQL AI Agent follows this flow:

1. User asks a question in natural language (any language)
2. Agent loads database schema to understand table structure and relationships
3. LLM generates SQL query based on the question and schema
4. Query optimizer analyzes the generated SQL for potential performance issues
5. Query executes against PostgreSQL database
6. If error occurs: LLM fixes the query and retries (up to 3 attempts)
7. If no data returned: Agent tries alternative query approaches
8. Results are collected and passed back to the LLM
9. LLM generates final answer in the same language as the question

The agent maintains context throughout the conversation, allowing it to:

• Learn from previous query errors
• Use results from earlier queries to inform subsequent ones
• Decompose complex questions into multiple simpler queries

Query Optimization and Protection

QueryExecutor provides:

• Query analysis for potential performance issues
• Detection of expensive operations (SELECT *, JOINs without WHERE)
• Optimization suggestions for better query performance
• Warnings when queries return large result sets
• EXPLAIN plan analysis capability

Development

Testing Individual Components

# Test database connection
python3 db_connector.py

# Test query execution
python3 query_executor.py

# Test LLM client
python3 llm_client.py

# Test SQL agent
python3 sql_agent.py

Code Structure

• agent.py - CLI interface and entry point
• sql_agent.py - Main agent logic (agentic loop)
• llm_client.py - LLM interaction via OpenRouter
• query_executor.py - SQL execution with protection
• db_connector.py - PostgreSQL connection
• setup_database.py - Automated database setup
• import_vacation_data.py - JSON data import

Requirements

• Python 3.8+
• PostgreSQL 12+
• OpenRouter API key

Dependencies

• litellm - universal LLM client
• psycopg2-binary - PostgreSQL driver
• pydantic - data validation
• python-dotenv - environment variable loading

License

Educational project.