diff --git a/core/llm/compileChatMessages.vitest.ts b/core/llm/compileChatMessages.vitest.ts new file mode 100644 index 00000000000..ce628ef1ded --- /dev/null +++ b/core/llm/compileChatMessages.vitest.ts @@ -0,0 +1,121 @@ +import { describe, expect, it } from "vitest"; +import { ChatMessage } from "../index.js"; +import { compileChatMessages } from "./countTokens.js"; + +describe("compileChatMessages", () => { + it("should preserve the latest user message when pruning during a tool sequence loop", () => { + const toolCall = { + id: "call_123", + type: "function" as const, + function: { + name: "test_tool", + arguments: "{}", + }, + }; + + const messages: ChatMessage[] = [ + { role: "system", content: "You are a helpful assistant." }, + { + role: "user", + content: "Initial user prompt that started the session.", + }, + { + role: "assistant", + content: "Calling tool...", + toolCalls: [toolCall], + }, + { + role: "tool", + content: "x".repeat(3000), // Large tool output causing context pressure + toolCallId: "call_123", + }, + ]; + + const result = compileChatMessages({ + modelName: "gpt-4", + msgs: messages, + knownContextLength: 4000, + maxTokens: 500, + supportsImages: false, + }); + + const userMessage = result.compiledChatMessages.find( + (m) => m.role === "user", + ); + expect(userMessage).toBeDefined(); + expect(userMessage?.content).toBe( + "Initial user prompt that started the session.", + ); + + // Ensure system message is preserved + const systemMessage = result.compiledChatMessages.find( + (m) => m.role === "system", + ); + expect(systemMessage).toBeDefined(); + + // Ensure tool sequence is preserved + const toolMessage = result.compiledChatMessages.find( + (m) => m.role === "tool", + ); + expect(toolMessage).toBeDefined(); + }); + + it("should preserve latest user message and prune older intermediate turns", () => { + const toolCall1 = { + id: "call_1", + type: "function" as const, + function: { name: "search", arguments: "{}" }, + }; + const toolCall2 = { + id: "call_2", + type: "function" as const, + function: { name: "search", arguments: "{}" }, + }; + + const messages: ChatMessage[] = [ + { role: "system", content: "You are a helpful assistant." }, + { + role: "user", + content: + "Turn 1 user request that is very old and long: " + "a".repeat(4000), + }, + { role: "assistant", content: "Turn 1 response." }, + { role: "user", content: "Turn 2 user request: do task B." }, + { + role: "assistant", + content: "Calling search...", + toolCalls: [toolCall2], + }, + { + role: "tool", + content: "Search results for task B: " + "b".repeat(500), + toolCallId: "call_2", + }, + ]; + + const result = compileChatMessages({ + modelName: "gpt-4", + msgs: messages, + knownContextLength: 800, + maxTokens: 100, + supportsImages: false, + }); + + // Old Turn 1 should be pruned under tight budget + const hasTurn1 = result.compiledChatMessages.some( + (m) => + typeof m.content === "string" && + m.content.includes("Turn 1 user request"), + ); + expect(hasTurn1).toBe(false); + + // Latest user message (Turn 2) MUST be preserved + const turn2User = result.compiledChatMessages.find( + (m) => + typeof m.content === "string" && + m.content.includes("Turn 2 user request"), + ); + expect(turn2User).toBeDefined(); + expect(result.didPrune).toBe(true); + }); +}); diff --git a/core/llm/countTokens.ts b/core/llm/countTokens.ts index b742d70b0f0..469f83d9fce 100644 --- a/core/llm/countTokens.ts +++ b/core/llm/countTokens.ts @@ -493,6 +493,17 @@ function compileChatMessages({ inputTokensAvailable -= systemMsgTokens; inputTokensAvailable -= lastMessagesTokens; + // Pin the latest user message if the trailing toolSequence does not already contain a user message + const hasUserInToolSequence = toolSequence.some((m) => m.role === "user"); + const pinnedUserMsg = hasUserInToolSequence + ? undefined + : [...msgsCopy].reverse().find((m) => m.role === "user"); + const pinnedUserTokens = pinnedUserMsg + ? countChatMessageTokens(modelName, pinnedUserMsg) + : 0; + + inputTokensAvailable -= pinnedUserTokens; + // Make sure there's enough context for the non-excludable items if (knownContextLength !== undefined && inputTokensAvailable < 0) { throw new Error( @@ -508,38 +519,56 @@ function compileChatMessages({ } // Now remove messages till we're under the limit + const pinnedIndex = pinnedUserMsg ? msgsCopy.indexOf(pinnedUserMsg) : -1; + const historyEntries: (ChatMessage & { tokens: number })[] = []; let currentTotal = 0; - const historyWithTokens = msgsCopy.map((message) => { + let pinnedEntry: (ChatMessage & { tokens: number }) | undefined; + + for (let i = 0; i < msgsCopy.length; i++) { + const message = msgsCopy[i]; const tokens = countChatMessageTokens(modelName, message); - currentTotal += tokens; - return { - ...message, - tokens, - }; - }); - - while (historyWithTokens.length > 0 && currentTotal > inputTokensAvailable) { - const message = historyWithTokens.shift()!; + if (i === pinnedIndex) { + pinnedEntry = { ...message, tokens }; + } else { + currentTotal += tokens; + historyEntries.push({ ...message, tokens }); + } + } + + let prunedCount = 0; + while (historyEntries.length > 0 && currentTotal > inputTokensAvailable) { + const message = historyEntries.shift()!; currentTotal -= message.tokens; didPrune = true; + prunedCount++; // At this point make sure no latent tool response without corresponding call - while (historyWithTokens[0]?.role === "tool") { - const message = historyWithTokens.shift()!; + while (historyEntries[0]?.role === "tool") { + const message = historyEntries.shift()!; currentTotal -= message.tokens; + prunedCount++; } } + if (pinnedEntry) { + const targetIndex = Math.max(0, pinnedIndex - prunedCount); + historyEntries.splice(targetIndex, 0, pinnedEntry); + } + // Now reassemble const reassembled: ChatMessage[] = []; if (systemMsg) { reassembled.push(systemMsg); } - reassembled.push(...historyWithTokens.map(({ tokens, ...rest }) => rest)); + reassembled.push(...historyEntries.map(({ tokens, ...rest }) => rest)); reassembled.push(...toolSequence); const inputTokens = - currentTotal + systemMsgTokens + toolTokens + lastMessagesTokens; + currentTotal + + pinnedUserTokens + + systemMsgTokens + + toolTokens + + lastMessagesTokens; const availableTokens = contextLength - countingSafetyBuffer - minOutputTokens; const contextPercentage = inputTokens / availableTokens;