Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving Tool Retrieval by Leveraging Large Language Models for Query Generation

Mohammad Kachuee, Sarthak Ahuja, Vaibhav Kumar, Puyang Xu, Xiaohu Liu

TL;DR

The paper tackles the challenge of retrieving relevant tools/APIs for large-scale LLM tool use. It proposes generating retrieval queries via LLMs and using dense retrieval with interleaving to select relevant APIs, reducing the need to complicate the retriever itself. Through a comparative study of zero-shot prompting, supervised fine-tuning, and alignment learning, the work finds that SFT excels in-domain while alignment learning provides the strongest out-of-domain performance, especially as new tools are introduced. The findings highlight a practical pathway for scalable tool integration, enabling better recall and ranking of appropriate APIs in dynamic tool ecosystems, with strong implications for real-world LLM-assisted workflows.

Abstract

Using tools by Large Language Models (LLMs) is a promising avenue to extend their reach beyond language or conversational settings. The number of tools can scale to thousands as they enable accessing sensory information, fetching updated factual knowledge, or taking actions in the real world. In such settings, in-context learning by providing a short list of relevant tools in the prompt is a viable approach. To retrieve relevant tools, various approaches have been suggested, ranging from simple frequency-based matching to dense embedding-based semantic retrieval. However, such approaches lack the contextual and common-sense understanding required to retrieve the right tools for complex user requests. Rather than increasing the complexity of the retrieval component itself, we propose leveraging LLM understanding to generate a retrieval query. Then, the generated query is embedded and used to find the most relevant tools via a nearest-neighbor search. We investigate three approaches for query generation: zero-shot prompting, supervised fine-tuning on tool descriptions, and alignment learning by iteratively optimizing a reward metric measuring retrieval performance. By conducting extensive experiments on a dataset covering complex and multi-tool scenarios, we show that leveraging LLMs for query generation improves the retrieval for in-domain (seen tools) and out-of-domain (unseen tools) settings.

Improving Tool Retrieval by Leveraging Large Language Models for Query Generation

TL;DR

The paper tackles the challenge of retrieving relevant tools/APIs for large-scale LLM tool use. It proposes generating retrieval queries via LLMs and using dense retrieval with interleaving to select relevant APIs, reducing the need to complicate the retriever itself. Through a comparative study of zero-shot prompting, supervised fine-tuning, and alignment learning, the work finds that SFT excels in-domain while alignment learning provides the strongest out-of-domain performance, especially as new tools are introduced. The findings highlight a practical pathway for scalable tool integration, enabling better recall and ranking of appropriate APIs in dynamic tool ecosystems, with strong implications for real-world LLM-assisted workflows.

Abstract

Using tools by Large Language Models (LLMs) is a promising avenue to extend their reach beyond language or conversational settings. The number of tools can scale to thousands as they enable accessing sensory information, fetching updated factual knowledge, or taking actions in the real world. In such settings, in-context learning by providing a short list of relevant tools in the prompt is a viable approach. To retrieve relevant tools, various approaches have been suggested, ranging from simple frequency-based matching to dense embedding-based semantic retrieval. However, such approaches lack the contextual and common-sense understanding required to retrieve the right tools for complex user requests. Rather than increasing the complexity of the retrieval component itself, we propose leveraging LLM understanding to generate a retrieval query. Then, the generated query is embedded and used to find the most relevant tools via a nearest-neighbor search. We investigate three approaches for query generation: zero-shot prompting, supervised fine-tuning on tool descriptions, and alignment learning by iteratively optimizing a reward metric measuring retrieval performance. By conducting extensive experiments on a dataset covering complex and multi-tool scenarios, we show that leveraging LLMs for query generation improves the retrieval for in-domain (seen tools) and out-of-domain (unseen tools) settings.

Paper Structure

This paper contains 24 sections, 3 equations, 4 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Problem Settings
  3. API Retrieval
  4. Retrieval Metrics
  5. LLM-Based Query Generation
  6. Zero-Shot Prompting
  7. Supervised fine-tuning
  8. Alignment Learning
  9. Experiments
  10. Dataset
  11. Retriever Setup
  12. Query Generation Setup
  13. Results
  14. Conclusion
  15. Hyperparameter Settings
  16. ...and 9 more sections

Figures (4)

  • Figure 1: An illustration of leveraging LLMs commonsense and contextual understanding to generate queries for tool retrieval. The steps before and after retrieval are similar to a typical in-context learning setup not shown here.
  • Figure 2: Prompt format used for the tool description generation experiments.
  • Figure 3: Comparison of Recall@5 performance for the zero-shot, SFT, and alignment iterations reported for the in-domain and out-of-domain evaluation sets.
  • Figure 5: Prompt format used for the intent generation experiments.