fix(langchain): remove _privateState (langchain-ai#9289)

Author: christian-bromann

libs/langchain/src/agents/ReactAgent.ts

@@ -46,11 +46,7 @@ import type {
   JumpTo,
   UserInput,
 } from "./types.js";
-import type {
-  PrivateState,
-  InvokeConfiguration,
-  StreamConfiguration,
-} from "./runtime.js";
+import type { InvokeConfiguration, StreamConfiguration } from "./runtime.js";
 import type {
   AgentMiddleware,
   InferMiddlewareContextInputs,
@@ -249,10 +245,20 @@ export class ReactAgent<
         throw new Error(`Middleware ${m.name} is defined multiple times`);
       }
 
+      const getState = () => {
+        return {
+          ...beforeAgentNode?.getState(),
+          ...beforeModelNode?.getState(),
+          ...afterModelNode?.getState(),
+          ...afterAgentNode?.getState(),
+          ...this.#agentNode.getState(),
+        };
+      };
+
       middlewareNames.add(m.name);
       if (m.beforeAgent) {
         beforeAgentNode = new BeforeAgentNode(m, {
-          getPrivateState: () => this.#agentNode.getState()._privateState,
+          getState,
         });
         const name = `${m.name}.before_agent`;
         beforeAgentNodes.push({
@@ -268,7 +274,7 @@ export class ReactAgent<
       }
       if (m.beforeModel) {
         beforeModelNode = new BeforeModelNode(m, {
-          getPrivateState: () => this.#agentNode.getState()._privateState,
+          getState,
         });
         const name = `${m.name}.before_model`;
         beforeModelNodes.push({
@@ -284,7 +290,7 @@ export class ReactAgent<
       }
       if (m.afterModel) {
         afterModelNode = new AfterModelNode(m, {
-          getPrivateState: () => this.#agentNode.getState()._privateState,
+          getState,
         });
         const name = `${m.name}.after_model`;
         afterModelNodes.push({
@@ -300,7 +306,7 @@ export class ReactAgent<
       }
       if (m.afterAgent) {
         afterAgentNode = new AfterAgentNode(m, {
-          getPrivateState: () => this.#agentNode.getState()._privateState,
+          getState,
         });
         const name = `${m.name}.after_agent`;
         afterAgentNodes.push({
@@ -316,15 +322,7 @@ export class ReactAgent<
       }
 
       if (m.wrapModelCall) {
-        wrapModelCallHookMiddleware.push([
-          m,
-          () => ({
-            ...beforeAgentNode?.getState(),
-            ...beforeModelNode?.getState(),
-            ...afterModelNode?.getState(),
-            ...afterAgentNode?.getState(),
-          }),
-        ]);
+        wrapModelCallHookMiddleware.push([m, getState]);
       }
     }
 
@@ -350,7 +348,6 @@ export class ReactAgent<
       const toolNode = new ToolNode(toolClasses.filter(isClientTool), {
         signal: this.options.signal,
         wrapToolCall: wrapToolCallHandler,
-        getPrivateState: () => this.#agentNode.getState()._privateState,
       });
       allNodeWorkflows.addNode("tools", toolNode);
     }
@@ -944,7 +941,8 @@ export class ReactAgent<
    * Initialize middleware states if not already present in the input state.
    */
   async #initializeMiddlewareStates(
-    state: InvokeStateParameter<StateSchema, TMiddleware>
+    state: InvokeStateParameter<StateSchema, TMiddleware>,
+    config: RunnableConfig
   ): Promise<InvokeStateParameter<StateSchema, TMiddleware>> {
     if (
       !this.options.middleware ||
@@ -959,10 +957,13 @@ export class ReactAgent<
       this.options.middleware,
       state
     );
-    const updatedState = { ...state } as InvokeStateParameter<
-      StateSchema,
-      TMiddleware
-    >;
+    const threadState = await this.#graph
+      .getState(config)
+      .catch(() => ({ values: {} }));
+    const updatedState = {
+      ...threadState.values,
+      ...state,
+    } as InvokeStateParameter<StateSchema, TMiddleware>;
     if (!updatedState) {
       return updatedState;
     }
@@ -977,35 +978,6 @@ export class ReactAgent<
     return updatedState;
   }
 
-  /**
-   * Populate the private state of the agent node from the previous state.
-   */
-  async #populatePrivateState(config?: RunnableConfig) {
-    /**
-     * not needed if thread_id is not provided
-     */
-    if (!config?.configurable?.thread_id) {
-      return;
-    }
-    const prevState = (await this.#graph.getState(config as any)) as {
-      values: {
-        _privateState: PrivateState;
-      };
-    };
-
-    /**
-     * not need if state is empty
-     */
-    if (!prevState.values._privateState) {
-      return;
-    }
-
-    this.#agentNode.setState({
-      structuredResponse: undefined,
-      _privateState: prevState.values._privateState,
-    });
-  }
-
   /**
    * Executes the agent with the given state and returns the final state after all processing.
    *
@@ -1061,8 +1033,10 @@ export class ReactAgent<
       StructuredResponseFormat,
       TMiddleware
     >;
-    const initializedState = await this.#initializeMiddlewareStates(state);
-    await this.#populatePrivateState(config);
+    const initializedState = await this.#initializeMiddlewareStates(
+      state,
+      config as RunnableConfig
+    );
 
     return this.#graph.invoke(
       initializedState,
@@ -1120,7 +1094,10 @@ export class ReactAgent<
         InferMiddlewareContextInputs<TMiddleware>
     >
   ): Promise<IterableReadableStream<any>> {
-    const initializedState = await this.#initializeMiddlewareStates(state);
+    const initializedState = await this.#initializeMiddlewareStates(
+      state,
+      config as RunnableConfig
+    );
     return this.#graph.stream(initializedState, config as Record<string, any>);
   }
 

libs/langchain/src/agents/annotation.ts

@@ -30,7 +30,12 @@ export function createAgentAnnotationConditional<
   const zodSchema: Record<string, any> = {
     messages: withLangGraph(z.custom<BaseMessage[]>(), MessagesZodMeta),
     jumpTo: z
-      .union([z.literal("model_request"), z.literal("tools"), z.undefined()])
+      .union([
+        z.literal("model_request"),
+        z.literal("tools"),
+        z.literal("end"),
+        z.undefined(),
+      ])
       .optional(),
   };
 

libs/langchain/src/agents/middleware/modelCallLimit.ts

@@ -27,6 +27,14 @@ export type ModelCallLimitMiddlewareConfig = Partial<
   InferInteropZodInput<typeof contextSchema>
 >;
 
+/**
+ * Middleware state schema to track the number of model calls made at the thread and run level.
+ */
+const stateSchema = z.object({
+  threadModelCallCount: z.number().default(0),
+  runModelCallCount: z.number().default(0),
+});
+
 /**
  * Error thrown when the model call limit is exceeded.
  *
@@ -133,6 +141,7 @@ export function modelCallLimitMiddleware(
   return createMiddleware({
     name: "ModelCallLimitMiddleware",
     contextSchema,
+    stateSchema,
     beforeModel: {
       canJumpTo: ["end"],
       hook: (state, runtime) => {
@@ -145,13 +154,12 @@ export function modelCallLimitMiddleware(
         const runLimit =
           runtime.context.runLimit ?? middlewareOptions?.runLimit;
 
-        if (
-          typeof threadLimit === "number" &&
-          threadLimit <= runtime.threadLevelCallCount
-        ) {
+        const threadCount = state.threadModelCallCount;
+        const runCount = state.runModelCallCount;
+        if (typeof threadLimit === "number" && threadLimit <= threadCount) {
           const error = new ModelCallLimitMiddlewareError({
             threadLimit,
-            threadCount: runtime.threadLevelCallCount,
+            threadCount,
           });
           if (exitBehavior === "end") {
             return {
@@ -162,13 +170,10 @@ export function modelCallLimitMiddleware(
 
           throw error;
         }
-        if (
-          typeof runLimit === "number" &&
-          runLimit <= runtime.runModelCallCount
-        ) {
+        if (typeof runLimit === "number" && runLimit <= runCount) {
           const error = new ModelCallLimitMiddlewareError({
             runLimit,
-            runCount: runtime.runModelCallCount,
+            runCount,
           });
           if (exitBehavior === "end") {
             return {
@@ -183,5 +188,9 @@ export function modelCallLimitMiddleware(
         return state;
       },
     },
+    afterModel: (state) => ({
+      runModelCallCount: state.runModelCallCount + 1,
+      threadModelCallCount: state.threadModelCallCount + 1,
+    }),
   });
 }

libs/langchain/src/agents/middleware/tests/modelCallLimit.test.ts

@@ -52,6 +52,10 @@ const tools = [
     name: "tool_1",
     description: "tool_1",
   }),
+  tool(() => "barfoo", {
+    name: "tool_2",
+    description: "tool_2",
+  }),
 ];
 
 describe("ModelCallLimitMiddleware", () => {
@@ -143,14 +147,19 @@ describe("ModelCallLimitMiddleware", () => {
           checkpointer,
         });
         if (exitBehavior === "throw") {
-          await expect(
-            agent2.invoke({ messages: ["Hello, world!"] }, config)
-          ).resolves.not.toThrow();
+          const result = await agent2.invoke(
+            { messages: ["Hello, world!"] },
+            config
+          );
+          await expect(result.runModelCallCount).toBe(3);
+          await expect(result.threadModelCallCount).toBe(3);
         } else {
           const result = await agent2.invoke(
             { messages: ["Hello, world!"] },
             config
           );
+          await expect(result.runModelCallCount).toBe(3);
+          await expect(result.threadModelCallCount).toBe(3);
           expect(result.messages.at(-1)?.content).not.toContain(
             "Model call limits exceeded"
           );