SQLAgent: Learning to Explore Before Generating as a Data Engineer

TL;DR

SQLAgent introduces a two-stage framework where an exploration phase builds a database-specific knowledge base, enabling more accurate SQL query generation in complex, real-world database scenarios.

Contribution

The paper presents a novel decoupled two-stage LLM-based approach that enhances generalization and accuracy in SQL query generation by proactive schema exploration and knowledge collection.

Findings

Significant accuracy improvements over baselines on large-scale benchmarks.

Effective handling of complex, multi-step reasoning in SQL queries.

Proactive schema exploration enhances generalization to unseen databases.

Abstract

Large Language Models have recently shown impressive capabilities in reasoning and code generation, making them promising tools for natural language interfaces to relational databases. However, existing approaches often fail to generalize in complex, real-world settings due to the highly database-specific nature of SQL reasoning, which requires deep familiarity with unique schemas, ambiguous semantics, and intricate join paths. To address this challenge, we introduce a novel two-stage LLM-based framework that decouples knowledge acquisition from query generation. In the Exploration Stage, the system autonomously constructs a database-specific knowledge base by navigating the schema with a Monte Carlo Tree Search-inspired strategy, generating triplets of schema fragments, executable queries, and natural language descriptions as usage examples. In the Deployment Stage, a dual-agent system leverages the collected knowledge as in-context examples to iteratively retrieve relevant information and generate accurate SQL queries in response to user questions. This design enables the agent to proactively familiarize itself with unseen databases and handle complex, multi-step reasoning. Extensive experiments on large-scale benchmarks demonstrate that our approach significantly improves accuracy over strong baselines, highlighting its effectiveness and generalizability.