/* 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 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: {},
        };
        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",
    })
        .addEdge("tool_two_fast", END)
        .addEdge("tool_two_slow", 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");
    async function last(iter) {
        // eslint-disable-next-line no-undef-init
        let value = undefined;
        for await (value of iter) {
            // do nothing
        }
        return value;
    }
    const thread1 = { configurable: { thread_id: "1" } };
    expect(await toolTwoWithCheckpointer.invoke({ my_key: "value", market: "DE" }, thread1)).toEqual({ my_key: "value", market: "DE" });
    expect(await toolTwoWithCheckpointer.getState(thread1)).toEqual({
        values: { my_key: "value", market: "DE" },
        next: ["tool_two_slow"],
        config: (await toolTwoWithCheckpointer.checkpointer.getTuple(thread1))
            .config,
        createdAt: (await toolTwoWithCheckpointer.checkpointer.getTuple(thread1))
            .checkpoint.ts,
        metadata: { source: "loop", step: 0, writes: null },
        parentConfig: (await last(toolTwoWithCheckpointer.checkpointer.list(thread1, 2))).config,
    });
    expect(await toolTwoWithCheckpointer.invoke(null, thread1)).toEqual({
        my_key: "value slow",
        market: "DE",
    });
    expect(await toolTwoWithCheckpointer.getState(thread1)).toEqual({
        values: { my_key: "value slow", market: "DE" },
        next: [],
        config: (await toolTwoWithCheckpointer.checkpointer.getTuple(thread1))
            .config,
        createdAt: (await toolTwoWithCheckpointer.checkpointer.getTuple(thread1))
            .checkpoint.ts,
        metadata: {
            source: "loop",
            step: 1,
            writes: { tool_two_slow: { my_key: " slow" } },
        },
        parentConfig: (await last(toolTwoWithCheckpointer.checkpointer.list(thread1, 2))).config,
    });
});
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",
    })
        .addEdge("tool_two_fast", "finish")
        .addEdge("tool_two_slow", "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",
    });
});
