Tool retrieval (selecting tools for a request)

Deep dive

How it works

Parts:

• Registry — ordered list of tools: name, description, parameter docs, example queries, required permission, token cost, core flag. Order is fixed.
• Index — inverted index over tool text for BM25, and/or embeddings in a vector index.
• Permission filter — drops tools the caller (user, tenant, role, device) may not use, before any scoring.
• Ranker — BM25 and/or dense scores, fused (hybrid rank fusion), optionally reranked (cross-encoder).
• Budgeter — keeps top results until a token budget is spent.
• Core set — tools always exposed (search, ask the user, tool-search itself).
• Assembler — emits the chosen tools in registry order.

Operations:

1. Build the query from the last user message plus task state (a bare "yes, do it" matches nothing).
2. Filter the registry by the caller's permissions. Only allowed tools reach scoring.
3. Score the allowed tools with BM25 and/or embeddings; fuse the lists.
4. Optionally rerank the top 20-50, or let a small model pick from the shortlist.
5. Add the core tools, then candidates in rank order while the token budget holds. Skip one that does not fit and try smaller ones.
6. Emit in registry order. The model call uses only these definitions.
7. Execute: re-check permissions on every call. Selection narrows the menu; it does not authorize.
8. Log query, chosen tools and the tool called, for evaluation.

request ──► query ──► [permission filter] ──► [BM25 + dense] ──► [fuse/rerank] ──► [token budget]
              ▲          registry (all)          index (allowed only)                     │
              │       core tools ─────────────────────────────────────────────────────────┤
    task state                                                          registry order ──► prompt

On-demand variant. Give the model a search-tools meta-tool; it calls it with a query and the matching definitions are loaded for the next step. Anthropic's tool search tool does this (defer_loading: true on the rest; regex or BM25 variant; 5 results by default, up to 10,000 deferred tools). Hierarchical variant. Pick a server or category first (ToolBench groups 16,464 APIs into 49 categories), then tools inside it.

Code

BM25 over tool descriptions with a permission filter before scoring, core tools, a token budget and registry-order output. Std only; tokens are estimated as bytes/4.

use std::collections::HashMap;

pub struct Tool { pub name: String, pub desc: String, pub perm: String, pub schema_len: usize, pub core: bool }

impl Tool {
    /// Rough prompt cost: ~4 bytes per token over name + description + schema.
    fn tokens(&self) -> usize { (self.name.len() + self.desc.len() + self.schema_len) / 4 + 1 }
}

fn tokenize(s: &str) -> Vec<String> {
    s.to_lowercase().split(|c: char| !c.is_alphanumeric()).filter(|w| w.len() > 1)
        .map(|w| w.strip_suffix('s').filter(|x| x.len() > 2).unwrap_or(w).to_string()).collect()
}

pub struct Index { tools: Vec<Tool>, post: HashMap<String, Vec<(u32, u32)>>, len: Vec<f32>, avg: f32 }

impl Index {
    pub fn new(tools: Vec<Tool>) -> Self {
        let (mut post, mut len): (HashMap<String, Vec<(u32, u32)>>, Vec<f32>) = (HashMap::new(), vec![]);
        for (i, t) in tools.iter().enumerate() {
            let words = tokenize(&format!("{} {} {}", t.name, t.name, t.desc)); // name counts double
            let mut tf: HashMap<&str, u32> = HashMap::new();
            for w in &words { *tf.entry(w).or_default() += 1; }
            for (w, n) in tf { post.entry(w.to_string()).or_default().push((i as u32, n)); }
            len.push(words.len() as f32);
        }
        let avg = len.iter().sum::<f32>() / len.len().max(1) as f32;
        Index { tools, post, len, avg }
    }

    /// Indices into the registry, in registry order (stable prefix for the prompt cache).
    pub fn select(&self, query: &str, caller_perms: &[&str], budget: usize) -> Vec<usize> {
        let allowed = |i: usize| { let p = &self.tools[i].perm; p.is_empty() || caller_perms.contains(&p.as_str()) };
        let (k1, b, n) = (1.2f32, 0.75f32, self.tools.len() as f32);
        let mut score = vec![0f32; self.tools.len()];
        for term in tokenize(query) {
            let Some(p) = self.post.get(&term) else { continue };
            let idf = ((n - p.len() as f32 + 0.5) / (p.len() as f32 + 0.5) + 1.0).ln();
            for &(i, tf) in p {
                let i = i as usize;
                if !allowed(i) { continue; }                    // permission filter BEFORE ranking
                let tf = tf as f32;
                score[i] += idf * tf * (k1 + 1.0) / (tf + k1 * (1.0 - b + b * self.len[i] / self.avg));
            }
        }
        let mut used = 0;
        let mut out: Vec<usize> = (0..self.tools.len()).filter(|&i| self.tools[i].core && allowed(i)).collect();
        for &i in &out { used += self.tools[i].tokens(); }
        let mut ranked: Vec<usize> = (0..self.tools.len()).filter(|&i| score[i] > 0.0 && !self.tools[i].core).collect();
        ranked.sort_by(|&a, &b| score[b].total_cmp(&score[a]).then(a.cmp(&b)));
        for i in ranked {
            let t = self.tools[i].tokens();
            if used + t <= budget { used += t; out.push(i); }   // skip what does not fit, keep trying smaller ones
        }
        out.sort_unstable();
        out
    }
    pub fn name(&self, i: usize) -> &str { &self.tools[i].name }
}

fn main() {
    let t = |n: &str, d: &str, p: &str, core| Tool { name: n.into(), desc: d.into(), perm: p.into(), schema_len: 300, core };
    let idx = Index::new(vec![
        t("search_help", "Search the help center articles", "", true),
        t("lookup_account", "Look up a customer account by email or id", "crm.read", false),
        t("order_status", "Get the shipping status of an order", "orders.read", false),
        t("issue_refund", "Refund a customer order payment to the card", "billing.refund", false),
        t("refund_policy", "Explain the refund policy to the customer", "", false),
        t("delete_account", "Delete a customer account permanently", "admin", false),
    ]);
    let names = |v: Vec<usize>| v.into_iter().map(|i| idx.name(i).to_string()).collect::<Vec<_>>();
    let agent = ["crm.read", "orders.read"];
    // best match is forbidden: never returned, the permitted near-match is
    let r = names(idx.select("refund the payment of this order", &agent, 400));
    assert!(!r.contains(&"issue_refund".into()) && r.contains(&"refund_policy".into()));
    assert!(names(idx.select("refund the payment of this order", &["billing.refund"], 400)).contains(&"issue_refund".into()));
    // right tool selected, core tool on, forbidden admin tool absent, registry order, deterministic
    let r = names(idx.select("where is my order shipping status", &agent, 400));
    assert_eq!(r, ["search_help", "order_status"]);
    assert_eq!(r, names(idx.select("where is my order shipping status", &agent, 400)));
    // tiny budget: only the core tool fits
    assert_eq!(names(idx.select("order status", &agent, 100)), ["search_help"]);
}

Compiled and run: all asserts pass.

Complexity and constants

• Index: O(total description terms). Select: O(sum of postings of the query terms) plus a sort of the scored tools.
• Measured here (Apple M4 Pro, release build, 1,000 synthetic tools, 4 rotating queries, 20,000 selections, 3 runs): ~90k selections/s, 11 µs each (one cold run gave ~50k/s). Synthetic descriptions share many words, so real catalogs with more distinct terms have shorter postings.
• Embedding variant: one embedding call per query (ms) plus ANN lookup (HNSW) over tool vectors, computed once per tool.

Efficiency

• Cost: selection is microseconds to milliseconds (measured above for BM25; dense adds an embedding call). The saving is in the prompt: input tokens fall from all definitions to the few chosen ones, per turn, e.g. ~55k down by over 85 % (Anthropic docs, their example). The bound is model input tokens and latency, not the selector.
• Levers:
  • defer loading / search-tool — over 85 % fewer definition tokens (Anthropic docs) — large multi-server catalogs.
  • permission pre-filter — smaller candidate set, no wasted slots on forbidden tools — any multi-tenant or role-based agent.
  • append-only tool list — the cached prefix survives; deferred tools are appended as tool_reference blocks (Anthropic docs) — long runs.
  • descriptions in the user's words — Anthropic advises adding common keywords; no figure — the main lever, all catalogs.
  • tight k / token budget — fewer distractors; tune on recall@k — when accuracy drops with menu size.
• Naive → standard → state of the art:
  • all tools in every prompt · O(tools x schema) tokens per turn, accuracy falls past ~30-50 tools (Anthropic).
  • BM25 or embedding top-k with permission filter and budget · microseconds, 3-5 tools.
  • hybrid retrieval, rerank, on-demand search tool plus learned retriever · extra ms to one model call, fewest misses.

Implementation notes

• Index name, description, parameter docs and a few example queries; weight name and examples up.
• The code takes IDF from the whole registry; with private per-tenant tools, compute statistics per tenant so one tenant's text cannot shift another's ranking.
• Keep k small (3-8) and the budget in tokens, not count: one huge schema can cost more than ten small ones.
• Put core tools first and the rest in registry order, with the cache breakpoint after the last stable tool (tool calling and caching).
• Keep tools already used in the conversation exposed: past calls in history refer to them.

Pitfalls

• Rank then filter. Top-k fills with forbidden tools, then shrinks to nothing, and tool names leak. Filter first.
• Selection is not authorization. Enforce permissions again at execution; high-risk actions still need approval.
• Reordering each turn. Any change before the end of the tools block invalidates the cached prefix (prompt caching).
• Query from the last message only. Multi-step tasks need the plan or recent tool names in the query.
• Near-duplicate tools. Three get_order* tools split scores and confuse the model: merge them or write contrastive descriptions.
• Untrusted descriptions. Third-party MCP descriptions go into the prompt: a poisoned one is an injection (defense).
• Index drift. An unindexed new tool is invisible; reindex on MCP list_changed notifications.

Variants and alternatives

In production

• Support and CRM agents expose only the lookup, order and refund tools the signed-in role allows.
• Coding agents defer MCP tools behind a search step; Anthropic's tool search tool is the hosted form; langgraph-bigtool is the LangGraph one.
• Evaluate with recall@k of the needed tool, end-task success and tokens per turn; ToolRet reports that weak retrieval lowers the pass rate of tool-using LLMs.

Implementations

Sources

• Anthropic, Tool search tool and Tool use with prompt caching, Claude API docs. search · caching
• Model Context Protocol, Tools specification. spec
• Yujia Qin et al., ToolLLM: Facilitating Large Language Models to Master 16000+ Real-world APIs, 2023. paper · code
• Zhengliang Shi et al., Retrieval Models Aren't Tool-Savvy: Benchmarking Tool Retrieval for Large Language Models, 2025. paper
• Shishir G. Patil et al., Gorilla: Large Language Model Connected with Massive APIs, 2023 (retriever-aided API calls). paper · code
• LangChain, langgraph-bigtool. code