Planning Agents on an Ego-Trip: Leveraging Hybrid Ego-Graph Ensembles for Improved Tool Retrieval in Enterprise Task Planning

TL;DR

This paper introduces a KG-based tool retrieval framework using ego-graph ensembles to improve multi-step tool selection accuracy for AI planning agents, outperforming traditional similarity-based methods.

Contribution

It proposes a novel knowledge graph approach with ego-graph ensembles for enhanced tool retrieval in complex, multi-step user queries, addressing limitations of existing similarity-based methods.

Findings

Achieves 91.85% tool coverage on the CompleteRecall metric.

Outperforms the strongest non-KG baseline with 89.26% coverage.

Demonstrates the effectiveness of structural graph information in tool retrieval.

Abstract

Effective tool pre-selection via retrieval is essential for AI agents to select from a vast array of tools when identifying and planning actions in the context of complex user queries. Despite its central role in planning, this aspect remains underexplored in the literature. Traditional approaches rely primarily on similarities between user queries and tool descriptions, which significantly limits retrieval accuracy, specifically when handling multi-step user requests. To address these limitations, we propose a Knowledge Graph (KG)-based tool retrieval framework that captures the semantic relationships between tools and their functional dependencies. Our retrieval algorithm leverages ensembles of 1-hop ego tool graphs to model direct and indirect connections between tools, enabling more comprehensive and contextual tool selection for multi-step tasks. We evaluate our approach on a synthetically generated internal dataset across six defined user classes, extending previous work on coherent dialogue synthesis and tool retrieval benchmarks. Results demonstrate that our tool graph-based method achieves 91.85% tool coverage on the micro-average CompleteRecall metric, compared to 89.26% for re-ranked semantic-lexical hybrid retrieval, the strongest non-KG baseline in our experiments. These findings support our hypothesis that the structural information modeled in the graph provides complementary signals to pure similarity matching, particularly for queries requiring sequential tool composition.