@langchain/langgraph 0.2.74 → 0.3.1

package/dist/pregel/index.cjs

@@ -22,9 +22,6 @@ const config_js_1 = require("./utils/config.cjs");
 const messages_js_1 = require("./messages.cjs");
 const runner_js_1 = require("./runner.cjs");
 const stream_js_1 = require("./stream.cjs");
-function isString(value) {
-    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.
@@ -43,7 +40,7 @@ class Channel {
             throw new Error("Can't specify a key when subscribing to multiple channels");
         }
         let channelMappingOrArray;
-        if (isString(channels)) {
+        if (typeof channels === "string") {
             if (key) {
                 channelMappingOrArray = { [key]: channels };
             }
@@ -116,6 +113,36 @@ class Channel {
     }
 }
 exports.Channel = Channel;
+// This is a workaround to allow Pregel to override `invoke` / `stream` and `withConfig`
+// without having to adhere to the types in the `Runnable` class (thanks to `any`).
+// Alternatively we could mark those methods with @ts-ignore / @ts-expect-error,
+// but these do not get carried over when building via `tsc`.
+class PartialRunnable extends runnables_1.Runnable {
+    constructor() {
+        super(...arguments);
+        Object.defineProperty(this, "lc_namespace", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: ["langgraph", "pregel"]
+        });
+    }
+    invoke(_input, _options
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    ) {
+        throw new Error("Not implemented");
+    }
+    // Overriden by `Pregel`
+    withConfig(_config) {
+        return super.withConfig(_config);
+    }
+    // Overriden by `Pregel`
+    stream(input, options
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    ) {
+        return super.stream(input, options);
+    }
+}
 /**
  * The Pregel class is the core runtime engine of LangGraph, implementing a message-passing graph computation model
  * inspired by [Google's Pregel system](https://research.google/pubs/pregel-a-system-for-large-scale-graph-processing/).
@@ -175,7 +202,7 @@ exports.Channel = Channel;
  * @typeParam InputType - Type of input values accepted by the graph
  * @typeParam OutputType - Type of output values produced by the graph
  */
-class Pregel extends runnables_1.Runnable {
+class Pregel extends PartialRunnable {
     /**
      * Name of the class when serialized
      * @internal
@@ -330,7 +357,7 @@ class Pregel extends runnables_1.Runnable {
             value: void 0
         });
         /**
-         * Optional long-term memory store for the graph, allows for persistance & retrieval of data across threads
+         * Optional long-term memory store for the graph, allows for persistence & retrieval of data across threads
          */
         Object.defineProperty(this, "store", {
             enumerable: true,
@@ -338,6 +365,21 @@ class Pregel extends runnables_1.Runnable {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "triggerToNodes", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: {}
+        });
+        /**
+         * Optional cache for the graph, useful for caching tasks.
+         */
+        Object.defineProperty(this, "cache", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         let { streamMode } = fields;
         if (streamMode != null && !Array.isArray(streamMode)) {
             streamMode = [streamMode];
@@ -357,6 +399,7 @@ class Pregel extends runnables_1.Runnable {
         this.retryPolicy = fields.retryPolicy;
         this.config = fields.config;
         this.store = fields.store;
+        this.cache = fields.cache;
         this.name = fields.name;
         if (this.autoValidate) {
             this.validate();
@@ -381,8 +424,6 @@ class Pregel extends runnables_1.Runnable {
      * @param config - The configuration to merge with the current configuration
      * @returns A new Pregel instance with the merged configuration
      */
-    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-    // @ts-ignore Remove ignore when we remove support for 0.2 versions of core
     withConfig(config) {
         const mergedConfig = (0, runnables_1.mergeConfigs)(this.config, config);
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -408,6 +449,12 @@ class Pregel extends runnables_1.Runnable {
             interruptAfterNodes: this.interruptAfter,
             interruptBeforeNodes: this.interruptBefore,
         });
+        for (const [name, node] of Object.entries(this.nodes)) {
+            for (const trigger of node.triggers) {
+                this.triggerToNodes[trigger] ??= [];
+                this.triggerToNodes[trigger].push(name);
+            }
+        }
         return this;
     }
     /**
@@ -549,7 +596,7 @@ class Pregel extends runnables_1.Runnable {
                         writes: nullWrites,
                         triggers: [],
                     },
-                ]);
+                ], undefined, this.triggerToNodes);
             }
         }
         // Prepare next tasks
@@ -613,7 +660,7 @@ class Pregel extends runnables_1.Runnable {
             // Apply writes from tasks that have writes
             const tasksWithWrites = nextTasks.filter((task) => task.writes.length > 0);
             if (tasksWithWrites.length > 0) {
-                (0, algo_js_1._applyWrites)(saved.checkpoint, channels, tasksWithWrites);
+                (0, algo_js_1._applyWrites)(saved.checkpoint, channels, tasksWithWrites, undefined, this.triggerToNodes);
             }
         }
         // Preserve thread_id from the config in metadata
@@ -829,7 +876,7 @@ class Pregel extends runnables_1.Runnable {
                                 writes: nullWrites,
                                 triggers: [],
                             },
-                        ]);
+                        ], undefined, this.triggerToNodes);
                     }
                     // apply writes from tasks that already ran
                     for (const [taskId, k, v] of saved.pendingWrites || []) {
@@ -842,7 +889,7 @@ class Pregel extends runnables_1.Runnable {
                         nextTasks[taskId].writes.push([k, v]);
                     }
                     // clear all current tasks
-                    (0, algo_js_1._applyWrites)(checkpoint, channels, Object.values(nextTasks));
+                    (0, algo_js_1._applyWrites)(checkpoint, channels, Object.values(nextTasks), undefined, this.triggerToNodes);
                 }
                 // save checkpoint
                 const nextConfig = await checkpointer.put(checkpointConfig, (0, base_js_1.createCheckpoint)(checkpoint, undefined, step), {
@@ -881,7 +928,7 @@ class Pregel extends runnables_1.Runnable {
                         writes: inputWrites,
                         triggers: [],
                     },
-                ], checkpointer.getNextVersion.bind(this.checkpointer));
+                ], checkpointer.getNextVersion.bind(this.checkpointer), this.triggerToNodes);
                 // apply input write to channels
                 const nextStep = saved?.metadata?.step != null ? saved.metadata.step + 1 : -1;
                 const nextConfig = await checkpointer.put(checkpointConfig, (0, base_js_1.createCheckpoint)(checkpoint, channels, nextStep), {
@@ -910,13 +957,7 @@ class Pregel extends runnables_1.Runnable {
                     .filter((w) => w[0] === constants_js_1.NULL_TASK_ID)
                     .map((w) => w.slice(1));
                 if (nullWrites.length > 0) {
-                    (0, algo_js_1._applyWrites)(saved.checkpoint, channels, [
-                        {
-                            name: constants_js_1.INPUT,
-                            writes: nullWrites,
-                            triggers: [],
-                        },
-                    ]);
+                    (0, algo_js_1._applyWrites)(saved.checkpoint, channels, [{ name: constants_js_1.INPUT, writes: nullWrites, triggers: [] }], undefined, this.triggerToNodes);
                 }
                 // apply writes
                 for (const [tid, k, v] of saved.pendingWrites) {
@@ -930,7 +971,7 @@ class Pregel extends runnables_1.Runnable {
                     return task.writes.length > 0;
                 });
                 if (tasks.length > 0) {
-                    (0, algo_js_1._applyWrites)(checkpoint, channels, tasks);
+                    (0, algo_js_1._applyWrites)(checkpoint, channels, tasks, undefined, this.triggerToNodes);
                 }
             }
             const nonNullVersion = Object.values(checkpoint.versions_seen)
@@ -1038,7 +1079,7 @@ class Pregel extends runnables_1.Runnable {
             }
             // apply to checkpoint
             // TODO: Why does keyof StrRecord allow number and symbol?
-            (0, algo_js_1._applyWrites)(checkpoint, channels, tasks, checkpointer.getNextVersion.bind(this.checkpointer));
+            (0, algo_js_1._applyWrites)(checkpoint, channels, tasks, checkpointer.getNextVersion.bind(this.checkpointer), this.triggerToNodes);
             const newVersions = (0, index_js_1.getNewChannelVersions)(checkpointPreviousVersions, checkpoint.channel_versions);
             const nextConfig = await checkpointer.put(checkpointConfig, (0, base_js_1.createCheckpoint)(checkpoint, channels, step + 1), {
                 source: "update",
@@ -1145,6 +1186,7 @@ class Pregel extends runnables_1.Runnable {
             defaultCheckpointer = this.checkpointer;
         }
         const defaultStore = config.store ?? this.store;
+        const defaultCache = config.cache ?? this.cache;
         return [
             defaultDebug,
             defaultStreamMode,
@@ -1156,6 +1198,7 @@ class Pregel extends runnables_1.Runnable {
             defaultCheckpointer,
             defaultStore,
             streamModeSingle,
+            defaultCache,
         ];
     }
     /**
@@ -1188,7 +1231,7 @@ class Pregel extends runnables_1.Runnable {
                 ? (0, index_js_1.combineAbortSignals)(options.signal, abortController.signal)
                 : abortController.signal,
         };
-        return new stream_js_1.IterableReadableStreamWithAbortSignal(await super.stream(input, config), abortController);
+        return new stream_js_1.IterableReadableStreamWithAbortSignal((await super.stream(input, config)), abortController);
     }
     streamEvents(input, options, streamOptions) {
         const abortController = new AbortController();
@@ -1301,7 +1344,7 @@ class Pregel extends runnables_1.Runnable {
         const validInput = await this._validateInput(input);
         const { runId, ...restConfig } = inputConfig;
         // assign defaults
-        const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer, store, streamModeSingle,] = this._defaults(restConfig);
+        const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer, store, streamModeSingle, cache,] = this._defaults(restConfig);
         config.configurable = await this._validateConfigurable(config.configurable);
         const stream = new stream_js_1.IterableReadableWritableStream({
             modes: new Set(streamMode),
@@ -1358,11 +1401,13 @@ class Pregel extends runnables_1.Runnable {
                     outputKeys,
                     streamKeys: this.streamChannelsAsIs,
                     store,
+                    cache: cache,
                     stream,
                     interruptAfter,
                     interruptBefore,
                     manager: runManager,
                     debug: this.debug,
+                    triggerToNodes: this.triggerToNodes,
                 });
                 const runner = new runner_js_1.PregelRunner({
                     loop,
@@ -1384,6 +1429,7 @@ class Pregel extends runnables_1.Runnable {
                     // Call `.stop()` again incase it was not called in the loop, e.g due to an error.
                     if (loop) {
                         await loop.store?.stop();
+                        await loop.cache?.stop();
                     }
                     await Promise.all([
                         ...(loop?.checkpointerPromises ?? []),
@@ -1458,11 +1504,31 @@ class Pregel extends runnables_1.Runnable {
         };
         const chunks = [];
         const stream = await this.stream(input, config);
+        const interruptChunks = [];
+        let latest;
         for await (const chunk of stream) {
-            chunks.push(chunk);
+            if (streamMode === "values") {
+                if ((0, constants_js_1.isInterrupted)(chunk)) {
+                    interruptChunks.push(chunk[constants_js_1.INTERRUPT]);
+                }
+                else {
+                    latest = chunk;
+                }
+            }
+            else {
+                chunks.push(chunk);
+            }
         }
         if (streamMode === "values") {
-            return chunks[chunks.length - 1];
+            if (interruptChunks.length > 0) {
+                const interrupts = interruptChunks.flat(1);
+                if (latest == null)
+                    return { [constants_js_1.INTERRUPT]: interrupts };
+                if (typeof latest === "object") {
+                    return { ...latest, [constants_js_1.INTERRUPT]: interrupts };
+                }
+            }
+            return latest;
         }
         return chunks;
     }
@@ -1470,9 +1536,10 @@ class Pregel extends runnables_1.Runnable {
         const { loop, runner, debug, config } = params;
         let tickError;
         try {
-            while (await loop.tick({
-                inputKeys: this.inputChannels,
-            })) {
+            while (await loop.tick({ inputKeys: this.inputChannels })) {
+                for (const { task } of await loop._matchCachedWrites()) {
+                    loop._outputWrites(task.id, task.writes, true);
+                }
                 if (debug) {
                     (0, debug_js_1.printStepCheckpoint)(loop.checkpointMetadata.step, loop.channels, this.streamChannelsList);
                 }
@@ -1514,6 +1581,9 @@ class Pregel extends runnables_1.Runnable {
             }
         }
     }
+    async clearCache() {
+        await this.cache?.clear([]);
+    }
 }
 exports.Pregel = Pregel;
 //# sourceMappingURL=index.js.map

package/dist/pregel/index.d.ts

@@ -1,12 +1,12 @@
 import { Runnable, RunnableConfig, RunnableFunc } from "@langchain/core/runnables";
 import { IterableReadableStream } from "@langchain/core/utils/stream";
-import { All, BaseCheckpointSaver, BaseStore, CheckpointListOptions, CheckpointTuple } from "@langchain/langgraph-checkpoint";
+import { All, BaseCheckpointSaver, BaseStore, CheckpointListOptions, CheckpointTuple, BaseCache } from "@langchain/langgraph-checkpoint";
 import type { StreamEvent } from "@langchain/core/tracers/log_stream";
 import { BaseChannel } from "../channels/base.js";
 import { PregelNode } from "./read.js";
 import { ChannelWrite } from "./write.js";
 import { Command } from "../constants.js";
-import { PregelInterface, PregelParams, StateSnapshot, StreamMode, PregelInputType, PregelOutputType, PregelOptions, SingleChannelSubscriptionOptions, MultipleChannelSubscriptionOptions, GetStateOptions } from "./types.js";
+import { PregelInterface, PregelParams, StateSnapshot, StreamMode, PregelInputType, PregelOutputType, PregelOptions, SingleChannelSubscriptionOptions, MultipleChannelSubscriptionOptions, GetStateOptions, type StreamOutputMap } from "./types.js";
 import { StrRecord } from "./algo.js";
 import { RetryPolicy } from "./utils/index.js";
 import { ManagedValueMapping, type ManagedValueSpec } from "../managed/base.js";
@@ -93,6 +93,12 @@ export declare class Channel {
     static writeTo(channels: string[], writes?: Record<string, WriteValue>): ChannelWrite;
 }
 export type { PregelInputType, PregelOutputType, PregelOptions };
+declare class PartialRunnable<RunInput, RunOutput, CallOptions extends RunnableConfig> extends Runnable<RunInput, RunOutput, CallOptions> {
+    lc_namespace: string[];
+    invoke(_input: RunInput, _options?: Partial<CallOptions>): Promise<any>;
+    withConfig(_config: CallOptions): typeof this;
+    stream(input: RunInput, options?: Partial<CallOptions>): Promise<IterableReadableStream<any>>;
+}
 /**
  * The Pregel class is the core runtime engine of LangGraph, implementing a message-passing graph computation model
  * inspired by [Google's Pregel system](https://research.google/pubs/pregel-a-system-for-large-scale-graph-processing/).
@@ -152,7 +158,7 @@ export type { PregelInputType, PregelOutputType, PregelOptions };
  * @typeParam InputType - Type of input values accepted by the graph
  * @typeParam OutputType - Type of output values produced by the graph
  */
-export declare class Pregel<Nodes extends StrRecord<string, PregelNode>, Channels extends StrRecord<string, BaseChannel | ManagedValueSpec>, ConfigurableFieldType extends Record<string, any> = StrRecord<string, any>, InputType = PregelInputType, OutputType = PregelOutputType> extends Runnable<InputType | Command | null, OutputType, PregelOptions<Nodes, Channels, ConfigurableFieldType>> implements PregelInterface<Nodes, Channels, ConfigurableFieldType>, PregelParams<Nodes, Channels> {
+export declare class Pregel<Nodes extends StrRecord<string, PregelNode>, Channels extends StrRecord<string, BaseChannel | ManagedValueSpec>, ConfigurableFieldType extends Record<string, any> = StrRecord<string, any>, InputType = PregelInputType, OutputType = PregelOutputType, StreamUpdatesType = InputType, StreamValuesType = OutputType> extends PartialRunnable<InputType | Command | null, OutputType, PregelOptions<Nodes, Channels, ConfigurableFieldType>> implements PregelInterface<Nodes, Channels, ConfigurableFieldType>, PregelParams<Nodes, Channels> {
     /**
      * Name of the class when serialized
      * @internal
@@ -222,9 +228,14 @@ export declare class Pregel<Nodes extends StrRecord<string, PregelNode>, Channel
     /** The default configuration for graph execution, can be overridden on a per-invocation basis */
     config?: LangGraphRunnableConfig;
     /**
-     * Optional long-term memory store for the graph, allows for persistance & retrieval of data across threads
+     * Optional long-term memory store for the graph, allows for persistence & retrieval of data across threads
      */
     store?: BaseStore;
+    triggerToNodes: Record<string, string[]>;
+    /**
+     * Optional cache for the graph, useful for caching tasks.
+     */
+    cache?: BaseCache;
     /**
      * Constructor for Pregel - meant for internal use only.
      *
@@ -418,7 +429,8 @@ export declare class Pregel<Nodes extends StrRecord<string, PregelNode>, Channel
         // interrupt after
         BaseCheckpointSaver | undefined,
         BaseStore | undefined,
-        boolean
+        boolean,
+        BaseCache | undefined
     ];
     /**
      * Streams the execution of the graph, emitting state updates as they occur.
@@ -436,7 +448,7 @@ export declare class Pregel<Nodes extends StrRecord<string, PregelNode>, Channel
      * @param options - Configuration options for streaming
      * @returns An async iterable stream of graph state updates
      */
-    stream(input: InputType | Command | null, options?: Partial<PregelOptions<Nodes, Channels, ConfigurableFieldType>>): Promise<IterableReadableStream<PregelOutputType>>;
+    stream<TStreamMode extends StreamMode | StreamMode[] | undefined, TSubgraphs extends boolean>(input: InputType | Command | null, options?: Partial<PregelOptions<Nodes, Channels, ConfigurableFieldType, TStreamMode, TSubgraphs>>): Promise<IterableReadableStream<StreamOutputMap<TStreamMode, TSubgraphs, StreamUpdatesType, StreamValuesType, keyof Nodes>>>;
     /**
      * @inheritdoc
      */
@@ -495,4 +507,5 @@ export declare class Pregel<Nodes extends StrRecord<string, PregelNode>, Channel
      */
     invoke(input: InputType | Command | null, options?: Partial<PregelOptions<Nodes, Channels, ConfigurableFieldType>>): Promise<OutputType>;
     private _runLoop;
+    clearCache(): Promise<void>;
 }