-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
ioredis.ts
135 lines (122 loc) Β· 3.61 KB
/
ioredis.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import type { Redis } from "ioredis";
import { BaseStore } from "../schema/storage.js";
/**
* Class that extends the BaseStore class to interact with a Redis
* database. It provides methods for getting, setting, and deleting data,
* as well as yielding keys from the database.
*/
export class RedisByteStore extends BaseStore<string, Uint8Array> {
lc_namespace = ["langchain", "storage"];
protected client: Redis;
protected ttl?: number;
protected namespace?: string;
protected yieldKeysScanBatchSize = 1000;
constructor(fields: {
client: Redis;
ttl?: number;
namespace?: string;
yieldKeysScanBatchSize?: number;
}) {
super(fields);
this.client = fields.client;
this.ttl = fields.ttl;
this.namespace = fields.namespace;
this.yieldKeysScanBatchSize =
fields.yieldKeysScanBatchSize ?? this.yieldKeysScanBatchSize;
}
_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 Redis database.
* @param keys Array of keys to be retrieved.
* @returns An array of retrieved values.
*/
async mget(keys: string[]) {
const prefixedKeys = keys.map(this._getPrefixedKey.bind(this));
const retrievedValues = await this.client.mgetBuffer(prefixedKeys);
return retrievedValues.map((value) => {
if (!value) {
return undefined;
} else {
return value;
}
});
}
/**
* Sets multiple keys in the 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.ttl) {
pipeline.set(key, value, "EX", this.ttl);
} else {
pipeline.set(key, value);
}
}
await pipeline.exec();
}
/**
* Deletes multiple keys from the 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 Redis database.
* @param prefix Optional prefix to filter the keys.
* @returns An AsyncGenerator that yields keys from the 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);
}
}
}
}