@langchain/langgraph 0.0.11 → 0.0.13

package/dist/tests/pregel.test.js

@@ -0,0 +1,1879 @@
+/* eslint-disable no-process-env */
+import { it, expect, jest, beforeAll, describe } from "@jest/globals";
+import { RunnableLambda, RunnablePassthrough, } from "@langchain/core/runnables";
+import { PromptTemplate } from "@langchain/core/prompts";
+import { FakeStreamingLLM } from "@langchain/core/utils/testing";
+import { Tool } from "@langchain/core/tools";
+import { z } from "zod";
+import { AIMessage, FunctionMessage, HumanMessage, } from "@langchain/core/messages";
+import { FakeChatModel, MemorySaverAssertImmutable } from "./utils.js";
+import { LastValue } from "../channels/last_value.js";
+import { END, Graph, START, StateGraph } from "../graph/index.js";
+import { Topic } from "../channels/topic.js";
+import { PregelNode } from "../pregel/read.js";
+import { MemorySaver } from "../checkpoint/memory.js";
+import { BinaryOperatorAggregate } from "../channels/binop.js";
+import { Channel, Pregel, _applyWrites, _localRead, _prepareNextTasks, _shouldInterrupt, } from "../pregel/index.js";
+import { ToolExecutor, createAgentExecutor } from "../prebuilt/index.js";
+import { MessageGraph } from "../graph/message.js";
+import { PASSTHROUGH } from "../pregel/write.js";
+import { GraphRecursionError, InvalidUpdateError } from "../errors.js";
+import { SqliteSaver } from "../checkpoint/sqlite.js";
+import { uuid6 } from "../checkpoint/id.js";
+// Tracing slows down the tests
+beforeAll(() => {
+    process.env.LANGCHAIN_TRACING_V2 = "false";
+    process.env.LANGCHAIN_ENDPOINT = "";
+    process.env.LANGCHAIN_ENDPOINT = "";
+    process.env.LANGCHAIN_API_KEY = "";
+    process.env.LANGCHAIN_PROJECT = "";
+});
+describe("Channel", () => {
+    describe("writeTo", () => {
+        it("should return a ChannelWrite instance with the expected writes", () => {
+            // call method / assertions
+            const channelWrite = Channel.writeTo(["foo", "bar"], {
+                fixed: 6,
+                func: () => 42,
+                runnable: new RunnablePassthrough(),
+            });
+            expect(channelWrite.writes.length).toBe(5);
+            expect(channelWrite.writes[0]).toEqual({
+                channel: "foo",
+                value: PASSTHROUGH,
+                skipNone: false,
+            });
+            expect(channelWrite.writes[1]).toEqual({
+                channel: "bar",
+                value: PASSTHROUGH,
+                skipNone: false,
+            });
+            expect(channelWrite.writes[2]).toEqual({
+                channel: "fixed",
+                value: 6,
+                skipNone: false,
+            });
+            // TODO: Figure out how to assert the mapper value
+            // expect(channelWrite.writes[3]).toEqual({
+            //   channel: "func",
+            //   value: PASSTHROUGH,
+            //   skipNone: true,
+            //   mapper: new RunnableLambda({ func: () => 42}),
+            // });
+            expect(channelWrite.writes[4]).toEqual({
+                channel: "runnable",
+                value: PASSTHROUGH,
+                skipNone: true,
+                mapper: new RunnablePassthrough(),
+            });
+        });
+    });
+});
+describe("Pregel", () => {
+    describe("streamChannelsList", () => {
+        it("should return the expected list of stream channels", () => {
+            // set up test
+            const chain = Channel.subscribeTo("input").pipe(Channel.writeTo(["output"]));
+            const pregel1 = new Pregel({
+                nodes: { one: chain },
+                channels: {
+                    input: new LastValue(),
+                    output: new LastValue(),
+                },
+                inputs: "input",
+                outputs: "output",
+                streamChannels: "output",
+            });
+            const pregel2 = new Pregel({
+                nodes: { one: chain },
+                channels: {
+                    input: new LastValue(),
+                    output: new LastValue(),
+                },
+                inputs: "input",
+                outputs: "output",
+                streamChannels: ["input", "output"],
+            });
+            const pregel3 = new Pregel({
+                nodes: { one: chain },
+                channels: {
+                    input: new LastValue(),
+                    output: new LastValue(),
+                },
+                inputs: "input",
+                outputs: "output",
+            });
+            // call method / assertions
+            expect(pregel1.streamChannelsList).toEqual(["output"]);
+            expect(pregel2.streamChannelsList).toEqual(["input", "output"]);
+            expect(pregel3.streamChannelsList).toEqual(["input", "output"]);
+            expect(pregel1.streamChannelsAsIs).toEqual("output");
+            expect(pregel2.streamChannelsAsIs).toEqual(["input", "output"]);
+            expect(pregel3.streamChannelsAsIs).toEqual(["input", "output"]);
+        });
+    });
+    describe("_defaults", () => {
+        it("should return the expected tuple of defaults", () => {
+            // Because the implementation of _defaults() contains independent
+            // if-else statements that determine that returned values in the tuple,
+            // this unit test can be separated into 2 parts. The first part of the
+            // test executes the "true" evaluation path of the if-else statements.
+            // The second part evaluates the "false" evaluation path.
+            // set up test
+            const channels = {
+                inputKey: new LastValue(),
+                outputKey: new LastValue(),
+                channel3: new LastValue(),
+            };
+            const nodes = {
+                one: new PregelNode({
+                    channels: ["channel3"],
+                    triggers: ["outputKey"],
+                }),
+            };
+            const config1 = {};
+            const config2 = {
+                streamMode: "updates",
+                inputKeys: "inputKey",
+                outputKeys: "outputKey",
+                interruptBefore: "*",
+                interruptAfter: ["one"],
+                debug: true,
+                tags: ["hello"],
+            };
+            // create Pregel class
+            const pregel = new Pregel({
+                nodes,
+                debug: false,
+                inputs: "outputKey",
+                outputs: "outputKey",
+                interruptBefore: ["one"],
+                interruptAfter: ["one"],
+                streamMode: "values",
+                channels,
+                checkpointer: new MemorySaver(),
+            });
+            // call method / assertions
+            const expectedDefaults1 = [
+                false, // debug
+                "values", // stream mode
+                "outputKey", // input keys
+                ["inputKey", "outputKey", "channel3"], // output keys,
+                {},
+                ["one"], // interrupt before
+                ["one"], // interrupt after
+            ];
+            const expectedDefaults2 = [
+                true, // debug
+                "updates", // stream mode
+                "inputKey", // input keys
+                "outputKey", // output keys
+                { tags: ["hello"] },
+                "*", // interrupt before
+                ["one"], // interrupt after
+            ];
+            expect(pregel._defaults(config1)).toEqual(expectedDefaults1);
+            expect(pregel._defaults(config2)).toEqual(expectedDefaults2);
+        });
+    });
+});
+describe("_shouldInterrupt", () => {
+    it("should return true if any snapshot channel has been updated since last interrupt and any channel written to is in interrupt nodes list", () => {
+        // set up test
+        const checkpoint = {
+            v: 1,
+            id: uuid6(-1),
+            ts: "2024-04-19T17:19:07.952Z",
+            channel_values: {
+                channel1: "channel1value",
+            },
+            channel_versions: {
+                channel1: 2, // current channel version is greater than last version seen
+            },
+            versions_seen: {
+                __interrupt__: {
+                    channel1: 1,
+                },
+            },
+        };
+        const interruptNodes = ["node1"];
+        const snapshotChannels = ["channel1"];
+        // call method / assertions
+        expect(_shouldInterrupt(checkpoint, interruptNodes, snapshotChannels, [
+            {
+                name: "node1",
+                input: undefined,
+                proc: new RunnablePassthrough(),
+                writes: [],
+                config: undefined,
+            },
+        ])).toBe(true);
+    });
+    it("should return false if all snapshot channels have not been updated", () => {
+        // set up test
+        const checkpoint = {
+            v: 1,
+            id: uuid6(-1),
+            ts: "2024-04-19T17:19:07.952Z",
+            channel_values: {
+                channel1: "channel1value",
+            },
+            channel_versions: {
+                channel1: 2, // current channel version is equal to last version seen
+            },
+            versions_seen: {
+                __interrupt__: {
+                    channel1: 2,
+                },
+            },
+        };
+        const interruptNodes = ["node1"];
+        const snapshotChannels = ["channel1"];
+        // call method / assertions
+        expect(_shouldInterrupt(checkpoint, interruptNodes, snapshotChannels, [
+            {
+                name: "node1",
+                input: undefined,
+                proc: new RunnablePassthrough(),
+                writes: [],
+                config: undefined,
+            },
+        ])).toBe(false);
+    });
+    it("should return false if all task nodes are not in interrupt nodes", () => {
+        // set up test
+        const checkpoint = {
+            v: 1,
+            id: uuid6(-1),
+            ts: "2024-04-19T17:19:07.952Z",
+            channel_values: {
+                channel1: "channel1value",
+            },
+            channel_versions: {
+                channel1: 2,
+            },
+            versions_seen: {
+                __interrupt__: {
+                    channel1: 1,
+                },
+            },
+        };
+        const interruptNodes = ["node1"];
+        const snapshotChannels = ["channel1"];
+        // call method / assertions
+        expect(_shouldInterrupt(checkpoint, interruptNodes, snapshotChannels, [
+            {
+                name: "node2", // node2 is not in interrupt nodes
+                input: undefined,
+                proc: new RunnablePassthrough(),
+                writes: [],
+                config: undefined,
+            },
+        ])).toBe(false);
+    });
+});
+describe("_localRead", () => {
+    it("should return the channel value when fresh is false", () => {
+        // set up test
+        const checkpoint = {
+            v: 0,
+            id: uuid6(-1),
+            ts: "",
+            channel_values: {},
+            channel_versions: {},
+            versions_seen: {},
+        };
+        const channel1 = new LastValue();
+        const channel2 = new LastValue();
+        channel1.update([1]);
+        channel2.update([2]);
+        const channels = {
+            channel1,
+            channel2,
+        };
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const writes = [];
+        // call method / assertions
+        expect(_localRead(checkpoint, channels, writes, "channel1", false)).toBe(1);
+        expect(_localRead(checkpoint, channels, writes, ["channel1", "channel2"], false)).toEqual({ channel1: 1, channel2: 2 });
+    });
+    it("should return the channel value after applying writes when fresh is true", () => {
+        // set up test
+        const checkpoint = {
+            v: 0,
+            id: uuid6(-1),
+            ts: "",
+            channel_values: {},
+            channel_versions: {},
+            versions_seen: {},
+        };
+        const channel1 = new LastValue();
+        const channel2 = new LastValue();
+        channel1.update([1]);
+        channel2.update([2]);
+        const channels = {
+            channel1,
+            channel2,
+        };
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const writes = [
+            ["channel1", 100],
+            ["channel2", 200],
+        ];
+        // call method / assertions
+        expect(_localRead(checkpoint, channels, writes, "channel1", true)).toBe(100);
+        expect(_localRead(checkpoint, channels, writes, ["channel1", "channel2"], true)).toEqual({ channel1: 100, channel2: 200 });
+    });
+});
+describe("_applyWrites", () => {
+    it("should update channels and checkpoints correctly (side effect)", () => {
+        // set up test
+        const checkpoint = {
+            v: 1,
+            id: uuid6(-1),
+            ts: "2024-04-19T17:19:07.952Z",
+            channel_values: {
+                channel1: "channel1value",
+            },
+            channel_versions: {
+                channel1: 2,
+                channel2: 5,
+            },
+            versions_seen: {
+                __interrupt__: {
+                    channel1: 1,
+                },
+            },
+        };
+        const lastValueChannel1 = new LastValue();
+        lastValueChannel1.update(["channel1value"]);
+        const lastValueChannel2 = new LastValue();
+        lastValueChannel2.update(["channel2value"]);
+        const channels = {
+            channel1: lastValueChannel1,
+            channel2: lastValueChannel2,
+        };
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const pendingWrites = [
+            ["channel1", "channel1valueUpdated!"],
+        ];
+        // call method / assertions
+        expect(channels.channel1.get()).toBe("channel1value");
+        expect(channels.channel2.get()).toBe("channel2value");
+        expect(checkpoint.channel_versions.channel1).toBe(2);
+        _applyWrites(checkpoint, channels, pendingWrites); // contains side effects
+        expect(channels.channel1.get()).toBe("channel1valueUpdated!");
+        expect(channels.channel2.get()).toBe("channel2value");
+        expect(checkpoint.channel_versions.channel1).toBe(6);
+    });
+    it("should throw an InvalidUpdateError if there are multiple updates to the same channel", () => {
+        // set up test
+        const checkpoint = {
+            v: 1,
+            id: uuid6(-1),
+            ts: "2024-04-19T17:19:07.952Z",
+            channel_values: {
+                channel1: "channel1value",
+            },
+            channel_versions: {
+                channel1: 2,
+            },
+            versions_seen: {
+                __interrupt__: {
+                    channel1: 1,
+                },
+            },
+        };
+        const lastValueChannel1 = new LastValue();
+        lastValueChannel1.update(["channel1value"]);
+        const channels = {
+            channel1: lastValueChannel1,
+        };
+        // LastValue channel can only be updated with one value at a time
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const pendingWrites = [
+            ["channel1", "channel1valueUpdated!"],
+            ["channel1", "channel1valueUpdatedAgain!"],
+        ];
+        // call method / assertions
+        expect(() => {
+            _applyWrites(checkpoint, channels, pendingWrites); // contains side effects
+        }).toThrow(InvalidUpdateError);
+    });
+});
+describe("_prepareNextTasks", () => {
+    it("should return an array of PregelTaskDescriptions", () => {
+        // set up test
+        const checkpoint = {
+            v: 1,
+            id: "123",
+            ts: "2024-04-19T17:19:07.952Z",
+            channel_values: {
+                channel1: 1,
+                channel2: 2,
+            },
+            channel_versions: {
+                channel1: 2,
+                channel2: 5,
+            },
+            versions_seen: {
+                node1: {
+                    channel1: 1,
+                },
+                node2: {
+                    channel2: 5,
+                },
+            },
+        };
+        const processes = {
+            node1: new PregelNode({
+                channels: ["channel1"],
+                triggers: ["channel1"],
+            }),
+            node2: new PregelNode({
+                channels: ["channel2"],
+                triggers: ["channel1", "channel2"],
+                mapper: () => 100, // return 100 no matter what
+            }),
+        };
+        const channel1 = new LastValue();
+        channel1.update([1]);
+        const channel2 = new LastValue();
+        channel2.update([2]);
+        const channels = {
+            channel1,
+            channel2,
+        };
+        // call method / assertions
+        const [newCheckpoint, taskDescriptions] = _prepareNextTasks(checkpoint, processes, channels, false);
+        expect(taskDescriptions.length).toBe(2);
+        expect(taskDescriptions[0]).toEqual({ name: "node1", input: 1 });
+        expect(taskDescriptions[1]).toEqual({ name: "node2", input: 100 });
+        // the returned checkpoint is a copy of the passed checkpoint without versionsSeen updated
+        expect(newCheckpoint.versions_seen.node1.channel1).toBe(1);
+        expect(newCheckpoint.versions_seen.node2.channel2).toBe(5);
+    });
+    it("should return an array of PregelExecutableTasks", () => {
+        const checkpoint = {
+            v: 1,
+            id: uuid6(-1),
+            ts: "2024-04-19T17:19:07.952Z",
+            channel_values: {
+                channel1: 1,
+                channel2: 2,
+            },
+            channel_versions: {
+                channel1: 2,
+                channel2: 5,
+                channel3: 4,
+                channel4: 4,
+                channel6: 4,
+            },
+            versions_seen: {
+                node1: {
+                    channel1: 1,
+                },
+                node2: {
+                    channel2: 5,
+                },
+                node3: {
+                    channel3: 4,
+                },
+                node4: {
+                    channel4: 3,
+                },
+                node6: {
+                    channel6: 3,
+                },
+            },
+        };
+        const processes = {
+            node1: new PregelNode({
+                channels: ["channel1"],
+                triggers: ["channel1"],
+                writers: [new RunnablePassthrough()],
+            }),
+            node2: new PregelNode({
+                channels: ["channel2"],
+                triggers: ["channel1", "channel2"],
+                writers: [new RunnablePassthrough()],
+                mapper: () => 100, // return 100 no matter what
+            }),
+            node3: new PregelNode({
+                // this task is filtered out because current version of channel3 matches version seen
+                channels: ["channel3"],
+                triggers: ["channel3"],
+            }),
+            node4: new PregelNode({
+                // this task is filtered out because channel5 is empty
+                channels: ["channel5"],
+                triggers: ["channel4"],
+            }),
+            node6: new PregelNode({
+                // this task is filtered out because channel5 is empty
+                channels: { channel5: "channel5" },
+                triggers: ["channel5", "channel6"],
+            }),
+        };
+        const channel1 = new LastValue();
+        channel1.update([1]);
+        const channel2 = new LastValue();
+        channel2.update([2]);
+        const channel3 = new LastValue();
+        channel3.update([3]);
+        const channel4 = new LastValue();
+        channel4.update([4]);
+        const channel5 = new LastValue();
+        const channel6 = new LastValue();
+        channel6.update([6]);
+        const channels = {
+            channel1,
+            channel2,
+            channel3,
+            channel4,
+            channel5,
+            channel6,
+        };
+        // call method / assertions
+        const [newCheckpoint, tasks] = _prepareNextTasks(checkpoint, processes, channels, true);
+        expect(tasks.length).toBe(2);
+        expect(tasks[0]).toEqual({
+            name: "node1",
+            input: 1,
+            proc: new RunnablePassthrough(),
+            writes: [],
+            config: { tags: [] },
+        });
+        expect(tasks[1]).toEqual({
+            name: "node2",
+            input: 100,
+            proc: new RunnablePassthrough(),
+            writes: [],
+            config: { tags: [] },
+        });
+        expect(newCheckpoint.versions_seen.node1.channel1).toBe(2);
+        expect(newCheckpoint.versions_seen.node2.channel1).toBe(2);
+        expect(newCheckpoint.versions_seen.node2.channel2).toBe(5);
+    });
+});
+it("can invoke pregel with a single process", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const chain = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: {
+            one: chain,
+        },
+        channels: {
+            input: new LastValue(),
+            output: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+    });
+    expect(await app.invoke(2)).toBe(3);
+    expect(await app.invoke(2, { outputKeys: ["output"] })).toEqual({
+        output: 3,
+    });
+    expect(() => app.toString()).not.toThrow();
+    // Verify the mock was called correctly
+    expect(addOne).toHaveBeenCalled();
+});
+it("can invoke graph with a single process", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const graph = new Graph()
+        .addNode("add_one", addOne)
+        .addEdge(START, "add_one")
+        .addEdge("add_one", END)
+        .compile();
+    expect(await graph.invoke(2)).toBe(3);
+});
+it("should process input and produce output with implicit channels", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const chain = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: { one: chain },
+        channels: {
+            input: new LastValue(),
+            output: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+    });
+    expect(await app.invoke(2)).toBe(3);
+    // Verify the mock was called correctly
+    expect(addOne).toHaveBeenCalled();
+});
+it("should process input and write kwargs correctly", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const chain = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"], {
+        fixed: 5,
+        outputPlusOne: (x) => x + 1,
+    }));
+    const app = new Pregel({
+        nodes: { one: chain },
+        channels: {
+            input: new LastValue(),
+            output: new LastValue(),
+            fixed: new LastValue(),
+            outputPlusOne: new LastValue(),
+        },
+        outputs: ["output", "fixed", "outputPlusOne"],
+        inputs: "input",
+    });
+    expect(await app.invoke(2)).toEqual({
+        output: 3,
+        fixed: 5,
+        outputPlusOne: 4,
+    });
+});
+it("should invoke single process in out objects", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const chain = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: {
+            one: chain,
+        },
+        channels: {
+            input: new LastValue(),
+            output: new LastValue(),
+        },
+        inputs: "input",
+        outputs: ["output"],
+    });
+    expect(await app.invoke(2)).toEqual({ output: 3 });
+});
+it("should process input and output as objects", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const chain = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: { one: chain },
+        channels: {
+            input: new LastValue(),
+            output: new LastValue(),
+        },
+        inputs: ["input"],
+        outputs: ["output"],
+    });
+    expect(await app.invoke({ input: 2 })).toEqual({ output: 3 });
+});
+it("should invoke two processes and get correct output", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const one = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["inbox"]));
+    const two = Channel.subscribeTo("inbox")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: { one, two },
+        channels: {
+            inbox: new LastValue(),
+            output: new LastValue(),
+            input: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+        streamChannels: ["inbox", "output"],
+    });
+    await expect(app.invoke(2, { recursionLimit: 1 })).rejects.toThrow(GraphRecursionError);
+    expect(await app.invoke(2)).toEqual(4);
+    const stream = await app.stream(2, { streamMode: "updates" });
+    let step = 0;
+    for await (const value of stream) {
+        if (step === 0) {
+            expect(value).toEqual({ one: { inbox: 3 } });
+        }
+        else if (step === 1) {
+            expect(value).toEqual({ two: { output: 4 } });
+        }
+        step += 1;
+    }
+    expect(step).toBe(2);
+});
+it("should process two processes with object input and output", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const one = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["inbox"]));
+    const two = Channel.subscribeTo("inbox")
+        .pipe(new RunnableLambda({ func: addOne }).map())
+        .pipe(Channel.writeTo(["output"]).map());
+    const app = new Pregel({
+        nodes: { one, two },
+        channels: {
+            inbox: new Topic(),
+            input: new LastValue(),
+            output: new LastValue(),
+        },
+        streamChannels: ["output", "inbox"],
+        inputs: ["input", "inbox"],
+        outputs: "output",
+    });
+    const streamResult = await app.stream({ input: 2, inbox: 12 }, { outputKeys: "output" });
+    const outputResults = [];
+    for await (const result of streamResult) {
+        outputResults.push(result);
+    }
+    expect(outputResults).toEqual([13, 4]); // [12 + 1, 2 + 1 + 1]
+    const fullStreamResult = await app.stream({ input: 2, inbox: 12 });
+    const fullOutputResults = [];
+    for await (const result of fullStreamResult) {
+        fullOutputResults.push(result);
+    }
+    expect(fullOutputResults).toEqual([
+        { inbox: [3], output: 13 },
+        { inbox: [], output: 4 },
+    ]);
+    const fullOutputResultsUpdates = [];
+    for await (const result of await app.stream({ input: 2, inbox: 12 }, { streamMode: "updates" })) {
+        fullOutputResultsUpdates.push(result);
+    }
+    expect(fullOutputResultsUpdates).toEqual([
+        {
+            one: {
+                inbox: 3,
+            },
+            two: {
+                output: 13,
+            },
+        },
+        { two: { output: 4 } },
+    ]);
+});
+it("should process batch with two processes and delays", async () => {
+    const addOneWithDelay = jest.fn((inp) => new Promise((resolve) => {
+        setTimeout(() => resolve(inp + 1), inp * 100);
+    }));
+    const one = Channel.subscribeTo("input")
+        .pipe(addOneWithDelay)
+        .pipe(Channel.writeTo(["one"]));
+    const two = Channel.subscribeTo("one")
+        .pipe(addOneWithDelay)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: { one, two },
+        channels: {
+            one: new LastValue(),
+            output: new LastValue(),
+            input: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+    });
+    expect(await app.batch([3, 2, 1, 3, 5])).toEqual([5, 4, 3, 5, 7]);
+    expect(await app.batch([3, 2, 1, 3, 5], { outputKeys: ["output"] })).toEqual([
+        { output: 5 },
+        { output: 4 },
+        { output: 3 },
+        { output: 5 },
+        { output: 7 },
+    ]);
+});
+it("should process batch with two processes and delays with graph", async () => {
+    const addOneWithDelay = jest.fn((inp) => new Promise((resolve) => {
+        setTimeout(() => resolve(inp + 1), inp * 100);
+    }));
+    const graph = new Graph()
+        .addNode("add_one", addOneWithDelay)
+        .addNode("add_one_more", addOneWithDelay)
+        .addEdge(START, "add_one")
+        .addEdge("add_one", "add_one_more")
+        .addEdge("add_one_more", END)
+        .compile();
+    expect(await graph.batch([3, 2, 1, 3, 5])).toEqual([5, 4, 3, 5, 7]);
+});
+it("should batch many processes with input and output", async () => {
+    const testSize = 100;
+    const addOne = jest.fn((x) => x + 1);
+    const channels = {
+        input: new LastValue(),
+        output: new LastValue(),
+        "-1": new LastValue(),
+    };
+    const nodes = {
+        "-1": Channel.subscribeTo("input")
+            .pipe(addOne)
+            .pipe(Channel.writeTo(["-1"])),
+    };
+    for (let i = 0; i < testSize - 2; i += 1) {
+        channels[String(i)] = new LastValue();
+        nodes[String(i)] = Channel.subscribeTo(String(i - 1))
+            .pipe(addOne)
+            .pipe(Channel.writeTo([String(i)]));
+    }
+    nodes.last = Channel.subscribeTo(String(testSize - 3))
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes,
+        channels,
+        inputs: "input",
+        outputs: "output",
+    });
+    for (let i = 0; i < 3; i += 1) {
+        await expect(app.batch([2, 1, 3, 4, 5], { recursionLimit: testSize })).resolves.toEqual([
+            2 + testSize,
+            1 + testSize,
+            3 + testSize,
+            4 + testSize,
+            5 + testSize,
+        ]);
+    }
+});
+it("should raise InvalidUpdateError when the same LastValue channel is updated twice in one iteration", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const one = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const two = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: { one, two },
+        channels: {
+            output: new LastValue(),
+            input: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+    });
+    await expect(app.invoke(2)).rejects.toThrow(InvalidUpdateError);
+});
+it("should process two inputs to two outputs validly", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const one = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const two = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: { one, two },
+        channels: {
+            output: new Topic(),
+            input: new LastValue(),
+            output2: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+    });
+    // An Inbox channel accumulates updates into a sequence
+    expect(await app.invoke(2)).toEqual([3, 3]);
+});
+it("should handle checkpoints correctly", async () => {
+    const inputPlusTotal = jest.fn((x) => x.total + x.input);
+    const raiseIfAbove10 = (input) => {
+        if (input > 10) {
+            throw new Error("Input is too large");
+        }
+        return input;
+    };
+    const one = Channel.subscribeTo(["input"])
+        .join(["total"])
+        .pipe(inputPlusTotal)
+        .pipe(Channel.writeTo(["output", "total"]))
+        .pipe(raiseIfAbove10);
+    const memory = new MemorySaverAssertImmutable();
+    const app = new Pregel({
+        nodes: { one },
+        channels: {
+            total: new BinaryOperatorAggregate((a, b) => a + b),
+            input: new LastValue(),
+            output: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+        checkpointer: memory,
+    });
+    // total starts out as 0, so output is 0+2=2
+    await expect(app.invoke(2, { configurable: { thread_id: "1" } })).resolves.toBe(2);
+    let checkpoint = await memory.get({ configurable: { thread_id: "1" } });
+    expect(checkpoint).not.toBeNull();
+    expect(checkpoint?.channel_values.total).toBe(2);
+    // total is now 2, so output is 2+3=5
+    await expect(app.invoke(3, { configurable: { thread_id: "1" } })).resolves.toBe(5);
+    checkpoint = await memory.get({ configurable: { thread_id: "1" } });
+    expect(checkpoint).not.toBeNull();
+    expect(checkpoint?.channel_values.total).toBe(7);
+    // total is now 2+5=7, so output would be 7+4=11, but raises Error
+    await expect(app.invoke(4, { configurable: { thread_id: "1" } })).rejects.toThrow("Input is too large");
+    // checkpoint is not updated
+    checkpoint = await memory.get({ configurable: { thread_id: "1" } });
+    expect(checkpoint).not.toBeNull();
+    expect(checkpoint?.channel_values.total).toBe(7);
+    // on a new thread, total starts out as 0, so output is 0+5=5
+    await expect(app.invoke(5, { configurable: { thread_id: "2" } })).resolves.toBe(5);
+    checkpoint = await memory.get({ configurable: { thread_id: "1" } });
+    expect(checkpoint).not.toBeNull();
+    expect(checkpoint?.channel_values.total).toBe(7);
+    checkpoint = await memory.get({ configurable: { thread_id: "2" } });
+    expect(checkpoint).not.toBeNull();
+    expect(checkpoint?.channel_values.total).toBe(5);
+});
+it("should process two inputs joined into one topic and produce two outputs", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const add10Each = jest.fn((x) => x.map((y) => y + 10).sort());
+    const one = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["inbox"]));
+    const chainThree = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["inbox"]));
+    const chainFour = Channel.subscribeTo("inbox")
+        .pipe(add10Each)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: {
+            one,
+            chainThree,
+            chainFour,
+        },
+        channels: {
+            inbox: new Topic(),
+            output: new LastValue(),
+            input: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+    });
+    // Invoke app and check results
+    for (let i = 0; i < 100; i += 1) {
+        expect(await app.invoke(2)).toEqual([13, 13]);
+    }
+    // Use Promise.all to simulate concurrent execution
+    const results = await Promise.all(Array(100)
+        .fill(null)
+        .map(async () => app.invoke(2)));
+    results.forEach((result) => {
+        expect(result).toEqual([13, 13]);
+    });
+});
+it("should invoke join then call other app", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const add10Each = jest.fn((x) => x.map((y) => y + 10));
+    const innerApp = new Pregel({
+        nodes: {
+            one: Channel.subscribeTo("input")
+                .pipe(addOne)
+                .pipe(Channel.writeTo(["output"])),
+        },
+        channels: {
+            output: new LastValue(),
+            input: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+    });
+    const one = Channel.subscribeTo("input")
+        .pipe(add10Each)
+        .pipe(Channel.writeTo(["inbox_one"]).map());
+    const two = Channel.subscribeTo("inbox_one")
+        .pipe(() => innerApp.map())
+        .pipe((x) => x.sort())
+        .pipe(Channel.writeTo(["outbox_one"]));
+    const chainThree = Channel.subscribeTo("outbox_one")
+        .pipe((x) => x.reduce((a, b) => a + b, 0))
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: {
+            one,
+            two,
+            chain_three: chainThree,
+        },
+        channels: {
+            inbox_one: new Topic(),
+            outbox_one: new Topic(),
+            output: new LastValue(),
+            input: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+    });
+    // Run the test 10 times sequentially
+    for (let i = 0; i < 10; i += 1) {
+        expect(await app.invoke([2, 3])).toEqual(27);
+    }
+    // Run the test 10 times in parallel
+    const results = await Promise.all(Array(10)
+        .fill(null)
+        .map(() => app.invoke([2, 3])));
+    expect(results).toEqual(Array(10).fill(27));
+});
+it("should handle two processes with one input and two outputs", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const one = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo([], {
+        output: new RunnablePassthrough(),
+        between: new RunnablePassthrough(),
+    }));
+    const two = Channel.subscribeTo("between")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const app = new Pregel({
+        nodes: { one, two },
+        channels: {
+            input: new LastValue(),
+            output: new LastValue(),
+            between: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+        streamChannels: ["output", "between"],
+    });
+    const results = await app.stream(2);
+    const streamResults = [];
+    for await (const chunk of results) {
+        streamResults.push(chunk);
+    }
+    expect(streamResults).toEqual([
+        { between: 3, output: 3 },
+        { between: 3, output: 4 },
+    ]);
+});
+it("should finish executing without output", async () => {
+    const addOne = jest.fn((x) => x + 1);
+    const one = Channel.subscribeTo("input")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["between"]));
+    const two = Channel.subscribeTo("between").pipe(addOne);
+    const app = new Pregel({
+        nodes: { one, two },
+        channels: {
+            input: new LastValue(),
+            between: new LastValue(),
+            output: new LastValue(),
+        },
+        inputs: "input",
+        outputs: "output",
+    });
+    // It finishes executing (once no more messages being published)
+    // but returns nothing, as nothing was published to OUT topic
+    expect(await app.invoke(2)).toBeUndefined();
+});
+it("should throw an error when no input channel is provided", () => {
+    const addOne = jest.fn((x) => x + 1);
+    const one = Channel.subscribeTo("between")
+        .pipe(addOne)
+        .pipe(Channel.writeTo(["output"]));
+    const two = Channel.subscribeTo("between").pipe(addOne);
+    // @ts-expect-error - this should throw an error
+    expect(() => new Pregel({ nodes: { one, two } })).toThrowError();
+});
+it("should type-error when Channel.subscribeTo would throw at runtime", () => {
+    expect(() => {
+        // @ts-expect-error - this would throw at runtime and thus we want it to become a type-error
+        Channel.subscribeTo(["input"], { key: "key" });
+    }).toThrow();
+});
+describe("StateGraph", () => {
+    class SearchAPI extends Tool {
+        constructor() {
+            super();
+            Object.defineProperty(this, "name", {
+                enumerable: true,
+                configurable: true,
+                writable: true,
+                value: "search_api"
+            });
+            Object.defineProperty(this, "description", {
+                enumerable: true,
+                configurable: true,
+                writable: true,
+                value: "A simple API that returns the input string."
+            });
+            Object.defineProperty(this, "schema", {
+                enumerable: true,
+                configurable: true,
+                writable: true,
+                value: z
+                    .object({
+                    input: z.string().optional(),
+                })
+                    .transform((data) => data.input)
+            });
+        }
+        async _call(query) {
+            return `result for ${query}`;
+        }
+    }
+    const tools = [new SearchAPI()];
+    const executeTools = async (data) => {
+        const newData = data;
+        const { agentOutcome } = newData;
+        delete newData.agentOutcome;
+        if (!agentOutcome || "returnValues" in agentOutcome) {
+            throw new Error("Agent has already finished.");
+        }
+        const observation = (await tools
+            .find((t) => t.name === agentOutcome.tool)
+            ?.invoke(agentOutcome.toolInput)) ?? "failed";
+        return {
+            steps: [[agentOutcome, observation]],
+        };
+    };
+    const shouldContinue = async (data) => {
+        if (data.agentOutcome && "returnValues" in data.agentOutcome) {
+            return "exit";
+        }
+        return "continue";
+    };
+    it("can invoke", async () => {
+        const prompt = PromptTemplate.fromTemplate("Hello!");
+        const llm = new FakeStreamingLLM({
+            responses: [
+                "tool:search_api:query",
+                "tool:search_api:another",
+                "finish:answer",
+            ],
+        });
+        const agentParser = (input) => {
+            if (input.startsWith("finish")) {
+                const answer = input.split(":")[1];
+                return {
+                    agentOutcome: {
+                        returnValues: { answer },
+                        log: input,
+                    },
+                };
+            }
+            const [, toolName, toolInput] = input.split(":");
+            return {
+                agentOutcome: {
+                    tool: toolName,
+                    toolInput,
+                    log: input,
+                },
+            };
+        };
+        const agent = async (state) => {
+            const chain = prompt.pipe(llm).pipe(agentParser);
+            const result = await chain.invoke({ input: state.input });
+            return {
+                ...result,
+            };
+        };
+        const graph = new StateGraph({
+            channels: {
+                input: null,
+                agentOutcome: null,
+                steps: {
+                    value: (x, y) => x.concat(y),
+                    default: () => [],
+                },
+            },
+        })
+            .addNode("agent", agent)
+            .addNode("tools", executeTools)
+            .addEdge(START, "agent")
+            .addConditionalEdges("agent", shouldContinue, {
+            continue: "tools",
+            exit: END,
+        })
+            .addEdge("tools", "agent")
+            .compile();
+        const result = await graph.invoke({ input: "what is the weather in sf?" });
+        expect(result).toEqual({
+            input: "what is the weather in sf?",
+            agentOutcome: {
+                returnValues: {
+                    answer: "answer",
+                },
+                log: "finish:answer",
+            },
+            steps: [
+                [
+                    {
+                        log: "tool:search_api:query",
+                        tool: "search_api",
+                        toolInput: "query",
+                    },
+                    "result for query",
+                ],
+                [
+                    {
+                        log: "tool:search_api:another",
+                        tool: "search_api",
+                        toolInput: "another",
+                    },
+                    "result for another",
+                ],
+            ],
+        });
+    });
+    it("can stream", async () => {
+        const prompt = PromptTemplate.fromTemplate("Hello!");
+        const llm = new FakeStreamingLLM({
+            responses: [
+                "tool:search_api:query",
+                "tool:search_api:another",
+                "finish:answer",
+            ],
+        });
+        const agentParser = (input) => {
+            if (input.startsWith("finish")) {
+                const answer = input.split(":")[1];
+                return {
+                    agentOutcome: {
+                        returnValues: { answer },
+                        log: input,
+                    },
+                };
+            }
+            const [, toolName, toolInput] = input.split(":");
+            return {
+                agentOutcome: {
+                    tool: toolName,
+                    toolInput,
+                    log: input,
+                },
+            };
+        };
+        const agent = async (state) => {
+            const chain = prompt.pipe(llm).pipe(agentParser);
+            const result = await chain.invoke({ input: state.input });
+            return {
+                ...result,
+            };
+        };
+        const app = new StateGraph({
+            channels: {
+                input: null,
+                agentOutcome: null,
+                steps: {
+                    value: (x, y) => x.concat(y),
+                    default: () => [],
+                },
+            },
+        })
+            .addNode("agent", agent)
+            .addNode("tools", executeTools)
+            .addEdge(START, "agent")
+            .addConditionalEdges("agent", shouldContinue, {
+            continue: "tools",
+            exit: END,
+        })
+            .addEdge("tools", "agent")
+            .compile();
+        const stream = await app.stream({ input: "what is the weather in sf?" });
+        const streamItems = [];
+        for await (const item of stream) {
+            streamItems.push(item);
+        }
+        expect(streamItems.length).toBe(5);
+        expect(streamItems[0]).toEqual({
+            agent: {
+                agentOutcome: {
+                    tool: "search_api",
+                    toolInput: "query",
+                    log: "tool:search_api:query",
+                },
+            },
+        });
+        // TODO: Need to rewrite this test.
+    });
+    it("can invoke a nested graph", async () => {
+        const innerGraph = new StateGraph({
+            channels: {
+                myKey: null,
+                myOtherKey: null,
+            },
+        })
+            .addNode("up", (state) => ({
+            myKey: `${state.myKey} there`,
+            myOtherKey: state.myOtherKey,
+        }))
+            .addEdge(START, "up")
+            .addEdge("up", END);
+        const graph = new StateGraph({
+            channels: {
+                myKey: null,
+                neverCalled: null,
+            },
+        })
+            .addNode("inner", innerGraph.compile())
+            .addNode("side", (state) => ({
+            myKey: `${state.myKey} and back again`,
+        }))
+            .addEdge("inner", "side")
+            .addEdge(START, "inner")
+            .addEdge("side", END)
+            .compile();
+        // call method / assertions
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const neverCalled = jest.fn((_) => {
+            throw new Error("This should never be called");
+        });
+        const result = await graph.invoke({
+            myKey: "my value",
+            neverCalled: new RunnableLambda({ func: neverCalled }),
+        });
+        expect(result).toEqual({
+            myKey: "my value there and back again",
+            neverCalled: new RunnableLambda({ func: neverCalled }),
+        });
+    });
+    it("can invoke a nested graph", async () => {
+        const innerGraph = new StateGraph({
+            channels: {
+                myKey: null,
+                myOtherKey: null,
+            },
+        })
+            .addNode("up", (state) => ({
+            myKey: `${state.myKey} there`,
+            myOtherKey: state.myOtherKey,
+        }))
+            .addEdge(START, "up")
+            .addEdge("up", END);
+        const graph = new StateGraph({
+            channels: {
+                myKey: null,
+                neverCalled: null,
+            },
+        })
+            .addNode("inner", innerGraph.compile())
+            .addNode("side", (state) => ({
+            myKey: `${state.myKey} and back again`,
+        }))
+            .addEdge("inner", "side")
+            .addEdge(START, "inner")
+            .addEdge("side", END)
+            .compile();
+        // call method / assertions
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const neverCalled = jest.fn((_) => {
+            throw new Error("This should never be called");
+        });
+        const result = await graph.invoke({
+            myKey: "my value",
+            neverCalled: new RunnableLambda({ func: neverCalled }),
+        });
+        expect(result).toEqual({
+            myKey: "my value there and back again",
+            neverCalled: new RunnableLambda({ func: neverCalled }),
+        });
+    });
+    it("Conditional edges is optional", async () => {
+        const nodeOne = (state) => {
+            const { keys } = state;
+            keys.value = 1;
+            return {
+                keys,
+            };
+        };
+        const nodeTwo = (state) => {
+            const { keys } = state;
+            keys.value = 2;
+            return {
+                keys,
+            };
+        };
+        const nodeThree = (state) => {
+            const { keys } = state;
+            keys.value = 3;
+            return {
+                keys,
+            };
+        };
+        const decideNext = (_) => "two";
+        const graph = new StateGraph({
+            channels: {
+                keys: null,
+            },
+        })
+            .addNode("one", nodeOne)
+            .addNode("two", nodeTwo)
+            .addNode("three", nodeThree)
+            .addEdge(START, "one")
+            .addConditionalEdges("one", decideNext)
+            .addEdge("two", "three")
+            .addEdge("three", END)
+            .compile();
+        // This will always return two, and two will always go to three
+        // meaning keys.value will always be 3
+        const result = await graph.invoke({ keys: { value: 0 } });
+        expect(result).toEqual({ keys: { value: 3 } });
+    });
+    it("In one fan out state graph waiting edge", async () => {
+        const sortedAdd = jest.fn((x, y) => [...x, ...y].sort());
+        function rewriteQuery(data) {
+            return { query: `query: ${data.query}` };
+        }
+        function analyzerOne(data) {
+            return { query: `analyzed: ${data.query}` };
+        }
+        function retrieverOne(_data) {
+            return { docs: ["doc1", "doc2"] };
+        }
+        function retrieverTwo(_data) {
+            return { docs: ["doc3", "doc4"] };
+        }
+        function qa(data) {
+            return { answer: data.docs?.join(",") };
+        }
+        const workflow = new StateGraph({
+            channels: {
+                query: null,
+                answer: null,
+                docs: { reducer: sortedAdd },
+            },
+        })
+            .addNode("rewrite_query", rewriteQuery)
+            .addNode("analyzer_one", analyzerOne)
+            .addNode("retriever_one", retrieverOne)
+            .addNode("retriever_two", retrieverTwo)
+            .addNode("qa", qa)
+            .addEdge(START, "rewrite_query")
+            .addEdge("rewrite_query", "analyzer_one")
+            .addEdge("analyzer_one", "retriever_one")
+            .addEdge("rewrite_query", "retriever_two")
+            .addEdge(["retriever_one", "retriever_two"], "qa")
+            .addEdge("qa", END);
+        const app = workflow.compile();
+        expect(await app.invoke({ query: "what is weather in sf" })).toEqual({
+            query: "analyzed: query: what is weather in sf",
+            docs: ["doc1", "doc2", "doc3", "doc4"],
+            answer: "doc1,doc2,doc3,doc4",
+        });
+    });
+});
+describe("PreBuilt", () => {
+    class SearchAPI extends Tool {
+        constructor() {
+            super();
+            Object.defineProperty(this, "name", {
+                enumerable: true,
+                configurable: true,
+                writable: true,
+                value: "search_api"
+            });
+            Object.defineProperty(this, "description", {
+                enumerable: true,
+                configurable: true,
+                writable: true,
+                value: "A simple API that returns the input string."
+            });
+        }
+        async _call(query) {
+            return `result for ${query}`;
+        }
+    }
+    const tools = [new SearchAPI()];
+    it("Can invoke createAgentExecutor", async () => {
+        const prompt = PromptTemplate.fromTemplate("Hello!");
+        const llm = new FakeStreamingLLM({
+            responses: [
+                "tool:search_api:query",
+                "tool:search_api:another",
+                "finish:answer",
+            ],
+        });
+        const agentParser = (input) => {
+            if (input.startsWith("finish")) {
+                const answer = input.split(":")[1];
+                return {
+                    returnValues: { answer },
+                    log: input,
+                };
+            }
+            const [, toolName, toolInput] = input.split(":");
+            return {
+                tool: toolName,
+                toolInput,
+                log: input,
+            };
+        };
+        const agent = prompt.pipe(llm).pipe(agentParser);
+        const agentExecutor = createAgentExecutor({
+            agentRunnable: agent,
+            tools,
+        });
+        const result = await agentExecutor.invoke({
+            input: "what is the weather in sf?",
+        });
+        expect(result).toEqual({
+            input: "what is the weather in sf?",
+            agentOutcome: {
+                returnValues: {
+                    answer: "answer",
+                },
+                log: "finish:answer",
+            },
+            steps: [
+                {
+                    action: {
+                        log: "tool:search_api:query",
+                        tool: "search_api",
+                        toolInput: "query",
+                    },
+                    observation: "result for query",
+                },
+                {
+                    action: {
+                        log: "tool:search_api:another",
+                        tool: "search_api",
+                        toolInput: "another",
+                    },
+                    observation: "result for another",
+                },
+            ],
+        });
+    });
+});
+describe("MessageGraph", () => {
+    class SearchAPI extends Tool {
+        constructor() {
+            super();
+            Object.defineProperty(this, "name", {
+                enumerable: true,
+                configurable: true,
+                writable: true,
+                value: "search_api"
+            });
+            Object.defineProperty(this, "description", {
+                enumerable: true,
+                configurable: true,
+                writable: true,
+                value: "A simple API that returns the input string."
+            });
+            Object.defineProperty(this, "schema", {
+                enumerable: true,
+                configurable: true,
+                writable: true,
+                value: z
+                    .object({
+                    input: z.string().optional(),
+                })
+                    .transform((data) => data.input)
+            });
+        }
+        async _call(query) {
+            return `result for ${query}`;
+        }
+    }
+    const tools = [new SearchAPI()];
+    it("can invoke a single message", async () => {
+        const model = new FakeChatModel({
+            responses: [
+                new AIMessage({
+                    content: "",
+                    additional_kwargs: {
+                        function_call: {
+                            name: "search_api",
+                            arguments: "query",
+                        },
+                    },
+                }),
+                new AIMessage({
+                    content: "",
+                    additional_kwargs: {
+                        function_call: {
+                            name: "search_api",
+                            arguments: "another",
+                        },
+                    },
+                }),
+                new AIMessage({
+                    content: "answer",
+                }),
+            ],
+        });
+        const toolExecutor = new ToolExecutor({ tools });
+        const shouldContinue = (data) => {
+            const lastMessage = data[data.length - 1];
+            // If there is no function call, then we finish
+            if (!("function_call" in lastMessage.additional_kwargs) ||
+                !lastMessage.additional_kwargs.function_call) {
+                return "end";
+            }
+            // Otherwise if there is, we continue
+            return "continue";
+        };
+        const callTool = async (data, options) => {
+            const lastMessage = data[data.length - 1];
+            const action = {
+                tool: lastMessage.additional_kwargs.function_call?.name ?? "",
+                toolInput: lastMessage.additional_kwargs.function_call?.arguments ?? "",
+                log: "",
+            };
+            const response = await toolExecutor.invoke(action, options?.config);
+            return new FunctionMessage({
+                content: JSON.stringify(response),
+                name: action.tool,
+            });
+        };
+        const app = new MessageGraph()
+            .addNode("agent", model)
+            .addNode("action", callTool)
+            .addEdge(START, "agent")
+            .addConditionalEdges("agent", shouldContinue, {
+            continue: "action",
+            end: END,
+        })
+            .addEdge("action", "agent")
+            .compile();
+        const result = await app.invoke(new HumanMessage("what is the weather in sf?"));
+        expect(result).toHaveLength(6);
+        expect(result).toStrictEqual([
+            new HumanMessage("what is the weather in sf?"),
+            new AIMessage({
+                content: "",
+                additional_kwargs: {
+                    function_call: {
+                        name: "search_api",
+                        arguments: "query",
+                    },
+                },
+            }),
+            new FunctionMessage({
+                content: '"result for query"',
+                name: "search_api",
+            }),
+            new AIMessage({
+                content: "",
+                additional_kwargs: {
+                    function_call: {
+                        name: "search_api",
+                        arguments: "another",
+                    },
+                },
+            }),
+            new FunctionMessage({
+                content: '"result for another"',
+                name: "search_api",
+            }),
+            new AIMessage("answer"),
+        ]);
+    });
+    it("can stream a list of messages", async () => {
+        const model = new FakeChatModel({
+            responses: [
+                new AIMessage({
+                    content: "",
+                    additional_kwargs: {
+                        function_call: {
+                            name: "search_api",
+                            arguments: "query",
+                        },
+                    },
+                }),
+                new AIMessage({
+                    content: "",
+                    additional_kwargs: {
+                        function_call: {
+                            name: "search_api",
+                            arguments: "another",
+                        },
+                    },
+                }),
+                new AIMessage({
+                    content: "answer",
+                }),
+            ],
+        });
+        const toolExecutor = new ToolExecutor({ tools });
+        const shouldContinue = (data) => {
+            const lastMessage = data[data.length - 1];
+            // If there is no function call, then we finish
+            if (!("function_call" in lastMessage.additional_kwargs) ||
+                !lastMessage.additional_kwargs.function_call) {
+                return "end";
+            }
+            // Otherwise if there is, we continue
+            return "continue";
+        };
+        const callTool = async (data, options) => {
+            const lastMessage = data[data.length - 1];
+            const action = {
+                tool: lastMessage.additional_kwargs.function_call?.name ?? "",
+                toolInput: lastMessage.additional_kwargs.function_call?.arguments ?? "",
+                log: "",
+            };
+            const response = await toolExecutor.invoke(action, options?.config);
+            return new FunctionMessage({
+                content: JSON.stringify(response),
+                name: action.tool,
+            });
+        };
+        const app = new MessageGraph()
+            .addNode("agent", model)
+            .addNode("action", callTool)
+            .addEdge(START, "agent")
+            .addConditionalEdges("agent", shouldContinue, {
+            continue: "action",
+            end: END,
+        })
+            .addEdge("action", "agent")
+            .compile();
+        const stream = await app.stream([
+            new HumanMessage("what is the weather in sf?"),
+        ]);
+        const streamItems = [];
+        for await (const item of stream) {
+            streamItems.push(item);
+        }
+        const lastItem = streamItems[streamItems.length - 1];
+        expect(Object.keys(lastItem)).toEqual(["agent"]);
+        expect(Object.values(lastItem)[0]).toEqual(new AIMessage("answer"));
+    });
+});
+it("StateGraph start branch then end", async () => {
+    const invalidBuilder = new StateGraph({
+        channels: {
+            my_key: { reducer: (x, y) => x + y },
+            market: null,
+        },
+    })
+        .addNode("tool_two_slow", (_) => ({ my_key: ` slow` }))
+        .addNode("tool_two_fast", (_) => ({ my_key: ` fast` }))
+        .addConditionalEdges(START, (state) => state.market === "DE" ? "tool_two_slow" : "tool_two_fast");
+    expect(() => invalidBuilder.compile()).toThrowError("Node `tool_two_slow` is a dead-end");
+    const toolTwoBuilder = new StateGraph({
+        channels: {
+            my_key: { reducer: (x, y) => x + y },
+            market: null,
+        },
+    })
+        .addNode("tool_two_slow", (_) => ({ my_key: ` slow` }))
+        .addNode("tool_two_fast", (_) => ({ my_key: ` fast` }))
+        .addConditionalEdges({
+        source: START,
+        path: (state) => state.market === "DE" ? "tool_two_slow" : "tool_two_fast",
+        then: END,
+    });
+    const toolTwo = toolTwoBuilder.compile();
+    expect(await toolTwo.invoke({ my_key: "value", market: "DE" })).toEqual({
+        my_key: "value slow",
+        market: "DE",
+    });
+    expect(await toolTwo.invoke({ my_key: "value", market: "US" })).toEqual({
+        my_key: "value fast",
+        market: "US",
+    });
+    const toolTwoWithCheckpointer = toolTwoBuilder.compile({
+        checkpointer: SqliteSaver.fromConnString(":memory:"),
+        interruptBefore: ["tool_two_fast", "tool_two_slow"],
+    });
+    await expect(() => toolTwoWithCheckpointer.invoke({ my_key: "value", market: "DE" })).rejects.toThrowError("thread_id");
+    // const thread1 = { configurable: { thread_id: "1" } }
+    // expect(toolTwoWithCheckpointer.invoke({ my_key: "value", market: "DE" }, thread1)).toEqual({ my_key: "value", market: "DE" })
+    // expect(toolTwoWithCheckpointer.getState(thread1)).toEqual({
+    //   values: { my_key: "value", market: "DE" },
+    //   next: ["tool_two_slow"],
+    //   config: toolTwoWithCheckpointer.checkpointer.getTuple(thread1).config,
+    //   metadata: { source: "loop", step: 0, writes: null },
+    //   parentConfig: [...toolTwoWithCheckpointer.checkpointer.list(thread1, { limit: 2 })].pop().config
+    // })
+    // expect(toolTwoWithCheckpointer.invoke(null, thread1, { debug: 1 })).toEqual({ my_key: "value slow", market: "DE" })
+    // expect(toolTwoWithCheckpointer.getState(thread1)).toEqual({
+    //   values: { my_key
+    //     : "value slow", market: "DE" },
+    //   next: [],
+    //   config: (await toolTwoWithCheckpointer.checkpointer!.getTuple(thread1))!.config,
+    //   metadata: { source: "loop", step: 1, writes: { tool_two_slow: { my_key: " slow" } } },
+    //   parentConfig: [...toolTwoWithCheckpointer.checkpointer!.list(thread1, { limit: 2 })].pop().config
+});
+/**
+ * def test_branch_then_node(snapshot: SnapshotAssertion) -> None:
+    class State(TypedDict):
+        my_key: Annotated[str, operator.add]
+        market: str
+
+    # this graph is invalid because there is no path to "finish"
+    invalid_graph = StateGraph(State)
+    invalid_graph.set_entry_point("prepare")
+    invalid_graph.set_finish_point("finish")
+    invalid_graph.add_conditional_edges(
+        source="prepare",
+        path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
+        path_map=["tool_two_slow", "tool_two_fast"],
+    )
+    invalid_graph.add_node("prepare", lambda s: {"my_key": " prepared"})
+    invalid_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
+    invalid_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
+    invalid_graph.add_node("finish", lambda s: {"my_key": " finished"})
+    with pytest.raises(ValueError):
+        invalid_graph.compile()
+
+    tool_two_graph = StateGraph(State)
+    tool_two_graph.set_entry_point("prepare")
+    tool_two_graph.set_finish_point("finish")
+    tool_two_graph.add_conditional_edges(
+        source="prepare",
+        path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
+        then="finish",
+    )
+    tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"})
+    tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
+    tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
+    tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"})
+    tool_two = tool_two_graph.compile()
+    assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
+    assert tool_two.get_graph().draw_mermaid() == snapshot
+
+    assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == {
+        "my_key": "value prepared slow finished",
+        "market": "DE",
+    }
+    assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
+        "my_key": "value prepared fast finished",
+        "market": "US",
+    }
+ */
+it("StateGraph branch then node", async () => {
+    const invalidBuilder = new StateGraph({
+        channels: {
+            my_key: { reducer: (x, y) => x + y },
+            market: null,
+        },
+    })
+        .addNode("prepare", (_) => ({ my_key: ` prepared` }))
+        .addNode("tool_two_slow", (_) => ({ my_key: ` slow` }))
+        .addNode("tool_two_fast", (_) => ({ my_key: ` fast` }))
+        .addNode("finish", (_) => ({ my_key: ` finished` }))
+        .addEdge(START, "prepare")
+        .addConditionalEdges({
+        source: "prepare",
+        path: (state) => state.market === "DE" ? "tool_two_slow" : "tool_two_fast",
+        pathMap: ["tool_two_slow", "tool_two_fast"],
+    })
+        .addEdge("finish", END);
+    expect(() => invalidBuilder.compile()).toThrowError();
+    const toolBuilder = new StateGraph({
+        channels: {
+            my_key: { reducer: (x, y) => x + y },
+            market: null,
+        },
+    })
+        .addNode("prepare", (_) => ({ my_key: ` prepared` }))
+        .addNode("tool_two_slow", (_) => ({ my_key: ` slow` }))
+        .addNode("tool_two_fast", (_) => ({ my_key: ` fast` }))
+        .addNode("finish", (_) => ({ my_key: ` finished` }))
+        .addEdge(START, "prepare")
+        .addConditionalEdges({
+        source: "prepare",
+        path: (state) => state.market === "DE" ? "tool_two_slow" : "tool_two_fast",
+        then: "finish",
+    })
+        .addEdge("finish", END);
+    const tool = toolBuilder.compile();
+    expect(await tool.invoke({ my_key: "value", market: "DE" })).toEqual({
+        my_key: "value prepared slow finished",
+        market: "DE",
+    });
+    expect(await tool.invoke({ my_key: "value", market: "FR" })).toEqual({
+        my_key: "value prepared fast finished",
+        market: "FR",
+    });
+});