Enhanced SQL Agent Architecture

Overview

The Enhanced SQL Agent is a state-of-the-art Text-to-SQL system that combines cutting-edge research insights with enterprise-grade reliability, safety, and performance. It implements a sophisticated multi-step reasoning workflow with intelligent schema analysis, advanced error handling, and comprehensive safety validation.

Research Foundation

This implementation is based on two key research papers:

1. “Automatic Metadata Extraction for Text-to-SQL” (arXiv:2505.19988)

  • Database Profiling: Systematic querying to gather statistics about tables and fields

  • Novel Schema Linking: Using LLM SQL generation capability for high-recall schema element identification

  • Query Log Analysis: Extracting metadata from historical SQL queries written by experts

  • SQL-to-Text Generation: Creating natural language questions from SQL for training data

2. “ReFoRCE: A Text-to-SQL Agent with Self-Refinement, Format Restriction, and Column Exploration”

  • Database Information Compression: Pattern-based table grouping to reduce schema size

  • Expected Answer Format Restriction: LLM-generated format specifications

  • Iterative Column Exploration: Execution-feedback driven column discovery

  • Self-Refinement Workflow: Multi-iteration query improvement with self-consistency

  • Parallelization with Consensus: Multiple candidates with majority-vote consensus

Architecture Components

Core Agents

EnhancedSqlAgent

The main orchestrator that implements the multi-step reasoning workflow.

Key Features:

  • Multi-step reasoning with state management

  • Intelligent schema linking and table selection

  • Comprehensive safety validation

  • Advanced error handling and recovery

  • Query optimization and performance analysis

ReFoRCESqlAgent

Advanced implementation incorporating ReFoRCE methodology for state-of-the-art performance.

Key Features:

  • Database information compression

  • Expected answer format restriction

  • Iterative column exploration with execution feedback

  • Self-refinement workflow with self-consistency

  • Parallelization with majority-vote consensus

Specialized Sub-Agents

SchemaLinkingAgent

Analyzes database schema and identifies relevant tables using intelligent algorithms.

Capabilities:

  • Schema relationship analysis

  • Table relevance scoring

  • Column significance assessment

  • Join path discovery

QueryPlannerAgent

Creates structured execution plans and analyzes query complexity.

Capabilities:

  • Query complexity analysis (Simple, Medium, Complex, Very Complex)

  • Multi-step query planning

  • Dependency identification

  • Performance estimation

SqlErrorHandler

Handles errors with intelligent recovery strategies.

Capabilities:

  • Error classification (Syntax, Semantic, Permission, Data, Performance)

  • Automatic query correction

  • Recovery strategy suggestions

  • Context-aware error analysis

SqlSafetyValidator

Validates queries for security and safety compliance.

Capabilities:

  • Multi-level security validation

  • SQL injection prevention

  • Resource limit enforcement

  • Configurable safety policies

Enhanced Tools

Core SQL Tools

  • SqlQueryTool: Executes SQL queries with safety checks

  • CreateTableTool: Creates database tables

  • QueryProfileTool: Profiles query performance

Advanced Analysis Tools

  • SchemaAnalysisTool: Analyzes database schema relationships

  • QueryValidationTool: Validates SQL syntax and semantics

  • TableSelectionTool: Suggests relevant tables for queries

  • QueryOptimizationTool: Provides performance optimization suggestions

  • QueryExplanationTool: Explains SQL queries in natural language

Workflow Architecture

Enhanced SQL Agent Workflow

Enhanced SQL Agent Workflow

The Enhanced SQL Agent follows an 8-step workflow:

  1. Question Analysis - Understands user intent and query complexity

  2. Schema Linking - Identifies relevant tables using intelligent algorithms

  3. Query Planning - Creates structured execution plans

  4. SQL Generation - Generates optimized SQL queries

  5. Safety Validation - Comprehensive security and safety checks

  6. Query Execution - Executes validated queries

  7. Error Handling - Intelligent error recovery and correction

  8. Response Formatting - User-friendly result presentation

ReFoRCE Workflow

ReFoRCE SQL Agent Workflow

The ReFoRCE agent implements an advanced workflow with additional optimizations:

  1. Database Information Compression - Pattern-based schema compression

  2. Format Restriction Generation - Expected answer format specification

  3. Column Exploration - Iterative column discovery with feedback

  4. Parallel Generation - Multiple SQL candidates in parallel

  5. Self-Refinement - Multi-iteration query improvement

  6. Consensus Voting - Majority-vote consensus mechanism

  7. Final Validation - Comprehensive result validation

Overall Architecture

SQL Agent Architecture Overview

State Management

The agents maintain comprehensive state throughout the workflow:

class SqlAgentState(TypedDict):
    messages: Annotated[List[BaseMessage], add_messages]
    original_question: str
    current_step: str
    schema_analysis: Optional[Dict]
    relevant_tables: Optional[List[Dict]]
    query_plan: Optional[Dict]
    generated_sql: Optional[str]
    validation_result: Optional[Dict]
    execution_result: Optional[str]
    error_info: Optional[Dict]
    safety_check: Optional[Dict]
    final_answer: Optional[str]
    retry_count: int
    max_retries: int

ReFoRCE Extended State

class ReFoRCESqlAgentState(SqlAgentState):
    compressed_schema: Optional[str]
    format_restriction: Optional[FormatRestriction]
    column_exploration: Optional[ColumnExplorationResult]
    self_refinement_iterations: int
    parallel_candidates: List[Dict[str, Any]]
    consensus_result: Optional[str]
    confidence_score: float

Safety and Security Features

Multi-Level Safety Validation

  1. Syntax Validation: Checks for valid SQL syntax

  2. Security Validation: Prevents dangerous operations

  3. Access Control: Enforces data access policies

  4. Resource Limits: Prevents resource-intensive queries

Configurable Safety Policies

safety_config = {
    "max_result_rows": 1000,
    "blocked_functions": ["DROP", "DELETE", "UPDATE"],
    "allowed_schemas": ["public", "analytics"],
    "query_timeout": 30,
    "max_joins": 5
}

Error Classification

  • Syntax Errors: Invalid SQL syntax

  • Semantic Errors: Valid syntax but logical errors

  • Permission Errors: Access control violations

  • Data Errors: Data-related issues

  • Performance Errors: Resource or timeout issues

Performance Optimizations

Database Schema Compression

  • Pattern-based table grouping

  • Representative table selection

  • Token limit management

  • Intelligent schema reduction

Parallel Processing

  • Multiple SQL candidate generation

  • Concurrent query execution

  • Majority-vote consensus

  • Confidence scoring

Query Optimization

  • Performance analysis

  • Index usage suggestions

  • Join optimization

  • Query rewriting recommendations

Usage Examples

Basic Enhanced Mode

from ryoma_ai.agent.sql import SqlAgent

# Initialize with enhanced capabilities
sql_agent = SqlAgent(
    model="gpt-4",
    datasource=datasource,
    use_enhanced_mode=True,
    safety_config={
        "max_result_rows": 1000,
        "blocked_functions": ["DROP", "DELETE"]
    }
)

# Process a question
result = sql_agent.invoke({
    "messages": [HumanMessage(content="Show top 10 customers by revenue")]
})

ReFoRCE Mode (State-of-the-Art)

# Initialize with ReFoRCE optimizations
sql_agent = SqlAgent(
    model="gpt-4",
    datasource=datasource,
    use_reforce_mode=True,
    max_parallel_threads=3,
    max_refinement_iterations=5,
    compression_threshold=30000
)

# Process complex queries with advanced reasoning
result = sql_agent.invoke({
    "messages": [HumanMessage(content="Calculate customer lifetime value by segment")]
})

Advanced Configuration

# Custom safety configuration
safety_config = {
    "max_result_rows": 5000,
    "blocked_functions": ["LOAD_FILE", "INTO OUTFILE"],
    "allowed_schemas": ["sales", "marketing"],
    "query_timeout": 60,
    "max_joins": 10,
    "safety_level": "STRICT"
}

# Initialize with custom settings
sql_agent = SqlAgent(
    model="gpt-4-turbo",
    datasource=datasource,
    use_reforce_mode=True,
    safety_config=safety_config,
    max_parallel_threads=5,
    max_refinement_iterations=3
)

Integration Points

Database Connections

The agents support multiple database types through the SqlDataSource interface:

  • PostgreSQL

  • MySQL

  • SQLite

  • SQL Server

  • Oracle

  • Snowflake

Model Integration

Compatible with various language models:

  • OpenAI GPT models (GPT-3.5, GPT-4)

  • Anthropic Claude models

  • Local models via Ollama

  • Azure OpenAI Service

Monitoring and Logging

Built-in monitoring capabilities:

  • Query execution metrics

  • Error tracking and analysis

  • Performance monitoring

  • Safety violation logging

Best Practices

Production Deployment

  1. Safety First: Always configure appropriate safety rules

  2. Resource Limits: Set reasonable query timeouts and result limits

  3. Access Control: Implement proper database access controls

  4. Monitoring: Set up comprehensive monitoring and alerting

  5. Testing: Thoroughly test with your specific database schema

Performance Optimization

  1. Schema Analysis: Regularly analyze and optimize database schema

  2. Index Management: Ensure proper indexing for common query patterns

  3. Query Caching: Implement query result caching where appropriate

  4. Resource Monitoring: Monitor database resource usage

Security Considerations

  1. Principle of Least Privilege: Grant minimal necessary database permissions

  2. Input Validation: Validate all user inputs

  3. Audit Logging: Log all query executions for audit purposes

  4. Regular Updates: Keep the system updated with latest security patches

Troubleshooting

Common Issues

  1. Schema Loading Errors: Check database connectivity and permissions

  2. Query Timeout: Adjust timeout settings or optimize queries

  3. Safety Violations: Review and adjust safety configuration

  4. Memory Issues: Consider schema compression for large databases

Debug Mode

Enable debug mode for detailed logging:

sql_agent = SqlAgent(
    model="gpt-4",
    datasource=datasource,
    use_enhanced_mode=True,
    debug=True,
    log_level="DEBUG"
)

Future Enhancements

Planned Features

  1. Advanced Caching: Intelligent query result caching

  2. Query Optimization: ML-based query optimization

  3. Multi-Database Support: Cross-database query capabilities

  4. Natural Language Explanations: Enhanced result explanations

  5. Interactive Query Building: Step-by-step query construction

Research Integration

Continuous integration of latest research:

  • New schema linking algorithms

  • Advanced error correction techniques

  • Improved safety validation methods

  • Performance optimization strategies

Contributing

For information on contributing to the Enhanced SQL Agent, please see the Contribution Guidelines.

References

  1. “Automatic Metadata Extraction for Text-to-SQL” (arXiv:2505.19988)

  2. “ReFoRCE: A Text-to-SQL Agent with Self-Refinement, Format Restriction, and Column Exploration”

  3. Snowflake Labs ReFoRCE Implementation