-
Notifications
You must be signed in to change notification settings - Fork 2k
/
in_memory.ts
73 lines (67 loc) · 2.28 KB
/
in_memory.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
import { BaseStore } from "@langchain/core/stores";
/**
* In-memory implementation of the BaseStore using a dictionary. Used for
* storing key-value pairs in memory.
* @example
* ```typescript
* const store = new InMemoryStore<BaseMessage>();
* await store.mset(
* Array.from({ length: 5 }).map((_, index) => [
* `message:id:${index}`,
* index % 2 === 0
* ? new AIMessage("ai stuff...")
* : new HumanMessage("human stuff..."),
* ]),
* );
*
* const retrievedMessages = await store.mget(["message:id:0", "message:id:1"]);
* await store.mdelete(await store.yieldKeys("message:id:").toArray());
* ```
*/
// eslint-disable-next-line @typescript-eslint/no-explicit-any
export class InMemoryStore<T = any> extends BaseStore<string, T> {
lc_namespace = ["langchain", "storage"];
protected store: Record<string, T> = {};
/**
* Retrieves the values associated with the given keys from the store.
* @param keys Keys to retrieve values for.
* @returns Array of values associated with the given keys.
*/
async mget(keys: string[]) {
return keys.map((key) => this.store[key]);
}
/**
* Sets the values for the given keys in the store.
* @param keyValuePairs Array of key-value pairs to set in the store.
* @returns Promise that resolves when all key-value pairs have been set.
*/
async mset(keyValuePairs: [string, T][]): Promise<void> {
for (const [key, value] of keyValuePairs) {
this.store[key] = value;
}
}
/**
* Deletes the given keys and their associated values from the store.
* @param keys Keys to delete from the store.
* @returns Promise that resolves when all keys have been deleted.
*/
async mdelete(keys: string[]): Promise<void> {
for (const key of keys) {
delete this.store[key];
}
}
/**
* Asynchronous generator that yields keys from the store. If a prefix is
* provided, it only yields keys that start with the prefix.
* @param prefix Optional prefix to filter keys.
* @returns AsyncGenerator that yields keys from the store.
*/
async *yieldKeys(prefix?: string | undefined): AsyncGenerator<string> {
const keys = Object.keys(this.store);
for (const key of keys) {
if (prefix === undefined || key.startsWith(prefix)) {
yield key;
}
}
}
}