SQLAgent: Learning to Explore Before Generating as a Data Engineer

2026-02-02

AI Summary

This work addresses the challenge that current large language models often fail to generate accurate SQL queries on complex real-world databases due to insufficient understanding of schema structure, semantic ambiguity, and multi-table join paths. To overcome this limitation, the authors propose a two-stage framework: first, a Monte Carlo Tree Search–based strategy autonomously explores the database to construct a structured knowledge base; then, a dual-agent collaborative mechanism leverages this knowledge base to iteratively produce high-quality SQL queries. By decoupling database exploration from SQL generation—a novel approach in this domain—the method significantly enhances the model’s adaptability to unseen databases and its multi-step reasoning capability. Extensive experiments on large-scale benchmarks demonstrate substantial performance gains over strong existing baselines, confirming the effectiveness and generalization ability of the proposed approach.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• SQL generation
• database schema
• natural language interface
• query generalization
• LLM reasoning

Innovation

• SQL generation
• schema exploration
• Monte Carlo Tree Search
• in-context learning
• dual-agent system