feat(langchain): Add type inference for ReactAgent stream method (#9368)

Author: sapenlei

libs/langchain/src/agents/ReactAgent.ts

@@ -11,6 +11,8 @@ import {
   CompiledStateGraph,
   type GetStateOptions,
   type LangGraphRunnableConfig,
+  type StreamMode,
+  type StreamOutputMap,
 } from "@langchain/langgraph";
 import type { CheckpointListOptions } from "@langchain/langgraph-checkpoint";
 import { ToolMessage, AIMessage } from "@langchain/core/messages";
@@ -1076,18 +1078,39 @@ export class ReactAgent<
    * }
    * ```
    */
-  async stream(
+  async stream<
+    TStreamMode extends StreamMode | StreamMode[] | undefined,
+    TEncoding extends "text/event-stream" | undefined
+  >(
     state: InvokeStateParameter<StateSchema, TMiddleware>,
     config?: StreamConfiguration<
       InferContextInput<ContextSchema> &
-        InferMiddlewareContextInputs<TMiddleware>
+        InferMiddlewareContextInputs<TMiddleware>,
+      TStreamMode,
+      TEncoding
     >
-  ): Promise<IterableReadableStream<any>> {
+  ) {
     const initializedState = await this.#initializeMiddlewareStates(
       state,
       config as RunnableConfig
     );
-    return this.#graph.stream(initializedState, config as Record<string, any>);
+    return this.#graph.stream(
+      initializedState,
+      config as Record<string, any>
+    ) as Promise<
+      IterableReadableStream<
+        StreamOutputMap<
+          TStreamMode,
+          false,
+          MergedAgentState<StateSchema, StructuredResponseFormat, TMiddleware>,
+          MergedAgentState<StateSchema, StructuredResponseFormat, TMiddleware>,
+          string,
+          unknown,
+          unknown,
+          TEncoding
+        >
+      >
+    >;
   }
 
   /**
@@ -1149,7 +1172,9 @@ export class ReactAgent<
     state: InvokeStateParameter<StateSchema, TMiddleware>,
     config?: StreamConfiguration<
       InferContextInput<ContextSchema> &
-        InferMiddlewareContextInputs<TMiddleware>
+        InferMiddlewareContextInputs<TMiddleware>,
+      StreamMode | StreamMode[] | undefined,
+      "text/event-stream" | undefined
     > & { version?: "v1" | "v2" },
     streamOptions?: Parameters<Runnable["streamEvents"]>[2]
   ): IterableReadableStream<StreamEvent> {

libs/langchain/src/agents/runtime.ts

@@ -4,6 +4,7 @@ import type { InteropZodDefault } from "@langchain/core/utils/types";
 import type {
   Runtime as LangGraphRuntime,
   PregelOptions,
+  StreamMode,
 } from "@langchain/langgraph";
 import type { BaseMessage } from "@langchain/core/messages";
 import type { BaseCallbackConfig } from "@langchain/core/callbacks/manager";
@@ -151,12 +152,21 @@ export type InvokeConfiguration<ContextSchema extends Record<string, any>> =
         Partial<Pick<PregelOptions<any, any, any>, CreateAgentPregelOptions>> &
         WithMaybeContext<ContextSchema>;
 
-export type StreamConfiguration<ContextSchema extends Record<string, any>> =
+export type StreamConfiguration<
+  ContextSchema extends Record<string, any>,
+  TStreamMode extends StreamMode | StreamMode[] | undefined,
+  TEncoding extends "text/event-stream" | undefined
+> =
   /**
    * If the context schema is a default object, `context` can be optional
    */
   ContextSchema extends InteropZodDefault<any>
-    ? Partial<Pick<PregelOptions<any, any, any>, CreateAgentPregelOptions>> & {
+    ? Partial<
+        Pick<
+          PregelOptions<any, any, any, TStreamMode, boolean, TEncoding>,
+          CreateAgentPregelOptions
+        >
+      > & {
         context?: Partial<ContextSchema>;
       }
     : /**
@@ -165,15 +175,15 @@ export type StreamConfiguration<ContextSchema extends Record<string, any>> =
     IsAllOptional<ContextSchema> extends true
     ? Partial<
         Pick<
-          PregelOptions<any, any, any>,
+          PregelOptions<any, any, any, TStreamMode, boolean, TEncoding>,
           CreateAgentPregelOptions | CreateAgentPregelStreamOptions
         >
       > & {
         context?: Partial<ContextSchema>;
       }
     : Partial<
         Pick<
-          PregelOptions<any, any, any>,
+          PregelOptions<any, any, any, TStreamMode, boolean, TEncoding>,
           CreateAgentPregelOptions | CreateAgentPregelStreamOptions
         >
       > &

libs/langchain/src/agents/tests/reactAgent.test-d.ts

@@ -4,7 +4,8 @@ import { LanguageModelLike } from "@langchain/core/language_models/base";
 import { describe, it, expectTypeOf } from "vitest";
 import type { IterableReadableStream } from "@langchain/core/utils/stream";
 
-import { createAgent } from "../index.js";
+import { type BuiltInState, createAgent } from "../index.js";
+import type { StreamOutputMap } from "@langchain/langgraph";
 
 describe("reactAgent", () => {
   it("should require model as only required property", async () => {
@@ -79,8 +80,48 @@ describe("reactAgent", () => {
         recursionLimit: 10,
       }
     );
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    expectTypeOf(stream).toEqualTypeOf<IterableReadableStream<any>>();
+    expectTypeOf(stream).toEqualTypeOf<
+      IterableReadableStream<
+        StreamOutputMap<
+          "values" | "updates" | "messages",
+          false,
+          Record<string, unknown>,
+          Record<string, unknown>,
+          string,
+          unknown,
+          unknown,
+          "text/event-stream"
+        >
+      >
+    >();
+
+    for await (const chunk of stream) {
+      expectTypeOf(chunk).toEqualTypeOf<Uint8Array>();
+    }
+
+    const multiModeStream = await agent.stream(
+      {
+        messages: [new HumanMessage("Hello, world!")],
+      },
+      {
+        streamMode: ["updates", "messages", "values"],
+      }
+    );
+
+    for await (const chunk of multiModeStream) {
+      const [mode, value] = chunk;
+      expectTypeOf(mode).toEqualTypeOf<"updates" | "messages" | "values">();
+      if (mode === "messages") {
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        expectTypeOf(value).toEqualTypeOf<[BaseMessage, Record<string, any>]>();
+      } else if (mode === "updates") {
+        expectTypeOf(value).toEqualTypeOf<
+          Record<string, Omit<BuiltInState, "jumpTo">>
+        >();
+      } else {
+        expectTypeOf(value.messages).toEqualTypeOf<BaseMessage[]>();
+      }
+    }
 
     await agent.invoke(
       {