Table of Contents
Fetching ...

ToolDreamer: Instilling LLM Reasoning Into Tool Retrievers

Saptarshi Sengupta, Zhengyu Zhou, Jun Araki, Xingbo Wang, Bingqing Wang, Suhang Wang, Zhe Feng

TL;DR

ToolDreamer tackles the problem of retrieving relevant tools when tool sets are too large for LLM context windows by conditioning retrievers on LLM-generated hypothetical tools (HTs). The method trains HT-aligned retrievers via a two-stage process: generate HTs, align them to gold tools with bipartite matching, and optimize an augmented InfoNCE loss using HT metadata. In inference, HTs are generated for a test query, top-K tools are retrieved per HT, and results are unified with Reciprocal Rank Fusion or LLM-based fusion; experiments on ToolRet show improvements for both sparse and dense retrievers, illustrating the framework’s effectiveness and flexibility. The work demonstrates that offloading reasoning to a guided retriever can enable scalable tool usage and improve tool-calling performance in large tool ecosystems.

Abstract

Tool calling has become increasingly popular for Large Language Models (LLMs). However, for large tool sets, the resulting tokens would exceed the LLM's context window limit, making it impossible to include every tool. Hence, an external retriever is used to provide LLMs with the most relevant tools for a query. Existing retrieval models rank tools based on the similarity between a user query and a tool description (TD). This leads to suboptimal retrieval as user requests are often poorly aligned with the language of TD. To remedy the issue, we propose ToolDreamer, a framework to condition retriever models to fetch tools based on hypothetical (synthetic) TD generated using an LLM, i.e., description of tools that the LLM feels will be potentially useful for the query. The framework enables a more natural alignment between queries and tools within the language space of TD's. We apply ToolDreamer on the ToolRet dataset and show that our method improves the performance of sparse and dense retrievers with and without training, thus showcasing its flexibility. Through our proposed framework, our aim is to offload a portion of the reasoning burden to the retriever so that the LLM may effectively handle a large collection of tools without inundating its context window.

ToolDreamer: Instilling LLM Reasoning Into Tool Retrievers

TL;DR

ToolDreamer tackles the problem of retrieving relevant tools when tool sets are too large for LLM context windows by conditioning retrievers on LLM-generated hypothetical tools (HTs). The method trains HT-aligned retrievers via a two-stage process: generate HTs, align them to gold tools with bipartite matching, and optimize an augmented InfoNCE loss using HT metadata. In inference, HTs are generated for a test query, top-K tools are retrieved per HT, and results are unified with Reciprocal Rank Fusion or LLM-based fusion; experiments on ToolRet show improvements for both sparse and dense retrievers, illustrating the framework’s effectiveness and flexibility. The work demonstrates that offloading reasoning to a guided retriever can enable scalable tool usage and improve tool-calling performance in large tool ecosystems.

Abstract

Tool calling has become increasingly popular for Large Language Models (LLMs). However, for large tool sets, the resulting tokens would exceed the LLM's context window limit, making it impossible to include every tool. Hence, an external retriever is used to provide LLMs with the most relevant tools for a query. Existing retrieval models rank tools based on the similarity between a user query and a tool description (TD). This leads to suboptimal retrieval as user requests are often poorly aligned with the language of TD. To remedy the issue, we propose ToolDreamer, a framework to condition retriever models to fetch tools based on hypothetical (synthetic) TD generated using an LLM, i.e., description of tools that the LLM feels will be potentially useful for the query. The framework enables a more natural alignment between queries and tools within the language space of TD's. We apply ToolDreamer on the ToolRet dataset and show that our method improves the performance of sparse and dense retrievers with and without training, thus showcasing its flexibility. Through our proposed framework, our aim is to offload a portion of the reasoning burden to the retriever so that the LLM may effectively handle a large collection of tools without inundating its context window.
Paper Structure (27 sections, 1 equation, 8 figures, 6 tables)

This paper contains 27 sections, 1 equation, 8 figures, 6 tables.

Figures (8)

  • Figure 1: Illustration of the problem statement/motivation. Retrievers conditioned to learn query-tool relationships struggle to locate tools that are not readily obvious from the question (top). LLM-generated hypothetical tools can potentially be used to locate the right tool but, it still misses the mark as it has not been conditioned to learn such tool-tool relationships (middle). Thus, if we first align the hypothetical tool with its gold tool counterpart, and incorporate the query background to train the retriever, it would lead to a far superior model capable of accurately locating the correct tool (bottom).
  • Figure 2: ToolDreamer overview. Training Phase (Top): A strong LLM first generates hypothetical tools (HTs) for a question. These tools are then aligned with their gold tool (GT) counterparts. Finally, the aligned tool pairs are used to train the retrieval model. Inference Phase (Bottom): At test time, an unseen question is provided to the same LLM which generates a set of hypothetical tools for it. For each tool, the trained retriever collects a top-K list of tools from the database. A unified list is obtained through reciprocal rank fusion (RRF).
  • Figure 3: Impact of scaling the number of tools on LLM hit-rate (correct tool call with appropriate arguments).
  • Figure 4: An example of hypothetical tools for a given query.
  • Figure 5: Prompt for generating inferior quality hypothetical tools for test data (ablation study).
  • ...and 3 more figures