Beyond Single-Shot: Multi-step Tool Retrieval via Query Planning

TL;DR

TOOLQP introduces an iterative query planning framework for tool retrieval in LLM agents, significantly improving accuracy and generalization over traditional single-shot methods by decomposing complex instructions into sub-tasks.

Contribution

The paper presents TOOLQP, a novel approach that models retrieval as iterative query planning, enabling better handling of complex requests and tool compositions in LLM agents.

Findings

Achieves state-of-the-art performance in tool retrieval tasks.

Demonstrates superior zero-shot generalization capabilities.

Shows robustness across diverse retriever architectures.

Abstract

LLM agents operating over massive, dynamic tool libraries rely on effective retrieval, yet standard single-shot dense retrievers struggle with complex requests. These failures primarily stem from the disconnect between abstract user goals and technical documentation, and the limited capacity of fixed-size embeddings to model combinatorial tool compositions. To address these challenges, we propose TOOLQP, a lightweight framework that models retrieval as iterative query planning. Instead of single-shot matching, TOOLQP decomposes instructions into sub-tasks and dynamically generates queries to interact with the retriever, effectively bridging the semantic gap by targeting the specific sub-tasks required for composition. We train TOOLQP using synthetic query trajectories followed by optimization via Reinforcement Learning with Verifiable Rewards (RLVR). Experiments demonstrate that TOOLQP achieves state-of-the-art performance, exhibiting superior zero-shot generalization, robustness across diverse retrievers, and significant improvements in downstream agentic execution.