diff --git a/.changeset/soft-tools-wonder.md b/.changeset/soft-tools-wonder.md new file mode 100644 index 0000000000000000000000000000000000000000..3eaae4f2c5b07d6a47bcddd5081bdf1dd20f9f8b --- /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 e89e4822dfc8945f12f87a5579492689e4c9d4be..ef468e5b9d12a653c686db3f0359a537186c4e01 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 9fd0673681e5d211c5129fcd8476dd292d766c19..68b0b6e2bbde7a8ed93b637261369f90c3c8c5b1 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 be19928f26b958c35f3e51520e8d7d335add1dfc..80e6295e81f21021d33d1c517592167420b445b0 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 35f4a50bbbdba1d041fb97a15fc07d098a02dd82..92f13c59bbb7f4abc46a81df4442b2b523dcf0f8 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", ); });