@langchain/langgraph 0.0.12 → 0.0.13

package/dist/graph/state.cjs

@@ -1,14 +1,18 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.StateGraph = exports.START = void 0;
-const runnables_1 = require("@langchain/core/runnables");
+exports.CompiledStateGraph = exports.StateGraph = void 0;
 const binop_js_1 = require("../channels/binop.cjs");
 const graph_js_1 = require("./graph.cjs");
 const last_value_js_1 = require("../channels/last_value.cjs");
 const write_js_1 = require("../pregel/write.cjs");
-const index_js_1 = require("../pregel/index.cjs");
 const read_js_1 = require("../pregel/read.cjs");
-exports.START = "__start__";
+const named_barrier_value_js_1 = require("../channels/named_barrier_value.cjs");
+const ephemeral_value_js_1 = require("../channels/ephemeral_value.cjs");
+const utils_js_1 = require("../utils.cjs");
+const constants_js_1 = require("../constants.cjs");
+const errors_js_1 = require("../errors.cjs");
+const dynamic_barrier_value_js_1 = require("../channels/dynamic_barrier_value.cjs");
+const ROOT = "__root__";
 class StateGraph extends graph_js_1.Graph {
     constructor(fields) {
         super();
@@ -18,129 +22,260 @@ class StateGraph extends graph_js_1.Graph {
             writable: true,
             value: void 0
         });
+        // TODO: this doesn't dedupe edges as in py, so worth fixing at some point
+        Object.defineProperty(this, "waitingEdges", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: new Set()
+        });
         this.channels = _getChannels(fields.channels);
+        for (const c of Object.values(this.channels)) {
+            if (c.lc_graph_name === "BinaryOperatorAggregate") {
+                this.supportMultipleEdges = true;
+                break;
+            }
+        }
+    }
+    get allEdges() {
+        return new Set([
+            ...this.edges,
+            ...Array.from(this.waitingEdges).flatMap(([starts, end]) => starts.map((start) => [start, end])),
+        ]);
     }
     addNode(key, action) {
-        if (Object.keys(this.nodes).some((key) => key in this.channels)) {
+        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.`);
         }
-        super.addNode(key, action);
+        return super.addNode(key, action);
     }
-    compile(checkpointer) {
-        this.validate();
-        if (Object.keys(this.nodes).some((key) => key in this.channels)) {
-            throw new Error("Cannot use channel names as node names");
+    addEdge(startKey, endKey) {
+        if (typeof startKey === "string") {
+            return super.addEdge(startKey, endKey);
         }
+        if (this.compiled) {
+            console.warn("Adding an edge to a graph that has already been compiled. This will " +
+                "not be reflected in the compiled graph.");
+        }
+        for (const start of startKey) {
+            if (start === graph_js_1.END) {
+                throw new Error("END cannot be a start node");
+            }
+            if (!Object.keys(this.nodes).some((node) => node === start)) {
+                throw new Error(`Need to add_node ${start} first`);
+            }
+        }
+        if (endKey === graph_js_1.END) {
+            throw new Error("END cannot be an end node");
+        }
+        if (!Object.keys(this.nodes).some((node) => node === endKey)) {
+            throw new Error(`Need to add_node ${endKey} first`);
+        }
+        this.waitingEdges.add([startKey, endKey]);
+        return this;
+    }
+    compile({ checkpointer, interruptBefore, interruptAfter, } = {}) {
+        // validate the graph
+        this.validate([
+            ...(Array.isArray(interruptBefore) ? interruptBefore : []),
+            ...(Array.isArray(interruptAfter) ? interruptAfter : []),
+        ]);
+        // prepare output channels
         const stateKeys = Object.keys(this.channels);
-        const stateKeysRead = stateKeys.length === 1 && stateKeys[0] === "__root__"
+        const outputs = stateKeys.length === 1 && stateKeys[0] === ROOT
             ? stateKeys[0]
             : stateKeys;
-        const updateState = Array.isArray(stateKeysRead)
-            ? (nodeName, 
-            // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            input, options) => _updateStateObject(stateKeys, nodeName, input, options)
-            : _updateStateRoot;
-        const outgoingEdges = {};
+        // create empty compiled graph
+        const compiled = new CompiledStateGraph({
+            builder: this,
+            checkpointer,
+            interruptAfter,
+            interruptBefore,
+            autoValidate: false,
+            nodes: {},
+            channels: {
+                ...this.channels,
+                [graph_js_1.START]: new ephemeral_value_js_1.EphemeralValue(),
+            },
+            inputs: graph_js_1.START,
+            outputs,
+            streamChannels: outputs,
+            streamMode: "updates",
+        });
+        // attach nodes, edges and branches
+        compiled.attachNode(graph_js_1.START);
+        for (const [key, node] of Object.entries(this.nodes)) {
+            compiled.attachNode(key, node);
+        }
         for (const [start, end] of this.edges) {
-            if (!outgoingEdges[start]) {
-                outgoingEdges[start] = [];
-            }
-            outgoingEdges[start].push(end !== graph_js_1.END ? `${end}:inbox` : graph_js_1.END);
+            compiled.attachEdge(start, end);
         }
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        const nodes = {};
-        for (const [key, node] of Object.entries(this.nodes)) {
-            nodes[key] = index_js_1.Channel.subscribeTo(`${key}:inbox`)
-                .pipe(node)
-                .pipe((
-            // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            input, options) => updateState(key, input, options))
-                .pipe(index_js_1.Channel.writeTo(key));
-        }
-        for (const key of Object.keys(this.nodes)) {
-            const outgoing = outgoingEdges[key];
-            const edgesKey = `${key}:edges`;
-            if (outgoing || this.branches[key]) {
-                nodes[edgesKey] = index_js_1.Channel.subscribeTo(key, {
-                    tags: ["langsmith:hidden"],
-                }).pipe(new read_js_1.ChannelRead(stateKeysRead));
-            }
-            if (outgoing) {
-                nodes[edgesKey] = nodes[edgesKey].pipe(index_js_1.Channel.writeTo(...outgoing));
-            }
-            if (this.branches[key]) {
-                for (const branch of this.branches[key]) {
-                    nodes[edgesKey] = nodes[edgesKey].pipe(new runnables_1.RunnableLambda({
-                        func: (i, c) => branch.runnable(i, c),
-                    }));
-                }
+        for (const [starts, end] of this.waitingEdges) {
+            compiled.attachEdge(starts, end);
+        }
+        for (const [start, branches] of Object.entries(this.branches)) {
+            for (const [name, branch] of Object.entries(branches)) {
+                compiled.attachBranch(start, name, branch);
             }
         }
-        nodes[exports.START] = index_js_1.Channel.subscribeTo(`${exports.START}:inbox`, {
-            tags: ["langsmith:hidden"],
-        })
-            .pipe((
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        input, options) => updateState(exports.START, input, options))
-            .pipe(index_js_1.Channel.writeTo(exports.START));
-        nodes[`${exports.START}:edges`] = index_js_1.Channel.subscribeTo(exports.START, {
-            tags: ["langsmith:hidden"],
-        })
-            .pipe(new read_js_1.ChannelRead(stateKeysRead))
-            .pipe(index_js_1.Channel.writeTo(`${this.entryPoint}:inbox`));
-        return new index_js_1.Pregel({
-            nodes,
-            channels: this.channels,
-            input: `${exports.START}:inbox`,
-            output: graph_js_1.END,
-            hidden: Object.keys(this.nodes)
-                .map((node) => `${node}:inbox`)
-                .concat(exports.START, stateKeys),
-            checkpointer,
-        });
+        return compiled.validate();
     }
 }
 exports.StateGraph = StateGraph;
-function _updateStateObject(stateKeys, nodeName, 
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-input, options
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-) {
-    if (!options?.config) {
-        throw new Error("Config not found when updating state.");
-    }
-    if (Object.keys(input).some((key) => !stateKeys.some((sk) => sk === key))) {
-        throw new Error(`Invalid state update from node ${nodeName}, expected object with one or more of ${stateKeys.join(", ")}, got ${Object.keys(input).join(",")}`);
+function _getChannels(schema) {
+    const channels = {};
+    for (const [name, val] of Object.entries(schema)) {
+        if (name === ROOT) {
+            channels[name] = getChannel(val);
+        }
+        else {
+            const key = name;
+            channels[name] = getChannel(val);
+        }
     }
-    write_js_1.ChannelWrite.doWrite(options.config, input);
-    return input;
+    return channels;
 }
-function _updateStateRoot(_nodeName, input, options) {
-    if (!options?.config) {
-        throw new Error("Config not found when updating state.");
+function getChannel(reducer) {
+    if (reducer && "reducer" in reducer && reducer.reducer) {
+        return new binop_js_1.BinaryOperatorAggregate(reducer.reducer, reducer.default);
+    }
+    if (reducer && "value" in reducer && reducer.value) {
+        return new binop_js_1.BinaryOperatorAggregate(reducer.value, reducer.default);
     }
-    write_js_1.ChannelWrite.doWrite(options.config, {
-        __root__: input,
-    });
-    return input;
+    return new last_value_js_1.LastValue();
 }
-function _getChannels(schema) {
-    if ("value" in schema && "default" in schema) {
-        if (!schema.value) {
-            throw new Error("Value is required for channels");
+class CompiledStateGraph extends graph_js_1.CompiledGraph {
+    attachNode(key, node) {
+        const stateKeys = Object.keys(this.builder.channels);
+        function getStateKey(key, input) {
+            if (!input) {
+                return write_js_1.SKIP_WRITE;
+            }
+            else if (typeof input !== "object" || Array.isArray(input)) {
+                throw new errors_js_1.InvalidUpdateError(`Expected dict, got ${typeof input}`);
+            }
+            else {
+                return key in input ? input[key] : write_js_1.SKIP_WRITE;
+            }
+        }
+        // state updaters
+        const stateWriteEntries = stateKeys.map((key) => key === ROOT
+            ? { channel: key, value: write_js_1.PASSTHROUGH, skipNone: true }
+            : {
+                channel: key,
+                value: write_js_1.PASSTHROUGH,
+                mapper: new utils_js_1.RunnableCallable({
+                    func: getStateKey.bind(null, key),
+                    trace: false,
+                    recurse: false,
+                }),
+            });
+        // add node and output channel
+        if (key === graph_js_1.START) {
+            this.nodes[key] = new read_js_1.PregelNode({
+                tags: [constants_js_1.TAG_HIDDEN],
+                triggers: [graph_js_1.START],
+                channels: [graph_js_1.START],
+                writers: [new write_js_1.ChannelWrite(stateWriteEntries, [constants_js_1.TAG_HIDDEN])],
+            });
+        }
+        else {
+            this.channels[key] = new ephemeral_value_js_1.EphemeralValue();
+            this.nodes[key] = new read_js_1.PregelNode({
+                triggers: [],
+                // read state keys
+                channels: stateKeys.length === 1 && stateKeys[0] === ROOT
+                    ? stateKeys
+                    : stateKeys.reduce((acc, k) => {
+                        acc[k] = k;
+                        return acc;
+                    }, {}),
+                // publish to this channel and state keys
+                writers: [
+                    new write_js_1.ChannelWrite(stateWriteEntries.concat({ channel: key, value: key }), [constants_js_1.TAG_HIDDEN]),
+                ],
+                bound: node,
+            });
         }
-        return {
-            __root__: new binop_js_1.BinaryOperatorAggregate(schema.value, schema.default),
-        };
     }
-    const channels = {};
-    for (const [name, values] of Object.entries(schema)) {
-        if (values.value) {
-            channels[name] = new binop_js_1.BinaryOperatorAggregate(values.value, values.default);
+    attachEdge(start, end) {
+        if (end === graph_js_1.END) {
+            return;
+        }
+        if (Array.isArray(start)) {
+            const channelName = `join:${start.join("+")}:${end}`;
+            // register channel
+            this.channels[channelName] =
+                new named_barrier_value_js_1.NamedBarrierValue(new Set(start));
+            // subscribe to channel
+            this.nodes[end].triggers.push(channelName);
+            // publish to channel
+            for (const s of start) {
+                this.nodes[s].writers.push(new write_js_1.ChannelWrite([{ channel: channelName, value: s }], [constants_js_1.TAG_HIDDEN]));
+            }
+        }
+        else if (start === graph_js_1.START) {
+            const channelName = `start:${end}`;
+            // register channel
+            this.channels[channelName] =
+                new ephemeral_value_js_1.EphemeralValue();
+            // subscribe to channel
+            this.nodes[end].triggers.push(channelName);
+            // publish to channel
+            this.nodes[graph_js_1.START].writers.push(new write_js_1.ChannelWrite([{ channel: channelName, value: graph_js_1.START }], [constants_js_1.TAG_HIDDEN]));
         }
         else {
-            channels[name] = new last_value_js_1.LastValue();
+            this.nodes[end].triggers.push(start);
+        }
+    }
+    attachBranch(start, name, branch) {
+        // attach branch publisher
+        this.nodes[start].writers.push(branch.compile(
+        // writer
+        (dests) => {
+            const filteredDests = dests.filter((dest) => dest !== graph_js_1.END);
+            if (!filteredDests.length) {
+                return;
+            }
+            const writes = filteredDests.map((dest) => ({
+                channel: `branch:${start}:${name}:${dest}`,
+                value: start,
+            }));
+            if (branch.then && branch.then !== graph_js_1.END) {
+                writes.push({
+                    channel: `branch:${start}:${name}:then`,
+                    value: { __names: filteredDests },
+                });
+            }
+            return new write_js_1.ChannelWrite(writes, [constants_js_1.TAG_HIDDEN]);
+        }, 
+        // reader
+        (config) => read_js_1.ChannelRead.doRead(config, this.outputs, true)));
+        // attach branch subscribers
+        const ends = branch.ends
+            ? Object.values(branch.ends)
+            : Object.keys(this.builder.nodes).filter((n) => n !== branch.then);
+        for (const end of ends) {
+            if (end === graph_js_1.END) {
+                continue;
+            }
+            const channelName = `branch:${start}:${name}:${end}`;
+            this.channels[channelName] =
+                new ephemeral_value_js_1.EphemeralValue();
+            this.nodes[end].triggers.push(channelName);
+        }
+        if (branch.then && branch.then !== graph_js_1.END) {
+            const channelName = `branch:${start}:${name}:then`;
+            this.channels[channelName] =
+                new dynamic_barrier_value_js_1.DynamicBarrierValue();
+            this.nodes[branch.then].triggers.push(channelName);
+            for (const end of ends) {
+                if (end === graph_js_1.END) {
+                    continue;
+                }
+                this.nodes[end].writers.push(new write_js_1.ChannelWrite([{ channel: channelName, value: end }], [constants_js_1.TAG_HIDDEN]));
+            }
         }
     }
-    return channels;
 }
+exports.CompiledStateGraph = CompiledStateGraph;

package/dist/graph/state.d.ts

@@ -1,24 +1,47 @@
-import { RunnableLike } from "@langchain/core/runnables";
+import { Runnable, RunnableConfig, RunnableLike } from "@langchain/core/runnables";
 import { BaseChannel } from "../channels/base.js";
 import { BinaryOperator } from "../channels/binop.js";
-import { Graph } from "./graph.js";
+import { END, CompiledGraph, Graph, START, Branch } from "./graph.js";
 import { BaseCheckpointSaver } from "../checkpoint/base.js";
-import { Pregel } from "../pregel/index.js";
-export declare const START = "__start__";
-export interface StateGraphArgs<Channels extends Record<string, any>> {
-    channels: {
-        [K in keyof Channels]: {
-            value: BinaryOperator<Channels[K]> | null;
-            default?: () => Channels[K];
-        } | string;
-    } | {
-        value: BinaryOperator<unknown> | null;
-        default?: () => unknown;
-    };
+import { All } from "../pregel/types.js";
+type SingleReducer<T> = {
+    reducer: BinaryOperator<T>;
+    default?: () => T;
+} | {
+    /**
+     * @deprecated Use `reducer` instead
+     */
+    value: BinaryOperator<T>;
+    default?: () => T;
+} | null;
+export type ChannelReducers<Channels extends object> = {
+    [K in keyof Channels]: SingleReducer<Channels[K]>;
+};
+export interface StateGraphArgs<Channels extends object | unknown> {
+    channels: Channels extends object ? Channels extends unknown[] ? ChannelReducers<{
+        __root__: Channels;
+    }> : ChannelReducers<Channels> : ChannelReducers<{
+        __root__: Channels;
+    }>;
 }
-export declare class StateGraph<Channels extends Record<string, any>> extends Graph<Channels> {
+export declare class StateGraph<const State extends object | unknown, const Update extends object | unknown = Partial<State>, const N extends string = typeof START> extends Graph<N, State, Update> {
     channels: Record<string, BaseChannel>;
-    constructor(fields: StateGraphArgs<Channels>);
-    addNode(key: string, action: RunnableLike): void;
-    compile(checkpointer?: BaseCheckpointSaver): Pregel;
+    waitingEdges: Set<[N[], N]>;
+    constructor(fields: StateGraphArgs<State>);
+    get allEdges(): Set<[string, string]>;
+    addNode<K extends string>(key: K, action: RunnableLike<State, Update>): StateGraph<State, Update, N | K>;
+    addEdge(startKey: typeof START | N | N[], endKey: N | typeof END): this;
+    compile({ checkpointer, interruptBefore, interruptAfter, }?: {
+        checkpointer?: BaseCheckpointSaver;
+        interruptBefore?: N[] | All;
+        interruptAfter?: N[] | All;
+    }): CompiledStateGraph<State, Update, N>;
 }
+export declare class CompiledStateGraph<const State extends object | unknown, const Update extends object | unknown = Partial<State>, const N extends string = typeof START> extends CompiledGraph<N, State, Update> {
+    builder: StateGraph<State, Update, N>;
+    attachNode(key: typeof START, node?: never): void;
+    attachNode(key: N, node: Runnable<State, Update, RunnableConfig>): void;
+    attachEdge(start: N | N[] | "__start__", end: N | "__end__"): void;
+    attachBranch(start: N | typeof START, name: string, branch: Branch<State, N>): void;
+}
+export {};