@langchain/langgraph 0.3.5 → 0.3.7

package/dist/pregel/loop.cjs

@@ -11,6 +11,7 @@ const errors_js_1 = require("../errors.cjs");
 const index_js_1 = require("./utils/index.cjs");
 const debug_js_1 = require("./debug.cjs");
 const stream_js_1 = require("./stream.cjs");
+const hash_js_1 = require("../hash.cjs");
 const INPUT_DONE = Symbol.for("INPUT_DONE");
 const INPUT_RESUMING = Symbol.for("INPUT_RESUMING");
 const DEFAULT_LOOP_LIMIT = 25;
@@ -129,6 +130,12 @@ class PregelLoop {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "checkpointIdSaved", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "checkpointConfig", {
             enumerable: true,
             configurable: true,
@@ -171,6 +178,12 @@ class PregelLoop {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "checkpointDuring", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "outputKeys", {
             enumerable: true,
             configurable: true,
@@ -201,6 +214,12 @@ class PregelLoop {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "prevCheckpoint", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "status", {
             enumerable: true,
             configurable: true,
@@ -320,6 +339,7 @@ class PregelLoop {
         this.prevCheckpointConfig = params.prevCheckpointConfig;
         this.interruptAfter = params.interruptAfter;
         this.interruptBefore = params.interruptBefore;
+        this.checkpointDuring = params.checkpointDuring;
         this.debug = params.debug;
         this.triggerToNodes = params.triggerToNodes;
     }
@@ -422,6 +442,7 @@ class PregelLoop {
             cache: params.cache,
             interruptAfter: params.interruptAfter,
             interruptBefore: params.interruptBefore,
+            checkpointDuring: params.checkpointDuring,
             debug: params.debug,
             triggerToNodes: params.triggerToNodes,
         });
@@ -446,22 +467,17 @@ class PregelLoop {
      */
     putWrites(taskId, writes) {
         let writesCopy = writes;
-        if (writesCopy.length === 0) {
+        if (writesCopy.length === 0)
             return;
-        }
         // deduplicate writes to special channels, last write wins
         if (writesCopy.every(([key]) => key in langgraph_checkpoint_1.WRITES_IDX_MAP)) {
             writesCopy = Array.from(new Map(writesCopy.map((w) => [w[0], w])).values());
         }
+        // remove existing writes for this task
+        this.checkpointPendingWrites = this.checkpointPendingWrites.filter((w) => w[0] !== taskId);
         // save writes
         for (const [c, v] of writesCopy) {
-            const idx = this.checkpointPendingWrites.findIndex((w) => w[0] === taskId && w[1] === c);
-            if (c in langgraph_checkpoint_1.WRITES_IDX_MAP && idx !== -1) {
-                this.checkpointPendingWrites[idx] = [taskId, c, v];
-            }
-            else {
-                this.checkpointPendingWrites.push([taskId, c, v]);
-            }
+            this.checkpointPendingWrites.push([taskId, c, v]);
         }
         const putWritePromise = this.checkpointer?.putWrites({
             ...this.checkpointConfig,
@@ -504,13 +520,27 @@ class PregelLoop {
                 (task.config.tags ?? []).includes(constants_js_1.TAG_HIDDEN)) {
                 return;
             }
-            if (writes.length > 0 &&
-                writes[0][0] !== constants_js_1.ERROR &&
-                writes[0][0] !== constants_js_1.INTERRUPT) {
-                this._emit((0, utils_js_1.gatherIteratorSync)((0, utils_js_1.prefixGenerator)((0, io_js_1.mapOutputUpdates)(this.outputKeys, [[task, writes]], cached), "updates")));
+            if (writes.length > 0) {
+                if (writes[0][0] === constants_js_1.INTERRUPT) {
+                    // in `algo.ts` we append a bool to the task path to indicate
+                    // whether or not a call was present. If so, we don't emit the
+                    // the interrupt as it'll be emitted by the parent.
+                    if (task.path?.[0] === constants_js_1.PUSH && task.path?.at(-1) === true)
+                        return;
+                    const interruptWrites = writes
+                        .filter((w) => w[0] === constants_js_1.INTERRUPT)
+                        .flatMap((w) => w[1]);
+                    this._emit([
+                        ["updates", { [constants_js_1.INTERRUPT]: interruptWrites }],
+                        ["values", { [constants_js_1.INTERRUPT]: interruptWrites }],
+                    ]);
+                }
+                else if (writes[0][0] !== constants_js_1.ERROR) {
+                    this._emit((0, utils_js_1.gatherIteratorSync)((0, utils_js_1.prefixGenerator)((0, io_js_1.mapOutputUpdates)(this.outputKeys, [[task, writes]], cached), "updates")));
+                }
             }
             if (!cached) {
-                this._emit((0, utils_js_1.gatherIteratorSync)((0, utils_js_1.prefixGenerator)((0, debug_js_1.mapDebugTaskResults)(this.step, [[task, writes]], this.streamKeys), "debug")));
+                this._emit((0, utils_js_1.gatherIteratorSync)((0, utils_js_1.prefixGenerator)((0, debug_js_1.mapDebugTaskResults)([[task, writes]], this.streamKeys), "tasks")));
             }
         }
     }
@@ -607,8 +637,7 @@ class PregelLoop {
         this.tasks = nextTasks;
         // Produce debug output
         if (this.checkpointer) {
-            this._emit(await (0, utils_js_1.gatherIterator)((0, utils_js_1.prefixGenerator)((0, debug_js_1.mapDebugCheckpoint)(this.step - 1, // printing checkpoint for previous step
-            this.checkpointConfig, this.channels, this.streamKeys, this.checkpointMetadata, Object.values(this.tasks), this.checkpointPendingWrites, this.prevCheckpointConfig, this.outputKeys), "debug")));
+            this._emit(await (0, utils_js_1.gatherIterator)((0, utils_js_1.prefixGenerator)((0, debug_js_1.mapDebugCheckpoint)(this.checkpointConfig, this.channels, this.streamKeys, this.checkpointMetadata, Object.values(this.tasks), this.checkpointPendingWrites, this.prevCheckpointConfig, this.outputKeys), "checkpoints")));
         }
         if (Object.values(this.tasks).length === 0) {
             this.status = "done";
@@ -641,11 +670,22 @@ class PregelLoop {
             throw new errors_js_1.GraphInterrupt();
         }
         // Produce debug output
-        const debugOutput = await (0, utils_js_1.gatherIterator)((0, utils_js_1.prefixGenerator)((0, debug_js_1.mapDebugTasks)(this.step, Object.values(this.tasks)), "debug"));
+        const debugOutput = await (0, utils_js_1.gatherIterator)((0, utils_js_1.prefixGenerator)((0, debug_js_1.mapDebugTasks)(Object.values(this.tasks)), "tasks"));
         this._emit(debugOutput);
         return true;
     }
     async finishAndHandleError(error) {
+        // persist current checkpoint and writes
+        if (!this.checkpointDuring &&
+            // if it's a top graph
+            (!this.isNested ||
+                // or a nested graph with error or interrupt
+                typeof error !== "undefined" ||
+                // or a nested graph with checkpointer: true
+                this.checkpointNamespace.every((part) => !part.includes(constants_js_1.CHECKPOINT_NAMESPACE_END)))) {
+            this._putCheckpoint(this.checkpointMetadata);
+            this._flushPendingWrites();
+        }
         const suppress = this._suppressInterrupt(error);
         if (suppress || error === undefined) {
             this.output = (0, io_js_1.readChannels)(this.channels, this.outputKeys);
@@ -662,11 +702,12 @@ class PregelLoop {
                 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", interrupts],
-                ["values", interrupts],
-            ]);
+            if ((0, errors_js_1.isGraphInterrupt)(error) && !error.interrupts.length) {
+                this._emit([
+                    ["updates", { [constants_js_1.INTERRUPT]: [] }],
+                    ["values", { [constants_js_1.INTERRUPT]: [] }],
+                ]);
+            }
         }
         return suppress;
     }
@@ -690,7 +731,7 @@ class PregelLoop {
             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")));
+        this._emit((0, utils_js_1.gatherIteratorSync)((0, utils_js_1.prefixGenerator)((0, debug_js_1.mapDebugTasks)([pushed]), "tasks")));
         if (this.debug)
             (0, debug_js_1.printStepTasks)(this.step, [pushed]);
         this.tasks[pushed.id] = pushed;
@@ -717,17 +758,22 @@ class PregelLoop {
         if (scratchpad && scratchpad.nullResume !== undefined) {
             this.putWrites(constants_js_1.NULL_TASK_ID, [[constants_js_1.RESUME, scratchpad.nullResume]]);
         }
+        // map command to writes
         if ((0, constants_js_1.isCommand)(this.input)) {
             const hasResume = this.input.resume != null;
+            if (this.input.resume != null &&
+                typeof this.input.resume === "object" &&
+                Object.keys(this.input.resume).every(hash_js_1.isXXH3)) {
+                this.config.configurable ??= {};
+                this.config.configurable[constants_js_1.CONFIG_KEY_RESUME_MAP] = this.input.resume;
+            }
             if (hasResume && this.checkpointer == null) {
                 throw new Error("Cannot use Command(resume=...) without checkpointer");
             }
             const writes = {};
             // group writes by task id
             for (const [tid, key, value] of (0, io_js_1.mapCommand)(this.input, this.checkpointPendingWrites)) {
-                if (writes[tid] === undefined) {
-                    writes[tid] = [];
-                }
+                writes[tid] ??= [];
                 writes[tid].push([key, value]);
             }
             if (Object.keys(writes).length === 0) {
@@ -810,33 +856,50 @@ class PregelLoop {
         }
     }
     _emit(values) {
-        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 },
+                ]);
             }
         }
     }
-    async _putCheckpoint(inputMetadata) {
-        // Assign step
-        const metadata = {
-            ...inputMetadata,
-            step: this.step,
-            parents: this.config.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINT_MAP] ?? {},
-        };
-        // Bail if no checkpointer
-        if (this.checkpointer !== undefined) {
+    _putCheckpoint(inputMetadata) {
+        const exiting = this.checkpointMetadata === inputMetadata;
+        const doCheckpoint = this.checkpointer != null && (this.checkpointDuring || exiting);
+        const storeCheckpoint = (checkpoint) => {
             // store the previous checkpoint config for debug events
             this.prevCheckpointConfig = this.checkpointConfig?.configurable
                 ?.checkpoint_id
                 ? this.checkpointConfig
                 : undefined;
-            // create new checkpoint
-            this.checkpointMetadata = metadata;
             // child graphs keep at most one checkpoint per parent checkpoint
             // this is achieved by writing child checkpoints as progress is made
             // (so that error recovery / resuming from interrupt don't lose work)
             // but doing so always with an id equal to that of the parent checkpoint
-            this.checkpoint = (0, base_js_1.createCheckpoint)(this.checkpoint, this.channels, this.step);
             this.checkpointConfig = {
                 ...this.checkpointConfig,
                 configurable: {
@@ -852,7 +915,7 @@ class PregelLoop {
             // ensuring checkpointers receive checkpoints in order
             void this._checkpointerPutAfterPrevious({
                 config: { ...this.checkpointConfig },
-                checkpoint: (0, langgraph_checkpoint_1.copyCheckpoint)(this.checkpoint),
+                checkpoint: (0, langgraph_checkpoint_1.copyCheckpoint)(checkpoint),
                 metadata: { ...this.checkpointMetadata },
                 newVersions,
             });
@@ -863,8 +926,59 @@ class PregelLoop {
                     checkpoint_id: this.checkpoint.id,
                 },
             };
+        };
+        // We need to retroactively store the previous checkpoint
+        // if it turns out that pending sends are scheduled.
+        // TODO: remove when `pending_sends` is removed from checkpoints
+        if (!exiting &&
+            !this.checkpointDuring &&
+            this.checkpointer != null &&
+            this.prevCheckpoint != null &&
+            this.checkpoint.pending_sends.length > 0) {
+            storeCheckpoint(this.prevCheckpoint);
+        }
+        if (!exiting) {
+            this.checkpointMetadata = {
+                ...inputMetadata,
+                step: this.step,
+                parents: this.config.configurable?.[constants_js_1.CONFIG_KEY_CHECKPOINT_MAP] ?? {},
+            };
+        }
+        // Store previous checkpoint in case pending_sends are scheduled
+        // for next checkpoint.
+        if (!this.checkpointDuring) {
+            this.prevCheckpoint = this.checkpoint;
+        }
+        // create new checkpoint
+        this.checkpoint = (0, base_js_1.createCheckpoint)(this.checkpoint, doCheckpoint ? this.channels : undefined, this.step, exiting ? { id: this.checkpoint.id } : undefined);
+        // Bail if no checkpointer
+        if (doCheckpoint)
+            storeCheckpoint(this.checkpoint);
+        if (!exiting) {
+            // increment step
+            this.step += 1;
+        }
+    }
+    _flushPendingWrites() {
+        if (this.checkpointer == null)
+            return;
+        if (this.checkpointPendingWrites.length === 0)
+            return;
+        // patch config
+        const config = (0, index_js_1.patchConfigurable)(this.checkpointConfig, {
+            [constants_js_1.CONFIG_KEY_CHECKPOINT_NS]: this.config.configurable?.checkpoint_ns ?? "",
+            [constants_js_1.CONFIG_KEY_CHECKPOINT_ID]: this.checkpoint.id,
+        });
+        // group writes by task id
+        const byTask = {};
+        for (const [tid, key, value] of this.checkpointPendingWrites) {
+            byTask[tid] ??= [];
+            byTask[tid].push([key, value]);
+        }
+        // submit writes to checkpointer
+        for (const [tid, ws] of Object.entries(byTask)) {
+            this.checkpointerPromises.push(this.checkpointer.putWrites(config, ws, tid));
         }
-        this.step += 1;
     }
     _matchWrites(tasks) {
         for (const [tid, k, v] of this.checkpointPendingWrites) {

package/dist/pregel/loop.d.ts

@@ -22,6 +22,7 @@ export type PregelLoopInitializeParams = {
     cache?: BaseCache<PendingWrite<string>[]>;
     interruptAfter: string[] | All;
     interruptBefore: string[] | All;
+    checkpointDuring: boolean;
     manager?: CallbackManagerForChainRun;
     debug: boolean;
     triggerToNodes: Record<string, string[]>;
@@ -52,6 +53,7 @@ type PregelLoopParams = {
     prevCheckpointConfig: RunnableConfig | undefined;
     interruptAfter: string[] | All;
     interruptBefore: string[] | All;
+    checkpointDuring: boolean;
     debug: boolean;
     triggerToNodes: Record<string, string[]>;
 };
@@ -81,6 +83,7 @@ export declare class PregelLoop {
     channels: Record<string, BaseChannel>;
     managed: ManagedValueMapping;
     protected checkpoint: Checkpoint;
+    protected checkpointIdSaved: string | undefined;
     protected checkpointConfig: RunnableConfig;
     checkpointMetadata: CheckpointMetadata;
     protected checkpointNamespace: string[];
@@ -88,11 +91,13 @@ export declare class PregelLoop {
     protected checkpointPreviousVersions: Record<string, string | number>;
     step: number;
     protected stop: number;
+    protected checkpointDuring: boolean;
     protected outputKeys: string | string[];
     protected streamKeys: string | string[];
     protected nodes: Record<string, PregelNode>;
     protected skipDoneTasks: boolean;
     protected prevCheckpointConfig: RunnableConfig | undefined;
+    protected prevCheckpoint: Checkpoint | undefined;
     status: "pending" | "done" | "interrupt_before" | "interrupt_after" | "out_of_steps";
     tasks: Record<string, PregelExecutableTask<any, any>>;
     stream: IterableReadableWritableStream;
@@ -141,7 +146,8 @@ export declare class PregelLoop {
     protected _suppressInterrupt(e?: Error): boolean;
     protected _first(inputKeys: string | string[]): Promise<void>;
     protected _emit(values: [StreamMode, unknown][]): void;
-    protected _putCheckpoint(inputMetadata: Omit<CheckpointMetadata, "step" | "parents">): Promise<void>;
+    protected _putCheckpoint(inputMetadata: Omit<CheckpointMetadata, "step" | "parents">): void;
+    protected _flushPendingWrites(): void;
     protected _matchWrites(tasks: Record<string, PregelExecutableTask<string, string>>): void;
 }
 export {};