new tests

Author: canrobins13

src/core/condense/__tests__/index.test.ts

@@ -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)
+	})
+})