Tool retrieval is a critical bottleneck for LLM-based agents operating over large, heterogeneous API ecosystems. Existing approaches face an inherent trade-off: semantic retrievers are fast but suffer from the semantic-functional gap, while execution-based validation improves precision at the cost of prohibitive latency. We propose Lookahead-R, a planning-based framework that reformulates tool retrieval as a resource-constrained sequential decision-making problem. At its core, Lookahead-R introduces a lightweight execution-aware surrogate world model that jointly predicts tool execution success, latency cost, and semantic utility---without invoking real APIs. This world model drives a cost-sensitive, uncertainty-guided Monte Carlo Tree Search that navigates the tool space under strict budget constraints. Evaluated on the large-scale ToolBench benchmark, Lookahead-R achieves a superior accuracy-efficiency trade-off across all test scenarios. On the most challenging I3 split, it attains an NDCG@5 of 91.40%, outperforming the state-of-the-art ToolGen (90.16%) by 1.24%. Ablation studies confirm that explicit latency modeling is the key discriminative signal for identifying high-quality tools under resource constraints.
Figures & tables
Figure 1. Comparison between previous paradigms and Lookahead-R. We illustrate the workflow in two scenarios: a tool retrieval task, where the model retrieves the correct tool for a given query, and an LLM-based agent task, where the model completes complex tasks involving real-world API calls. Unlike methods relying on static similarity matching, Lookahead-R reformulates retrieval as a planning problem. A learned world model enables virtual simulation and uncertainty-aware decision-making to select tools based on execution feasibility rather than just semantics.
Figure 2. Overview of Lookahead-R. (A) Data Synthesis: A Student-Teacher protocol constructs execution traces (latency, errors) to bridge data gaps. (B) Model Training: A Llama-3 backbone is fine-tuned as a Predictive World Model. (C) MCTS Inference: The planner utilizes structured JSON predictions to perform virtual rollouts, aggressively pruning high-cost branches (e.g., timeouts) to ensure budget adherence.
Method
Backbone
I1: Single-Tool
I2: Intra-Category
I3: Intra-Collection
Overall Performance
@1
@3
@5
@1
@3
@5
@1
@3
@5
Avg. PR
Win Rate
Sparse & Dense Retrievers
BM25
-
29.46
31.12
33.27
24.13
25.29
27.65
32.00
25.88
29.78
30.5
25.4
EmbSim (OpenAI)
Ada-002
63.67
61.03
65.37
49.11
42.27
46.56
53.00
46.40
52.73
51.2
40.8
ToolRetriever ( Qin et al., 2024 )
RoBERTa
80.50
79.55
84.39
71.18
64.81
70.35
70.00
60.44
64.70
62.5
53.0
Re-Invoke ( Chen et al., 2024b )
LLaMA-2
69.47
-
61.10
54.56
-
53.79
59.65
-
59.55
58.0
48.5
Table 1. Main Results on ToolBench. We report the retrieval performance using NDCG@k (k=1,3,5) across three instruction complexity levels: I1 (Single-tool), I2 (Intra-category Multi-tool), and I3 (Intra-collection Multi-tool). The last columns report the Average Pass Rate (PR) and Win Rate compared to the standard ChatGPT-ReAct baseline. Bold indicates the best performance, and underlined indicates the second best. Lookahead-R demonstrates superior ranking capability, especially in complex I3 scenarios, significantly outperforming advanced iterative and generative retrieval baselines.
Figure 3. Anytime Performance Profile. Lookahead-R progressively refines tool selection accuracy as the latency budget increases.
Variant
Success Rate ↑
Failure Rate ↓
Avg. Latency ↓
Diversity ↑
Lookahead-R (Full)
78.5%
4.1%
680 ms
High
w/o Exec. Prediction
61.5%
28.3%
1,640 ms
Med
Cost-Shuffled
75.2%
5.8%
1,120 ms
High
Entropy-only Expl.
69.2%
12.4%
920 ms
Low
Table 2. Mechanism-level ablation results on I1 dataset (N=1,000). Latency is measured using ground-truth execution logs. The high Success Rate of Cost-Shuffled confirms it still identifies valid tools, but fails to optimize for efficiency.
Stage
Lookahead-R Execution Trace
1. Query
[Image: Movie Poster] + “Who is the director?”
2. Recall
Retriever (System 1) finds candidates by semantics: 1. Search_Director_Text (Rank 1, High Text Similarity) 2. Visual_Movie_Lookup (Rank 2, Multi-modal)
3. Mental Sim.
Generative World Model (System 2) rollouts: ∙ Search_Director_Text : ↪ Gen. Resp: "Error: Missing arg ’movie_name’. Cannot process image." ↪ Metrics: Risk r^=0.95→ Penalize ( Q↓ ) ∙ Visual_Movie_Lookup : ↪ Gen. Resp: "[Success] Recognized: Dune (2021). Director: D. Villeneuve" ↪ Metrics: Utility v^=0.9 , Risk r^=0.1→ Boost ( Q↑ )
School of Mathematical Sciences, Shanghai Jiao Tong University · Department of Data Science and Artificial Intelligence, The Hong Kong Polytechnic University