Skip to content

Commit 01bdc53

Browse files
committed
update sliding-window tests
1 parent 9765612 commit 01bdc53

File tree

1 file changed

+199
-19
lines changed

1 file changed

+199
-19
lines changed

src/core/sliding-window/__tests__/sliding-window.test.ts

Lines changed: 199 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
truncateConversationIfNeeded,
1212
} from "../index"
1313
import { ApiMessage } from "../../task-persistence/apiMessages"
14+
import * as condenseModule from "../../condense"
1415

1516
// Create a mock ApiHandler for testing
1617
class MockApiHandler extends BaseProvider {
@@ -248,7 +249,14 @@ describe("truncateConversationIfNeeded", () => {
248249
maxTokens: modelInfo.maxTokens,
249250
apiHandler: mockApiHandler,
250251
})
251-
expect(result).toEqual(messagesWithSmallContent) // No truncation occurs
252+
253+
// Check the new return type
254+
expect(result).toEqual({
255+
messages: messagesWithSmallContent,
256+
summary: "",
257+
cost: 0,
258+
prevContextTokens: totalTokens,
259+
})
252260
})
253261

254262
it("should truncate if tokens are above max tokens threshold", async () => {
@@ -260,7 +268,7 @@ describe("truncateConversationIfNeeded", () => {
260268

261269
// When truncating, always uses 0.5 fraction
262270
// With 4 messages after the first, 0.5 fraction means remove 2 messages
263-
const expectedResult = [messagesWithSmallContent[0], messagesWithSmallContent[3], messagesWithSmallContent[4]]
271+
const expectedMessages = [messagesWithSmallContent[0], messagesWithSmallContent[3], messagesWithSmallContent[4]]
264272

265273
const result = await truncateConversationIfNeeded({
266274
messages: messagesWithSmallContent,
@@ -269,7 +277,13 @@ describe("truncateConversationIfNeeded", () => {
269277
maxTokens: modelInfo.maxTokens,
270278
apiHandler: mockApiHandler,
271279
})
272-
expect(result).toEqual(expectedResult)
280+
281+
expect(result).toEqual({
282+
messages: expectedMessages,
283+
summary: "",
284+
cost: 0,
285+
prevContextTokens: totalTokens,
286+
})
273287
})
274288

275289
it("should work with non-prompt caching models the same as prompt caching models", async () => {
@@ -298,7 +312,10 @@ describe("truncateConversationIfNeeded", () => {
298312
apiHandler: mockApiHandler,
299313
})
300314

301-
expect(result1).toEqual(result2)
315+
expect(result1.messages).toEqual(result2.messages)
316+
expect(result1.summary).toEqual(result2.summary)
317+
expect(result1.cost).toEqual(result2.cost)
318+
expect(result1.prevContextTokens).toEqual(result2.prevContextTokens)
302319

303320
// Test above threshold
304321
const aboveThreshold = 70001
@@ -318,7 +335,10 @@ describe("truncateConversationIfNeeded", () => {
318335
apiHandler: mockApiHandler,
319336
})
320337

321-
expect(result3).toEqual(result4)
338+
expect(result3.messages).toEqual(result4.messages)
339+
expect(result3.summary).toEqual(result4.summary)
340+
expect(result3.cost).toEqual(result4.cost)
341+
expect(result3.prevContextTokens).toEqual(result4.prevContextTokens)
322342
})
323343

324344
it("should consider incoming content when deciding to truncate", async () => {
@@ -344,7 +364,12 @@ describe("truncateConversationIfNeeded", () => {
344364
maxTokens,
345365
apiHandler: mockApiHandler,
346366
})
347-
expect(resultWithSmall).toEqual(messagesWithSmallContent) // No truncation
367+
expect(resultWithSmall).toEqual({
368+
messages: messagesWithSmallContent,
369+
summary: "",
370+
cost: 0,
371+
prevContextTokens: baseTokensForSmall + smallContentTokens,
372+
}) // No truncation
348373

349374
// Test case 2: Large content that will push us over the threshold
350375
const largeContent = [
@@ -368,7 +393,10 @@ describe("truncateConversationIfNeeded", () => {
368393
maxTokens,
369394
apiHandler: mockApiHandler,
370395
})
371-
expect(resultWithLarge).not.toEqual(messagesWithLargeContent) // Should truncate
396+
expect(resultWithLarge.messages).not.toEqual(messagesWithLargeContent) // Should truncate
397+
expect(resultWithLarge.summary).toBe("")
398+
expect(resultWithLarge.cost).toBe(0)
399+
expect(resultWithLarge.prevContextTokens).toBe(baseTokensForLarge + largeContentTokens)
372400

373401
// Test case 3: Very large content that will definitely exceed threshold
374402
const veryLargeContent = [{ type: "text" as const, text: "X".repeat(1000) }]
@@ -387,7 +415,10 @@ describe("truncateConversationIfNeeded", () => {
387415
maxTokens,
388416
apiHandler: mockApiHandler,
389417
})
390-
expect(resultWithVeryLarge).not.toEqual(messagesWithVeryLargeContent) // Should truncate
418+
expect(resultWithVeryLarge.messages).not.toEqual(messagesWithVeryLargeContent) // Should truncate
419+
expect(resultWithVeryLarge.summary).toBe("")
420+
expect(resultWithVeryLarge.cost).toBe(0)
421+
expect(resultWithVeryLarge.prevContextTokens).toBe(baseTokensForVeryLarge + veryLargeContentTokens)
391422
})
392423

393424
it("should truncate if tokens are within TOKEN_BUFFER_PERCENTAGE of the threshold", async () => {
@@ -409,7 +440,140 @@ describe("truncateConversationIfNeeded", () => {
409440
maxTokens: modelInfo.maxTokens,
410441
apiHandler: mockApiHandler,
411442
})
412-
expect(result).toEqual(expectedResult)
443+
expect(result).toEqual({
444+
messages: expectedResult,
445+
summary: "",
446+
cost: 0,
447+
prevContextTokens: totalTokens,
448+
})
449+
})
450+
451+
it("should use summarizeConversation when autoCondenseContext is true and tokens exceed threshold", async () => {
452+
// Mock the summarizeConversation function
453+
const mockSummary = "This is a summary of the conversation"
454+
const mockCost = 0.05
455+
const mockSummarizeResponse: condenseModule.SummarizeResponse = {
456+
messages: [
457+
{ role: "user", content: "First message" },
458+
{ role: "assistant", content: mockSummary, isSummary: true },
459+
{ role: "user", content: "Last message" },
460+
],
461+
summary: mockSummary,
462+
cost: mockCost,
463+
newContextTokens: 100,
464+
}
465+
466+
const summarizeSpy = jest
467+
.spyOn(condenseModule, "summarizeConversation")
468+
.mockResolvedValue(mockSummarizeResponse)
469+
470+
const modelInfo = createModelInfo(100000, 30000)
471+
const totalTokens = 70001 // Above threshold
472+
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
473+
474+
const result = await truncateConversationIfNeeded({
475+
messages: messagesWithSmallContent,
476+
totalTokens,
477+
contextWindow: modelInfo.contextWindow,
478+
maxTokens: modelInfo.maxTokens,
479+
apiHandler: mockApiHandler,
480+
autoCondenseContext: true,
481+
systemPrompt: "System prompt",
482+
})
483+
484+
// Verify summarizeConversation was called with the right parameters
485+
expect(summarizeSpy).toHaveBeenCalledWith(messagesWithSmallContent, mockApiHandler, "System prompt")
486+
487+
// Verify the result contains the summary information
488+
expect(result).toMatchObject({
489+
messages: mockSummarizeResponse.messages,
490+
summary: mockSummary,
491+
cost: mockCost,
492+
prevContextTokens: totalTokens,
493+
})
494+
// newContextTokens might be present, but we don't need to verify its exact value
495+
496+
// Clean up
497+
summarizeSpy.mockRestore()
498+
})
499+
500+
it("should fall back to truncateConversation when autoCondenseContext is true but summarization fails", async () => {
501+
// Mock the summarizeConversation function to return empty summary
502+
const mockSummarizeResponse: condenseModule.SummarizeResponse = {
503+
messages: messages, // Original messages unchanged
504+
summary: "", // Empty summary indicates failure
505+
cost: 0.01,
506+
}
507+
508+
const summarizeSpy = jest
509+
.spyOn(condenseModule, "summarizeConversation")
510+
.mockResolvedValue(mockSummarizeResponse)
511+
512+
const modelInfo = createModelInfo(100000, 30000)
513+
const totalTokens = 70001 // Above threshold
514+
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
515+
516+
// When truncating, always uses 0.5 fraction
517+
// With 4 messages after the first, 0.5 fraction means remove 2 messages
518+
const expectedMessages = [messagesWithSmallContent[0], messagesWithSmallContent[3], messagesWithSmallContent[4]]
519+
520+
const result = await truncateConversationIfNeeded({
521+
messages: messagesWithSmallContent,
522+
totalTokens,
523+
contextWindow: modelInfo.contextWindow,
524+
maxTokens: modelInfo.maxTokens,
525+
apiHandler: mockApiHandler,
526+
autoCondenseContext: true,
527+
})
528+
529+
// Verify summarizeConversation was called
530+
expect(summarizeSpy).toHaveBeenCalled()
531+
532+
// Verify it fell back to truncation
533+
expect(result.messages).toEqual(expectedMessages)
534+
expect(result.summary).toBe("")
535+
expect(result.prevContextTokens).toBe(totalTokens)
536+
// The cost might be different than expected, so we don't check it
537+
538+
// Clean up
539+
summarizeSpy.mockRestore()
540+
})
541+
542+
it("should not call summarizeConversation when autoCondenseContext is false", async () => {
543+
// Reset any previous mock calls
544+
jest.clearAllMocks()
545+
const summarizeSpy = jest.spyOn(condenseModule, "summarizeConversation")
546+
547+
const modelInfo = createModelInfo(100000, 30000)
548+
const totalTokens = 70001 // Above threshold
549+
const messagesWithSmallContent = [...messages.slice(0, -1), { ...messages[messages.length - 1], content: "" }]
550+
551+
// When truncating, always uses 0.5 fraction
552+
// With 4 messages after the first, 0.5 fraction means remove 2 messages
553+
const expectedMessages = [messagesWithSmallContent[0], messagesWithSmallContent[3], messagesWithSmallContent[4]]
554+
555+
const result = await truncateConversationIfNeeded({
556+
messages: messagesWithSmallContent,
557+
totalTokens,
558+
contextWindow: modelInfo.contextWindow,
559+
maxTokens: modelInfo.maxTokens,
560+
apiHandler: mockApiHandler,
561+
autoCondenseContext: false,
562+
})
563+
564+
// Verify summarizeConversation was not called
565+
expect(summarizeSpy).not.toHaveBeenCalled()
566+
567+
// Verify it used truncation
568+
expect(result).toEqual({
569+
messages: expectedMessages,
570+
summary: "",
571+
cost: 0,
572+
prevContextTokens: totalTokens,
573+
})
574+
575+
// Clean up
576+
summarizeSpy.mockRestore()
413577
})
414578
})
415579

@@ -449,7 +613,12 @@ describe("getMaxTokens", () => {
449613
maxTokens: modelInfo.maxTokens,
450614
apiHandler: mockApiHandler,
451615
})
452-
expect(result1).toEqual(messagesWithSmallContent)
616+
expect(result1).toEqual({
617+
messages: messagesWithSmallContent,
618+
summary: "",
619+
cost: 0,
620+
prevContextTokens: 39999,
621+
})
453622

454623
// Above max tokens - truncate
455624
const result2 = await truncateConversationIfNeeded({
@@ -459,8 +628,11 @@ describe("getMaxTokens", () => {
459628
maxTokens: modelInfo.maxTokens,
460629
apiHandler: mockApiHandler,
461630
})
462-
expect(result2).not.toEqual(messagesWithSmallContent)
463-
expect(result2.length).toBe(3) // Truncated with 0.5 fraction
631+
expect(result2.messages).not.toEqual(messagesWithSmallContent)
632+
expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction
633+
expect(result2.summary).toBe("")
634+
expect(result2.cost).toBe(0)
635+
expect(result2.prevContextTokens).toBe(50001)
464636
})
465637

466638
it("should use 20% of context window as buffer when maxTokens is undefined", async () => {
@@ -479,7 +651,12 @@ describe("getMaxTokens", () => {
479651
maxTokens: modelInfo.maxTokens,
480652
apiHandler: mockApiHandler,
481653
})
482-
expect(result1).toEqual(messagesWithSmallContent)
654+
expect(result1).toEqual({
655+
messages: messagesWithSmallContent,
656+
summary: "",
657+
cost: 0,
658+
prevContextTokens: 69999,
659+
})
483660

484661
// Above max tokens - truncate
485662
const result2 = await truncateConversationIfNeeded({
@@ -489,8 +666,11 @@ describe("getMaxTokens", () => {
489666
maxTokens: modelInfo.maxTokens,
490667
apiHandler: mockApiHandler,
491668
})
492-
expect(result2).not.toEqual(messagesWithSmallContent)
493-
expect(result2.length).toBe(3) // Truncated with 0.5 fraction
669+
expect(result2.messages).not.toEqual(messagesWithSmallContent)
670+
expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction
671+
expect(result2.summary).toBe("")
672+
expect(result2.cost).toBe(0)
673+
expect(result2.prevContextTokens).toBe(80001)
494674
})
495675

496676
it("should handle small context windows appropriately", async () => {
@@ -508,7 +688,7 @@ describe("getMaxTokens", () => {
508688
maxTokens: modelInfo.maxTokens,
509689
apiHandler: mockApiHandler,
510690
})
511-
expect(result1).toEqual(messagesWithSmallContent)
691+
expect(result1.messages).toEqual(messagesWithSmallContent)
512692

513693
// Above max tokens - truncate
514694
const result2 = await truncateConversationIfNeeded({
@@ -519,7 +699,7 @@ describe("getMaxTokens", () => {
519699
apiHandler: mockApiHandler,
520700
})
521701
expect(result2).not.toEqual(messagesWithSmallContent)
522-
expect(result2.length).toBe(3) // Truncated with 0.5 fraction
702+
expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction
523703
})
524704

525705
it("should handle large context windows appropriately", async () => {
@@ -538,7 +718,7 @@ describe("getMaxTokens", () => {
538718
maxTokens: modelInfo.maxTokens,
539719
apiHandler: mockApiHandler,
540720
})
541-
expect(result1).toEqual(messagesWithSmallContent)
721+
expect(result1.messages).toEqual(messagesWithSmallContent)
542722

543723
// Above max tokens - truncate
544724
const result2 = await truncateConversationIfNeeded({
@@ -549,6 +729,6 @@ describe("getMaxTokens", () => {
549729
apiHandler: mockApiHandler,
550730
})
551731
expect(result2).not.toEqual(messagesWithSmallContent)
552-
expect(result2.length).toBe(3) // Truncated with 0.5 fraction
732+
expect(result2.messages.length).toBe(3) // Truncated with 0.5 fraction
553733
})
554734
})

0 commit comments

Comments
 (0)