|
|
@@ -0,0 +1,228 @@
|
|
|
+import { describe, expect, it, jest, beforeEach } from "@jest/globals"
|
|
|
+import { ApiHandler } from "../../../api"
|
|
|
+import { ApiMessage } from "../../task-persistence/apiMessages"
|
|
|
+import { maybeRemoveImageBlocks } from "../../../api/transform/image-cleaning"
|
|
|
+import { summarizeConversation, getMessagesSinceLastSummary, N_MESSAGES_TO_KEEP } from "../index"
|
|
|
+
|
|
|
+// Mock dependencies
|
|
|
+jest.mock("../../../api/transform/image-cleaning", () => ({
|
|
|
+ maybeRemoveImageBlocks: jest.fn((messages: ApiMessage[], _apiHandler: ApiHandler) => [...messages]),
|
|
|
+}))
|
|
|
+
|
|
|
+describe("getMessagesSinceLastSummary", () => {
|
|
|
+ it("should return all messages when there is no summary", () => {
|
|
|
+ const messages: ApiMessage[] = [
|
|
|
+ { role: "user", content: "Hello", ts: 1 },
|
|
|
+ { role: "assistant", content: "Hi there", ts: 2 },
|
|
|
+ { role: "user", content: "How are you?", ts: 3 },
|
|
|
+ ]
|
|
|
+
|
|
|
+ const result = getMessagesSinceLastSummary(messages)
|
|
|
+ expect(result).toEqual(messages)
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should return messages since the last summary", () => {
|
|
|
+ const messages: ApiMessage[] = [
|
|
|
+ { role: "user", content: "Hello", ts: 1 },
|
|
|
+ { role: "assistant", content: "Hi there", ts: 2 },
|
|
|
+ { role: "assistant", content: "Summary of conversation", ts: 3, isSummary: true },
|
|
|
+ { role: "user", content: "How are you?", ts: 4 },
|
|
|
+ { role: "assistant", content: "I'm good", ts: 5 },
|
|
|
+ ]
|
|
|
+
|
|
|
+ const result = getMessagesSinceLastSummary(messages)
|
|
|
+ expect(result).toEqual([
|
|
|
+ { role: "assistant", content: "Summary of conversation", ts: 3, isSummary: true },
|
|
|
+ { role: "user", content: "How are you?", ts: 4 },
|
|
|
+ { role: "assistant", content: "I'm good", ts: 5 },
|
|
|
+ ])
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should handle multiple summary messages and return since the last one", () => {
|
|
|
+ const messages: ApiMessage[] = [
|
|
|
+ { role: "user", content: "Hello", ts: 1 },
|
|
|
+ { role: "assistant", content: "First summary", ts: 2, isSummary: true },
|
|
|
+ { role: "user", content: "How are you?", ts: 3 },
|
|
|
+ { role: "assistant", content: "Second summary", ts: 4, isSummary: true },
|
|
|
+ { role: "user", content: "What's new?", ts: 5 },
|
|
|
+ ]
|
|
|
+
|
|
|
+ const result = getMessagesSinceLastSummary(messages)
|
|
|
+ expect(result).toEqual([
|
|
|
+ { role: "assistant", content: "Second summary", ts: 4, isSummary: true },
|
|
|
+ { role: "user", content: "What's new?", ts: 5 },
|
|
|
+ ])
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should handle empty messages array", () => {
|
|
|
+ const result = getMessagesSinceLastSummary([])
|
|
|
+ expect(result).toEqual([])
|
|
|
+ })
|
|
|
+})
|
|
|
+
|
|
|
+describe("summarizeConversation", () => {
|
|
|
+ // Mock ApiHandler
|
|
|
+ let mockApiHandler: ApiHandler
|
|
|
+ let mockStream: AsyncGenerator<any, void, unknown>
|
|
|
+
|
|
|
+ beforeEach(() => {
|
|
|
+ // Reset mocks
|
|
|
+ jest.clearAllMocks()
|
|
|
+
|
|
|
+ // Setup mock stream
|
|
|
+ mockStream = (async function* () {
|
|
|
+ yield { type: "text" as const, text: "This is " }
|
|
|
+ yield { type: "text" as const, text: "a summary" }
|
|
|
+ })()
|
|
|
+
|
|
|
+ // Setup mock API handler
|
|
|
+ mockApiHandler = {
|
|
|
+ createMessage: jest.fn().mockReturnValue(mockStream),
|
|
|
+ countTokens: jest.fn().mockImplementation(() => Promise.resolve(100)),
|
|
|
+ getModel: jest.fn().mockReturnValue({
|
|
|
+ id: "test-model",
|
|
|
+ info: {
|
|
|
+ contextWindow: 8000,
|
|
|
+ supportsImages: true,
|
|
|
+ supportsComputerUse: true,
|
|
|
+ supportsVision: true,
|
|
|
+ maxTokens: 4000,
|
|
|
+ supportsPromptCache: true,
|
|
|
+ maxCachePoints: 10,
|
|
|
+ minTokensPerCachePoint: 100,
|
|
|
+ cachableFields: ["system", "messages"],
|
|
|
+ },
|
|
|
+ }),
|
|
|
+ } as unknown as ApiHandler
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should not summarize when there are not enough messages", async () => {
|
|
|
+ const messages: ApiMessage[] = [
|
|
|
+ { role: "user", content: "Hello", ts: 1 },
|
|
|
+ { role: "assistant", content: "Hi there", ts: 2 },
|
|
|
+ ]
|
|
|
+
|
|
|
+ const result = await summarizeConversation(messages, mockApiHandler)
|
|
|
+ expect(result).toEqual(messages)
|
|
|
+ expect(mockApiHandler.createMessage).not.toHaveBeenCalled()
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should not summarize when there was a recent summary", async () => {
|
|
|
+ const messages: ApiMessage[] = [
|
|
|
+ { role: "user", content: "Hello", ts: 1 },
|
|
|
+ { role: "assistant", content: "Hi there", ts: 2 },
|
|
|
+ { role: "user", content: "How are you?", ts: 3 },
|
|
|
+ { role: "assistant", content: "I'm good", ts: 4 },
|
|
|
+ { role: "user", content: "What's new?", ts: 5 },
|
|
|
+ { role: "assistant", content: "Not much", ts: 6, isSummary: true }, // Recent summary
|
|
|
+ { role: "user", content: "Tell me more", ts: 7 },
|
|
|
+ ]
|
|
|
+
|
|
|
+ const result = await summarizeConversation(messages, mockApiHandler)
|
|
|
+ expect(result).toEqual(messages)
|
|
|
+ expect(mockApiHandler.createMessage).not.toHaveBeenCalled()
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should summarize conversation and insert summary message", async () => {
|
|
|
+ const messages: ApiMessage[] = [
|
|
|
+ { role: "user", content: "Hello", ts: 1 },
|
|
|
+ { role: "assistant", content: "Hi there", ts: 2 },
|
|
|
+ { role: "user", content: "How are you?", ts: 3 },
|
|
|
+ { role: "assistant", content: "I'm good", ts: 4 },
|
|
|
+ { role: "user", content: "What's new?", ts: 5 },
|
|
|
+ { role: "assistant", content: "Not much", ts: 6 },
|
|
|
+ { role: "user", content: "Tell me more", ts: 7 },
|
|
|
+ ]
|
|
|
+
|
|
|
+ const result = await summarizeConversation(messages, mockApiHandler)
|
|
|
+
|
|
|
+ // Check that the API was called correctly
|
|
|
+ expect(mockApiHandler.createMessage).toHaveBeenCalled()
|
|
|
+ expect(maybeRemoveImageBlocks).toHaveBeenCalled()
|
|
|
+
|
|
|
+ // Verify the structure of the result
|
|
|
+ // The result should be: original messages (except last N) + summary + last N messages
|
|
|
+ expect(result.length).toBe(messages.length + 1) // Original + summary
|
|
|
+
|
|
|
+ // Check that the summary message was inserted correctly
|
|
|
+ const summaryMessage = result[result.length - N_MESSAGES_TO_KEEP - 1]
|
|
|
+ expect(summaryMessage.role).toBe("assistant")
|
|
|
+ expect(summaryMessage.content).toBe("This is a summary")
|
|
|
+ expect(summaryMessage.isSummary).toBe(true)
|
|
|
+
|
|
|
+ // Check that the last N_MESSAGES_TO_KEEP messages are preserved
|
|
|
+ const lastMessages = messages.slice(-N_MESSAGES_TO_KEEP)
|
|
|
+ expect(result.slice(-N_MESSAGES_TO_KEEP)).toEqual(lastMessages)
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should handle empty summary response", async () => {
|
|
|
+ // We need enough messages to trigger summarization
|
|
|
+ const messages: ApiMessage[] = [
|
|
|
+ { role: "user", content: "Hello", ts: 1 },
|
|
|
+ { role: "assistant", content: "Hi there", ts: 2 },
|
|
|
+ { role: "user", content: "How are you?", ts: 3 },
|
|
|
+ { role: "assistant", content: "I'm good", ts: 4 },
|
|
|
+ { role: "user", content: "What's new?", ts: 5 },
|
|
|
+ { role: "assistant", content: "Not much", ts: 6 },
|
|
|
+ { role: "user", content: "Tell me more", ts: 7 },
|
|
|
+ ]
|
|
|
+
|
|
|
+ // Mock console.warn before we call the function
|
|
|
+ const originalWarn = console.warn
|
|
|
+ const mockWarn = jest.fn()
|
|
|
+ console.warn = mockWarn
|
|
|
+
|
|
|
+ // Setup empty summary response
|
|
|
+ const emptyStream = (async function* () {
|
|
|
+ yield { type: "text" as const, text: "" }
|
|
|
+ })()
|
|
|
+
|
|
|
+ // Create a new mock for createMessage that returns empty stream
|
|
|
+ const createMessageMock = jest.fn().mockReturnValue(emptyStream)
|
|
|
+ mockApiHandler.createMessage = createMessageMock as any
|
|
|
+
|
|
|
+ // We need to mock maybeRemoveImageBlocks to return the expected messages
|
|
|
+ ;(maybeRemoveImageBlocks as jest.Mock).mockImplementationOnce((messages: any) => {
|
|
|
+ return messages.map(({ role, content }: { role: string; content: any }) => ({ role, content }))
|
|
|
+ })
|
|
|
+
|
|
|
+ const result = await summarizeConversation(messages, mockApiHandler)
|
|
|
+
|
|
|
+ // Should return original messages when summary is empty
|
|
|
+ expect(result).toEqual(messages)
|
|
|
+ expect(mockWarn).toHaveBeenCalledWith("Received empty summary from API")
|
|
|
+
|
|
|
+ // Restore console.warn
|
|
|
+ console.warn = originalWarn
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should correctly format the request to the API", async () => {
|
|
|
+ const messages: ApiMessage[] = [
|
|
|
+ { role: "user", content: "Hello", ts: 1 },
|
|
|
+ { role: "assistant", content: "Hi there", ts: 2 },
|
|
|
+ { role: "user", content: "How are you?", ts: 3 },
|
|
|
+ { role: "assistant", content: "I'm good", ts: 4 },
|
|
|
+ { role: "user", content: "What's new?", ts: 5 },
|
|
|
+ { role: "assistant", content: "Not much", ts: 6 },
|
|
|
+ { role: "user", content: "Tell me more", ts: 7 },
|
|
|
+ ]
|
|
|
+
|
|
|
+ await summarizeConversation(messages, mockApiHandler)
|
|
|
+
|
|
|
+ // Verify the final request message
|
|
|
+ const expectedFinalMessage = {
|
|
|
+ role: "user",
|
|
|
+ content: "Summarize the conversation so far, as described in the prompt instructions.",
|
|
|
+ }
|
|
|
+
|
|
|
+ // Verify that createMessage was called with the correct prompt
|
|
|
+ expect(mockApiHandler.createMessage).toHaveBeenCalledWith(
|
|
|
+ expect.stringContaining("Your task is to create a detailed summary of the conversation"),
|
|
|
+ expect.any(Array),
|
|
|
+ )
|
|
|
+
|
|
|
+ // Check that maybeRemoveImageBlocks was called with the correct messages
|
|
|
+ const mockCallArgs = (maybeRemoveImageBlocks as jest.Mock).mock.calls[0][0] as any[]
|
|
|
+ expect(mockCallArgs[mockCallArgs.length - 1]).toEqual(expectedFinalMessage)
|
|
|
+ })
|
|
|
+})
|