Custom Checkpointer / MemorySaver - is there a comprehensive saver with a Test file to use as a guide? (LangGraph JS)

hgoona commented 3 days ago

Hi team, @jacoblee93 @bracesproul Is there a MemorySaver for a database that has a complete or "comprehensive" .test.ts written for it?

I've spent the greater part of the last 7 days trying to comprehend what on earth is happening inside the Memory Saver adapter (or whatever this "thing" is called, in order to try creating a custom version that works with SurrealDB. I've been looking at the Sqlite version: "SqliteSaver" and the test file, here: https://github.com/langchain-ai/langgraphjs/blob/main/libs/checkpoint-sqlite/src/index.ts#L45 but I'm having errors translating it for SurrealDB. While testing my SdbSaver against the test.ts I'm getting errors here:

✓ should get undefined checkpoint
     ✓ should save first checkpoint
     ✓ should add some writes
     × should get first checkpoint tuple
     × should save second checkpoint
     × should list checkpoints

In looking at the .getTuple method, the Type specified in the Sqlite version, and the corresponding test I'm wondering if the test file is actually correctly written, or if it's incomplete? It could very well be my mistake but I'd greatly appreciate a pointer!

If there is a "complete" reference Saver with an accompanying test file can you please recommend? I realise I'm trying to reverse engineer this, fairly blindly, and now realise the test I'm looking at may have been misleading me.

jacoblee93 commented 3 days ago

Have overhauling docs/reference examples as one of the next things I'll look at - stay tuned!

hgoona commented 2 days ago

@jacoblee93 looking forward to this, but would it be at all possible to get your review of my (draft) adapter for SurrealDB?

Here's what I have so far: sdbSaver.ts sdbSaver.test.ts - tested with Vitest ☝🏾these are both based off the Sqlite saver (where I also looked other "savers?" to try fill in the blanks.

NOTES: The SdbSaver is currently failing these tests:

❯ SdbSaver - should save and retrieve checkpoints correctly (8)
     ✅ confirm TestSdb - is connected
     ✅ get undefined checkpoint
     ✅ save first checkpoint
     ✅ add some writes
     ❌ get first checkpoint tuple
     ❌ save second checkpoint
     ✅ list checkpoints
     ➖ KEEP_TEST_DB_ALIVE

CURRENT UNDERSTANDING OF "THE SAVER??":

db storage has 2 tables: checkpoints,writes
What is writes? What does it do? Is it a temporary file or permanent thing to be stored in the DB?
If "state" is stored within checkpoints, are messages (AI, Human, Tool Call, Tool Messages) also stored within checkpoints? <<< if Yes, shouldn't these messages be a separate table of unique messages that is then referenced by the specific checkpoint:record to avoid duplication of identical data that would otherwise be copied inside each consecutive checkpoints record that is part of the same "thread" ? Therefore, when a thread is branched, the next checkpoint can reference the existing messages from before the "branching-point" and append new messages for states that continue the threads beyond the "branching-point". 👆🏾at least this is how I was planning to maintain history before "Checkpointers" were created by LangChain.
Also, what is the checkpoint_ns ??

hgoona commented 2 days ago

WORK-IN-PROGRESS for SdbSaver: I have kept the SqliteSaver code commented out as a direct reference to hopeful come up with an equivalent query in SurrealDB🤞🏾 - I've attempted to NOT deviate from the SqliteSaver query structure as much as possible (given my current lack of clarity on the above).

// sdbSaver.ts

// import Database, { Database as DatabaseType } from "better-sqlite3";
import { debugLog } from "$lib/debugLog";
import type { SdbQueryClass } from "$lib/server/sdbUtils";
// import { SdbQuery } from "$lib/server/sdbUtils";
import type { RunnableConfig } from "@langchain/core/runnables";
import {
    BaseCheckpointSaver,
    type Checkpoint,
    type CheckpointListOptions,
    type CheckpointTuple,
    type SerializerProtocol,
    type PendingWrite,
    type CheckpointMetadata,
} from "@langchain/langgraph-checkpoint";
// import type Surreal from "surrealdb";

interface CheckpointRow {
    checkpoint: string;
    metadata: string;
    parent_checkpoint_id?: string;
    thread_id: string;
    checkpoint_id: string;
    checkpoint_ns?: string;
    cpType?: string; //fka: 'type'
}

interface WritesRow {
    thread_id: string;
    checkpoint_ns: string;
    checkpoint_id: string;
    task_id: string;
    idx: number;
    channel: string;
    cpType?: string; //fka: 'type'
    cpValue?: string; //fka: 'value
}

// DEBUGS
const debug_checkpointer = true
// const debug_checkpointer = false

// const debug_output = true
// // const debug_output = false

const serviceName = "Memory5b_sdb"

/**
 * Based on SqliteSaver
 */
export class SdbSaver extends BaseCheckpointSaver {
    // constructor(serde?: SerializerProtocol) {
    //  super(serde);
    // }

    // public db: Surreal;

    // constructor(db: Surreal, serde?: SerializerProtocol) {
    public db: SdbQueryClass;

    constructor(db: SdbQueryClass, serde?: SerializerProtocol) {
        super(serde);
        this.db = db;
    }
    // //   constructor(db: DatabaseType, serde?: SerializerProtocol) {
    // //     super(serde);
    // //     this.db = db;
    // //     this.isSetup = false;
    // //   }

    //   static fromConnString(connStringOrLocalPath: string): SqliteSaver {
    //     return new SqliteSaver(new Database(connStringOrLocalPath));
    //   }

    //   protected setup(): void {
    //     if (this.isSetup) {
    //       return;
    //     }

    //     this.db.pragma("journal_mode=WAL");
    //     this.db.exec(`
    // CREATE TABLE IF NOT EXISTS checkpoints (
    //   thread_id TEXT NOT NULL,
    //   checkpoint_ns TEXT NOT NULL DEFAULT '',
    //   checkpoint_id TEXT NOT NULL,
    //   parent_checkpoint_id TEXT,
    //   type TEXT,
    //   checkpoint BLOB,
    //   metadata BLOB,
    //   PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
    // );`);
    //     this.db.exec(`
    // CREATE TABLE IF NOT EXISTS writes (
    //   thread_id TEXT NOT NULL,
    //   checkpoint_ns TEXT NOT NULL DEFAULT '',
    //   checkpoint_id TEXT NOT NULL,
    //   task_id TEXT NOT NULL,
    //   idx INTEGER NOT NULL,
    //   channel TEXT NOT NULL,
    //   type TEXT,
    //   value BLOB,
    //   PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
    // );`);

    //     this.isSetup = true;
    //   }

    // NOTE
    //   PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
    //   PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)

    async getTuple(config: RunnableConfig): Promise<CheckpointTuple | undefined> {
        debugLog(debug_checkpointer, `${serviceName} getTuple▶`)

        const {
            thread_id,
            // checkpoint_ns = "", // why is this here??
            checkpoint_ns, // why is this here??
            checkpoint_id,
        } = config.configurable ?? {};

        console.log("getTuple:", {
            thread_id,
            checkpoint_ns,
            checkpoint_id
        });

        // let row: CheckpointRow;
        let rowSDB: CheckpointRow;

        if (checkpoint_id) {
            console.log("getTuple▶1");
            // row = this.db
            //  .prepare(
            //      `SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?`
            //  )
            //  .get(thread_id, checkpoint_ns, checkpoint_id) as CheckpointRow;

            const query = /*surql*/ `
                    SELECT * FROM checkpoints
                    WHERE thread_id = $thread_id
                    AND checkpoint_ns = $checkpoint_ns
                    AND checkpoint_id = $checkpoint_id
                `
            // rowSDB = await SdbQuery.query(query, { thread_id, checkpoint_ns, checkpoint_id }) as CheckpointRow
            rowSDB = await this.db.query(query, { thread_id, checkpoint_ns, checkpoint_id }) as CheckpointRow
        } else { // no checkpoint_id
            console.log("getTuple▶2");
            // row = this.db
            //  .prepare(
            //      `SELECT thread_id, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? AND checkpoint_ns = ? ORDER BY checkpoint_id DESC LIMIT 1`
            //  )
            //  .get(thread_id, checkpoint_ns) as CheckpointRow;

            const query = /*surql*/ `
                SELECT * FROM checkpoints
                WHERE thread_id = $thread_id
                AND checkpoint_ns = $checkpoint_ns
                ORDER BY checkpoint_id DESC 
                LIMIT 1
            `
            rowSDB = await this.db.query(query, { thread_id, checkpoint_ns }) as CheckpointRow
        }

        console.log("getTuple:", { rowSDB });

        // if (row === undefined) {
        if (rowSDB === undefined) {
            return undefined;
        }

        let finalConfig = config;
        if (!checkpoint_id) {
            console.log("getTuple▶3");

            finalConfig = {
                configurable: {
                    // thread_id: row.thread_id,
                    thread_id: rowSDB.thread_id,
                    checkpoint_ns: checkpoint_ns ? checkpoint_ns : "", //prevent returning 'undefined'
                    // checkpoint_id: row.checkpoint_id,
                    checkpoint_id: rowSDB.checkpoint_id,
                },
            };
        }
        if (
            finalConfig.configurable?.thread_id === undefined ||
            finalConfig.configurable?.checkpoint_id === undefined
        ) {
            console.log("getTuple▶4");

            throw new Error("Missing thread_id or checkpoint_id");
        }
        console.log("getTuple▶5");

        // find any pending writes
        // const pendingWritesRows = this.db
        //  .prepare(
        //      `SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?`
        //  )
        //  .all(
        //      finalConfig.configurable.thread_id.toString(),
        //      checkpoint_ns,
        //      finalConfig.configurable.checkpoint_id.toString()
        //  ) as WritesRow[];

        const query = /*surql*/ `
            SELECT * FROM writes
            WHERE thread_id = $thread_id
            AND checkpoint_ns = $checkpoint_ns
            AND checkpoint_id = $checkpoint_id
        `
        const bindings = {
            thread_id: finalConfig.configurable.thread_id.toString(),
            checkpoint_ns,
            checkpoint_id: finalConfig.configurable.checkpoint_id.toString()
        }

        const pendingWritesRowsSDB = await this.db.query(query, {
            bindings
        }, true) as WritesRow[]

        const pendingWrites = await Promise.all(
            // pendingWritesRows.map(async (row) => {
            pendingWritesRowsSDB.map(async (rowSDB) => {
                return [
                    // row.task_id,
                    // row.channel,
                    // await this.serde.loadsTyped(row.type ?? "json", row.value ?? ""),
                    rowSDB.task_id,
                    rowSDB.channel,
                    await this.serde.loadsTyped(rowSDB.cpType ?? "json", rowSDB.cpValue ?? ""),
                ] as [string, string, unknown];
            })
        );

        console.log("getTuple▶6");

        // const checkpointSerde = (await this.serde.loadsTyped(
        //  // row.type ?? "json",
        //  // row.checkpoint
        //  rowSDB.cpType ?? "json",
        //  rowSDB.checkpoint
        // )) as Checkpoint

        // const metadataSerde = (await this.serde.loadsTyped(
        //  // row.type ?? "json",
        //  // row.metadata
        //  rowSDB.cpType ?? "json",
        //  rowSDB.metadata
        // )) as CheckpointMetadata

        let checkpointSerde: Checkpoint = {} as Checkpoint
        let metadataSerde

        try {
            checkpointSerde = await this.serde.loadsTyped(
                rowSDB.cpType ?? "json",
                rowSDB.checkpoint ?? "{}"
            ) as Checkpoint

            metadataSerde = (await this.serde.loadsTyped(
                rowSDB.cpType ?? "json",
                rowSDB.metadata ?? "{}"
            )) as CheckpointMetadata
        } catch (error) {
            console.error(error)
            // handle the error as needed
        }

        const finalOutput: CheckpointTuple = {
            config: finalConfig,
            checkpoint: checkpointSerde,
            metadata: metadataSerde,
            // parentConfig: row.parent_checkpoint_id
            parentConfig: rowSDB.parent_checkpoint_id
                ? {
                    configurable: {
                        // thread_id: row.thread_id,
                        thread_id: rowSDB.thread_id,
                        checkpoint_ns,
                        // checkpoint_id: row.parent_checkpoint_id,
                        checkpoint_id: rowSDB.parent_checkpoint_id,
                    },
                }
                : undefined,
            pendingWrites,
        };

        console.log("getTuple▶7:", { finalOutput });
        return finalOutput
    }

    async *list(
        config: RunnableConfig,
        options?: CheckpointListOptions
    ): AsyncGenerator<CheckpointTuple> {
        const { limit, before } = options ?? {};
        debugLog(debug_checkpointer, `${serviceName} list▶`)

        const thread_id = config.configurable?.thread_id;
        // let sql = `SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata FROM checkpoints WHERE thread_id = ? ${before ? "AND checkpoint_id < ?" : ""
        //     } ORDER BY checkpoint_id DESC`;
        // if (limit) {
        //     sql += ` LIMIT ${limit}`;
        // }
        // const args = [thread_id, before?.configurable?.checkpoint_id].filter(
        //     Boolean
        // );

        // const rows: CheckpointRow[] = this.db
        //     .prepare(sql)
        //     .all(...args) as CheckpointRow[];

        const binding = {
            thread_id,
            beforeId: before?.configurable?.checkpoint_id || null
        }
        const query = /*surql*/ `
            SELECT * FROM checkpoints
            WHERE thread_id = $thread_id
            ${before ? `AND checkpoint_id < $beforeId` : ""}
            ORDER BY checkpoint_id DESC
            ${limit ? `LIMIT $limit` : ""}
        `
        const rows = await this.db.query(query, binding) as CheckpointRow[]

        // if (rows) {
        if (Array.isArray(rows) && rows.length > 0) {
            for (const row of rows) {
                yield {
                    config: {
                        configurable: {
                            thread_id: row.thread_id,
                            checkpoint_ns: row.checkpoint_ns,
                            checkpoint_id: row.checkpoint_id,
                        },
                    },
                    checkpoint: (await this.serde.loadsTyped(
                        row.cpType ?? "json",
                        row.checkpoint
                    )) as Checkpoint,
                    metadata: (await this.serde.loadsTyped(
                        row.cpType ?? "json",
                        row.metadata
                    )) as CheckpointMetadata,
                    parentConfig: row.parent_checkpoint_id
                        ? {
                            configurable: {
                                thread_id: row.thread_id,
                                checkpoint_ns: row.checkpoint_ns,
                                checkpoint_id: row.parent_checkpoint_id,
                            },
                        }
                        : undefined,
                };
            }
        }
    }

    /**
     * 
     * @note PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
     * 
     * @param config 
     * @param checkpoint 
     * @param metadata 
     * @returns 
     */
    async put(
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata
    ): Promise<RunnableConfig> {
        debugLog(debug_checkpointer, `${serviceName} put▶`)

        const [type1, serializedCheckpoint] = this.serde.dumpsTyped(checkpoint);
        const [type2, serializedMetadata] = this.serde.dumpsTyped(metadata);
        if (type1 !== type2) {
            throw new Error(
                "Failed to serialized checkpoint and metadata to the same type."
            );
        }
        // const row = [
        //  config.configurable?.thread_id?.toString(),
        //  config.configurable?.checkpoint_ns,
        //  checkpoint.id,
        //  config.configurable?.checkpoint_id,
        //  type1,
        //  serializedCheckpoint,
        //  serializedMetadata,
        // ];

        // this.db
        //  .prepare(
        //      `INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)`
        //  )
        //  .run(...row);

        const binding: CheckpointRow = {
            thread_id: config.configurable?.thread_id?.toString(),
            checkpoint_ns: config.configurable?.checkpoint_ns,
            checkpoint_id: checkpoint.id,
            parent_checkpoint_id: config.configurable?.checkpoint_id,
            cpType: type1,
            checkpoint: Buffer.from(serializedCheckpoint).toString("utf-8"),
            metadata: Buffer.from(serializedMetadata).toString("utf-8"),
            // checkpoint: serializedCheckpoint,
            // metadata: serializedMetadata,
        }

        // TODO: review if this is the same as Sqlite `INSERT OR REPLACE INTO`
        const query = /*surql*/ `
            INSERT INTO checkpoints {
                id: [$thread_id, $checkpoint_ns, $checkpoint_id],
                thread_id: $thread_id,
                checkpoint_ns: $checkpoint_ns,
                checkpoint_id: $checkpoint_id,
                parent_checkpoint_id: $parent_checkpoint_id,
                cpType: $cpType,
                checkpoint: $checkpoint,
                metadata: $metadata
            } 
            ON DUPLICATE KEY UPDATE
                parent_checkpoint_id = $Input.parent_checkpoint_id,
                cpType = $Input.cpType,
                checkpoint = $Input.checkpoint,
                metadata = $Input.metadata
            ;
        `
        // const query = /*surql*/ `
        //  INSERT INTO checkpoints {
        //      thread_id: $thread_id
        //      checkpoint_ns: $checkpoint_ns
        //      checkpoint_id: $checkpoint_id
        //      parent_checkpoint_id: $parent_checkpoint_id
        //      cpType: $cpType
        //      checkpoint: $checkpoint
        //      metadata: $metadata
        //  }
        // `
        // WHERE thread_id = $thread_id
        // AND checkpoint_ns = $checkpoint_ns
        // AND checkpoint_id = $checkpoint_id

        await this.db.query(query, binding as unknown as SurrealBindings)
        // const results = await SdbQuery.query(query, binding as unknown as SurrealBindings)

        const configurable = {
            thread_id: config.configurable?.thread_id,
            checkpoint_ns: config.configurable?.checkpoint_ns,
            checkpoint_id: checkpoint.id,
        }

        return {
            configurable: configurable,
        };
    }

    /**
     * @note PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
     * 
     * @param config 
     * @param writes 
     * @param taskId 
     */
    async putWrites(
        config: RunnableConfig,
        writes: PendingWrite[],
        taskId: string
    ): Promise<void> {
        debugLog(debug_checkpointer, `${serviceName} putWrites▶`)

        // const stmt = this.db.prepare(`
        //  INSERT OR REPLACE INTO writes 
        //  (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) 
        //  VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        // `);

        // const transaction = this.db.transaction((rows) => {
        //  for (const row of rows) {
        //      stmt.run(...row);
        //  }
        // });

        /** get the rows from the writes */
        // const rows: WritesRow[][] = writes.map((write, idx) => {
        const rows = writes.map((write, idx) => {
            const [cpType, serializedWrite] = this.serde.dumpsTyped(write[1]);
            return [
                config.configurable?.thread_id,
                config.configurable?.checkpoint_ns,
                config.configurable?.checkpoint_id,
                taskId,
                idx,
                write[0],
                cpType,
                serializedWrite,
            ];
        });

        console.log("writes:", writes)
        console.log("rows:", rows)

        // transaction(rows);

        // // TODO loop over the `rows` and run the SdbQuery fn using `query` and spread the variables from each `row` in `rows` as a `binding` in the SdbQuery

        // // const binding: WritesRow = {
        // //   thread_id,
        // //   checkpoint_ns,
        // //   checkpoint_id,
        // //   task_id,
        // //   idx,
        // //   channel,
        // //   type,
        // //   value
        // // }
        // // const query = /*surql*/ `
        // //   INSERT INTO writes {
        // //       thread_id: $thread_id,
        // //       checkpoint_ns: $checkpoint_ns,
        // //       checkpoint_id: $checkpoint_id,
        // //       task_id: $task_id,
        // //       idx: $idx,
        // //       channel: $channel,
        // //       type: $type,
        // //       value: $value               
        // //   }
        // // `

        // SurrealDB QUERY
        const query = /*surql*/ `
            INSERT INTO writes {
                id: [$thread_id, $checkpoint_ns, $checkpoint_id, $task_id, $idx],
                thread_id: $thread_id,
                checkpoint_ns: $checkpoint_ns,
                checkpoint_id: $checkpoint_id,
                task_id: $task_id,
                idx: $idx,
                channel: $channel,
                cpType: $cpType,
                cpValue: $cpValue
            } 
            ON DUPLICATE KEY UPDATE
                channel = $channel,
                cpType = $cpType,
                cpValue = $cpValue              
            ;
        `

        for (const row of rows) {
            const [thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, cpType, cpValue] = row
            // const binding: WritesRow = {
            const binding = {
                thread_id,
                checkpoint_ns,
                checkpoint_id,
                task_id,
                idx,
                channel,
                cpType,
                cpValue,
            }
            await this.db.query(query, binding as unknown as SurrealBindings)
        }

        // // const results = await SdbQuery.query(query, binding as unknown as SurrealBindings)
    }
}

hgoona commented 2 days ago

Test file, tested with Vitest:

// sdbSaver.test.ts

import { describe, it, expect, beforeAll, type TestOptions } from "vitest";
// import { describe, it, expect } from "@jest/globals";
import {
    type Checkpoint,
    type CheckpointTuple,
    uuid6,
} from "@langchain/langgraph-checkpoint";
import { TestSdb1 } from "$lib/server/sdbTestDb";
import { SdbQueryClass } from "$lib/server/sdbUtils";
import { SdbSaver } from "./sdbSaver";
// import { SqliteSaver } from "../index.js";

const checkpoint1: Checkpoint = {
    v: 1,
    id: uuid6(-1),
    ts: "2024-04-19T17:19:07.952Z",
    channel_values: {
        someKey1: "someValue1",
    },
    channel_versions: {
        someKey2: 1,
    },
    versions_seen: {
        someKey3: {
            someKey4: 1,
        },
    },
    pending_sends: [],
};
console.log({ checkpoint1 });

const checkpoint2: Checkpoint = {
    v: 1,
    id: uuid6(1),
    ts: "2024-04-20T17:19:07.952Z",
    channel_values: {
        someKey1: "someValue2",
    },
    channel_versions: {
        someKey2: 2,
    },
    versions_seen: {
        someKey3: {
            someKey4: 2,
        },
    },
    pending_sends: [],
};

describe("SdbSaver - should save and retrieve checkpoints correctly", () => {
    /** 
     * New instance of SdbSaver with (test) DB attached 
     */
    let sdbSaver: SdbSaver

    /**
     * SDBTester is a SdbQueryClass 
     * that connects to the TestSdb1 database (in-Memory DB).
     */
    let SDBTester: SdbQueryClass

    beforeAll(() => {
        SDBTester = new SdbQueryClass(TestSdb1)

        // initialize SdbSaver
        sdbSaver = new SdbSaver(SDBTester); // must be an SdbQueryClass
    })

    it("confirm TestSdb - is connected", async () => {
        const info = await SDBTester.query("INFO FOR ROOT;")
        // console.log({ info })

        expect(info).not.toBeNull()
    })

    // CHECKPOINT SAVER TESTS
    // it("should save and retrieve checkpoints correctly", async () => {
    // const sqliteSaver = SqliteSaver.fromConnString(":memory:");
    // const sdbSaver = new SdbSaver(TestSdb1);

    it("get undefined checkpoint", async () => {
        // get undefined checkpoint
        const undefinedCheckpoint = await sdbSaver.getTuple({
            configurable: { thread_id: "1" },
        });
        expect(undefinedCheckpoint).toBeUndefined();
    });

    it("save first checkpoint", async () => {
        // save first checkpoint
        const runnableConfig = await sdbSaver.put(
            { configurable: { thread_id: "1" } },
            checkpoint1,
            { source: "update", step: -1, writes: null, parents: {} }
        );
        expect(runnableConfig).toEqual({
            configurable: {
                thread_id: "1",
                checkpoint_id: checkpoint1.id,
            },
        });
    });

    it("add some writes", async () => {
        // add some writes
        await sdbSaver.putWrites(
            {
                configurable: {
                    checkpoint_id: checkpoint1.id,
                    checkpoint_ns: "",
                    thread_id: "1",
                },
            },
            [["bar", "baz"]],
            "foo"
        );
    });

    it("get first checkpoint tuple", async () => {
        // get first checkpoint tuple
        const firstCheckpointTuple = await sdbSaver.getTuple({
            configurable: { thread_id: "1" },
        });
        console.log("firstCheckpointTuple:", JSON.stringify(firstCheckpointTuple, null, 2));
        console.log("checkpoint1.id:", checkpoint1.id);

        expect(firstCheckpointTuple?.config).toEqual({
            configurable: {
                thread_id: "1",
                checkpoint_ns: "",
                checkpoint_id: checkpoint1.id,
            },
        });
        expect(firstCheckpointTuple?.checkpoint).toEqual(checkpoint1);
        expect(firstCheckpointTuple?.parentConfig).toBeUndefined();
        expect(firstCheckpointTuple?.pendingWrites).toEqual([
            ["foo", "bar", "baz"],
        ]);
    });

    it("save second checkpoint", async () => {
        // save second checkpoint
        await sdbSaver.put(
            {
                configurable: {
                    thread_id: "1",
                    checkpoint_id: "2024-04-18T17:19:07.952Z",
                },
            },
            checkpoint2,
            { source: "update", step: -1, writes: null, parents: {} }
        );

        // verify that parentTs is set and retrieved correctly for second checkpoint
        const secondCheckpointTuple = await sdbSaver.getTuple({
            configurable: { thread_id: "1" },
        });
        expect(secondCheckpointTuple?.parentConfig).toEqual({
            configurable: {
                thread_id: "1",
                checkpoint_ns: "",
                checkpoint_id: "2024-04-18T17:19:07.952Z",
            },
        });
    });

    it("list checkpoints", async () => {
        // list checkpoints
        const checkpointTupleGenerator = await sdbSaver.list({
            configurable: { thread_id: "1" },
        });
        const checkpointTuples: CheckpointTuple[] = [];
        for await (const checkpoint of checkpointTupleGenerator) {
            checkpointTuples.push(checkpoint);
        }
        expect(checkpointTuples.length).toBe(2);

        const checkpointTuple1 = checkpointTuples[0];
        const checkpointTuple2 = checkpointTuples[1];
        expect(checkpointTuple1.checkpoint.ts).toBe("2024-04-20T17:19:07.952Z");
        expect(checkpointTuple2.checkpoint.ts).toBe("2024-04-19T17:19:07.952Z");
    });

    //keep the Vitest running to inspect the in-memory DB
    const options: TestOptions = {
        timeout: 2147483647, //max❗❗❗ 24.8 days
    }
    it("KEEP_TEST_DB_ALIVE",
        options,
        async () => {
            await new Promise(resolve => { })
        }
    );

});

langchain-ai / langgraphjs

Custom Checkpointer / MemorySaver - is there a comprehensive saver with a Test file to use as a guide? (LangGraph JS) #545