feat(langgraph): use Send for each of the tool calls (#1498)

.changeset/cold-walls-burn.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+feat(langgraph): createReactAgent v2: use Send for each of the tool calls

libs/langgraph/src/prebuilt/react_agent_executor.ts

@@ -41,7 +41,7 @@ import { ToolNode } from "./tool_node.js";
 import { LangGraphRunnableConfig, Runtime } from "../pregel/runnable_types.js";
 import { Annotation } from "../graph/annotation.js";
 import { Messages, messagesStateReducer } from "../graph/message.js";
-import { END, START } from "../constants.js";
+import { END, Send, START } from "../constants.js";
 import { withAgentName } from "./agentName.js";
 import type { InteropZodToStateDefinition } from "../graph/zod/meta.js";
 
@@ -591,6 +591,19 @@ export type CreateReactAgentParams<
     ToAnnotationRoot<A>["Update"],
     LangGraphRunnableConfig
   >;
+
+  /**
+   * Determines the version of the graph to create.
+   *
+   * Can be one of
+   * - `"v1"`: The tool node processes a single message. All tool calls in the message are
+   *           executed in parallel within the tool node.
+   * - `"v2"`: The tool node processes a single tool call. Tool calls are distributed across
+   *           multiple instances of the tool node using the Send API.
+   *
+   * @default `"v1"`
+   */
+  version?: "v1" | "v2";
 };
 
 /**
@@ -670,6 +683,7 @@ export function createReactAgent<
     preModelHook,
     postModelHook,
     name,
+    version = "v1",
     includeAgentName,
   } = params;
 
@@ -861,6 +875,12 @@ export function createReactAgent<
           const lastMessage = messages[messages.length - 1];
 
           if (isAIMessage(lastMessage) && lastMessage.tool_calls?.length) {
+            if (version === "v2") {
+              return lastMessage.tool_calls.map(
+                (toolCall) =>
+                  new Send("tools", { ...state, lg_tool_call: toolCall })
+              );
+            }
             return "tools";
           }
 
@@ -898,6 +918,13 @@ export function createReactAgent<
         }
 
         // there are function calls, we continue
+        if (version === "v2") {
+          return lastMessage.tool_calls.map(
+            (toolCall) =>
+              new Send("tools", { ...state, lg_tool_call: toolCall })
+          );
+        }
+
         return "tools";
       },
       conditionalMap({

libs/langgraph/src/prebuilt/tool_node.ts

@@ -6,6 +6,7 @@ import {
 } from "@langchain/core/messages";
 import { RunnableConfig, RunnableToolLike } from "@langchain/core/runnables";
 import { DynamicTool, StructuredToolInterface } from "@langchain/core/tools";
+import type { ToolCall } from "@langchain/core/messages/tool";
 import { RunnableCallable } from "../utils.js";
 import { MessagesAnnotation } from "../graph/messages_annotation.js";
 import { isGraphInterrupt } from "../errors.js";
@@ -17,6 +18,20 @@ export type ToolNodeOptions = {
   handleToolErrors?: boolean;
 };
 
+const isBaseMessageArray = (input: unknown): input is BaseMessage[] =>
+  Array.isArray(input);
+
+const isMessagesState = (
+  input: unknown
+): input is { messages: BaseMessage[] } =>
+  typeof input === "object" &&
+  input != null &&
+  "messages" in input &&
+  Array.isArray(input.messages);
+
+const isSendInput = (input: unknown): input is { lg_tool_call: ToolCall } =>
+  typeof input === "object" && input != null && "lg_tool_call" in input;
+
 /**
  * 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
@@ -152,59 +167,70 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     this.handleToolErrors = handleToolErrors ?? this.handleToolErrors;
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  protected async run(input: any, config: RunnableConfig): Promise<T> {
-    const message = Array.isArray(input)
-      ? input[input.length - 1]
-      : input.messages[input.messages.length - 1];
+  protected async runTool(
+    call: ToolCall,
+    config: RunnableConfig
+  ): Promise<ToolMessage | Command> {
+    const tool = this.tools.find((tool) => tool.name === call.name);
+    try {
+      if (tool === undefined) {
+        throw new Error(`Tool "${call.name}" not found.`);
+      }
+      const output = await tool.invoke({ ...call, type: "tool_call" }, config);
 
-    if (message?._getType() !== "ai") {
-      throw new Error("ToolNode only accepts AIMessages as input.");
+      if (
+        (isBaseMessage(output) && output.getType() === "tool") ||
+        isCommand(output)
+      ) {
+        return output as ToolMessage | Command;
+      }
+
+      return new ToolMessage({
+        name: tool.name,
+        content: typeof output === "string" ? output : JSON.stringify(output),
+        tool_call_id: call.id!,
+      });
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    } catch (e: any) {
+      if (!this.handleToolErrors) throw e;
+
+      if (isGraphInterrupt(e)) {
+        // `NodeInterrupt` errors are a breakpoint to bring a human into the loop.
+        // As such, they are not recoverable by the agent and shouldn't be fed
+        // back. Instead, re-throw these errors even when `handleToolErrors = true`.
+        throw e;
+      }
+
+      return new ToolMessage({
+        content: `Error: ${e.message}\n Please fix your mistakes.`,
+        name: call.name,
+        tool_call_id: call.id ?? "",
+      });
     }
+  }
 
-    const outputs = await Promise.all(
-      (message as AIMessage).tool_calls?.map(async (call) => {
-        const tool = this.tools.find((tool) => tool.name === call.name);
-        try {
-          if (tool === undefined) {
-            throw new Error(`Tool "${call.name}" not found.`);
-          }
-          const output = await tool.invoke(
-            { ...call, type: "tool_call" },
-            config
-          );
-          if (
-            (isBaseMessage(output) && output._getType() === "tool") ||
-            isCommand(output)
-          ) {
-            return output;
-          } else {
-            return new ToolMessage({
-              name: tool.name,
-              content:
-                typeof output === "string" ? output : JSON.stringify(output),
-              tool_call_id: call.id!,
-            });
-          }
-          // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        } catch (e: any) {
-          if (!this.handleToolErrors) {
-            throw e;
-          }
-          if (isGraphInterrupt(e)) {
-            // `NodeInterrupt` errors are a breakpoint to bring a human into the loop.
-            // As such, they are not recoverable by the agent and shouldn't be fed
-            // back. Instead, re-throw these errors even when `handleToolErrors = true`.
-            throw e;
-          }
-          return new ToolMessage({
-            content: `Error: ${e.message}\n Please fix your mistakes.`,
-            name: call.name,
-            tool_call_id: call.id ?? "",
-          });
-        }
-      }) ?? []
-    );
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  protected async run(input: unknown, config: RunnableConfig): Promise<T> {
+    let outputs: (ToolMessage | Command)[];
+
+    if (isSendInput(input)) {
+      outputs = [await this.runTool(input.lg_tool_call, config)];
+    } else {
+      let message: AIMessage | undefined;
+      if (isBaseMessageArray(input)) {
+        message = input.at(-1);
+      } else if (isMessagesState(input)) {
+        message = input.messages.at(-1);
+      }
+
+      if (message?.getType() !== "ai") {
+        throw new Error("ToolNode only accepts AIMessages as input.");
+      }
+
+      outputs = await Promise.all(
+        message.tool_calls?.map((call) => this.runTool(call, config)) ?? []
+      );
+    }
 
     // Preserve existing behavior for non-command tool outputs for backwards compatibility
     if (!outputs.some(isCommand)) {