@langchain/langgraph 0.2.3 → 0.2.4

package/dist/pregel/index.cjs

@@ -16,6 +16,8 @@ const algo_js_1 = require("./algo.cjs");
 const utils_js_1 = require("./utils.cjs");
 const loop_js_1 = require("./loop.cjs");
 const retry_js_1 = require("./retry.cjs");
+const base_js_2 = require("../managed/base.cjs");
+const utils_js_2 = require("../utils.cjs");
 function isString(value) {
     return typeof value === "string";
 }
@@ -165,6 +167,12 @@ class Pregel extends runnables_1.Runnable {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "store", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         let { streamMode } = fields;
         if (streamMode != null && !Array.isArray(streamMode)) {
             streamMode = [streamMode];
@@ -182,6 +190,7 @@ class Pregel extends runnables_1.Runnable {
         this.debug = fields.debug ?? this.debug;
         this.checkpointer = fields.checkpointer;
         this.retryPolicy = fields.retryPolicy;
+        this.store = fields.store;
         if (this.autoValidate) {
             this.validate();
         }
@@ -227,7 +236,8 @@ class Pregel extends runnables_1.Runnable {
         const saved = await this.checkpointer.getTuple(config);
         const checkpoint = saved ? saved.checkpoint : (0, langgraph_checkpoint_1.emptyCheckpoint)();
         const channels = (0, base_js_1.emptyChannels)(this.channels, checkpoint);
-        const nextTasks = (0, algo_js_1._prepareNextTasks)(checkpoint, this.nodes, channels, saved !== undefined ? saved.config : config, false, { step: saved ? (saved.metadata?.step ?? -1) + 1 : -1 });
+        const { managed } = await this.prepareSpecs(config);
+        const nextTasks = (0, algo_js_1._prepareNextTasks)(checkpoint, this.nodes, channels, managed, saved !== undefined ? saved.config : config, false, { step: saved ? (saved.metadata?.step ?? -1) + 1 : -1 });
         return {
             values: (0, io_js_1.readChannels)(channels, this.streamChannelsAsIs),
             next: nextTasks.map((task) => task.name),
@@ -245,9 +255,10 @@ class Pregel extends runnables_1.Runnable {
         if (!this.checkpointer) {
             throw new errors_js_1.GraphValueError("No checkpointer set");
         }
+        const { managed } = await this.prepareSpecs(config);
         for await (const saved of this.checkpointer.list(config, options)) {
             const channels = (0, base_js_1.emptyChannels)(this.channels, saved.checkpoint);
-            const nextTasks = (0, algo_js_1._prepareNextTasks)(saved.checkpoint, this.nodes, channels, saved.config, false, { step: -1 });
+            const nextTasks = (0, algo_js_1._prepareNextTasks)(saved.checkpoint, this.nodes, channels, managed, saved.config, false, { step: -1 });
             yield {
                 values: (0, io_js_1.readChannels)(channels, this.streamChannelsAsIs),
                 next: nextTasks.map((task) => task.name),
@@ -336,6 +347,7 @@ class Pregel extends runnables_1.Runnable {
         }
         // update channels
         const channels = (0, base_js_1.emptyChannels)(this.channels, checkpoint);
+        const { managed } = await this.prepareSpecs(config);
         // run all writers of the chosen node
         const writers = this.nodes[asNode].getWriters();
         if (!writers.length) {
@@ -356,9 +368,9 @@ class Pregel extends runnables_1.Runnable {
             runName: config.runName ?? `${this.getName()}UpdateState`,
             configurable: {
                 [constants_js_1.CONFIG_KEY_SEND]: (items) => task.writes.push(...items),
-                [constants_js_1.CONFIG_KEY_READ]: algo_js_1._localRead.bind(undefined, checkpoint, channels, 
+                [constants_js_1.CONFIG_KEY_READ]: (select_, fresh_ = false) => (0, algo_js_1._localRead)(step, checkpoint, channels, managed, 
                 // TODO: Why does keyof StrRecord allow number and symbol?
-                task),
+                task, select_, fresh_),
             },
         }));
         if (saved !== undefined) {
@@ -442,6 +454,43 @@ class Pregel extends runnables_1.Runnable {
     async stream(input, options) {
         return super.stream(input, options);
     }
+    async prepareSpecs(config) {
+        const configForManaged = (0, utils_js_2.patchConfigurable)(config, {
+            [constants_js_1.CONFIG_KEY_STORE]: this.store,
+        });
+        const channelSpecs = {};
+        const managedSpecs = {};
+        for (const [name, spec] of Object.entries(this.channels)) {
+            if ((0, base_js_1.isBaseChannel)(spec)) {
+                channelSpecs[name] = spec;
+            }
+            else {
+                managedSpecs[name] = spec;
+            }
+        }
+        const managed = new base_js_2.ManagedValueMapping(await Object.entries(managedSpecs).reduce(async (accPromise, [key, value]) => {
+            const acc = await accPromise;
+            let initializedValue;
+            if ((0, base_js_2.isConfiguredManagedValue)(value)) {
+                if ("key" in value.params &&
+                    value.params.key === base_js_2.ChannelKeyPlaceholder) {
+                    value.params.key = key;
+                }
+                initializedValue = await value.cls.initialize(configForManaged, value.params);
+            }
+            else {
+                initializedValue = await value.initialize(configForManaged);
+            }
+            if (initializedValue !== undefined) {
+                acc.push([key, initializedValue]);
+            }
+            return acc;
+        }, Promise.resolve([])));
+        return {
+            channelSpecs,
+            managed,
+        };
+    }
     async *_streamIterator(input, options) {
         const inputConfig = (0, runnables_1.ensureConfig)(options);
         if (inputConfig.recursionLimit === undefined ||
@@ -457,6 +506,7 @@ class Pregel extends runnables_1.Runnable {
         delete inputConfig.runId;
         // assign defaults
         const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer,] = this._defaults(inputConfig);
+        const { channelSpecs, managed } = await this.prepareSpecs(config);
         let loop;
         try {
             loop = await loop_js_1.PregelLoop.initialize({
@@ -464,9 +514,11 @@ class Pregel extends runnables_1.Runnable {
                 config,
                 checkpointer,
                 nodes: this.nodes,
-                channelSpecs: this.channels,
+                channelSpecs,
+                managed,
                 outputKeys,
                 streamKeys: this.streamChannelsAsIs,
+                store: this.store,
             });
             while (await loop.tick({
                 inputKeys: this.inputChannels,
@@ -568,7 +620,14 @@ class Pregel extends runnables_1.Runnable {
             throw e;
         }
         finally {
-            await Promise.all(loop?.checkpointerPromises ?? []);
+            // Call `.stop()` again incase it was not called in the loop, e.g due to an error.
+            if (loop) {
+                loop.store?.stop();
+            }
+            await Promise.all([
+                loop?.checkpointerPromises ?? [],
+                ...Array.from(managed.values()).map((mv) => mv.promises()),
+            ]);
         }
     }
     /**

package/dist/pregel/index.d.ts

@@ -7,6 +7,8 @@ import { ChannelWrite } from "./write.js";
 import { PregelInterface, PregelParams, StateSnapshot, StreamMode } from "./types.js";
 import { StrRecord } from "./algo.js";
 import { RetryPolicy } from "./utils.js";
+import { BaseStore } from "../store/base.js";
+import { ManagedValueMapping, type ManagedValueSpec } from "../managed/base.js";
 type WriteValue = Runnable | RunnableFunc<unknown, unknown> | unknown;
 export declare class Channel {
     static subscribeTo(channels: string, options?: {
@@ -21,7 +23,7 @@ export declare class Channel {
 /**
  * Config for executing the graph.
  */
-export interface PregelOptions<Nn extends StrRecord<string, PregelNode>, Cc extends StrRecord<string, BaseChannel>> extends RunnableConfig {
+export interface PregelOptions<Nn extends StrRecord<string, PregelNode>, Cc extends StrRecord<string, BaseChannel | ManagedValueSpec>> extends RunnableConfig {
     /** The stream mode for the graph run. Default is ["values"]. */
     streamMode?: StreamMode | StreamMode[];
     inputKeys?: keyof Cc | Array<keyof Cc>;
@@ -36,7 +38,7 @@ export interface PregelOptions<Nn extends StrRecord<string, PregelNode>, Cc exte
 }
 export type PregelInputType = any;
 export type PregelOutputType = any;
-export declare class Pregel<Nn extends StrRecord<string, PregelNode>, Cc extends StrRecord<string, BaseChannel>> extends Runnable<PregelInputType, PregelOutputType, PregelOptions<Nn, Cc>> implements PregelInterface<Nn, Cc> {
+export declare class Pregel<Nn extends StrRecord<string, PregelNode>, Cc extends StrRecord<string, BaseChannel | ManagedValueSpec>> extends Runnable<PregelInputType, PregelOutputType, PregelOptions<Nn, Cc>> implements PregelInterface<Nn, Cc> {
     static lc_name(): string;
     lc_namespace: string[];
     nodes: Nn;
@@ -52,6 +54,7 @@ export declare class Pregel<Nn extends StrRecord<string, PregelNode>, Cc extends
     debug: boolean;
     checkpointer?: BaseCheckpointSaver;
     retryPolicy?: RetryPolicy;
+    store?: BaseStore;
     constructor(fields: PregelParams<Nn, Cc>);
     validate(): this;
     get streamChannelsList(): Array<keyof Cc>;
@@ -101,6 +104,10 @@ export declare class Pregel<Nn extends StrRecord<string, PregelNode>, Cc extends
      * @param options.debug Whether to print debug information during execution.
      */
     stream(input: PregelInputType, options?: Partial<PregelOptions<Nn, Cc>>): Promise<IterableReadableStream<PregelOutputType>>;
+    protected prepareSpecs(config: RunnableConfig): Promise<{
+        channelSpecs: Record<string, BaseChannel<unknown, unknown, unknown>>;
+        managed: ManagedValueMapping;
+    }>;
     _streamIterator(input: PregelInputType, options?: Partial<PregelOptions<Nn, Cc>>): AsyncGenerator<PregelOutputType>;
     /**
      * Run the graph with a single input and config.

package/dist/pregel/index.js

@@ -1,18 +1,20 @@
 /* eslint-disable no-param-reassign */
 import { Runnable, RunnableSequence, _coerceToRunnable, ensureConfig, getCallbackManagerForConfig, patchConfig, } from "@langchain/core/runnables";
 import { compareChannelVersions, copyCheckpoint, emptyCheckpoint, uuid5, } from "@langchain/langgraph-checkpoint";
-import { createCheckpoint, emptyChannels, } from "../channels/base.js";
+import { createCheckpoint, emptyChannels, isBaseChannel, } from "../channels/base.js";
 import { PregelNode } from "./read.js";
 import { validateGraph, validateKeys } from "./validate.js";
 import { readChannels } from "./io.js";
 import { printStepCheckpoint, printStepTasks, printStepWrites, tasksWithWrites, } from "./debug.js";
 import { ChannelWrite, PASSTHROUGH } from "./write.js";
-import { CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, CONFIG_KEY_SEND, ERROR, INTERRUPT, } from "../constants.js";
+import { CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STORE, ERROR, INTERRUPT, } from "../constants.js";
 import { GraphRecursionError, GraphValueError, InvalidUpdateError, isGraphInterrupt, } from "../errors.js";
 import { _prepareNextTasks, _localRead, _applyWrites, } from "./algo.js";
 import { _coerceToDict, getNewChannelVersions } from "./utils.js";
 import { PregelLoop } from "./loop.js";
 import { executeTasksWithRetry } from "./retry.js";
+import { ChannelKeyPlaceholder, isConfiguredManagedValue, ManagedValueMapping, } from "../managed/base.js";
+import { patchConfigurable } from "../utils.js";
 function isString(value) {
     return typeof value === "string";
 }
@@ -161,6 +163,12 @@ export class Pregel extends Runnable {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "store", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         let { streamMode } = fields;
         if (streamMode != null && !Array.isArray(streamMode)) {
             streamMode = [streamMode];
@@ -178,6 +186,7 @@ export class Pregel extends Runnable {
         this.debug = fields.debug ?? this.debug;
         this.checkpointer = fields.checkpointer;
         this.retryPolicy = fields.retryPolicy;
+        this.store = fields.store;
         if (this.autoValidate) {
             this.validate();
         }
@@ -223,7 +232,8 @@ export class Pregel extends Runnable {
         const saved = await this.checkpointer.getTuple(config);
         const checkpoint = saved ? saved.checkpoint : emptyCheckpoint();
         const channels = emptyChannels(this.channels, checkpoint);
-        const nextTasks = _prepareNextTasks(checkpoint, this.nodes, channels, saved !== undefined ? saved.config : config, false, { step: saved ? (saved.metadata?.step ?? -1) + 1 : -1 });
+        const { managed } = await this.prepareSpecs(config);
+        const nextTasks = _prepareNextTasks(checkpoint, this.nodes, channels, managed, saved !== undefined ? saved.config : config, false, { step: saved ? (saved.metadata?.step ?? -1) + 1 : -1 });
         return {
             values: readChannels(channels, this.streamChannelsAsIs),
             next: nextTasks.map((task) => task.name),
@@ -241,9 +251,10 @@ export class Pregel extends Runnable {
         if (!this.checkpointer) {
             throw new GraphValueError("No checkpointer set");
         }
+        const { managed } = await this.prepareSpecs(config);
         for await (const saved of this.checkpointer.list(config, options)) {
             const channels = emptyChannels(this.channels, saved.checkpoint);
-            const nextTasks = _prepareNextTasks(saved.checkpoint, this.nodes, channels, saved.config, false, { step: -1 });
+            const nextTasks = _prepareNextTasks(saved.checkpoint, this.nodes, channels, managed, saved.config, false, { step: -1 });
             yield {
                 values: readChannels(channels, this.streamChannelsAsIs),
                 next: nextTasks.map((task) => task.name),
@@ -332,6 +343,7 @@ export class Pregel extends Runnable {
         }
         // update channels
         const channels = emptyChannels(this.channels, checkpoint);
+        const { managed } = await this.prepareSpecs(config);
         // run all writers of the chosen node
         const writers = this.nodes[asNode].getWriters();
         if (!writers.length) {
@@ -352,9 +364,9 @@ export class Pregel extends Runnable {
             runName: config.runName ?? `${this.getName()}UpdateState`,
             configurable: {
                 [CONFIG_KEY_SEND]: (items) => task.writes.push(...items),
-                [CONFIG_KEY_READ]: _localRead.bind(undefined, checkpoint, channels, 
+                [CONFIG_KEY_READ]: (select_, fresh_ = false) => _localRead(step, checkpoint, channels, managed, 
                 // TODO: Why does keyof StrRecord allow number and symbol?
-                task),
+                task, select_, fresh_),
             },
         }));
         if (saved !== undefined) {
@@ -438,6 +450,43 @@ export class Pregel extends Runnable {
     async stream(input, options) {
         return super.stream(input, options);
     }
+    async prepareSpecs(config) {
+        const configForManaged = patchConfigurable(config, {
+            [CONFIG_KEY_STORE]: this.store,
+        });
+        const channelSpecs = {};
+        const managedSpecs = {};
+        for (const [name, spec] of Object.entries(this.channels)) {
+            if (isBaseChannel(spec)) {
+                channelSpecs[name] = spec;
+            }
+            else {
+                managedSpecs[name] = spec;
+            }
+        }
+        const managed = new ManagedValueMapping(await Object.entries(managedSpecs).reduce(async (accPromise, [key, value]) => {
+            const acc = await accPromise;
+            let initializedValue;
+            if (isConfiguredManagedValue(value)) {
+                if ("key" in value.params &&
+                    value.params.key === ChannelKeyPlaceholder) {
+                    value.params.key = key;
+                }
+                initializedValue = await value.cls.initialize(configForManaged, value.params);
+            }
+            else {
+                initializedValue = await value.initialize(configForManaged);
+            }
+            if (initializedValue !== undefined) {
+                acc.push([key, initializedValue]);
+            }
+            return acc;
+        }, Promise.resolve([])));
+        return {
+            channelSpecs,
+            managed,
+        };
+    }
     async *_streamIterator(input, options) {
         const inputConfig = ensureConfig(options);
         if (inputConfig.recursionLimit === undefined ||
@@ -453,6 +502,7 @@ export class Pregel extends Runnable {
         delete inputConfig.runId;
         // assign defaults
         const [debug, streamMode, , outputKeys, config, interruptBefore, interruptAfter, checkpointer,] = this._defaults(inputConfig);
+        const { channelSpecs, managed } = await this.prepareSpecs(config);
         let loop;
         try {
             loop = await PregelLoop.initialize({
@@ -460,9 +510,11 @@ export class Pregel extends Runnable {
                 config,
                 checkpointer,
                 nodes: this.nodes,
-                channelSpecs: this.channels,
+                channelSpecs,
+                managed,
                 outputKeys,
                 streamKeys: this.streamChannelsAsIs,
+                store: this.store,
             });
             while (await loop.tick({
                 inputKeys: this.inputChannels,
@@ -564,7 +616,14 @@ export class Pregel extends Runnable {
             throw e;
         }
         finally {
-            await Promise.all(loop?.checkpointerPromises ?? []);
+            // Call `.stop()` again incase it was not called in the loop, e.g due to an error.
+            if (loop) {
+                loop.store?.stop();
+            }
+            await Promise.all([
+                loop?.checkpointerPromises ?? [],
+                ...Array.from(managed.values()).map((mv) => mv.promises()),
+            ]);
         }
     }
     /**

package/dist/pregel/loop.cjs

@@ -14,6 +14,7 @@ const io_js_1 = require("./io.cjs");
 const errors_js_1 = require("../errors.cjs");
 const utils_js_2 = require("./utils.cjs");
 const debug_js_1 = require("./debug.cjs");
+const batch_js_1 = require("../store/batch.cjs");
 const INPUT_DONE = Symbol.for("INPUT_DONE");
 const INPUT_RESUMING = Symbol.for("INPUT_RESUMING");
 const DEFAULT_LOOP_LIMIT = 25;
@@ -50,6 +51,12 @@ class PregelLoop {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "managed", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "checkpoint", {
             enumerable: true,
             configurable: true,
@@ -154,6 +161,12 @@ class PregelLoop {
             writable: true,
             value: Promise.resolve()
         });
+        Object.defineProperty(this, "store", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         this.input = params.input;
         this.config = params.config;
         this.checkpointer = params.checkpointer;
@@ -170,6 +183,7 @@ class PregelLoop {
         this.checkpointMetadata = params.checkpointMetadata;
         this.checkpointPreviousVersions = params.checkpointPreviousVersions;
         this.channels = params.channels;
+        this.managed = params.managed;
         this.checkpointPendingWrites = params.checkpointPendingWrites;
         this.step = params.step;
         this.stop = params.stop;
@@ -178,6 +192,7 @@ class PregelLoop {
         this.streamKeys = params.streamKeys;
         this.nodes = params.nodes;
         this.skipDoneTasks = this.config.configurable?.checkpoint_id === undefined;
+        this.store = params.store;
     }
     static async initialize(params) {
         const saved = (await params.checkpointer?.getTuple(params.config)) ?? {
@@ -205,6 +220,13 @@ class PregelLoop {
         const step = (checkpointMetadata.step ?? 0) + 1;
         const stop = step + (params.config.recursionLimit ?? DEFAULT_LOOP_LIMIT) + 1;
         const checkpointPreviousVersions = { ...checkpoint.channel_versions };
+        const store = params.store
+            ? new batch_js_1.AsyncBatchedStore(params.store)
+            : undefined;
+        if (store) {
+            // Start the store. This is a batch store, so it will run continuously
+            store.start();
+        }
         return new PregelLoop({
             input: params.input,
             config: params.config,
@@ -213,6 +235,7 @@ class PregelLoop {
             checkpointMetadata,
             checkpointConfig,
             channels,
+            managed: params.managed,
             step,
             stop,
             checkpointPreviousVersions,
@@ -220,6 +243,7 @@ class PregelLoop {
             outputKeys: params.outputKeys ?? [],
             streamKeys: params.streamKeys ?? [],
             nodes: params.nodes,
+            store,
         });
     }
     _checkpointerPutAfterPrevious(input) {
@@ -228,6 +252,13 @@ class PregelLoop {
         });
         this.checkpointerPromises.push(this._checkpointerChainedPromise);
     }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    async updateManagedValues(key, values) {
+        const mv = this.managed.get(key);
+        if (mv && "update" in mv && typeof mv.update === "function") {
+            await mv.update(values);
+        }
+    }
     /**
      * Put writes for a task, to be read by the next tick.
      * @param taskId
@@ -261,6 +292,9 @@ class PregelLoop {
      * @param params
      */
     async tick(params) {
+        if (this.store && !this.store.isRunning) {
+            this.store?.start();
+        }
         const { inputKeys = [], interruptAfter = [], interruptBefore = [], manager, } = params;
         if (this.status !== "pending") {
             throw new Error(`Cannot tick when status is no longer "pending". Current status: "${this.status}"`);
@@ -271,7 +305,10 @@ class PregelLoop {
         else if (this.tasks.every((task) => task.writes.length > 0)) {
             const writes = this.tasks.flatMap((t) => t.writes);
             // All tasks have finished
-            (0, algo_js_1._applyWrites)(this.checkpoint, this.channels, this.tasks, this.checkpointerGetNextVersion);
+            const myWrites = (0, algo_js_1._applyWrites)(this.checkpoint, this.channels, this.tasks, this.checkpointerGetNextVersion);
+            for (const [key, values] of Object.entries(myWrites)) {
+                await this.updateManagedValues(key, values);
+            }
             // produce values output
             const valuesOutput = await (0, utils_js_1.gatherIterator)((0, utils_js_1.prefixGenerator)((0, io_js_1.mapOutputValues)(this.outputKeys, writes, this.channels), "values"));
             this.stream.push(...valuesOutput);
@@ -299,7 +336,7 @@ class PregelLoop {
             this.status = "out_of_steps";
             return false;
         }
-        const nextTasks = (0, algo_js_1._prepareNextTasks)(this.checkpoint, this.nodes, this.channels, this.config, true, {
+        const nextTasks = (0, algo_js_1._prepareNextTasks)(this.checkpoint, this.nodes, this.channels, this.managed, this.config, true, {
             step: this.step,
             checkpointer: this.checkpointer,
             isResuming: this.input === INPUT_RESUMING,
@@ -377,7 +414,7 @@ class PregelLoop {
             if (inputWrites.length === 0) {
                 throw new errors_js_1.EmptyInputError(`Received no input writes for ${JSON.stringify(inputKeys, null, 2)}`);
             }
-            const discardTasks = (0, algo_js_1._prepareNextTasks)(this.checkpoint, this.nodes, this.channels, this.config, true, { step: this.step });
+            const discardTasks = (0, algo_js_1._prepareNextTasks)(this.checkpoint, this.nodes, this.channels, this.managed, this.config, true, { step: this.step });
             (0, algo_js_1._applyWrites)(this.checkpoint, this.channels, discardTasks.concat([
                 {
                     name: constants_js_1.INPUT,

package/dist/pregel/loop.d.ts

@@ -5,6 +5,9 @@ import { BaseCheckpointSaver, Checkpoint, PendingWrite, CheckpointPendingWrite,
 import { BaseChannel } from "../channels/base.js";
 import { PregelExecutableTask, StreamMode } from "./types.js";
 import { PregelNode } from "./read.js";
+import { BaseStore } from "../store/base.js";
+import { AsyncBatchedStore } from "../store/batch.js";
+import { ManagedValueMapping } from "../managed/base.js";
 export type PregelLoopInitializeParams = {
     input?: any;
     config: RunnableConfig;
@@ -13,6 +16,8 @@ export type PregelLoopInitializeParams = {
     streamKeys: string | string[];
     nodes: Record<string, PregelNode>;
     channelSpecs: Record<string, BaseChannel>;
+    managed: ManagedValueMapping;
+    store?: BaseStore;
 };
 type PregelLoopParams = {
     input?: any;
@@ -24,11 +29,13 @@ type PregelLoopParams = {
     checkpointPendingWrites: CheckpointPendingWrite[];
     checkpointConfig: RunnableConfig;
     channels: Record<string, BaseChannel>;
+    managed: ManagedValueMapping;
     step: number;
     stop: number;
     outputKeys: string | string[];
     streamKeys: string | string[];
     nodes: Record<string, PregelNode>;
+    store?: AsyncBatchedStore;
 };
 export declare class PregelLoop {
     protected input?: any;
@@ -36,6 +43,7 @@ export declare class PregelLoop {
     protected checkpointer?: BaseCheckpointSaver;
     protected checkpointerGetNextVersion: (current: number | undefined, channel: BaseChannel) => number;
     channels: Record<string, BaseChannel>;
+    managed: ManagedValueMapping;
     protected checkpoint: Checkpoint;
     protected checkpointConfig: RunnableConfig;
     checkpointMetadata: CheckpointMetadata;
@@ -53,6 +61,7 @@ export declare class PregelLoop {
     checkpointerPromises: Promise<unknown>[];
     protected isNested: boolean;
     protected _checkpointerChainedPromise: Promise<unknown>;
+    store?: AsyncBatchedStore;
     constructor(params: PregelLoopParams);
     static initialize(params: PregelLoopInitializeParams): Promise<PregelLoop>;
     protected _checkpointerPutAfterPrevious(input: {
@@ -61,6 +70,7 @@ export declare class PregelLoop {
         metadata: CheckpointMetadata;
         newVersions: Record<string, string | number>;
     }): void;
+    protected updateManagedValues(key: string, values: any[]): Promise<void>;
     /**
      * Put writes for a task, to be read by the next tick.
      * @param taskId

package/dist/pregel/loop.js

@@ -8,6 +8,7 @@ import { mapInput, mapOutputUpdates, mapOutputValues } from "./io.js";
 import { EmptyInputError, GraphInterrupt } from "../errors.js";
 import { getNewChannelVersions } from "./utils.js";
 import { mapDebugTasks, mapDebugCheckpoint, mapDebugTaskResults, } from "./debug.js";
+import { AsyncBatchedStore } from "../store/batch.js";
 const INPUT_DONE = Symbol.for("INPUT_DONE");
 const INPUT_RESUMING = Symbol.for("INPUT_RESUMING");
 const DEFAULT_LOOP_LIMIT = 25;
@@ -44,6 +45,12 @@ export class PregelLoop {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "managed", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "checkpoint", {
             enumerable: true,
             configurable: true,
@@ -148,6 +155,12 @@ export class PregelLoop {
             writable: true,
             value: Promise.resolve()
         });
+        Object.defineProperty(this, "store", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         this.input = params.input;
         this.config = params.config;
         this.checkpointer = params.checkpointer;
@@ -164,6 +177,7 @@ export class PregelLoop {
         this.checkpointMetadata = params.checkpointMetadata;
         this.checkpointPreviousVersions = params.checkpointPreviousVersions;
         this.channels = params.channels;
+        this.managed = params.managed;
         this.checkpointPendingWrites = params.checkpointPendingWrites;
         this.step = params.step;
         this.stop = params.stop;
@@ -172,6 +186,7 @@ export class PregelLoop {
         this.streamKeys = params.streamKeys;
         this.nodes = params.nodes;
         this.skipDoneTasks = this.config.configurable?.checkpoint_id === undefined;
+        this.store = params.store;
     }
     static async initialize(params) {
         const saved = (await params.checkpointer?.getTuple(params.config)) ?? {
@@ -199,6 +214,13 @@ export class PregelLoop {
         const step = (checkpointMetadata.step ?? 0) + 1;
         const stop = step + (params.config.recursionLimit ?? DEFAULT_LOOP_LIMIT) + 1;
         const checkpointPreviousVersions = { ...checkpoint.channel_versions };
+        const store = params.store
+            ? new AsyncBatchedStore(params.store)
+            : undefined;
+        if (store) {
+            // Start the store. This is a batch store, so it will run continuously
+            store.start();
+        }
         return new PregelLoop({
             input: params.input,
             config: params.config,
@@ -207,6 +229,7 @@ export class PregelLoop {
             checkpointMetadata,
             checkpointConfig,
             channels,
+            managed: params.managed,
             step,
             stop,
             checkpointPreviousVersions,
@@ -214,6 +237,7 @@ export class PregelLoop {
             outputKeys: params.outputKeys ?? [],
             streamKeys: params.streamKeys ?? [],
             nodes: params.nodes,
+            store,
         });
     }
     _checkpointerPutAfterPrevious(input) {
@@ -222,6 +246,13 @@ export class PregelLoop {
         });
         this.checkpointerPromises.push(this._checkpointerChainedPromise);
     }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    async updateManagedValues(key, values) {
+        const mv = this.managed.get(key);
+        if (mv && "update" in mv && typeof mv.update === "function") {
+            await mv.update(values);
+        }
+    }
     /**
      * Put writes for a task, to be read by the next tick.
      * @param taskId
@@ -255,6 +286,9 @@ export class PregelLoop {
      * @param params
      */
     async tick(params) {
+        if (this.store && !this.store.isRunning) {
+            this.store?.start();
+        }
         const { inputKeys = [], interruptAfter = [], interruptBefore = [], manager, } = params;
         if (this.status !== "pending") {
             throw new Error(`Cannot tick when status is no longer "pending". Current status: "${this.status}"`);
@@ -265,7 +299,10 @@ export class PregelLoop {
         else if (this.tasks.every((task) => task.writes.length > 0)) {
             const writes = this.tasks.flatMap((t) => t.writes);
             // All tasks have finished
-            _applyWrites(this.checkpoint, this.channels, this.tasks, this.checkpointerGetNextVersion);
+            const myWrites = _applyWrites(this.checkpoint, this.channels, this.tasks, this.checkpointerGetNextVersion);
+            for (const [key, values] of Object.entries(myWrites)) {
+                await this.updateManagedValues(key, values);
+            }
             // produce values output
             const valuesOutput = await gatherIterator(prefixGenerator(mapOutputValues(this.outputKeys, writes, this.channels), "values"));
             this.stream.push(...valuesOutput);
@@ -293,7 +330,7 @@ export class PregelLoop {
             this.status = "out_of_steps";
             return false;
         }
-        const nextTasks = _prepareNextTasks(this.checkpoint, this.nodes, this.channels, this.config, true, {
+        const nextTasks = _prepareNextTasks(this.checkpoint, this.nodes, this.channels, this.managed, this.config, true, {
             step: this.step,
             checkpointer: this.checkpointer,
             isResuming: this.input === INPUT_RESUMING,
@@ -371,7 +408,7 @@ export class PregelLoop {
             if (inputWrites.length === 0) {
                 throw new EmptyInputError(`Received no input writes for ${JSON.stringify(inputKeys, null, 2)}`);
             }
-            const discardTasks = _prepareNextTasks(this.checkpoint, this.nodes, this.channels, this.config, true, { step: this.step });
+            const discardTasks = _prepareNextTasks(this.checkpoint, this.nodes, this.channels, this.managed, this.config, true, { step: this.step });
             _applyWrites(this.checkpoint, this.channels, discardTasks.concat([
                 {
                     name: INPUT,