-
Notifications
You must be signed in to change notification settings - Fork 1.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* feat: upstash redis storage BaseStore * nit: rename file * nit: rename class & input interface * nit: remove unused input * feat: tests * fix: yield & tests yield * chore: lint files * nit * chore: added entrypoint * nit: tests * fix: added to requiresOptionalDependency
- Loading branch information
1 parent
d8da36a
commit e87ee52
Showing
9 changed files
with
266 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,87 @@ | ||
| /* eslint-disable no-process-env */ | ||
|
|
||
| import { test } from "@jest/globals"; | ||
| import { Redis as UpstashRedis } from "@upstash/redis"; | ||
| import { UpstashRedisStore } from "../upstash_redis.js"; | ||
|
|
||
| const getClient = () => { | ||
| if ( | ||
| !process.env.UPSTASH_REDIS_REST_URL || | ||
| !process.env.UPSTASH_REDIS_REST_TOKEN | ||
| ) { | ||
| throw new Error("Missing Upstash Redis env variables."); | ||
| } | ||
|
|
||
| const config = { | ||
| url: process.env.UPSTASH_REDIS_REST_URL, | ||
| token: process.env.UPSTASH_REDIS_REST_TOKEN, | ||
| }; | ||
|
|
||
| return new UpstashRedis(config); | ||
| }; | ||
|
|
||
| describe("UpstashRedisStore", () => { | ||
| const keys = ["key1", "key2"]; | ||
| const client = getClient(); | ||
|
|
||
| afterEach(async () => { | ||
| await client.del(...keys); | ||
| }); | ||
|
|
||
| test("UpstashRedis can write & read values", async () => { | ||
| const encoder = new TextEncoder(); | ||
| const decoder = new TextDecoder(); | ||
| const store = new UpstashRedisStore({ | ||
| client, | ||
| }); | ||
| const value1 = new Date().toISOString(); | ||
| const value2 = new Date().toISOString() + new Date().toISOString(); | ||
| await store.mset([ | ||
| [keys[0], encoder.encode(value1)], | ||
| [keys[1], encoder.encode(value2)], | ||
| ]); | ||
| const retrievedValues = await store.mget([keys[0], keys[1]]); | ||
| const everyValueDefined = retrievedValues.every((v) => v !== undefined); | ||
| expect(everyValueDefined).toBe(true); | ||
| expect(retrievedValues.map((v) => decoder.decode(v))).toEqual([ | ||
| value1, | ||
| value2, | ||
| ]); | ||
| }); | ||
|
|
||
| test("UpstashRedis can delete values", async () => { | ||
| const encoder = new TextEncoder(); | ||
| const store = new UpstashRedisStore({ | ||
| client, | ||
| }); | ||
| const value1 = new Date().toISOString(); | ||
| const value2 = new Date().toISOString() + new Date().toISOString(); | ||
| await store.mset([ | ||
| [keys[0], encoder.encode(value1)], | ||
| [keys[1], encoder.encode(value2)], | ||
| ]); | ||
| await store.mdelete(keys); | ||
| const retrievedValues = await store.mget([keys[0], keys[1]]); | ||
| const everyValueUndefined = retrievedValues.every((v) => v === undefined); | ||
| expect(everyValueUndefined).toBe(true); | ||
| }); | ||
|
|
||
| test("UpstashRedis can yield keys with prefix", async () => { | ||
| const prefix = "prefix_"; | ||
| const keysWithPrefix = keys.map((key) => `${prefix}${key}`); | ||
| const encoder = new TextEncoder(); | ||
| const store = new UpstashRedisStore({ | ||
| client, | ||
| }); | ||
| const value = new Date().toISOString(); | ||
| await store.mset(keysWithPrefix.map((key) => [key, encoder.encode(value)])); | ||
| const yieldedKeys = []; | ||
| for await (const key of store.yieldKeys(prefix)) { | ||
| yieldedKeys.push(key); | ||
| } | ||
| console.log("Yielded keys:", yieldedKeys); | ||
| expect(yieldedKeys.sort()).toEqual(keysWithPrefix.sort()); | ||
| // afterEach won't automatically delete these since we're applying a prefix. | ||
| await store.mdelete(keysWithPrefix); | ||
| }); | ||
| }); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,157 @@ | ||
| import { Redis as UpstashRedis, type RedisConfigNodejs } from "@upstash/redis"; | ||
|
|
||
| import { BaseStore } from "../schema/storage.js"; | ||
|
|
||
| /** | ||
| * Type definition for the input parameters required to initialize an | ||
| * instance of the UpstashStoreInput class. | ||
| */ | ||
| export interface UpstashRedisStoreInput { | ||
| sessionTTL?: number; | ||
| config?: RedisConfigNodejs; | ||
| client?: UpstashRedis; | ||
| /** | ||
| * The amount of keys to retrieve per batch when yielding keys. | ||
| * @default 1000 | ||
| */ | ||
| yieldKeysScanBatchSize?: number; | ||
| /** | ||
| * The namespace to use for the keys in the database. | ||
| */ | ||
| namespace?: string; | ||
| } | ||
|
|
||
| /** | ||
| * Class that extends the BaseStore class to interact with an Upstash Redis | ||
| * database. It provides methods for getting, setting, and deleting data, | ||
| * as well as yielding keys from the database. | ||
| */ | ||
| export class UpstashRedisStore extends BaseStore<string, Uint8Array> { | ||
| lc_namespace = ["langchain", "storage"]; | ||
|
|
||
| protected client: UpstashRedis; | ||
|
|
||
| protected namespace?: string; | ||
|
|
||
| protected yieldKeysScanBatchSize = 1000; | ||
|
|
||
| private sessionTTL?: number; | ||
|
|
||
| constructor(fields: UpstashRedisStoreInput) { | ||
| super(fields); | ||
| if (fields.client) { | ||
| this.client = fields.client; | ||
| } else if (fields.config) { | ||
| this.client = new UpstashRedis(fields.config); | ||
| } else { | ||
| throw new Error( | ||
| `Upstash Redis store requires either a config object or a pre-configured client.` | ||
| ); | ||
| } | ||
| this.sessionTTL = fields.sessionTTL; | ||
| this.yieldKeysScanBatchSize = | ||
| fields.yieldKeysScanBatchSize ?? this.yieldKeysScanBatchSize; | ||
| this.namespace = fields.namespace; | ||
| } | ||
|
|
||
| _getPrefixedKey(key: string) { | ||
| if (this.namespace) { | ||
| const delimiter = "/"; | ||
| return `${this.namespace}${delimiter}${key}`; | ||
| } | ||
| return key; | ||
| } | ||
|
|
||
| _getDeprefixedKey(key: string) { | ||
| if (this.namespace) { | ||
| const delimiter = "/"; | ||
| return key.slice(this.namespace.length + delimiter.length); | ||
| } | ||
| return key; | ||
| } | ||
|
|
||
| /** | ||
| * Gets multiple keys from the Upstash Redis database. | ||
| * @param keys Array of keys to be retrieved. | ||
| * @returns An array of retrieved values. | ||
| */ | ||
| async mget(keys: string[]) { | ||
| const encoder = new TextEncoder(); | ||
|
|
||
| const prefixedKeys = keys.map(this._getPrefixedKey.bind(this)); | ||
| const retrievedValues = await this.client.mget<Uint8Array[]>( | ||
| ...prefixedKeys | ||
| ); | ||
| return retrievedValues.map((value) => { | ||
| if (!value) { | ||
| return undefined; | ||
| } else if (typeof value === "object") { | ||
| return encoder.encode(JSON.stringify(value)); | ||
| } else { | ||
| return encoder.encode(value); | ||
| } | ||
| }); | ||
| } | ||
|
|
||
| /** | ||
| * Sets multiple keys in the Upstash Redis database. | ||
| * @param keyValuePairs Array of key-value pairs to be set. | ||
| * @returns Promise that resolves when all keys have been set. | ||
| */ | ||
| async mset(keyValuePairs: [string, Uint8Array][]): Promise<void> { | ||
| const decoder = new TextDecoder(); | ||
| const encodedKeyValuePairs = keyValuePairs.map(([key, value]) => [ | ||
| this._getPrefixedKey(key), | ||
| decoder.decode(value), | ||
| ]); | ||
| const pipeline = this.client.pipeline(); | ||
| for (const [key, value] of encodedKeyValuePairs) { | ||
| if (this.sessionTTL) { | ||
| pipeline.setex(key, this.sessionTTL, value); | ||
| } else { | ||
| pipeline.set(key, value); | ||
| } | ||
| } | ||
| await pipeline.exec(); | ||
| } | ||
|
|
||
| /** | ||
| * Deletes multiple keys from the Upstash Redis database. | ||
| * @param keys Array of keys to be deleted. | ||
| * @returns Promise that resolves when all keys have been deleted. | ||
| */ | ||
| async mdelete(keys: string[]): Promise<void> { | ||
| await this.client.del(...keys.map(this._getPrefixedKey.bind(this))); | ||
| } | ||
|
|
||
| /** | ||
| * Yields keys from the Upstash Redis database. | ||
| * @param prefix Optional prefix to filter the keys. A wildcard (*) is always appended to the end. | ||
| * @returns An AsyncGenerator that yields keys from the Upstash Redis database. | ||
| */ | ||
| async *yieldKeys(prefix?: string): AsyncGenerator<string> { | ||
| let pattern; | ||
| if (prefix) { | ||
| const wildcardPrefix = prefix.endsWith("*") ? prefix : `${prefix}*`; | ||
| pattern = `${this._getPrefixedKey(wildcardPrefix)}*`; | ||
| } else { | ||
| pattern = this._getPrefixedKey("*"); | ||
| } | ||
| let [cursor, batch] = await this.client.scan(0, { | ||
| match: pattern, | ||
| count: this.yieldKeysScanBatchSize, | ||
| }); | ||
| for (const key of batch) { | ||
| yield this._getDeprefixedKey(key); | ||
| } | ||
| while (cursor !== 0) { | ||
| [cursor, batch] = await this.client.scan(cursor, { | ||
| match: pattern, | ||
| count: this.yieldKeysScanBatchSize, | ||
| }); | ||
| for (const key of batch) { | ||
| yield this._getDeprefixedKey(key); | ||
| } | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters