AI agents increasingly rely on external tools for task completion. When retrieval selects an unavailable or broken tool, the agent fails silently or loops unnecessarily. Existing methods prioritize semantic similarity but treat tool availability as a static property, missing real-time failures.
Why Tool Retrieval Fails Without Runtime Awareness
Traditional tool retrievers operate on text embeddings alone. They match a user query to tool descriptions without considering whether the tool's underlying API is reachable, rate-limited, or actively returning errors. This gap means a highly relevant tool can still cause agent failure if it is offline. Prior work has explored agent benchmarks and tool selection, but few integrate runtime feedback into the retrieval step itself.
Frozen Dual-Encoder with Contrastive Fine-Tuning
RATED introduces two components. The first is a frozen dual-encoder architecture with contrastive fine-tuning that produces robust semantic matches. The encoder stays fixed after fine-tuning, which keeps inference cheap and stable. This design allows the system to match queries to tool descriptions using established similarity metrics without recomputing embeddings when the bandit module adjusts rankings.
Contextual Bandit Re-Ranking via Execution Dynamics
The second component is a lightweight contextual bandit module that re-ranks the top candidates by observing execution outcomes in real time. When a tool call returns an error, the bandit lowers that tool's selection probability and surfaces alternatives. This re-ranking happens in under 1 ms because it operates on already-computed embeddings and a simple bandit update rule, not on model recomputation.
EcoAgent: 5,172 Tools and 210,000 Query-Tool Pairs
The authors release EcoAgent, a benchmark comprising 5,172 real-world tools and over 210,000 query-tool pairs. The tools span common agent use cases such as file system operations, web searches, and data formatting. Each query-tool pair carries labels indicating whether the tool succeeded or failed during execution, enabling offline retrieval metrics as well as dynamic reliability evaluation under simulated API outages.
MRR Improvement and Fault Selection Under Controlled API Outage
On real-world queries, RATED achieves a 52.2% relative MRR improvement over the ToolLLM API retriever. Under a 30% simulated API failure rate, fault selection drops from 86.2% with the baseline to 2.9% with RATED. Latency impact is negligible: less than 1 ms added to encode-only inference, confirming the framework's efficiency for production environments.
Limitations in Sustained Execution Environments
The bandit re-ranking depends on observed execution feedback, which means it requires at least one attempt per query before adaptation takes effect. In scenarios where the very first tool call must succeed, the initial semantic selection still matters. Also, the evaluation uses simulated outages; real-world network partitions may introduce variables not captured in the benchmark.
Practical Integration for Agent Workflows
Developers building AI agents can integrate RATED with minimal changes. The frozen encoder means no additional training cost, and the bandit update rule can be swapped or calibrated per deployment. The 52% MRR gain and near-elimination of fault selection under outage conditions suggest that reliability-aware retrieval becomes a practical necessity for agentic systems.
As agents expand into more tool-heavy workflows, the cost of missed or broken selections grows proportionally. RATED demonstrates that a lightweight, reliability-aware layer can dramatically improve robustness without sacrificing speed or requiring full model retraining. Future work could explore multi-step bandit policies or integration with broader agent planning modules. Read the paper on arXiv