TDI300K: Dependency Dataset for Tool Retrieval

• TDI300K is a dataset classifying tool relationships into forward dependency, reverse dependency, and no dependency for effective workflow-sensitive retrieval.
• It is constructed via a two-stage pipeline combining synthetic pretraining with GPT-3.5-based agents and manually annotated finetuning from open-source sources.
• The dataset underpins the Tool Graph Retriever architecture, where dependency predictions form a directed graph that improves retrieval through graph convolution.

TDI300K is a dataset introduced in "Tool Graph Retriever: Exploring Dependency Graph-based Tool Retrieval for LLMs" as the training substrate for a tool-dependency discriminator used in dependency-aware tool retrieval (Gao et al., 7 Aug 2025). Its purpose is to model relations that are typically missed by retrieval methods based only on semantic similarity between tool descriptions and queries. In the formulation used by the paper, TDI300K supports a 3-class classification task over pairs of tool documents, distinguishing forward dependency, reverse dependency, and no dependency. The dataset combines balanced synthetic pretraining data with manually annotated, realistically imbalanced finetuning and test data, and it functions as the principal supervision source for constructing the directed tool dependency graph used by Tool Graph Retriever (TGR) (Gao et al., 7 Aug 2025).

1. Definition and conceptual scope

TDI300K is designed around a specific notion of tool dependency. The paper defines tool dependency as follows: a tool $t_a$ depends on tool $t_b$ if either $t_a$ requires the result from $t_b$ as input, or $t_a$ requires $t_b$ for prior verification (Gao et al., 7 Aug 2025). This definition targets execution structure rather than purely lexical relatedness.

The paper motivates this design by noting that dominant tool-retrieval methods primarily rely on semantic similarities between tool descriptions and user queries, and that such methods often consider each tool independently. According to the paper, this can omit prerequisite tools needed for successful task execution. TDI300K is therefore intended to supervise a model that detects dependencies not directly recoverable from query-document similarity alone (Gao et al., 7 Aug 2025).

Several examples are used to illustrate the dependency concept. The paper explicitly gives UpdateEmail $\rightarrow$ Login, Login $\rightarrow$ Validate, DeleteAccount $\rightarrow$ GetUserToken, and Get Next Predictions $\rightarrow$ Get Today’s Predictions. These examples show that the prerequisite tool may be operationally necessary even when its description is not the strongest semantic match to the query (Gao et al., 7 Aug 2025).

A plausible implication is that TDI300K formalizes retrieval as a workflow-sensitive problem rather than a purely relevance-ranking problem. That interpretation follows from the dataset’s role in training a discriminator whose outputs are later transformed into graph edges between tools.

2. Construction pipeline and corpus composition

TDI300K is built in two stages: pretraining and finetuning. The paper presents this two-stage design as a response to the contrast between synthetically balanced data and real-world imbalanced dependency distributions (Gao et al., 7 Aug 2025).

The pretraining portion is derived from CodeSearchNet, which the paper describes as containing 1.78 million real function implementations across multiple programming languages. Construction proceeds through three steps implemented with three GPT-3.5-turbo-based agents. First, one agent extracts tool documentation from a function implementation. The extracted document is in JSON format and includes function name, description, input parameters, and output parameters. Second, a dependent tool document is synthesized so that it depends on the first tool. Third, the pair is validated against the dependency criteria, and invalid tool pairs are discarded (Gao et al., 7 Aug 2025).

Once a valid dependency instance is obtained, the authors augment the label space in two ways. They swap positions to produce the opposite-direction category, and they create non-dependency examples by “breaking up and shuffling” tool pairs so that they become independent. The pretraining dataset is explicitly balanced across all three classes (Gao et al., 7 Aug 2025).

For finetuning, the paper states that the data are manually constructed from open-source datasets, projects, and libraries. It specifically mentions the training set of ToolBench, projects like online shopping, and libraries like OpenGL. In this stage, tool documents are again written in the same JSON style, tools are organized into tool sets, and dependency labels are manually annotated between tool pairs within each set (Gao et al., 7 Aug 2025).

Validation and test handling are also specified. 20\% of the finetuning data is split off as validation to select the best checkpoint. The test set is manually built from API-Bank tools because ToolBench is described as too large and difficult to annotate exhaustively (Gao et al., 7 Aug 2025).

Category	Pretraining	Finetuning
$t_b$0	92,000	1,029
$t_b$1	92,000	33,365
$t_b$2	92,000	1,056

From the counts reported in the paper, the pretraining split contributes 276,000 instances, the finetuning split 35,450, and the test split 620. The paper states that the name TDI300K reflects the fact that the dataset is on the order of 300K examples overall (Gao et al., 7 Aug 2025).

3. Label space, input format, and optimization objective

The dependency-discrimination problem supported by TDI300K is framed as a 3-class classification task over a pair of tool documents. The three labels are $t_b$3, $t_b$4, and $t_b$5, corresponding respectively to forward dependency, no dependency, and reverse dependency (Gao et al., 7 Aug 2025).

Given a pair $t_b$6, the input is formed by concatenating the two documents with special tokens:

$t_b$7

The final hidden state of $t_b$8 is used for classification (Gao et al., 7 Aug 2025).

For the base training objective, the paper uses standard cross-entropy loss:

$t_b$9

where $t_a$0 is the gold one-hot label distribution (Gao et al., 7 Aug 2025).

Because real dependency categories are imbalanced, the finetuning stage uses a category-specific average loss:

$t_a$1

where $t_a$2 is the number of instances in the $t_a$3-th dependency category (Gao et al., 7 Aug 2025). The paper states that this is meant to reduce bias toward the majority no-dependency class.

The test set is intentionally imbalanced and contains exactly 60 samples for $t_a$4, 500 samples for $t_a$5, and 60 samples for $t_a$6. This design is presented as more realistic than a balanced evaluation regime (Gao et al., 7 Aug 2025).

4. Function within Tool Graph Retriever

TDI300K is not an isolated benchmark; it is the enabling dataset for the dependency-aware retrieval architecture called Tool Graph Retriever (TGR). The dataset trains the discriminator that predicts whether one tool depends on another. Once trained, that discriminator is applied to a candidate tool set to build a directed dependency graph

$t_a$7

where $t_a$8 denotes tools as nodes and $t_a$9 denotes directed edges (Gao et al., 7 Aug 2025).

If tool $t_b$0 depends on tool $t_b$1, the graph contains a directed edge linking the node for $t_b$2 to the node for $t_b$3. The paper’s example graph includes edges such as Login $t_b$4 Validate and UpdateEmail $t_b$5 Login (Gao et al., 7 Aug 2025).

The graph is then processed with graph convolution. The paper gives the following expression:

$t_b$6

where $t_b$7 is the tool embedding matrix, $t_b$8 is the adjacency matrix, $t_b$9 is the degree matrix, and $t_a$0 is the identity matrix (Gao et al., 7 Aug 2025). The stated intent is to propagate dependency information into tool embeddings before retrieval.

After graph encoding, query-tool similarity is computed between the query embedding and the updated tool embeddings, and the top-$t_a$1 tools are returned. The paper reports that cosine similarity is the best-performing similarity function in an appendix evaluation (Gao et al., 7 Aug 2025).

This architecture gives TDI300K a dual role. At the immediate level, it is a supervised classification dataset. At the system level, it is the source of edge predictions that determine