LLM Database Agent

A AI-powered database querying system that transforms natural language questions into SQL queries and provides insights with visualizations.

Quick Start

Prerequisites

• Python 3.12+
• Ollama with Code Llama 7B model

Running the Application

1. Start Ollama and pull the model:

   ollama pull codellama:7b-code

2. Launch the Streamlit app:

   streamlit run app.py

3. Open your browser to http://localhost:8501

LLM agent/
├── agent_main.py            # Main LLM agent
├── test_correction_direct.py # Self-correction testing script
├── test_self_correction.py  # Additional test suite
├── sales_data.db            # SQLite database
└── README.md                # This documentation

Features

• Agent Control Panel: Interactive sidebar with database selection and agent management
• Database Selection: Switch between data sources with automatic schema reloading
• Schema Discovery: Automatically reads and understands database structure
• Natural Language Queries: Uses Code Llama to convert questions to SQL
• Self-Correction Loop: Automatically retries and learns from SQL errors
• Conversational Memory: Remembers previous questions for context-aware follow-ups
• Disambiguation Routing: Identifies ambiguous questions and asks for clarification
• Secure Execution: Validates and safely executes SQL queries
• Intelligent Responses: Provides natural language answers
• Visualization: Generates matplotlib code for data visualization
• Interactive Web Interface: Professional Streamlit application with real-time updates
• Error Recovery: Handles LLM hallucinations and schema misunderstandings

Prerequisites

1. Python Dependencies

pip install requests pandas matplotlib

2. Ollama with Code Llama

# Install Ollama (if not already installed)
# Download from: https://ollama.ai/

# Pull Code Llama model
ollama pull codellama:7b-code

# Start Ollama server (keep running in background)
ollama serve

3. Database Setup

python sales_database.py

Usage

Quick Start

python agent_main.py

Example Output

Starting LLM Database Agent
==================================================
Agent initialized with database: sales_data.db
Database schema loaded: ['Users', 'Products', 'Orders'] tables

Available Tables:
  • Users: user_id, username, join_date
  • Products: product_id, product_name, category, price
  • Orders: order_id, user_id, product_id, quantity, order_date

Example Question: How many orders does user with ID 1 have?

Calling Code Llama to generate SQL...
Generated SQL: SELECT COUNT(*) FROM Orders WHERE user_id = 1;
Executing SQL: SELECT COUNT(*) FROM Orders WHERE user_id = 1;
Execution successful!

[AGENT NATURAL LANGUAGE RESPONSE]
Based on your question 'How many orders does user with ID 1 have?', I found 1 result(s).
The answer is: 2

[AGENT VISUALIZATION CODE]
import matplotlib.pyplot as plt
import pandas as pd
# ... (visualization code generated)

Database Schema

Users Table

• user_id (INTEGER PRIMARY KEY)
• username (TEXT)
• join_date (TEXT)

Products Table

• product_id (INTEGER PRIMARY KEY)
• product_name (TEXT)
• category (TEXT)
• price (REAL)

Orders Table

• order_id (INTEGER PRIMARY KEY)
• user_id (INTEGER, FK → Users.user_id)
• product_id (INTEGER, FK → Products.product_id)
• quantity (INTEGER)
• order_date (TEXT)

Security Features

• SQL Injection Protection: Validates queries for dangerous keywords
• Read-Only Operations: Blocks INSERT, UPDATE, DELETE, DROP operations
• Safe Execution: Uses pandas for secure query execution
• Error Handling: Comprehensive exception handling

Configuration

LLM Settings

• Model: codellama:7b-code
• Endpoint: http://localhost:11434/api/generate
• Temperature: 0.1 (for consistent SQL generation)

Database Settings

• File: sales_data.db
• Type: SQLite 3
• Connection: Local file-based

Troubleshooting

Common Issues

1. "LLM connection failed"

   • Ensure Ollama is running: ollama serve
   • Check if Code Llama is installed: ollama list

2. "Database not found"

   • Run the database setup: python sales_database.py

3. "Package import errors"

   • Install dependencies: pip install requests pandas matplotlib

Performance

• Query Response Time: ~2-5 seconds (depending on LLM)
• Database Size: ~50KB (sample data)
• Memory Usage: ~100MB (including LLM client)

Agent Workflow

SQL Generation with Self-Correction

1. Schema Discovery: Read database structure and relationships
2. Query Processing: Convert natural language to SQL using LLM with chat history context
3. Self-Correction Loop:
   • Try to execute SQL
   • If error occurs, pass error context to LLM for correction
   • Retry up to 3 times with improved prompts
   • Learn from previous failures

Secure Execution & Validation

4. Security Validation: Check for dangerous SQL operations
5. Execution: Run query safely and capture results

User Experience

6. Natural Language Synthesis: Use LLM to generate conversational response
7. Intelligent Visualization: Use LLM to create appropriate chart code
8. Final Output: Provide complete answer with visualization

Conversational Memory

9. Chat History Management: Store successful interactions for context
10. Context-Aware Processing: Use conversation history in future queries
11. Follow-up Understanding: Handle implicit references to previous topics

Complete Flow Example with Memory

User: "What are the total orders for user 1?"
   → SQL: SELECT COUNT(*) FROM Orders WHERE user_id = 1;
   → Result: 2 orders
   → Saved to chat history

User: "What about user 2?" (Follow-up question)
   → Context: Previous question was about total orders for users
   → SQL: SELECT COUNT(*) FROM Orders WHERE user_id = 2;
   → Result: 1 order
   → Agent understands "user 2" refers to same query type

Technical Details

Core Components

• DatabaseAgent: Main agent class handling all operations
• Schema Reader: Extracts table structures and relationships
• LLM Interface: Communicates with Code Llama via REST API
• Security Validator: Prevents harmful SQL execution
• Response Synthesizer: Converts results to natural language
• Visualization Generator: Creates plot code based on data types

Error Handling

• Connection timeouts and retries
• SQL syntax error recovery
• Graceful LLM failure handling
• Database lock prevention

---