Skip to content

Commit 2b453ee

Browse files
canrobins13canrobins13
committed
add a test for condensing
1 parent 547d4d7 commit 2b453ee

File tree

2 files changed

+278
-2
lines changed

2 files changed

+278
-2
lines changed

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

Lines changed: 276 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,276 @@
1+
import { jest } from "@jest/globals"
2+
import { ApiHandler } from "../../../api"
3+
import { ApiMessage } from "../../task-persistence/apiMessages"
4+
import { maybeRemoveImageBlocks } from "../../../api/transform/image-cleaning"
5+
import { summarizeConversationIfNeeded, getMessagesSinceLastSummary } from "../index"
6+
import { ApiStream, ApiStreamChunk } from "../../../api/transform/stream"
7+
import { CONTEXT_FRAC_FOR_SUMMARY, N_MESSAGES_TO_KEEP } from "../index"
8+
9+
const CONTEXT_WINDOW_SIZE = 1000
10+
const OVER_THRESHOLD_TOTAL_TOKENS = Math.ceil(CONTEXT_WINDOW_SIZE * CONTEXT_FRAC_FOR_SUMMARY) + 1
11+
12+
// Mock dependencies
13+
jest.mock("../../../api/transform/image-cleaning", () => ({
14+
maybeRemoveImageBlocks: jest.fn((messages) => messages),
15+
}))
16+
17+
// Mock Anthropic SDK
18+
jest.mock("@anthropic-ai/sdk", () => {
19+
return {
20+
default: jest.fn().mockImplementation(() => ({
21+
messages: {
22+
create: jest.fn(),
23+
},
24+
})),
25+
}
26+
})
27+
28+
describe("Conversation Condensing", () => {
29+
// Mock API handler
30+
const mockApiHandler: jest.Mocked<ApiHandler> = {
31+
createMessage: jest.fn(),
32+
} as unknown as jest.Mocked<ApiHandler>
33+
34+
// Reset mocks before each test
35+
beforeEach(() => {
36+
jest.clearAllMocks()
37+
38+
// Setup default mock for createMessage
39+
mockApiHandler.createMessage.mockImplementation((): ApiStream => {
40+
return (async function* (): AsyncGenerator<ApiStreamChunk> {
41+
yield { type: "text", text: "This is a summary of the conversation." }
42+
})()
43+
})
44+
})
45+
46+
describe("getMessagesSinceLastSummary", () => {
47+
it("should return all messages if there is no summary", () => {
48+
const messages: ApiMessage[] = [
49+
{ role: "user", content: "Hello", ts: 1 },
50+
{ role: "assistant", content: "Hi there", ts: 2 },
51+
]
52+
53+
const result = getMessagesSinceLastSummary(messages)
54+
expect(result).toEqual(messages)
55+
})
56+
57+
it("should return messages since the last summary", () => {
58+
const messages: ApiMessage[] = [
59+
{ role: "user", content: "First message", ts: 1 },
60+
{ role: "assistant", content: "Summary of conversation", ts: 2, isSummary: true },
61+
{ role: "user", content: "New message", ts: 3 },
62+
{ role: "assistant", content: "Response", ts: 4 },
63+
]
64+
65+
const result = getMessagesSinceLastSummary(messages)
66+
expect(result).toEqual([
67+
{ role: "assistant", content: "Summary of conversation", ts: 2, isSummary: true },
68+
{ role: "user", content: "New message", ts: 3 },
69+
{ role: "assistant", content: "Response", ts: 4 },
70+
])
71+
})
72+
73+
it("should handle multiple summary messages and return since the last one", () => {
74+
const messages: ApiMessage[] = [
75+
{ role: "user", content: "First message", ts: 1 },
76+
{ role: "assistant", content: "First summary", ts: 2, isSummary: true },
77+
{ role: "user", content: "Second message", ts: 3 },
78+
{ role: "assistant", content: "Second summary", ts: 4, isSummary: true },
79+
{ role: "user", content: "New message", ts: 5 },
80+
]
81+
82+
const result = getMessagesSinceLastSummary(messages)
83+
expect(result).toEqual([
84+
{ role: "assistant", content: "Second summary", ts: 4, isSummary: true },
85+
{ role: "user", content: "New message", ts: 5 },
86+
])
87+
})
88+
89+
it("should handle empty message array", () => {
90+
const messages: ApiMessage[] = []
91+
const result = getMessagesSinceLastSummary(messages)
92+
expect(result).toEqual([])
93+
})
94+
})
95+
96+
describe("summarizeConversationIfNeeded", () => {
97+
it("should not summarize when below token threshold", async () => {
98+
const messages: ApiMessage[] = [
99+
{ role: "user", content: "Hello", ts: 1 },
100+
{ role: "assistant", content: "Hi there", ts: 2 },
101+
]
102+
103+
const totalTokens = 100
104+
const result = await summarizeConversationIfNeeded(
105+
messages,
106+
totalTokens,
107+
CONTEXT_WINDOW_SIZE,
108+
mockApiHandler,
109+
)
110+
111+
expect(result).toBe(messages)
112+
expect(mockApiHandler.createMessage).not.toHaveBeenCalled()
113+
})
114+
115+
it("should summarize when above token threshold", async () => {
116+
const messages: ApiMessage[] = [
117+
{ role: "user", content: "Message 1", ts: 1 },
118+
{ role: "assistant", content: "Response 1", ts: 2 },
119+
{ role: "user", content: "Message 2", ts: 3 },
120+
{ role: "assistant", content: "Response 2", ts: 4 },
121+
{ role: "user", content: "Message 3", ts: 5 },
122+
{ role: "assistant", content: "Response 3", ts: 6 },
123+
{ role: "user", content: "Message 4", ts: 7 },
124+
]
125+
126+
const result = await summarizeConversationIfNeeded(
127+
messages,
128+
OVER_THRESHOLD_TOTAL_TOKENS,
129+
CONTEXT_WINDOW_SIZE,
130+
mockApiHandler,
131+
)
132+
133+
// Should have called createMessage
134+
expect(mockApiHandler.createMessage).toHaveBeenCalled()
135+
136+
// Should have a summary message inserted
137+
expect(result.some((msg) => msg.isSummary)).toBe(true)
138+
139+
// Should preserve the last N_MESSAGES_TO_KEEP messages
140+
for (let i = 1; i <= N_MESSAGES_TO_KEEP; i++) {
141+
expect(result).toContainEqual(messages[messages.length - i])
142+
}
143+
})
144+
145+
it("should not summarize if there are not enough messages", async () => {
146+
const messages: ApiMessage[] = [{ role: "user", content: "Hello", ts: 1 }]
147+
148+
const result = await summarizeConversationIfNeeded(
149+
messages,
150+
OVER_THRESHOLD_TOTAL_TOKENS,
151+
CONTEXT_WINDOW_SIZE,
152+
mockApiHandler,
153+
)
154+
155+
expect(result).toBe(messages)
156+
expect(mockApiHandler.createMessage).not.toHaveBeenCalled()
157+
})
158+
159+
it("should not summarize if we recently summarized", async () => {
160+
const messages: ApiMessage[] = [
161+
{ role: "user", content: "Message 1", ts: 1 },
162+
{ role: "assistant", content: "Response 1", ts: 2 },
163+
{ role: "user", content: "Message 2", ts: 3 },
164+
{ role: "assistant", content: "Summary", ts: 4, isSummary: true },
165+
{ role: "user", content: "Message 3", ts: 5 },
166+
]
167+
168+
const result = await summarizeConversationIfNeeded(
169+
messages,
170+
OVER_THRESHOLD_TOTAL_TOKENS,
171+
CONTEXT_WINDOW_SIZE,
172+
mockApiHandler,
173+
)
174+
175+
// Should not have called createMessage because one of the last 3 messages is already a summary
176+
expect(mockApiHandler.createMessage).not.toHaveBeenCalled()
177+
expect(result).toBe(messages)
178+
})
179+
180+
it("should handle empty API response", async () => {
181+
// Setup mock to return empty summary
182+
mockApiHandler.createMessage.mockImplementation((): ApiStream => {
183+
return (async function* (): AsyncGenerator<ApiStreamChunk> {
184+
yield { type: "text", text: "" }
185+
})()
186+
})
187+
188+
const messages: ApiMessage[] = [
189+
{ role: "user", content: "Message 1", ts: 1 },
190+
{ role: "assistant", content: "Response 1", ts: 2 },
191+
{ role: "user", content: "Message 2", ts: 3 },
192+
{ role: "assistant", content: "Response 2", ts: 4 },
193+
{ role: "user", content: "Message 3", ts: 5 },
194+
{ role: "assistant", content: "Response 3", ts: 6 },
195+
{ role: "user", content: "Message 4", ts: 7 },
196+
]
197+
198+
const consoleSpy = jest.spyOn(console, "warn").mockImplementation(() => {})
199+
200+
const result = await summarizeConversationIfNeeded(
201+
messages,
202+
OVER_THRESHOLD_TOTAL_TOKENS,
203+
CONTEXT_WINDOW_SIZE,
204+
mockApiHandler,
205+
)
206+
207+
// Should have called createMessage
208+
expect(mockApiHandler.createMessage).toHaveBeenCalled()
209+
210+
// Should have logged a warning
211+
expect(consoleSpy).toHaveBeenCalledWith("Received empty summary from API")
212+
213+
// Should return original messages
214+
expect(result).toBe(messages)
215+
216+
consoleSpy.mockRestore()
217+
})
218+
219+
it("should correctly handle non-text chunks in API response", async () => {
220+
// Setup mock to return mixed chunks
221+
mockApiHandler.createMessage.mockImplementation((): ApiStream => {
222+
return (async function* (): AsyncGenerator<ApiStreamChunk> {
223+
yield { type: "text", text: "This is " } as ApiStreamChunk
224+
yield { type: "text", text: "a summary." } as ApiStreamChunk
225+
})()
226+
})
227+
228+
const messages: ApiMessage[] = [
229+
{ role: "user", content: "Message 1", ts: 1 },
230+
{ role: "assistant", content: "Response 1", ts: 2 },
231+
{ role: "user", content: "Message 2", ts: 3 },
232+
{ role: "assistant", content: "Response 2", ts: 4 },
233+
{ role: "user", content: "Message 3", ts: 5 },
234+
{ role: "assistant", content: "Response 3", ts: 6 },
235+
{ role: "user", content: "Message 4", ts: 7 },
236+
]
237+
238+
const result = await summarizeConversationIfNeeded(
239+
messages,
240+
OVER_THRESHOLD_TOTAL_TOKENS,
241+
CONTEXT_WINDOW_SIZE,
242+
mockApiHandler,
243+
)
244+
245+
// Should have called createMessage
246+
expect(mockApiHandler.createMessage).toHaveBeenCalled()
247+
248+
// Should have a summary message with the correct content
249+
const summaryMessage = result.find((msg) => msg.isSummary)
250+
expect(summaryMessage).toBeDefined()
251+
expect(summaryMessage?.content).toBe("This is a summary.")
252+
})
253+
254+
it("should use maybeRemoveImageBlocks when preparing messages for summarization", async () => {
255+
const messages: ApiMessage[] = [
256+
{ role: "user", content: "Message 1", ts: 1 },
257+
{ role: "assistant", content: "Response 1", ts: 2 },
258+
{ role: "user", content: "Message 2", ts: 3 },
259+
{ role: "assistant", content: "Response 2", ts: 4 },
260+
{ role: "user", content: "Message 3", ts: 5 },
261+
{ role: "assistant", content: "Response 3", ts: 6 },
262+
{ role: "user", content: "Message 4", ts: 7 },
263+
]
264+
265+
await summarizeConversationIfNeeded(
266+
messages,
267+
OVER_THRESHOLD_TOTAL_TOKENS,
268+
CONTEXT_WINDOW_SIZE,
269+
mockApiHandler,
270+
)
271+
272+
// Should have called maybeRemoveImageBlocks
273+
expect(maybeRemoveImageBlocks).toHaveBeenCalled()
274+
})
275+
})
276+
})

‎src/core/condense/index.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ import { ApiHandler } from "../../api"
33
import { ApiMessage } from "../task-persistence/apiMessages"
44
import { maybeRemoveImageBlocks } from "../../api/transform/image-cleaning"
55

6-
const CONTEXT_FRAC_FOR_SUMMARY = 0.5 // TODO(canyon): make this configurable
7-
const N_MESSAGES_TO_KEEP = 3
6+
export const CONTEXT_FRAC_FOR_SUMMARY = 0.5
7+
export const N_MESSAGES_TO_KEEP = 3
88

99
const SUMMARY_PROMPT = `\
1010
Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions.

0 commit comments

Comments
 (0)