toolLoop

fun <State> StateGraph<State>.toolLoop(model: ChatModel, tools: List<Tool> = emptyList(), messages: (State) -> List<ChatMessage>, append: (State, List<ChatMessage>) -> State, firstMessage: (State) -> String? = null, system: String? = null, then: NodeRef = NodeRef.END, modelNode: String = "model", toolsNode: String = "tools", stateTools: List<StateTool<State>> = emptyList(), inOrder: Boolean = false, countTokens: (State, TokenUsage) -> State? = null, stopWhen: (State) -> Boolean? = null, retry: Retry? = null): NodeRef

Adds the loop of a tool-calling agent to this graph and returns its first node.

The loop has two nodes. The model node sends the conversation to model. When the answer asks for tools, the tools node runs them, at the same time unless inOrder or stateTools say otherwise, and the model is asked again with their results. When the answer asks for none, the graph continues at then.

data class Ticket(val messages: List<ChatMessage>, val reply: String = "")

val graph = StateGraph<Ticket> {
val send = node("send") { it.copy(reply = it.messages.last().text) }
val agent = toolLoop(
model = model,
tools = listOf(orderStatus, menuPrice),
messages = { it.messages },
append = { ticket, new -> ticket.copy(messages = ticket.messages + new) },
then = send,
)
START then agent
send then END
}.compile()

graph.invoke(Ticket(listOf(ChatMessage.User("Where is my pizza?"))))

Both nodes have a work and an update, so the loop can run next to other branches without a Reducer; a loop with stateTools is the exception, see below. Every round of tools takes two steps of the run, so raise GraphConfig.maxIterations for an agent that needs many rounds.

messages and append read and write the same list: messages must return what append stored. When the conversation starts from other fields of the state, do not build the first message in messages, because an append that adds to the stored list would then lose it. Give it to firstMessage instead, and the loop stores it with the model's first answer:

data class Order(val customer: String, val question: String, val messages: List<ChatMessage> = emptyList())

toolLoop(
model = model,
tools = tools,
messages = { it.messages },
append = { order, new -> order.copy(messages = order.messages + new) },
firstMessage = { "${it.customer} writes: ${it.question}" },
)

When the run is collected with stream, the model node streams: each piece of text the model writes arrives as a GraphEvent.NodeProgress of the model node, and textDelta reads it. With invoke, the model is asked for the whole answer at once.

When the model reaches its output limit in the text of its answer, the loop ends with that answer and its ChatMessage.Assistant.truncated is true, so append and the code that reads the state can tell. When it reaches the limit in a tool call, the model node fails with a ChatModelException, because a part of a call cannot be run.

To approve tool calls by hand, run the graph with interruptBefore = setOf("tools"). The calls that wait are messages.pendingToolCalls(). Resume the run to let them through, or resume it with a state that has a ChatMessage.ToolResult for each call you reject: a call that already has a result is not run.

A tool that acts on the state of the run is a StateTool, and goes into stateTools:

val addLine = StateTool<Order, NewLine>("add_line", "Adds a line to the order.") { order, line ->
ToolOutcome(order.copy(lines = order.lines + line.item), "Added ${line.item}.")
}

toolLoop(
model = model,
tools = listOf(menuPrice),
stateTools = listOf(addLine),
messages = { it.messages },
append = { order, new -> order.copy(messages = order.messages + new) },
)

A loop with state tools works differently in three ways:

  • The calls of an answer run one after the other, in the order the model wrote them, and each call gets the state the call before it left. A call that fails changes nothing, and the calls after it still run.

  • append is called once for each result, before the next call runs. A state tool that calls interrupt with the state it received therefore saves what the calls before it did, with their results, and resume continues with the call that paused.

  • The tools node returns a whole state. When it can run in the same step as another node that returns a whole state, compile() asks for a Reducer, as for any two such nodes.

To know what the loop costs, give it countTokens, which adds the tokens of each answer of the model to the state. To limit it, give it stopWhen. It is asked before each call to the model, and when it returns true the model is not asked and the graph continues at then, as after a last answer:

data class Order(val messages: List<ChatMessage>, val tokens: TokenUsage = TokenUsage())

toolLoop(
model = model,
tools = tools,
messages = { it.messages },
append = { order, new -> order.copy(messages = order.messages + new) },
countTokens = { order, used -> order.copy(tokens = order.tokens + used) },
stopWhen = { order -> order.tokens.total >= 50_000 },
)

The count has the numbers the provider reported for the calls so far, so a loop with a limit on its tokens ends up to one call above it. A loop that was stopped ends with the results of its last tools, or with the message of the user, and not with an answer of the model.

A call to the model that fails ends the run. Give the loop a retry to ask the model again first. Tools need none: a tool that fails returns an error to the model, which can call it again.

toolLoop(model = model, tools = tools, messages = { it.messages }, append = ::append, retry = Retry(attempts = 3))

In a streamed run, the text of a try that failed has already arrived. The stream then has a GraphEvent.NodeRetry for the model node, and the text of the next try starts from its beginning.

Parameters

model

the model to ask.

tools

the tools the model may call that only need their input.

messages

reads the conversation from the state.

append

returns the state with the given messages added to its conversation. It is called with the model's answer, and with the results of the tools.

firstMessage

returns the user message that starts the conversation. It is used only while messages returns an empty list, and append receives it together with the model's first answer. Leave it out when the state already holds the first message.

system

instructions for the model.

then

where the graph continues when the model is done. The default ends the branch.

modelNode

the name of the model node.

toolsNode

the name of the tools node.

stateTools

the tools the model may call that read and change the state of the run.

inOrder

true to run the calls of an answer one after the other, in the order the model wrote them, for tools that depend on each other. A loop with stateTools always does.

countTokens

returns the state with the tokens of one answer of the model added to its count. It is not called for an answer whose provider reports no usage.

stopWhen

returns true when the model must not be asked again, for example because the state says that a budget is used.

retry

how the model node is tried again when a call to the model fails. null ends the run at the first failure.

Throws

if two tools have the same name, or a node name is already used.