Tool Graph Retriever: Exploring Dependency Graph-based Tool Retrieval for Large Language Models

📝 Paper Summary

Multi-call tool use with flexible plan Retrieval-Augmented Generation

ToolGraphRetriever improves tool retrieval by modeling prerequisites as a dependency graph and propagating embeddings via graph convolution, ensuring necessary but semantically distinct tools are retrieved.

Core Problem

Existing tool retrievers rely on semantic similarity between queries and tool descriptions, failing to retrieve prerequisite tools that are necessary for execution but semantically unrelated to the user's specific request.

Why it matters:

Complex tasks often require a sequence of tools (e.g., validation before action), where missing the first step causes total execution failure
LLM context windows are finite, necessitating accurate retrieval rather than ingesting all available tool documentation
Current semantic retrievers treat tools as independent items, ignoring the structural logic of API dependencies inherent in software systems

Concrete Example: For the query 'Update my email to new@domain.com', a semantic retriever finds 'UpdateInfo' but misses 'Login' and 'Validate' because 'Login' is not semantically similar to 'Update email'. Without 'Login', the 'UpdateInfo' tool cannot be invoked, causing the agent to fail.

Key Novelty

ToolGraphRetriever (TGR)

Constructs a directed graph where nodes are tools and edges represent functional dependencies (e.g., Tool A requires output from Tool B)
Uses a trained BERT-based discriminator to automatically identify these dependencies from tool documentation
Applies a graph convolution operation to update tool embeddings, mixing information from prerequisite tools into dependent tools to aid retrieval

Architecture

The three-step process of Tool Graph Retriever: Dependency Identification, Graph-Based Tool Encoding, and Online Retrieval.

Evaluation Highlights

+0.065 Recall@5 improvement on API-Bank when adding TGR to the ToolBench-IR baseline (0.697 -> 0.762)
+0.062 NDCG@5 improvement on API-Bank when adding TGR to ToolBench-IR (0.722 -> 0.784)
+0.077 Pass Rate@5 improvement on API-Bank when adding TGR to ToolBench-IR (0.463 -> 0.540)

Breakthrough Assessment

6/10

Offers a logical and effective solution (dependency graphs) to a specific flaw in semantic retrieval (missing prerequisites). While the method is sound and results are positive, it relies on standard components (BERT, GCN) and the improvement on ToolBench is more modest than on API-Bank.

⚙️ Technical Details

Problem Definition

Setting: Retrieving a subset of k tools from a large candidate set T based on a user query q

Inputs: User query q, Set of candidate tools T (with descriptions)

Outputs: Ranked list of k tools

Pipeline Flow

Preprocessing: Dependency Identification (BERT Discriminator)
Preprocessing: Graph Construction
Inference: Graph-Based Tool Encoding (GCN Propagation)
Inference: Online Retrieval (Similarity Search)

System Modules

Dependency Discriminator

Classify whether a dependency exists between two tools (A depends on B, B depends on A, or Independent)

Model or implementation: BERT-base-uncased

Graph Encoder (Inference)

Update tool embeddings by aggregating information from connected neighbors (dependencies)

Model or implementation: Graph Convolution (non-trainable propagation)

Similarity Scorer (Inference)

Rank tools by similarity to the user query

Model or implementation: Dot product / Cosine similarity

Novel Architectural Elements

Integration of a pre-computed dependency graph into the embedding space via a parameter-free graph convolution step prior to similarity search

Modeling

Base Model: BERT-base-uncased (for the discriminator)

Training Method: Supervised Classification (for Discriminator)

Objective Functions:

Purpose: Train discriminator to predict dependency class.

Formally: Cross-entropy loss L = -sum(y_k * log(y_hat_k))
Purpose: Handle class imbalance during finetuning.

Formally: Category-specific average training loss L = sum(1/N_k * sum(L_i))

Training Data:

Pretraining: TDI300K (synthetic, derived from CodeSearchNet via GPT-3.5)
Finetuning: Manually annotated dataset from open-source projects
Testing: Derived from API-Bank tools (manually annotated)

Key Hyperparameters:

base_model: BERT-base-uncased

Compute: Not reported in the paper

Comparison to Prior Work

vs. ToolBench-IR: TGR adds a graph propagation step to ToolBench-IR embeddings to capture dependencies
vs. Standard Dense Retrieval (MiniLM, Ada-002): TGR explicitly models tool-tool relations, whereas standard retrieval treats tools as independent documents

Limitations

Relies on the quality of the dependency graph; errors in the discriminator propagate to retrieval
Graph construction via discriminator is O(N^2) complexity for all tool pairs if not optimized (though paper notes only N tools in graph)
Experiments limited to ToolBench I1 level due to time complexity of constructing graphs for larger sets
Requires manual annotation for high-quality finetuning data for the discriminator

Reproducibility

The paper mentions constructing a dataset 'TDI300K', but no URL is provided in the text. The method relies on a trained discriminator and a graph construction step. Code availability is not explicitly stated.

📊 Experiments & Results

Evaluation Setup

Retrieve top-k tools given a query from API-Bank and ToolBench datasets

Benchmarks:

API-Bank (Tool Retrieval (Level 1-3))
ToolBench (Tool Retrieval (Instruction Level I1))

Metrics:

Recall@K
NDCG@K
Pass Rate@K (percentage of queries where ALL required tools are retrieved)
Statistical methodology: Not explicitly reported in the paper

Key Results

Benchmark	Metric	Baseline	This Paper	Δ
Main results on API-Bank showing TGR improves upon base embeddings (MiniLM and ToolBench-IR).
API-Bank	Recall@5	0.697	0.762	+0.065
API-Bank	NDCG@5	0.722	0.784	+0.062
API-Bank	Pass Rate@5	0.463	0.540	+0.077
Main results on ToolBench (I1 level) showing consistent but smaller gains compared to API-Bank.
ToolBench	Recall@5	0.548	0.563	+0.015
ToolBench	Pass Rate@5	0.260	0.337	+0.077
Ablation study comparing discriminator-based graph (+TGR-d) vs manually annotated graph (+TGR-m) on API-Bank.
API-Bank	Recall@5	0.650	0.672	+0.022

Experiment Figures

The recall increment of TGR over the baseline relative to the proportion of connected nodes (graph density) in tool categories.

Main Takeaways

Incorporating tool dependencies significantly improves Pass Rate, which is critical for multi-step tasks where missing one tool causes failure.
The method is compatible with various embedding backends (MiniLM, ToolBench-IR), serving as a plug-and-play enhancement module.
Performance gains are correlated with graph density; tools with denser dependency networks benefit more from the graph propagation.
Manual graphs outperform discriminator-generated graphs, highlighting that accurate dependency identification is the current bottleneck.

📚 Prerequisite Knowledge

Prerequisites

Understanding of dense retrieval (embedding-based search)
Basic Graph Neural Networks (specifically Graph Convolutional Networks)
Knowledge of API/Tool use in LLM agents

Key Terms

GCN: Graph Convolutional Network—a neural network architecture that operates on graphs, allowing nodes to aggregate information from their neighbors

NDCG: Normalized Discounted Cumulative Gain—a measure of ranking quality that considers the position of relevant items in the list

Pass Rate: The proportion of test queries where *all* ground-truth required tools are successfully retrieved within the top-k results

ToolBench-IR: A specific BERT-based dense retriever fine-tuned on the ToolBench dataset for retrieving APIs

TDI300K: Tool Dependency Identification 300K—a dataset constructed by the authors containing pairs of tools and their dependency labels

BM25: A probabilistic information retrieval function that ranks documents based on the query terms appearing in each document

Graph Convolution: An operation where a node's feature vector is updated by aggregating feature vectors from its connected neighbors in the graph