feat(langgraph): Allow partially applying tool calls via postModelHook (#1499)

.changeset/moody-bags-take.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+feat(langgraph): Allow partially applying tool calls via postModelHook

libs/langgraph/src/prebuilt/react_agent_executor.ts

@@ -10,6 +10,7 @@ import {
   isBaseMessage,
   isToolMessage,
   SystemMessage,
+  type AIMessage,
 } from "@langchain/core/messages";
 import {
   Runnable,
@@ -872,11 +873,29 @@ export function createReactAgent<
         "post_model_hook",
         (state: AgentState<StructuredResponseFormat>) => {
           const { messages } = state;
-          const lastMessage = messages[messages.length - 1];
 
-          if (isAIMessage(lastMessage) && lastMessage.tool_calls?.length) {
+          const toolMessageIds: Set<string> = new Set(
+            messages.filter(isToolMessage).map((msg) => msg.tool_call_id)
+          );
+
+          let lastAiMessage: AIMessage | undefined;
+          for (let i = messages.length - 1; i >= 0; i -= 1) {
+            const message = messages[i];
+            if (isAIMessage(message)) {
+              lastAiMessage = message;
+              break;
+            }
+          }
+
+          const pendingToolCalls =
+            lastAiMessage?.tool_calls?.filter(
+              (i) => i.id == null || !toolMessageIds.has(i.id)
+            ) ?? [];
+
+          const lastMessage = messages.at(-1);
+          if (pendingToolCalls.length > 0) {
             if (version === "v2") {
-              return lastMessage.tool_calls.map(
+              return pendingToolCalls.map(
                 (toolCall) =>
                   new Send("tools", { ...state, lg_tool_call: toolCall })
               );
@@ -884,7 +903,7 @@ export function createReactAgent<
             return "tools";
           }
 
-          if (isToolMessage(lastMessage)) return entrypoint;
+          if (lastMessage && isToolMessage(lastMessage)) return entrypoint;
           if (responseFormat != null) return "generate_structured_response";
           return END;
         },

libs/langgraph/src/prebuilt/tool_node.ts

@@ -19,7 +19,7 @@ export type ToolNodeOptions = {
 };
 
 const isBaseMessageArray = (input: unknown): input is BaseMessage[] =>
-  Array.isArray(input);
+  Array.isArray(input) && input.every(isBaseMessage);
 
 const isMessagesState = (
   input: unknown
@@ -27,7 +27,7 @@ const isMessagesState = (
   typeof input === "object" &&
   input != null &&
   "messages" in input &&
-  Array.isArray(input.messages);
+  isBaseMessageArray(input.messages);
 
 const isSendInput = (input: unknown): input is { lg_tool_call: ToolCall } =>
   typeof input === "object" && input != null && "lg_tool_call" in input;
@@ -216,19 +216,40 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     if (isSendInput(input)) {
       outputs = [await this.runTool(input.lg_tool_call, config)];
     } else {
-      let message: AIMessage | undefined;
+      let messages: BaseMessage[];
       if (isBaseMessageArray(input)) {
-        message = input.at(-1);
+        messages = input;
       } else if (isMessagesState(input)) {
-        message = input.messages.at(-1);
+        messages = input.messages;
+      } else {
+        throw new Error(
+          "ToolNode only accepts BaseMessage[] or { messages: BaseMessage[] } as input."
+        );
       }
 
-      if (message?.getType() !== "ai") {
+      const toolMessageIds: Set<string> = new Set(
+        messages
+          .filter((msg) => msg.getType() === "tool")
+          .map((msg) => (msg as ToolMessage).tool_call_id)
+      );
+
+      let aiMessage: AIMessage | undefined;
+      for (let i = messages.length - 1; i >= 0; i -= 1) {
+        const message = messages[i];
+        if (message.getType() === "ai") {
+          aiMessage = message;
+          break;
+        }
+      }
+
+      if (aiMessage?.getType() !== "ai") {
         throw new Error("ToolNode only accepts AIMessages as input.");
       }
 
       outputs = await Promise.all(
-        message.tool_calls?.map((call) => this.runTool(call, config)) ?? []
+        aiMessage.tool_calls
+          ?.filter((call) => call.id == null || !toolMessageIds.has(call.id))
+          .map((call) => this.runTool(call, config)) ?? []
       );
     }
 

libs/langgraph/src/tests/prebuilt.test.ts

@@ -8,6 +8,7 @@ import {
   AIMessage,
   BaseMessage,
   HumanMessage,
+  isAIMessage,
   RemoveMessage,
   SystemMessage,
   ToolMessage,
@@ -1542,6 +1543,68 @@ describe.each([["v1" as const], ["v2" as const]])(
         ]);
       });
 
+      it("postModelHook + partial tool call application", async () => {
+        const llm = new FakeToolCallingChatModel({
+          responses: [
+            new AIMessage({
+              content: "result1",
+              tool_calls: [
+                { name: "search_api", id: "tool_a", args: { query: "foo" } },
+                { name: "search_api", id: "tool_b", args: { query: "bar" } },
+              ],
+            }),
+            new AIMessage("done"),
+          ],
+        });
+
+        const agent = createReactAgent({
+          llm,
+          tools: [new SearchAPI()],
+          version,
+          postModelHook: (state) => {
+            const lastMessage = state.messages.at(-1);
+            if (
+              lastMessage != null &&
+              isAIMessage(lastMessage) &&
+              lastMessage.tool_calls?.length
+            ) {
+              // apply only the first tool call
+              const firstToolCall = lastMessage.tool_calls[0]!;
+
+              return {
+                messages: [
+                  new ToolMessage({
+                    content: "post-model-hook",
+                    tool_call_id: firstToolCall.id!,
+                  }),
+                ],
+              };
+            }
+
+            return {};
+          },
+        });
+
+        const result = await agent.invoke({
+          messages: [new HumanMessage("What's the weather?")],
+        });
+
+        expect(result).toMatchObject({
+          messages: [
+            { text: "What's the weather?" },
+            {
+              tool_calls: [
+                { name: "search_api", id: "tool_a", args: { query: "foo" } },
+                { name: "search_api", id: "tool_b", args: { query: "bar" } },
+              ],
+            },
+            { text: "post-model-hook" },
+            { text: "result for bar" },
+            { text: "done" },
+          ],
+        });
+      });
+
       it.each([
         [
           {