Feature or docs request: Better/more explicit support for stateful tools #74

Closed
opened 2026-02-15 17:15:25 -05:00 by yindo · 5 comments
Owner

Originally created by @benjamincburns on GitHub (Aug 27, 2024).

I may be missing something, but at present it doesn't look like LangGraph.js has addressed the concern of stateful tools, or if it has, I don't think this is well covered in documentation yet.

For example, let's say that I have an agent with query and analysis functionality. Perhaps it queries a database for some data, and then analyzes that data in a code execution environment. How do I plumb that together via the state management system?

Let's say that I use a ToolNode for this. The query tool can be written to return an artifact, which arguably solves the first step of getting the data from the query tool into the state manager. However I'm not seeing a clear path for how I'd then go about providing this artifact as part of the input to tool calls to the code execution tool.

I suppose I could intercept the artifact via a reducer on the message state, and then provide it to the code execution environment at that stage, but it feels quite wrong to have side-effects in a reducer object.

Is there an intended approach for something like this? One way I could see to do it without modifications to LangGraph or LangChain would be to define a node for each tool handler that needs to manage state, and then handle the routing to tool handlers via conditional edges, but this seems rather cumbersome, especially if you may want to dynamically manage which tools are available, or if you want to allow for fully modular tool definitions.

Edit: the idea in the paragraph below was flawed. See this comment for details.

If there isn't an intended approach, perhaps a solution would be to make it possible to specify custom input state transform logic for each tool in the ToolNode, defaulting to the current state transform logic in the case (to avoid breaking existing code). That way I can specify an input state transformer as part of my tool module, and autowire that up to the ToolNode on application initialization, without having to have tool-specific logic defined in my ToolNode initialization logic.

I'm not particularly convinced that this is the only or best way to solve this, though - and if there's already a solution for this, I'd be very happy if the only outcome of this issue were for it to be discussed more explicitly/clearly in the docs.

Originally created by @benjamincburns on GitHub (Aug 27, 2024). I may be missing something, but at present it doesn't look like LangGraph.js has addressed the concern of stateful tools, or if it has, I don't think this is well covered in documentation yet. For example, let's say that I have an agent with query and analysis functionality. Perhaps it queries a database for some data, and then analyzes that data in a code execution environment. How do I plumb that together via the state management system? Let's say that I use a `ToolNode` for this. The query tool can be written to return an artifact, which arguably solves the first step of getting the data from the query tool into the state manager. However I'm not seeing a clear path for how I'd then go about providing this artifact as part of the input to tool calls to the code execution tool. I suppose I could intercept the artifact via a reducer on the message state, and then provide it to the code execution environment at that stage, but it feels quite wrong to have side-effects in a reducer object. Is there an intended approach for something like this? One way I could see to do it without modifications to LangGraph or LangChain would be to define a node for each tool handler that needs to manage state, and then handle the routing to tool handlers via conditional edges, but this seems rather cumbersome, especially if you may want to dynamically manage which tools are available, or if you want to allow for fully modular tool definitions. Edit: the idea in the paragraph below was flawed. See [this comment](https://github.com/langchain-ai/langgraphjs/issues/401#issuecomment-2316401990) for details. ~~If there isn't an intended approach, perhaps a solution would be to make it possible to specify custom input state transform logic for each tool in the `ToolNode`, defaulting to the current state transform logic in the case (to avoid breaking existing code). That way I can specify an input state transformer as part of my tool module, and autowire that up to the ToolNode on application initialization, without having to have tool-specific logic defined in my `ToolNode` initialization logic.~~ I'm not particularly convinced that this is the only or best way to solve this, though - and if there's already a solution for this, I'd be very happy if the only outcome of this issue were for it to be discussed more explicitly/clearly in the docs.
yindo closed this issue 2026-02-15 17:15:25 -05:00
Author
Owner

@benjamincburns commented on GitHub (Aug 27, 2024):

Edit: this strategy was flawed - see this comment for an explanation.

Happy to submit this as a PR if you guys like it, but here's the custom ToolNode implementation that I wrote for my own project. Just attaching it here to see if it passes the sniff test before I refactor it for inclusion into your repo. Note that this is not yet tested in my application, but hopefully it gets the gist across.

import {
  type BaseMessage,
  type ToolCall,
  ToolMessage,
  isAIMessage,
  isBaseMessage,
} from '@langchain/core/messages';
import type { RunnableConfig } from '@langchain/core/runnables';
import type { StructuredToolInterface } from '@langchain/core/tools';
import { END } from '@langchain/langgraph/dist/graph/graph.js';
import type { MessagesState } from '@langchain/langgraph/dist/graph/message.js';
import { RunnableCallable } from '@langchain/langgraph/dist/utils.js';

export type ToolStateTransform<
  StateT extends BaseMessage[] | MessagesState,
  OutputT,
> = (state: StateT) => OutputT;

export interface ToolNodeOptions<StateT extends BaseMessage[] | MessagesState> {
  name?: string;
  tags?: string[];
  handleToolErrors?: boolean;
  // biome-ignore lint/suspicious/noExplicitAny: <explanation>
  stateTransforms?: Record<string, ToolStateTransform<StateT, any>>;
}

export class ToolNode<
  StateT extends BaseMessage[] | MessagesState,
> extends RunnableCallable<StateT> {
  /**
    A node that runs the tools requested in the last AIMessage. It can be used
    either in StateGraph with a "messages" key or in MessageGraph. If multiple
    tool calls are requested, they will be run in parallel. The output will be
    a list of ToolMessages, one for each tool call.
    */
  tools: StructuredToolInterface[];
  handleToolErrors = true;
  // biome-ignore lint/suspicious/noExplicitAny: <explanation>
  #stateTransforms: Record<string, ToolStateTransform<StateT, any>> = {};

  constructor(
    tools: StructuredToolInterface[],
    options: ToolNodeOptions<StateT>,
  ) {
    const { name, tags, handleToolErrors, stateTransforms } = options ?? {};
    super({ name, tags, func: (input, config) => this.run(input, config) });

    this.tools = tools;
    this.handleToolErrors = handleToolErrors ?? this.handleToolErrors;
    this.#stateTransforms = stateTransforms ?? this.#stateTransforms;
  }

  async run(input: StateT, config: RunnableConfig) {
    const message = Array.isArray(input)
      ? input[input.length - 1]
      : input.messages[input.messages.length - 1];
    if (!message || !isAIMessage(message)) {
      throw new Error('ToolNode only accepts AIMessages as input.');
    }

    const outputs = await Promise.all(
      message.tool_calls?.map(async (call) => {
        const tool = this.tools.find((tool) => tool.name === call.name);
        const stateTransform = this.#stateTransforms[call.name];

        try {
          if (tool === undefined) {
            throw new Error(`Tool "${call.name}" not found.`);
          }
          let output: ReturnType<typeof tool.invoke>;
          if (stateTransform) {
            const transformedInput = await stateTransform(input);
            output = await tool.invoke(
              { ...call, args: transformedInput, type: 'tool_call' },
              config,
            );
          } else {
            output = await tool.invoke({ ...call, type: 'tool_call' }, config);
          }
          if (isBaseMessage(output) && output._getType() === 'tool') {
            return output;
          }
          return new ToolMessage({
            name: tool.name,
            content:
              typeof output === 'string' ? output : JSON.stringify(output),
            tool_call_id: call.id!,
          });
        } catch (e) {
          if (!this.handleToolErrors) {
            throw e;
          }
          return new ToolMessage({
            content: `Error: ${(e as Error).message}\n Please fix your mistakes.`,
            name: call.name,
            tool_call_id: call.id ?? '',
          });
        }
      }) ?? [],
    );
    return Array.isArray(input) ? outputs : { messages: outputs };
  }
}

export function toolsCondition<StateT extends BaseMessage[] | MessagesState>(
  state: StateT,
) {
  const message = Array.isArray(state)
    ? state[state.length - 1]
    : state.messages[state.messages.length - 1];
  if (
    message &&
    'tool_calls' in message &&
    (message.tool_calls as ToolCall[] | null)?.length
  ) {
    return 'tools';
  }
  return END;
}
@benjamincburns commented on GitHub (Aug 27, 2024): Edit: this strategy was flawed - see [this comment](https://github.com/langchain-ai/langgraphjs/issues/401#issuecomment-2316401990) for an explanation. ~~Happy to submit this as a PR if you guys like it, but here's the custom `ToolNode` implementation that I wrote for my own project. Just attaching it here to see if it passes the sniff test before I refactor it for inclusion into your repo. Note that this is not yet tested in my application, but hopefully it gets the gist across.~~ ```typescript import { type BaseMessage, type ToolCall, ToolMessage, isAIMessage, isBaseMessage, } from '@langchain/core/messages'; import type { RunnableConfig } from '@langchain/core/runnables'; import type { StructuredToolInterface } from '@langchain/core/tools'; import { END } from '@langchain/langgraph/dist/graph/graph.js'; import type { MessagesState } from '@langchain/langgraph/dist/graph/message.js'; import { RunnableCallable } from '@langchain/langgraph/dist/utils.js'; export type ToolStateTransform< StateT extends BaseMessage[] | MessagesState, OutputT, > = (state: StateT) => OutputT; export interface ToolNodeOptions<StateT extends BaseMessage[] | MessagesState> { name?: string; tags?: string[]; handleToolErrors?: boolean; // biome-ignore lint/suspicious/noExplicitAny: <explanation> stateTransforms?: Record<string, ToolStateTransform<StateT, any>>; } export class ToolNode< StateT extends BaseMessage[] | MessagesState, > extends RunnableCallable<StateT> { /** A node that runs the tools requested in the last AIMessage. It can be used either in StateGraph with a "messages" key or in MessageGraph. If multiple tool calls are requested, they will be run in parallel. The output will be a list of ToolMessages, one for each tool call. */ tools: StructuredToolInterface[]; handleToolErrors = true; // biome-ignore lint/suspicious/noExplicitAny: <explanation> #stateTransforms: Record<string, ToolStateTransform<StateT, any>> = {}; constructor( tools: StructuredToolInterface[], options: ToolNodeOptions<StateT>, ) { const { name, tags, handleToolErrors, stateTransforms } = options ?? {}; super({ name, tags, func: (input, config) => this.run(input, config) }); this.tools = tools; this.handleToolErrors = handleToolErrors ?? this.handleToolErrors; this.#stateTransforms = stateTransforms ?? this.#stateTransforms; } async run(input: StateT, config: RunnableConfig) { const message = Array.isArray(input) ? input[input.length - 1] : input.messages[input.messages.length - 1]; if (!message || !isAIMessage(message)) { throw new Error('ToolNode only accepts AIMessages as input.'); } const outputs = await Promise.all( message.tool_calls?.map(async (call) => { const tool = this.tools.find((tool) => tool.name === call.name); const stateTransform = this.#stateTransforms[call.name]; try { if (tool === undefined) { throw new Error(`Tool "${call.name}" not found.`); } let output: ReturnType<typeof tool.invoke>; if (stateTransform) { const transformedInput = await stateTransform(input); output = await tool.invoke( { ...call, args: transformedInput, type: 'tool_call' }, config, ); } else { output = await tool.invoke({ ...call, type: 'tool_call' }, config); } if (isBaseMessage(output) && output._getType() === 'tool') { return output; } return new ToolMessage({ name: tool.name, content: typeof output === 'string' ? output : JSON.stringify(output), tool_call_id: call.id!, }); } catch (e) { if (!this.handleToolErrors) { throw e; } return new ToolMessage({ content: `Error: ${(e as Error).message}\n Please fix your mistakes.`, name: call.name, tool_call_id: call.id ?? '', }); } }) ?? [], ); return Array.isArray(input) ? outputs : { messages: outputs }; } } export function toolsCondition<StateT extends BaseMessage[] | MessagesState>( state: StateT, ) { const message = Array.isArray(state) ? state[state.length - 1] : state.messages[state.messages.length - 1]; if ( message && 'tool_calls' in message && (message.tool_calls as ToolCall[] | null)?.length ) { return 'tools'; } return END; } ```
Author
Owner

@benjamincburns commented on GitHub (Aug 28, 2024):

Sigh - it just occurred to me that I've written the above against 0.0.34, not 0.1.x.

I'll have a crack at refactoring it for 0.1.4 tomorrow.

@benjamincburns commented on GitHub (Aug 28, 2024): Sigh - it just occurred to me that I've written the above against 0.0.34, not 0.1.x. I'll have a crack at refactoring it for 0.1.4 tomorrow.
Author
Owner

@benjamincburns commented on GitHub (Aug 28, 2024):

On further reflection, I don't think the strategy above will work, due to the need to bind the tools to the model. The input schema that is used during binding is used to check the inputs passed to StructuredTool.invoke.

That is, if I modify the input type of my tool to be something like { args: LLMGeneratedToolArgsT, artifacts: ArtifactT }, the LLM will be prompted to generate inputs that match that type, not the LLMGeneratedToolArgsT type. Back to the drawing board, I guess.

@benjamincburns commented on GitHub (Aug 28, 2024): On further reflection, I don't think the strategy above will work, due to the need to bind the tools to the model. The input schema that is used during binding is used to check the inputs passed to `StructuredTool.invoke`. That is, if I modify the input type of my tool to be something like `{ args: LLMGeneratedToolArgsT, artifacts: ArtifactT }`, the LLM will be prompted to generate inputs that match that type, not the `LLMGeneratedToolArgsT` type. Back to the drawing board, I guess.
Author
Owner

@benjamincburns commented on GitHub (Aug 29, 2024):

I think I've reached the conclusion that ToolNode is really only the right approach for stateless tools. The moment you add any stateful tools, you'll either need to roll your own ToolNode, or implement a graph structure where you have a node per tool handler, and you use conditional branching for parallel execution in the case when the LLM makes multiple tool calls.

Rolling your own ToolNode is probably fine for smaller projects, but as the complexity of the project goes it seems likely to become an incohesive dumping ground for state transform logic for a bunch of unrelated tools.

Personally I'd like it better if LangGraph could ship some prebuilt components to make the latter approach (a node per tool handler) the easy default, as that would make tools feel more like first-class citizens w.r.t. state management data flows, and make it easier to implement tools in a more modular, pluggable fashion. The major benefit there would be that it would allow me to point a new team member at some well-defined existing structures and just say "use that and fill in the blanks."

In the meantime I'll be reworking my own project for this approach, and if I happen to come up with some nice reusable structure for this, I'll be happy to share details here and/or submit it as a PR. That said, I expect that the LangGraph team would want some degree of parity between the Python and JS projects for this, and I'm not sure that I can address both sides of that equation here, as this is day job work for me, and it'll be difficult to explain to my bosses why I'm submitting PRs for libraries/frameworks that we're not using in our product 😅.

@benjamincburns commented on GitHub (Aug 29, 2024): I think I've reached the conclusion that `ToolNode` is really only the right approach for stateless tools. The moment you add any stateful tools, you'll either need to roll your own `ToolNode`, or implement a graph structure where you have a node per tool handler, and you use [conditional branching](https://langchain-ai.github.io/langgraphjs/how-tos/branching/#conditional-branching) for parallel execution in the case when the LLM makes multiple tool calls. Rolling your own `ToolNode` is probably fine for smaller projects, but as the complexity of the project goes it seems likely to become an [incohesive](https://en.wikipedia.org/wiki/Cohesion_(computer_science)) dumping ground for state transform logic for a bunch of unrelated tools. Personally I'd like it better if LangGraph could ship some prebuilt components to make the latter approach (a node per tool handler) the easy default, as that would make tools feel more like first-class citizens w.r.t. state management data flows, and make it easier to implement tools in a more modular, pluggable fashion. The major benefit there would be that it would allow me to point a new team member at some well-defined existing structures and just say "use that and fill in the blanks." In the meantime I'll be reworking my own project for this approach, and if I happen to come up with some nice reusable structure for this, I'll be happy to share details here and/or submit it as a PR. That said, I expect that the LangGraph team would want some degree of parity between the Python and JS projects for this, and I'm not sure that I can address both sides of that equation here, as this is day job work for me, and it'll be difficult to explain to my bosses why I'm submitting PRs for libraries/frameworks that we're not using in our product 😅.
Author
Owner

@benjamincburns commented on GitHub (Sep 29, 2024):

Closing this, as I'm not sure it's in a good state for discussion - I won't be offended if the maintainers decide to reopen, however.

My solution in the end was a bit convoluted. I had to decouple the type expected as the input argument for the StructuredTool._call method from the type described by the schema that gets used during the llm.bindTools call.

To do this, I wrote an abstract StateAugmentedTool class that extends StructuredTool and allows child classes to define a augmentedInputSchema and a augmentInput method. I'm using zod in my application, so the schema defines the actual input argument type for the _call method, and the augmentInput method gets called by a conditional edge. It gets the application state and the input arg that the LLM generated as input, and returns an object that conforms to the autmentedInputSchema.

The conditional edge that calls the augmentInput method creates a fake/transformed AIMessage from the augmented input objects, and sends this fake message to the ToolNode rather than the real state by returning something like new Send('tools', { messages: [ aiMessageWithAugmentedToolCalls ] }).

Hope this helps, in case there's anyone who needs something similar.

Personally I think there needs to be some consideration given as to how to officially support this sort of pattern, as I think the current ToolNode is too "tight" of an encapsulation, in that there's no direct way to allow individual tools to inspect application state. Or at least - no direct way that I'm currently aware of, but this project is moving fast, so I wouldn't be surprised if I missed something.

@benjamincburns commented on GitHub (Sep 29, 2024): Closing this, as I'm not sure it's in a good state for discussion - I won't be offended if the maintainers decide to reopen, however. My solution in the end was a bit convoluted. I had to decouple the type expected as the input argument for the `StructuredTool._call` method from the type described by the schema that gets used during the `llm.bindTools` call. To do this, I wrote an abstract `StateAugmentedTool` class that extends `StructuredTool` and allows child classes to define a `augmentedInputSchema` and a `augmentInput` method. I'm using zod in my application, so the schema defines the actual input argument type for the `_call` method, and the `augmentInput` method gets called by a conditional edge. It gets the application state and the input arg that the LLM generated as input, and returns an object that conforms to the `autmentedInputSchema`. The conditional edge that calls the `augmentInput` method creates a fake/transformed `AIMessage` from the augmented input objects, and sends this fake message to the `ToolNode` rather than the real state by returning something like `new Send('tools', { messages: [ aiMessageWithAugmentedToolCalls ] })`. Hope this helps, in case there's anyone who needs something similar. Personally I think there needs to be some consideration given as to how to officially support this sort of pattern, as I think the current `ToolNode` is too "tight" of an encapsulation, in that there's no direct way to allow individual tools to inspect application state. Or at least - no direct way that I'm currently aware of, but this project is moving _fast_, so I wouldn't be surprised if I missed something.
Sign in to join this conversation.
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: langchain-ai/langgraphjs#74