Retrieval-Augmented Tool Selection for Long-Horizon Language Agents

1.Introduction

Language models can now call external tools such as search engines, code interpreters, calendars and databases [1, 2]. In practice the number of available tools grows quickly: an enterprise assistant may expose thousands of functions through an API gateway. The common approach of placing every tool description in the prompt does not scale, since long prompts raise latency and cost and, as we show, reduce the accuracy with which the right tool is chosen.

We therefore ask whether tool selection can be treated as retrieval, as is done for documents in retrieval-augmented generation [3, 4]. Two properties of agentic tasks make this harder than document retrieval. First, the information need changes as the agent proceeds: the tool needed at step 12 is rarely implied by the initial instruction. Second, tools are used, not read, so a near-miss is costly.

Our contributions are as follows:

• a formulation of tool selection as retrieval conditioned on the evolving task state;
• a two-stage selector combining dense retrieval with a constrained model re-rank;
• an evaluation harness with 4,000 synthetic tools and three task families, released with this article.

2.Related work

Tool-using language models have been trained with supervised signals [1] or prompted to interleave reasoning and acting [2]. These approaches assume that a small, fixed set of tools is available. Dense retrieval [4] has been used to select passages; we adapt the idea to tool descriptions and, importantly, to a query that changes during an episode. Our study is related to benchmark efforts that evaluate agents across environments [5], but we isolate the selection step.

3.Method

Let T = {t₁, …, tₙ} be the tool set, where each tool tⱼ has a natural-language description dⱼ. At step k the agent has a task state sₖ consisting of the instruction, a short summary of previous actions and the most recent observation.

Stage 1: state-conditioned retrieval

We embed the state and each description with the same encoder and score tools by cosine similarity. The shortlist Cₖ contains the top m tools:

Stage 2: constrained re-ranking

The model receives only the shortlist and must output the identifier of one tool or the token NONE. Because the output space is restricted to |Cₖ|+1 symbols, we decode with a constrained grammar, which removes invalid tool names. The full procedure is summarised below.

def select_tool(state, tools, encoder, model, m=8):
    q = encoder(state.summary())
    shortlist = top_k(tools, key=lambda t: cosine(q, t.embedding), k=m)
    return model.choose(state, options=shortlist + [NONE], constrained=True)

The training objective for the encoder is a contrastive loss over (state, correct tool) pairs, with the other shortlisted tools as hard negatives [4].

4.Experiments

Setup

We generate 4,000 tools with natural-language descriptions, grouped into 120 families with deliberately overlapping descriptions, and three task families (data analysis, scheduling and document workflows) with 8 to 40 steps each. Each configuration is run with five seeds. We compare against (a) all tools in the prompt, (b) instruction-only retrieval, and (c) an oracle that is shown only the correct tools.

Results

Table 1 reports task success. State-conditioned selection is within 2.1 points of the oracle, and substantially better than instruction-only retrieval. Figure 1 shows that the advantage grows with the number of available tools: with all tools in the prompt, success falls steadily, whereas the selector stays nearly flat.

Ablations

Removing the re-ranking stage lowers success by 3.4 points, and shrinking the shortlist from 8 to 3 candidates lowers it by 5.2 points, mostly through missed retrievals. Using the instruction alone as the query reduces success by 8.9 points on the longest tasks.

5.Discussion and limitations

The dominant failure mode is confusion between tools in the same family: when two descriptions differ only in a parameter, both retrieval and re-ranking select the wrong one in 31% of failures. Better descriptions, or retrieval over usage examples, are plausible remedies. Our tools are synthetic, and real APIs contain noisier documentation; the results should be read as controlled evidence about the selection mechanism, not as a deployment estimate.

6.Conclusion

Treating tool selection as state-conditioned retrieval lets agents work with very large tool sets at a fraction of the prompt cost. We release the harness so that selection can be evaluated in isolation from the rest of the agent.