feat(langgraph): add dynamic model choice to createReactAgent (#1493)

.changeset/quick-spoons-shop.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+feat(langgraph): add dynamic model choice to createReactAgent

libs/langgraph/src/prebuilt/react_agent_executor.ts

@@ -14,7 +14,6 @@ import {
 import {
   Runnable,
   RunnableConfig,
-  RunnableInterface,
   RunnableLambda,
   RunnableToolLike,
   RunnableSequence,
@@ -39,7 +38,7 @@ import {
 } from "../graph/index.js";
 import { MessagesAnnotation } from "../graph/messages_annotation.js";
 import { ToolNode } from "./tool_node.js";
-import { LangGraphRunnableConfig } from "../pregel/runnable_types.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";
@@ -101,8 +100,8 @@ function _convertMessageModifierToPrompt(
 
 const PROMPT_RUNNABLE_NAME = "prompt";
 
-function _getPromptRunnable(prompt?: Prompt): RunnableInterface {
-  let promptRunnable: RunnableInterface;
+function _getPromptRunnable(prompt?: Prompt): Runnable {
+  let promptRunnable: Runnable;
 
   if (prompt == null) {
     promptRunnable = RunnableLambda.from(
@@ -143,7 +142,7 @@ function _getPrompt(
   prompt?: Prompt,
   stateModifier?: CreateReactAgentParams["stateModifier"],
   messageModifier?: CreateReactAgentParams["messageModifier"]
-) {
+): Runnable {
   // Check if multiple modifiers exist
   const definedCount = [prompt, stateModifier, messageModifier].filter(
     (x) => x != null
@@ -376,7 +375,7 @@ export async function _bindTools(
 
 export async function _getModel(
   llm: LanguageModelLike | ConfigurableModelInterface
-): Promise<BaseChatModel> {
+): Promise<LanguageModelLike> {
   // If model is a RunnableSequence, find a RunnableBinding or BaseChatModel in its steps
   let model = llm;
   if (RunnableSequence.isRunnableSequence(model)) {
@@ -395,7 +394,7 @@ export async function _getModel(
 
   // Get the underlying model from a RunnableBinding
   if (RunnableBinding.isRunnableBinding(model)) {
-    model = model.bound as BaseChatModel;
+    model = model.bound;
   }
 
   if (!_isBaseChatModel(model)) {
@@ -404,7 +403,7 @@ export async function _getModel(
     );
   }
 
-  return model as BaseChatModel;
+  return model;
 }
 
 export type Prompt =
@@ -477,10 +476,16 @@ type ToAnnotationRoot<A extends AnyAnnotationRoot | InteropZodObject> =
 export type CreateReactAgentParams<
   A extends AnyAnnotationRoot | InteropZodObject = AnyAnnotationRoot,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  StructuredResponseType = Record<string, any>
+  StructuredResponseType extends Record<string, any> = Record<string, any>,
+  C extends AnyAnnotationRoot | InteropZodObject = AnyAnnotationRoot
 > = {
   /** The chat model that can utilize OpenAI-style tool calling. */
-  llm: LanguageModelLike;
+  llm:
+    | LanguageModelLike
+    | ((
+        state: ToAnnotationRoot<A>["State"] & PreHookAnnotation["State"],
+        runtime: Runtime<ToAnnotationRoot<C>["State"]>
+      ) => Promise<LanguageModelLike> | LanguageModelLike);
 
   /** A list of tools or a ToolNode. */
   tools: ToolNode | (ServerTool | ClientTool)[];
@@ -510,7 +515,16 @@ export type CreateReactAgentParams<
    * This is now deprecated and will be removed in a future release.
    */
   prompt?: Prompt;
+
+  /**
+   * Additional state schema for the agent.
+   */
   stateSchema?: A;
+
+  /**
+   * An optional schema for the context.
+   */
+  contextSchema?: C;
   /** An optional checkpoint saver to persist the agent's state. */
   checkpointSaver?: BaseCheckpointSaver | boolean;
   /** An optional checkpoint saver to persist the agent's state. Alias of "checkpointSaver". */
@@ -624,9 +638,10 @@ export type CreateReactAgentParams<
 export function createReactAgent<
   A extends AnyAnnotationRoot | InteropZodObject = typeof MessagesAnnotation,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  StructuredResponseFormat extends Record<string, any> = Record<string, any>
+  StructuredResponseFormat extends Record<string, any> = Record<string, any>,
+  C extends AnyAnnotationRoot | InteropZodObject = AnyAnnotationRoot
 >(
-  params: CreateReactAgentParams<A, StructuredResponseFormat>
+  params: CreateReactAgentParams<A, StructuredResponseFormat, C>
 ): CompiledStateGraph<
   ToAnnotationRoot<A>["State"],
   ToAnnotationRoot<A>["Update"],
@@ -645,6 +660,7 @@ export function createReactAgent<
     stateModifier,
     prompt,
     stateSchema,
+    contextSchema,
     checkpointSaver,
     checkpointer,
     interruptBefore,
@@ -668,12 +684,10 @@ export function createReactAgent<
     toolNode = new ToolNode(toolClasses.filter(isClientTool));
   }
 
-  let cachedModelRunnable: Runnable | null = null;
+  let cachedStaticModel: Runnable | null = null;
 
-  const getModelRunnable = async (llm: LanguageModelLike) => {
-    if (cachedModelRunnable) {
-      return cachedModelRunnable;
-    }
+  const _getStaticModel = async (llm: LanguageModelLike): Promise<Runnable> => {
+    if (cachedStaticModel) return cachedStaticModel;
 
     let modelWithTools: LanguageModelLike;
     if (await _shouldBindTools(llm, toolClasses)) {
@@ -682,19 +696,31 @@ export function createReactAgent<
       modelWithTools = llm;
     }
 
-    const promptRunnable = _getPrompt(
-      prompt,
-      stateModifier,
-      messageModifier
-    ) as Runnable;
-
+    const promptRunnable = _getPrompt(prompt, stateModifier, messageModifier);
     const modelRunnable =
       includeAgentName === "inline"
-        ? promptRunnable.pipe(withAgentName(modelWithTools, includeAgentName))
-        : promptRunnable.pipe(modelWithTools);
+        ? withAgentName(modelWithTools, includeAgentName)
+        : modelWithTools;
 
-    cachedModelRunnable = modelRunnable;
-    return modelRunnable;
+    cachedStaticModel = promptRunnable.pipe(modelRunnable);
+    return cachedStaticModel;
+  };
+
+  const _getDynamicModel = async (
+    llm: (
+      state: AgentState<StructuredResponseFormat> & PreHookAnnotation["State"],
+      runtime: LangGraphRunnableConfig
+    ) => Promise<LanguageModelLike> | LanguageModelLike,
+    state: AgentState<StructuredResponseFormat> & PreHookAnnotation["State"],
+    config: LangGraphRunnableConfig
+  ) => {
+    const model = await llm(state, config);
+
+    return _getPrompt(prompt, stateModifier, messageModifier).pipe(
+      includeAgentName === "inline"
+        ? withAgentName(model, includeAgentName)
+        : model
+    );
   };
 
   // If any of the tools are configured to return_directly after running,
@@ -717,8 +743,8 @@ export function createReactAgent<
   }
 
   const generateStructuredResponse = async (
-    state: AgentState<StructuredResponseFormat>,
-    config?: RunnableConfig
+    state: AgentState<StructuredResponseFormat> & PreHookAnnotation["State"],
+    config: RunnableConfig
   ) => {
     if (responseFormat == null) {
       throw new Error(
@@ -728,7 +754,16 @@ export function createReactAgent<
     const messages = [...state.messages];
     let modelWithStructuredOutput;
 
-    const model = await _getModel(llm);
+    const model: LanguageModelLike =
+      typeof llm === "function"
+        ? await llm(state, config)
+        : await _getModel(llm);
+
+    if (!_isBaseChatModel(model)) {
+      throw new Error(
+        `Expected \`llm\` to be a ChatModel with .withStructuredOutput() method, got ${model.constructor.name}`
+      );
+    }
 
     if (typeof responseFormat === "object" && "schema" in responseFormat) {
       const { prompt, schema, ...options } =
@@ -748,11 +783,15 @@ export function createReactAgent<
 
   const callModel = async (
     state: AgentState<StructuredResponseFormat> & PreHookAnnotation["State"],
-    config?: RunnableConfig
+    config: RunnableConfig
   ) => {
     // NOTE: we're dynamically creating the model runnable here
     // to ensure that we can validate ConfigurableModel properly
-    const modelRunnable = await getModelRunnable(llm);
+    const modelRunnable: Runnable =
+      typeof llm === "function"
+        ? await _getDynamicModel(llm, state, config)
+        : await _getStaticModel(llm);
+
     // TODO: Auto-promote streaming.
     const response = (await modelRunnable.invoke(
       getModelInputState(state),
@@ -768,10 +807,10 @@ export function createReactAgent<
   const schema =
     stateSchema ?? createReactAgentAnnotation<StructuredResponseFormat>();
 
-  const workflow = new StateGraph(schema as AnyAnnotationRoot).addNode(
-    "tools",
-    toolNode
-  );
+  const workflow = new StateGraph(
+    schema as AnyAnnotationRoot,
+    contextSchema
+  ).addNode("tools", toolNode);
 
   if (!("messages" in workflow._schemaDefinition)) {
     throw new Error("Missing required `messages` key in state schema.");

libs/langgraph/src/pregel/runnable_types.ts

@@ -35,9 +35,7 @@ export interface LangGraphRunnableConfig<
   writer?: (chunk: unknown) => void;
 }
 
-export interface Runtime<
-  ContextType extends Record<string, unknown> = Record<string, unknown>
-> {
+export interface Runtime<ContextType = Record<string, unknown>> {
   context?: ContextType;
 
   store?: BaseStore;

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

@@ -45,6 +45,7 @@ import {
   MemorySaver,
   messagesStateReducer,
   REMOVE_ALL_MESSAGES,
+  Runtime,
   Send,
   StateGraph,
 } from "../index.js";
@@ -2881,6 +2882,7 @@ describe("ToolNode with Commands", () => {
     ]);
   });
 });
+
 describe("ToolNode should raise GraphInterrupt", () => {
   it("should raise GraphInterrupt", async () => {
     const toolWithInterrupt = tool(
@@ -2908,3 +2910,415 @@ describe("ToolNode should raise GraphInterrupt", () => {
     ).rejects.toThrow(GraphInterrupt);
   });
 });
+
+describe("Dynamic Model", () => {
+  it("should handle basic dynamic model functionality", async () => {
+    const dynamicModel = (state: typeof MessagesAnnotation.State) => {
+      // Return different models based on state
+      if (state.messages.at(-1)?.text.includes("urgent")) {
+        return new FakeToolCallingChatModel({
+          responses: [new AIMessage("urgent called")],
+        });
+      }
+
+      return new FakeToolCallingChatModel({ responses: [] });
+    };
+
+    const agent = createReactAgent({ llm: dynamicModel, tools: [] });
+
+    const result = await agent.invoke({ messages: "hello" });
+    expect(result.messages.at(-1)?.text).toBe("hello");
+
+    const result2 = await agent.invoke({ messages: "urgent help" });
+    expect(result2.messages.at(-1)?.text).toBe("urgent called");
+  });
+
+  it("should handle dynamic model with tool calling", async () => {
+    const basicModel = new FakeToolCallingChatModel({
+      responses: [
+        new AIMessage({
+          content: "",
+          tool_calls: [{ args: { x: 1 }, id: "1", name: "basic_tool" }],
+        }),
+        new AIMessage("basic request"),
+      ],
+    });
+
+    const basicTool = tool(async (args: { x: number }) => `basic: ${args.x}`, {
+      name: "basic_tool",
+      description: "Basic tool.",
+      schema: z.object({ x: z.number() }),
+    });
+
+    const advancedModel = new FakeToolCallingChatModel({
+      responses: [
+        new AIMessage({
+          content: "",
+          tool_calls: [{ args: { x: 1 }, id: "1", name: "advanced_tool" }],
+        }),
+        new AIMessage("advanced request"),
+      ],
+    });
+
+    const advancedTool = tool(
+      async (args: { x: number }) => `advanced: ${args.x}`,
+      {
+        name: "advanced_tool",
+        description: "Advanced tool.",
+        schema: z.object({ x: z.number() }),
+      }
+    );
+
+    const dynamicModel = (state: typeof MessagesAnnotation.State) => {
+      // Return model with different behaviors based on message content
+      if (state.messages.at(-1)?.text.includes("advanced")) {
+        return advancedModel;
+      }
+
+      return basicModel;
+    };
+
+    const agent = createReactAgent({
+      llm: dynamicModel,
+      tools: [basicTool, advancedTool],
+    });
+
+    // Test basic tool usage
+    const result = await agent.invoke({ messages: "basic request" });
+    expect(result.messages.slice(-2)).toMatchObject([
+      { text: "basic: 1", name: "basic_tool" },
+      { text: "basic request" },
+    ]);
+
+    // Test advanced tool usage
+    const result2 = await agent.invoke({ messages: "advanced request" });
+    expect(result2.messages.slice(-2)).toMatchObject([
+      { text: "advanced: 1", name: "advanced_tool" },
+      { text: "advanced request" },
+    ]);
+  });
+
+  it("should handle dynamic model using config parameters", async () => {
+    const context = z.object({ user_id: z.string() });
+
+    const dynamicModel = (
+      _: typeof MessagesAnnotation.State,
+      runtime: Runtime<z.infer<typeof context>>
+    ) => {
+      // Use context to determine model behavior
+      const user_id = runtime.context?.user_id;
+      if (user_id === "user_premium") {
+        return new FakeToolCallingChatModel({
+          responses: [new AIMessage("premium")],
+        });
+      }
+
+      return new FakeToolCallingChatModel({
+        responses: [new AIMessage("basic")],
+      });
+    };
+
+    const agent = createReactAgent({
+      llm: dynamicModel,
+      tools: [],
+      contextSchema: context,
+    });
+
+    // Test with basic user
+    expect(
+      await agent.invoke(
+        { messages: "hello" },
+        { context: { user_id: "user_basic" } }
+      )
+    ).toMatchObject({
+      messages: [{ text: "hello" }, { text: "basic" }],
+    });
+
+    // Test with premium user
+    expect(
+      await agent.invoke(
+        { messages: "hello" },
+        { context: { user_id: "user_premium" } }
+      )
+    ).toMatchObject({
+      messages: [{ text: "hello" }, { text: "premium" }],
+    });
+  });
+
+  it("should handle dynamic model with custom state schema", async () => {
+    const CustomDynamicState = Annotation.Root({
+      messages: MessagesAnnotation.spec.messages,
+      model_preference: Annotation<"basic" | "advanced">({
+        reducer: (_, next) => next,
+        default: () => "basic",
+      }),
+    });
+
+    const dynamicModel = (state: typeof CustomDynamicState.State) => {
+      // Use custom state field to determine model
+      if (state.model_preference === "advanced") {
+        return new FakeToolCallingChatModel({
+          responses: [new AIMessage("advanced")],
+        });
+      }
+
+      return new FakeToolCallingChatModel({
+        responses: [new AIMessage("basic")],
+      });
+    };
+
+    const agent = createReactAgent({
+      llm: dynamicModel,
+      tools: [],
+      stateSchema: CustomDynamicState,
+    });
+
+    expect(
+      await agent.invoke({
+        messages: [new HumanMessage("hello")],
+        model_preference: "advanced",
+      })
+    ).toMatchObject({
+      messages: [{ text: "hello" }, { text: "advanced" }],
+    });
+  });
+
+  it("should handle dynamic model with different prompt types", async () => {
+    const model = new FakeToolCallingChatModel({
+      responses: [new AIMessage("ai response")],
+    });
+    const spyInvoke = vi.spyOn(model, "invoke");
+
+    // Test with string prompt
+    const agent = createReactAgent({
+      llm: () => model,
+      tools: [],
+      prompt: "system_msg",
+    });
+
+    expect(await agent.invoke({ messages: "human_msg" })).toMatchObject({
+      messages: [{ text: "human_msg" }, { text: "ai response" }],
+    });
+
+    expect(spyInvoke).toHaveBeenCalledWith(
+      expect.arrayContaining([
+        expect.objectContaining({ text: "system_msg" }),
+        expect.objectContaining({ text: "human_msg" }),
+      ]),
+      expect.any(Object)
+    );
+
+    spyInvoke.mockClear();
+
+    // Test with callable prompt
+    const dynamicPrompt = (state: typeof MessagesAnnotation.State) => {
+      return [new SystemMessage("system_msg"), ...state.messages];
+    };
+
+    const agent2 = createReactAgent({
+      llm: () => model,
+      tools: [],
+      prompt: dynamicPrompt,
+    });
+
+    expect(await agent2.invoke({ messages: "human_msg" })).toMatchObject({
+      messages: [{ text: "human_msg" }, { text: "ai response" }],
+    });
+
+    expect(spyInvoke).toHaveBeenCalledWith(
+      expect.arrayContaining([
+        expect.objectContaining({ text: "system_msg" }),
+        expect.objectContaining({ text: "human_msg" }),
+      ]),
+      expect.any(Object)
+    );
+  });
+
+  it("should handle dynamic model with structured response format", async () => {
+    const TestResponse = z.object({
+      message: z.string(),
+      confidence: z.number(),
+    });
+
+    const dynamicModel = () => {
+      const expectedResponse = { message: "dynamic response", confidence: 0.9 };
+
+      return new FakeToolCallingChatModel({
+        responses: [new AIMessage("dynamic response")],
+        structuredResponse: expectedResponse,
+      });
+    };
+
+    const agent = createReactAgent({
+      llm: dynamicModel,
+      tools: [],
+      responseFormat: TestResponse,
+    });
+
+    expect(await agent.invoke({ messages: "hello" })).toMatchObject({
+      messages: [{ text: "hello" }, { text: "dynamic response" }],
+      structuredResponse: {
+        message: "dynamic response",
+        confidence: 0.9,
+      },
+    });
+  });
+
+  it("should handle dynamic model that changes available tools based on state", async () => {
+    const toolA = tool(async (args: { x: number }) => `A: ${args.x}`, {
+      name: "tool_a",
+      description: "Tool A.",
+      schema: z.object({ x: z.number() }),
+    });
+
+    const modelA = new FakeToolCallingChatModel({
+      responses: [
+        new AIMessage({
+          content: "",
+          tool_calls: [{ args: { x: 1 }, id: "1", name: "tool_a" }],
+        }),
+        new AIMessage({ content: "use_a please" }),
+      ],
+    });
+
+    const toolB = tool(async (args: { x: number }) => `B: ${args.x}`, {
+      name: "tool_b",
+      description: "Tool B.",
+      schema: z.object({ x: z.number() }),
+    });
+
+    const modelB = new FakeToolCallingChatModel({
+      responses: [
+        new AIMessage({
+          content: "",
+          tool_calls: [{ args: { x: 2 }, id: "1", name: "tool_b" }],
+        }),
+        new AIMessage({ content: "use_b please" }),
+      ],
+    });
+
+    const dynamicModel = (state: typeof MessagesAnnotation.State) => {
+      // Switch tools based on message history
+      if (state.messages.some((msg) => msg.text.includes("use_b"))) {
+        return modelB;
+      }
+
+      return modelA;
+    };
+
+    const agent = createReactAgent({
+      llm: dynamicModel,
+      tools: [toolA, toolB],
+    });
+
+    // Ask to use tool A
+    expect(await agent.invoke({ messages: "use_a" })).toMatchObject({
+      messages: [
+        { text: "use_a" },
+        { tool_calls: [{ args: { x: 1 }, name: "tool_a" }] },
+        { text: "A: 1" },
+        { text: "use_a please" },
+      ],
+    });
+
+    // Ask to use tool B
+    expect(await agent.invoke({ messages: "use_b" })).toMatchObject({
+      messages: [
+        { text: "use_b" },
+        { tool_calls: [{ args: { x: 2 }, name: "tool_b" }] },
+        { text: "B: 2" },
+        { text: "use_b please" },
+      ],
+    });
+  });
+
+  it("should handle error handling in dynamic model", async () => {
+    const failingDynamicModel = (state: typeof MessagesAnnotation.State) => {
+      if (state.messages.at(-1)?.text.includes("fail")) {
+        throw new Error("Dynamic model failed");
+      }
+
+      return new FakeToolCallingChatModel({
+        responses: [new AIMessage("ai response")],
+      });
+    };
+
+    const agent = createReactAgent({
+      llm: failingDynamicModel,
+      tools: [],
+    });
+
+    // Normal operation should work
+    expect(await agent.invoke({ messages: "hello" })).toMatchObject({
+      messages: [{ text: "hello" }, { text: "ai response" }],
+    });
+
+    // Should propagate the error
+    await expect(
+      agent.invoke({ messages: [new HumanMessage("fail now")] })
+    ).rejects.toThrow("Dynamic model failed");
+  });
+
+  it("should produce equivalent results when configured the same", async () => {
+    // Static model
+    const staticAgent = createReactAgent({
+      llm: new FakeToolCallingChatModel({
+        responses: [new AIMessage("ai response")],
+      }),
+      tools: [],
+    });
+
+    // Dynamic model returning the same model
+    const dynamicAgent = createReactAgent({
+      llm: () =>
+        new FakeToolCallingChatModel({
+          responses: [new AIMessage("ai response")],
+        }),
+      tools: [],
+    });
+
+    const inputMsg = { messages: "test message" };
+
+    const staticResult = await staticAgent.invoke(inputMsg);
+    const dynamicResult = await dynamicAgent.invoke(inputMsg);
+
+    // Results should be equivalent (content-wise, IDs may differ)
+    expect(staticResult.messages.length).toEqual(dynamicResult.messages.length);
+    expect(staticResult.messages[0].text).toEqual(
+      dynamicResult.messages[0].text
+    );
+    expect(staticResult.messages[1].text).toEqual(
+      dynamicResult.messages[1].text
+    );
+  });
+
+  it("should receive correct state, not the model input", async () => {
+    const CustomAgentState = Annotation.Root({
+      messages: MessagesAnnotation.spec.messages,
+      custom_field: Annotation<string>,
+    });
+
+    const dynamicModel = vi.fn(
+      () =>
+        new FakeToolCallingChatModel({
+          responses: [new AIMessage("ai response")],
+        })
+    );
+
+    const agent = createReactAgent({
+      llm: dynamicModel,
+      tools: [],
+      stateSchema: CustomAgentState,
+    });
+
+    // Test with initial state
+    const inputState = {
+      messages: [new HumanMessage("hello")],
+      custom_field: "test_value",
+    };
+    await agent.invoke(inputState);
+
+    // The dynamic model function should receive the original state, not the processed model input
+    expect(dynamicModel).toHaveBeenCalledWith(inputState, expect.any(Object));
+  });
+});