From f865c984d3d18445f65650360b3fb8001c163131 Mon Sep 17 00:00:00 2001
From: Alex Yang <himself65@outlook.com>
Date: Tue, 12 Nov 2024 10:59:44 -0800
Subject: [PATCH] feat: async get message on chat store (#1470)
---
.changeset/soft-tools-wonder.md | 5 +++++
packages/core/src/memory/base.ts | 4 +++-
packages/core/src/memory/chat-memory-buffer.ts | 4 ++--
.../core/src/storage/chat-store/base-chat-store.ts | 6 +++++-
.../core/tests/memory/chat-memory-buffer.test.ts | 14 +++++++-------
5 files changed, 22 insertions(+), 11 deletions(-)
create mode 100644 .changeset/soft-tools-wonder.md
diff --git a/.changeset/soft-tools-wonder.md b/.changeset/soft-tools-wonder.md
new file mode 100644
index 000000000..3eaae4f2c
--- /dev/null
+++ b/.changeset/soft-tools-wonder.md
@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+feat: async get message on chat store
diff --git a/packages/core/src/memory/base.ts b/packages/core/src/memory/base.ts
index e89e4822d..ef468e5b9 100644
--- a/packages/core/src/memory/base.ts
+++ b/packages/core/src/memory/base.ts
@@ -65,7 +65,9 @@ export abstract class BaseChatStoreMemory<
super();
}
- getAllMessages(): ChatMessage<AdditionalMessageOptions>[] {
+ getAllMessages():
+ | ChatMessage<AdditionalMessageOptions>[]
+ | Promise<ChatMessage<AdditionalMessageOptions>[]> {
return this.chatStore.getMessages(this.chatStoreKey);
}
diff --git a/packages/core/src/memory/chat-memory-buffer.ts b/packages/core/src/memory/chat-memory-buffer.ts
index 9fd067368..68b0b6e2b 100644
--- a/packages/core/src/memory/chat-memory-buffer.ts
+++ b/packages/core/src/memory/chat-memory-buffer.ts
@@ -33,11 +33,11 @@ export class ChatMemoryBuffer<
}
}
- getMessages(
+ async getMessages(
transientMessages?: ChatMessage<AdditionalMessageOptions>[] | undefined,
initialTokenCount: number = 0,
) {
- const messages = this.getAllMessages();
+ const messages = await this.getAllMessages();
if (initialTokenCount > this.tokenLimit) {
throw new Error("Initial token count exceeds token limit");
diff --git a/packages/core/src/storage/chat-store/base-chat-store.ts b/packages/core/src/storage/chat-store/base-chat-store.ts
index be19928f2..80e6295e8 100644
--- a/packages/core/src/storage/chat-store/base-chat-store.ts
+++ b/packages/core/src/storage/chat-store/base-chat-store.ts
@@ -7,7 +7,11 @@ export abstract class BaseChatStore<
key: string,
messages: ChatMessage<AdditionalMessageOptions>[],
): void;
- abstract getMessages(key: string): ChatMessage<AdditionalMessageOptions>[];
+ abstract getMessages(
+ key: string,
+ ):
+ | ChatMessage<AdditionalMessageOptions>[]
+ | Promise<ChatMessage<AdditionalMessageOptions>[]>;
abstract addMessage(
key: string,
message: ChatMessage<AdditionalMessageOptions>,
diff --git a/packages/core/tests/memory/chat-memory-buffer.test.ts b/packages/core/tests/memory/chat-memory-buffer.test.ts
index 35f4a50bb..92f13c59b 100644
--- a/packages/core/tests/memory/chat-memory-buffer.test.ts
+++ b/packages/core/tests/memory/chat-memory-buffer.test.ts
@@ -19,7 +19,7 @@ describe("ChatMemoryBuffer", () => {
expect(buffer.tokenLimit).toBe(500);
});
- test("getMessages returns all messages when under token limit", () => {
+ test("getMessages returns all messages when under token limit", async () => {
const messages: ChatMessage[] = [
{ role: "user", content: "Hello" },
{ role: "assistant", content: "Hi there!" },
@@ -30,11 +30,11 @@ describe("ChatMemoryBuffer", () => {
chatHistory: messages,
});
- const result = buffer.getMessages();
+ const result = await buffer.getMessages();
expect(result).toEqual(messages);
});
- test("getMessages truncates messages when over token limit", () => {
+ test("getMessages truncates messages when over token limit", async () => {
const messages: ChatMessage[] = [
{ role: "user", content: "This is a long message" },
{ role: "assistant", content: "This is also a long reply" },
@@ -45,11 +45,11 @@ describe("ChatMemoryBuffer", () => {
chatHistory: messages,
});
- const result = buffer.getMessages();
+ const result = await buffer.getMessages();
expect(result).toEqual([{ role: "user", content: "Short" }]);
});
- test("getMessages handles input messages", () => {
+ test("getMessages handles input messages", async () => {
const storedMessages: ChatMessage[] = [
{ role: "user", content: "Hello" },
{ role: "assistant", content: "Hi there!" },
@@ -62,13 +62,13 @@ describe("ChatMemoryBuffer", () => {
const inputMessages: ChatMessage[] = [
{ role: "user", content: "New message" },
];
- const result = buffer.getMessages(inputMessages);
+ const result = await buffer.getMessages(inputMessages);
expect(result).toEqual([...inputMessages, ...storedMessages]);
});
test("getMessages throws error when initial token count exceeds limit", () => {
const buffer = new ChatMemoryBuffer({ tokenLimit: 10 });
- expect(() => buffer.getMessages(undefined, 20)).toThrow(
+ expect(async () => buffer.getMessages(undefined, 20)).rejects.toThrow(
"Initial token count exceeds token limit",
);
});
--
GitLab