@langchain/langgraph 0.1.9 → 0.1.10-rc.0

package/dist/graph/graph.d.ts

@@ -1,11 +1,10 @@
 import { Runnable, RunnableConfig, RunnableLike } from "@langchain/core/runnables";
 import { Graph as RunnableGraph } from "@langchain/core/runnables/graph";
-import { BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
+import { All, BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
 import { PregelNode } from "../pregel/read.js";
 import { Pregel } from "../pregel/index.js";
 import type { PregelParams } from "../pregel/types.js";
 import { BaseChannel } from "../channels/base.js";
-import { All } from "../pregel/types.js";
 import { Send } from "../constants.js";
 import { RunnableCallable } from "../utils.js";
 /** Special reserved node name denoting the start of a graph. */

package/dist/graph/state.cjs

@@ -82,7 +82,7 @@ class StateGraph extends graph_js_1.Graph {
             enumerable: true,
             configurable: true,
             writable: true,
-            value: void 0
+            value: {}
         });
         // TODO: this doesn't dedupe edges as in py, so worth fixing at some point
         Object.defineProperty(this, "waitingEdges", {
@@ -91,21 +91,50 @@ class StateGraph extends graph_js_1.Graph {
             writable: true,
             value: new Set()
         });
-        if (isStateDefinition(fields) || isAnnotationRoot(fields)) {
+        Object.defineProperty(this, "schema", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "input", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "output", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        if (isStateGraphArgsWithInputOutputSchemas(fields)) {
+            this.schema = fields.input.spec;
+            this.input = fields.input.spec;
+            this.output = fields.output.spec;
+        }
+        else if (isStateGraphArgsWithStateSchema(fields)) {
+            this.schema = fields.stateSchema.spec;
+            this.input = (fields.input?.spec ?? this.schema);
+            this.output = (fields.output?.spec ?? this.schema);
+        }
+        else if (isStateDefinition(fields) || isAnnotationRoot(fields)) {
             const spec = isAnnotationRoot(fields) ? fields.spec : fields;
-            this.channels = {};
-            for (const [key, val] of Object.entries(spec)) {
-                if (typeof val === "function") {
-                    this.channels[key] = val();
-                }
-                else {
-                    this.channels[key] = val;
-                }
-            }
+            this.schema = spec;
+        }
+        else if (isStateGraphArgs(fields)) {
+            const spec = _getChannels(fields.channels);
+            this.schema = spec;
         }
         else {
-            this.channels = _getChannels(fields.channels);
+            throw new Error("Invalid StateGraph input.");
         }
+        this.input = this.input ?? this.schema;
+        this.output = this.output ?? this.schema;
+        this._addSchema(this.schema);
+        this._addSchema(this.input);
+        this._addSchema(this.output);
         for (const c of Object.values(this.channels)) {
             if (c.lc_graph_name === "BinaryOperatorAggregate") {
                 this.supportMultipleEdges = true;
@@ -119,6 +148,27 @@ class StateGraph extends graph_js_1.Graph {
             ...Array.from(this.waitingEdges).flatMap(([starts, end]) => starts.map((start) => [start, end])),
         ]);
     }
+    _addSchema(stateDefinition) {
+        for (const [key, val] of Object.entries(stateDefinition)) {
+            let channel;
+            if (typeof val === "function") {
+                channel = val();
+            }
+            else {
+                channel = val;
+            }
+            if (this.channels[key] !== undefined) {
+                if (this.channels[key].lc_graph_name !== channel.lc_graph_name) {
+                    if (channel.lc_graph_name !== "LastValue") {
+                        throw new Error(`Channel "${key}" already exists with a different type.`);
+                    }
+                }
+            }
+            else {
+                this.channels[key] = channel;
+            }
+        }
+    }
     addNode(key, action, options) {
         if (key in this.channels) {
             throw new Error(`${key} is already being used as a state attribute (a.k.a. a channel), cannot also be used as a node name.`);
@@ -173,10 +223,10 @@ class StateGraph extends graph_js_1.Graph {
             ...(Array.isArray(interruptAfter) ? interruptAfter : []),
         ]);
         // prepare output channels
-        const stateKeys = Object.keys(this.channels);
-        const outputChannels = stateKeys.length === 1 && stateKeys[0] === ROOT
-            ? stateKeys[0]
-            : stateKeys;
+        const stateKeys = Object.keys(this.schema);
+        const outputChannels = stateKeys.length === 1 && stateKeys[0] === ROOT ? ROOT : stateKeys;
+        const outputKeys = Object.keys(this.output);
+        const streamChannels = outputKeys.length === 1 && outputKeys[0] === ROOT ? ROOT : outputKeys;
         // create empty compiled graph
         const compiled = new CompiledStateGraph({
             builder: this,
@@ -191,7 +241,7 @@ class StateGraph extends graph_js_1.Graph {
             },
             inputChannels: graph_js_1.START,
             outputChannels,
-            streamChannels: outputChannels,
+            streamChannels,
             streamMode: "updates",
         });
         // attach nodes, edges and branches
@@ -229,7 +279,7 @@ function _getChannels(schema) {
 }
 class CompiledStateGraph extends graph_js_1.CompiledGraph {
     attachNode(key, node) {
-        const stateKeys = Object.keys(this.builder.channels);
+        const stateKeys = Object.keys(this.builder.input);
         function getStateKey(key, input) {
             if (!input) {
                 return write_js_1.SKIP_WRITE;
@@ -367,3 +417,20 @@ function isAnnotationRoot(obj) {
         "lc_graph_name" in obj &&
         obj.lc_graph_name === "AnnotationRoot");
 }
+function isStateGraphArgs(obj) {
+    return (typeof obj === "object" &&
+        obj !== null &&
+        obj.channels !== undefined);
+}
+function isStateGraphArgsWithStateSchema(obj) {
+    return (typeof obj === "object" &&
+        obj !== null &&
+        obj.stateSchema !== undefined);
+}
+function isStateGraphArgsWithInputOutputSchemas(obj) {
+    return (typeof obj === "object" &&
+        obj !== null &&
+        obj.stateSchema === undefined &&
+        obj.input !== undefined &&
+        obj.output !== undefined);
+}

package/dist/graph/state.d.ts

@@ -1,8 +1,7 @@
 import { RunnableLike } from "@langchain/core/runnables";
-import { BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
+import { All, BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
 import { BaseChannel } from "../channels/base.js";
 import { END, CompiledGraph, Graph, START, Branch, AddNodeOptions, NodeSpec } from "./graph.js";
-import { All } from "../pregel/types.js";
 import { AnnotationRoot, SingleReducer, StateDefinition, StateType, UpdateType } from "./annotation.js";
 import type { RetryPolicy } from "../pregel/utils.js";
 export type ChannelReducers<Channels extends object> = {
@@ -22,6 +21,15 @@ export type StateGraphNodeSpec<RunInput, RunOutput> = NodeSpec<RunInput, RunOutp
 export type StateGraphAddNodeOptions = {
     retryPolicy?: RetryPolicy;
 } & AddNodeOptions;
+export type StateGraphArgsWithStateSchema<SD extends StateDefinition, I extends StateDefinition, O extends StateDefinition> = {
+    stateSchema: AnnotationRoot<SD>;
+    input?: AnnotationRoot<I>;
+    output?: AnnotationRoot<O>;
+};
+export type StateGraphArgsWithInputOutputSchemas<SD extends StateDefinition, O extends StateDefinition = SD> = {
+    input: AnnotationRoot<SD>;
+    output: AnnotationRoot<O>;
+};
 /**
  * A graph whose nodes communicate by reading and writing to a shared state.
  * Each node takes a defined `State` as input and returns a `Partial<State>`.
@@ -84,21 +92,25 @@ export type StateGraphAddNodeOptions = {
  * // }
  * ```
  */
-export declare class StateGraph<SD extends StateDefinition | unknown, S = SD extends StateDefinition ? StateType<SD> : SD, U = SD extends StateDefinition ? UpdateType<SD> : Partial<S>, N extends string = typeof START> extends Graph<N, S, U, StateGraphNodeSpec<S, U>> {
+export declare class StateGraph<SD extends StateDefinition | unknown, S = SD extends StateDefinition ? StateType<SD> : SD, U = SD extends StateDefinition ? UpdateType<SD> : Partial<S>, N extends string = typeof START, I extends StateDefinition = SD extends StateDefinition ? SD : StateDefinition, O extends StateDefinition = SD extends StateDefinition ? SD : StateDefinition> extends Graph<N, S, U, StateGraphNodeSpec<S, U>> {
     channels: Record<string, BaseChannel>;
     waitingEdges: Set<[N[], N]>;
-    constructor(fields: SD extends StateDefinition ? SD | AnnotationRoot<SD> | StateGraphArgs<S> : StateGraphArgs<S>);
+    schema: StateDefinition;
+    input: I;
+    output: O;
+    constructor(fields: SD extends StateDefinition ? SD | AnnotationRoot<SD> | StateGraphArgs<S> | StateGraphArgsWithStateSchema<SD, I, O> | StateGraphArgsWithInputOutputSchemas<SD, O> : StateGraphArgs<S>);
     get allEdges(): Set<[string, string]>;
-    addNode<K extends string, NodeInput = S>(key: K, action: RunnableLike<NodeInput, U>, options?: StateGraphAddNodeOptions): StateGraph<SD, S, U, N | K>;
+    _addSchema(stateDefinition: StateDefinition): void;
+    addNode<K extends string, NodeInput = S>(key: K, action: RunnableLike<NodeInput, U>, options?: StateGraphAddNodeOptions): StateGraph<SD, S, U, N | K, I, O>;
     addEdge(startKey: typeof START | N | N[], endKey: N | typeof END): this;
     compile({ checkpointer, interruptBefore, interruptAfter, }?: {
         checkpointer?: BaseCheckpointSaver;
         interruptBefore?: N[] | All;
         interruptAfter?: N[] | All;
-    }): CompiledStateGraph<S, U, N>;
+    }): CompiledStateGraph<S, U, N, I, O>;
 }
-export declare class CompiledStateGraph<S, U, N extends string = typeof START> extends CompiledGraph<N, S, U> {
-    builder: StateGraph<unknown, S, U, N>;
+export declare class CompiledStateGraph<S, U, N extends string = typeof START, I extends StateDefinition = StateDefinition, O extends StateDefinition = StateDefinition> extends CompiledGraph<N, S, U> {
+    builder: StateGraph<unknown, S, U, N, I, O>;
     attachNode(key: typeof START, node?: never): void;
     attachNode(key: N, node: StateGraphNodeSpec<S, U>): void;
     attachEdge(start: N | N[] | "__start__", end: N | "__end__"): void;

package/dist/graph/state.js

@@ -79,7 +79,7 @@ export class StateGraph extends Graph {
             enumerable: true,
             configurable: true,
             writable: true,
-            value: void 0
+            value: {}
         });
         // TODO: this doesn't dedupe edges as in py, so worth fixing at some point
         Object.defineProperty(this, "waitingEdges", {
@@ -88,21 +88,50 @@ export class StateGraph extends Graph {
             writable: true,
             value: new Set()
         });
-        if (isStateDefinition(fields) || isAnnotationRoot(fields)) {
+        Object.defineProperty(this, "schema", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "input", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "output", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        if (isStateGraphArgsWithInputOutputSchemas(fields)) {
+            this.schema = fields.input.spec;
+            this.input = fields.input.spec;
+            this.output = fields.output.spec;
+        }
+        else if (isStateGraphArgsWithStateSchema(fields)) {
+            this.schema = fields.stateSchema.spec;
+            this.input = (fields.input?.spec ?? this.schema);
+            this.output = (fields.output?.spec ?? this.schema);
+        }
+        else if (isStateDefinition(fields) || isAnnotationRoot(fields)) {
             const spec = isAnnotationRoot(fields) ? fields.spec : fields;
-            this.channels = {};
-            for (const [key, val] of Object.entries(spec)) {
-                if (typeof val === "function") {
-                    this.channels[key] = val();
-                }
-                else {
-                    this.channels[key] = val;
-                }
-            }
+            this.schema = spec;
+        }
+        else if (isStateGraphArgs(fields)) {
+            const spec = _getChannels(fields.channels);
+            this.schema = spec;
         }
         else {
-            this.channels = _getChannels(fields.channels);
+            throw new Error("Invalid StateGraph input.");
         }
+        this.input = this.input ?? this.schema;
+        this.output = this.output ?? this.schema;
+        this._addSchema(this.schema);
+        this._addSchema(this.input);
+        this._addSchema(this.output);
         for (const c of Object.values(this.channels)) {
             if (c.lc_graph_name === "BinaryOperatorAggregate") {
                 this.supportMultipleEdges = true;
@@ -116,6 +145,27 @@ export class StateGraph extends Graph {
             ...Array.from(this.waitingEdges).flatMap(([starts, end]) => starts.map((start) => [start, end])),
         ]);
     }
+    _addSchema(stateDefinition) {
+        for (const [key, val] of Object.entries(stateDefinition)) {
+            let channel;
+            if (typeof val === "function") {
+                channel = val();
+            }
+            else {
+                channel = val;
+            }
+            if (this.channels[key] !== undefined) {
+                if (this.channels[key].lc_graph_name !== channel.lc_graph_name) {
+                    if (channel.lc_graph_name !== "LastValue") {
+                        throw new Error(`Channel "${key}" already exists with a different type.`);
+                    }
+                }
+            }
+            else {
+                this.channels[key] = channel;
+            }
+        }
+    }
     addNode(key, action, options) {
         if (key in this.channels) {
             throw new Error(`${key} is already being used as a state attribute (a.k.a. a channel), cannot also be used as a node name.`);
@@ -170,10 +220,10 @@ export class StateGraph extends Graph {
             ...(Array.isArray(interruptAfter) ? interruptAfter : []),
         ]);
         // prepare output channels
-        const stateKeys = Object.keys(this.channels);
-        const outputChannels = stateKeys.length === 1 && stateKeys[0] === ROOT
-            ? stateKeys[0]
-            : stateKeys;
+        const stateKeys = Object.keys(this.schema);
+        const outputChannels = stateKeys.length === 1 && stateKeys[0] === ROOT ? ROOT : stateKeys;
+        const outputKeys = Object.keys(this.output);
+        const streamChannels = outputKeys.length === 1 && outputKeys[0] === ROOT ? ROOT : outputKeys;
         // create empty compiled graph
         const compiled = new CompiledStateGraph({
             builder: this,
@@ -188,7 +238,7 @@ export class StateGraph extends Graph {
             },
             inputChannels: START,
             outputChannels,
-            streamChannels: outputChannels,
+            streamChannels,
             streamMode: "updates",
         });
         // attach nodes, edges and branches
@@ -225,7 +275,7 @@ function _getChannels(schema) {
 }
 export class CompiledStateGraph extends CompiledGraph {
     attachNode(key, node) {
-        const stateKeys = Object.keys(this.builder.channels);
+        const stateKeys = Object.keys(this.builder.input);
         function getStateKey(key, input) {
             if (!input) {
                 return SKIP_WRITE;
@@ -362,3 +412,20 @@ function isAnnotationRoot(obj) {
         "lc_graph_name" in obj &&
         obj.lc_graph_name === "AnnotationRoot");
 }
+function isStateGraphArgs(obj) {
+    return (typeof obj === "object" &&
+        obj !== null &&
+        obj.channels !== undefined);
+}
+function isStateGraphArgsWithStateSchema(obj) {
+    return (typeof obj === "object" &&
+        obj !== null &&
+        obj.stateSchema !== undefined);
+}
+function isStateGraphArgsWithInputOutputSchemas(obj) {
+    return (typeof obj === "object" &&
+        obj !== null &&
+        obj.stateSchema === undefined &&
+        obj.input !== undefined &&
+        obj.output !== undefined);
+}

package/dist/prebuilt/agent_executor.d.ts

@@ -18,5 +18,5 @@ export interface AgentExecutorState {
 export declare function createAgentExecutor({ agentRunnable, tools, }: {
     agentRunnable: Runnable;
     tools: Array<Tool> | ToolExecutor;
-}): import("../graph/state.js").CompiledStateGraph<AgentExecutorState, Partial<AgentExecutorState>, "__start__" | "agent" | "action">;
+}): import("../graph/state.js").CompiledStateGraph<AgentExecutorState, Partial<AgentExecutorState>, "__start__" | "agent" | "action", import("../graph/annotation.js").StateDefinition, import("../graph/annotation.js").StateDefinition>;
 export {};

package/dist/prebuilt/react_agent_executor.d.ts

@@ -2,11 +2,10 @@ import { BaseChatModel } from "@langchain/core/language_models/chat_models";
 import { BaseMessage, SystemMessage } from "@langchain/core/messages";
 import { Runnable, RunnableToolLike } from "@langchain/core/runnables";
 import { StructuredToolInterface } from "@langchain/core/tools";
-import { BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
+import { All, BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
 import { START } from "../graph/index.js";
 import { MessagesAnnotation } from "../graph/messages_annotation.js";
 import { CompiledStateGraph } from "../graph/state.js";
-import { All } from "../pregel/types.js";
 import { ToolNode } from "./tool_node.js";
 export interface AgentState {
     messages: BaseMessage[];

package/dist/pregel/algo.d.ts

@@ -1,9 +1,9 @@
 import { RunnableConfig } from "@langchain/core/runnables";
 import { CallbackManagerForChainRun } from "@langchain/core/callbacks/manager";
-import { BaseCheckpointSaver, Checkpoint, ReadonlyCheckpoint, type PendingWrite } from "@langchain/langgraph-checkpoint";
+import { All, BaseCheckpointSaver, Checkpoint, ReadonlyCheckpoint, type PendingWrite } from "@langchain/langgraph-checkpoint";
 import { BaseChannel } from "../channels/base.js";
 import { PregelNode } from "./read.js";
-import { All, PregelExecutableTask, PregelTaskDescription } from "./types.js";
+import { PregelExecutableTask, PregelTaskDescription } from "./types.js";
 /**
  * Construct a type with a set of properties K of type T
  */

package/dist/pregel/index.cjs

@@ -464,13 +464,15 @@ class Pregel extends runnables_1.Runnable {
                 input,
                 config,
                 checkpointer,
-                graph: this,
+                nodes: this.nodes,
+                channelSpecs: this.channels,
                 onBackgroundError,
                 outputKeys,
                 streamKeys: this.streamChannelsAsIs,
             });
             while (backgroundError === undefined &&
                 (await loop.tick({
+                    inputKeys: this.inputChannels,
                     interruptAfter,
                     interruptBefore,
                     manager: runManager,

package/dist/pregel/index.d.ts

@@ -1,10 +1,10 @@
 import { Runnable, RunnableConfig, RunnableFunc } from "@langchain/core/runnables";
 import { IterableReadableStream } from "@langchain/core/utils/stream";
-import { BaseCheckpointSaver, CheckpointListOptions } from "@langchain/langgraph-checkpoint";
+import { All, BaseCheckpointSaver, CheckpointListOptions } from "@langchain/langgraph-checkpoint";
 import { BaseChannel } from "../channels/base.js";
 import { PregelNode } from "./read.js";
 import { ChannelWrite } from "./write.js";
-import { All, PregelInterface, PregelParams, StateSnapshot, StreamMode } from "./types.js";
+import { PregelInterface, PregelParams, StateSnapshot, StreamMode } from "./types.js";
 import { StrRecord } from "./algo.js";
 import { RetryPolicy } from "./utils.js";
 type WriteValue = Runnable | RunnableFunc<unknown, unknown> | unknown;

package/dist/pregel/index.js

@@ -460,13 +460,15 @@ export class Pregel extends Runnable {
                 input,
                 config,
                 checkpointer,
-                graph: this,
+                nodes: this.nodes,
+                channelSpecs: this.channels,
                 onBackgroundError,
                 outputKeys,
                 streamKeys: this.streamChannelsAsIs,
             });
             while (backgroundError === undefined &&
                 (await loop.tick({
+                    inputKeys: this.inputChannels,
                     interruptAfter,
                     interruptBefore,
                     manager: runManager,

package/dist/pregel/loop.cjs

@@ -47,14 +47,6 @@ class PregelLoop {
             writable: true,
             value: void 0
         });
-        // TODO: Fix typing
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        Object.defineProperty(this, "graph", {
-            enumerable: true,
-            configurable: true,
-            writable: true,
-            value: void 0
-        });
         Object.defineProperty(this, "channels", {
             enumerable: true,
             configurable: true,
@@ -103,6 +95,24 @@ class PregelLoop {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "outputKeys", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "streamKeys", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "nodes", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "status", {
             enumerable: true,
             configurable: true,
@@ -135,18 +145,6 @@ class PregelLoop {
             writable: true,
             value: Promise.resolve()
         });
-        Object.defineProperty(this, "outputKeys", {
-            enumerable: true,
-            configurable: true,
-            writable: true,
-            value: void 0
-        });
-        Object.defineProperty(this, "streamKeys", {
-            enumerable: true,
-            configurable: true,
-            writable: true,
-            value: void 0
-        });
         Object.defineProperty(this, "onBackgroundError", {
             enumerable: true,
             configurable: true,
@@ -164,7 +162,6 @@ class PregelLoop {
         else {
             this.checkpointerGetNextVersion = algo_js_1.increment;
         }
-        this.graph = params.graph;
         this.checkpoint = params.checkpoint;
         this.checkpointConfig = params.checkpointConfig;
         this.checkpointMetadata = params.checkpointMetadata;
@@ -176,6 +173,7 @@ class PregelLoop {
         this.isNested = constants_js_1.CONFIG_KEY_READ in (this.config.configurable ?? {});
         this.outputKeys = params.outputKeys;
         this.streamKeys = params.streamKeys;
+        this.nodes = params.nodes;
         this.onBackgroundError = params.onBackgroundError;
     }
     static async initialize(params) {
@@ -200,7 +198,7 @@ class PregelLoop {
         const checkpoint = (0, langgraph_checkpoint_1.copyCheckpoint)(saved.checkpoint);
         const checkpointMetadata = { ...saved.metadata };
         const checkpointPendingWrites = saved.pendingWrites ?? [];
-        const channels = (0, base_js_1.emptyChannels)(params.graph.channels, checkpoint);
+        const channels = (0, base_js_1.emptyChannels)(params.channelSpecs, checkpoint);
         const step = (checkpointMetadata.step ?? 0) + 1;
         const stop = step + (params.config.recursionLimit ?? DEFAULT_LOOP_LIMIT) + 1;
         const checkpointPreviousVersions = { ...checkpoint.channel_versions };
@@ -208,7 +206,6 @@ class PregelLoop {
             input: params.input,
             config: params.config,
             checkpointer: params.checkpointer,
-            graph: params.graph,
             checkpoint,
             checkpointMetadata,
             checkpointConfig,
@@ -217,8 +214,9 @@ class PregelLoop {
             stop,
             checkpointPreviousVersions,
             checkpointPendingWrites,
-            outputKeys: params.outputKeys,
-            streamKeys: params.streamKeys,
+            outputKeys: params.outputKeys ?? [],
+            streamKeys: params.streamKeys ?? [],
+            nodes: params.nodes,
             onBackgroundError: params.onBackgroundError,
         });
     }
@@ -267,12 +265,12 @@ class PregelLoop {
      * @param params
      */
     async tick(params) {
-        const { interruptAfter = [], interruptBefore = [], manager } = params;
+        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}"`);
         }
         if (![INPUT_DONE, INPUT_RESUMING].includes(this.input)) {
-            await this._first();
+            await this._first(inputKeys);
         }
         else if (this.tasks.every((task) => task.writes.length > 0)) {
             const writes = this.tasks.flatMap((t) => t.writes);
@@ -283,10 +281,9 @@ class PregelLoop {
             this.stream.push(...valuesOutput);
             // clear pending writes
             this.checkpointPendingWrites = [];
-            const metadataWrites = (0, io_js_1.mapOutputUpdates)(this.outputKeys, this.tasks).next().value;
             await this._putCheckpoint({
                 source: "loop",
-                writes: metadataWrites ?? null,
+                writes: (0, io_js_1.mapOutputUpdates)(this.outputKeys, this.tasks).next().value ?? null,
             });
             // after execution, check if we should interrupt
             if ((0, algo_js_1.shouldInterrupt)(this.checkpoint, interruptAfter, this.tasks)) {
@@ -306,7 +303,7 @@ class PregelLoop {
             this.status = "out_of_steps";
             return false;
         }
-        const nextTasks = (0, algo_js_1._prepareNextTasks)(this.checkpoint, this.graph.nodes, this.channels, this.config, true, {
+        const nextTasks = (0, algo_js_1._prepareNextTasks)(this.checkpoint, this.nodes, this.channels, this.config, true, {
             step: this.step,
             checkpointer: this.checkpointer,
             isResuming: this.input === INPUT_RESUMING,
@@ -316,7 +313,7 @@ class PregelLoop {
         // Produce debug output
         if (this.checkpointer) {
             this.stream.push(...(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.graph.streamChannelsAsIs, this.checkpointMetadata, this.tasks, this.checkpointPendingWrites), "debug"))));
+            this.checkpointConfig, this.channels, this.streamKeys, this.checkpointMetadata, this.tasks, this.checkpointPendingWrites), "debug"))));
         }
         if (this.tasks.length === 0) {
             this.status = "done";
@@ -337,6 +334,7 @@ class PregelLoop {
         // if all tasks have finished, re-tick
         if (this.tasks.every((task) => task.writes.length > 0)) {
             return this.tick({
+                inputKeys,
                 interruptAfter,
                 interruptBefore,
                 manager,
@@ -362,7 +360,7 @@ class PregelLoop {
      * - finding a previous checkpoint
      * - receiving None input (outer graph) or RESUMING flag (subgraph)
      */
-    async _first() {
+    async _first(inputKeys) {
         const isResuming = (Object.keys(this.checkpoint.channel_versions).length !== 0 &&
             this.config.configurable?.[constants_js_1.CONFIG_KEY_RESUMING] !== undefined) ||
             this.input === null;
@@ -379,11 +377,11 @@ class PregelLoop {
             // map inputs to channel updates
         }
         else {
-            const inputWrites = await (0, utils_js_1.gatherIterator)((0, io_js_1.mapInput)(this.graph.inputChannels, this.input));
+            const inputWrites = await (0, utils_js_1.gatherIterator)((0, io_js_1.mapInput)(inputKeys, this.input));
             if (inputWrites.length === 0) {
-                throw new errors_js_1.EmptyInputError(`Received no input writes for ${JSON.stringify(this.graph.inputChannels, null, 2)}`);
+                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.graph.nodes, this.channels, this.config, true, { step: this.step });
+            const discardTasks = (0, algo_js_1._prepareNextTasks)(this.checkpoint, this.nodes, this.channels, 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

@@ -3,21 +3,22 @@ import type { RunnableConfig } from "@langchain/core/runnables";
 import type { CallbackManagerForChainRun } from "@langchain/core/callbacks/manager";
 import { BaseCheckpointSaver, Checkpoint, PendingWrite, CheckpointPendingWrite, CheckpointMetadata, All } from "@langchain/langgraph-checkpoint";
 import { BaseChannel } from "../channels/base.js";
-import { PregelExecutableTask, PregelInterface, StreamMode } from "./types.js";
+import { PregelExecutableTask, StreamMode } from "./types.js";
+import { PregelNode } from "./read.js";
 export type PregelLoopInitializeParams = {
     input?: any;
     config: RunnableConfig;
     checkpointer?: BaseCheckpointSaver;
-    graph: PregelInterface<any, any>;
+    onBackgroundError: (e: Error) => void;
     outputKeys: string | string[];
     streamKeys: string | string[];
-    onBackgroundError: (e: Error) => void;
+    nodes: Record<string, PregelNode>;
+    channelSpecs: Record<string, BaseChannel>;
 };
 type PregelLoopParams = {
     input?: any;
     config: RunnableConfig;
     checkpointer?: BaseCheckpointSaver;
-    graph: PregelInterface<any, any>;
     checkpoint: Checkpoint;
     checkpointMetadata: CheckpointMetadata;
     checkpointPreviousVersions: Record<string, string | number>;
@@ -29,13 +30,13 @@ type PregelLoopParams = {
     outputKeys: string | string[];
     streamKeys: string | string[];
     onBackgroundError: (e: Error) => void;
+    nodes: Record<string, PregelNode>;
 };
 export declare class PregelLoop {
     protected input?: any;
     config: RunnableConfig;
     protected checkpointer?: BaseCheckpointSaver;
     protected checkpointerGetNextVersion: (current: number | undefined, channel: BaseChannel) => number;
-    protected graph: PregelInterface<any, any>;
     channels: Record<string, BaseChannel>;
     protected checkpoint: Checkpoint;
     protected checkpointConfig: RunnableConfig;
@@ -44,13 +45,14 @@ export declare class PregelLoop {
     protected checkpointPreviousVersions: Record<string, string | number>;
     step: number;
     protected stop: number;
+    protected outputKeys: string | string[];
+    protected streamKeys: string | string[];
+    protected nodes: Record<string, PregelNode>;
     status: "pending" | "done" | "interrupt_before" | "interrupt_after" | "out_of_steps";
     tasks: PregelExecutableTask<any, any>[];
     stream: Deque<[StreamMode, any]>;
     protected isNested: boolean;
     protected _putCheckpointPromise: Promise<unknown>;
-    outputKeys: string | string[];
-    streamKeys: string | string[];
     onBackgroundError: (e: Error) => void;
     get backgroundTasksPromise(): Promise<unknown>;
     constructor(params: PregelLoopParams);
@@ -73,6 +75,7 @@ export declare class PregelLoop {
      * @param params
      */
     tick(params: {
+        inputKeys?: string | string[];
         interruptAfter: string[] | All;
         interruptBefore: string[] | All;
         manager?: CallbackManagerForChainRun;
@@ -82,7 +85,7 @@ export declare class PregelLoop {
      * - finding a previous checkpoint
      * - receiving None input (outer graph) or RESUMING flag (subgraph)
      */
-    protected _first(): Promise<void>;
+    protected _first(inputKeys: string | string[]): Promise<void>;
     protected _putCheckpoint(inputMetadata: Omit<CheckpointMetadata, "step">): Promise<void>;
 }
 export {};

package/dist/pregel/loop.js

@@ -41,14 +41,6 @@ export class PregelLoop {
             writable: true,
             value: void 0
         });
-        // TODO: Fix typing
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        Object.defineProperty(this, "graph", {
-            enumerable: true,
-            configurable: true,
-            writable: true,
-            value: void 0
-        });
         Object.defineProperty(this, "channels", {
             enumerable: true,
             configurable: true,
@@ -97,6 +89,24 @@ export class PregelLoop {
             writable: true,
             value: void 0
         });
+        Object.defineProperty(this, "outputKeys", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "streamKeys", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        Object.defineProperty(this, "nodes", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
         Object.defineProperty(this, "status", {
             enumerable: true,
             configurable: true,
@@ -129,18 +139,6 @@ export class PregelLoop {
             writable: true,
             value: Promise.resolve()
         });
-        Object.defineProperty(this, "outputKeys", {
-            enumerable: true,
-            configurable: true,
-            writable: true,
-            value: void 0
-        });
-        Object.defineProperty(this, "streamKeys", {
-            enumerable: true,
-            configurable: true,
-            writable: true,
-            value: void 0
-        });
         Object.defineProperty(this, "onBackgroundError", {
             enumerable: true,
             configurable: true,
@@ -158,7 +156,6 @@ export class PregelLoop {
         else {
             this.checkpointerGetNextVersion = increment;
         }
-        this.graph = params.graph;
         this.checkpoint = params.checkpoint;
         this.checkpointConfig = params.checkpointConfig;
         this.checkpointMetadata = params.checkpointMetadata;
@@ -170,6 +167,7 @@ export class PregelLoop {
         this.isNested = CONFIG_KEY_READ in (this.config.configurable ?? {});
         this.outputKeys = params.outputKeys;
         this.streamKeys = params.streamKeys;
+        this.nodes = params.nodes;
         this.onBackgroundError = params.onBackgroundError;
     }
     static async initialize(params) {
@@ -194,7 +192,7 @@ export class PregelLoop {
         const checkpoint = copyCheckpoint(saved.checkpoint);
         const checkpointMetadata = { ...saved.metadata };
         const checkpointPendingWrites = saved.pendingWrites ?? [];
-        const channels = emptyChannels(params.graph.channels, checkpoint);
+        const channels = emptyChannels(params.channelSpecs, checkpoint);
         const step = (checkpointMetadata.step ?? 0) + 1;
         const stop = step + (params.config.recursionLimit ?? DEFAULT_LOOP_LIMIT) + 1;
         const checkpointPreviousVersions = { ...checkpoint.channel_versions };
@@ -202,7 +200,6 @@ export class PregelLoop {
             input: params.input,
             config: params.config,
             checkpointer: params.checkpointer,
-            graph: params.graph,
             checkpoint,
             checkpointMetadata,
             checkpointConfig,
@@ -211,8 +208,9 @@ export class PregelLoop {
             stop,
             checkpointPreviousVersions,
             checkpointPendingWrites,
-            outputKeys: params.outputKeys,
-            streamKeys: params.streamKeys,
+            outputKeys: params.outputKeys ?? [],
+            streamKeys: params.streamKeys ?? [],
+            nodes: params.nodes,
             onBackgroundError: params.onBackgroundError,
         });
     }
@@ -261,12 +259,12 @@ export class PregelLoop {
      * @param params
      */
     async tick(params) {
-        const { interruptAfter = [], interruptBefore = [], manager } = params;
+        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}"`);
         }
         if (![INPUT_DONE, INPUT_RESUMING].includes(this.input)) {
-            await this._first();
+            await this._first(inputKeys);
         }
         else if (this.tasks.every((task) => task.writes.length > 0)) {
             const writes = this.tasks.flatMap((t) => t.writes);
@@ -277,10 +275,9 @@ export class PregelLoop {
             this.stream.push(...valuesOutput);
             // clear pending writes
             this.checkpointPendingWrites = [];
-            const metadataWrites = mapOutputUpdates(this.outputKeys, this.tasks).next().value;
             await this._putCheckpoint({
                 source: "loop",
-                writes: metadataWrites ?? null,
+                writes: mapOutputUpdates(this.outputKeys, this.tasks).next().value ?? null,
             });
             // after execution, check if we should interrupt
             if (shouldInterrupt(this.checkpoint, interruptAfter, this.tasks)) {
@@ -300,7 +297,7 @@ export class PregelLoop {
             this.status = "out_of_steps";
             return false;
         }
-        const nextTasks = _prepareNextTasks(this.checkpoint, this.graph.nodes, this.channels, this.config, true, {
+        const nextTasks = _prepareNextTasks(this.checkpoint, this.nodes, this.channels, this.config, true, {
             step: this.step,
             checkpointer: this.checkpointer,
             isResuming: this.input === INPUT_RESUMING,
@@ -310,7 +307,7 @@ export class PregelLoop {
         // Produce debug output
         if (this.checkpointer) {
             this.stream.push(...(await gatherIterator(prefixGenerator(mapDebugCheckpoint(this.step - 1, // printing checkpoint for previous step
-            this.checkpointConfig, this.channels, this.graph.streamChannelsAsIs, this.checkpointMetadata, this.tasks, this.checkpointPendingWrites), "debug"))));
+            this.checkpointConfig, this.channels, this.streamKeys, this.checkpointMetadata, this.tasks, this.checkpointPendingWrites), "debug"))));
         }
         if (this.tasks.length === 0) {
             this.status = "done";
@@ -331,6 +328,7 @@ export class PregelLoop {
         // if all tasks have finished, re-tick
         if (this.tasks.every((task) => task.writes.length > 0)) {
             return this.tick({
+                inputKeys,
                 interruptAfter,
                 interruptBefore,
                 manager,
@@ -356,7 +354,7 @@ export class PregelLoop {
      * - finding a previous checkpoint
      * - receiving None input (outer graph) or RESUMING flag (subgraph)
      */
-    async _first() {
+    async _first(inputKeys) {
         const isResuming = (Object.keys(this.checkpoint.channel_versions).length !== 0 &&
             this.config.configurable?.[CONFIG_KEY_RESUMING] !== undefined) ||
             this.input === null;
@@ -373,11 +371,11 @@ export class PregelLoop {
             // map inputs to channel updates
         }
         else {
-            const inputWrites = await gatherIterator(mapInput(this.graph.inputChannels, this.input));
+            const inputWrites = await gatherIterator(mapInput(inputKeys, this.input));
             if (inputWrites.length === 0) {
-                throw new EmptyInputError(`Received no input writes for ${JSON.stringify(this.graph.inputChannels, null, 2)}`);
+                throw new EmptyInputError(`Received no input writes for ${JSON.stringify(inputKeys, null, 2)}`);
             }
-            const discardTasks = _prepareNextTasks(this.checkpoint, this.graph.nodes, this.channels, this.config, true, { step: this.step });
+            const discardTasks = _prepareNextTasks(this.checkpoint, this.nodes, this.channels, this.config, true, { step: this.step });
             _applyWrites(this.checkpoint, this.channels, discardTasks.concat([
                 {
                     name: INPUT,

package/dist/pregel/types.d.ts

@@ -1,5 +1,5 @@
 import type { Runnable, RunnableConfig } from "@langchain/core/runnables";
-import type { PendingWrite, CheckpointMetadata, BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
+import type { All, PendingWrite, CheckpointMetadata, BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
 import type { BaseChannel } from "../channels/base.js";
 import type { PregelNode } from "./read.js";
 import { RetryPolicy } from "./utils.js";
@@ -93,5 +93,4 @@ export interface StateSnapshot {
      */
     readonly tasks: PregelTaskDescription[];
 }
-export type All = "*";
 export {};

package/dist/pregel/validate.d.ts

@@ -1,6 +1,6 @@
+import { All } from "@langchain/langgraph-checkpoint";
 import { BaseChannel } from "../channels/index.js";
 import { PregelNode } from "./read.js";
-import { All } from "./types.js";
 export declare class GraphValidationError extends Error {
     constructor(message?: string);
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@langchain/langgraph",
-  "version": "0.1.9",
+  "version": "0.1.10-rc.0",
   "description": "LangGraph",
   "type": "module",
   "engines": {