Tools are under-documented: Simple Document Expansion Boosts Tool Retrieval

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

Document expansion significantly boosts tool retrieval for LLM agents

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Tool documentation is incomplete and inconsistent, creating a semantic gap that bottlenecks tool retrieval in LLM agent systems. This paper introduces TOOL-REX, a benchmark and pipeline that systematically expands raw tool docs with structured fields (function description, when-to-use, limitations, tags), then trains dedicated retriever (Tool-Embed) and reranker (Tool-Rank) models on the enriched corpus. Document expansion alone boosts existing models, and the specialized models achieve new state-of-the-art on both TOOL-REX and the prior ToolRet benchmark.

Problem

Tool retrieval for LLM agents suffers from a persistent semantic gap between natural-language user queries and formal API documentation. Existing tool docs are heterogeneous (at least 7 distinct phrasings for equivalent functions in ToolRet), incomplete (missing description fields, no usage scenarios, no limitations), and unstandardized across 35+ datasets. Prior work like Re-Invoke and EASYTOOL works around poor documentation rather than fixing it; there is no large-scale training corpus tailored to tool retrieval, creating a performance ceiling for both retrievers and rerankers.

Method

Document expansion pipeline (4 stages): 1. Expansion: Qwen3-32B (reasoning mode) generates structured profiles \(d_{\text{profile}}\) from raw documentation \(d_{\text{original}}\), adding fields: function_description, tags, when_to_use, limitations, example_usage. The expansion is query-independent and strictly grounded—unsupported fields are omitted. The expanded document is \(d_{\text{expansion}} = d_{\text{original}} \cup d_{\text{profile}}\). 2. Judgement: LLaMA-3.1-70B verifies field faithfulness after rule-based JSON/format checks; ~1.5% of cases fail. 3. Refinement: GPT-4o regenerates failed cases (~600 examples) using the same prompt. 4. Human validation: 100 randomly sampled profiles manually verified for faithfulness.

An ablation study (Add-One / One-Out protocols with BM25 and GritLM) finds example_usage provides minimal or negative gains; it is excluded from final profiles.

Training corpora: 50k instances for embedding retrievers and 200k for rerankers, drawn exclusively from Web-category tools (APIGen + ToolBench), enabling OOD evaluation on Code and Customized categories.

Tool-Embed: Fine-tunes Qwen3-Embedding-0.6B/4B with InfoNCE loss, 5 in-batch negatives per query, 1 epoch, full-parameter tuning via DeepSpeed ZeRO-3.

Tool-Rank: Fine-tunes Qwen3-Reranker-4B with LoRA (\(r=32\), \(\alpha=64\), dropout 0.1) and cross-entropy loss. Relevance score: \(P(\text{relevant}|q,d) = \frac{\exp(\ell_{\text{true}})}{\exp(\ell_{\text{true}}) + \exp(\ell_{\text{false}})}\).

Key Contributions

• Scalable, low-cost LLM-based document expansion pipeline that enriches tool docs with structured semantic fields.
• TOOL-REX benchmark: 7,615 tasks over 43,215 tools across Web/Code/Customized domains, built on 35 datasets from ToolRet.
• Two large-scale training corpora: Tool-Embed-Train (50k) and Tool-Rank-Train (200k).
• Tool-Embed (dense retriever) and Tool-Rank (LLM reranker) setting new SoTA on both TOOL-REX and ToolRet.
• Systematic ablation quantifying per-field contribution to retrieval performance.

Results

On TOOL-REX (expanded benchmark), Tool-Embed-4B vs. best prior open-source (Qwen3-Embedding-8B): - NDCG@10: 52.23 vs. 46.23 (+6.00) - Recall@10: 63.13 vs. 56.83 (+6.30) - Completeness@10: 51.61 vs. 46.70 (+4.91)

Tool-Rank-4B (reranking Tool-Embed-4B top-100) on TOOL-REX: - NDCG@10: 56.44 (+4.21 over first-stage retriever) - Recall@10: 67.81 (+4.68) - Completeness@10: 56.60 (+4.99) - Improvement over MTEB SoTA open-source (Qwen3-Embedding-8B): +10.23 N@10, +10.29 R@10, +9.08 C@10

Impact of expansion on training (vs. non-expansion counterparts): - Tool-Embed-0.6B/4B: +4.97/+6.69 N@10 vs. baseline Qwen3-Embedding-0.6B; non-expansion trained models gain only +3.69/+3.67 - Tool-Rank: +4.21 N@10 gain over Tool-Embed-4B; Tool-Rankoriginal gains only +2.00

OOD generalization: Both models outperform all baselines on Code and Customized categories despite being trained exclusively on Web-category tools.

Document expansion on existing models (ToolRet→TOOL-REX): Most dense retrievers improve; GritLM gains +2.41 N@10; BM25 improves consistently; E5-Mistral is the sole exception (slight decrease).

Limitations

• example_usage field, though generated, degrades retrieval performance in some settings—suggesting LLM-generated augmentations are not uniformly beneficial and field selection must be empirically validated.
• Training data is restricted to Web-category APIs; Code and Customized categories are evaluated zero-shot, which, while demonstrating generalization, leaves domain-specific performance gains unexplored.
• The expansion pipeline relies on closed-source GPT-4o for the ~1.5% refinement cases, introducing a dependency on proprietary models for full-quality data production.
• Longer expanded documents risk truncation under fixed token budgets of retrieval models, potentially diluting salient signals (though the average expansion adds ~46 tokens per document).
• Human validation covers only 100 sampled profiles; hallucination rates at scale are estimated but not exhaustively verified.

Relevance to Agentic AI / LLM Agents

Tool retrieval is a foundational bottleneck for tool-augmented LLM agents operating over large API ecosystems—when agents cannot reliably surface the right tool for a user intent, downstream execution fails regardless of planning quality. This work directly addresses a data-quality root cause (poor documentation) rather than patching retrieval models, providing a reusable pipeline applicable to any tool corpus an agent might use. The TOOL-REX benchmark and training corpora fill a critical infrastructure gap for researchers building scalable agent tool-selection systems, and the OOD generalization results suggest the approach can handle the open-ended, evolving tool libraries characteristic of real deployments. The finding that when_to_use and function_description fields drive most retrieval gains has direct implications for how developers should structure tool documentation in agent frameworks (e.g., MCP servers, plugin registries).