@langchain/langgraph 0.2.54 → 0.2.56

package/dist/pregel/index.js

@@ -673,89 +673,162 @@ export class Pregel extends Runnable {
         }
     }
     /**
-     * Updates the state of the graph with new values.
+     * Apply updates to the graph state in bulk.
      * Requires a checkpointer to be configured.
      *
-     * This method can be used for:
-     * - Implementing human-in-the-loop workflows
-     * - Modifying graph state during breakpoints
-     * - Integrating external inputs into the graph
+     * This method is useful for recreating a thread
+     * from a list of updates, especially if a checkpoint
+     * is created as a result of multiple tasks.
      *
-     * @param inputConfig - Configuration for the update
-     * @param values - The values to update the state with
-     * @param asNode - Optional node name to attribute the update to
+     * @param startConfig - Configuration for the update
+     * @param updates - The list of updates to apply to graph state
      * @returns Updated configuration
      * @throws {GraphValueError} If no checkpointer is configured
-     * @throws {InvalidUpdateError} If the update cannot be attributed to a node
+     * @throws {InvalidUpdateError} If the update cannot be attributed to a node or an update can be only applied in sequence.
      */
-    async updateState(inputConfig, values, asNode) {
-        const checkpointer = inputConfig.configurable?.[CONFIG_KEY_CHECKPOINTER] ?? this.checkpointer;
+    async bulkUpdateState(startConfig, supersteps) {
+        const checkpointer = startConfig.configurable?.[CONFIG_KEY_CHECKPOINTER] ?? this.checkpointer;
         if (!checkpointer) {
             throw new GraphValueError("No checkpointer set");
         }
+        if (supersteps.length === 0) {
+            throw new Error("No supersteps provided");
+        }
+        if (supersteps.some((s) => s.updates.length === 0)) {
+            throw new Error("No updates provided");
+        }
         // delegate to subgraph
-        const checkpointNamespace = inputConfig.configurable?.checkpoint_ns ?? "";
+        const checkpointNamespace = startConfig.configurable?.checkpoint_ns ?? "";
         if (checkpointNamespace !== "" &&
-            inputConfig.configurable?.[CONFIG_KEY_CHECKPOINTER] === undefined) {
+            startConfig.configurable?.[CONFIG_KEY_CHECKPOINTER] === undefined) {
             // remove task_ids from checkpoint_ns
             const recastNamespace = recastCheckpointNamespace(checkpointNamespace);
             // find the subgraph with the matching name
             // eslint-disable-next-line no-unreachable-loop
             for await (const [, pregel] of this.getSubgraphsAsync(recastNamespace, true)) {
-                return await pregel.updateState(patchConfigurable(inputConfig, {
+                return await pregel.bulkUpdateState(patchConfigurable(startConfig, {
                     [CONFIG_KEY_CHECKPOINTER]: checkpointer,
-                }), values, asNode);
+                }), supersteps);
             }
             throw new Error(`Subgraph "${recastNamespace}" not found`);
         }
-        // get last checkpoint
-        const config = this.config
-            ? mergeConfigs(this.config, inputConfig)
-            : inputConfig;
-        const saved = await checkpointer.getTuple(config);
-        const checkpoint = saved !== undefined
-            ? copyCheckpoint(saved.checkpoint)
-            : emptyCheckpoint();
-        const checkpointPreviousVersions = {
-            ...saved?.checkpoint.channel_versions,
-        };
-        const step = saved?.metadata?.step ?? -1;
-        // merge configurable fields with previous checkpoint config
-        let checkpointConfig = patchConfigurable(config, {
-            checkpoint_ns: config.configurable?.checkpoint_ns ?? "",
-        });
-        let checkpointMetadata = config.metadata ?? {};
-        if (saved?.config.configurable) {
-            checkpointConfig = patchConfigurable(config, saved.config.configurable);
-            checkpointMetadata = {
-                ...saved.metadata,
-                ...checkpointMetadata,
+        const updateSuperStep = async (inputConfig, updates) => {
+            // get last checkpoint
+            const config = this.config
+                ? mergeConfigs(this.config, inputConfig)
+                : inputConfig;
+            const saved = await checkpointer.getTuple(config);
+            const checkpoint = saved !== undefined
+                ? copyCheckpoint(saved.checkpoint)
+                : emptyCheckpoint();
+            const checkpointPreviousVersions = {
+                ...saved?.checkpoint.channel_versions,
             };
-        }
-        // Find last node that updated the state, if not provided
-        if (values == null && asNode === undefined) {
-            const nextConfig = await checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, undefined, step), {
-                source: "update",
-                step: step + 1,
-                writes: {},
-                parents: saved?.metadata?.parents ?? {},
-            }, {});
-            return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
-        }
-        // update channels
-        const channels = emptyChannels(this.channels, checkpoint);
-        // Pass `skipManaged: true` as managed values are not used/relevant in update state calls.
-        const { managed } = await this.prepareSpecs(config, { skipManaged: true });
-        if (values === null && asNode === "__end__") {
-            if (saved) {
+            const step = saved?.metadata?.step ?? -1;
+            // merge configurable fields with previous checkpoint config
+            let checkpointConfig = patchConfigurable(config, {
+                checkpoint_ns: config.configurable?.checkpoint_ns ?? "",
+            });
+            let checkpointMetadata = config.metadata ?? {};
+            if (saved?.config.configurable) {
+                checkpointConfig = patchConfigurable(config, saved.config.configurable);
+                checkpointMetadata = {
+                    ...saved.metadata,
+                    ...checkpointMetadata,
+                };
+            }
+            // Find last node that updated the state, if not provided
+            const { values, asNode } = updates[0];
+            if (values == null && asNode === undefined) {
+                if (updates.length > 1) {
+                    throw new InvalidUpdateError(`Cannot create empty checkpoint with multiple updates`);
+                }
+                const nextConfig = await checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, undefined, step), {
+                    source: "update",
+                    step: step + 1,
+                    writes: {},
+                    parents: saved?.metadata?.parents ?? {},
+                }, {});
+                return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
+            }
+            // update channels
+            const channels = emptyChannels(this.channels, checkpoint);
+            // Pass `skipManaged: true` as managed values are not used/relevant in update state calls.
+            const { managed } = await this.prepareSpecs(config, {
+                skipManaged: true,
+            });
+            if (values === null && asNode === "__end__") {
+                if (updates.length > 1) {
+                    throw new InvalidUpdateError(`Cannot apply multiple updates when clearing state`);
+                }
+                if (saved) {
+                    // tasks for this checkpoint
+                    const nextTasks = _prepareNextTasks(checkpoint, saved.pendingWrites || [], this.nodes, channels, managed, saved.config, true, {
+                        step: (saved.metadata?.step ?? -1) + 1,
+                        checkpointer: this.checkpointer || undefined,
+                        store: this.store,
+                    });
+                    // apply null writes
+                    const nullWrites = (saved.pendingWrites || [])
+                        .filter((w) => w[0] === NULL_TASK_ID)
+                        .map((w) => w.slice(1));
+                    if (nullWrites.length > 0) {
+                        _applyWrites(saved.checkpoint, channels, [
+                            {
+                                name: INPUT,
+                                writes: nullWrites,
+                                triggers: [],
+                            },
+                        ]);
+                    }
+                    // apply writes from tasks that already ran
+                    for (const [taskId, k, v] of saved.pendingWrites || []) {
+                        if ([ERROR, INTERRUPT, SCHEDULED].includes(k)) {
+                            continue;
+                        }
+                        if (!(taskId in nextTasks)) {
+                            continue;
+                        }
+                        nextTasks[taskId].writes.push([k, v]);
+                    }
+                    // clear all current tasks
+                    _applyWrites(checkpoint, channels, Object.values(nextTasks));
+                }
+                // save checkpoint
+                const nextConfig = await checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, undefined, step), {
+                    ...checkpointMetadata,
+                    source: "update",
+                    step: step + 1,
+                    writes: {},
+                    parents: saved?.metadata?.parents ?? {},
+                }, {});
+                return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
+            }
+            if (values == null && asNode === "__copy__") {
+                if (updates.length > 1) {
+                    throw new InvalidUpdateError(`Cannot copy checkpoint with multiple updates`);
+                }
+                const nextConfig = await checkpointer.put(saved?.parentConfig ?? checkpointConfig, createCheckpoint(checkpoint, undefined, step), {
+                    source: "fork",
+                    step: step + 1,
+                    writes: {},
+                    parents: saved?.metadata?.parents ?? {},
+                }, {});
+                return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
+            }
+            // apply pending writes, if not on specific checkpoint
+            if (config.configurable?.checkpoint_id === undefined &&
+                saved?.pendingWrites !== undefined &&
+                saved.pendingWrites.length > 0) {
                 // tasks for this checkpoint
-                const nextTasks = _prepareNextTasks(checkpoint, saved.pendingWrites || [], this.nodes, channels, managed, saved.config, true, {
-                    step: (saved.metadata?.step ?? -1) + 1,
-                    checkpointer: this.checkpointer || undefined,
+                const nextTasks = _prepareNextTasks(checkpoint, saved.pendingWrites, this.nodes, channels, managed, saved.config, true, {
                     store: this.store,
+                    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                    checkpointer: this.checkpointer,
+                    step: (saved.metadata?.step ?? -1) + 1,
                 });
                 // apply null writes
-                const nullWrites = (saved.pendingWrites || [])
+                const nullWrites = (saved.pendingWrites ?? [])
                     .filter((w) => w[0] === NULL_TASK_ID)
                     .map((w) => w.slice(1));
                 if (nullWrites.length > 0) {
@@ -767,172 +840,165 @@ export class Pregel extends Runnable {
                         },
                     ]);
                 }
-                // apply writes from tasks that already ran
-                for (const [taskId, k, v] of saved.pendingWrites || []) {
-                    if ([ERROR, INTERRUPT, SCHEDULED].includes(k)) {
-                        continue;
-                    }
-                    if (!(taskId in nextTasks)) {
+                // apply writes
+                for (const [tid, k, v] of saved.pendingWrites) {
+                    if ([ERROR, INTERRUPT, SCHEDULED].includes(k) ||
+                        nextTasks[tid] === undefined) {
                         continue;
                     }
-                    nextTasks[taskId].writes.push([k, v]);
+                    nextTasks[tid].writes.push([k, v]);
+                }
+                const tasks = Object.values(nextTasks).filter((task) => {
+                    return task.writes.length > 0;
+                });
+                if (tasks.length > 0) {
+                    _applyWrites(checkpoint, channels, tasks);
                 }
-                // clear all current tasks
-                _applyWrites(checkpoint, channels, Object.values(nextTasks));
             }
-            // save checkpoint
-            const nextConfig = await checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, undefined, step), {
-                ...checkpointMetadata,
-                source: "update",
-                step: step + 1,
-                writes: {},
-                parents: saved?.metadata?.parents ?? {},
-            }, {});
-            return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
-        }
-        if (values == null && asNode === "__copy__") {
-            const nextConfig = await checkpointer.put(saved?.parentConfig ?? checkpointConfig, createCheckpoint(checkpoint, undefined, step), {
-                source: "fork",
-                step: step + 1,
-                writes: {},
-                parents: saved?.metadata?.parents ?? {},
-            }, {});
-            return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
-        }
-        // apply pending writes, if not on specific checkpoint
-        if (config.configurable?.checkpoint_id === undefined &&
-            saved?.pendingWrites !== undefined &&
-            saved.pendingWrites.length > 0) {
-            // tasks for this checkpoint
-            const nextTasks = _prepareNextTasks(checkpoint, saved.pendingWrites, this.nodes, channels, managed, saved.config, true, {
-                store: this.store,
-                // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                checkpointer: this.checkpointer,
-                step: (saved.metadata?.step ?? -1) + 1,
-            });
-            // apply null writes
-            const nullWrites = (saved.pendingWrites ?? [])
-                .filter((w) => w[0] === NULL_TASK_ID)
-                .map((w) => w.slice(1));
-            if (nullWrites.length > 0) {
-                _applyWrites(saved.checkpoint, channels, [
-                    {
-                        name: INPUT,
-                        writes: nullWrites,
-                        triggers: [],
-                    },
-                ]);
+            const nonNullVersion = Object.values(checkpoint.versions_seen)
+                .map((seenVersions) => {
+                return Object.values(seenVersions);
+            })
+                .flat()
+                .find((v) => !!v);
+            const validUpdates = [];
+            if (updates.length === 1) {
+                // eslint-disable-next-line prefer-const
+                let { values, asNode } = updates[0];
+                if (asNode === undefined && nonNullVersion === undefined) {
+                    if (typeof this.inputChannels === "string" &&
+                        this.nodes[this.inputChannels] !== undefined) {
+                        asNode = this.inputChannels;
+                    }
+                }
+                else if (asNode === undefined) {
+                    const lastSeenByNode = Object.entries(checkpoint.versions_seen)
+                        .map(([n, seen]) => {
+                        return Object.values(seen).map((v) => {
+                            return [v, n];
+                        });
+                    })
+                        .flat()
+                        .sort(([aNumber], [bNumber]) => compareChannelVersions(aNumber, bNumber));
+                    // if two nodes updated the state at the same time, it's ambiguous
+                    if (lastSeenByNode) {
+                        if (lastSeenByNode.length === 1) {
+                            // eslint-disable-next-line prefer-destructuring
+                            asNode = lastSeenByNode[0][1];
+                        }
+                        else if (lastSeenByNode[lastSeenByNode.length - 1][0] !==
+                            lastSeenByNode[lastSeenByNode.length - 2][0]) {
+                            // eslint-disable-next-line prefer-destructuring
+                            asNode = lastSeenByNode[lastSeenByNode.length - 1][1];
+                        }
+                    }
+                }
+                if (asNode === undefined) {
+                    throw new InvalidUpdateError(`Ambiguous update, specify "asNode"`);
+                }
+                validUpdates.push({ values, asNode });
             }
-            // apply writes
-            for (const [tid, k, v] of saved.pendingWrites) {
-                if ([ERROR, INTERRUPT, SCHEDULED].includes(k) ||
-                    nextTasks[tid] === undefined) {
-                    continue;
+            else {
+                for (const { asNode, values } of updates) {
+                    if (asNode == null) {
+                        throw new InvalidUpdateError(`"asNode" is required when applying multiple updates`);
+                    }
+                    validUpdates.push({ values, asNode });
                 }
-                nextTasks[tid].writes.push([k, v]);
             }
-            const tasks = Object.values(nextTasks).filter((task) => {
-                return task.writes.length > 0;
-            });
-            if (tasks.length > 0) {
-                _applyWrites(checkpoint, channels, tasks);
+            const tasks = [];
+            for (const { asNode, values } of validUpdates) {
+                if (this.nodes[asNode] === undefined) {
+                    throw new InvalidUpdateError(`Node "${asNode.toString()}" does not exist`);
+                }
+                // run all writers of the chosen node
+                const writers = this.nodes[asNode].getWriters();
+                if (!writers.length) {
+                    throw new InvalidUpdateError(`No writers found for node "${asNode.toString()}"`);
+                }
+                tasks.push({
+                    name: asNode,
+                    input: values,
+                    proc: writers.length > 1
+                        ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                            RunnableSequence.from(writers, {
+                                omitSequenceTags: true,
+                            })
+                        : writers[0],
+                    writes: [],
+                    triggers: [INTERRUPT],
+                    id: uuid5(INTERRUPT, checkpoint.id),
+                    writers: [],
+                });
             }
-        }
-        const nonNullVersion = Object.values(checkpoint.versions_seen)
-            .map((seenVersions) => {
-            return Object.values(seenVersions);
-        })
-            .flat()
-            .find((v) => !!v);
-        if (asNode === undefined && nonNullVersion === undefined) {
-            if (typeof this.inputChannels === "string" &&
-                this.nodes[this.inputChannels] !== undefined) {
-                asNode = this.inputChannels;
+            for (const task of tasks) {
+                // execute task
+                await task.proc.invoke(task.input, patchConfig({
+                    ...config,
+                    store: config?.store ?? this.store,
+                }, {
+                    runName: config.runName ?? `${this.getName()}UpdateState`,
+                    configurable: {
+                        [CONFIG_KEY_SEND]: (items) => task.writes.push(...items),
+                        [CONFIG_KEY_READ]: (select_, fresh_ = false) => _localRead(step, checkpoint, channels, managed, 
+                        // TODO: Why does keyof StrRecord allow number and symbol?
+                        task, select_, fresh_),
+                    },
+                }));
             }
-        }
-        else if (asNode === undefined) {
-            const lastSeenByNode = Object.entries(checkpoint.versions_seen)
-                .map(([n, seen]) => {
-                return Object.values(seen).map((v) => {
-                    return [v, n];
-                });
-            })
-                .flat()
-                .sort(([aNumber], [bNumber]) => compareChannelVersions(aNumber, bNumber));
-            // if two nodes updated the state at the same time, it's ambiguous
-            if (lastSeenByNode) {
-                if (lastSeenByNode.length === 1) {
-                    // eslint-disable-next-line prefer-destructuring
-                    asNode = lastSeenByNode[0][1];
+            for (const task of tasks) {
+                // channel writes are saved to current checkpoint
+                const channelWrites = task.writes.filter((w) => w[0] !== PUSH);
+                // save task writes
+                if (saved !== undefined && channelWrites.length > 0) {
+                    await checkpointer.putWrites(checkpointConfig, channelWrites, task.id);
                 }
-                else if (lastSeenByNode[lastSeenByNode.length - 1][0] !==
-                    lastSeenByNode[lastSeenByNode.length - 2][0]) {
-                    // eslint-disable-next-line prefer-destructuring
-                    asNode = lastSeenByNode[lastSeenByNode.length - 1][1];
+            }
+            // apply to checkpoint
+            // TODO: Why does keyof StrRecord allow number and symbol?
+            _applyWrites(checkpoint, channels, tasks, checkpointer.getNextVersion.bind(this.checkpointer));
+            const newVersions = getNewChannelVersions(checkpointPreviousVersions, checkpoint.channel_versions);
+            const nextConfig = await checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, channels, step + 1), {
+                source: "update",
+                step: step + 1,
+                writes: Object.fromEntries(validUpdates.map((update) => [update.asNode, update.values])),
+                parents: saved?.metadata?.parents ?? {},
+            }, newVersions);
+            for (const task of tasks) {
+                // push writes are saved to next checkpoint
+                const pushWrites = task.writes.filter((w) => w[0] === PUSH);
+                if (pushWrites.length > 0) {
+                    await checkpointer.putWrites(nextConfig, pushWrites, task.id);
                 }
             }
-        }
-        if (asNode === undefined) {
-            throw new InvalidUpdateError(`Ambiguous update, specify "asNode"`);
-        }
-        if (this.nodes[asNode] === undefined) {
-            throw new InvalidUpdateError(`Node "${asNode.toString()}" does not exist`);
-        }
-        // run all writers of the chosen node
-        const writers = this.nodes[asNode].getWriters();
-        if (!writers.length) {
-            throw new InvalidUpdateError(`No writers found for node "${asNode.toString()}"`);
-        }
-        const task = {
-            name: asNode,
-            input: values,
-            proc: writers.length > 1
-                ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                    RunnableSequence.from(writers, { omitSequenceTags: true })
-                : writers[0],
-            writes: [],
-            triggers: [INTERRUPT],
-            id: uuid5(INTERRUPT, checkpoint.id),
-            writers: [],
+            return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
         };
-        // execute task
-        await task.proc.invoke(task.input, patchConfig({
-            ...config,
-            store: config?.store ?? this.store,
-        }, {
-            runName: config.runName ?? `${this.getName()}UpdateState`,
-            configurable: {
-                [CONFIG_KEY_SEND]: (items) => task.writes.push(...items),
-                [CONFIG_KEY_READ]: (select_, fresh_ = false) => _localRead(step, checkpoint, channels, managed, 
-                // TODO: Why does keyof StrRecord allow number and symbol?
-                task, select_, fresh_),
-            },
-        }));
-        // save task writes
-        // channel writes are saved to current checkpoint
-        // push writes are saved to next checkpoint
-        const [channelWrites, pushWrites] = [
-            task.writes.filter((w) => w[0] !== PUSH),
-            task.writes.filter((w) => w[0] === PUSH),
-        ];
-        // save task writes
-        if (saved !== undefined && channelWrites.length > 0) {
-            await checkpointer.putWrites(checkpointConfig, channelWrites, task.id);
+        let currentConfig = startConfig;
+        for (const { updates } of supersteps) {
+            currentConfig = await updateSuperStep(currentConfig, updates);
         }
-        // apply to checkpoint
-        // TODO: Why does keyof StrRecord allow number and symbol?
-        _applyWrites(checkpoint, channels, [task], checkpointer.getNextVersion.bind(this.checkpointer));
-        const newVersions = getNewChannelVersions(checkpointPreviousVersions, checkpoint.channel_versions);
-        const nextConfig = await checkpointer.put(checkpointConfig, createCheckpoint(checkpoint, channels, step + 1), {
-            source: "update",
-            step: step + 1,
-            writes: { [asNode]: values },
-            parents: saved?.metadata?.parents ?? {},
-        }, newVersions);
-        if (pushWrites.length > 0) {
-            await checkpointer.putWrites(nextConfig, pushWrites, task.id);
-        }
-        return patchCheckpointMap(nextConfig, saved ? saved.metadata : undefined);
+        return currentConfig;
+    }
+    /**
+     * Updates the state of the graph with new values.
+     * Requires a checkpointer to be configured.
+     *
+     * This method can be used for:
+     * - Implementing human-in-the-loop workflows
+     * - Modifying graph state during breakpoints
+     * - Integrating external inputs into the graph
+     *
+     * @param inputConfig - Configuration for the update
+     * @param values - The values to update the state with
+     * @param asNode - Optional node name to attribute the update to
+     * @returns Updated configuration
+     * @throws {GraphValueError} If no checkpointer is configured
+     * @throws {InvalidUpdateError} If the update cannot be attributed to a node
+     */
+    async updateState(inputConfig, values, asNode) {
+        return this.bulkUpdateState(inputConfig, [
+            { updates: [{ values, asNode }] },
+        ]);
     }
     /**
      * Gets the default values for various graph configuration options.