SQLAgent: Learning to Explore Before Generating as a Data Engineer

TL;DR

SQLAgent addresses the limited generalization of Text-to-SQL systems to unseen, enterprise-scale databases by decoupling knowledge acquisition from query generation. It uses an Exploration Stage where an LLM-guided MCTS-like process constructs a database-specific knowledge base of schema triplets $(S, Q, U)$, and a Deployment Stage where a dual-agent loop InfoAgent/GenAgent leverages these triplets as in-context examples to iteratively synthesize accurate SQL. The approach yields significant accuracy gains on large benchmarks like Spider 2.0-Snow, with ablations showing the value of the exploration knowledge and the collaborative synthesis. Its offline knowledge-building plus inference-time retrieval offers improved robustness and efficiency for complex, multi-step SQL reasoning in real-world 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.

Figures (3)

• Figure 1: Databases Representation and Exploration Phase. The left side of this diagram illustrates our representation and processing of the database structure. The right side of this diagram shows a schematic of our exploration using Tree Search on the existing structure. This includes four phases: selection, expansion, simulation, and backpropagation. This approach enables the collection of a series of triplets.
• Figure 2: SQL Deployment Stage. In this stage, the database information obtained in the exploration stage and the user's actual query are used to generate the SQL. The dual-agent architecture controls the information acquisition context and SQL generation context of the agent, enabling the system to process complex SQL statements while maintaining high query accuracy.
• Figure 3: Analysis of the Exploration Stage and Key Hyperparameters. The left figure shows Execution Accuracy as a function of agent iterations, demonstrating that the Exploration Stage consistently improves performance. The middle figure presents the effect of LLM temperature on different models, with the GPT model's accuracy peaking at a temperature of 0.7. The right figure evaluates the top-k schema retrieval parameter, showing our method (Ours) outperforms the Baseline across all settings and that its performance gain saturates for k values greater than 3.