feat(langgraph): introduce strict typing for .stream() methods

Supersedes #512

Author: dqbd

libs/langgraph/.eslintrc.cjs

@@ -39,7 +39,7 @@ module.exports = {
     "import/extensions": [2, "ignorePackages"],
     "import/no-extraneous-dependencies": [
       "error",
-      { devDependencies: ["**/*.test.ts"] },
+      { devDependencies: ["**/*.test.ts", "**/*.test-d.ts"] },
     ],
     "import/no-unresolved": 0,
     "import/prefer-default-export": 0,

libs/langgraph/src/func/index.ts

@@ -146,7 +146,10 @@ export interface EntrypointFunction {
     },
     Record<string, unknown>,
     InputT,
-    EntrypointReturnT<OutputT>
+    EntrypointReturnT<OutputT>,
+    // Because the update type is an return type union of tasks + entrypoint,
+    // thus we can't type it properly.
+    any // eslint-disable-line @typescript-eslint/no-explicit-any
   >;
 
   /**

libs/langgraph/src/pregel/index.ts

@@ -72,6 +72,7 @@ import {
   SingleChannelSubscriptionOptions,
   MultipleChannelSubscriptionOptions,
   GetStateOptions,
+  type StreamOutputMap,
 } from "./types.js";
 import {
   GraphRecursionError,
@@ -118,10 +119,6 @@ import {
 type WriteValue = Runnable | RunnableFunc<unknown, unknown> | unknown;
 type StreamEventsOptions = Parameters<Runnable["streamEvents"]>[2];
 
-function isString(value: unknown): value is string {
-  return typeof value === "string";
-}
-
 /**
  * Utility class for working with channels in the Pregel system.
  * Provides static methods for subscribing to channels and writing to them.
@@ -201,7 +198,7 @@ export class Channel {
 
     let channelMappingOrArray: string[] | Record<string, string>;
 
-    if (isString(channels)) {
+    if (typeof channels === "string") {
       if (key) {
         channelMappingOrArray = { [key]: channels };
       } else {
@@ -349,7 +346,9 @@ export class Pregel<
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     ConfigurableFieldType extends Record<string, any> = StrRecord<string, any>,
     InputType = PregelInputType,
-    OutputType = PregelOutputType
+    OutputType = PregelOutputType,
+    StreamUpdatesType = InputType,
+    StreamValuesType = OutputType
   >
   extends Runnable<
     InputType | Command | null,
@@ -1645,10 +1644,31 @@ export class Pregel<
    * @param options - Configuration options for streaming
    * @returns An async iterable stream of graph state updates
    */
-  override async stream(
+  // @ts-expect-error Return type of `stream()` differs from `invoke()`, which is expected.
+  override async stream<
+    TStreamMode extends StreamMode | StreamMode[] | undefined,
+    TSubgraphs extends boolean
+  >(
     input: InputType | Command | null,
-    options?: Partial<PregelOptions<Nodes, Channels, ConfigurableFieldType>>
-  ): Promise<IterableReadableStream<PregelOutputType>> {
+    options?: Partial<
+      PregelOptions<
+        Nodes,
+        Channels,
+        ConfigurableFieldType,
+        TStreamMode,
+        TSubgraphs
+      >
+    >
+  ): Promise<
+    IterableReadableStream<
+      StreamOutputMap<
+        TStreamMode,
+        TSubgraphs,
+        StreamUpdatesType,
+        StreamValuesType
+      >
+    >
+  > {
     // The ensureConfig method called internally defaults recursionLimit to 25 if not
     // passed directly in `options`.
     // There is currently no way in _streamIterator to determine whether this was
@@ -1665,7 +1685,14 @@ export class Pregel<
     };
 
     return new IterableReadableStreamWithAbortSignal(
-      await super.stream(input, config),
+      (await super.stream(input, config)) as IterableReadableStream<
+        StreamOutputMap<
+          TStreamMode,
+          TSubgraphs,
+          StreamUpdatesType,
+          StreamValuesType
+        >
+      >,
       abortController
     );
   }
@@ -2025,7 +2052,7 @@ export class Pregel<
       chunks.push(chunk);
     }
     if (streamMode === "values") {
-      return chunks[chunks.length - 1];
+      return chunks[chunks.length - 1] as OutputType;
     }
     return chunks as OutputType;
   }

libs/langgraph/src/pregel/types.ts

@@ -9,6 +9,7 @@ import type {
 } from "@langchain/langgraph-checkpoint";
 import { Graph as DrawableGraph } from "@langchain/core/runnables/graph";
 import { IterableReadableStream } from "@langchain/core/utils/stream";
+import type { BaseMessage } from "@langchain/core/messages";
 import type { BaseChannel } from "../channels/base.js";
 import type { PregelNode } from "./read.js";
 import { RetryPolicy } from "./utils/index.js";
@@ -27,6 +28,68 @@ export type PregelInputType = any;
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export type PregelOutputType = any;
 
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type StreamMessageOutput = [BaseMessage, Record<string, any>];
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type StreamCustomOutput = any;
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type StreamDebugOutput = Record<string, any>;
+
+type DefaultStreamMode = "updates";
+
+export type StreamOutputMap<
+  TStreamMode extends StreamMode | StreamMode[] | undefined,
+  TStreamSubgraphs extends boolean,
+  StreamUpdates,
+  StreamValues
+> = (
+  undefined extends TStreamMode
+    ? []
+    : StreamMode | StreamMode[] extends TStreamMode
+    ? TStreamMode extends StreamMode[]
+      ? TStreamMode[number]
+      : TStreamMode
+    : TStreamMode extends StreamMode[]
+    ? TStreamMode[number]
+    : []
+) extends infer Multiple extends StreamMode
+  ? [TStreamSubgraphs] extends [true]
+    ? {
+        values: [string[], "values", StreamValues];
+        updates: [string[], "updates", Record<string, StreamUpdates>];
+        messages: [string[], "messages", StreamMessageOutput];
+        custom: [string[], "custom", StreamCustomOutput];
+        debug: [string[], "debug", StreamDebugOutput];
+      }[Multiple]
+    : {
+        values: ["values", StreamValues];
+        updates: ["updates", Record<string, StreamUpdates>];
+        messages: ["messages", StreamMessageOutput];
+        custom: ["custom", StreamCustomOutput];
+        debug: ["debug", StreamDebugOutput];
+      }[Multiple]
+  : (
+      undefined extends TStreamMode ? DefaultStreamMode : TStreamMode
+    ) extends infer Single extends StreamMode
+  ? [TStreamSubgraphs] extends [true]
+    ? {
+        values: [string[], StreamValues];
+        updates: [string[], Record<string, StreamUpdates>];
+        messages: [string[], StreamMessageOutput];
+        custom: [string[], StreamCustomOutput];
+        debug: [string[], StreamDebugOutput];
+      }[Single]
+    : {
+        values: StreamValues;
+        updates: Record<string, StreamUpdates>;
+        messages: StreamMessageOutput;
+        custom: StreamCustomOutput;
+        debug: StreamDebugOutput;
+      }[Single]
+  : never;
+
 /**
  * Configuration options for executing a Pregel graph.
  * These options control how the graph executes, what data is streamed, and how interrupts are handled.
@@ -39,7 +102,12 @@ export interface PregelOptions<
   Nodes extends StrRecord<string, PregelNode>,
   Channels extends StrRecord<string, BaseChannel | ManagedValueSpec>,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  ConfigurableFieldType extends Record<string, any> = Record<string, any>
+  ConfigurableFieldType extends Record<string, any> = Record<string, any>,
+  TStreamMode extends StreamMode | StreamMode[] | undefined =
+    | StreamMode
+    | StreamMode[]
+    | undefined,
+  TSubgraphs extends boolean = boolean
 > extends RunnableConfig<ConfigurableFieldType> {
   /**
    * Controls what information is streamed during graph execution.
@@ -63,7 +131,7 @@ export interface PregelOptions<
    *
    * @default ["values"]
    */
-  streamMode?: StreamMode | StreamMode[];
+  streamMode?: TStreamMode;
 
   /**
    * Specifies which channel keys to retrieve from the checkpoint when resuming execution.
@@ -140,7 +208,7 @@ export interface PregelOptions<
    *
    * @default false
    */
-  subgraphs?: boolean;
+  subgraphs?: TSubgraphs;
 
   /**
    * A shared value store that allows you to store and retrieve state across

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

@@ -208,9 +208,7 @@ describe("createReactAgent", () => {
     });
 
     const stream = await reactAgent.stream(
-      {
-        messages: [new HumanMessage("What's the weather like in SF?")],
-      },
+      { messages: [new HumanMessage("What's the weather like in SF?")] },
       { configurable: { thread_id: "foo" }, streamMode: "values" }
     );
     const fullResponse = [];
@@ -226,7 +224,7 @@ describe("createReactAgent", () => {
 
     const lastMessage = endState.messages[endState.messages.length - 1];
     expect(lastMessage._getType()).toBe("ai");
-    expect(lastMessage.content.toLowerCase()).toContain("not too cold");
+    expect(lastMessage.text.toLowerCase()).toContain("not too cold");
     const stream2 = await reactAgent.stream(
       {
         messages: [new HumanMessage("What about NYC?")],
@@ -245,6 +243,6 @@ describe("createReactAgent", () => {
 
     const lastMessage2 = endState.messages[endState.messages.length - 1];
     expect(lastMessage2._getType()).toBe("ai");
-    expect(lastMessage2.content.toLowerCase()).toContain("not too cold");
+    expect(lastMessage2.text.toLowerCase()).toContain("not too cold");
   });
 });