|
|
@@ -3,12 +3,19 @@ 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"
|
|
|
+import { telemetryService } from "../../../services/telemetry/TelemetryService"
|
|
|
|
|
|
// Mock dependencies
|
|
|
jest.mock("../../../api/transform/image-cleaning", () => ({
|
|
|
maybeRemoveImageBlocks: jest.fn((messages: ApiMessage[], _apiHandler: ApiHandler) => [...messages]),
|
|
|
}))
|
|
|
|
|
|
+jest.mock("../../../services/telemetry/TelemetryService", () => ({
|
|
|
+ telemetryService: {
|
|
|
+ captureContextCondensed: jest.fn(),
|
|
|
+ },
|
|
|
+}))
|
|
|
+
|
|
|
const taskId = "test-task-id"
|
|
|
|
|
|
describe("getMessagesSinceLastSummary", () => {
|
|
|
@@ -279,3 +286,279 @@ describe("summarizeConversation", () => {
|
|
|
expect(result.summary).toBe("This is a summary with system prompt")
|
|
|
})
|
|
|
})
|
|
|
+
|
|
|
+describe("summarizeConversation with custom settings", () => {
|
|
|
+ // Mock necessary dependencies
|
|
|
+ let mockMainApiHandler: ApiHandler
|
|
|
+ let mockCondensingApiHandler: ApiHandler
|
|
|
+ const defaultSystemPrompt = "Default prompt"
|
|
|
+ const taskId = "test-task"
|
|
|
+
|
|
|
+ // Sample messages for testing
|
|
|
+ const sampleMessages: 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 },
|
|
|
+ ]
|
|
|
+
|
|
|
+ beforeEach(() => {
|
|
|
+ // Reset mocks
|
|
|
+ jest.clearAllMocks()
|
|
|
+
|
|
|
+ // Reset telemetry mock
|
|
|
+ ;(telemetryService.captureContextCondensed as jest.Mock).mockClear()
|
|
|
+
|
|
|
+ // Setup mock API handlers
|
|
|
+ mockMainApiHandler = {
|
|
|
+ createMessage: jest.fn().mockImplementation(() => {
|
|
|
+ return (async function* () {
|
|
|
+ yield { type: "text" as const, text: "Summary from main handler" }
|
|
|
+ yield { type: "usage" as const, totalCost: 0.05, outputTokens: 100 }
|
|
|
+ })()
|
|
|
+ }),
|
|
|
+ countTokens: jest.fn().mockImplementation(() => Promise.resolve(50)),
|
|
|
+ getModel: jest.fn().mockReturnValue({
|
|
|
+ id: "main-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
|
|
|
+
|
|
|
+ mockCondensingApiHandler = {
|
|
|
+ createMessage: jest.fn().mockImplementation(() => {
|
|
|
+ return (async function* () {
|
|
|
+ yield { type: "text" as const, text: "Summary from condensing handler" }
|
|
|
+ yield { type: "usage" as const, totalCost: 0.03, outputTokens: 80 }
|
|
|
+ })()
|
|
|
+ }),
|
|
|
+ countTokens: jest.fn().mockImplementation(() => Promise.resolve(40)),
|
|
|
+ getModel: jest.fn().mockReturnValue({
|
|
|
+ id: "condensing-model",
|
|
|
+ info: {
|
|
|
+ contextWindow: 4000,
|
|
|
+ supportsImages: true,
|
|
|
+ supportsComputerUse: false,
|
|
|
+ supportsVision: false,
|
|
|
+ maxTokens: 2000,
|
|
|
+ supportsPromptCache: false,
|
|
|
+ maxCachePoints: 0,
|
|
|
+ minTokensPerCachePoint: 0,
|
|
|
+ cachableFields: [],
|
|
|
+ },
|
|
|
+ }),
|
|
|
+ } as unknown as ApiHandler
|
|
|
+ })
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Test that custom prompt is used when provided
|
|
|
+ */
|
|
|
+ it("should use custom prompt when provided", async () => {
|
|
|
+ const customPrompt = "Custom summarization prompt"
|
|
|
+
|
|
|
+ await summarizeConversation(
|
|
|
+ sampleMessages,
|
|
|
+ mockMainApiHandler,
|
|
|
+ defaultSystemPrompt,
|
|
|
+ taskId,
|
|
|
+ false,
|
|
|
+ customPrompt,
|
|
|
+ )
|
|
|
+
|
|
|
+ // Verify the custom prompt was used
|
|
|
+ const createMessageCalls = (mockMainApiHandler.createMessage as jest.Mock).mock.calls
|
|
|
+ expect(createMessageCalls.length).toBe(1)
|
|
|
+ expect(createMessageCalls[0][0]).toBe(customPrompt)
|
|
|
+ })
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Test that default system prompt is used when custom prompt is empty
|
|
|
+ */
|
|
|
+ it("should use default systemPrompt when custom prompt is empty or not provided", async () => {
|
|
|
+ // Test with empty string
|
|
|
+ await summarizeConversation(
|
|
|
+ sampleMessages,
|
|
|
+ mockMainApiHandler,
|
|
|
+ defaultSystemPrompt,
|
|
|
+ taskId,
|
|
|
+ false,
|
|
|
+ " ", // Empty custom prompt
|
|
|
+ )
|
|
|
+
|
|
|
+ // Verify the default prompt was used
|
|
|
+ let createMessageCalls = (mockMainApiHandler.createMessage as jest.Mock).mock.calls
|
|
|
+ expect(createMessageCalls.length).toBe(1)
|
|
|
+ expect(createMessageCalls[0][0]).toContain("Your task is to create a detailed summary")
|
|
|
+
|
|
|
+ // Reset mock and test with undefined
|
|
|
+ jest.clearAllMocks()
|
|
|
+ await summarizeConversation(
|
|
|
+ sampleMessages,
|
|
|
+ mockMainApiHandler,
|
|
|
+ defaultSystemPrompt,
|
|
|
+ taskId,
|
|
|
+ false,
|
|
|
+ undefined, // No custom prompt
|
|
|
+ )
|
|
|
+
|
|
|
+ // Verify the default prompt was used again
|
|
|
+ createMessageCalls = (mockMainApiHandler.createMessage as jest.Mock).mock.calls
|
|
|
+ expect(createMessageCalls.length).toBe(1)
|
|
|
+ expect(createMessageCalls[0][0]).toContain("Your task is to create a detailed summary")
|
|
|
+ })
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Test that condensing API handler is used when provided and valid
|
|
|
+ */
|
|
|
+ it("should use condensingApiHandler when provided and valid", async () => {
|
|
|
+ await summarizeConversation(
|
|
|
+ sampleMessages,
|
|
|
+ mockMainApiHandler,
|
|
|
+ defaultSystemPrompt,
|
|
|
+ taskId,
|
|
|
+ false,
|
|
|
+ undefined,
|
|
|
+ mockCondensingApiHandler,
|
|
|
+ )
|
|
|
+
|
|
|
+ // Verify the condensing handler was used
|
|
|
+ expect((mockCondensingApiHandler.createMessage as jest.Mock).mock.calls.length).toBe(1)
|
|
|
+ expect((mockMainApiHandler.createMessage as jest.Mock).mock.calls.length).toBe(0)
|
|
|
+ })
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Test fallback to main API handler when condensing handler is not provided
|
|
|
+ */
|
|
|
+ it("should fall back to mainApiHandler if condensingApiHandler is not provided", async () => {
|
|
|
+ await summarizeConversation(
|
|
|
+ sampleMessages,
|
|
|
+ mockMainApiHandler,
|
|
|
+ defaultSystemPrompt,
|
|
|
+ taskId,
|
|
|
+ false,
|
|
|
+ undefined,
|
|
|
+ undefined,
|
|
|
+ )
|
|
|
+
|
|
|
+ // Verify the main handler was used
|
|
|
+ expect((mockMainApiHandler.createMessage as jest.Mock).mock.calls.length).toBe(1)
|
|
|
+ })
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Test fallback to main API handler when condensing handler is invalid
|
|
|
+ */
|
|
|
+ it("should fall back to mainApiHandler if condensingApiHandler is invalid", async () => {
|
|
|
+ // Create an invalid handler (missing createMessage)
|
|
|
+ const invalidHandler = {
|
|
|
+ countTokens: jest.fn(),
|
|
|
+ getModel: jest.fn(),
|
|
|
+ // createMessage is missing
|
|
|
+ } as unknown as ApiHandler
|
|
|
+
|
|
|
+ // Mock console.warn to verify warning message
|
|
|
+ const originalWarn = console.warn
|
|
|
+ const mockWarn = jest.fn()
|
|
|
+ console.warn = mockWarn
|
|
|
+
|
|
|
+ await summarizeConversation(
|
|
|
+ sampleMessages,
|
|
|
+ mockMainApiHandler,
|
|
|
+ defaultSystemPrompt,
|
|
|
+ taskId,
|
|
|
+ false,
|
|
|
+ undefined,
|
|
|
+ invalidHandler,
|
|
|
+ )
|
|
|
+
|
|
|
+ // Verify the main handler was used as fallback
|
|
|
+ expect((mockMainApiHandler.createMessage as jest.Mock).mock.calls.length).toBe(1)
|
|
|
+
|
|
|
+ // Verify warning was logged
|
|
|
+ expect(mockWarn).toHaveBeenCalledWith(
|
|
|
+ expect.stringContaining("Chosen API handler for condensing does not support message creation"),
|
|
|
+ )
|
|
|
+
|
|
|
+ // Restore console.warn
|
|
|
+ console.warn = originalWarn
|
|
|
+ })
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Test that telemetry is called for custom prompt usage
|
|
|
+ */
|
|
|
+ it("should capture telemetry when using custom prompt", async () => {
|
|
|
+ await summarizeConversation(
|
|
|
+ sampleMessages,
|
|
|
+ mockMainApiHandler,
|
|
|
+ defaultSystemPrompt,
|
|
|
+ taskId,
|
|
|
+ false,
|
|
|
+ "Custom prompt",
|
|
|
+ )
|
|
|
+
|
|
|
+ // Verify telemetry was called with custom prompt flag
|
|
|
+ expect(telemetryService.captureContextCondensed).toHaveBeenCalledWith(
|
|
|
+ taskId,
|
|
|
+ false,
|
|
|
+ true, // usedCustomPrompt
|
|
|
+ false, // usedCustomApiHandler
|
|
|
+ )
|
|
|
+ })
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Test that telemetry is called for custom API handler usage
|
|
|
+ */
|
|
|
+ it("should capture telemetry when using custom API handler", async () => {
|
|
|
+ await summarizeConversation(
|
|
|
+ sampleMessages,
|
|
|
+ mockMainApiHandler,
|
|
|
+ defaultSystemPrompt,
|
|
|
+ taskId,
|
|
|
+ false,
|
|
|
+ undefined,
|
|
|
+ mockCondensingApiHandler,
|
|
|
+ )
|
|
|
+
|
|
|
+ // Verify telemetry was called with custom API handler flag
|
|
|
+ expect(telemetryService.captureContextCondensed).toHaveBeenCalledWith(
|
|
|
+ taskId,
|
|
|
+ false,
|
|
|
+ false, // usedCustomPrompt
|
|
|
+ true, // usedCustomApiHandler
|
|
|
+ )
|
|
|
+ })
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Test that telemetry is called with both custom prompt and API handler
|
|
|
+ */
|
|
|
+ it("should capture telemetry when using both custom prompt and API handler", async () => {
|
|
|
+ await summarizeConversation(
|
|
|
+ sampleMessages,
|
|
|
+ mockMainApiHandler,
|
|
|
+ defaultSystemPrompt,
|
|
|
+ taskId,
|
|
|
+ true, // isAutomaticTrigger
|
|
|
+ "Custom prompt",
|
|
|
+ mockCondensingApiHandler,
|
|
|
+ )
|
|
|
+
|
|
|
+ // Verify telemetry was called with both flags
|
|
|
+ expect(telemetryService.captureContextCondensed).toHaveBeenCalledWith(
|
|
|
+ taskId,
|
|
|
+ true, // isAutomaticTrigger
|
|
|
+ true, // usedCustomPrompt
|
|
|
+ true, // usedCustomApiHandler
|
|
|
+ )
|
|
|
+ })
|
|
|
+})
|