@langchain/langgraph 0.2.19 → 0.2.20

package/dist/pregel/index.js

@@ -1,7 +1,6 @@
 /* eslint-disable no-param-reassign */
 import { Runnable, RunnableSequence, getCallbackManagerForConfig, mergeConfigs, patchConfig, _coerceToRunnable, } from "@langchain/core/runnables";
 import { compareChannelVersions, copyCheckpoint, emptyCheckpoint, uuid5, } from "@langchain/langgraph-checkpoint";
-import Deque from "double-ended-queue";
 import { createCheckpoint, emptyChannels, isBaseChannel, } from "../channels/base.js";
 import { PregelNode } from "./read.js";
 import { validateGraph, validateKeys } from "./validate.js";
@@ -13,11 +12,12 @@ import { GraphRecursionError, GraphValueError, InvalidUpdateError, isGraphInterr
 import { _prepareNextTasks, _localRead, _applyWrites, } from "./algo.js";
 import { _coerceToDict, getNewChannelVersions, patchCheckpointMap, } from "./utils/index.js";
 import { findSubgraphPregel } from "./utils/subgraph.js";
-import { PregelLoop, StreamProtocol } from "./loop.js";
+import { PregelLoop, IterableReadableWritableStream } from "./loop.js";
 import { executeTasksWithRetry } from "./retry.js";
 import { ChannelKeyPlaceholder, isConfiguredManagedValue, ManagedValueMapping, NoopManagedValue, } from "../managed/base.js";
 import { gatherIterator, patchConfigurable } from "../utils.js";
 import { ensureLangGraphConfig } from "./utils/config.js";
+import { StreamMessagesHandler } from "./messages.js";
 function isString(value) {
     return typeof value === "string";
 }
@@ -572,6 +572,7 @@ export class Pregel extends Runnable {
     }
     _defaults(config) {
         const { debug, streamMode, inputKeys, outputKeys, interruptAfter, interruptBefore, ...rest } = config;
+        let streamModeSingle = true;
         const defaultDebug = debug !== undefined ? debug : this.debug;
         let defaultOutputKeys = outputKeys;
         if (defaultOutputKeys === undefined) {
@@ -592,9 +593,11 @@ export class Pregel extends Runnable {
         let defaultStreamMode;
         if (streamMode !== undefined) {
             defaultStreamMode = Array.isArray(streamMode) ? streamMode : [streamMode];
+            streamModeSingle = typeof streamMode === "string";
         }
         else {
             defaultStreamMode = this.streamMode;
+            streamModeSingle = true;
         }
         // if being called as a node in another graph, always use values mode
         if (config.configurable?.[CONFIG_KEY_TASK_ID] !== undefined) {
@@ -622,6 +625,7 @@ export class Pregel extends Runnable {
             defaultInterruptAfter,
             defaultCheckpointer,
             defaultStore,
+            streamModeSingle,
         ];
     }
     /**
@@ -698,135 +702,180 @@ export class Pregel extends Runnable {
             inputConfig.configurable === undefined) {
             throw new Error(`Checkpointer requires one or more of the following "configurable" keys: "thread_id", "checkpoint_ns", "checkpoint_id"`);
         }
-        const callbackManager = await getCallbackManagerForConfig(inputConfig);
-        const runManager = await callbackManager?.handleChainStart(this.toJSON(), _coerceToDict(input, "input"), inputConfig.runId, undefined, undefined, undefined, inputConfig?.runName ?? this.getName());
-        delete inputConfig.runId;
+        const { runId, ...restConfig } = inputConfig;
         // assign defaults
-        const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer, store,] = this._defaults(inputConfig);
+        const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer, store, streamModeSingle,] = this._defaults(restConfig);
+        const stream = new IterableReadableWritableStream({
+            modes: new Set(streamMode),
+        });
+        // set up messages stream mode
+        if (streamMode.includes("messages")) {
+            const messageStreamer = new StreamMessagesHandler((chunk) => stream.push(chunk));
+            const { callbacks } = config;
+            if (callbacks === undefined) {
+                config.callbacks = [messageStreamer];
+            }
+            else if (Array.isArray(callbacks)) {
+                config.callbacks = callbacks.concat(messageStreamer);
+            }
+            else {
+                const copiedCallbacks = callbacks.copy();
+                copiedCallbacks.addHandler(messageStreamer, true);
+                config.callbacks = copiedCallbacks;
+            }
+        }
+        // setup custom stream mode
+        if (streamMode.includes("custom")) {
+            config.writer = (chunk) => stream.push([[], "custom", chunk]);
+        }
+        const callbackManager = await getCallbackManagerForConfig(config);
+        const runManager = await callbackManager?.handleChainStart(this.toJSON(), _coerceToDict(input, "input"), runId, undefined, undefined, undefined, config?.runName ?? this.getName());
         const { channelSpecs, managed } = await this.prepareSpecs(config);
         let loop;
-        const stream = new Deque();
-        function* emitCurrentLoopOutputs() {
-            while (loop !== undefined && stream.length > 0) {
-                const nextItem = stream.shift();
-                if (nextItem === undefined) {
-                    throw new Error("Data structure error.");
+        let loopError;
+        const runLoop = async () => {
+            try {
+                loop = await PregelLoop.initialize({
+                    input,
+                    config,
+                    checkpointer,
+                    nodes: this.nodes,
+                    channelSpecs,
+                    managed,
+                    outputKeys,
+                    streamKeys: this.streamChannelsAsIs,
+                    store,
+                    stream,
+                });
+                if (options?.subgraphs) {
+                    loop.config.configurable = {
+                        ...loop.config.configurable,
+                        [CONFIG_KEY_STREAM]: loop.stream,
+                    };
                 }
-                const [namespace, mode, payload] = nextItem;
-                if (streamMode.includes(mode)) {
-                    if (streamSubgraphs && streamMode.length > 1) {
-                        yield [namespace, mode, payload];
-                    }
-                    else if (streamMode.length > 1) {
-                        yield [mode, payload];
-                    }
-                    else if (streamSubgraphs) {
-                        yield [namespace, payload];
+                while (await loop.tick({
+                    inputKeys: this.inputChannels,
+                    interruptAfter,
+                    interruptBefore,
+                    manager: runManager,
+                })) {
+                    if (debug) {
+                        printStepCheckpoint(loop.checkpointMetadata.step, loop.channels, this.streamChannelsList);
                     }
-                    else {
-                        yield payload;
+                    if (debug) {
+                        printStepTasks(loop.step, Object.values(loop.tasks));
                     }
-                }
-            }
-        }
-        try {
-            loop = await PregelLoop.initialize({
-                input,
-                config,
-                checkpointer,
-                nodes: this.nodes,
-                channelSpecs,
-                managed,
-                outputKeys,
-                streamKeys: this.streamChannelsAsIs,
-                store,
-                stream: new StreamProtocol((chunk) => stream.push(chunk), new Set(streamMode)),
-            });
-            if (options?.subgraphs) {
-                loop.config.configurable = {
-                    ...loop.config.configurable,
-                    [CONFIG_KEY_STREAM]: loop.stream,
-                };
-            }
-            while (await loop.tick({
-                inputKeys: this.inputChannels,
-                interruptAfter,
-                interruptBefore,
-                manager: runManager,
-            })) {
-                if (debug) {
-                    printStepCheckpoint(loop.checkpointMetadata.step, loop.channels, this.streamChannelsList);
-                }
-                yield* emitCurrentLoopOutputs();
-                if (debug) {
-                    printStepTasks(loop.step, Object.values(loop.tasks));
-                }
-                // execute tasks, and wait for one to fail or all to finish.
-                // each task is independent from all other concurrent tasks
-                // yield updates/debug output as each task finishes
-                const taskStream = executeTasksWithRetry(Object.values(loop.tasks).filter((task) => task.writes.length === 0), {
-                    stepTimeout: this.stepTimeout,
-                    signal: config.signal,
-                    retryPolicy: this.retryPolicy,
-                });
-                // Timeouts will be thrown
-                for await (const { task, error } of taskStream) {
-                    if (error !== undefined) {
-                        if (isGraphInterrupt(error)) {
-                            if (loop.isNested) {
-                                throw error;
+                    // execute tasks, and wait for one to fail or all to finish.
+                    // each task is independent from all other concurrent tasks
+                    // yield updates/debug output as each task finishes
+                    const taskStream = executeTasksWithRetry(Object.values(loop.tasks).filter((task) => task.writes.length === 0), {
+                        stepTimeout: this.stepTimeout,
+                        signal: config.signal,
+                        retryPolicy: this.retryPolicy,
+                    });
+                    // Timeouts will be thrown
+                    for await (const { task, error } of taskStream) {
+                        if (error !== undefined) {
+                            if (isGraphInterrupt(error)) {
+                                if (loop.isNested) {
+                                    throw error;
+                                }
+                                if (error.interrupts.length) {
+                                    loop.putWrites(task.id, error.interrupts.map((interrupt) => [INTERRUPT, interrupt]));
+                                }
                             }
-                            if (error.interrupts.length) {
-                                loop.putWrites(task.id, error.interrupts.map((interrupt) => [INTERRUPT, interrupt]));
+                            else {
+                                loop.putWrites(task.id, [
+                                    [ERROR, { message: error.message, name: error.name }],
+                                ]);
                             }
                         }
                         else {
-                            loop.putWrites(task.id, [
-                                [ERROR, { message: error.message, name: error.name }],
-                            ]);
+                            loop.putWrites(task.id, task.writes);
+                        }
+                        if (error !== undefined && !isGraphInterrupt(error)) {
+                            throw error;
                         }
                     }
-                    else {
-                        loop.putWrites(task.id, task.writes);
+                    if (debug) {
+                        printStepWrites(loop.step, Object.values(loop.tasks)
+                            .map((task) => task.writes)
+                            .flat(), this.streamChannelsList);
                     }
-                    yield* emitCurrentLoopOutputs();
-                    if (error !== undefined && !isGraphInterrupt(error)) {
-                        throw error;
+                }
+                if (loop.status === "out_of_steps") {
+                    throw new GraphRecursionError([
+                        `Recursion limit of ${config.recursionLimit} reached`,
+                        "without hitting a stop condition. You can increase the",
+                        `limit by setting the "recursionLimit" config key.`,
+                    ].join(" "), {
+                        lc_error_code: "GRAPH_RECURSION_LIMIT",
+                    });
+                }
+            }
+            catch (e) {
+                loopError = e;
+            }
+            finally {
+                try {
+                    // Call `.stop()` again incase it was not called in the loop, e.g due to an error.
+                    if (loop) {
+                        await loop.store?.stop();
                     }
+                    await Promise.all([
+                        ...(loop?.checkpointerPromises ?? []),
+                        ...Array.from(managed.values()).map((mv) => mv.promises()),
+                    ]);
                 }
-                if (debug) {
-                    printStepWrites(loop.step, Object.values(loop.tasks)
-                        .map((task) => task.writes)
-                        .flat(), this.streamChannelsList);
+                catch (e) {
+                    loopError = loopError ?? e;
+                }
+                if (loopError) {
+                    // "Causes any future interactions with the associated stream to error".
+                    // Wraps ReadableStreamDefaultController#error:
+                    // https://developer.mozilla.org/en-US/docs/Web/API/ReadableStreamDefaultController/error
+                    stream.error(loopError);
+                }
+                else {
+                    // Will end the iterator outside of this method,
+                    // keeping previously enqueued chunks.
+                    // Wraps ReadableStreamDefaultController#close:
+                    // https://developer.mozilla.org/en-US/docs/Web/API/ReadableStreamDefaultController/close
+                    stream.close();
                 }
             }
-            yield* emitCurrentLoopOutputs();
-            if (loop.status === "out_of_steps") {
-                throw new GraphRecursionError([
-                    `Recursion limit of ${config.recursionLimit} reached`,
-                    "without hitting a stop condition. You can increase the",
-                    `limit by setting the "recursionLimit" config key.`,
-                ].join(" "), {
-                    lc_error_code: "GRAPH_RECURSION_LIMIT",
-                });
+        };
+        const runLoopPromise = runLoop();
+        try {
+            for await (const chunk of stream) {
+                if (chunk === undefined) {
+                    throw new Error("Data structure error.");
+                }
+                const [namespace, mode, payload] = chunk;
+                if (streamMode.includes(mode)) {
+                    if (streamSubgraphs && !streamModeSingle) {
+                        yield [namespace, mode, payload];
+                    }
+                    else if (!streamModeSingle) {
+                        yield [mode, payload];
+                    }
+                    else if (streamSubgraphs) {
+                        yield [namespace, payload];
+                    }
+                    else {
+                        yield payload;
+                    }
+                }
             }
-            await Promise.all(loop?.checkpointerPromises ?? []);
-            await runManager?.handleChainEnd(loop.output);
         }
         catch (e) {
-            await runManager?.handleChainError(e);
+            await runManager?.handleChainError(loopError);
             throw e;
         }
         finally {
-            // Call `.stop()` again incase it was not called in the loop, e.g due to an error.
-            if (loop) {
-                await loop.store?.stop();
-            }
-            await Promise.all([
-                ...(loop?.checkpointerPromises ?? []),
-                ...Array.from(managed.values()).map((mv) => mv.promises()),
-            ]);
+            await runLoopPromise;
         }
+        await runManager?.handleChainEnd(loop?.output ?? {});
     }
     /**
      * Run the graph with a single input and config.

package/dist/pregel/loop.cjs

@@ -1,6 +1,7 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.PregelLoop = exports.StreamProtocol = void 0;
+exports.PregelLoop = exports.IterableReadableWritableStream = void 0;
+const stream_1 = require("@langchain/core/utils/stream");
 const langgraph_checkpoint_1 = require("@langchain/langgraph-checkpoint");
 const base_js_1 = require("../channels/base.cjs");
 const constants_js_1 = require("../constants.cjs");
@@ -14,33 +15,72 @@ const INPUT_DONE = Symbol.for("INPUT_DONE");
 const INPUT_RESUMING = Symbol.for("INPUT_RESUMING");
 const DEFAULT_LOOP_LIMIT = 25;
 const SPECIAL_CHANNELS = [constants_js_1.ERROR, constants_js_1.INTERRUPT];
-class StreamProtocol {
-    constructor(pushFn, modes) {
-        Object.defineProperty(this, "push", {
+class IterableReadableWritableStream extends stream_1.IterableReadableStream {
+    constructor(params) {
+        let streamControllerPromiseResolver;
+        const streamControllerPromise = new Promise((resolve) => {
+            streamControllerPromiseResolver = resolve;
+        });
+        super({
+            start: (controller) => {
+                streamControllerPromiseResolver(controller);
+            },
+        });
+        Object.defineProperty(this, "modes", {
             enumerable: true,
             configurable: true,
             writable: true,
             value: void 0
         });
-        Object.defineProperty(this, "modes", {
+        Object.defineProperty(this, "controller", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "passthroughFn", {
             enumerable: true,
             configurable: true,
             writable: true,
             value: void 0
         });
-        this.push = pushFn;
-        this.modes = modes;
+        // .start() will always be called before the stream can be interacted
+        // with anyway
+        void streamControllerPromise.then((controller) => {
+            this.controller = controller;
+        });
+        this.passthroughFn = params.passthroughFn;
+        this.modes = params.modes;
+    }
+    push(chunk) {
+        this.passthroughFn?.(chunk);
+        this.controller.enqueue(chunk);
+    }
+    close() {
+        try {
+            this.controller.close();
+        }
+        catch (e) {
+            // pass
+        }
+    }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    error(e) {
+        this.controller.error(e);
     }
 }
-exports.StreamProtocol = StreamProtocol;
+exports.IterableReadableWritableStream = IterableReadableWritableStream;
 function createDuplexStream(...streams) {
-    return new StreamProtocol((value) => {
-        for (const stream of streams) {
-            if (stream.modes.has(value[1])) {
-                stream.push(value);
+    return new IterableReadableWritableStream({
+        passthroughFn: (value) => {
+            for (const stream of streams) {
+                if (stream.modes.has(value[1])) {
+                    stream.push(value);
+                }
             }
-        }
-    }, new Set(streams.flatMap((s) => Array.from(s.modes))));
+        },
+        modes: new Set(streams.flatMap((s) => Array.from(s.modes))),
+    });
 }
 class PregelLoop {
     constructor(params) {

package/dist/pregel/loop.d.ts

@@ -1,5 +1,6 @@
 import type { RunnableConfig } from "@langchain/core/runnables";
 import type { CallbackManagerForChainRun } from "@langchain/core/callbacks/manager";
+import { IterableReadableStream } from "@langchain/core/utils/stream";
 import { BaseCheckpointSaver, Checkpoint, PendingWrite, CheckpointPendingWrite, CheckpointMetadata, All, BaseStore, AsyncBatchedStore } from "@langchain/langgraph-checkpoint";
 import { BaseChannel } from "../channels/base.js";
 import { PregelExecutableTask, StreamMode } from "./types.js";
@@ -16,7 +17,7 @@ export type PregelLoopInitializeParams = {
     nodes: Record<string, PregelNode>;
     channelSpecs: Record<string, BaseChannel>;
     managed: ManagedValueMapping;
-    stream: StreamProtocol;
+    stream: IterableReadableWritableStream;
     store?: BaseStore;
     checkSubgraphs?: boolean;
 };
@@ -39,14 +40,21 @@ type PregelLoopParams = {
     checkpointNamespace: string[];
     skipDoneTasks: boolean;
     isNested: boolean;
-    stream: StreamProtocol;
+    stream: IterableReadableWritableStream;
     store?: AsyncBatchedStore;
     prevCheckpointConfig: RunnableConfig | undefined;
 };
-export declare class StreamProtocol {
-    push: (chunk: StreamChunk) => void;
+export declare class IterableReadableWritableStream extends IterableReadableStream<StreamChunk> {
     modes: Set<StreamMode>;
-    constructor(pushFn: (chunk: StreamChunk) => void, modes: Set<StreamMode>);
+    private controller;
+    private passthroughFn?;
+    constructor(params: {
+        passthroughFn?: (chunk: StreamChunk) => void;
+        modes: Set<StreamMode>;
+    });
+    push(chunk: StreamChunk): void;
+    close(): void;
+    error(e: any): void;
 }
 export declare class PregelLoop {
     protected input?: any;
@@ -72,7 +80,7 @@ export declare class PregelLoop {
     protected prevCheckpointConfig: RunnableConfig | undefined;
     status: "pending" | "done" | "interrupt_before" | "interrupt_after" | "out_of_steps";
     tasks: Record<string, PregelExecutableTask<any, any>>;
-    stream: StreamProtocol;
+    stream: IterableReadableWritableStream;
     checkpointerPromises: Promise<unknown>[];
     isNested: boolean;
     protected _checkpointerChainedPromise: Promise<unknown>;

package/dist/pregel/loop.js

@@ -1,3 +1,4 @@
+import { IterableReadableStream } from "@langchain/core/utils/stream";
 import { copyCheckpoint, emptyCheckpoint, AsyncBatchedStore, } from "@langchain/langgraph-checkpoint";
 import { createCheckpoint, emptyChannels, } from "../channels/base.js";
 import { CHECKPOINT_NAMESPACE_SEPARATOR, CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_READ, CONFIG_KEY_RESUMING, CONFIG_KEY_STREAM, ERROR, INPUT, INTERRUPT, TAG_HIDDEN, TASKS, } from "../constants.js";
@@ -11,32 +12,71 @@ const INPUT_DONE = Symbol.for("INPUT_DONE");
 const INPUT_RESUMING = Symbol.for("INPUT_RESUMING");
 const DEFAULT_LOOP_LIMIT = 25;
 const SPECIAL_CHANNELS = [ERROR, INTERRUPT];
-export class StreamProtocol {
-    constructor(pushFn, modes) {
-        Object.defineProperty(this, "push", {
+export class IterableReadableWritableStream extends IterableReadableStream {
+    constructor(params) {
+        let streamControllerPromiseResolver;
+        const streamControllerPromise = new Promise((resolve) => {
+            streamControllerPromiseResolver = resolve;
+        });
+        super({
+            start: (controller) => {
+                streamControllerPromiseResolver(controller);
+            },
+        });
+        Object.defineProperty(this, "modes", {
             enumerable: true,
             configurable: true,
             writable: true,
             value: void 0
         });
-        Object.defineProperty(this, "modes", {
+        Object.defineProperty(this, "controller", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "passthroughFn", {
             enumerable: true,
             configurable: true,
             writable: true,
             value: void 0
         });
-        this.push = pushFn;
-        this.modes = modes;
+        // .start() will always be called before the stream can be interacted
+        // with anyway
+        void streamControllerPromise.then((controller) => {
+            this.controller = controller;
+        });
+        this.passthroughFn = params.passthroughFn;
+        this.modes = params.modes;
+    }
+    push(chunk) {
+        this.passthroughFn?.(chunk);
+        this.controller.enqueue(chunk);
+    }
+    close() {
+        try {
+            this.controller.close();
+        }
+        catch (e) {
+            // pass
+        }
+    }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    error(e) {
+        this.controller.error(e);
     }
 }
 function createDuplexStream(...streams) {
-    return new StreamProtocol((value) => {
-        for (const stream of streams) {
-            if (stream.modes.has(value[1])) {
-                stream.push(value);
+    return new IterableReadableWritableStream({
+        passthroughFn: (value) => {
+            for (const stream of streams) {
+                if (stream.modes.has(value[1])) {
+                    stream.push(value);
+                }
             }
-        }
-    }, new Set(streams.flatMap((s) => Array.from(s.modes))));
+        },
+        modes: new Set(streams.flatMap((s) => Array.from(s.modes))),
+    });
 }
 export class PregelLoop {
     constructor(params) {