feat(langgraph): add "tasks" and "checkpoints" stream mode

libs/langgraph/src/pregel/debug.ts

@@ -88,10 +88,8 @@ export function* _readChannels<Value>(
 }
 
 export function* mapDebugTasks<N extends PropertyKey, C extends PropertyKey>(
-  step: number,
   tasks: readonly PregelExecutableTask<N, C>[]
 ) {
-  const ts = new Date().toISOString();
   for (const { id, name, input, config, triggers, writes } of tasks) {
     if (config?.tags?.includes(TAG_HIDDEN)) continue;
 
@@ -102,46 +100,28 @@ export function* mapDebugTasks<N extends PropertyKey, C extends PropertyKey>(
       .map(([, v]) => {
         return v;
       });
-    yield {
-      type: "task",
-      timestamp: ts,
-      step,
-      payload: {
-        id,
-        name,
-        input,
-        triggers,
-        interrupts,
-      },
-    };
+    yield { id, name, input, triggers, interrupts };
   }
 }
 
 export function* mapDebugTaskResults<
   N extends PropertyKey,
   C extends PropertyKey
 >(
-  step: number,
   tasks: readonly [PregelExecutableTask<N, C>, PendingWrite<C>[]][],
   streamChannels: PropertyKey | Array<PropertyKey>
 ) {
-  const ts = new Date().toISOString();
   for (const [{ id, name, config }, writes] of tasks) {
     if (config?.tags?.includes(TAG_HIDDEN)) continue;
     yield {
-      type: "task_result",
-      timestamp: ts,
-      step,
-      payload: {
-        id,
-        name,
-        result: writes.filter(([channel]) => {
-          return Array.isArray(streamChannels)
-            ? streamChannels.includes(channel)
-            : channel === streamChannels;
-        }),
-        interrupts: writes.filter((w) => w[0] === INTERRUPT).map((w) => w[1]),
-      },
+      id,
+      name,
+      result: writes.filter(([channel]) => {
+        return Array.isArray(streamChannels)
+          ? streamChannels.includes(channel)
+          : channel === streamChannels;
+      }),
+      interrupts: writes.filter((w) => w[0] === INTERRUPT).map((w) => w[1]),
     };
   }
 }
@@ -152,7 +132,6 @@ export function* mapDebugCheckpoint<
   N extends PropertyKey,
   C extends PropertyKey
 >(
-  step: number,
   config: RunnableConfig,
   channels: Record<string, BaseChannel>,
   streamChannels: string | string[],
@@ -213,19 +192,13 @@ export function* mapDebugCheckpoint<
     };
   }
 
-  const ts = new Date().toISOString();
   yield {
-    type: "checkpoint",
-    timestamp: ts,
-    step,
-    payload: {
-      config: formatConfig(config),
-      values: readChannels(channels, streamChannels),
-      metadata,
-      next: tasks.map((task) => task.name),
-      tasks: tasksWithWrites(tasks, pendingWrites, taskStates, outputKeys),
-      parentConfig: parentConfig ? formatConfig(parentConfig) : undefined,
-    },
+    config: formatConfig(config),
+    values: readChannels(channels, streamChannels),
+    metadata,
+    next: tasks.map((task) => task.name),
+    tasks: tasksWithWrites(tasks, pendingWrites, taskStates, outputKeys),
+    parentConfig: parentConfig ? formatConfig(parentConfig) : undefined,
   };
 }
 

libs/langgraph/src/pregel/loop.ts

@@ -633,8 +633,8 @@ export class PregelLoop {
         this._emit(
           gatherIteratorSync(
             prefixGenerator(
-              mapDebugTaskResults(this.step, [[task, writes]], this.streamKeys),
-              "debug"
+              mapDebugTaskResults([[task, writes]], this.streamKeys),
+              "tasks"
             )
           )
         );
@@ -782,7 +782,6 @@ export class PregelLoop {
         await gatherIterator(
           prefixGenerator(
             mapDebugCheckpoint(
-              this.step - 1, // printing checkpoint for previous step
               this.checkpointConfig,
               this.channels,
               this.streamKeys,
@@ -792,7 +791,7 @@ export class PregelLoop {
               this.prevCheckpointConfig,
               this.outputKeys
             ),
-            "debug"
+            "checkpoints"
           )
         )
       );
@@ -838,10 +837,7 @@ export class PregelLoop {
 
     // Produce debug output
     const debugOutput = await gatherIterator(
-      prefixGenerator(
-        mapDebugTasks(this.step, Object.values(this.tasks)),
-        "debug"
-      )
+      prefixGenerator(mapDebugTasks(Object.values(this.tasks)), "tasks")
     );
     this._emit(debugOutput);
 
@@ -952,9 +948,7 @@ export class PregelLoop {
     }
 
     this._emit(
-      gatherIteratorSync(
-        prefixGenerator(mapDebugTasks(this.step, [pushed]), "debug")
-      )
+      gatherIteratorSync(prefixGenerator(mapDebugTasks([pushed]), "tasks"))
     );
 
     if (this.debug) printStepTasks(this.step, [pushed]);
@@ -1127,9 +1121,38 @@ export class PregelLoop {
   }
 
   protected _emit(values: [StreamMode, unknown][]) {
-    for (const chunk of values) {
-      if (this.stream.modes.has(chunk[0])) {
-        this.stream.push([this.checkpointNamespace, ...chunk]);
+    for (const [mode, payload] of values) {
+      if (this.stream.modes.has(mode)) {
+        this.stream.push([this.checkpointNamespace, mode, payload]);
+      }
+
+      // debug mode is a "checkpoints" or "tasks" wrapped in an object
+      // TODO: consider deprecating this in 1.x
+      if (
+        (mode === "checkpoints" || mode === "tasks") &&
+        this.stream.modes.has("debug")
+      ) {
+        const step = mode === "checkpoints" ? this.step - 1 : this.step;
+        const timestamp = new Date().toISOString();
+        const type = (() => {
+          if (mode === "checkpoints") {
+            return "checkpoint";
+          } else if (
+            typeof payload === "object" &&
+            payload != null &&
+            "result" in payload
+          ) {
+            return "task_result";
+          } else {
+            return "task";
+          }
+        })();
+
+        this.stream.push([
+          this.checkpointNamespace,
+          "debug",
+          { step, type, timestamp, payload },
+        ]);
       }
     }
   }

libs/langgraph/src/pregel/types.ts

@@ -21,7 +21,14 @@ import { LangGraphRunnableConfig } from "./runnable_types.js";
 /**
  * Selects the type of output you'll receive when streaming from the graph. See [Streaming](/langgraphjs/how-tos/#streaming) for more details.
  */
-export type StreamMode = "values" | "updates" | "debug" | "messages" | "custom";
+export type StreamMode =
+  | "values"
+  | "updates"
+  | "debug"
+  | "messages"
+  | "checkpoints"
+  | "tasks"
+  | "custom";
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export type PregelInputType = any;
@@ -38,14 +45,43 @@ type StreamCustomOutput = any;
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type StreamDebugOutput = Record<string, any>;
 
+type StreamCheckpointsOutput<StreamValues> = {
+  values: StreamValues;
+  next: string[];
+  config: RunnableConfig;
+  metadata?: CheckpointMetadata;
+  parentConfig?: RunnableConfig | undefined;
+  tasks: PregelTaskDescription[];
+};
+
+interface StreamTasksOutputBase {
+  id: string;
+  name: string;
+  interrupts: Interrupt[];
+}
+
+interface StreamTasksCreateOutput<StreamValues> extends StreamTasksOutputBase {
+  input: StreamValues;
+  triggers: string[];
+}
+
+interface StreamTasksResultOutput<Keys, StreamUpdates>
+  extends StreamTasksOutputBase {
+  result: [Keys, StreamUpdates][];
+}
+
+type StreamTasksOutput<StreamUpdates, StreamValues, Nodes = string> =
+  | StreamTasksCreateOutput<StreamValues>
+  | StreamTasksResultOutput<Nodes, StreamUpdates>;
+
 type DefaultStreamMode = "updates";
 
 export type StreamOutputMap<
   TStreamMode extends StreamMode | StreamMode[] | undefined,
   TStreamSubgraphs extends boolean,
   StreamUpdates,
   StreamValues,
-  Nodes = string
+  Nodes
 > = (
   undefined extends TStreamMode
     ? []
@@ -67,6 +103,16 @@ export type StreamOutputMap<
         ];
         messages: [string[], "messages", StreamMessageOutput];
         custom: [string[], "custom", StreamCustomOutput];
+        checkpoints: [
+          string[],
+          "checkpoints",
+          StreamCheckpointsOutput<StreamValues>
+        ];
+        tasks: [
+          string[],
+          "tasks",
+          StreamTasksOutput<StreamUpdates, StreamValues>
+        ];
         debug: [string[], "debug", StreamDebugOutput];
       }[Multiple]
     : {
@@ -77,6 +123,8 @@ export type StreamOutputMap<
         ];
         messages: ["messages", StreamMessageOutput];
         custom: ["custom", StreamCustomOutput];
+        checkpoints: ["checkpoints", StreamCheckpointsOutput<StreamValues>];
+        tasks: ["tasks", StreamTasksOutput<StreamUpdates, StreamValues, Nodes>];
         debug: ["debug", StreamDebugOutput];
       }[Multiple]
   : (
@@ -91,13 +139,20 @@ export type StreamOutputMap<
         ];
         messages: [string[], StreamMessageOutput];
         custom: [string[], StreamCustomOutput];
+        checkpoints: [string[], StreamCheckpointsOutput<StreamValues>];
+        tasks: [
+          string[],
+          StreamTasksOutput<StreamUpdates, StreamValues, Nodes>
+        ];
         debug: [string[], StreamDebugOutput];
       }[Single]
     : {
         values: StreamValues;
         updates: Record<Nodes extends string ? Nodes : string, StreamUpdates>;
         messages: StreamMessageOutput;
         custom: StreamCustomOutput;
+        checkpoints: StreamCheckpointsOutput<StreamValues>;
+        tasks: StreamTasksOutput<StreamUpdates, StreamValues, Nodes>;
         debug: StreamDebugOutput;
       }[Single]
   : never;

libs/langgraph/src/tests/pregel.test-d.ts

@@ -67,7 +67,10 @@ it("state graph annotation", async () => {
 
   expectTypeOf(
     await gatherIterator(
-      graph.stream(input, { streamMode: ["values"], subgraphs: true })
+      graph.stream(input, {
+        streamMode: ["values"],
+        subgraphs: true,
+      })
     )
   ).toExtend<[string[], "values", { foo: string[] }][]>();
 
@@ -83,7 +86,10 @@ it("state graph annotation", async () => {
 
   expectTypeOf(
     await gatherIterator(
-      graph.stream(input, { streamMode: ["updates"], subgraphs: true })
+      graph.stream(input, {
+        streamMode: ["updates"],
+        subgraphs: true,
+      })
     )
   ).toExtend<
     [
@@ -145,30 +151,14 @@ it("state graph annotation", async () => {
       | ["debug", Record<string, any>]
       | ["messages", [BaseMessage, Record<string, any>]]
       | ["custom", any]
-    )[]
-  >();
-
-  expectTypeOf(
-    await gatherIterator(
-      graph.stream(input, {
-        streamMode: ["updates", "values"] as
-          | StreamMode
-          | StreamMode[]
-          | undefined,
-        subgraphs: true,
-      })
-    )
-  ).toExtend<
-    (
+      | ["checkpoints", { values: { foo: string[] } }]
       | [
-          string[],
-          "updates",
-          Record<"one" | "two" | "three", { foo?: string[] | string }>
+          "tasks",
+          { id: string; name: string } & (
+            | { input: unknown }
+            | { result: [string, unknown][] }
+          )
         ]
-      | [string[], "values", { foo: string[] }]
-      | [string[], "debug", Record<string, any>]
-      | [string[], "messages", [BaseMessage, Record<string, any>]]
-      | [string[], "custom", any]
     )[]
   >();
 });
@@ -296,6 +286,14 @@ it("state graph zod", async () => {
       | ["debug", Record<string, any>]
       | ["messages", [BaseMessage, Record<string, any>]]
       | ["custom", any]
+      | ["checkpoints", { values: { foo: string[] } }]
+      | [
+          "tasks",
+          { id: string; name: string } & (
+            | { input: unknown }
+            | { result: [string, unknown][] }
+          )
+        ]
     )[]
   >();
 
@@ -320,6 +318,15 @@ it("state graph zod", async () => {
       | [string[], "debug", Record<string, any>]
       | [string[], "messages", [BaseMessage, Record<string, any>]]
       | [string[], "custom", any]
+      | [string[], "checkpoints", { values: { foo: string[] } }]
+      | [
+          string[],
+          "tasks",
+          { id: string; name: string } & (
+            | { input: unknown }
+            | { result: [string, unknown][] }
+          )
+        ]
     )[]
   >();
 });