diff --git a/libs/checkpoint-vercel-kv/.env.example b/libs/checkpoint-vercel-kv/.env.example new file mode 100644 index 00000000..aea660a4 --- /dev/null +++ b/libs/checkpoint-vercel-kv/.env.example @@ -0,0 +1,6 @@ +# ------------------LangSmith tracing------------------ +LANGCHAIN_TRACING_V2=true +LANGCHAIN_ENDPOINT="https://api.smith.langchain.com" +LANGCHAIN_API_KEY= +LANGCHAIN_PROJECT= +# ----------------------------------------------------- \ No newline at end of file diff --git a/libs/checkpoint-vercel-kv/.eslintrc.cjs b/libs/checkpoint-vercel-kv/.eslintrc.cjs new file mode 100644 index 00000000..02711dad --- /dev/null +++ b/libs/checkpoint-vercel-kv/.eslintrc.cjs @@ -0,0 +1,69 @@ +module.exports = { + extends: [ + "airbnb-base", + "eslint:recommended", + "prettier", + "plugin:@typescript-eslint/recommended", + ], + parserOptions: { + ecmaVersion: 12, + parser: "@typescript-eslint/parser", + project: "./tsconfig.json", + sourceType: "module", + }, + plugins: ["@typescript-eslint", "no-instanceof", "eslint-plugin-jest"], + ignorePatterns: [ + ".eslintrc.cjs", + "scripts", + "node_modules", + "dist", + "dist-cjs", + "*.js", + "*.cjs", + "*.d.ts", + ], + rules: { + "no-process-env": 2, + "no-instanceof/no-instanceof": 2, + "@typescript-eslint/explicit-module-boundary-types": 0, + "@typescript-eslint/no-empty-function": 0, + "@typescript-eslint/no-shadow": 0, + "@typescript-eslint/no-empty-interface": 0, + "@typescript-eslint/no-use-before-define": ["error", "nofunc"], + "@typescript-eslint/no-unused-vars": ["warn", { args: "none" }], + "@typescript-eslint/no-floating-promises": "error", + "@typescript-eslint/no-misused-promises": "error", + "arrow-body-style": 0, + camelcase: 0, + "class-methods-use-this": 0, + "import/extensions": [2, "ignorePackages"], + "import/no-extraneous-dependencies": [ + "error", + { devDependencies: ["**/*.test.ts"] }, + ], + "import/no-unresolved": 0, + "import/prefer-default-export": 0, + 'jest/no-focused-tests': 'error', + "keyword-spacing": "error", + "max-classes-per-file": 0, + "max-len": 0, + "no-await-in-loop": 0, + "no-bitwise": 0, + "no-console": 0, + "no-empty-function": 0, + "no-restricted-syntax": 0, + "no-shadow": 0, + "no-continue": 0, + "no-void": 0, + "no-underscore-dangle": 0, + "no-use-before-define": 0, + "no-useless-constructor": 0, + "no-return-await": 0, + "consistent-return": 0, + "no-else-return": 0, + "func-names": 0, + "no-lonely-if": 0, + "prefer-rest-params": 0, + "new-cap": ["error", { properties: false, capIsNew: false }], + }, +}; diff --git a/libs/checkpoint-vercel-kv/.gitignore b/libs/checkpoint-vercel-kv/.gitignore new file mode 100644 index 00000000..c10034e2 --- /dev/null +++ b/libs/checkpoint-vercel-kv/.gitignore @@ -0,0 +1,7 @@ +index.cjs +index.js +index.d.ts +index.d.cts +node_modules +dist +.yarn diff --git a/libs/checkpoint-vercel-kv/.prettierrc b/libs/checkpoint-vercel-kv/.prettierrc new file mode 100644 index 00000000..ba08ff04 --- /dev/null +++ b/libs/checkpoint-vercel-kv/.prettierrc @@ -0,0 +1,19 @@ +{ + "$schema": "https://json.schemastore.org/prettierrc", + "printWidth": 80, + "tabWidth": 2, + "useTabs": false, + "semi": true, + "singleQuote": false, + "quoteProps": "as-needed", + "jsxSingleQuote": false, + "trailingComma": "es5", + "bracketSpacing": true, + "arrowParens": "always", + "requirePragma": false, + "insertPragma": false, + "proseWrap": "preserve", + "htmlWhitespaceSensitivity": "css", + "vueIndentScriptAndStyle": false, + "endOfLine": "lf" +} diff --git a/libs/checkpoint-vercel-kv/.release-it.json b/libs/checkpoint-vercel-kv/.release-it.json new file mode 100644 index 00000000..a1236e8d --- /dev/null +++ b/libs/checkpoint-vercel-kv/.release-it.json @@ -0,0 +1,13 @@ +{ + "github": { + "release": true, + "autoGenerate": true, + "tokenRef": "GITHUB_TOKEN_RELEASE" + }, + "npm": { + "publish": true, + "versionArgs": [ + "--workspaces-update=false" + ] + } +} diff --git a/libs/checkpoint-vercel-kv/LICENSE b/libs/checkpoint-vercel-kv/LICENSE new file mode 100644 index 00000000..e7530f5e --- /dev/null +++ b/libs/checkpoint-vercel-kv/LICENSE @@ -0,0 +1,21 @@ +The MIT License + +Copyright (c) 2024 LangChain + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. \ No newline at end of file diff --git a/libs/checkpoint-vercel-kv/README.md b/libs/checkpoint-vercel-kv/README.md new file mode 100644 index 00000000..213e111f --- /dev/null +++ b/libs/checkpoint-vercel-kv/README.md @@ -0,0 +1,62 @@ +# @langchain/langgraph-checkpoint-vercel-kv + +Implementation of a [LangGraph.js](https://github.com/langchain-ai/langgraphjs) CheckpointSaver that uses a Vercel KV instance. + +## Usage + +```ts +import { VercelKVSaver } from "@langchain/langgraph-checkpoint-vercel-kv"; + +const writeConfig = { + configurable: { + thread_id: "1", + } +}; +const readConfig = { + configurable: { + thread_id: "1" + } +}; + +const checkpointer = new VercelKVSaver({ + url: process.env.VERCEL_KV_URL!, + token: process.env.VERCEL_KV_TOKEN!, +}); + +const checkpoint = { + v: 1, + ts: "2024-07-31T20:14:19.804150+00:00", + id: "1ef4f797-8335-6428-8001-8a1503f9b875", + channel_values: { + my_key: "meow", + node: "node" + }, + channel_versions: { + __start__: 2, + my_key: 3, + start:node: 3, + node: 3 + }, + versions_seen: { + __input__: {}, + __start__: { + __start__: 1 + }, + node: { + start:node: 2 + } + }, + pending_sends: [], +} + +// store checkpoint +await checkpointer.put(writeConfig, checkpoint, {}, {}); + +// load checkpoint +await checkpointer.get(readConfig); + +// list checkpoints +for await (const checkpoint of checkpointer.list(readConfig)) { + console.log(checkpoint); +} +``` diff --git a/libs/checkpoint-vercel-kv/jest.config.cjs b/libs/checkpoint-vercel-kv/jest.config.cjs new file mode 100644 index 00000000..385d19f6 --- /dev/null +++ b/libs/checkpoint-vercel-kv/jest.config.cjs @@ -0,0 +1,20 @@ +/** @type {import('ts-jest').JestConfigWithTsJest} */ +module.exports = { + preset: "ts-jest/presets/default-esm", + testEnvironment: "./jest.env.cjs", + modulePathIgnorePatterns: ["dist/"], + moduleNameMapper: { + "^(\\.{1,2}/.*)\\.js$": "$1", + }, + transform: { + "^.+\\.tsx?$": ["@swc/jest"], + }, + transformIgnorePatterns: [ + "/node_modules/", + "\\.pnp\\.[^\\/]+$", + "./scripts/jest-setup-after-env.js", + ], + setupFiles: ["dotenv/config"], + testTimeout: 20_000, + passWithNoTests: true, +}; diff --git a/libs/checkpoint-vercel-kv/jest.env.cjs b/libs/checkpoint-vercel-kv/jest.env.cjs new file mode 100644 index 00000000..2ccedccb --- /dev/null +++ b/libs/checkpoint-vercel-kv/jest.env.cjs @@ -0,0 +1,12 @@ +const { TestEnvironment } = require("jest-environment-node"); + +class AdjustedTestEnvironmentToSupportFloat32Array extends TestEnvironment { + constructor(config, context) { + // Make `instanceof Float32Array` return true in tests + // to avoid https://github.com/xenova/transformers.js/issues/57 and https://github.com/jestjs/jest/issues/2549 + super(config, context); + this.global.Float32Array = Float32Array; + } +} + +module.exports = AdjustedTestEnvironmentToSupportFloat32Array; diff --git a/libs/checkpoint-vercel-kv/langchain.config.js b/libs/checkpoint-vercel-kv/langchain.config.js new file mode 100644 index 00000000..fe70c345 --- /dev/null +++ b/libs/checkpoint-vercel-kv/langchain.config.js @@ -0,0 +1,21 @@ +import { resolve, dirname } from "node:path"; +import { fileURLToPath } from "node:url"; + +/** + * @param {string} relativePath + * @returns {string} + */ +function abs(relativePath) { + return resolve(dirname(fileURLToPath(import.meta.url)), relativePath); +} + +export const config = { + internals: [/node\:/, /@langchain\/core\//, /async_hooks/], + entrypoints: { + index: "index" + }, + tsConfigPath: resolve("./tsconfig.json"), + cjsSource: "./dist-cjs", + cjsDestination: "./dist", + abs, +}; diff --git a/libs/checkpoint-vercel-kv/package.json b/libs/checkpoint-vercel-kv/package.json new file mode 100644 index 00000000..c72af421 --- /dev/null +++ b/libs/checkpoint-vercel-kv/package.json @@ -0,0 +1,92 @@ +{ + "name": "@langchain/langgraph-checkpoint-vercel-kv", + "version": "0.0.1", + "description": "LangGraph", + "type": "module", + "engines": { + "node": ">=18" + }, + "main": "./index.js", + "types": "./index.d.ts", + "repository": { + "type": "git", + "url": "git@github.com:langchain-ai/langgraphjs.git" + }, + "scripts": { + "build": "yarn turbo:command build:internal --filter=@langchain/langgraph-checkpoint-vercel-kv", + "build:internal": "yarn clean && yarn lc_build --create-entrypoints --pre --tree-shaking", + "clean": "rm -rf dist/ dist-cjs/ .turbo/", + "lint:eslint": "NODE_OPTIONS=--max-old-space-size=4096 eslint --cache --ext .ts,.js src/", + "lint:dpdm": "dpdm --exit-code circular:1 --no-warning --no-tree src/*.ts src/**/*.ts", + "lint": "yarn lint:eslint && yarn lint:dpdm", + "lint:fix": "yarn lint:eslint --fix && yarn lint:dpdm", + "prepack": "yarn build", + "test": "NODE_OPTIONS=--experimental-vm-modules jest --testPathIgnorePatterns=\\.int\\.test.ts --testTimeout 30000 --maxWorkers=50%", + "test:watch": "NODE_OPTIONS=--experimental-vm-modules jest --watch --testPathIgnorePatterns=\\.int\\.test.ts", + "test:single": "NODE_OPTIONS=--experimental-vm-modules yarn run jest --config jest.config.cjs --testTimeout 100000", + "test:int": "NODE_OPTIONS=--experimental-vm-modules jest --testPathPattern=\\.int\\.test.ts --testTimeout 100000 --maxWorkers=50%", + "format": "prettier --config .prettierrc --write \"src\"", + "format:check": "prettier --config .prettierrc --check \"src\"" + }, + "author": "LangChain", + "license": "MIT", + "dependencies": { + "@vercel/kv": "^3.0.0" + }, + "peerDependencies": { + "@langchain/core": ">=0.2.31 <0.4.0", + "@langchain/langgraph-checkpoint": "~0.0.6" + }, + "devDependencies": { + "@jest/globals": "^29.5.0", + "@langchain/langgraph-checkpoint": "workspace:*", + "@langchain/scripts": ">=0.1.3 <0.2.0", + "@swc/core": "^1.3.90", + "@swc/jest": "^0.2.29", + "@tsconfig/recommended": "^1.0.3", + "@types/better-sqlite3": "^7.6.9", + "@types/uuid": "^10", + "@typescript-eslint/eslint-plugin": "^6.12.0", + "@typescript-eslint/parser": "^6.12.0", + "dotenv": "^16.3.1", + "dpdm": "^3.12.0", + "eslint": "^8.33.0", + "eslint-config-airbnb-base": "^15.0.0", + "eslint-config-prettier": "^8.6.0", + "eslint-plugin-import": "^2.29.1", + "eslint-plugin-jest": "^28.8.0", + "eslint-plugin-no-instanceof": "^1.0.1", + "eslint-plugin-prettier": "^4.2.1", + "jest": "^29.5.0", + "jest-environment-node": "^29.6.4", + "prettier": "^2.8.3", + "release-it": "^17.6.0", + "rollup": "^4.23.0", + "ts-jest": "^29.1.0", + "tsx": "^4.7.0", + "typescript": "^4.9.5 || ^5.4.5" + }, + "publishConfig": { + "access": "public", + "registry": "https://registry.npmjs.org/" + }, + "exports": { + ".": { + "types": { + "import": "./index.d.ts", + "require": "./index.d.cts", + "default": "./index.d.ts" + }, + "import": "./index.js", + "require": "./index.cjs" + }, + "./package.json": "./package.json" + }, + "files": [ + "dist/", + "index.cjs", + "index.js", + "index.d.ts", + "index.d.cts" + ] +} diff --git a/libs/checkpoint-vercel-kv/src/index.ts b/libs/checkpoint-vercel-kv/src/index.ts new file mode 100644 index 00000000..ea4fd4dc --- /dev/null +++ b/libs/checkpoint-vercel-kv/src/index.ts @@ -0,0 +1,388 @@ +import { get } from "@vercel/edge-config"; +import { type VercelKV, createClient } from "@vercel/kv"; +import type { RunnableConfig } from "@langchain/core/runnables"; +import { + BaseCheckpointSaver, + type Checkpoint, + type CheckpointListOptions, + type CheckpointTuple, + type PendingWrite, + type SerializerProtocol, + type CheckpointMetadata, + CheckpointPendingWrite, +} from "@langchain/langgraph-checkpoint"; + +// snake_case is used to match Python implementation +interface KVRow { + parent_checkpoint_id: string; + type: string; + checkpoint: string; + metadata: string; +} + +interface KVConfig { + url: string; + token: string; +} + +interface KVPendingWrite { + type: string; + channel: string; + task_id: string; + value: string; +} + +/** + * LangGraph checkpointer that uses a Vercel KV instance as the backing store. + * + * @example + * ``` + * import { ChatOpenAI } from "@langchain/openai"; + * import { VercelKVSaver } from "@langchain/langgraph-checkpoint-vercel-kv"; + * import { createReactAgent } from "@langchain/langgraph/prebuilt"; + * + * const checkpointer = new VercelKVSaver({ + * url: "https://your-vercel-project.vercel.app", + * token: "your-vercel-token" + * }); + * + * const graph = createReactAgent({ + * tools: [getWeather], + * llm: new ChatOpenAI({ + * model: "gpt-4o-mini", + * }), + * checkpointSaver: checkpointer, + * }); + * const config = { configurable: { thread_id: "1" } }; + * + * await graph.invoke({ + * messages: [{ + * role: "user", + * content: "what's the weather in sf" + * }], + * }, config); + * ``` + */ +export class VercelKVSaver extends BaseCheckpointSaver { + private kv: VercelKV; + + constructor(config: KVConfig, serde?: SerializerProtocol) { + super(serde); + this.kv = createClient(config); + } + + /** + * Retrieves a checkpoint from the Vercel KV database based on the + * provided config. If the config contains a "checkpoint_id" key, the checkpoint with + * the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint + * for the given thread ID is retrieved. + */ + async getTuple(config: RunnableConfig): Promise { + const { + thread_id, + checkpoint_ns = "", + checkpoint_id = "", + } = config.configurable ?? {}; + + if (!thread_id) { + return undefined; + } + + const key = checkpoint_id + ? `${thread_id}:${checkpoint_ns}:${checkpoint_id}` + : `${thread_id}:${checkpoint_ns}:last`; + + const row: KVRow | null = await this.kv.get(key); + + if (!row) { + return undefined; + } + + const [cType, cDumpedValue] = this.serde.dumpsTyped(row.checkpoint); + const checkpointPromise = this.serde.loadsTyped(cType, cDumpedValue); + + const [mType, mDumpedValue] = this.serde.dumpsTyped(row.metadata); + const metadataPromise = this.serde.loadsTyped(mType, mDumpedValue); + + const [checkpoint, metadata] = await Promise.all([ + checkpointPromise as Checkpoint, + metadataPromise as CheckpointMetadata, + ]); + + // PENDING WRITES + let pendingWrites: CheckpointPendingWrite[] = []; + + const pendingWPrefix = `PENDING_WRITES:${thread_id}:${checkpoint_ns}:${checkpoint_id}`; + + const pendingWLua = ` + local prefix = ARGV[1] + local cursor = '0' + local result = {} + repeat + local scanResult = redis.call('SCAN', cursor, 'MATCH', prefix .. '*', 'COUNT', 1000) + cursor = scanResult[1] + local keys = scanResult[2] + for _, key in ipairs(keys) do + table.insert(result, key) + end + until cursor == '0' + return result + `; + + const pendingWriteKeys: string[] = await this.kv.eval( + pendingWLua, + [], + [pendingWPrefix] + ); + + if (pendingWriteKeys.length) { + const serializedWrites: (KVPendingWrite | null)[] = await this.kv.mget( + ...pendingWriteKeys + ); + pendingWrites = await Promise.all( + serializedWrites + .filter((write): write is KVPendingWrite => write !== null) + .map(async (serializedWrite) => { + const [sType, sDumpedValue] = this.serde.dumpsTyped( + serializedWrite.value + ); + const unserializedValue = await this.serde.loadsTyped( + sType, + sDumpedValue + ); + return [ + serializedWrite.task_id, + serializedWrite.channel, + unserializedValue, + ] as CheckpointPendingWrite; + }) + ); + } + return { + checkpoint, + metadata, + pendingWrites, + config: { + configurable: { + thread_id, + checkpoint_ns, + checkpoint_id: checkpoint.id, + }, + }, + parentConfig: row.parent_checkpoint_id + ? { + configurable: { + thread_id, + checkpoint_ns, + checkpoint_id: row.parent_checkpoint_id, + }, + } + : undefined, + }; + } + + async *list( + config: RunnableConfig, + options?: CheckpointListOptions + ): AsyncGenerator { + const { thread_id, checkpoint_ns = "" } = config.configurable ?? {}; + const { limit, before, filter } = options ?? {}; + + // LUA script to get keys excluding those starting with "last" + const luaScript = ` + local prefix = ARGV[1] + local cursor = '0' + local result = {} + repeat + local scanResult = redis.call('SCAN', cursor, 'MATCH', prefix .. '*', 'COUNT', 1000) + cursor = scanResult[1] + local keys = scanResult[2] + for _, key in ipairs(keys) do + if key:sub(-5) ~= ':last' then + table.insert(result, key) + end + end + until cursor == '0' + return result + `; + + // Execute the LUA script with the prefix as an argument + const prefix = `${thread_id}:${checkpoint_ns}`; + const keys: string[] = await this.kv.eval(luaScript, [], [prefix]); + + // Filter keys based on the before parameter + const filteredKeys = keys.filter((key: string) => { + const [, , checkpoint_id] = key.split(":"); + return !before || checkpoint_id < before?.configurable?.checkpoint_id; + }); + + const sortedKeys = filteredKeys.sort((a: string, b: string) => + b.localeCompare(a) + ); + + const rows: (KVRow | null)[] = await this.kv.mget(...sortedKeys); + + let limitCount = 0; + + for (const row of rows) { + if (row) { + const [cType, cDumpedValue] = this.serde.dumpsTyped(row.checkpoint); + const checkpointPromise = this.serde.loadsTyped(cType, cDumpedValue); + + const [mType, mDumpedValue] = this.serde.dumpsTyped(row.metadata); + const metadataPromise = this.serde.loadsTyped(mType, mDumpedValue); + + const [checkpoint, metadata] = await Promise.all([ + checkpointPromise as Checkpoint, + metadataPromise as CheckpointMetadata, + ]); + + // filter by metadata + if (filter && Object.keys(filter).length > 0) { + const matches = Object.keys(filter).every( + (key) => filter[key] === metadata[key as keyof CheckpointMetadata] + ); + if (!matches) { + continue; + } + } + + yield { + config: { + configurable: { + thread_id, + checkpoint_ns, + checkpoint_id: checkpoint.id, + }, + }, + checkpoint: checkpoint, + metadata: metadata, + parentConfig: row.parent_checkpoint_id + ? { + configurable: { + thread_id, + checkpoint_ns: checkpoint_ns, + checkpoint_id: row.parent_checkpoint_id, + }, + } + : undefined, + }; + if (limit && ++limitCount >= limit) { + break; + } + } + } + } + + /** + * Saves a checkpoint. The checkpoint is associated + * with the provided config and its parent config (if any). + */ + async put( + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata + ): Promise { + if (config.configurable === undefined) { + throw new Error(`Missing "configurable" field in "config" param`); + } + + const { + thread_id, + checkpoint_ns = "", + checkpoint_id = "", + } = config.configurable; + + if (!thread_id || !checkpoint.id) { + throw new Error("Thread ID and Checkpoint ID must be defined"); + } + + const [checkpointType, checkpointValue] = this.serde.dumpsTyped(checkpoint); + const [metadataType, metadataValue] = this.serde.dumpsTyped(metadata); + + if (checkpointType !== metadataType) { + throw new Error("Mismatched checkpoint and metadata types."); + } + + const row: KVRow = { + type: checkpointType, + checkpoint: JSON.parse( + new TextDecoder().decode(checkpointValue).replace(/\0/g, "") + ), + metadata: JSON.parse( + new TextDecoder().decode(metadataValue).replace(/\0/g, "") + ), + parent_checkpoint_id: checkpoint_id, + }; + + // LUA script to set checkpoint data atomically + const luaScript = ` + local prefix = ARGV[1] + local checkpoint_id = ARGV[2] + local row = ARGV[3] + + redis.call('SET', prefix .. ':' .. checkpoint_id, row) + redis.call('SET', prefix .. ':last', row) + `; + + // Save the checkpoint and the last checkpoint + const prefix = `${thread_id}:${checkpoint_ns}`; + await this.kv.eval(luaScript, [], [prefix, checkpoint.id, row]); + + return { + configurable: { + thread_id, + checkpoint_ns, + checkpoint_id: checkpoint.id, + }, + }; + } + + /** + * Saves intermediate writes associated with a checkpoint. + */ + async putWrites( + config: RunnableConfig, + writes: PendingWrite[], + taskId: string + ): Promise { + const thread_id = config.configurable?.thread_id; + const checkpoint_ns = config.configurable?.checkpoint_ns; + const checkpoint_id = config.configurable?.checkpoint_id; + if ( + thread_id === undefined || + checkpoint_ns === undefined || + checkpoint_id === undefined + ) { + throw new Error( + `The provided config must contain a configurable field with "thread_id", "checkpoint_ns" and "checkpoint_id" fields.` + ); + } + + const prefix = `PENDING_WRITES:${thread_id}:${checkpoint_ns}:${checkpoint_id}`; + + const values: Record = writes.reduce( + (acc, [channel, value], idx) => { + const key = `${prefix}:${taskId}:${idx}`; + const [type, serializedValue] = this.serde.dumpsTyped(value); + return { + ...acc, + [key]: { + type, + channel, + task_id: taskId, + value: JSON.parse( + new TextDecoder().decode(serializedValue).replace(/\0/g, "") + ), + }, + }; + }, + {} + ); + await this.kv.mset(values); + } + + async flush(): Promise { + await this.kv.flushall(); + } +} diff --git a/libs/checkpoint-vercel-kv/src/tests/checkpoints.int.test.ts b/libs/checkpoint-vercel-kv/src/tests/checkpoints.int.test.ts new file mode 100644 index 00000000..fc7c2081 --- /dev/null +++ b/libs/checkpoint-vercel-kv/src/tests/checkpoints.int.test.ts @@ -0,0 +1,147 @@ +import { describe, test, expect } from "@jest/globals"; +import { + Checkpoint, + CheckpointTuple, + uuid6, +} from "@langchain/langgraph-checkpoint"; + +import { getEnvironmentVariable } from "@langchain/core/utils/env"; +import { VercelKVSaver } 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: [], +}; + +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("VercelKVSaver", () => { + const vercelSaver = new VercelKVSaver({ + url: getEnvironmentVariable("KV_REST_API_URL")!, + token: getEnvironmentVariable("KV_REST_API_TOKEN")!, + }); + + test("should save and retrieve checkpoints correctly", async () => { + await vercelSaver.flush(); + + // get undefined checkpoint + const undefinedCheckpoint = await vercelSaver.getTuple({ + configurable: { thread_id: "1" }, + }); + expect(undefinedCheckpoint).toBeUndefined(); + + // save first checkpoint + const runnableConfig = await vercelSaver.put( + { configurable: { thread_id: "1" } }, + checkpoint1, + { source: "update", step: -1, writes: null, parents: {} } + ); + expect(runnableConfig).toEqual({ + configurable: { + thread_id: "1", + checkpoint_ns: "", + checkpoint_id: checkpoint1.id, + }, + }); + + // add writes + await vercelSaver.putWrites( + { + configurable: { + checkpoint_id: checkpoint1.id, + checkpoint_ns: "", + thread_id: "1", + }, + }, + [["bar", "baz"]], + "foo" + ); + // get first checkpoint tuple + const firstCheckpointTuple = await vercelSaver.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 vercelSaver.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 vercelSaver.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 = vercelSaver.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"); + }); +}); diff --git a/libs/checkpoint-vercel-kv/tsconfig.cjs.json b/libs/checkpoint-vercel-kv/tsconfig.cjs.json new file mode 100644 index 00000000..3b7026ea --- /dev/null +++ b/libs/checkpoint-vercel-kv/tsconfig.cjs.json @@ -0,0 +1,8 @@ +{ + "extends": "./tsconfig.json", + "compilerOptions": { + "module": "commonjs", + "declaration": false + }, + "exclude": ["node_modules", "dist", "docs", "**/tests"] +} diff --git a/libs/checkpoint-vercel-kv/tsconfig.json b/libs/checkpoint-vercel-kv/tsconfig.json new file mode 100644 index 00000000..bc85d83b --- /dev/null +++ b/libs/checkpoint-vercel-kv/tsconfig.json @@ -0,0 +1,23 @@ +{ + "extends": "@tsconfig/recommended", + "compilerOptions": { + "outDir": "../dist", + "rootDir": "./src", + "target": "ES2021", + "lib": ["ES2021", "ES2022.Object", "DOM"], + "module": "ES2020", + "moduleResolution": "nodenext", + "esModuleInterop": true, + "declaration": true, + "noImplicitReturns": true, + "noFallthroughCasesInSwitch": true, + "noUnusedLocals": true, + "noUnusedParameters": true, + "useDefineForClassFields": true, + "strictPropertyInitialization": false, + "allowJs": true, + "strict": true + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "dist", "docs"] +} diff --git a/libs/checkpoint-vercel-kv/turbo.json b/libs/checkpoint-vercel-kv/turbo.json new file mode 100644 index 00000000..d1bb60a7 --- /dev/null +++ b/libs/checkpoint-vercel-kv/turbo.json @@ -0,0 +1,11 @@ +{ + "extends": ["//"], + "tasks": { + "build": { + "outputs": ["**/dist/**"] + }, + "build:internal": { + "dependsOn": ["^build:internal"] + } + } +} diff --git a/yarn.lock b/yarn.lock index ca19e8fb..ceb8b996 100644 --- a/yarn.lock +++ b/yarn.lock @@ -1789,6 +1789,44 @@ __metadata: languageName: unknown linkType: soft +"@langchain/langgraph-checkpoint-vercel-kv@workspace:libs/checkpoint-vercel-kv": + version: 0.0.0-use.local + resolution: "@langchain/langgraph-checkpoint-vercel-kv@workspace:libs/checkpoint-vercel-kv" + dependencies: + "@jest/globals": ^29.5.0 + "@langchain/langgraph-checkpoint": "workspace:*" + "@langchain/scripts": ">=0.1.3 <0.2.0" + "@swc/core": ^1.3.90 + "@swc/jest": ^0.2.29 + "@tsconfig/recommended": ^1.0.3 + "@types/better-sqlite3": ^7.6.9 + "@types/uuid": ^10 + "@typescript-eslint/eslint-plugin": ^6.12.0 + "@typescript-eslint/parser": ^6.12.0 + "@vercel/kv": ^3.0.0 + dotenv: ^16.3.1 + dpdm: ^3.12.0 + eslint: ^8.33.0 + eslint-config-airbnb-base: ^15.0.0 + eslint-config-prettier: ^8.6.0 + eslint-plugin-import: ^2.29.1 + eslint-plugin-jest: ^28.8.0 + eslint-plugin-no-instanceof: ^1.0.1 + eslint-plugin-prettier: ^4.2.1 + jest: ^29.5.0 + jest-environment-node: ^29.6.4 + prettier: ^2.8.3 + release-it: ^17.6.0 + rollup: ^4.23.0 + ts-jest: ^29.1.0 + tsx: ^4.7.0 + typescript: ^4.9.5 || ^5.4.5 + peerDependencies: + "@langchain/core": ">=0.2.31 <0.4.0" + "@langchain/langgraph-checkpoint": ~0.0.6 + languageName: unknown + linkType: soft + "@langchain/langgraph-checkpoint@workspace:*, @langchain/langgraph-checkpoint@workspace:libs/checkpoint, @langchain/langgraph-checkpoint@~0.0.10": version: 0.0.0-use.local resolution: "@langchain/langgraph-checkpoint@workspace:libs/checkpoint" @@ -3738,6 +3776,24 @@ __metadata: languageName: node linkType: hard +"@upstash/redis@npm:^1.34.0": + version: 1.34.3 + resolution: "@upstash/redis@npm:1.34.3" + dependencies: + crypto-js: ^4.2.0 + checksum: 4c738eb82e3d8906999345379ff0c96913627cdf96ec54b6880fa8fb5616331a92228b889c076b4a93bee4bb31324f5afa68aa927fb751d308fc83464de68e79 + languageName: node + linkType: hard + +"@vercel/kv@npm:^3.0.0": + version: 3.0.0 + resolution: "@vercel/kv@npm:3.0.0" + dependencies: + "@upstash/redis": ^1.34.0 + checksum: 5bebab15c770e32409c03054dbd58e3e5a4d77d1a813c91e6613a9382a96b391878b67a1e2fbe8f7e197899a2dc589212ee6f6b11ab272798954e63816c6eba1 + languageName: node + linkType: hard + "@xenova/transformers@npm:2.17.2, @xenova/transformers@npm:^2.17.2": version: 2.17.2 resolution: "@xenova/transformers@npm:2.17.2" @@ -5011,6 +5067,13 @@ __metadata: languageName: node linkType: hard +"crypto-js@npm:^4.2.0": + version: 4.2.0 + resolution: "crypto-js@npm:4.2.0" + checksum: f051666dbc077c8324777f44fbd3aaea2986f198fe85092535130d17026c7c2ccf2d23ee5b29b36f7a4a07312db2fae23c9094b644cc35f7858b1b4fcaf27774 + languageName: node + linkType: hard + "crypto-random-string@npm:^4.0.0": version: 4.0.0 resolution: "crypto-random-string@npm:4.0.0"