-
Notifications
You must be signed in to change notification settings - Fork 295
/
SimpleChatStore.ts
65 lines (55 loc) · 1.55 KB
/
SimpleChatStore.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
import type { ChatMessage } from "../../llm/index.js";
import type { BaseChatStore } from "./types.js";
/**
* fixme: User could carry object references in the messages.
* This could lead to memory leaks if the messages are not properly cleaned up.
*/
export class SimpleChatStore<
AdditionalMessageOptions extends object = Record<string, unknown>,
> implements BaseChatStore<AdditionalMessageOptions>
{
store: { [key: string]: ChatMessage<AdditionalMessageOptions>[] } = {};
public setMessages(
key: string,
messages: ChatMessage<AdditionalMessageOptions>[],
) {
this.store[key] = messages;
}
public getMessages(key: string): ChatMessage<AdditionalMessageOptions>[] {
return this.store[key] || [];
}
public addMessage(
key: string,
message: ChatMessage<AdditionalMessageOptions>,
) {
this.store[key] = this.store[key] || [];
this.store[key].push(message);
}
public deleteMessages(key: string) {
if (!(key in this.store)) {
return null;
}
const messages = this.store[key];
delete this.store[key];
return messages;
}
public deleteMessage(key: string, idx: number) {
if (!(key in this.store)) {
return null;
}
if (idx >= this.store[key].length) {
return null;
}
return this.store[key].splice(idx, 1)[0];
}
public deleteLastMessage(key: string) {
if (!(key in this.store)) {
return null;
}
const lastMessage = this.store[key].pop();
return lastMessage || null;
}
public getKeys(): string[] {
return Object.keys(this.store);
}
}