fix(langgraph): handle wrapped LLM models in createReactAgent (#1369)

.changeset/quick-colts-rhyme.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+Handle wrapped LLM models in createReactAgent (RunnableSequence, withConfig, ...etc)

libs/langgraph/src/prebuilt/react_agent_executor.ts

@@ -180,6 +180,13 @@ function _isConfigurableModel(model: any): model is ConfigurableModelInterface {
   );
 }
 
+function _isChatModelWithBindTools(
+  llm: LanguageModelLike
+): llm is BaseChatModel & Required<Pick<BaseChatModel, "bindTools">> {
+  if (!_isBaseChatModel(llm)) return false;
+  return "bindTools" in llm && typeof llm.bindTools === "function";
+}
+
 export async function _shouldBindTools(
   llm: LanguageModelLike,
   tools: (ClientTool | ServerTool)[]
@@ -205,21 +212,42 @@ export async function _shouldBindTools(
     return true;
   }
 
-  // If no tools in kwargs, we should bind tools
-  if (
-    !model.kwargs ||
-    typeof model.kwargs !== "object" ||
-    !("tools" in model.kwargs)
-  ) {
-    return true;
-  }
+  let boundTools = (() => {
+    // check if model.kwargs contain the tools key
+    if (
+      model.kwargs != null &&
+      typeof model.kwargs === "object" &&
+      "tools" in model.kwargs &&
+      Array.isArray(model.kwargs.tools)
+    ) {
+      return (model.kwargs.tools ?? null) as BindToolsInput[] | null;
+    }
+
+    // Some models can bind the tools via `withConfig()` instead of `bind()`
+    if (
+      model.config != null &&
+      typeof model.config === "object" &&
+      "tools" in model.config &&
+      Array.isArray(model.config.tools)
+    ) {
+      return (model.config.tools ?? null) as BindToolsInput[] | null;
+    }
+
+    return null;
+  })();
 
-  let boundTools = model.kwargs.tools as BindToolsInput[];
   // google-style
-  if (boundTools.length === 1 && "functionDeclarations" in boundTools[0]) {
+  if (
+    boundTools != null &&
+    boundTools.length === 1 &&
+    "functionDeclarations" in boundTools[0]
+  ) {
     boundTools = boundTools[0].functionDeclarations;
   }
 
+  // If no tools in kwargs, we should bind tools
+  if (boundTools == null) return true;
+
   // Check if tools count matches
   if (tools.length !== boundTools.length) {
     throw new Error(
@@ -269,6 +297,76 @@ export async function _shouldBindTools(
   return false;
 }
 
+const _simpleBindTools = (
+  llm: LanguageModelLike,
+  toolClasses: (ClientTool | ServerTool)[]
+) => {
+  if (_isChatModelWithBindTools(llm)) {
+    return llm.bindTools(toolClasses);
+  }
+
+  if (
+    RunnableBinding.isRunnableBinding(llm) &&
+    _isChatModelWithBindTools(llm.bound)
+  ) {
+    const newBound = llm.bound.bindTools(toolClasses);
+
+    if (RunnableBinding.isRunnableBinding(newBound)) {
+      return new RunnableBinding({
+        bound: newBound.bound,
+        config: { ...llm.config, ...newBound.config },
+        kwargs: { ...llm.kwargs, ...newBound.kwargs },
+        configFactories: newBound.configFactories ?? llm.configFactories,
+      });
+    }
+
+    return new RunnableBinding({
+      bound: newBound,
+      config: llm.config,
+      kwargs: llm.kwargs,
+      configFactories: llm.configFactories,
+    });
+  }
+
+  return null;
+};
+
+export async function _bindTools(
+  llm: LanguageModelLike,
+  toolClasses: (ClientTool | ServerTool)[]
+) {
+  const model = _simpleBindTools(llm, toolClasses);
+  if (model) return model;
+
+  if (_isConfigurableModel(llm)) {
+    const model = _simpleBindTools(await llm._model(), toolClasses);
+    if (model) return model;
+  }
+
+  if (RunnableSequence.isRunnableSequence(llm)) {
+    const modelStep = llm.steps.findIndex(
+      (step) =>
+        RunnableBinding.isRunnableBinding(step) ||
+        _isBaseChatModel(step) ||
+        _isConfigurableModel(step)
+    );
+
+    if (modelStep >= 0) {
+      const model = _simpleBindTools(llm.steps[modelStep], toolClasses);
+      if (model) {
+        const nextSteps: unknown[] = llm.steps.slice();
+        nextSteps.splice(modelStep, 1, model);
+
+        return RunnableSequence.from(
+          nextSteps as [RunnableLike, ...RunnableLike[], RunnableLike]
+        );
+      }
+    }
+  }
+
+  throw new Error(`llm ${llm} must define bindTools method.`);
+}
+
 export async function _getModel(
   llm: LanguageModelLike | ConfigurableModelInterface
 ): Promise<BaseChatModel> {
@@ -562,10 +660,7 @@ export function createReactAgent<
 
     let modelWithTools: LanguageModelLike;
     if (await _shouldBindTools(llm, toolClasses)) {
-      if (!("bindTools" in llm) || typeof llm.bindTools !== "function") {
-        throw new Error(`llm ${llm} must define bindTools method.`);
-      }
-      modelWithTools = llm.bindTools(toolClasses);
+      modelWithTools = await _bindTools(llm, toolClasses);
     } else {
       modelWithTools = llm;
     }

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

@@ -13,7 +13,11 @@ import {
   ToolMessage,
 } from "@langchain/core/messages";
 import { z } from "zod";
-import { RunnableLambda, RunnableSequence } from "@langchain/core/runnables";
+import {
+  Runnable,
+  RunnableLambda,
+  RunnableSequence,
+} from "@langchain/core/runnables";
 import { CallbackManager } from "@langchain/core/callbacks/manager";
 import { BaseCallbackHandler } from "@langchain/core/callbacks/base";
 import {
@@ -29,6 +33,7 @@ import { ToolNode, createReactAgent } from "../prebuilt/index.js";
 import {
   _shouldBindTools,
   _getModel,
+  _bindTools,
 } from "../prebuilt/react_agent_executor.js";
 // Enable automatic config passing
 import {
@@ -2009,6 +2014,128 @@ describe("_shouldBindTools", () => {
     }
   );
 
+  it("should bind model with bindTools", async () => {
+    const tool1 = tool((input) => `Tool 1: ${input.someVal}`, {
+      name: "tool1",
+      description: "Tool 1 docstring.",
+      schema: z.object({
+        someVal: z.number().describe("Input value"),
+      }),
+    });
+
+    const model = new FakeToolCallingChatModel({
+      responses: [new AIMessage("test")],
+      toolStyle: "openai",
+    });
+
+    const confModel = new FakeConfigurableModel({ model });
+
+    async function serialize(runnable: Runnable | Promise<Runnable>) {
+      return JSON.parse(JSON.stringify(await runnable));
+    }
+
+    // Should bind when a regular model
+    expect(await serialize(_bindTools(model, [tool1]))).toEqual(
+      await serialize(model.bindTools([tool1]))
+    );
+
+    // Should bind when model wrapped in `withConfig`
+    expect(
+      await serialize(
+        _bindTools(model.withConfig({ tags: ["nostream"] }), [tool1])
+      )
+    ).toEqual(
+      await serialize(
+        model.bindTools([tool1]).withConfig({ tags: ["nostream"] })
+      )
+    );
+
+    // Should bind when model wrapped in multiple `withConfig`
+    expect(
+      await serialize(
+        _bindTools(
+          model
+            .withConfig({ tags: ["nostream"] })
+            .withConfig({ metadata: { hello: "world" } }),
+          [tool1]
+        )
+      )
+    ).toEqual(
+      await serialize(
+        model
+          .bindTools([tool1])
+          .withConfig({ tags: ["nostream"], metadata: { hello: "world" } })
+      )
+    );
+
+    // Should bind when a configurable model
+    expect(await serialize(_bindTools(confModel, [tool1]))).toEqual(
+      await serialize(confModel.bindTools([tool1]))
+    );
+
+    // Should bind when a seq
+    expect(
+      await serialize(
+        _bindTools(
+          RunnableSequence.from([
+            model,
+            RunnableLambda.from((message) => message),
+          ]),
+          [tool1]
+        )
+      )
+    ).toEqual(
+      await serialize(
+        RunnableSequence.from([
+          model.bindTools([tool1]),
+          RunnableLambda.from((message) => message),
+        ])
+      )
+    );
+
+    // Should bind when a seq with configurable model
+    expect(
+      await serialize(
+        _bindTools(
+          RunnableSequence.from([
+            confModel,
+            RunnableLambda.from((message) => message),
+          ]),
+          [tool1]
+        )
+      )
+    ).toEqual(
+      await serialize(
+        RunnableSequence.from([
+          confModel.bindTools([tool1]),
+          RunnableLambda.from((message) => message),
+        ])
+      )
+    );
+
+    // Should bind when a seq with config model
+    expect(
+      await serialize(
+        _bindTools(
+          RunnableSequence.from([
+            confModel.withConfig({ tags: ["nostream"] }),
+            RunnableLambda.from((message) => message),
+          ]),
+          [tool1]
+        )
+      )
+    ).toEqual(
+      await serialize(
+        RunnableSequence.from([
+          confModel.bindTools([tool1]).withConfig({
+            tags: ["nostream"],
+          }),
+          RunnableLambda.from((message) => message),
+        ])
+      )
+    );
+  });
+
   it("should handle bindTool with server tools", async () => {
     const tool1 = tool((input) => `Tool 1: ${input.someVal}`, {
       name: "tool1",