GRAFT: Graph-Tokenized LLMs for Tool Planning

TL;DR

GRAFT is a novel framework that internalizes tool dependency graphs within language models to improve the accuracy and reliability of multi-step tool planning in complex workflows.

Contribution

It introduces a graph-tokenized approach and on-policy distillation to better align tool plans with subtask structures, surpassing existing external graph-use methods.

Findings

GRAFT achieves state-of-the-art exact sequence matching.

It significantly improves dependency legality in tool planning.

The method enhances reliability in complex LLM workflows.

Abstract

Large language models (LLMs) are increasingly used to complete complex tasks by selecting and coordinating external tools across multiple steps. This requires aligning tool choices with subtask intent while satisfying directional execution dependencies among tools. To do this, existing methods model these dependencies as tool graphs and incorporate the graphs with LLMs through retrieval, serialization, or prompt-level injection. However, these external graph-use strategies all follow a matching paradigm, which often fails to align tool choices with the underlying subtask structure, producing semantically plausible plans that violate graph constraints. This issue is further exacerbated by error accumulation, where an early incorrect tool selection shifts the plan into an invalid graph state and causes subsequent predictions to drift away from the valid execution path. To address these challenges, we propose GRAFT, a graph-tokenized language model framework for dependency-aware tool planning. GRAFT internalizes the tool graph by mapping each tool node to a dedicated special token and learning directed tool dependencies within the representation space. It further introduces on-policy tool context distillation, training the model on its own sampled trajectories while distilling stepwise planning signals. Experiments show that GRAFT achieves state-of-the-art performance in exact sequence matching and dependency legality, supporting more reliable LLM tool planning in complex workflows.