Risk-Aware Reranking for Agentic Tool Retrieval
2026-08-24 • Information Retrieval
Information Retrieval
AI summaryⓘ
The authors study how to choose external tools for AI language models in a way that balances usefulness and safety, since some tools might be risky to use in certain situations. They propose a method that re-ranks tools based on both how relevant they are to a query and how risky they are, allowing control over this tradeoff. They also created a risk rating for over 6,000 tools and new ways to measure how risky the selected tools are. Their experiments show improved safety without losing much usefulness, especially when combining ranking with rule-based filters. This helps prevent exposing AI models to dangerous tools before they use them.
tool retrievallanguage modelsrisk-aware retrievalrerankingsafety in AItool exposure riskquery relevanceToolGraphrule-based filtersUltraTool
Authors
Qinfei Li, Xiaoxuan Dong, Jin Zhang, Dexu Yu, Wenhao Deng, Junchen Fu, Youhua Li, Hanwen Du, Chunxiao Li
Abstract
Tool retrieval determines which external tools are exposed to an LLM agent for a user query or task, making retrieval a critical pre-execution safety boundary. Unlike document retrieval, tool retrieval exposes executable actions: a tool that is useful for one task may be unnecessary or risky for another. However, existing tool-retrieval methods primarily optimize semantic relevance, and safety evaluations often focus on failures after tool execution rather than risks introduced during retrieval. We study risk-aware tool retrieval, where the goal is to retrieve useful tools while reducing exposure to higher-risk tools. We propose a lightweight reranking framework on top of a frozen first-stage retriever. The framework models query-conditioned relevance and tool-level exposure risk separately, combines them through an explicit parameter controlling the tradeoff between safety and utility, smooths scores over a ToolGraph, and optionally applies rule-based safety constraints. To support retrieval-time safety evaluation, we annotate 6,108 tools across UltraTool and Seal-Tools with five ordinal risk levels and define metrics that measure risky-tool exposure in the top-$k$ results. Experiments on UltraTool and Seal-Tools show that our approach improves the relevance--safety tradeoff over relevance-only retrievers and reranking baselines, with the rule-filtered variant providing a conservative operating point for safety-critical deployments. These findings indicate that retrieval-stage filtering can reduce the candidate action space exposed to agents before execution, complementing downstream tool-use safeguards. The code and supplementary materials are available at: https://github.com/qli447/risk-aware-tool-retrieval-release.