@ixo/sqlite-saver 1.0.4

package/src/tests/checkpointer.test.ts

@@ -0,0 +1,628 @@
+import type { RunnableConfig } from '@langchain/core/runnables';
+import {
+  Checkpoint,
+  CheckpointTuple,
+  emptyCheckpoint,
+  TASKS,
+  uuid6,
+} from '@langchain/langgraph-checkpoint';
+import { SqliteSaver } from '../index';
+
+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,
+    },
+  },
+};
+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,
+    },
+  },
+};
+
+describe('SqliteSaver', () => {
+  it('should save and retrieve checkpoints correctly', async () => {
+    const sqliteSaver = SqliteSaver.fromConnString(':memory:');
+
+    // get undefined checkpoint
+    const undefinedCheckpoint = await sqliteSaver.getTuple({
+      configurable: { thread_id: '1' },
+    });
+    expect(undefinedCheckpoint).toBeUndefined();
+
+    // save first checkpoint
+    const runnableConfig = await sqliteSaver.put(
+      { configurable: { thread_id: '1' } },
+      checkpoint1,
+      { source: 'update', step: -1, parents: {} },
+    );
+    expect(runnableConfig).toEqual({
+      configurable: {
+        thread_id: '1',
+        checkpoint_ns: '',
+        checkpoint_id: checkpoint1.id,
+      },
+    });
+
+    // add some writes
+    await sqliteSaver.putWrites(
+      {
+        configurable: {
+          checkpoint_id: checkpoint1.id,
+          checkpoint_ns: '',
+          thread_id: '1',
+        },
+      },
+      [['bar', 'baz']],
+      'foo',
+    );
+
+    // get first checkpoint tuple
+    const firstCheckpointTuple = await sqliteSaver.getTuple({
+      configurable: { thread_id: '1' },
+    });
+    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'],
+    ]);
+
+    // save second checkpoint
+    await sqliteSaver.put(
+      {
+        configurable: {
+          thread_id: '1',
+          checkpoint_id: '2024-04-18T17:19:07.952Z',
+        },
+      },
+      checkpoint2,
+      {
+        source: 'update',
+        step: -1,
+        parents: { '': checkpoint1.id },
+      },
+    );
+
+    // verify that parentTs is set and retrieved correctly for second checkpoint
+    const secondCheckpointTuple = await sqliteSaver.getTuple({
+      configurable: { thread_id: '1' },
+    });
+    expect(secondCheckpointTuple?.parentConfig).toEqual({
+      configurable: {
+        thread_id: '1',
+        checkpoint_ns: '',
+        checkpoint_id: '2024-04-18T17:19:07.952Z',
+      },
+    });
+
+    // list checkpoints
+    const checkpointTupleGenerator = await sqliteSaver.list(
+      {
+        configurable: { thread_id: '1' },
+      },
+      {
+        filter: {
+          source: 'update',
+          step: -1,
+          parents: { '': checkpoint1.id },
+        },
+      },
+    );
+    const checkpointTuples: CheckpointTuple[] = [];
+    for await (const checkpoint of checkpointTupleGenerator) {
+      checkpointTuples.push(checkpoint);
+    }
+    expect(checkpointTuples.length).toBe(1);
+
+    const checkpointTuple1 = checkpointTuples[0];
+    expect(checkpointTuple1?.checkpoint.ts).toBe('2024-04-20T17:19:07.952Z');
+  });
+
+  it('should delete thread', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+    await saver.put({ configurable: { thread_id: '1' } }, emptyCheckpoint(), {
+      source: 'update',
+      step: -1,
+      parents: {},
+    });
+
+    await saver.put({ configurable: { thread_id: '2' } }, emptyCheckpoint(), {
+      source: 'update',
+      step: -1,
+      parents: {},
+    });
+
+    await saver.deleteThread('1');
+
+    expect(
+      await saver.getTuple({ configurable: { thread_id: '1' } }),
+    ).toBeUndefined();
+
+    expect(
+      await saver.getTuple({ configurable: { thread_id: '2' } }),
+    ).toBeDefined();
+  });
+
+  it('pending sends migration', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+
+    let config: RunnableConfig = {
+      configurable: { thread_id: 'thread-1', checkpoint_ns: '' },
+    };
+
+    const checkpoint0 = emptyCheckpoint();
+
+    config = await saver.put(config, checkpoint0, {
+      source: 'loop',
+      parents: {},
+      step: 0,
+    });
+
+    await saver.putWrites(
+      config,
+      [
+        [TASKS, 'send-1'],
+        [TASKS, 'send-2'],
+      ],
+      'task-1',
+    );
+    await saver.putWrites(config, [[TASKS, 'send-3']], 'task-2');
+
+    // check that fetching checkpount 0 doesn't attach pending sends
+    // (they should be attached to the next checkpoint)
+    const tuple0 = await saver.getTuple(config);
+    expect(tuple0?.checkpoint.channel_values).toEqual({});
+    expect(tuple0?.checkpoint.channel_versions).toEqual({});
+
+    // create second checkpoint
+    const checkpoint1: Checkpoint = {
+      v: 1,
+      id: uuid6(1),
+      ts: '2024-04-20T17:19:07.952Z',
+      channel_values: {},
+      channel_versions: checkpoint0.channel_versions,
+      versions_seen: checkpoint0.versions_seen,
+    };
+    config = await saver.put(config, checkpoint1, {
+      source: 'loop',
+      parents: {},
+      step: 1,
+    });
+
+    // check that pending sends are attached to checkpoint1
+    const checkpoint1Tuple = await saver.getTuple(config);
+    expect(checkpoint1Tuple?.checkpoint.channel_values).toEqual({
+      [TASKS]: ['send-1', 'send-2', 'send-3'],
+    });
+    expect(checkpoint1Tuple?.checkpoint.channel_versions[TASKS]).toBeDefined();
+
+    // check that the list also applies the migration
+    const checkpointTupleGenerator = saver.list({
+      configurable: { thread_id: 'thread-1' },
+    });
+
+    const checkpointTuples: CheckpointTuple[] = [];
+    for await (const checkpoint of checkpointTupleGenerator) {
+      checkpointTuples.push(checkpoint);
+    }
+    expect(checkpointTuples.length).toBe(2);
+    expect(checkpointTuples[0]?.checkpoint.channel_values).toEqual({
+      [TASKS]: ['send-1', 'send-2', 'send-3'],
+    });
+    expect(
+      checkpointTuples[0]?.checkpoint.channel_versions[TASKS],
+    ).toBeDefined();
+  });
+
+  it('should filter list by checkpoint_ns', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+
+    // Create checkpoints in different namespaces
+    const checkpoint1 = emptyCheckpoint();
+    const checkpoint2 = emptyCheckpoint();
+    const checkpoint3 = emptyCheckpoint();
+
+    await saver.put(
+      { configurable: { thread_id: 'thread-1', checkpoint_ns: 'ns1' } },
+      checkpoint1,
+      { source: 'update', step: -1, parents: {} },
+    );
+
+    await saver.put(
+      { configurable: { thread_id: 'thread-1', checkpoint_ns: 'ns2' } },
+      checkpoint2,
+      { source: 'update', step: -1, parents: {} },
+    );
+
+    await saver.put(
+      { configurable: { thread_id: 'thread-1', checkpoint_ns: 'ns1' } },
+      checkpoint3,
+      { source: 'update', step: -1, parents: {} },
+    );
+
+    // List checkpoints in ns1
+    const ns1Generator = saver.list({
+      configurable: { thread_id: 'thread-1', checkpoint_ns: 'ns1' },
+    });
+
+    const ns1Tuples: CheckpointTuple[] = [];
+    for await (const checkpoint of ns1Generator) {
+      ns1Tuples.push(checkpoint);
+    }
+    expect(ns1Tuples.length).toBe(2);
+    expect(
+      ns1Tuples.every((t) => t.config.configurable?.checkpoint_ns === 'ns1'),
+    ).toBe(true);
+
+    // List checkpoints in ns2
+    const ns2Generator = saver.list({
+      configurable: { thread_id: 'thread-1', checkpoint_ns: 'ns2' },
+    });
+
+    const ns2Tuples: CheckpointTuple[] = [];
+    for await (const checkpoint of ns2Generator) {
+      ns2Tuples.push(checkpoint);
+    }
+    expect(ns2Tuples.length).toBe(1);
+    expect(ns2Tuples[0]?.config.configurable?.checkpoint_ns).toBe('ns2');
+  });
+
+  it('should paginate list with before cursor', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+
+    // Create multiple checkpoints with sequential timestamps
+    const checkpoint1: Checkpoint = {
+      v: 1,
+      id: uuid6(-2),
+      ts: '2024-04-19T17:19:07.952Z',
+      channel_values: {},
+      channel_versions: {},
+      versions_seen: {},
+    };
+    const checkpoint2: Checkpoint = {
+      v: 1,
+      id: uuid6(-1),
+      ts: '2024-04-20T17:19:07.952Z',
+      channel_values: {},
+      channel_versions: {},
+      versions_seen: {},
+    };
+    const checkpoint3: Checkpoint = {
+      v: 1,
+      id: uuid6(0),
+      ts: '2024-04-21T17:19:07.952Z',
+      channel_values: {},
+      channel_versions: {},
+      versions_seen: {},
+    };
+
+    await saver.put({ configurable: { thread_id: 'thread-1' } }, checkpoint1, {
+      source: 'update',
+      step: -1,
+      parents: {},
+    });
+    await saver.put({ configurable: { thread_id: 'thread-1' } }, checkpoint2, {
+      source: 'update',
+      step: -1,
+      parents: {},
+    });
+    const config3 = await saver.put(
+      { configurable: { thread_id: 'thread-1' } },
+      checkpoint3,
+      { source: 'update', step: -1, parents: {} },
+    );
+
+    // List all checkpoints
+    const allGenerator = saver.list({
+      configurable: { thread_id: 'thread-1' },
+    });
+    const allTuples: CheckpointTuple[] = [];
+    for await (const checkpoint of allGenerator) {
+      allTuples.push(checkpoint);
+    }
+    expect(allTuples.length).toBe(3);
+
+    // List with before cursor (should return checkpoints before checkpoint3)
+    const beforeGenerator = saver.list(
+      {
+        configurable: { thread_id: 'thread-1' },
+      },
+      {
+        before: {
+          configurable: {
+            checkpoint_id: config3.configurable?.checkpoint_id,
+          },
+        },
+      },
+    );
+
+    const beforeTuples: CheckpointTuple[] = [];
+    for await (const checkpoint of beforeGenerator) {
+      beforeTuples.push(checkpoint);
+    }
+    expect(beforeTuples.length).toBe(2);
+    expect(
+      beforeTuples.every(
+        (t) =>
+          t.config.configurable?.checkpoint_id !==
+          config3.configurable?.checkpoint_id,
+      ),
+    ).toBe(true);
+  });
+
+  it('should limit list results', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+
+    // Create 5 checkpoints
+    for (let i = 0; i < 5; i++) {
+      const checkpoint: Checkpoint = {
+        v: 1,
+        id: uuid6(i - 3),
+        ts: `2024-04-${19 + i}T17:19:07.952Z`,
+        channel_values: {},
+        channel_versions: {},
+        versions_seen: {},
+      };
+      await saver.put({ configurable: { thread_id: 'thread-1' } }, checkpoint, {
+        source: 'update',
+        step: -1,
+        parents: {},
+      });
+    }
+
+    // List without limit
+    const allGenerator = saver.list({
+      configurable: { thread_id: 'thread-1' },
+    });
+    const allTuples: CheckpointTuple[] = [];
+    for await (const checkpoint of allGenerator) {
+      allTuples.push(checkpoint);
+    }
+    expect(allTuples.length).toBe(5);
+
+    // List with limit of 2
+    const limitedGenerator = saver.list(
+      {
+        configurable: { thread_id: 'thread-1' },
+      },
+      { limit: 2 },
+    );
+
+    const limitedTuples: CheckpointTuple[] = [];
+    for await (const checkpoint of limitedGenerator) {
+      limitedTuples.push(checkpoint);
+    }
+    expect(limitedTuples.length).toBe(2);
+  });
+
+  it('should filter list by metadata fields', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+
+    const checkpoint1 = emptyCheckpoint();
+    const checkpoint2 = emptyCheckpoint();
+    const checkpoint3 = emptyCheckpoint();
+
+    await saver.put({ configurable: { thread_id: 'thread-1' } }, checkpoint1, {
+      source: 'update',
+      step: 1,
+      parents: {},
+    });
+
+    await saver.put({ configurable: { thread_id: 'thread-1' } }, checkpoint2, {
+      source: 'loop',
+      step: 2,
+      parents: {},
+    });
+
+    await saver.put({ configurable: { thread_id: 'thread-1' } }, checkpoint3, {
+      source: 'update',
+      step: 3,
+      parents: {},
+    });
+
+    // Filter by source
+    const updateGenerator = saver.list(
+      {
+        configurable: { thread_id: 'thread-1' },
+      },
+      {
+        filter: { source: 'update' },
+      },
+    );
+
+    const updateTuples: CheckpointTuple[] = [];
+    for await (const checkpoint of updateGenerator) {
+      updateTuples.push(checkpoint);
+    }
+    expect(updateTuples.length).toBe(2);
+    expect(updateTuples.every((t) => t.metadata?.source === 'update')).toBe(
+      true,
+    );
+
+    // Filter by step
+    const stepGenerator = saver.list(
+      {
+        configurable: { thread_id: 'thread-1' },
+      },
+      {
+        filter: { step: 2 },
+      },
+    );
+
+    const stepTuples: CheckpointTuple[] = [];
+    for await (const checkpoint of stepGenerator) {
+      stepTuples.push(checkpoint);
+    }
+    expect(stepTuples.length).toBe(1);
+    expect(stepTuples[0]?.metadata?.step).toBe(2);
+  });
+
+  it('should throw error when put() is called without config.configurable', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+    const checkpoint = emptyCheckpoint();
+
+    await expect(
+      saver.put({} as RunnableConfig, checkpoint, {
+        source: 'update',
+        step: -1,
+        parents: {},
+      }),
+    ).rejects.toThrow('Empty configuration supplied.');
+  });
+
+  it('should throw error when put() is called without thread_id', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+    const checkpoint = emptyCheckpoint();
+
+    await expect(
+      saver.put({ configurable: {} } as RunnableConfig, checkpoint, {
+        source: 'update',
+        step: -1,
+        parents: {},
+      }),
+    ).rejects.toThrow(
+      'Missing "thread_id" field in passed "config.configurable".',
+    );
+  });
+
+  it('should throw error when putWrites() is called without config.configurable', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+
+    await expect(
+      saver.putWrites({} as RunnableConfig, [['channel', 'value']], 'task-1'),
+    ).rejects.toThrow('Empty configuration supplied.');
+  });
+
+  it('should throw error when putWrites() is called without thread_id', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+
+    await expect(
+      saver.putWrites(
+        { configurable: {} } as RunnableConfig,
+        [['channel', 'value']],
+        'task-1',
+      ),
+    ).rejects.toThrow('Missing thread_id field in config.configurable.');
+  });
+
+  it('should throw error when putWrites() is called without checkpoint_id', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+
+    await expect(
+      saver.putWrites(
+        { configurable: { thread_id: 'thread-1' } } as RunnableConfig,
+        [['channel', 'value']],
+        'task-1',
+      ),
+    ).rejects.toThrow('Missing checkpoint_id field in config.configurable.');
+  });
+
+  it('should migrate pending sends for v<4 checkpoints in list()', async () => {
+    const saver = SqliteSaver.fromConnString(':memory:');
+
+    // Create v<4 checkpoint (v: 1)
+    const checkpoint0: Checkpoint = {
+      v: 1,
+      id: uuid6(-1),
+      ts: '2024-04-19T17:19:07.952Z',
+      channel_values: {},
+      channel_versions: {},
+      versions_seen: {},
+    };
+
+    let config = await saver.put(
+      { configurable: { thread_id: 'thread-1' } },
+      checkpoint0,
+      {
+        source: 'loop',
+        parents: {},
+        step: 0,
+      },
+    );
+
+    // Add pending sends to checkpoint0
+    await saver.putWrites(
+      config,
+      [
+        [TASKS, 'send-1'],
+        [TASKS, 'send-2'],
+      ],
+      'task-1',
+    );
+
+    // Create v<4 child checkpoint that should inherit pending sends
+    const checkpoint1: Checkpoint = {
+      v: 1,
+      id: uuid6(0),
+      ts: '2024-04-20T17:19:07.952Z',
+      channel_values: {},
+      channel_versions: {},
+      versions_seen: {},
+    };
+
+    config = await saver.put(
+      {
+        configurable: {
+          thread_id: 'thread-1',
+          checkpoint_id: checkpoint0.id,
+        },
+      },
+      checkpoint1,
+      {
+        source: 'loop',
+        parents: { '': checkpoint0.id },
+        step: 1,
+      },
+    );
+
+    // List should migrate pending sends for v<4 checkpoint
+    const listGenerator = saver.list({
+      configurable: { thread_id: 'thread-1' },
+    });
+
+    const listTuples: CheckpointTuple[] = [];
+    for await (const checkpoint of listGenerator) {
+      listTuples.push(checkpoint);
+    }
+
+    // Find checkpoint1 in the list
+    const checkpoint1Tuple = listTuples.find(
+      (t) => t.config.configurable?.checkpoint_id === checkpoint1.id,
+    );
+
+    expect(checkpoint1Tuple).toBeDefined();
+    expect(checkpoint1Tuple?.checkpoint.channel_values[TASKS]).toEqual([
+      'send-1',
+      'send-2',
+    ]);
+    expect(checkpoint1Tuple?.checkpoint.channel_versions[TASKS]).toBeDefined();
+  });
+});