@langchain/langgraph 0.2.74 → 0.3.1

package/dist/pregel/loop.cjs

@@ -26,6 +26,44 @@ function createDuplexStream(...streams) {
         modes: new Set(streams.flatMap((s) => Array.from(s.modes))),
     });
 }
+class AsyncBatchedCache extends langgraph_checkpoint_1.BaseCache {
+    constructor(cache) {
+        super();
+        Object.defineProperty(this, "cache", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "queue", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: Promise.resolve()
+        });
+        this.cache = cache;
+    }
+    async get(keys) {
+        return this.enqueueOperation("get", keys);
+    }
+    async set(pairs) {
+        return this.enqueueOperation("set", pairs);
+    }
+    async clear(namespaces) {
+        return this.enqueueOperation("clear", namespaces);
+    }
+    async stop() {
+        await this.queue;
+    }
+    enqueueOperation(type, ...args) {
+        const newPromise = this.queue.then(() => {
+            // @ts-expect-error Tuple type warning
+            return this.cache[type](...args);
+        });
+        this.queue = newPromise.then(() => void 0, () => void 0);
+        return newPromise;
+    }
+}
 class PregelLoop {
     get isResuming() {
         const hasChannelVersions = Object.keys(this.checkpoint.channel_versions).length !== 0;
@@ -207,6 +245,12 @@ class PregelLoop {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "cache", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "manager", {
             enumerable: true,
             configurable: true,
@@ -237,6 +281,12 @@ class PregelLoop {
             writable: true,
             value: false
         });
+        Object.defineProperty(this, "triggerToNodes", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         this.input = params.input;
         this.checkpointer = params.checkpointer;
         // TODO: if managed values no longer needs graph we can replace with
@@ -264,12 +314,14 @@ class PregelLoop {
         this.nodes = params.nodes;
         this.skipDoneTasks = params.skipDoneTasks;
         this.store = params.store;
+        this.cache = params.cache ? new AsyncBatchedCache(params.cache) : undefined;
         this.stream = params.stream;
         this.checkpointNamespace = params.checkpointNamespace;
         this.prevCheckpointConfig = params.prevCheckpointConfig;
         this.interruptAfter = params.interruptAfter;
         this.interruptBefore = params.interruptBefore;
         this.debug = params.debug;
+        this.triggerToNodes = params.triggerToNodes;
     }
     static async initialize(params) {
         let { config, stream } = params;
@@ -367,9 +419,11 @@ class PregelLoop {
             nodes: params.nodes,
             stream,
             store,
+            cache: params.cache,
             interruptAfter: params.interruptAfter,
             interruptBefore: params.interruptBefore,
             debug: params.debug,
+            triggerToNodes: params.triggerToNodes,
         });
     }
     _checkpointerPutAfterPrevious(input) {
@@ -423,6 +477,25 @@ class PregelLoop {
         if (this.tasks) {
             this._outputWrites(taskId, writesCopy);
         }
+        if (!writes.length || !this.cache || !this.tasks) {
+            return;
+        }
+        // only cache tasks with a cache key
+        const task = this.tasks[taskId];
+        if (task == null || task.cache_key == null) {
+            return;
+        }
+        // only cache successful tasks
+        if (writes[0][0] === constants_js_1.ERROR || writes[0][0] === constants_js_1.INTERRUPT) {
+            return;
+        }
+        void this.cache.set([
+            {
+                key: [task.cache_key.ns, task.cache_key.key],
+                value: task.writes,
+                ttl: task.cache_key.ttl,
+            },
+        ]);
     }
     _outputWrites(taskId, writes, cached = false) {
         const task = this.tasks[taskId];
@@ -441,6 +514,34 @@ class PregelLoop {
             }
         }
     }
+    async _matchCachedWrites() {
+        if (!this.cache)
+            return [];
+        const matched = [];
+        const serializeKey = ([ns, key]) => {
+            return `ns:${ns.join(",")}|key:${key}`;
+        };
+        const keys = [];
+        const keyMap = {};
+        for (const task of Object.values(this.tasks)) {
+            if (task.cache_key != null && !task.writes.length) {
+                keys.push([task.cache_key.ns, task.cache_key.key]);
+                keyMap[serializeKey([task.cache_key.ns, task.cache_key.key])] = task;
+            }
+        }
+        if (keys.length === 0)
+            return [];
+        const cache = await this.cache.get(keys);
+        for (const { key, value } of cache) {
+            const task = keyMap[serializeKey(key)];
+            if (task != null) {
+                // update the task with the cached writes
+                task.writes.push(...value);
+                matched.push({ task, result: value });
+            }
+        }
+        return matched;
+    }
     /**
      * Execute a single iteration of the Pregel loop.
      * Returns true if more iterations are needed.
@@ -465,7 +566,7 @@ class PregelLoop {
             // finish superstep
             const writes = Object.values(this.tasks).flatMap((t) => t.writes);
             // All tasks have finished
-            const managedValueWrites = (0, algo_js_1._applyWrites)(this.checkpoint, this.channels, Object.values(this.tasks), this.checkpointerGetNextVersion);
+            const managedValueWrites = (0, algo_js_1._applyWrites)(this.checkpoint, this.channels, Object.values(this.tasks), this.checkpointerGetNextVersion, this.triggerToNodes);
             for (const [key, values] of Object.entries(managedValueWrites)) {
                 await this.updateManagedValues(key, values);
             }
@@ -554,25 +655,22 @@ class PregelLoop {
             if (this.tasks !== undefined &&
                 this.checkpointPendingWrites.length > 0 &&
                 Object.values(this.tasks).some((task) => task.writes.length > 0)) {
-                const managedValueWrites = (0, algo_js_1._applyWrites)(this.checkpoint, this.channels, Object.values(this.tasks), this.checkpointerGetNextVersion);
+                const managedValueWrites = (0, algo_js_1._applyWrites)(this.checkpoint, this.channels, Object.values(this.tasks), this.checkpointerGetNextVersion, this.triggerToNodes);
                 for (const [key, values] of Object.entries(managedValueWrites)) {
                     await this.updateManagedValues(key, values);
                 }
                 this._emit((0, utils_js_1.gatherIteratorSync)((0, utils_js_1.prefixGenerator)((0, io_js_1.mapOutputValues)(this.outputKeys, Object.values(this.tasks).flatMap((t) => t.writes), this.channels), "values")));
             }
             // Emit INTERRUPT event
+            const interrupts = { [constants_js_1.INTERRUPT]: error.interrupts };
             this._emit([
-                [
-                    "updates",
-                    {
-                        [constants_js_1.INTERRUPT]: error.interrupts,
-                    },
-                ],
+                ["updates", interrupts],
+                ["values", interrupts],
             ]);
         }
         return suppress;
     }
-    acceptPush(task, writeIdx, call) {
+    async acceptPush(task, writeIdx, call) {
         if (this.interruptAfter?.length > 0 &&
             (0, algo_js_1.shouldInterrupt)(this.checkpoint, this.interruptAfter, [task])) {
             this.toInterrupt.push(task);
@@ -585,22 +683,24 @@ class PregelLoop {
             store: this.store,
             stream: this.stream,
         });
-        if (pushed) {
-            if (this.interruptBefore?.length > 0 &&
-                (0, algo_js_1.shouldInterrupt)(this.checkpoint, this.interruptBefore, [pushed])) {
-                this.toInterrupt.push(pushed);
-                return;
-            }
-            this._emit((0, utils_js_1.gatherIteratorSync)((0, utils_js_1.prefixGenerator)((0, debug_js_1.mapDebugTasks)(this.step, [pushed]), "debug")));
-            if (this.debug) {
-                (0, debug_js_1.printStepTasks)(this.step, [pushed]);
-            }
-            this.tasks[pushed.id] = pushed;
-            if (this.skipDoneTasks) {
-                this._matchWrites({ [pushed.id]: pushed });
-            }
-            return pushed;
+        if (!pushed)
+            return;
+        if (this.interruptBefore?.length > 0 &&
+            (0, algo_js_1.shouldInterrupt)(this.checkpoint, this.interruptBefore, [pushed])) {
+            this.toInterrupt.push(pushed);
+            return;
         }
+        this._emit((0, utils_js_1.gatherIteratorSync)((0, utils_js_1.prefixGenerator)((0, debug_js_1.mapDebugTasks)(this.step, [pushed]), "debug")));
+        if (this.debug)
+            (0, debug_js_1.printStepTasks)(this.step, [pushed]);
+        this.tasks[pushed.id] = pushed;
+        if (this.skipDoneTasks)
+            this._matchWrites({ [pushed.id]: pushed });
+        const tasks = await this._matchCachedWrites();
+        for (const { task } of tasks) {
+            this._outputWrites(task.id, task.writes, true);
+        }
+        return pushed;
     }
     _suppressInterrupt(e) {
         return (0, errors_js_1.isGraphInterrupt)(e) && !this.isNested;
@@ -649,7 +749,7 @@ class PregelLoop {
                     writes: nullWrites,
                     triggers: [],
                 },
-            ], this.checkpointerGetNextVersion);
+            ], this.checkpointerGetNextVersion, this.triggerToNodes);
         }
         const isCommandUpdateOrGoto = (0, constants_js_1.isCommand)(this.input) && nullWrites.length > 0;
         if (this.isResuming || isCommandUpdateOrGoto) {
@@ -687,7 +787,7 @@ class PregelLoop {
                         writes: inputWrites,
                         triggers: [],
                     },
-                ]), this.checkpointerGetNextVersion);
+                ]), this.checkpointerGetNextVersion, this.triggerToNodes);
                 // save input checkpoint
                 await this._putCheckpoint({
                     source: "input",

package/dist/pregel/loop.d.ts

@@ -1,6 +1,6 @@
 import type { RunnableConfig } from "@langchain/core/runnables";
 import type { CallbackManagerForChainRun } from "@langchain/core/callbacks/manager";
-import { BaseCheckpointSaver, Checkpoint, PendingWrite, CheckpointPendingWrite, CheckpointMetadata, All, BaseStore, AsyncBatchedStore } from "@langchain/langgraph-checkpoint";
+import { BaseCheckpointSaver, Checkpoint, PendingWrite, CheckpointPendingWrite, CheckpointMetadata, All, BaseStore, AsyncBatchedStore, BaseCache, CacheFullKey, CacheNamespace } from "@langchain/langgraph-checkpoint";
 import { BaseChannel } from "../channels/base.js";
 import { Call, PregelExecutableTask, StreamMode } from "./types.js";
 import { Command } from "../constants.js";
@@ -19,10 +19,12 @@ export type PregelLoopInitializeParams = {
     managed: ManagedValueMapping;
     stream: IterableReadableWritableStream;
     store?: BaseStore;
+    cache?: BaseCache<PendingWrite<string>[]>;
     interruptAfter: string[] | All;
     interruptBefore: string[] | All;
     manager?: CallbackManagerForChainRun;
     debug: boolean;
+    triggerToNodes: Record<string, string[]>;
 };
 type PregelLoopParams = {
     input?: any | Command;
@@ -46,11 +48,30 @@ type PregelLoopParams = {
     manager?: CallbackManagerForChainRun;
     stream: IterableReadableWritableStream;
     store?: AsyncBatchedStore;
+    cache?: BaseCache<PendingWrite<string>[]>;
     prevCheckpointConfig: RunnableConfig | undefined;
     interruptAfter: string[] | All;
     interruptBefore: string[] | All;
     debug: boolean;
+    triggerToNodes: Record<string, string[]>;
 };
+declare class AsyncBatchedCache extends BaseCache<PendingWrite<string>[]> {
+    protected cache: BaseCache<PendingWrite<string>[]>;
+    private queue;
+    constructor(cache: BaseCache<unknown>);
+    get(keys: CacheFullKey[]): Promise<{
+        key: CacheFullKey;
+        value: PendingWrite<string>[];
+    }[]>;
+    set(pairs: {
+        key: CacheFullKey;
+        value: PendingWrite<string>[];
+        ttl?: number;
+    }[]): Promise<void>;
+    clear(namespaces: CacheNamespace[]): Promise<void>;
+    stop(): Promise<void>;
+    private enqueueOperation;
+}
 export declare class PregelLoop {
     protected input?: any | Command;
     output: any;
@@ -79,11 +100,13 @@ export declare class PregelLoop {
     isNested: boolean;
     protected _checkpointerChainedPromise: Promise<unknown>;
     store?: AsyncBatchedStore;
+    cache?: AsyncBatchedCache;
     manager?: CallbackManagerForChainRun;
     interruptAfter: string[] | All;
     interruptBefore: string[] | All;
     toInterrupt: PregelExecutableTask<string, string>[];
     debug: boolean;
+    triggerToNodes: Record<string, string[]>;
     get isResuming(): any;
     constructor(params: PregelLoopParams);
     static initialize(params: PregelLoopInitializeParams): Promise<PregelLoop>;
@@ -101,6 +124,10 @@ export declare class PregelLoop {
      */
     putWrites(taskId: string, writes: PendingWrite<string>[]): void;
     _outputWrites(taskId: string, writes: [string, unknown][], cached?: boolean): void;
+    _matchCachedWrites(): Promise<{
+        task: PregelExecutableTask<string, string>;
+        result: unknown;
+    }[]>;
     /**
      * Execute a single iteration of the Pregel loop.
      * Returns true if more iterations are needed.
@@ -110,7 +137,7 @@ export declare class PregelLoop {
         inputKeys?: string | string[];
     }): Promise<boolean>;
     finishAndHandleError(error?: Error): Promise<boolean>;
-    acceptPush(task: PregelExecutableTask<string, string>, writeIdx: number, call?: Call): PregelExecutableTask<string, string> | void;
+    acceptPush(task: PregelExecutableTask<string, string>, writeIdx: number, call?: Call): Promise<PregelExecutableTask<string, string> | void>;
     protected _suppressInterrupt(e?: Error): boolean;
     protected _first(inputKeys: string | string[]): Promise<void>;
     protected _emit(values: [StreamMode, unknown][]): void;

package/dist/pregel/loop.js

@@ -1,4 +1,4 @@
-import { copyCheckpoint, emptyCheckpoint, AsyncBatchedStore, WRITES_IDX_MAP, } from "@langchain/langgraph-checkpoint";
+import { copyCheckpoint, emptyCheckpoint, AsyncBatchedStore, WRITES_IDX_MAP, BaseCache, } from "@langchain/langgraph-checkpoint";
 import { createCheckpoint, emptyChannels, } from "../channels/base.js";
 import { isCommand, CHECKPOINT_NAMESPACE_SEPARATOR, CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_READ, CONFIG_KEY_RESUMING, CONFIG_KEY_STREAM, ERROR, INPUT, INTERRUPT, NULL_TASK_ID, RESUME, TAG_HIDDEN, PUSH, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_CHECKPOINT_NS, } from "../constants.js";
 import { _applyWrites, _prepareNextTasks, _prepareSingleTask, increment, shouldInterrupt, } from "./algo.js";
@@ -23,6 +23,44 @@ function createDuplexStream(...streams) {
         modes: new Set(streams.flatMap((s) => Array.from(s.modes))),
     });
 }
+class AsyncBatchedCache extends BaseCache {
+    constructor(cache) {
+        super();
+        Object.defineProperty(this, "cache", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "queue", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: Promise.resolve()
+        });
+        this.cache = cache;
+    }
+    async get(keys) {
+        return this.enqueueOperation("get", keys);
+    }
+    async set(pairs) {
+        return this.enqueueOperation("set", pairs);
+    }
+    async clear(namespaces) {
+        return this.enqueueOperation("clear", namespaces);
+    }
+    async stop() {
+        await this.queue;
+    }
+    enqueueOperation(type, ...args) {
+        const newPromise = this.queue.then(() => {
+            // @ts-expect-error Tuple type warning
+            return this.cache[type](...args);
+        });
+        this.queue = newPromise.then(() => void 0, () => void 0);
+        return newPromise;
+    }
+}
 export class PregelLoop {
     get isResuming() {
         const hasChannelVersions = Object.keys(this.checkpoint.channel_versions).length !== 0;
@@ -204,6 +242,12 @@ export class PregelLoop {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "cache", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "manager", {
             enumerable: true,
             configurable: true,
@@ -234,6 +278,12 @@ export class PregelLoop {
             writable: true,
             value: false
         });
+        Object.defineProperty(this, "triggerToNodes", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         this.input = params.input;
         this.checkpointer = params.checkpointer;
         // TODO: if managed values no longer needs graph we can replace with
@@ -261,12 +311,14 @@ export class PregelLoop {
         this.nodes = params.nodes;
         this.skipDoneTasks = params.skipDoneTasks;
         this.store = params.store;
+        this.cache = params.cache ? new AsyncBatchedCache(params.cache) : undefined;
         this.stream = params.stream;
         this.checkpointNamespace = params.checkpointNamespace;
         this.prevCheckpointConfig = params.prevCheckpointConfig;
         this.interruptAfter = params.interruptAfter;
         this.interruptBefore = params.interruptBefore;
         this.debug = params.debug;
+        this.triggerToNodes = params.triggerToNodes;
     }
     static async initialize(params) {
         let { config, stream } = params;
@@ -364,9 +416,11 @@ export class PregelLoop {
             nodes: params.nodes,
             stream,
             store,
+            cache: params.cache,
             interruptAfter: params.interruptAfter,
             interruptBefore: params.interruptBefore,
             debug: params.debug,
+            triggerToNodes: params.triggerToNodes,
         });
     }
     _checkpointerPutAfterPrevious(input) {
@@ -420,6 +474,25 @@ export class PregelLoop {
         if (this.tasks) {
             this._outputWrites(taskId, writesCopy);
         }
+        if (!writes.length || !this.cache || !this.tasks) {
+            return;
+        }
+        // only cache tasks with a cache key
+        const task = this.tasks[taskId];
+        if (task == null || task.cache_key == null) {
+            return;
+        }
+        // only cache successful tasks
+        if (writes[0][0] === ERROR || writes[0][0] === INTERRUPT) {
+            return;
+        }
+        void this.cache.set([
+            {
+                key: [task.cache_key.ns, task.cache_key.key],
+                value: task.writes,
+                ttl: task.cache_key.ttl,
+            },
+        ]);
     }
     _outputWrites(taskId, writes, cached = false) {
         const task = this.tasks[taskId];
@@ -438,6 +511,34 @@ export class PregelLoop {
             }
         }
     }
+    async _matchCachedWrites() {
+        if (!this.cache)
+            return [];
+        const matched = [];
+        const serializeKey = ([ns, key]) => {
+            return `ns:${ns.join(",")}|key:${key}`;
+        };
+        const keys = [];
+        const keyMap = {};
+        for (const task of Object.values(this.tasks)) {
+            if (task.cache_key != null && !task.writes.length) {
+                keys.push([task.cache_key.ns, task.cache_key.key]);
+                keyMap[serializeKey([task.cache_key.ns, task.cache_key.key])] = task;
+            }
+        }
+        if (keys.length === 0)
+            return [];
+        const cache = await this.cache.get(keys);
+        for (const { key, value } of cache) {
+            const task = keyMap[serializeKey(key)];
+            if (task != null) {
+                // update the task with the cached writes
+                task.writes.push(...value);
+                matched.push({ task, result: value });
+            }
+        }
+        return matched;
+    }
     /**
      * Execute a single iteration of the Pregel loop.
      * Returns true if more iterations are needed.
@@ -462,7 +563,7 @@ export class PregelLoop {
             // finish superstep
             const writes = Object.values(this.tasks).flatMap((t) => t.writes);
             // All tasks have finished
-            const managedValueWrites = _applyWrites(this.checkpoint, this.channels, Object.values(this.tasks), this.checkpointerGetNextVersion);
+            const managedValueWrites = _applyWrites(this.checkpoint, this.channels, Object.values(this.tasks), this.checkpointerGetNextVersion, this.triggerToNodes);
             for (const [key, values] of Object.entries(managedValueWrites)) {
                 await this.updateManagedValues(key, values);
             }
@@ -551,25 +652,22 @@ export class PregelLoop {
             if (this.tasks !== undefined &&
                 this.checkpointPendingWrites.length > 0 &&
                 Object.values(this.tasks).some((task) => task.writes.length > 0)) {
-                const managedValueWrites = _applyWrites(this.checkpoint, this.channels, Object.values(this.tasks), this.checkpointerGetNextVersion);
+                const managedValueWrites = _applyWrites(this.checkpoint, this.channels, Object.values(this.tasks), this.checkpointerGetNextVersion, this.triggerToNodes);
                 for (const [key, values] of Object.entries(managedValueWrites)) {
                     await this.updateManagedValues(key, values);
                 }
                 this._emit(gatherIteratorSync(prefixGenerator(mapOutputValues(this.outputKeys, Object.values(this.tasks).flatMap((t) => t.writes), this.channels), "values")));
             }
             // Emit INTERRUPT event
+            const interrupts = { [INTERRUPT]: error.interrupts };
             this._emit([
-                [
-                    "updates",
-                    {
-                        [INTERRUPT]: error.interrupts,
-                    },
-                ],
+                ["updates", interrupts],
+                ["values", interrupts],
             ]);
         }
         return suppress;
     }
-    acceptPush(task, writeIdx, call) {
+    async acceptPush(task, writeIdx, call) {
         if (this.interruptAfter?.length > 0 &&
             shouldInterrupt(this.checkpoint, this.interruptAfter, [task])) {
             this.toInterrupt.push(task);
@@ -582,22 +680,24 @@ export class PregelLoop {
             store: this.store,
             stream: this.stream,
         });
-        if (pushed) {
-            if (this.interruptBefore?.length > 0 &&
-                shouldInterrupt(this.checkpoint, this.interruptBefore, [pushed])) {
-                this.toInterrupt.push(pushed);
-                return;
-            }
-            this._emit(gatherIteratorSync(prefixGenerator(mapDebugTasks(this.step, [pushed]), "debug")));
-            if (this.debug) {
-                printStepTasks(this.step, [pushed]);
-            }
-            this.tasks[pushed.id] = pushed;
-            if (this.skipDoneTasks) {
-                this._matchWrites({ [pushed.id]: pushed });
-            }
-            return pushed;
+        if (!pushed)
+            return;
+        if (this.interruptBefore?.length > 0 &&
+            shouldInterrupt(this.checkpoint, this.interruptBefore, [pushed])) {
+            this.toInterrupt.push(pushed);
+            return;
         }
+        this._emit(gatherIteratorSync(prefixGenerator(mapDebugTasks(this.step, [pushed]), "debug")));
+        if (this.debug)
+            printStepTasks(this.step, [pushed]);
+        this.tasks[pushed.id] = pushed;
+        if (this.skipDoneTasks)
+            this._matchWrites({ [pushed.id]: pushed });
+        const tasks = await this._matchCachedWrites();
+        for (const { task } of tasks) {
+            this._outputWrites(task.id, task.writes, true);
+        }
+        return pushed;
     }
     _suppressInterrupt(e) {
         return isGraphInterrupt(e) && !this.isNested;
@@ -646,7 +746,7 @@ export class PregelLoop {
                     writes: nullWrites,
                     triggers: [],
                 },
-            ], this.checkpointerGetNextVersion);
+            ], this.checkpointerGetNextVersion, this.triggerToNodes);
         }
         const isCommandUpdateOrGoto = isCommand(this.input) && nullWrites.length > 0;
         if (this.isResuming || isCommandUpdateOrGoto) {
@@ -684,7 +784,7 @@ export class PregelLoop {
                         writes: inputWrites,
                         triggers: [],
                     },
-                ]), this.checkpointerGetNextVersion);
+                ]), this.checkpointerGetNextVersion, this.triggerToNodes);
                 // save input checkpoint
                 await this._putCheckpoint({
                     source: "input",