@librechat/agents 2.2.9 → 2.3.1
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
|
@@ -64,8 +64,15 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
|
|
|
64
64
|
context = context.slice(requiredTypeIndex);
|
|
65
65
|
}
|
|
66
66
|
}
|
|
67
|
+
// Add system message if it exists
|
|
68
|
+
if (instructions && _messages.length > 0) {
|
|
69
|
+
context.push(_messages[0]);
|
|
70
|
+
messages$1.shift();
|
|
71
|
+
}
|
|
67
72
|
// Handle thinking mode requirement for Anthropic
|
|
68
73
|
if (thinkingEnabled && context.length > 0 && tokenCounter) {
|
|
74
|
+
// Check if the latest message is an assistant message
|
|
75
|
+
const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';
|
|
69
76
|
// Process only if we have an assistant message in the context
|
|
70
77
|
const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');
|
|
71
78
|
if (firstAssistantIndex >= 0) {
|
|
@@ -75,10 +82,11 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
|
|
|
75
82
|
firstAssistantMsg.content.some(item => item && typeof item === 'object' && item.type === 'thinking');
|
|
76
83
|
// Only proceed if we need to add thinking blocks
|
|
77
84
|
if (!hasThinkingBlock) {
|
|
78
|
-
// Collect thinking blocks from pruned assistant messages
|
|
85
|
+
// Collect thinking blocks from pruned assistant messages, starting from the most recent
|
|
79
86
|
const thinkingBlocks = [];
|
|
80
|
-
// Look through pruned messages for thinking blocks
|
|
81
|
-
for (
|
|
87
|
+
// Look through pruned messages for thinking blocks, starting from the end (most recent)
|
|
88
|
+
for (let i = messages$1.length - 1; i >= 0; i--) {
|
|
89
|
+
const msg = messages$1[i];
|
|
82
90
|
if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
|
|
83
91
|
for (const item of msg.content) {
|
|
84
92
|
if (item && typeof item === 'object' && item.type === 'thinking') {
|
|
@@ -125,28 +133,124 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
|
|
|
125
133
|
context[firstAssistantIndex] = newMessage;
|
|
126
134
|
// If we've exceeded the token limit, we need to prune more messages
|
|
127
135
|
if (currentTokenCount > remainingContextTokens) {
|
|
128
|
-
//
|
|
129
|
-
|
|
136
|
+
// Build a map of tool call IDs to track AI <--> tool message correspondences
|
|
137
|
+
const toolCallIdMap = new Map();
|
|
138
|
+
// Identify tool call IDs in the context
|
|
139
|
+
for (let i = 0; i < context.length; i++) {
|
|
140
|
+
const msg = context[i];
|
|
141
|
+
// Check for tool calls in AI messages
|
|
142
|
+
if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
|
|
143
|
+
for (const item of msg.content) {
|
|
144
|
+
if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
|
|
145
|
+
toolCallIdMap.set(item.id, i);
|
|
146
|
+
}
|
|
147
|
+
}
|
|
148
|
+
}
|
|
149
|
+
// Check for tool messages
|
|
150
|
+
if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {
|
|
151
|
+
toolCallIdMap.set(msg.tool_call_id, i);
|
|
152
|
+
}
|
|
153
|
+
}
|
|
154
|
+
// Track which messages to remove
|
|
155
|
+
const indicesToRemove = new Set();
|
|
156
|
+
// Start removing messages from the end, but preserve AI <--> tool message correspondences
|
|
130
157
|
let i = context.length - 1;
|
|
131
158
|
while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {
|
|
132
159
|
const msgToRemove = context[i];
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
160
|
+
// Check if this is a tool message or has tool calls
|
|
161
|
+
let canRemove = true;
|
|
162
|
+
if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {
|
|
163
|
+
// If this is a tool message, check if we need to keep its corresponding AI message
|
|
164
|
+
const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);
|
|
165
|
+
if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {
|
|
166
|
+
// We need to remove both the tool message and its corresponding AI message
|
|
167
|
+
indicesToRemove.add(i);
|
|
168
|
+
indicesToRemove.add(aiIndex);
|
|
169
|
+
currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));
|
|
170
|
+
canRemove = false;
|
|
171
|
+
}
|
|
172
|
+
}
|
|
173
|
+
else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {
|
|
174
|
+
// If this is an AI message with tool calls, check if we need to keep its corresponding tool messages
|
|
175
|
+
for (const item of msgToRemove.content) {
|
|
176
|
+
if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
|
|
177
|
+
const toolIndex = toolCallIdMap.get(item.id);
|
|
178
|
+
if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {
|
|
179
|
+
// We need to remove both the AI message and its corresponding tool message
|
|
180
|
+
indicesToRemove.add(i);
|
|
181
|
+
indicesToRemove.add(toolIndex);
|
|
182
|
+
currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));
|
|
183
|
+
canRemove = false;
|
|
184
|
+
break;
|
|
185
|
+
}
|
|
186
|
+
}
|
|
187
|
+
}
|
|
188
|
+
}
|
|
189
|
+
// If we can remove this message individually
|
|
190
|
+
if (canRemove && !indicesToRemove.has(i)) {
|
|
191
|
+
indicesToRemove.add(i);
|
|
192
|
+
currentTokenCount -= tokenCounter(msgToRemove);
|
|
193
|
+
}
|
|
136
194
|
i--;
|
|
137
195
|
}
|
|
196
|
+
// Remove messages in reverse order to avoid index shifting
|
|
197
|
+
const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);
|
|
198
|
+
for (const index of sortedIndices) {
|
|
199
|
+
context.splice(index, 1);
|
|
200
|
+
}
|
|
138
201
|
// Update remainingContextTokens to reflect the new token count
|
|
139
202
|
remainingContextTokens = maxContextTokens - currentTokenCount;
|
|
140
203
|
}
|
|
141
204
|
}
|
|
142
205
|
}
|
|
143
206
|
}
|
|
207
|
+
// If the latest message is an assistant message, ensure an assistant message with thinking appears at the end
|
|
208
|
+
// of the context (which will become the beginning after reversal)
|
|
209
|
+
// but maintain system message precedence after reversal
|
|
210
|
+
if (latestMessageIsAssistant && context.length > 0) {
|
|
211
|
+
// Find assistant messages with thinking blocks
|
|
212
|
+
const assistantIndices = [];
|
|
213
|
+
for (let i = 0; i < context.length; i++) {
|
|
214
|
+
const msg = context[i];
|
|
215
|
+
if (msg.getType() === 'ai') {
|
|
216
|
+
const hasThinking = Array.isArray(msg.content) &&
|
|
217
|
+
msg.content.some(item => item && typeof item === 'object' && item.type === 'thinking');
|
|
218
|
+
if (hasThinking) {
|
|
219
|
+
assistantIndices.push(i);
|
|
220
|
+
}
|
|
221
|
+
}
|
|
222
|
+
}
|
|
223
|
+
// If we found assistant messages with thinking blocks
|
|
224
|
+
if (assistantIndices.length > 0) {
|
|
225
|
+
// Get the first assistant message with thinking
|
|
226
|
+
const assistantWithThinkingIndex = assistantIndices[0];
|
|
227
|
+
const assistantWithThinking = context[assistantWithThinkingIndex];
|
|
228
|
+
// Remove it from its current position
|
|
229
|
+
context.splice(assistantWithThinkingIndex, 1);
|
|
230
|
+
// Check if there's a system message in the context
|
|
231
|
+
const systemIndex = context.findIndex(msg => msg.getType() === 'system');
|
|
232
|
+
const hasSystem = systemIndex !== -1;
|
|
233
|
+
if (hasSystem) {
|
|
234
|
+
// We want the system message to be first after reversal
|
|
235
|
+
// This means we need to put it at the end position before reversal
|
|
236
|
+
// And the assistant message should be second after reversal
|
|
237
|
+
// This means we need to put it at the end - 1 position before reversal
|
|
238
|
+
// First, ensure the system message is at the end (will be first after reversal)
|
|
239
|
+
const systemMsg = context[systemIndex];
|
|
240
|
+
context.splice(systemIndex, 1);
|
|
241
|
+
context.push(systemMsg);
|
|
242
|
+
// Then, put the assistant message right before the system message (will be second after reversal)
|
|
243
|
+
context.splice(context.length - 1, 0, assistantWithThinking);
|
|
244
|
+
}
|
|
245
|
+
else {
|
|
246
|
+
// No system message, so we want assistant to be first after reversal
|
|
247
|
+
// This means we need to put it at the end position before reversal
|
|
248
|
+
context.push(assistantWithThinking);
|
|
249
|
+
}
|
|
250
|
+
}
|
|
251
|
+
}
|
|
144
252
|
}
|
|
145
253
|
}
|
|
146
|
-
if (instructions && _messages.length > 0) {
|
|
147
|
-
context.push(_messages[0]);
|
|
148
|
-
messages$1.shift();
|
|
149
|
-
}
|
|
150
254
|
const prunedMemory = messages$1;
|
|
151
255
|
summaryIndex = prunedMemory.length - 1;
|
|
152
256
|
remainingContextTokens -= currentTokenCount;
|
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"prune.cjs","sources":["../../../src/messages/prune.ts"],"sourcesContent":["import { AIMessage } from '@langchain/core/messages';\nimport type { BaseMessage, UsageMetadata } from '@langchain/core/messages';\nimport type { TokenCounter } from '@/types/run';\nexport type PruneMessagesFactoryParams = {\n maxTokens: number;\n startIndex: number;\n tokenCounter: TokenCounter;\n indexTokenCountMap: Record<string, number>;\n thinkingEnabled?: boolean;\n};\nexport type PruneMessagesParams = {\n messages: BaseMessage[];\n usageMetadata?: Partial<UsageMetadata>;\n startOnMessageType?: ReturnType<BaseMessage['getType']>;\n}\n\n/**\n * Calculates the total tokens from a single usage object\n * \n * @param usage The usage metadata object containing token information\n * @returns An object containing the total input and output tokens\n */\nfunction calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetadata {\n const baseInputTokens = Number(usage.input_tokens) || 0;\n const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;\n const cacheRead = Number(usage.input_token_details?.cache_read) || 0;\n \n const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;\n const totalOutputTokens = Number(usage.output_tokens) || 0;\n\n return {\n input_tokens: totalInputTokens,\n output_tokens: totalOutputTokens,\n total_tokens: totalInputTokens + totalOutputTokens\n };\n}\n\n/**\n * Processes an array of messages and returns a context of messages that fit within a specified token limit.\n * It iterates over the messages from newest to oldest, adding them to the context until the token limit is reached.\n * \n * @param options Configuration options for processing messages\n * @returns Object containing the message context, remaining tokens, messages not included, and summary index\n */\nfunction getMessagesWithinTokenLimit({\n messages: _messages,\n maxContextTokens,\n indexTokenCountMap,\n startOnMessageType,\n thinkingEnabled,\n tokenCounter,\n}: {\n messages: BaseMessage[];\n maxContextTokens: number;\n indexTokenCountMap: Record<string, number>;\n startOnMessageType?: string;\n thinkingEnabled?: boolean;\n tokenCounter?: TokenCounter;\n}): {\n context: BaseMessage[];\n remainingContextTokens: number;\n messagesToRefine: BaseMessage[];\n summaryIndex: number;\n} {\n // Every reply is primed with <|start|>assistant<|message|>, so we\n // start with 3 tokens for the label after all messages have been counted.\n let summaryIndex = -1;\n let currentTokenCount = 3;\n const instructions = _messages?.[0]?.getType() === 'system' ? _messages[0] : undefined;\n const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;\n let remainingContextTokens = maxContextTokens - instructionsTokenCount;\n const messages = [..._messages];\n let context: BaseMessage[] = [];\n\n if (currentTokenCount < remainingContextTokens) {\n let currentIndex = messages.length;\n while (messages.length > 0 && currentTokenCount < remainingContextTokens && currentIndex > 1) {\n currentIndex--;\n if (messages.length === 1 && instructions) {\n break;\n }\n const poppedMessage = messages.pop();\n if (!poppedMessage) continue;\n \n const tokenCount = indexTokenCountMap[currentIndex] || 0;\n\n if ((currentTokenCount + tokenCount) <= remainingContextTokens) {\n context.push(poppedMessage);\n currentTokenCount += tokenCount;\n } else {\n messages.push(poppedMessage);\n break;\n }\n }\n \n // Handle startOnMessageType requirement\n if (startOnMessageType && context.length > 0) {\n const requiredTypeIndex = context.findIndex(msg => msg.getType() === startOnMessageType);\n \n if (requiredTypeIndex > 0) {\n context = context.slice(requiredTypeIndex);\n }\n }\n \n // Handle thinking mode requirement for Anthropic\n if (thinkingEnabled && context.length > 0 && tokenCounter) {\n // Process only if we have an assistant message in the context\n const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');\n if (firstAssistantIndex >= 0) {\n const firstAssistantMsg = context[firstAssistantIndex];\n \n // Check if the first assistant message already has a thinking block\n const hasThinkingBlock = Array.isArray(firstAssistantMsg.content) && \n firstAssistantMsg.content.some(item => \n item && typeof item === 'object' && item.type === 'thinking');\n \n // Only proceed if we need to add thinking blocks\n if (!hasThinkingBlock) {\n // Collect thinking blocks from pruned assistant messages\n const thinkingBlocks: any[] = [];\n \n // Look through pruned messages for thinking blocks\n for (const msg of messages) {\n if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n for (const item of msg.content) {\n if (item && typeof item === 'object' && item.type === 'thinking') {\n thinkingBlocks.push(item);\n // We only need one thinking block\n break;\n }\n }\n if (thinkingBlocks.length > 0) break; // Stop after finding one thinking block\n }\n }\n \n // If we found thinking blocks, add them to the first assistant message\n if (thinkingBlocks.length > 0) {\n // Calculate token count of original message\n const originalTokenCount = tokenCounter(firstAssistantMsg);\n \n // Create a new content array with thinking blocks at the beginning\n let newContent: any[];\n \n if (Array.isArray(firstAssistantMsg.content)) {\n // Keep the original content (excluding any existing thinking blocks)\n const originalContent = firstAssistantMsg.content.filter(item => \n !(item && typeof item === 'object' && item.type === 'thinking'));\n \n newContent = [...thinkingBlocks, ...originalContent];\n } else if (typeof firstAssistantMsg.content === 'string') {\n newContent = [\n ...thinkingBlocks,\n { type: 'text', text: firstAssistantMsg.content }\n ];\n } else {\n newContent = thinkingBlocks;\n }\n \n // Create a new message with the updated content\n const newMessage = new AIMessage({\n content: newContent,\n additional_kwargs: firstAssistantMsg.additional_kwargs,\n response_metadata: firstAssistantMsg.response_metadata,\n });\n \n // Calculate token count of new message\n const newTokenCount = tokenCounter(newMessage);\n \n // Adjust current token count\n currentTokenCount += (newTokenCount - originalTokenCount);\n \n // Replace the first assistant message\n context[firstAssistantIndex] = newMessage;\n \n // If we've exceeded the token limit, we need to prune more messages\n if (currentTokenCount > remainingContextTokens) {\n // Remove messages from the end of the context until we're under the token limit\n // But make sure to keep the first assistant message with thinking block\n let i = context.length - 1;\n while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {\n const msgToRemove = context[i];\n const msgTokenCount = tokenCounter(msgToRemove);\n context.splice(i, 1);\n currentTokenCount -= msgTokenCount;\n i--;\n }\n \n // Update remainingContextTokens to reflect the new token count\n remainingContextTokens = maxContextTokens - currentTokenCount;\n }\n }\n }\n }\n }\n }\n\n if (instructions && _messages.length > 0) {\n context.push(_messages[0] as BaseMessage);\n messages.shift();\n }\n\n const prunedMemory = messages;\n summaryIndex = prunedMemory.length - 1;\n remainingContextTokens -= currentTokenCount;\n\n return {\n summaryIndex,\n remainingContextTokens,\n context: context.reverse(),\n messagesToRefine: prunedMemory,\n };\n}\n\nfunction checkValidNumber(value: unknown): value is number {\n return typeof value === 'number' && !isNaN(value) && value > 0;\n}\n\nexport function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {\n const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };\n let lastTurnStartIndex = factoryParams.startIndex;\n let totalTokens = (Object.values(indexTokenCountMap)).reduce((a, b) => a + b, 0);\n \n return function pruneMessages(params: PruneMessagesParams): {\n context: BaseMessage[];\n indexTokenCountMap: Record<string, number>;\n } {\n let currentUsage: UsageMetadata | undefined;\n if (params.usageMetadata && (\n checkValidNumber(params.usageMetadata.input_tokens)\n || (\n checkValidNumber(params.usageMetadata.input_token_details)\n && (\n checkValidNumber(params.usageMetadata.input_token_details.cache_creation)\n || checkValidNumber(params.usageMetadata.input_token_details.cache_read)\n )\n )\n ) && checkValidNumber(params.usageMetadata.output_tokens)) {\n currentUsage = calculateTotalTokens(params.usageMetadata);\n totalTokens = currentUsage.total_tokens;\n }\n\n for (let i = lastTurnStartIndex; i < params.messages.length; i++) {\n const message = params.messages[i];\n if (i === lastTurnStartIndex && indexTokenCountMap[i] === undefined && currentUsage) {\n indexTokenCountMap[i] = currentUsage.output_tokens;\n } else if (indexTokenCountMap[i] === undefined) {\n indexTokenCountMap[i] = factoryParams.tokenCounter(message);\n totalTokens += indexTokenCountMap[i];\n }\n }\n\n // If `currentUsage` is defined, we need to distribute the current total tokensto our `indexTokenCountMap`,\n // for all message index keys before `lastTurnStartIndex`, as it has the most accurate count for those messages.\n // We must distribute it in a weighted manner, so that the total token count is equal to `currentUsage.total_tokens`,\n // relative the manually counted tokens in `indexTokenCountMap`.\n if (currentUsage) {\n const totalIndexTokens = Object.values(indexTokenCountMap).reduce((a, b) => a + b, 0);\n const ratio = currentUsage.total_tokens / totalIndexTokens;\n for (const key in indexTokenCountMap) {\n indexTokenCountMap[key] = Math.round(indexTokenCountMap[key] * ratio);\n }\n }\n\n lastTurnStartIndex = params.messages.length;\n if (totalTokens <= factoryParams.maxTokens) {\n return { context: params.messages, indexTokenCountMap };\n }\n\n // Pass the tokenCounter to getMessagesWithinTokenLimit for token recalculation\n const { context } = getMessagesWithinTokenLimit({\n maxContextTokens: factoryParams.maxTokens,\n messages: params.messages,\n indexTokenCountMap,\n startOnMessageType: params.startOnMessageType,\n thinkingEnabled: factoryParams.thinkingEnabled,\n tokenCounter: factoryParams.tokenCounter,\n });\n\n return { context, indexTokenCountMap };\n }\n}\n"],"names":["messages","AIMessage"],"mappings":";;;;AAgBA;;;;;AAKG;AACH,SAAS,oBAAoB,CAAC,KAA6B,EAAA;IACzD,MAAM,eAAe,GAAG,MAAM,CAAC,KAAK,CAAC,YAAY,CAAC,IAAI,CAAC;AACvD,IAAA,MAAM,aAAa,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,cAAc,CAAC,IAAI,CAAC;AAC5E,IAAA,MAAM,SAAS,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,UAAU,CAAC,IAAI,CAAC;AAEpE,IAAA,MAAM,gBAAgB,GAAG,eAAe,GAAG,aAAa,GAAG,SAAS;IACpE,MAAM,iBAAiB,GAAG,MAAM,CAAC,KAAK,CAAC,aAAa,CAAC,IAAI,CAAC;IAE1D,OAAO;AACL,QAAA,YAAY,EAAE,gBAAgB;AAC9B,QAAA,aAAa,EAAE,iBAAiB;QAChC,YAAY,EAAE,gBAAgB,GAAG;KAClC;AACH;AAEA;;;;;;AAMG;AACH,SAAS,2BAA2B,CAAC,EACnC,QAAQ,EAAE,SAAS,EACnB,gBAAgB,EAChB,kBAAkB,EAClB,kBAAkB,EAClB,eAAe,EACf,YAAY,GAQb,EAAA;;;AAQC,IAAA,IAAI,YAAY,GAAG,EAAE;IACrB,IAAI,iBAAiB,GAAG,CAAC;IACzB,MAAM,YAAY,GAAG,SAAS,GAAG,CAAC,CAAC,EAAE,OAAO,EAAE,KAAK,QAAQ,GAAG,SAAS,CAAC,CAAC,CAAC,GAAG,SAAS;AACtF,IAAA,MAAM,sBAAsB,GAAG,YAAY,IAAI,IAAI,GAAG,kBAAkB,CAAC,CAAC,CAAC,GAAG,CAAC;AAC/E,IAAA,IAAI,sBAAsB,GAAG,gBAAgB,GAAG,sBAAsB;AACtE,IAAA,MAAMA,UAAQ,GAAG,CAAC,GAAG,SAAS,CAAC;IAC/B,IAAI,OAAO,GAAkB,EAAE;AAE/B,IAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC9C,QAAA,IAAI,YAAY,GAAGA,UAAQ,CAAC,MAAM;AAClC,QAAA,OAAOA,UAAQ,CAAC,MAAM,GAAG,CAAC,IAAI,iBAAiB,GAAG,sBAAsB,IAAI,YAAY,GAAG,CAAC,EAAE;AAC5F,YAAA,YAAY,EAAE;YACd,IAAIA,UAAQ,CAAC,MAAM,KAAK,CAAC,IAAI,YAAY,EAAE;gBACzC;;AAEF,YAAA,MAAM,aAAa,GAAGA,UAAQ,CAAC,GAAG,EAAE;AACpC,YAAA,IAAI,CAAC,aAAa;gBAAE;YAEpB,MAAM,UAAU,GAAG,kBAAkB,CAAC,YAAY,CAAC,IAAI,CAAC;YAExD,IAAI,CAAC,iBAAiB,GAAG,UAAU,KAAK,sBAAsB,EAAE;AAC9D,gBAAA,OAAO,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC3B,iBAAiB,IAAI,UAAU;;iBAC1B;AACL,gBAAAA,UAAQ,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC5B;;;;QAKN,IAAI,kBAAkB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;AAC5C,YAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,kBAAkB,CAAC;AAExF,YAAA,IAAI,iBAAiB,GAAG,CAAC,EAAE;AACzB,gBAAA,OAAO,GAAG,OAAO,CAAC,KAAK,CAAC,iBAAiB,CAAC;;;;QAK9C,IAAI,eAAe,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,IAAI,YAAY,EAAE;;AAEzD,YAAA,MAAM,mBAAmB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,CAAC;AAC5E,YAAA,IAAI,mBAAmB,IAAI,CAAC,EAAE;AAC5B,gBAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,mBAAmB,CAAC;;gBAGtD,MAAM,gBAAgB,GAAG,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC;oBAC/D,iBAAiB,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,IACjC,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC;;gBAGjE,IAAI,CAAC,gBAAgB,EAAE;;oBAErB,MAAM,cAAc,GAAU,EAAE;;AAGhC,oBAAA,KAAK,MAAM,GAAG,IAAIA,UAAQ,EAAE;AAC1B,wBAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,4BAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,gCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,EAAE;AAChE,oCAAA,cAAc,CAAC,IAAI,CAAC,IAAI,CAAC;;oCAEzB;;;AAGJ,4BAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC;AAAE,gCAAA,MAAM;;;;AAKzC,oBAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC,EAAE;;AAE7B,wBAAA,MAAM,kBAAkB,GAAG,YAAY,CAAC,iBAAiB,CAAC;;AAG1D,wBAAA,IAAI,UAAiB;wBAErB,IAAI,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC,EAAE;;AAE5C,4BAAA,MAAM,eAAe,GAAG,iBAAiB,CAAC,OAAO,CAAC,MAAM,CAAC,IAAI,IAC3D,EAAE,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC,CAAC;4BAElE,UAAU,GAAG,CAAC,GAAG,cAAc,EAAE,GAAG,eAAe,CAAC;;AAC/C,6BAAA,IAAI,OAAO,iBAAiB,CAAC,OAAO,KAAK,QAAQ,EAAE;AACxD,4BAAA,UAAU,GAAG;AACX,gCAAA,GAAG,cAAc;gCACjB,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,iBAAiB,CAAC,OAAO;6BAChD;;6BACI;4BACL,UAAU,GAAG,cAAc;;;AAI7B,wBAAA,MAAM,UAAU,GAAG,IAAIC,kBAAS,CAAC;AAC/B,4BAAA,OAAO,EAAE,UAAU;4BACnB,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;4BACtD,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;AACvD,yBAAA,CAAC;;AAGF,wBAAA,MAAM,aAAa,GAAG,YAAY,CAAC,UAAU,CAAC;;AAG9C,wBAAA,iBAAiB,KAAK,aAAa,GAAG,kBAAkB,CAAC;;AAGzD,wBAAA,OAAO,CAAC,mBAAmB,CAAC,GAAG,UAAU;;AAGzC,wBAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;;;AAG9C,4BAAA,IAAI,CAAC,GAAG,OAAO,CAAC,MAAM,GAAG,CAAC;4BAC1B,OAAO,CAAC,GAAG,mBAAmB,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC5E,gCAAA,MAAM,WAAW,GAAG,OAAO,CAAC,CAAC,CAAC;AAC9B,gCAAA,MAAM,aAAa,GAAG,YAAY,CAAC,WAAW,CAAC;AAC/C,gCAAA,OAAO,CAAC,MAAM,CAAC,CAAC,EAAE,CAAC,CAAC;gCACpB,iBAAiB,IAAI,aAAa;AAClC,gCAAA,CAAC,EAAE;;;AAIL,4BAAA,sBAAsB,GAAG,gBAAgB,GAAG,iBAAiB;;;;;;;IAQvE,IAAI,YAAY,IAAI,SAAS,CAAC,MAAM,GAAG,CAAC,EAAE;QACxC,OAAO,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,CAAgB,CAAC;QACzCD,UAAQ,CAAC,KAAK,EAAE;;IAGlB,MAAM,YAAY,GAAGA,UAAQ;AAC7B,IAAA,YAAY,GAAG,YAAY,CAAC,MAAM,GAAG,CAAC;IACtC,sBAAsB,IAAI,iBAAiB;IAE3C,OAAO;QACL,YAAY;QACZ,sBAAsB;AACtB,QAAA,OAAO,EAAE,OAAO,CAAC,OAAO,EAAE;AAC1B,QAAA,gBAAgB,EAAE,YAAY;KAC/B;AACH;AAEA,SAAS,gBAAgB,CAAC,KAAc,EAAA;AACtC,IAAA,OAAO,OAAO,KAAK,KAAK,QAAQ,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,KAAK,GAAG,CAAC;AAChE;AAEM,SAAU,mBAAmB,CAAC,aAAyC,EAAA;IAC3E,MAAM,kBAAkB,GAAG,EAAE,GAAG,aAAa,CAAC,kBAAkB,EAAE;AAClE,IAAA,IAAI,kBAAkB,GAAG,aAAa,CAAC,UAAU;IACjD,IAAI,WAAW,GAAG,CAAC,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;IAEhF,OAAO,SAAS,aAAa,CAAC,MAA2B,EAAA;AAIvD,QAAA,IAAI,YAAuC;AAC3C,QAAA,IAAI,MAAM,CAAC,aAAa,KACtB,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,YAAY;AAC/C,gBACD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB;oBAEvD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,cAAc;uBACrE,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,UAAU,CAAC,CACzE,CACF,CACF,IAAI,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,aAAa,CAAC,EAAE;AACzD,YAAA,YAAY,GAAG,oBAAoB,CAAC,MAAM,CAAC,aAAa,CAAC;AACzD,YAAA,WAAW,GAAG,YAAY,CAAC,YAAY;;AAGzC,QAAA,KAAK,IAAI,CAAC,GAAG,kBAAkB,EAAE,CAAC,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;YAChE,MAAM,OAAO,GAAG,MAAM,CAAC,QAAQ,CAAC,CAAC,CAAC;AAClC,YAAA,IAAI,CAAC,KAAK,kBAAkB,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,IAAI,YAAY,EAAE;AACnF,gBAAA,kBAAkB,CAAC,CAAC,CAAC,GAAG,YAAY,CAAC,aAAa;;AAC7C,iBAAA,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,EAAE;gBAC9C,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,YAAY,CAAC,OAAO,CAAC;AAC3D,gBAAA,WAAW,IAAI,kBAAkB,CAAC,CAAC,CAAC;;;;;;;QAQxC,IAAI,YAAY,EAAE;YAChB,MAAM,gBAAgB,GAAG,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;AACrF,YAAA,MAAM,KAAK,GAAG,YAAY,CAAC,YAAY,GAAG,gBAAgB;AAC1D,YAAA,KAAK,MAAM,GAAG,IAAI,kBAAkB,EAAE;AACpC,gBAAA,kBAAkB,CAAC,GAAG,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,GAAG,CAAC,GAAG,KAAK,CAAC;;;AAIzE,QAAA,kBAAkB,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM;AAC3C,QAAA,IAAI,WAAW,IAAI,aAAa,CAAC,SAAS,EAAE;YAC1C,OAAO,EAAE,OAAO,EAAE,MAAM,CAAC,QAAQ,EAAE,kBAAkB,EAAE;;;AAIzD,QAAA,MAAM,EAAE,OAAO,EAAE,GAAG,2BAA2B,CAAC;YAC9C,gBAAgB,EAAE,aAAa,CAAC,SAAS;YACzC,QAAQ,EAAE,MAAM,CAAC,QAAQ;YACzB,kBAAkB;YAClB,kBAAkB,EAAE,MAAM,CAAC,kBAAkB;YAC7C,eAAe,EAAE,aAAa,CAAC,eAAe;YAC9C,YAAY,EAAE,aAAa,CAAC,YAAY;AACzC,SAAA,CAAC;AAEF,QAAA,OAAO,EAAE,OAAO,EAAE,kBAAkB,EAAE;AACxC,KAAC;AACH;;;;"}
|
|
1
|
+
{"version":3,"file":"prune.cjs","sources":["../../../src/messages/prune.ts"],"sourcesContent":["import { AIMessage } from '@langchain/core/messages';\nimport type { BaseMessage, UsageMetadata } from '@langchain/core/messages';\nimport type { TokenCounter } from '@/types/run';\nexport type PruneMessagesFactoryParams = {\n maxTokens: number;\n startIndex: number;\n tokenCounter: TokenCounter;\n indexTokenCountMap: Record<string, number>;\n thinkingEnabled?: boolean;\n};\nexport type PruneMessagesParams = {\n messages: BaseMessage[];\n usageMetadata?: Partial<UsageMetadata>;\n startOnMessageType?: ReturnType<BaseMessage['getType']>;\n}\n\n/**\n * Calculates the total tokens from a single usage object\n * \n * @param usage The usage metadata object containing token information\n * @returns An object containing the total input and output tokens\n */\nfunction calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetadata {\n const baseInputTokens = Number(usage.input_tokens) || 0;\n const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;\n const cacheRead = Number(usage.input_token_details?.cache_read) || 0;\n \n const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;\n const totalOutputTokens = Number(usage.output_tokens) || 0;\n\n return {\n input_tokens: totalInputTokens,\n output_tokens: totalOutputTokens,\n total_tokens: totalInputTokens + totalOutputTokens\n };\n}\n\n/**\n * Processes an array of messages and returns a context of messages that fit within a specified token limit.\n * It iterates over the messages from newest to oldest, adding them to the context until the token limit is reached.\n * \n * @param options Configuration options for processing messages\n * @returns Object containing the message context, remaining tokens, messages not included, and summary index\n */\nfunction getMessagesWithinTokenLimit({\n messages: _messages,\n maxContextTokens,\n indexTokenCountMap,\n startOnMessageType,\n thinkingEnabled,\n tokenCounter,\n}: {\n messages: BaseMessage[];\n maxContextTokens: number;\n indexTokenCountMap: Record<string, number>;\n startOnMessageType?: string;\n thinkingEnabled?: boolean;\n tokenCounter?: TokenCounter;\n}): {\n context: BaseMessage[];\n remainingContextTokens: number;\n messagesToRefine: BaseMessage[];\n summaryIndex: number;\n} {\n // Every reply is primed with <|start|>assistant<|message|>, so we\n // start with 3 tokens for the label after all messages have been counted.\n let summaryIndex = -1;\n let currentTokenCount = 3;\n const instructions = _messages?.[0]?.getType() === 'system' ? _messages[0] : undefined;\n const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;\n let remainingContextTokens = maxContextTokens - instructionsTokenCount;\n const messages = [..._messages];\n let context: BaseMessage[] = [];\n\n if (currentTokenCount < remainingContextTokens) {\n let currentIndex = messages.length;\n while (messages.length > 0 && currentTokenCount < remainingContextTokens && currentIndex > 1) {\n currentIndex--;\n if (messages.length === 1 && instructions) {\n break;\n }\n const poppedMessage = messages.pop();\n if (!poppedMessage) continue;\n \n const tokenCount = indexTokenCountMap[currentIndex] || 0;\n\n if ((currentTokenCount + tokenCount) <= remainingContextTokens) {\n context.push(poppedMessage);\n currentTokenCount += tokenCount;\n } else {\n messages.push(poppedMessage);\n break;\n }\n }\n \n // Handle startOnMessageType requirement\n if (startOnMessageType && context.length > 0) {\n const requiredTypeIndex = context.findIndex(msg => msg.getType() === startOnMessageType);\n \n if (requiredTypeIndex > 0) {\n context = context.slice(requiredTypeIndex);\n }\n }\n \n // Add system message if it exists\n if (instructions && _messages.length > 0) {\n context.push(_messages[0] as BaseMessage);\n messages.shift();\n }\n \n // Handle thinking mode requirement for Anthropic\n if (thinkingEnabled && context.length > 0 && tokenCounter) {\n // Check if the latest message is an assistant message\n const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';\n \n // Process only if we have an assistant message in the context\n const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');\n \n if (firstAssistantIndex >= 0) {\n const firstAssistantMsg = context[firstAssistantIndex];\n \n // Check if the first assistant message already has a thinking block\n const hasThinkingBlock = Array.isArray(firstAssistantMsg.content) && \n firstAssistantMsg.content.some(item => \n item && typeof item === 'object' && item.type === 'thinking');\n \n // Only proceed if we need to add thinking blocks\n if (!hasThinkingBlock) {\n // Collect thinking blocks from pruned assistant messages, starting from the most recent\n const thinkingBlocks: any[] = [];\n \n // Look through pruned messages for thinking blocks, starting from the end (most recent)\n for (let i = messages.length - 1; i >= 0; i--) {\n const msg = messages[i];\n if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n for (const item of msg.content) {\n if (item && typeof item === 'object' && item.type === 'thinking') {\n thinkingBlocks.push(item);\n // We only need one thinking block\n break;\n }\n }\n if (thinkingBlocks.length > 0) break; // Stop after finding one thinking block\n }\n }\n \n // If we found thinking blocks, add them to the first assistant message\n if (thinkingBlocks.length > 0) {\n // Calculate token count of original message\n const originalTokenCount = tokenCounter(firstAssistantMsg);\n \n // Create a new content array with thinking blocks at the beginning\n let newContent: any[];\n \n if (Array.isArray(firstAssistantMsg.content)) {\n // Keep the original content (excluding any existing thinking blocks)\n const originalContent = firstAssistantMsg.content.filter(item => \n !(item && typeof item === 'object' && item.type === 'thinking'));\n \n newContent = [...thinkingBlocks, ...originalContent];\n } else if (typeof firstAssistantMsg.content === 'string') {\n newContent = [\n ...thinkingBlocks,\n { type: 'text', text: firstAssistantMsg.content }\n ];\n } else {\n newContent = thinkingBlocks;\n }\n \n // Create a new message with the updated content\n const newMessage = new AIMessage({\n content: newContent,\n additional_kwargs: firstAssistantMsg.additional_kwargs,\n response_metadata: firstAssistantMsg.response_metadata,\n });\n \n // Calculate token count of new message\n const newTokenCount = tokenCounter(newMessage);\n \n // Adjust current token count\n currentTokenCount += (newTokenCount - originalTokenCount);\n \n // Replace the first assistant message\n context[firstAssistantIndex] = newMessage;\n \n // If we've exceeded the token limit, we need to prune more messages\n if (currentTokenCount > remainingContextTokens) {\n // Build a map of tool call IDs to track AI <--> tool message correspondences\n const toolCallIdMap = new Map<string, number>();\n \n // Identify tool call IDs in the context\n for (let i = 0; i < context.length; i++) {\n const msg = context[i];\n \n // Check for tool calls in AI messages\n if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n for (const item of msg.content) {\n if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {\n toolCallIdMap.set(item.id, i);\n }\n }\n }\n \n // Check for tool messages\n if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {\n toolCallIdMap.set(msg.tool_call_id, i);\n }\n }\n \n // Track which messages to remove\n const indicesToRemove = new Set<number>();\n \n // Start removing messages from the end, but preserve AI <--> tool message correspondences\n let i = context.length - 1;\n while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {\n const msgToRemove = context[i];\n \n // Check if this is a tool message or has tool calls\n let canRemove = true;\n \n if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {\n // If this is a tool message, check if we need to keep its corresponding AI message\n const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);\n if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {\n // We need to remove both the tool message and its corresponding AI message\n indicesToRemove.add(i);\n indicesToRemove.add(aiIndex);\n currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));\n canRemove = false;\n }\n } else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {\n // If this is an AI message with tool calls, check if we need to keep its corresponding tool messages\n for (const item of msgToRemove.content) {\n if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {\n const toolIndex = toolCallIdMap.get(item.id as string);\n if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {\n // We need to remove both the AI message and its corresponding tool message\n indicesToRemove.add(i);\n indicesToRemove.add(toolIndex);\n currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));\n canRemove = false;\n break;\n }\n }\n }\n }\n \n // If we can remove this message individually\n if (canRemove && !indicesToRemove.has(i)) {\n indicesToRemove.add(i);\n currentTokenCount -= tokenCounter(msgToRemove);\n }\n \n i--;\n }\n \n // Remove messages in reverse order to avoid index shifting\n const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);\n for (const index of sortedIndices) {\n context.splice(index, 1);\n }\n \n // Update remainingContextTokens to reflect the new token count\n remainingContextTokens = maxContextTokens - currentTokenCount;\n }\n }\n }\n }\n \n // If the latest message is an assistant message, ensure an assistant message with thinking appears at the end\n // of the context (which will become the beginning after reversal)\n // but maintain system message precedence after reversal\n if (latestMessageIsAssistant && context.length > 0) {\n // Find assistant messages with thinking blocks\n const assistantIndices: number[] = [];\n for (let i = 0; i < context.length; i++) {\n const msg = context[i];\n if (msg.getType() === 'ai') {\n const hasThinking = Array.isArray(msg.content) && \n msg.content.some(item => item && typeof item === 'object' && item.type === 'thinking');\n \n if (hasThinking) {\n assistantIndices.push(i);\n }\n }\n }\n \n // If we found assistant messages with thinking blocks\n if (assistantIndices.length > 0) {\n // Get the first assistant message with thinking\n const assistantWithThinkingIndex = assistantIndices[0];\n const assistantWithThinking = context[assistantWithThinkingIndex];\n \n // Remove it from its current position\n context.splice(assistantWithThinkingIndex, 1);\n \n // Check if there's a system message in the context\n const systemIndex = context.findIndex(msg => msg.getType() === 'system');\n const hasSystem = systemIndex !== -1;\n \n if (hasSystem) {\n // We want the system message to be first after reversal\n // This means we need to put it at the end position before reversal\n // And the assistant message should be second after reversal\n // This means we need to put it at the end - 1 position before reversal\n \n // First, ensure the system message is at the end (will be first after reversal)\n const systemMsg = context[systemIndex];\n context.splice(systemIndex, 1);\n context.push(systemMsg);\n \n // Then, put the assistant message right before the system message (will be second after reversal)\n context.splice(context.length - 1, 0, assistantWithThinking);\n } else {\n // No system message, so we want assistant to be first after reversal\n // This means we need to put it at the end position before reversal\n context.push(assistantWithThinking);\n }\n }\n }\n }\n }\n\n const prunedMemory = messages;\n summaryIndex = prunedMemory.length - 1;\n remainingContextTokens -= currentTokenCount;\n\n return {\n summaryIndex,\n remainingContextTokens,\n context: context.reverse(),\n messagesToRefine: prunedMemory,\n };\n}\n\nfunction checkValidNumber(value: unknown): value is number {\n return typeof value === 'number' && !isNaN(value) && value > 0;\n}\n\nexport function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {\n const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };\n let lastTurnStartIndex = factoryParams.startIndex;\n let totalTokens = (Object.values(indexTokenCountMap)).reduce((a, b) => a + b, 0);\n \n return function pruneMessages(params: PruneMessagesParams): {\n context: BaseMessage[];\n indexTokenCountMap: Record<string, number>;\n } {\n let currentUsage: UsageMetadata | undefined;\n if (params.usageMetadata && (\n checkValidNumber(params.usageMetadata.input_tokens)\n || (\n checkValidNumber(params.usageMetadata.input_token_details)\n && (\n checkValidNumber(params.usageMetadata.input_token_details.cache_creation)\n || checkValidNumber(params.usageMetadata.input_token_details.cache_read)\n )\n )\n ) && checkValidNumber(params.usageMetadata.output_tokens)) {\n currentUsage = calculateTotalTokens(params.usageMetadata);\n totalTokens = currentUsage.total_tokens;\n }\n\n for (let i = lastTurnStartIndex; i < params.messages.length; i++) {\n const message = params.messages[i];\n if (i === lastTurnStartIndex && indexTokenCountMap[i] === undefined && currentUsage) {\n indexTokenCountMap[i] = currentUsage.output_tokens;\n } else if (indexTokenCountMap[i] === undefined) {\n indexTokenCountMap[i] = factoryParams.tokenCounter(message);\n totalTokens += indexTokenCountMap[i];\n }\n }\n\n // If `currentUsage` is defined, we need to distribute the current total tokensto our `indexTokenCountMap`,\n // for all message index keys before `lastTurnStartIndex`, as it has the most accurate count for those messages.\n // We must distribute it in a weighted manner, so that the total token count is equal to `currentUsage.total_tokens`,\n // relative the manually counted tokens in `indexTokenCountMap`.\n if (currentUsage) {\n const totalIndexTokens = Object.values(indexTokenCountMap).reduce((a, b) => a + b, 0);\n const ratio = currentUsage.total_tokens / totalIndexTokens;\n for (const key in indexTokenCountMap) {\n indexTokenCountMap[key] = Math.round(indexTokenCountMap[key] * ratio);\n }\n }\n\n lastTurnStartIndex = params.messages.length;\n if (totalTokens <= factoryParams.maxTokens) {\n return { context: params.messages, indexTokenCountMap };\n }\n\n // Pass the tokenCounter to getMessagesWithinTokenLimit for token recalculation\n const { context } = getMessagesWithinTokenLimit({\n maxContextTokens: factoryParams.maxTokens,\n messages: params.messages,\n indexTokenCountMap,\n startOnMessageType: params.startOnMessageType,\n thinkingEnabled: factoryParams.thinkingEnabled,\n tokenCounter: factoryParams.tokenCounter,\n });\n\n return { context, indexTokenCountMap };\n }\n}\n"],"names":["messages","AIMessage"],"mappings":";;;;AAgBA;;;;;AAKG;AACH,SAAS,oBAAoB,CAAC,KAA6B,EAAA;IACzD,MAAM,eAAe,GAAG,MAAM,CAAC,KAAK,CAAC,YAAY,CAAC,IAAI,CAAC;AACvD,IAAA,MAAM,aAAa,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,cAAc,CAAC,IAAI,CAAC;AAC5E,IAAA,MAAM,SAAS,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,UAAU,CAAC,IAAI,CAAC;AAEpE,IAAA,MAAM,gBAAgB,GAAG,eAAe,GAAG,aAAa,GAAG,SAAS;IACpE,MAAM,iBAAiB,GAAG,MAAM,CAAC,KAAK,CAAC,aAAa,CAAC,IAAI,CAAC;IAE1D,OAAO;AACL,QAAA,YAAY,EAAE,gBAAgB;AAC9B,QAAA,aAAa,EAAE,iBAAiB;QAChC,YAAY,EAAE,gBAAgB,GAAG;KAClC;AACH;AAEA;;;;;;AAMG;AACH,SAAS,2BAA2B,CAAC,EACnC,QAAQ,EAAE,SAAS,EACnB,gBAAgB,EAChB,kBAAkB,EAClB,kBAAkB,EAClB,eAAe,EACf,YAAY,GAQb,EAAA;;;AAQC,IAAA,IAAI,YAAY,GAAG,EAAE;IACrB,IAAI,iBAAiB,GAAG,CAAC;IACzB,MAAM,YAAY,GAAG,SAAS,GAAG,CAAC,CAAC,EAAE,OAAO,EAAE,KAAK,QAAQ,GAAG,SAAS,CAAC,CAAC,CAAC,GAAG,SAAS;AACtF,IAAA,MAAM,sBAAsB,GAAG,YAAY,IAAI,IAAI,GAAG,kBAAkB,CAAC,CAAC,CAAC,GAAG,CAAC;AAC/E,IAAA,IAAI,sBAAsB,GAAG,gBAAgB,GAAG,sBAAsB;AACtE,IAAA,MAAMA,UAAQ,GAAG,CAAC,GAAG,SAAS,CAAC;IAC/B,IAAI,OAAO,GAAkB,EAAE;AAE/B,IAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC9C,QAAA,IAAI,YAAY,GAAGA,UAAQ,CAAC,MAAM;AAClC,QAAA,OAAOA,UAAQ,CAAC,MAAM,GAAG,CAAC,IAAI,iBAAiB,GAAG,sBAAsB,IAAI,YAAY,GAAG,CAAC,EAAE;AAC5F,YAAA,YAAY,EAAE;YACd,IAAIA,UAAQ,CAAC,MAAM,KAAK,CAAC,IAAI,YAAY,EAAE;gBACzC;;AAEF,YAAA,MAAM,aAAa,GAAGA,UAAQ,CAAC,GAAG,EAAE;AACpC,YAAA,IAAI,CAAC,aAAa;gBAAE;YAEpB,MAAM,UAAU,GAAG,kBAAkB,CAAC,YAAY,CAAC,IAAI,CAAC;YAExD,IAAI,CAAC,iBAAiB,GAAG,UAAU,KAAK,sBAAsB,EAAE;AAC9D,gBAAA,OAAO,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC3B,iBAAiB,IAAI,UAAU;;iBAC1B;AACL,gBAAAA,UAAQ,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC5B;;;;QAKN,IAAI,kBAAkB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;AAC5C,YAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,kBAAkB,CAAC;AAExF,YAAA,IAAI,iBAAiB,GAAG,CAAC,EAAE;AACzB,gBAAA,OAAO,GAAG,OAAO,CAAC,KAAK,CAAC,iBAAiB,CAAC;;;;QAK9C,IAAI,YAAY,IAAI,SAAS,CAAC,MAAM,GAAG,CAAC,EAAE;YACxC,OAAO,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,CAAgB,CAAC;YACzCA,UAAQ,CAAC,KAAK,EAAE;;;QAIlB,IAAI,eAAe,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,IAAI,YAAY,EAAE;;YAEzD,MAAM,wBAAwB,GAAG,SAAS,CAAC,MAAM,GAAG,CAAC,IAAI,SAAS,CAAC,SAAS,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC,OAAO,EAAE,KAAK,IAAI;;AAG3G,YAAA,MAAM,mBAAmB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,CAAC;AAE5E,YAAA,IAAI,mBAAmB,IAAI,CAAC,EAAE;AAC5B,gBAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,mBAAmB,CAAC;;gBAGtD,MAAM,gBAAgB,GAAG,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC;oBAC/D,iBAAiB,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,IACjC,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC;;gBAGjE,IAAI,CAAC,gBAAgB,EAAE;;oBAErB,MAAM,cAAc,GAAU,EAAE;;AAGhC,oBAAA,KAAK,IAAI,CAAC,GAAGA,UAAQ,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,EAAE;AAC7C,wBAAA,MAAM,GAAG,GAAGA,UAAQ,CAAC,CAAC,CAAC;AACvB,wBAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,4BAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,gCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,EAAE;AAChE,oCAAA,cAAc,CAAC,IAAI,CAAC,IAAI,CAAC;;oCAEzB;;;AAGJ,4BAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC;AAAE,gCAAA,MAAM;;;;AAKzC,oBAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC,EAAE;;AAE7B,wBAAA,MAAM,kBAAkB,GAAG,YAAY,CAAC,iBAAiB,CAAC;;AAG1D,wBAAA,IAAI,UAAiB;wBAErB,IAAI,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC,EAAE;;AAE5C,4BAAA,MAAM,eAAe,GAAG,iBAAiB,CAAC,OAAO,CAAC,MAAM,CAAC,IAAI,IAC3D,EAAE,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC,CAAC;4BAElE,UAAU,GAAG,CAAC,GAAG,cAAc,EAAE,GAAG,eAAe,CAAC;;AAC/C,6BAAA,IAAI,OAAO,iBAAiB,CAAC,OAAO,KAAK,QAAQ,EAAE;AACxD,4BAAA,UAAU,GAAG;AACX,gCAAA,GAAG,cAAc;gCACjB,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,iBAAiB,CAAC,OAAO;6BAChD;;6BACI;4BACL,UAAU,GAAG,cAAc;;;AAI7B,wBAAA,MAAM,UAAU,GAAG,IAAIC,kBAAS,CAAC;AAC/B,4BAAA,OAAO,EAAE,UAAU;4BACnB,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;4BACtD,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;AACvD,yBAAA,CAAC;;AAGF,wBAAA,MAAM,aAAa,GAAG,YAAY,CAAC,UAAU,CAAC;;AAG9C,wBAAA,iBAAiB,KAAK,aAAa,GAAG,kBAAkB,CAAC;;AAGzD,wBAAA,OAAO,CAAC,mBAAmB,CAAC,GAAG,UAAU;;AAGzC,wBAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;;AAE9C,4BAAA,MAAM,aAAa,GAAG,IAAI,GAAG,EAAkB;;AAG/C,4BAAA,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;AACvC,gCAAA,MAAM,GAAG,GAAG,OAAO,CAAC,CAAC,CAAC;;AAGtB,gCAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,oCAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,wCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,IAAI,IAAI,CAAC,EAAE,EAAE;4CAC3E,aAAa,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,CAAC,CAAC;;;;;AAMnC,gCAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,MAAM,IAAI,cAAc,IAAI,GAAG,IAAI,OAAO,GAAG,CAAC,YAAY,KAAK,QAAQ,EAAE;oCAC7F,aAAa,CAAC,GAAG,CAAC,GAAG,CAAC,YAAY,EAAE,CAAC,CAAC;;;;AAK1C,4BAAA,MAAM,eAAe,GAAG,IAAI,GAAG,EAAU;;AAGzC,4BAAA,IAAI,CAAC,GAAG,OAAO,CAAC,MAAM,GAAG,CAAC;4BAC1B,OAAO,CAAC,GAAG,mBAAmB,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC5E,gCAAA,MAAM,WAAW,GAAG,OAAO,CAAC,CAAC,CAAC;;gCAG9B,IAAI,SAAS,GAAG,IAAI;AAEpB,gCAAA,IAAI,WAAW,CAAC,OAAO,EAAE,KAAK,MAAM,IAAI,cAAc,IAAI,WAAW,IAAI,OAAO,WAAW,CAAC,YAAY,KAAK,QAAQ,EAAE;;oCAErH,MAAM,OAAO,GAAG,aAAa,CAAC,GAAG,CAAC,WAAW,CAAC,YAAY,CAAC;AAC3D,oCAAA,IAAI,OAAO,KAAK,SAAS,IAAI,OAAO,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;;AAE3E,wCAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,wCAAA,eAAe,CAAC,GAAG,CAAC,OAAO,CAAC;AAC5B,wCAAA,iBAAiB,KAAK,YAAY,CAAC,WAAW,CAAC,GAAG,YAAY,CAAC,OAAO,CAAC,OAAO,CAAC,CAAC,CAAC;wCACjF,SAAS,GAAG,KAAK;;;AAEd,qCAAA,IAAI,WAAW,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,CAAC,EAAE;;AAE/E,oCAAA,KAAK,MAAM,IAAI,IAAI,WAAW,CAAC,OAAO,EAAE;AACtC,wCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,IAAI,IAAI,CAAC,EAAE,EAAE;4CAC3E,MAAM,SAAS,GAAG,aAAa,CAAC,GAAG,CAAC,IAAI,CAAC,EAAY,CAAC;AACtD,4CAAA,IAAI,SAAS,KAAK,SAAS,IAAI,SAAS,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,SAAS,CAAC,EAAE;;AAEjF,gDAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,gDAAA,eAAe,CAAC,GAAG,CAAC,SAAS,CAAC;AAC9B,gDAAA,iBAAiB,KAAK,YAAY,CAAC,WAAW,CAAC,GAAG,YAAY,CAAC,OAAO,CAAC,SAAS,CAAC,CAAC,CAAC;gDACnF,SAAS,GAAG,KAAK;gDACjB;;;;;;gCAOR,IAAI,SAAS,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE;AACxC,oCAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,oCAAA,iBAAiB,IAAI,YAAY,CAAC,WAAW,CAAC;;AAGhD,gCAAA,CAAC,EAAE;;;4BAIL,MAAM,aAAa,GAAG,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC;AACvE,4BAAA,KAAK,MAAM,KAAK,IAAI,aAAa,EAAE;AACjC,gCAAA,OAAO,CAAC,MAAM,CAAC,KAAK,EAAE,CAAC,CAAC;;;AAI1B,4BAAA,sBAAsB,GAAG,gBAAgB,GAAG,iBAAiB;;;;;;;;YASrE,IAAI,wBAAwB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;;gBAElD,MAAM,gBAAgB,GAAa,EAAE;AACrC,gBAAA,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;AACvC,oBAAA,MAAM,GAAG,GAAG,OAAO,CAAC,CAAC,CAAC;AACtB,oBAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,EAAE;wBAC1B,MAAM,WAAW,GAAG,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC;4BAC5C,GAAG,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC;wBAExF,IAAI,WAAW,EAAE;AACf,4BAAA,gBAAgB,CAAC,IAAI,CAAC,CAAC,CAAC;;;;;AAM9B,gBAAA,IAAI,gBAAgB,CAAC,MAAM,GAAG,CAAC,EAAE;;AAE/B,oBAAA,MAAM,0BAA0B,GAAG,gBAAgB,CAAC,CAAC,CAAC;AACtD,oBAAA,MAAM,qBAAqB,GAAG,OAAO,CAAC,0BAA0B,CAAC;;AAGjE,oBAAA,OAAO,CAAC,MAAM,CAAC,0BAA0B,EAAE,CAAC,CAAC;;AAG7C,oBAAA,MAAM,WAAW,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,QAAQ,CAAC;AACxE,oBAAA,MAAM,SAAS,GAAG,WAAW,KAAK,EAAE;oBAEpC,IAAI,SAAS,EAAE;;;;;;AAOb,wBAAA,MAAM,SAAS,GAAG,OAAO,CAAC,WAAW,CAAC;AACtC,wBAAA,OAAO,CAAC,MAAM,CAAC,WAAW,EAAE,CAAC,CAAC;AAC9B,wBAAA,OAAO,CAAC,IAAI,CAAC,SAAS,CAAC;;AAGvB,wBAAA,OAAO,CAAC,MAAM,CAAC,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC,EAAE,qBAAqB,CAAC;;yBACvD;;;AAGL,wBAAA,OAAO,CAAC,IAAI,CAAC,qBAAqB,CAAC;;;;;;IAO3C,MAAM,YAAY,GAAGD,UAAQ;AAC7B,IAAA,YAAY,GAAG,YAAY,CAAC,MAAM,GAAG,CAAC;IACtC,sBAAsB,IAAI,iBAAiB;IAE3C,OAAO;QACL,YAAY;QACZ,sBAAsB;AACtB,QAAA,OAAO,EAAE,OAAO,CAAC,OAAO,EAAE;AAC1B,QAAA,gBAAgB,EAAE,YAAY;KAC/B;AACH;AAEA,SAAS,gBAAgB,CAAC,KAAc,EAAA;AACtC,IAAA,OAAO,OAAO,KAAK,KAAK,QAAQ,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,KAAK,GAAG,CAAC;AAChE;AAEM,SAAU,mBAAmB,CAAC,aAAyC,EAAA;IAC3E,MAAM,kBAAkB,GAAG,EAAE,GAAG,aAAa,CAAC,kBAAkB,EAAE;AAClE,IAAA,IAAI,kBAAkB,GAAG,aAAa,CAAC,UAAU;IACjD,IAAI,WAAW,GAAG,CAAC,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;IAEhF,OAAO,SAAS,aAAa,CAAC,MAA2B,EAAA;AAIvD,QAAA,IAAI,YAAuC;AAC3C,QAAA,IAAI,MAAM,CAAC,aAAa,KACtB,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,YAAY;AAC/C,gBACD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB;oBAEvD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,cAAc;uBACrE,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,UAAU,CAAC,CACzE,CACF,CACF,IAAI,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,aAAa,CAAC,EAAE;AACzD,YAAA,YAAY,GAAG,oBAAoB,CAAC,MAAM,CAAC,aAAa,CAAC;AACzD,YAAA,WAAW,GAAG,YAAY,CAAC,YAAY;;AAGzC,QAAA,KAAK,IAAI,CAAC,GAAG,kBAAkB,EAAE,CAAC,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;YAChE,MAAM,OAAO,GAAG,MAAM,CAAC,QAAQ,CAAC,CAAC,CAAC;AAClC,YAAA,IAAI,CAAC,KAAK,kBAAkB,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,IAAI,YAAY,EAAE;AACnF,gBAAA,kBAAkB,CAAC,CAAC,CAAC,GAAG,YAAY,CAAC,aAAa;;AAC7C,iBAAA,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,EAAE;gBAC9C,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,YAAY,CAAC,OAAO,CAAC;AAC3D,gBAAA,WAAW,IAAI,kBAAkB,CAAC,CAAC,CAAC;;;;;;;QAQxC,IAAI,YAAY,EAAE;YAChB,MAAM,gBAAgB,GAAG,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;AACrF,YAAA,MAAM,KAAK,GAAG,YAAY,CAAC,YAAY,GAAG,gBAAgB;AAC1D,YAAA,KAAK,MAAM,GAAG,IAAI,kBAAkB,EAAE;AACpC,gBAAA,kBAAkB,CAAC,GAAG,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,GAAG,CAAC,GAAG,KAAK,CAAC;;;AAIzE,QAAA,kBAAkB,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM;AAC3C,QAAA,IAAI,WAAW,IAAI,aAAa,CAAC,SAAS,EAAE;YAC1C,OAAO,EAAE,OAAO,EAAE,MAAM,CAAC,QAAQ,EAAE,kBAAkB,EAAE;;;AAIzD,QAAA,MAAM,EAAE,OAAO,EAAE,GAAG,2BAA2B,CAAC;YAC9C,gBAAgB,EAAE,aAAa,CAAC,SAAS;YACzC,QAAQ,EAAE,MAAM,CAAC,QAAQ;YACzB,kBAAkB;YAClB,kBAAkB,EAAE,MAAM,CAAC,kBAAkB;YAC7C,eAAe,EAAE,aAAa,CAAC,eAAe;YAC9C,YAAY,EAAE,aAAa,CAAC,YAAY;AACzC,SAAA,CAAC;AAEF,QAAA,OAAO,EAAE,OAAO,EAAE,kBAAkB,EAAE;AACxC,KAAC;AACH;;;;"}
|
|
@@ -62,8 +62,15 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
|
|
|
62
62
|
context = context.slice(requiredTypeIndex);
|
|
63
63
|
}
|
|
64
64
|
}
|
|
65
|
+
// Add system message if it exists
|
|
66
|
+
if (instructions && _messages.length > 0) {
|
|
67
|
+
context.push(_messages[0]);
|
|
68
|
+
messages.shift();
|
|
69
|
+
}
|
|
65
70
|
// Handle thinking mode requirement for Anthropic
|
|
66
71
|
if (thinkingEnabled && context.length > 0 && tokenCounter) {
|
|
72
|
+
// Check if the latest message is an assistant message
|
|
73
|
+
const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';
|
|
67
74
|
// Process only if we have an assistant message in the context
|
|
68
75
|
const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');
|
|
69
76
|
if (firstAssistantIndex >= 0) {
|
|
@@ -73,10 +80,11 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
|
|
|
73
80
|
firstAssistantMsg.content.some(item => item && typeof item === 'object' && item.type === 'thinking');
|
|
74
81
|
// Only proceed if we need to add thinking blocks
|
|
75
82
|
if (!hasThinkingBlock) {
|
|
76
|
-
// Collect thinking blocks from pruned assistant messages
|
|
83
|
+
// Collect thinking blocks from pruned assistant messages, starting from the most recent
|
|
77
84
|
const thinkingBlocks = [];
|
|
78
|
-
// Look through pruned messages for thinking blocks
|
|
79
|
-
for (
|
|
85
|
+
// Look through pruned messages for thinking blocks, starting from the end (most recent)
|
|
86
|
+
for (let i = messages.length - 1; i >= 0; i--) {
|
|
87
|
+
const msg = messages[i];
|
|
80
88
|
if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
|
|
81
89
|
for (const item of msg.content) {
|
|
82
90
|
if (item && typeof item === 'object' && item.type === 'thinking') {
|
|
@@ -123,28 +131,124 @@ function getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, in
|
|
|
123
131
|
context[firstAssistantIndex] = newMessage;
|
|
124
132
|
// If we've exceeded the token limit, we need to prune more messages
|
|
125
133
|
if (currentTokenCount > remainingContextTokens) {
|
|
126
|
-
//
|
|
127
|
-
|
|
134
|
+
// Build a map of tool call IDs to track AI <--> tool message correspondences
|
|
135
|
+
const toolCallIdMap = new Map();
|
|
136
|
+
// Identify tool call IDs in the context
|
|
137
|
+
for (let i = 0; i < context.length; i++) {
|
|
138
|
+
const msg = context[i];
|
|
139
|
+
// Check for tool calls in AI messages
|
|
140
|
+
if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
|
|
141
|
+
for (const item of msg.content) {
|
|
142
|
+
if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
|
|
143
|
+
toolCallIdMap.set(item.id, i);
|
|
144
|
+
}
|
|
145
|
+
}
|
|
146
|
+
}
|
|
147
|
+
// Check for tool messages
|
|
148
|
+
if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {
|
|
149
|
+
toolCallIdMap.set(msg.tool_call_id, i);
|
|
150
|
+
}
|
|
151
|
+
}
|
|
152
|
+
// Track which messages to remove
|
|
153
|
+
const indicesToRemove = new Set();
|
|
154
|
+
// Start removing messages from the end, but preserve AI <--> tool message correspondences
|
|
128
155
|
let i = context.length - 1;
|
|
129
156
|
while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {
|
|
130
157
|
const msgToRemove = context[i];
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
158
|
+
// Check if this is a tool message or has tool calls
|
|
159
|
+
let canRemove = true;
|
|
160
|
+
if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {
|
|
161
|
+
// If this is a tool message, check if we need to keep its corresponding AI message
|
|
162
|
+
const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);
|
|
163
|
+
if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {
|
|
164
|
+
// We need to remove both the tool message and its corresponding AI message
|
|
165
|
+
indicesToRemove.add(i);
|
|
166
|
+
indicesToRemove.add(aiIndex);
|
|
167
|
+
currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));
|
|
168
|
+
canRemove = false;
|
|
169
|
+
}
|
|
170
|
+
}
|
|
171
|
+
else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {
|
|
172
|
+
// If this is an AI message with tool calls, check if we need to keep its corresponding tool messages
|
|
173
|
+
for (const item of msgToRemove.content) {
|
|
174
|
+
if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
|
|
175
|
+
const toolIndex = toolCallIdMap.get(item.id);
|
|
176
|
+
if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {
|
|
177
|
+
// We need to remove both the AI message and its corresponding tool message
|
|
178
|
+
indicesToRemove.add(i);
|
|
179
|
+
indicesToRemove.add(toolIndex);
|
|
180
|
+
currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));
|
|
181
|
+
canRemove = false;
|
|
182
|
+
break;
|
|
183
|
+
}
|
|
184
|
+
}
|
|
185
|
+
}
|
|
186
|
+
}
|
|
187
|
+
// If we can remove this message individually
|
|
188
|
+
if (canRemove && !indicesToRemove.has(i)) {
|
|
189
|
+
indicesToRemove.add(i);
|
|
190
|
+
currentTokenCount -= tokenCounter(msgToRemove);
|
|
191
|
+
}
|
|
134
192
|
i--;
|
|
135
193
|
}
|
|
194
|
+
// Remove messages in reverse order to avoid index shifting
|
|
195
|
+
const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);
|
|
196
|
+
for (const index of sortedIndices) {
|
|
197
|
+
context.splice(index, 1);
|
|
198
|
+
}
|
|
136
199
|
// Update remainingContextTokens to reflect the new token count
|
|
137
200
|
remainingContextTokens = maxContextTokens - currentTokenCount;
|
|
138
201
|
}
|
|
139
202
|
}
|
|
140
203
|
}
|
|
141
204
|
}
|
|
205
|
+
// If the latest message is an assistant message, ensure an assistant message with thinking appears at the end
|
|
206
|
+
// of the context (which will become the beginning after reversal)
|
|
207
|
+
// but maintain system message precedence after reversal
|
|
208
|
+
if (latestMessageIsAssistant && context.length > 0) {
|
|
209
|
+
// Find assistant messages with thinking blocks
|
|
210
|
+
const assistantIndices = [];
|
|
211
|
+
for (let i = 0; i < context.length; i++) {
|
|
212
|
+
const msg = context[i];
|
|
213
|
+
if (msg.getType() === 'ai') {
|
|
214
|
+
const hasThinking = Array.isArray(msg.content) &&
|
|
215
|
+
msg.content.some(item => item && typeof item === 'object' && item.type === 'thinking');
|
|
216
|
+
if (hasThinking) {
|
|
217
|
+
assistantIndices.push(i);
|
|
218
|
+
}
|
|
219
|
+
}
|
|
220
|
+
}
|
|
221
|
+
// If we found assistant messages with thinking blocks
|
|
222
|
+
if (assistantIndices.length > 0) {
|
|
223
|
+
// Get the first assistant message with thinking
|
|
224
|
+
const assistantWithThinkingIndex = assistantIndices[0];
|
|
225
|
+
const assistantWithThinking = context[assistantWithThinkingIndex];
|
|
226
|
+
// Remove it from its current position
|
|
227
|
+
context.splice(assistantWithThinkingIndex, 1);
|
|
228
|
+
// Check if there's a system message in the context
|
|
229
|
+
const systemIndex = context.findIndex(msg => msg.getType() === 'system');
|
|
230
|
+
const hasSystem = systemIndex !== -1;
|
|
231
|
+
if (hasSystem) {
|
|
232
|
+
// We want the system message to be first after reversal
|
|
233
|
+
// This means we need to put it at the end position before reversal
|
|
234
|
+
// And the assistant message should be second after reversal
|
|
235
|
+
// This means we need to put it at the end - 1 position before reversal
|
|
236
|
+
// First, ensure the system message is at the end (will be first after reversal)
|
|
237
|
+
const systemMsg = context[systemIndex];
|
|
238
|
+
context.splice(systemIndex, 1);
|
|
239
|
+
context.push(systemMsg);
|
|
240
|
+
// Then, put the assistant message right before the system message (will be second after reversal)
|
|
241
|
+
context.splice(context.length - 1, 0, assistantWithThinking);
|
|
242
|
+
}
|
|
243
|
+
else {
|
|
244
|
+
// No system message, so we want assistant to be first after reversal
|
|
245
|
+
// This means we need to put it at the end position before reversal
|
|
246
|
+
context.push(assistantWithThinking);
|
|
247
|
+
}
|
|
248
|
+
}
|
|
249
|
+
}
|
|
142
250
|
}
|
|
143
251
|
}
|
|
144
|
-
if (instructions && _messages.length > 0) {
|
|
145
|
-
context.push(_messages[0]);
|
|
146
|
-
messages.shift();
|
|
147
|
-
}
|
|
148
252
|
const prunedMemory = messages;
|
|
149
253
|
summaryIndex = prunedMemory.length - 1;
|
|
150
254
|
remainingContextTokens -= currentTokenCount;
|
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"prune.mjs","sources":["../../../src/messages/prune.ts"],"sourcesContent":["import { AIMessage } from '@langchain/core/messages';\nimport type { BaseMessage, UsageMetadata } from '@langchain/core/messages';\nimport type { TokenCounter } from '@/types/run';\nexport type PruneMessagesFactoryParams = {\n maxTokens: number;\n startIndex: number;\n tokenCounter: TokenCounter;\n indexTokenCountMap: Record<string, number>;\n thinkingEnabled?: boolean;\n};\nexport type PruneMessagesParams = {\n messages: BaseMessage[];\n usageMetadata?: Partial<UsageMetadata>;\n startOnMessageType?: ReturnType<BaseMessage['getType']>;\n}\n\n/**\n * Calculates the total tokens from a single usage object\n * \n * @param usage The usage metadata object containing token information\n * @returns An object containing the total input and output tokens\n */\nfunction calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetadata {\n const baseInputTokens = Number(usage.input_tokens) || 0;\n const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;\n const cacheRead = Number(usage.input_token_details?.cache_read) || 0;\n \n const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;\n const totalOutputTokens = Number(usage.output_tokens) || 0;\n\n return {\n input_tokens: totalInputTokens,\n output_tokens: totalOutputTokens,\n total_tokens: totalInputTokens + totalOutputTokens\n };\n}\n\n/**\n * Processes an array of messages and returns a context of messages that fit within a specified token limit.\n * It iterates over the messages from newest to oldest, adding them to the context until the token limit is reached.\n * \n * @param options Configuration options for processing messages\n * @returns Object containing the message context, remaining tokens, messages not included, and summary index\n */\nfunction getMessagesWithinTokenLimit({\n messages: _messages,\n maxContextTokens,\n indexTokenCountMap,\n startOnMessageType,\n thinkingEnabled,\n tokenCounter,\n}: {\n messages: BaseMessage[];\n maxContextTokens: number;\n indexTokenCountMap: Record<string, number>;\n startOnMessageType?: string;\n thinkingEnabled?: boolean;\n tokenCounter?: TokenCounter;\n}): {\n context: BaseMessage[];\n remainingContextTokens: number;\n messagesToRefine: BaseMessage[];\n summaryIndex: number;\n} {\n // Every reply is primed with <|start|>assistant<|message|>, so we\n // start with 3 tokens for the label after all messages have been counted.\n let summaryIndex = -1;\n let currentTokenCount = 3;\n const instructions = _messages?.[0]?.getType() === 'system' ? _messages[0] : undefined;\n const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;\n let remainingContextTokens = maxContextTokens - instructionsTokenCount;\n const messages = [..._messages];\n let context: BaseMessage[] = [];\n\n if (currentTokenCount < remainingContextTokens) {\n let currentIndex = messages.length;\n while (messages.length > 0 && currentTokenCount < remainingContextTokens && currentIndex > 1) {\n currentIndex--;\n if (messages.length === 1 && instructions) {\n break;\n }\n const poppedMessage = messages.pop();\n if (!poppedMessage) continue;\n \n const tokenCount = indexTokenCountMap[currentIndex] || 0;\n\n if ((currentTokenCount + tokenCount) <= remainingContextTokens) {\n context.push(poppedMessage);\n currentTokenCount += tokenCount;\n } else {\n messages.push(poppedMessage);\n break;\n }\n }\n \n // Handle startOnMessageType requirement\n if (startOnMessageType && context.length > 0) {\n const requiredTypeIndex = context.findIndex(msg => msg.getType() === startOnMessageType);\n \n if (requiredTypeIndex > 0) {\n context = context.slice(requiredTypeIndex);\n }\n }\n \n // Handle thinking mode requirement for Anthropic\n if (thinkingEnabled && context.length > 0 && tokenCounter) {\n // Process only if we have an assistant message in the context\n const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');\n if (firstAssistantIndex >= 0) {\n const firstAssistantMsg = context[firstAssistantIndex];\n \n // Check if the first assistant message already has a thinking block\n const hasThinkingBlock = Array.isArray(firstAssistantMsg.content) && \n firstAssistantMsg.content.some(item => \n item && typeof item === 'object' && item.type === 'thinking');\n \n // Only proceed if we need to add thinking blocks\n if (!hasThinkingBlock) {\n // Collect thinking blocks from pruned assistant messages\n const thinkingBlocks: any[] = [];\n \n // Look through pruned messages for thinking blocks\n for (const msg of messages) {\n if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n for (const item of msg.content) {\n if (item && typeof item === 'object' && item.type === 'thinking') {\n thinkingBlocks.push(item);\n // We only need one thinking block\n break;\n }\n }\n if (thinkingBlocks.length > 0) break; // Stop after finding one thinking block\n }\n }\n \n // If we found thinking blocks, add them to the first assistant message\n if (thinkingBlocks.length > 0) {\n // Calculate token count of original message\n const originalTokenCount = tokenCounter(firstAssistantMsg);\n \n // Create a new content array with thinking blocks at the beginning\n let newContent: any[];\n \n if (Array.isArray(firstAssistantMsg.content)) {\n // Keep the original content (excluding any existing thinking blocks)\n const originalContent = firstAssistantMsg.content.filter(item => \n !(item && typeof item === 'object' && item.type === 'thinking'));\n \n newContent = [...thinkingBlocks, ...originalContent];\n } else if (typeof firstAssistantMsg.content === 'string') {\n newContent = [\n ...thinkingBlocks,\n { type: 'text', text: firstAssistantMsg.content }\n ];\n } else {\n newContent = thinkingBlocks;\n }\n \n // Create a new message with the updated content\n const newMessage = new AIMessage({\n content: newContent,\n additional_kwargs: firstAssistantMsg.additional_kwargs,\n response_metadata: firstAssistantMsg.response_metadata,\n });\n \n // Calculate token count of new message\n const newTokenCount = tokenCounter(newMessage);\n \n // Adjust current token count\n currentTokenCount += (newTokenCount - originalTokenCount);\n \n // Replace the first assistant message\n context[firstAssistantIndex] = newMessage;\n \n // If we've exceeded the token limit, we need to prune more messages\n if (currentTokenCount > remainingContextTokens) {\n // Remove messages from the end of the context until we're under the token limit\n // But make sure to keep the first assistant message with thinking block\n let i = context.length - 1;\n while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {\n const msgToRemove = context[i];\n const msgTokenCount = tokenCounter(msgToRemove);\n context.splice(i, 1);\n currentTokenCount -= msgTokenCount;\n i--;\n }\n \n // Update remainingContextTokens to reflect the new token count\n remainingContextTokens = maxContextTokens - currentTokenCount;\n }\n }\n }\n }\n }\n }\n\n if (instructions && _messages.length > 0) {\n context.push(_messages[0] as BaseMessage);\n messages.shift();\n }\n\n const prunedMemory = messages;\n summaryIndex = prunedMemory.length - 1;\n remainingContextTokens -= currentTokenCount;\n\n return {\n summaryIndex,\n remainingContextTokens,\n context: context.reverse(),\n messagesToRefine: prunedMemory,\n };\n}\n\nfunction checkValidNumber(value: unknown): value is number {\n return typeof value === 'number' && !isNaN(value) && value > 0;\n}\n\nexport function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {\n const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };\n let lastTurnStartIndex = factoryParams.startIndex;\n let totalTokens = (Object.values(indexTokenCountMap)).reduce((a, b) => a + b, 0);\n \n return function pruneMessages(params: PruneMessagesParams): {\n context: BaseMessage[];\n indexTokenCountMap: Record<string, number>;\n } {\n let currentUsage: UsageMetadata | undefined;\n if (params.usageMetadata && (\n checkValidNumber(params.usageMetadata.input_tokens)\n || (\n checkValidNumber(params.usageMetadata.input_token_details)\n && (\n checkValidNumber(params.usageMetadata.input_token_details.cache_creation)\n || checkValidNumber(params.usageMetadata.input_token_details.cache_read)\n )\n )\n ) && checkValidNumber(params.usageMetadata.output_tokens)) {\n currentUsage = calculateTotalTokens(params.usageMetadata);\n totalTokens = currentUsage.total_tokens;\n }\n\n for (let i = lastTurnStartIndex; i < params.messages.length; i++) {\n const message = params.messages[i];\n if (i === lastTurnStartIndex && indexTokenCountMap[i] === undefined && currentUsage) {\n indexTokenCountMap[i] = currentUsage.output_tokens;\n } else if (indexTokenCountMap[i] === undefined) {\n indexTokenCountMap[i] = factoryParams.tokenCounter(message);\n totalTokens += indexTokenCountMap[i];\n }\n }\n\n // If `currentUsage` is defined, we need to distribute the current total tokensto our `indexTokenCountMap`,\n // for all message index keys before `lastTurnStartIndex`, as it has the most accurate count for those messages.\n // We must distribute it in a weighted manner, so that the total token count is equal to `currentUsage.total_tokens`,\n // relative the manually counted tokens in `indexTokenCountMap`.\n if (currentUsage) {\n const totalIndexTokens = Object.values(indexTokenCountMap).reduce((a, b) => a + b, 0);\n const ratio = currentUsage.total_tokens / totalIndexTokens;\n for (const key in indexTokenCountMap) {\n indexTokenCountMap[key] = Math.round(indexTokenCountMap[key] * ratio);\n }\n }\n\n lastTurnStartIndex = params.messages.length;\n if (totalTokens <= factoryParams.maxTokens) {\n return { context: params.messages, indexTokenCountMap };\n }\n\n // Pass the tokenCounter to getMessagesWithinTokenLimit for token recalculation\n const { context } = getMessagesWithinTokenLimit({\n maxContextTokens: factoryParams.maxTokens,\n messages: params.messages,\n indexTokenCountMap,\n startOnMessageType: params.startOnMessageType,\n thinkingEnabled: factoryParams.thinkingEnabled,\n tokenCounter: factoryParams.tokenCounter,\n });\n\n return { context, indexTokenCountMap };\n }\n}\n"],"names":[],"mappings":";;AAgBA;;;;;AAKG;AACH,SAAS,oBAAoB,CAAC,KAA6B,EAAA;IACzD,MAAM,eAAe,GAAG,MAAM,CAAC,KAAK,CAAC,YAAY,CAAC,IAAI,CAAC;AACvD,IAAA,MAAM,aAAa,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,cAAc,CAAC,IAAI,CAAC;AAC5E,IAAA,MAAM,SAAS,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,UAAU,CAAC,IAAI,CAAC;AAEpE,IAAA,MAAM,gBAAgB,GAAG,eAAe,GAAG,aAAa,GAAG,SAAS;IACpE,MAAM,iBAAiB,GAAG,MAAM,CAAC,KAAK,CAAC,aAAa,CAAC,IAAI,CAAC;IAE1D,OAAO;AACL,QAAA,YAAY,EAAE,gBAAgB;AAC9B,QAAA,aAAa,EAAE,iBAAiB;QAChC,YAAY,EAAE,gBAAgB,GAAG;KAClC;AACH;AAEA;;;;;;AAMG;AACH,SAAS,2BAA2B,CAAC,EACnC,QAAQ,EAAE,SAAS,EACnB,gBAAgB,EAChB,kBAAkB,EAClB,kBAAkB,EAClB,eAAe,EACf,YAAY,GAQb,EAAA;;;AAQC,IAAA,IAAI,YAAY,GAAG,EAAE;IACrB,IAAI,iBAAiB,GAAG,CAAC;IACzB,MAAM,YAAY,GAAG,SAAS,GAAG,CAAC,CAAC,EAAE,OAAO,EAAE,KAAK,QAAQ,GAAG,SAAS,CAAC,CAAC,CAAC,GAAG,SAAS;AACtF,IAAA,MAAM,sBAAsB,GAAG,YAAY,IAAI,IAAI,GAAG,kBAAkB,CAAC,CAAC,CAAC,GAAG,CAAC;AAC/E,IAAA,IAAI,sBAAsB,GAAG,gBAAgB,GAAG,sBAAsB;AACtE,IAAA,MAAM,QAAQ,GAAG,CAAC,GAAG,SAAS,CAAC;IAC/B,IAAI,OAAO,GAAkB,EAAE;AAE/B,IAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC9C,QAAA,IAAI,YAAY,GAAG,QAAQ,CAAC,MAAM;AAClC,QAAA,OAAO,QAAQ,CAAC,MAAM,GAAG,CAAC,IAAI,iBAAiB,GAAG,sBAAsB,IAAI,YAAY,GAAG,CAAC,EAAE;AAC5F,YAAA,YAAY,EAAE;YACd,IAAI,QAAQ,CAAC,MAAM,KAAK,CAAC,IAAI,YAAY,EAAE;gBACzC;;AAEF,YAAA,MAAM,aAAa,GAAG,QAAQ,CAAC,GAAG,EAAE;AACpC,YAAA,IAAI,CAAC,aAAa;gBAAE;YAEpB,MAAM,UAAU,GAAG,kBAAkB,CAAC,YAAY,CAAC,IAAI,CAAC;YAExD,IAAI,CAAC,iBAAiB,GAAG,UAAU,KAAK,sBAAsB,EAAE;AAC9D,gBAAA,OAAO,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC3B,iBAAiB,IAAI,UAAU;;iBAC1B;AACL,gBAAA,QAAQ,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC5B;;;;QAKN,IAAI,kBAAkB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;AAC5C,YAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,kBAAkB,CAAC;AAExF,YAAA,IAAI,iBAAiB,GAAG,CAAC,EAAE;AACzB,gBAAA,OAAO,GAAG,OAAO,CAAC,KAAK,CAAC,iBAAiB,CAAC;;;;QAK9C,IAAI,eAAe,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,IAAI,YAAY,EAAE;;AAEzD,YAAA,MAAM,mBAAmB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,CAAC;AAC5E,YAAA,IAAI,mBAAmB,IAAI,CAAC,EAAE;AAC5B,gBAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,mBAAmB,CAAC;;gBAGtD,MAAM,gBAAgB,GAAG,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC;oBAC/D,iBAAiB,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,IACjC,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC;;gBAGjE,IAAI,CAAC,gBAAgB,EAAE;;oBAErB,MAAM,cAAc,GAAU,EAAE;;AAGhC,oBAAA,KAAK,MAAM,GAAG,IAAI,QAAQ,EAAE;AAC1B,wBAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,4BAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,gCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,EAAE;AAChE,oCAAA,cAAc,CAAC,IAAI,CAAC,IAAI,CAAC;;oCAEzB;;;AAGJ,4BAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC;AAAE,gCAAA,MAAM;;;;AAKzC,oBAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC,EAAE;;AAE7B,wBAAA,MAAM,kBAAkB,GAAG,YAAY,CAAC,iBAAiB,CAAC;;AAG1D,wBAAA,IAAI,UAAiB;wBAErB,IAAI,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC,EAAE;;AAE5C,4BAAA,MAAM,eAAe,GAAG,iBAAiB,CAAC,OAAO,CAAC,MAAM,CAAC,IAAI,IAC3D,EAAE,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC,CAAC;4BAElE,UAAU,GAAG,CAAC,GAAG,cAAc,EAAE,GAAG,eAAe,CAAC;;AAC/C,6BAAA,IAAI,OAAO,iBAAiB,CAAC,OAAO,KAAK,QAAQ,EAAE;AACxD,4BAAA,UAAU,GAAG;AACX,gCAAA,GAAG,cAAc;gCACjB,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,iBAAiB,CAAC,OAAO;6BAChD;;6BACI;4BACL,UAAU,GAAG,cAAc;;;AAI7B,wBAAA,MAAM,UAAU,GAAG,IAAI,SAAS,CAAC;AAC/B,4BAAA,OAAO,EAAE,UAAU;4BACnB,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;4BACtD,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;AACvD,yBAAA,CAAC;;AAGF,wBAAA,MAAM,aAAa,GAAG,YAAY,CAAC,UAAU,CAAC;;AAG9C,wBAAA,iBAAiB,KAAK,aAAa,GAAG,kBAAkB,CAAC;;AAGzD,wBAAA,OAAO,CAAC,mBAAmB,CAAC,GAAG,UAAU;;AAGzC,wBAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;;;AAG9C,4BAAA,IAAI,CAAC,GAAG,OAAO,CAAC,MAAM,GAAG,CAAC;4BAC1B,OAAO,CAAC,GAAG,mBAAmB,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC5E,gCAAA,MAAM,WAAW,GAAG,OAAO,CAAC,CAAC,CAAC;AAC9B,gCAAA,MAAM,aAAa,GAAG,YAAY,CAAC,WAAW,CAAC;AAC/C,gCAAA,OAAO,CAAC,MAAM,CAAC,CAAC,EAAE,CAAC,CAAC;gCACpB,iBAAiB,IAAI,aAAa;AAClC,gCAAA,CAAC,EAAE;;;AAIL,4BAAA,sBAAsB,GAAG,gBAAgB,GAAG,iBAAiB;;;;;;;IAQvE,IAAI,YAAY,IAAI,SAAS,CAAC,MAAM,GAAG,CAAC,EAAE;QACxC,OAAO,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,CAAgB,CAAC;QACzC,QAAQ,CAAC,KAAK,EAAE;;IAGlB,MAAM,YAAY,GAAG,QAAQ;AAC7B,IAAA,YAAY,GAAG,YAAY,CAAC,MAAM,GAAG,CAAC;IACtC,sBAAsB,IAAI,iBAAiB;IAE3C,OAAO;QACL,YAAY;QACZ,sBAAsB;AACtB,QAAA,OAAO,EAAE,OAAO,CAAC,OAAO,EAAE;AAC1B,QAAA,gBAAgB,EAAE,YAAY;KAC/B;AACH;AAEA,SAAS,gBAAgB,CAAC,KAAc,EAAA;AACtC,IAAA,OAAO,OAAO,KAAK,KAAK,QAAQ,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,KAAK,GAAG,CAAC;AAChE;AAEM,SAAU,mBAAmB,CAAC,aAAyC,EAAA;IAC3E,MAAM,kBAAkB,GAAG,EAAE,GAAG,aAAa,CAAC,kBAAkB,EAAE;AAClE,IAAA,IAAI,kBAAkB,GAAG,aAAa,CAAC,UAAU;IACjD,IAAI,WAAW,GAAG,CAAC,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;IAEhF,OAAO,SAAS,aAAa,CAAC,MAA2B,EAAA;AAIvD,QAAA,IAAI,YAAuC;AAC3C,QAAA,IAAI,MAAM,CAAC,aAAa,KACtB,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,YAAY;AAC/C,gBACD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB;oBAEvD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,cAAc;uBACrE,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,UAAU,CAAC,CACzE,CACF,CACF,IAAI,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,aAAa,CAAC,EAAE;AACzD,YAAA,YAAY,GAAG,oBAAoB,CAAC,MAAM,CAAC,aAAa,CAAC;AACzD,YAAA,WAAW,GAAG,YAAY,CAAC,YAAY;;AAGzC,QAAA,KAAK,IAAI,CAAC,GAAG,kBAAkB,EAAE,CAAC,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;YAChE,MAAM,OAAO,GAAG,MAAM,CAAC,QAAQ,CAAC,CAAC,CAAC;AAClC,YAAA,IAAI,CAAC,KAAK,kBAAkB,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,IAAI,YAAY,EAAE;AACnF,gBAAA,kBAAkB,CAAC,CAAC,CAAC,GAAG,YAAY,CAAC,aAAa;;AAC7C,iBAAA,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,EAAE;gBAC9C,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,YAAY,CAAC,OAAO,CAAC;AAC3D,gBAAA,WAAW,IAAI,kBAAkB,CAAC,CAAC,CAAC;;;;;;;QAQxC,IAAI,YAAY,EAAE;YAChB,MAAM,gBAAgB,GAAG,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;AACrF,YAAA,MAAM,KAAK,GAAG,YAAY,CAAC,YAAY,GAAG,gBAAgB;AAC1D,YAAA,KAAK,MAAM,GAAG,IAAI,kBAAkB,EAAE;AACpC,gBAAA,kBAAkB,CAAC,GAAG,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,GAAG,CAAC,GAAG,KAAK,CAAC;;;AAIzE,QAAA,kBAAkB,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM;AAC3C,QAAA,IAAI,WAAW,IAAI,aAAa,CAAC,SAAS,EAAE;YAC1C,OAAO,EAAE,OAAO,EAAE,MAAM,CAAC,QAAQ,EAAE,kBAAkB,EAAE;;;AAIzD,QAAA,MAAM,EAAE,OAAO,EAAE,GAAG,2BAA2B,CAAC;YAC9C,gBAAgB,EAAE,aAAa,CAAC,SAAS;YACzC,QAAQ,EAAE,MAAM,CAAC,QAAQ;YACzB,kBAAkB;YAClB,kBAAkB,EAAE,MAAM,CAAC,kBAAkB;YAC7C,eAAe,EAAE,aAAa,CAAC,eAAe;YAC9C,YAAY,EAAE,aAAa,CAAC,YAAY;AACzC,SAAA,CAAC;AAEF,QAAA,OAAO,EAAE,OAAO,EAAE,kBAAkB,EAAE;AACxC,KAAC;AACH;;;;"}
|
|
1
|
+
{"version":3,"file":"prune.mjs","sources":["../../../src/messages/prune.ts"],"sourcesContent":["import { AIMessage } from '@langchain/core/messages';\nimport type { BaseMessage, UsageMetadata } from '@langchain/core/messages';\nimport type { TokenCounter } from '@/types/run';\nexport type PruneMessagesFactoryParams = {\n maxTokens: number;\n startIndex: number;\n tokenCounter: TokenCounter;\n indexTokenCountMap: Record<string, number>;\n thinkingEnabled?: boolean;\n};\nexport type PruneMessagesParams = {\n messages: BaseMessage[];\n usageMetadata?: Partial<UsageMetadata>;\n startOnMessageType?: ReturnType<BaseMessage['getType']>;\n}\n\n/**\n * Calculates the total tokens from a single usage object\n * \n * @param usage The usage metadata object containing token information\n * @returns An object containing the total input and output tokens\n */\nfunction calculateTotalTokens(usage: Partial<UsageMetadata>): UsageMetadata {\n const baseInputTokens = Number(usage.input_tokens) || 0;\n const cacheCreation = Number(usage.input_token_details?.cache_creation) || 0;\n const cacheRead = Number(usage.input_token_details?.cache_read) || 0;\n \n const totalInputTokens = baseInputTokens + cacheCreation + cacheRead;\n const totalOutputTokens = Number(usage.output_tokens) || 0;\n\n return {\n input_tokens: totalInputTokens,\n output_tokens: totalOutputTokens,\n total_tokens: totalInputTokens + totalOutputTokens\n };\n}\n\n/**\n * Processes an array of messages and returns a context of messages that fit within a specified token limit.\n * It iterates over the messages from newest to oldest, adding them to the context until the token limit is reached.\n * \n * @param options Configuration options for processing messages\n * @returns Object containing the message context, remaining tokens, messages not included, and summary index\n */\nfunction getMessagesWithinTokenLimit({\n messages: _messages,\n maxContextTokens,\n indexTokenCountMap,\n startOnMessageType,\n thinkingEnabled,\n tokenCounter,\n}: {\n messages: BaseMessage[];\n maxContextTokens: number;\n indexTokenCountMap: Record<string, number>;\n startOnMessageType?: string;\n thinkingEnabled?: boolean;\n tokenCounter?: TokenCounter;\n}): {\n context: BaseMessage[];\n remainingContextTokens: number;\n messagesToRefine: BaseMessage[];\n summaryIndex: number;\n} {\n // Every reply is primed with <|start|>assistant<|message|>, so we\n // start with 3 tokens for the label after all messages have been counted.\n let summaryIndex = -1;\n let currentTokenCount = 3;\n const instructions = _messages?.[0]?.getType() === 'system' ? _messages[0] : undefined;\n const instructionsTokenCount = instructions != null ? indexTokenCountMap[0] : 0;\n let remainingContextTokens = maxContextTokens - instructionsTokenCount;\n const messages = [..._messages];\n let context: BaseMessage[] = [];\n\n if (currentTokenCount < remainingContextTokens) {\n let currentIndex = messages.length;\n while (messages.length > 0 && currentTokenCount < remainingContextTokens && currentIndex > 1) {\n currentIndex--;\n if (messages.length === 1 && instructions) {\n break;\n }\n const poppedMessage = messages.pop();\n if (!poppedMessage) continue;\n \n const tokenCount = indexTokenCountMap[currentIndex] || 0;\n\n if ((currentTokenCount + tokenCount) <= remainingContextTokens) {\n context.push(poppedMessage);\n currentTokenCount += tokenCount;\n } else {\n messages.push(poppedMessage);\n break;\n }\n }\n \n // Handle startOnMessageType requirement\n if (startOnMessageType && context.length > 0) {\n const requiredTypeIndex = context.findIndex(msg => msg.getType() === startOnMessageType);\n \n if (requiredTypeIndex > 0) {\n context = context.slice(requiredTypeIndex);\n }\n }\n \n // Add system message if it exists\n if (instructions && _messages.length > 0) {\n context.push(_messages[0] as BaseMessage);\n messages.shift();\n }\n \n // Handle thinking mode requirement for Anthropic\n if (thinkingEnabled && context.length > 0 && tokenCounter) {\n // Check if the latest message is an assistant message\n const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';\n \n // Process only if we have an assistant message in the context\n const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');\n \n if (firstAssistantIndex >= 0) {\n const firstAssistantMsg = context[firstAssistantIndex];\n \n // Check if the first assistant message already has a thinking block\n const hasThinkingBlock = Array.isArray(firstAssistantMsg.content) && \n firstAssistantMsg.content.some(item => \n item && typeof item === 'object' && item.type === 'thinking');\n \n // Only proceed if we need to add thinking blocks\n if (!hasThinkingBlock) {\n // Collect thinking blocks from pruned assistant messages, starting from the most recent\n const thinkingBlocks: any[] = [];\n \n // Look through pruned messages for thinking blocks, starting from the end (most recent)\n for (let i = messages.length - 1; i >= 0; i--) {\n const msg = messages[i];\n if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n for (const item of msg.content) {\n if (item && typeof item === 'object' && item.type === 'thinking') {\n thinkingBlocks.push(item);\n // We only need one thinking block\n break;\n }\n }\n if (thinkingBlocks.length > 0) break; // Stop after finding one thinking block\n }\n }\n \n // If we found thinking blocks, add them to the first assistant message\n if (thinkingBlocks.length > 0) {\n // Calculate token count of original message\n const originalTokenCount = tokenCounter(firstAssistantMsg);\n \n // Create a new content array with thinking blocks at the beginning\n let newContent: any[];\n \n if (Array.isArray(firstAssistantMsg.content)) {\n // Keep the original content (excluding any existing thinking blocks)\n const originalContent = firstAssistantMsg.content.filter(item => \n !(item && typeof item === 'object' && item.type === 'thinking'));\n \n newContent = [...thinkingBlocks, ...originalContent];\n } else if (typeof firstAssistantMsg.content === 'string') {\n newContent = [\n ...thinkingBlocks,\n { type: 'text', text: firstAssistantMsg.content }\n ];\n } else {\n newContent = thinkingBlocks;\n }\n \n // Create a new message with the updated content\n const newMessage = new AIMessage({\n content: newContent,\n additional_kwargs: firstAssistantMsg.additional_kwargs,\n response_metadata: firstAssistantMsg.response_metadata,\n });\n \n // Calculate token count of new message\n const newTokenCount = tokenCounter(newMessage);\n \n // Adjust current token count\n currentTokenCount += (newTokenCount - originalTokenCount);\n \n // Replace the first assistant message\n context[firstAssistantIndex] = newMessage;\n \n // If we've exceeded the token limit, we need to prune more messages\n if (currentTokenCount > remainingContextTokens) {\n // Build a map of tool call IDs to track AI <--> tool message correspondences\n const toolCallIdMap = new Map<string, number>();\n \n // Identify tool call IDs in the context\n for (let i = 0; i < context.length; i++) {\n const msg = context[i];\n \n // Check for tool calls in AI messages\n if (msg.getType() === 'ai' && Array.isArray(msg.content)) {\n for (const item of msg.content) {\n if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {\n toolCallIdMap.set(item.id, i);\n }\n }\n }\n \n // Check for tool messages\n if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {\n toolCallIdMap.set(msg.tool_call_id, i);\n }\n }\n \n // Track which messages to remove\n const indicesToRemove = new Set<number>();\n \n // Start removing messages from the end, but preserve AI <--> tool message correspondences\n let i = context.length - 1;\n while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {\n const msgToRemove = context[i];\n \n // Check if this is a tool message or has tool calls\n let canRemove = true;\n \n if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {\n // If this is a tool message, check if we need to keep its corresponding AI message\n const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);\n if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {\n // We need to remove both the tool message and its corresponding AI message\n indicesToRemove.add(i);\n indicesToRemove.add(aiIndex);\n currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));\n canRemove = false;\n }\n } else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {\n // If this is an AI message with tool calls, check if we need to keep its corresponding tool messages\n for (const item of msgToRemove.content) {\n if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {\n const toolIndex = toolCallIdMap.get(item.id as string);\n if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {\n // We need to remove both the AI message and its corresponding tool message\n indicesToRemove.add(i);\n indicesToRemove.add(toolIndex);\n currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));\n canRemove = false;\n break;\n }\n }\n }\n }\n \n // If we can remove this message individually\n if (canRemove && !indicesToRemove.has(i)) {\n indicesToRemove.add(i);\n currentTokenCount -= tokenCounter(msgToRemove);\n }\n \n i--;\n }\n \n // Remove messages in reverse order to avoid index shifting\n const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);\n for (const index of sortedIndices) {\n context.splice(index, 1);\n }\n \n // Update remainingContextTokens to reflect the new token count\n remainingContextTokens = maxContextTokens - currentTokenCount;\n }\n }\n }\n }\n \n // If the latest message is an assistant message, ensure an assistant message with thinking appears at the end\n // of the context (which will become the beginning after reversal)\n // but maintain system message precedence after reversal\n if (latestMessageIsAssistant && context.length > 0) {\n // Find assistant messages with thinking blocks\n const assistantIndices: number[] = [];\n for (let i = 0; i < context.length; i++) {\n const msg = context[i];\n if (msg.getType() === 'ai') {\n const hasThinking = Array.isArray(msg.content) && \n msg.content.some(item => item && typeof item === 'object' && item.type === 'thinking');\n \n if (hasThinking) {\n assistantIndices.push(i);\n }\n }\n }\n \n // If we found assistant messages with thinking blocks\n if (assistantIndices.length > 0) {\n // Get the first assistant message with thinking\n const assistantWithThinkingIndex = assistantIndices[0];\n const assistantWithThinking = context[assistantWithThinkingIndex];\n \n // Remove it from its current position\n context.splice(assistantWithThinkingIndex, 1);\n \n // Check if there's a system message in the context\n const systemIndex = context.findIndex(msg => msg.getType() === 'system');\n const hasSystem = systemIndex !== -1;\n \n if (hasSystem) {\n // We want the system message to be first after reversal\n // This means we need to put it at the end position before reversal\n // And the assistant message should be second after reversal\n // This means we need to put it at the end - 1 position before reversal\n \n // First, ensure the system message is at the end (will be first after reversal)\n const systemMsg = context[systemIndex];\n context.splice(systemIndex, 1);\n context.push(systemMsg);\n \n // Then, put the assistant message right before the system message (will be second after reversal)\n context.splice(context.length - 1, 0, assistantWithThinking);\n } else {\n // No system message, so we want assistant to be first after reversal\n // This means we need to put it at the end position before reversal\n context.push(assistantWithThinking);\n }\n }\n }\n }\n }\n\n const prunedMemory = messages;\n summaryIndex = prunedMemory.length - 1;\n remainingContextTokens -= currentTokenCount;\n\n return {\n summaryIndex,\n remainingContextTokens,\n context: context.reverse(),\n messagesToRefine: prunedMemory,\n };\n}\n\nfunction checkValidNumber(value: unknown): value is number {\n return typeof value === 'number' && !isNaN(value) && value > 0;\n}\n\nexport function createPruneMessages(factoryParams: PruneMessagesFactoryParams) {\n const indexTokenCountMap = { ...factoryParams.indexTokenCountMap };\n let lastTurnStartIndex = factoryParams.startIndex;\n let totalTokens = (Object.values(indexTokenCountMap)).reduce((a, b) => a + b, 0);\n \n return function pruneMessages(params: PruneMessagesParams): {\n context: BaseMessage[];\n indexTokenCountMap: Record<string, number>;\n } {\n let currentUsage: UsageMetadata | undefined;\n if (params.usageMetadata && (\n checkValidNumber(params.usageMetadata.input_tokens)\n || (\n checkValidNumber(params.usageMetadata.input_token_details)\n && (\n checkValidNumber(params.usageMetadata.input_token_details.cache_creation)\n || checkValidNumber(params.usageMetadata.input_token_details.cache_read)\n )\n )\n ) && checkValidNumber(params.usageMetadata.output_tokens)) {\n currentUsage = calculateTotalTokens(params.usageMetadata);\n totalTokens = currentUsage.total_tokens;\n }\n\n for (let i = lastTurnStartIndex; i < params.messages.length; i++) {\n const message = params.messages[i];\n if (i === lastTurnStartIndex && indexTokenCountMap[i] === undefined && currentUsage) {\n indexTokenCountMap[i] = currentUsage.output_tokens;\n } else if (indexTokenCountMap[i] === undefined) {\n indexTokenCountMap[i] = factoryParams.tokenCounter(message);\n totalTokens += indexTokenCountMap[i];\n }\n }\n\n // If `currentUsage` is defined, we need to distribute the current total tokensto our `indexTokenCountMap`,\n // for all message index keys before `lastTurnStartIndex`, as it has the most accurate count for those messages.\n // We must distribute it in a weighted manner, so that the total token count is equal to `currentUsage.total_tokens`,\n // relative the manually counted tokens in `indexTokenCountMap`.\n if (currentUsage) {\n const totalIndexTokens = Object.values(indexTokenCountMap).reduce((a, b) => a + b, 0);\n const ratio = currentUsage.total_tokens / totalIndexTokens;\n for (const key in indexTokenCountMap) {\n indexTokenCountMap[key] = Math.round(indexTokenCountMap[key] * ratio);\n }\n }\n\n lastTurnStartIndex = params.messages.length;\n if (totalTokens <= factoryParams.maxTokens) {\n return { context: params.messages, indexTokenCountMap };\n }\n\n // Pass the tokenCounter to getMessagesWithinTokenLimit for token recalculation\n const { context } = getMessagesWithinTokenLimit({\n maxContextTokens: factoryParams.maxTokens,\n messages: params.messages,\n indexTokenCountMap,\n startOnMessageType: params.startOnMessageType,\n thinkingEnabled: factoryParams.thinkingEnabled,\n tokenCounter: factoryParams.tokenCounter,\n });\n\n return { context, indexTokenCountMap };\n }\n}\n"],"names":[],"mappings":";;AAgBA;;;;;AAKG;AACH,SAAS,oBAAoB,CAAC,KAA6B,EAAA;IACzD,MAAM,eAAe,GAAG,MAAM,CAAC,KAAK,CAAC,YAAY,CAAC,IAAI,CAAC;AACvD,IAAA,MAAM,aAAa,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,cAAc,CAAC,IAAI,CAAC;AAC5E,IAAA,MAAM,SAAS,GAAG,MAAM,CAAC,KAAK,CAAC,mBAAmB,EAAE,UAAU,CAAC,IAAI,CAAC;AAEpE,IAAA,MAAM,gBAAgB,GAAG,eAAe,GAAG,aAAa,GAAG,SAAS;IACpE,MAAM,iBAAiB,GAAG,MAAM,CAAC,KAAK,CAAC,aAAa,CAAC,IAAI,CAAC;IAE1D,OAAO;AACL,QAAA,YAAY,EAAE,gBAAgB;AAC9B,QAAA,aAAa,EAAE,iBAAiB;QAChC,YAAY,EAAE,gBAAgB,GAAG;KAClC;AACH;AAEA;;;;;;AAMG;AACH,SAAS,2BAA2B,CAAC,EACnC,QAAQ,EAAE,SAAS,EACnB,gBAAgB,EAChB,kBAAkB,EAClB,kBAAkB,EAClB,eAAe,EACf,YAAY,GAQb,EAAA;;;AAQC,IAAA,IAAI,YAAY,GAAG,EAAE;IACrB,IAAI,iBAAiB,GAAG,CAAC;IACzB,MAAM,YAAY,GAAG,SAAS,GAAG,CAAC,CAAC,EAAE,OAAO,EAAE,KAAK,QAAQ,GAAG,SAAS,CAAC,CAAC,CAAC,GAAG,SAAS;AACtF,IAAA,MAAM,sBAAsB,GAAG,YAAY,IAAI,IAAI,GAAG,kBAAkB,CAAC,CAAC,CAAC,GAAG,CAAC;AAC/E,IAAA,IAAI,sBAAsB,GAAG,gBAAgB,GAAG,sBAAsB;AACtE,IAAA,MAAM,QAAQ,GAAG,CAAC,GAAG,SAAS,CAAC;IAC/B,IAAI,OAAO,GAAkB,EAAE;AAE/B,IAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC9C,QAAA,IAAI,YAAY,GAAG,QAAQ,CAAC,MAAM;AAClC,QAAA,OAAO,QAAQ,CAAC,MAAM,GAAG,CAAC,IAAI,iBAAiB,GAAG,sBAAsB,IAAI,YAAY,GAAG,CAAC,EAAE;AAC5F,YAAA,YAAY,EAAE;YACd,IAAI,QAAQ,CAAC,MAAM,KAAK,CAAC,IAAI,YAAY,EAAE;gBACzC;;AAEF,YAAA,MAAM,aAAa,GAAG,QAAQ,CAAC,GAAG,EAAE;AACpC,YAAA,IAAI,CAAC,aAAa;gBAAE;YAEpB,MAAM,UAAU,GAAG,kBAAkB,CAAC,YAAY,CAAC,IAAI,CAAC;YAExD,IAAI,CAAC,iBAAiB,GAAG,UAAU,KAAK,sBAAsB,EAAE;AAC9D,gBAAA,OAAO,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC3B,iBAAiB,IAAI,UAAU;;iBAC1B;AACL,gBAAA,QAAQ,CAAC,IAAI,CAAC,aAAa,CAAC;gBAC5B;;;;QAKN,IAAI,kBAAkB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;AAC5C,YAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,kBAAkB,CAAC;AAExF,YAAA,IAAI,iBAAiB,GAAG,CAAC,EAAE;AACzB,gBAAA,OAAO,GAAG,OAAO,CAAC,KAAK,CAAC,iBAAiB,CAAC;;;;QAK9C,IAAI,YAAY,IAAI,SAAS,CAAC,MAAM,GAAG,CAAC,EAAE;YACxC,OAAO,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,CAAgB,CAAC;YACzC,QAAQ,CAAC,KAAK,EAAE;;;QAIlB,IAAI,eAAe,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,IAAI,YAAY,EAAE;;YAEzD,MAAM,wBAAwB,GAAG,SAAS,CAAC,MAAM,GAAG,CAAC,IAAI,SAAS,CAAC,SAAS,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC,OAAO,EAAE,KAAK,IAAI;;AAG3G,YAAA,MAAM,mBAAmB,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,CAAC;AAE5E,YAAA,IAAI,mBAAmB,IAAI,CAAC,EAAE;AAC5B,gBAAA,MAAM,iBAAiB,GAAG,OAAO,CAAC,mBAAmB,CAAC;;gBAGtD,MAAM,gBAAgB,GAAG,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC;oBAC/D,iBAAiB,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,IACjC,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC;;gBAGjE,IAAI,CAAC,gBAAgB,EAAE;;oBAErB,MAAM,cAAc,GAAU,EAAE;;AAGhC,oBAAA,KAAK,IAAI,CAAC,GAAG,QAAQ,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,EAAE;AAC7C,wBAAA,MAAM,GAAG,GAAG,QAAQ,CAAC,CAAC,CAAC;AACvB,wBAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,4BAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,gCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,EAAE;AAChE,oCAAA,cAAc,CAAC,IAAI,CAAC,IAAI,CAAC;;oCAEzB;;;AAGJ,4BAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC;AAAE,gCAAA,MAAM;;;;AAKzC,oBAAA,IAAI,cAAc,CAAC,MAAM,GAAG,CAAC,EAAE;;AAE7B,wBAAA,MAAM,kBAAkB,GAAG,YAAY,CAAC,iBAAiB,CAAC;;AAG1D,wBAAA,IAAI,UAAiB;wBAErB,IAAI,KAAK,CAAC,OAAO,CAAC,iBAAiB,CAAC,OAAO,CAAC,EAAE;;AAE5C,4BAAA,MAAM,eAAe,GAAG,iBAAiB,CAAC,OAAO,CAAC,MAAM,CAAC,IAAI,IAC3D,EAAE,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC,CAAC;4BAElE,UAAU,GAAG,CAAC,GAAG,cAAc,EAAE,GAAG,eAAe,CAAC;;AAC/C,6BAAA,IAAI,OAAO,iBAAiB,CAAC,OAAO,KAAK,QAAQ,EAAE;AACxD,4BAAA,UAAU,GAAG;AACX,gCAAA,GAAG,cAAc;gCACjB,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,iBAAiB,CAAC,OAAO;6BAChD;;6BACI;4BACL,UAAU,GAAG,cAAc;;;AAI7B,wBAAA,MAAM,UAAU,GAAG,IAAI,SAAS,CAAC;AAC/B,4BAAA,OAAO,EAAE,UAAU;4BACnB,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;4BACtD,iBAAiB,EAAE,iBAAiB,CAAC,iBAAiB;AACvD,yBAAA,CAAC;;AAGF,wBAAA,MAAM,aAAa,GAAG,YAAY,CAAC,UAAU,CAAC;;AAG9C,wBAAA,iBAAiB,KAAK,aAAa,GAAG,kBAAkB,CAAC;;AAGzD,wBAAA,OAAO,CAAC,mBAAmB,CAAC,GAAG,UAAU;;AAGzC,wBAAA,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;;AAE9C,4BAAA,MAAM,aAAa,GAAG,IAAI,GAAG,EAAkB;;AAG/C,4BAAA,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;AACvC,gCAAA,MAAM,GAAG,GAAG,OAAO,CAAC,CAAC,CAAC;;AAGtB,gCAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;AACxD,oCAAA,KAAK,MAAM,IAAI,IAAI,GAAG,CAAC,OAAO,EAAE;AAC9B,wCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,IAAI,IAAI,CAAC,EAAE,EAAE;4CAC3E,aAAa,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,CAAC,CAAC;;;;;AAMnC,gCAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,MAAM,IAAI,cAAc,IAAI,GAAG,IAAI,OAAO,GAAG,CAAC,YAAY,KAAK,QAAQ,EAAE;oCAC7F,aAAa,CAAC,GAAG,CAAC,GAAG,CAAC,YAAY,EAAE,CAAC,CAAC;;;;AAK1C,4BAAA,MAAM,eAAe,GAAG,IAAI,GAAG,EAAU;;AAGzC,4BAAA,IAAI,CAAC,GAAG,OAAO,CAAC,MAAM,GAAG,CAAC;4BAC1B,OAAO,CAAC,GAAG,mBAAmB,IAAI,iBAAiB,GAAG,sBAAsB,EAAE;AAC5E,gCAAA,MAAM,WAAW,GAAG,OAAO,CAAC,CAAC,CAAC;;gCAG9B,IAAI,SAAS,GAAG,IAAI;AAEpB,gCAAA,IAAI,WAAW,CAAC,OAAO,EAAE,KAAK,MAAM,IAAI,cAAc,IAAI,WAAW,IAAI,OAAO,WAAW,CAAC,YAAY,KAAK,QAAQ,EAAE;;oCAErH,MAAM,OAAO,GAAG,aAAa,CAAC,GAAG,CAAC,WAAW,CAAC,YAAY,CAAC;AAC3D,oCAAA,IAAI,OAAO,KAAK,SAAS,IAAI,OAAO,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE;;AAE3E,wCAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,wCAAA,eAAe,CAAC,GAAG,CAAC,OAAO,CAAC;AAC5B,wCAAA,iBAAiB,KAAK,YAAY,CAAC,WAAW,CAAC,GAAG,YAAY,CAAC,OAAO,CAAC,OAAO,CAAC,CAAC,CAAC;wCACjF,SAAS,GAAG,KAAK;;;AAEd,qCAAA,IAAI,WAAW,CAAC,OAAO,EAAE,KAAK,IAAI,IAAI,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,OAAO,CAAC,EAAE;;AAE/E,oCAAA,KAAK,MAAM,IAAI,IAAI,WAAW,CAAC,OAAO,EAAE;AACtC,wCAAA,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,IAAI,IAAI,CAAC,EAAE,EAAE;4CAC3E,MAAM,SAAS,GAAG,aAAa,CAAC,GAAG,CAAC,IAAI,CAAC,EAAY,CAAC;AACtD,4CAAA,IAAI,SAAS,KAAK,SAAS,IAAI,SAAS,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,SAAS,CAAC,EAAE;;AAEjF,gDAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,gDAAA,eAAe,CAAC,GAAG,CAAC,SAAS,CAAC;AAC9B,gDAAA,iBAAiB,KAAK,YAAY,CAAC,WAAW,CAAC,GAAG,YAAY,CAAC,OAAO,CAAC,SAAS,CAAC,CAAC,CAAC;gDACnF,SAAS,GAAG,KAAK;gDACjB;;;;;;gCAOR,IAAI,SAAS,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE;AACxC,oCAAA,eAAe,CAAC,GAAG,CAAC,CAAC,CAAC;AACtB,oCAAA,iBAAiB,IAAI,YAAY,CAAC,WAAW,CAAC;;AAGhD,gCAAA,CAAC,EAAE;;;4BAIL,MAAM,aAAa,GAAG,KAAK,CAAC,IAAI,CAAC,eAAe,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC;AACvE,4BAAA,KAAK,MAAM,KAAK,IAAI,aAAa,EAAE;AACjC,gCAAA,OAAO,CAAC,MAAM,CAAC,KAAK,EAAE,CAAC,CAAC;;;AAI1B,4BAAA,sBAAsB,GAAG,gBAAgB,GAAG,iBAAiB;;;;;;;;YASrE,IAAI,wBAAwB,IAAI,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE;;gBAElD,MAAM,gBAAgB,GAAa,EAAE;AACrC,gBAAA,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;AACvC,oBAAA,MAAM,GAAG,GAAG,OAAO,CAAC,CAAC,CAAC;AACtB,oBAAA,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,IAAI,EAAE;wBAC1B,MAAM,WAAW,GAAG,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,OAAO,CAAC;4BAC5C,GAAG,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,IAAI,IAAI,IAAI,OAAO,IAAI,KAAK,QAAQ,IAAI,IAAI,CAAC,IAAI,KAAK,UAAU,CAAC;wBAExF,IAAI,WAAW,EAAE;AACf,4BAAA,gBAAgB,CAAC,IAAI,CAAC,CAAC,CAAC;;;;;AAM9B,gBAAA,IAAI,gBAAgB,CAAC,MAAM,GAAG,CAAC,EAAE;;AAE/B,oBAAA,MAAM,0BAA0B,GAAG,gBAAgB,CAAC,CAAC,CAAC;AACtD,oBAAA,MAAM,qBAAqB,GAAG,OAAO,CAAC,0BAA0B,CAAC;;AAGjE,oBAAA,OAAO,CAAC,MAAM,CAAC,0BAA0B,EAAE,CAAC,CAAC;;AAG7C,oBAAA,MAAM,WAAW,GAAG,OAAO,CAAC,SAAS,CAAC,GAAG,IAAI,GAAG,CAAC,OAAO,EAAE,KAAK,QAAQ,CAAC;AACxE,oBAAA,MAAM,SAAS,GAAG,WAAW,KAAK,EAAE;oBAEpC,IAAI,SAAS,EAAE;;;;;;AAOb,wBAAA,MAAM,SAAS,GAAG,OAAO,CAAC,WAAW,CAAC;AACtC,wBAAA,OAAO,CAAC,MAAM,CAAC,WAAW,EAAE,CAAC,CAAC;AAC9B,wBAAA,OAAO,CAAC,IAAI,CAAC,SAAS,CAAC;;AAGvB,wBAAA,OAAO,CAAC,MAAM,CAAC,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC,EAAE,qBAAqB,CAAC;;yBACvD;;;AAGL,wBAAA,OAAO,CAAC,IAAI,CAAC,qBAAqB,CAAC;;;;;;IAO3C,MAAM,YAAY,GAAG,QAAQ;AAC7B,IAAA,YAAY,GAAG,YAAY,CAAC,MAAM,GAAG,CAAC;IACtC,sBAAsB,IAAI,iBAAiB;IAE3C,OAAO;QACL,YAAY;QACZ,sBAAsB;AACtB,QAAA,OAAO,EAAE,OAAO,CAAC,OAAO,EAAE;AAC1B,QAAA,gBAAgB,EAAE,YAAY;KAC/B;AACH;AAEA,SAAS,gBAAgB,CAAC,KAAc,EAAA;AACtC,IAAA,OAAO,OAAO,KAAK,KAAK,QAAQ,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,KAAK,GAAG,CAAC;AAChE;AAEM,SAAU,mBAAmB,CAAC,aAAyC,EAAA;IAC3E,MAAM,kBAAkB,GAAG,EAAE,GAAG,aAAa,CAAC,kBAAkB,EAAE;AAClE,IAAA,IAAI,kBAAkB,GAAG,aAAa,CAAC,UAAU;IACjD,IAAI,WAAW,GAAG,CAAC,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;IAEhF,OAAO,SAAS,aAAa,CAAC,MAA2B,EAAA;AAIvD,QAAA,IAAI,YAAuC;AAC3C,QAAA,IAAI,MAAM,CAAC,aAAa,KACtB,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,YAAY;AAC/C,gBACD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB;oBAEvD,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,cAAc;uBACrE,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,mBAAmB,CAAC,UAAU,CAAC,CACzE,CACF,CACF,IAAI,gBAAgB,CAAC,MAAM,CAAC,aAAa,CAAC,aAAa,CAAC,EAAE;AACzD,YAAA,YAAY,GAAG,oBAAoB,CAAC,MAAM,CAAC,aAAa,CAAC;AACzD,YAAA,WAAW,GAAG,YAAY,CAAC,YAAY;;AAGzC,QAAA,KAAK,IAAI,CAAC,GAAG,kBAAkB,EAAE,CAAC,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;YAChE,MAAM,OAAO,GAAG,MAAM,CAAC,QAAQ,CAAC,CAAC,CAAC;AAClC,YAAA,IAAI,CAAC,KAAK,kBAAkB,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,IAAI,YAAY,EAAE;AACnF,gBAAA,kBAAkB,CAAC,CAAC,CAAC,GAAG,YAAY,CAAC,aAAa;;AAC7C,iBAAA,IAAI,kBAAkB,CAAC,CAAC,CAAC,KAAK,SAAS,EAAE;gBAC9C,kBAAkB,CAAC,CAAC,CAAC,GAAG,aAAa,CAAC,YAAY,CAAC,OAAO,CAAC;AAC3D,gBAAA,WAAW,IAAI,kBAAkB,CAAC,CAAC,CAAC;;;;;;;QAQxC,IAAI,YAAY,EAAE;YAChB,MAAM,gBAAgB,GAAG,MAAM,CAAC,MAAM,CAAC,kBAAkB,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,KAAK,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;AACrF,YAAA,MAAM,KAAK,GAAG,YAAY,CAAC,YAAY,GAAG,gBAAgB;AAC1D,YAAA,KAAK,MAAM,GAAG,IAAI,kBAAkB,EAAE;AACpC,gBAAA,kBAAkB,CAAC,GAAG,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,kBAAkB,CAAC,GAAG,CAAC,GAAG,KAAK,CAAC;;;AAIzE,QAAA,kBAAkB,GAAG,MAAM,CAAC,QAAQ,CAAC,MAAM;AAC3C,QAAA,IAAI,WAAW,IAAI,aAAa,CAAC,SAAS,EAAE;YAC1C,OAAO,EAAE,OAAO,EAAE,MAAM,CAAC,QAAQ,EAAE,kBAAkB,EAAE;;;AAIzD,QAAA,MAAM,EAAE,OAAO,EAAE,GAAG,2BAA2B,CAAC;YAC9C,gBAAgB,EAAE,aAAa,CAAC,SAAS;YACzC,QAAQ,EAAE,MAAM,CAAC,QAAQ;YACzB,kBAAkB;YAClB,kBAAkB,EAAE,MAAM,CAAC,kBAAkB;YAC7C,eAAe,EAAE,aAAa,CAAC,eAAe;YAC9C,YAAY,EAAE,aAAa,CAAC,YAAY;AACzC,SAAA,CAAC;AAEF,QAAA,OAAO,EAAE,OAAO,EAAE,kBAAkB,EAAE;AACxC,KAAC;AACH;;;;"}
|
package/package.json
CHANGED
package/src/messages/prune.ts
CHANGED
|
@@ -102,10 +102,20 @@ function getMessagesWithinTokenLimit({
|
|
|
102
102
|
}
|
|
103
103
|
}
|
|
104
104
|
|
|
105
|
+
// Add system message if it exists
|
|
106
|
+
if (instructions && _messages.length > 0) {
|
|
107
|
+
context.push(_messages[0] as BaseMessage);
|
|
108
|
+
messages.shift();
|
|
109
|
+
}
|
|
110
|
+
|
|
105
111
|
// Handle thinking mode requirement for Anthropic
|
|
106
112
|
if (thinkingEnabled && context.length > 0 && tokenCounter) {
|
|
113
|
+
// Check if the latest message is an assistant message
|
|
114
|
+
const latestMessageIsAssistant = _messages.length > 0 && _messages[_messages.length - 1].getType() === 'ai';
|
|
115
|
+
|
|
107
116
|
// Process only if we have an assistant message in the context
|
|
108
117
|
const firstAssistantIndex = context.findIndex(msg => msg.getType() === 'ai');
|
|
118
|
+
|
|
109
119
|
if (firstAssistantIndex >= 0) {
|
|
110
120
|
const firstAssistantMsg = context[firstAssistantIndex];
|
|
111
121
|
|
|
@@ -116,11 +126,12 @@ function getMessagesWithinTokenLimit({
|
|
|
116
126
|
|
|
117
127
|
// Only proceed if we need to add thinking blocks
|
|
118
128
|
if (!hasThinkingBlock) {
|
|
119
|
-
// Collect thinking blocks from pruned assistant messages
|
|
129
|
+
// Collect thinking blocks from pruned assistant messages, starting from the most recent
|
|
120
130
|
const thinkingBlocks: any[] = [];
|
|
121
131
|
|
|
122
|
-
// Look through pruned messages for thinking blocks
|
|
123
|
-
for (
|
|
132
|
+
// Look through pruned messages for thinking blocks, starting from the end (most recent)
|
|
133
|
+
for (let i = messages.length - 1; i >= 0; i--) {
|
|
134
|
+
const msg = messages[i];
|
|
124
135
|
if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
|
|
125
136
|
for (const item of msg.content) {
|
|
126
137
|
if (item && typeof item === 'object' && item.type === 'thinking') {
|
|
@@ -174,31 +185,142 @@ function getMessagesWithinTokenLimit({
|
|
|
174
185
|
|
|
175
186
|
// If we've exceeded the token limit, we need to prune more messages
|
|
176
187
|
if (currentTokenCount > remainingContextTokens) {
|
|
177
|
-
//
|
|
178
|
-
|
|
188
|
+
// Build a map of tool call IDs to track AI <--> tool message correspondences
|
|
189
|
+
const toolCallIdMap = new Map<string, number>();
|
|
190
|
+
|
|
191
|
+
// Identify tool call IDs in the context
|
|
192
|
+
for (let i = 0; i < context.length; i++) {
|
|
193
|
+
const msg = context[i];
|
|
194
|
+
|
|
195
|
+
// Check for tool calls in AI messages
|
|
196
|
+
if (msg.getType() === 'ai' && Array.isArray(msg.content)) {
|
|
197
|
+
for (const item of msg.content) {
|
|
198
|
+
if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
|
|
199
|
+
toolCallIdMap.set(item.id, i);
|
|
200
|
+
}
|
|
201
|
+
}
|
|
202
|
+
}
|
|
203
|
+
|
|
204
|
+
// Check for tool messages
|
|
205
|
+
if (msg.getType() === 'tool' && 'tool_call_id' in msg && typeof msg.tool_call_id === 'string') {
|
|
206
|
+
toolCallIdMap.set(msg.tool_call_id, i);
|
|
207
|
+
}
|
|
208
|
+
}
|
|
209
|
+
|
|
210
|
+
// Track which messages to remove
|
|
211
|
+
const indicesToRemove = new Set<number>();
|
|
212
|
+
|
|
213
|
+
// Start removing messages from the end, but preserve AI <--> tool message correspondences
|
|
179
214
|
let i = context.length - 1;
|
|
180
215
|
while (i > firstAssistantIndex && currentTokenCount > remainingContextTokens) {
|
|
181
216
|
const msgToRemove = context[i];
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
217
|
+
|
|
218
|
+
// Check if this is a tool message or has tool calls
|
|
219
|
+
let canRemove = true;
|
|
220
|
+
|
|
221
|
+
if (msgToRemove.getType() === 'tool' && 'tool_call_id' in msgToRemove && typeof msgToRemove.tool_call_id === 'string') {
|
|
222
|
+
// If this is a tool message, check if we need to keep its corresponding AI message
|
|
223
|
+
const aiIndex = toolCallIdMap.get(msgToRemove.tool_call_id);
|
|
224
|
+
if (aiIndex !== undefined && aiIndex !== i && !indicesToRemove.has(aiIndex)) {
|
|
225
|
+
// We need to remove both the tool message and its corresponding AI message
|
|
226
|
+
indicesToRemove.add(i);
|
|
227
|
+
indicesToRemove.add(aiIndex);
|
|
228
|
+
currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[aiIndex]));
|
|
229
|
+
canRemove = false;
|
|
230
|
+
}
|
|
231
|
+
} else if (msgToRemove.getType() === 'ai' && Array.isArray(msgToRemove.content)) {
|
|
232
|
+
// If this is an AI message with tool calls, check if we need to keep its corresponding tool messages
|
|
233
|
+
for (const item of msgToRemove.content) {
|
|
234
|
+
if (item && typeof item === 'object' && item.type === 'tool_use' && item.id) {
|
|
235
|
+
const toolIndex = toolCallIdMap.get(item.id as string);
|
|
236
|
+
if (toolIndex !== undefined && toolIndex !== i && !indicesToRemove.has(toolIndex)) {
|
|
237
|
+
// We need to remove both the AI message and its corresponding tool message
|
|
238
|
+
indicesToRemove.add(i);
|
|
239
|
+
indicesToRemove.add(toolIndex);
|
|
240
|
+
currentTokenCount -= (tokenCounter(msgToRemove) + tokenCounter(context[toolIndex]));
|
|
241
|
+
canRemove = false;
|
|
242
|
+
break;
|
|
243
|
+
}
|
|
244
|
+
}
|
|
245
|
+
}
|
|
246
|
+
}
|
|
247
|
+
|
|
248
|
+
// If we can remove this message individually
|
|
249
|
+
if (canRemove && !indicesToRemove.has(i)) {
|
|
250
|
+
indicesToRemove.add(i);
|
|
251
|
+
currentTokenCount -= tokenCounter(msgToRemove);
|
|
252
|
+
}
|
|
253
|
+
|
|
185
254
|
i--;
|
|
186
255
|
}
|
|
187
256
|
|
|
257
|
+
// Remove messages in reverse order to avoid index shifting
|
|
258
|
+
const sortedIndices = Array.from(indicesToRemove).sort((a, b) => b - a);
|
|
259
|
+
for (const index of sortedIndices) {
|
|
260
|
+
context.splice(index, 1);
|
|
261
|
+
}
|
|
262
|
+
|
|
188
263
|
// Update remainingContextTokens to reflect the new token count
|
|
189
264
|
remainingContextTokens = maxContextTokens - currentTokenCount;
|
|
190
265
|
}
|
|
191
266
|
}
|
|
192
267
|
}
|
|
193
268
|
}
|
|
269
|
+
|
|
270
|
+
// If the latest message is an assistant message, ensure an assistant message with thinking appears at the end
|
|
271
|
+
// of the context (which will become the beginning after reversal)
|
|
272
|
+
// but maintain system message precedence after reversal
|
|
273
|
+
if (latestMessageIsAssistant && context.length > 0) {
|
|
274
|
+
// Find assistant messages with thinking blocks
|
|
275
|
+
const assistantIndices: number[] = [];
|
|
276
|
+
for (let i = 0; i < context.length; i++) {
|
|
277
|
+
const msg = context[i];
|
|
278
|
+
if (msg.getType() === 'ai') {
|
|
279
|
+
const hasThinking = Array.isArray(msg.content) &&
|
|
280
|
+
msg.content.some(item => item && typeof item === 'object' && item.type === 'thinking');
|
|
281
|
+
|
|
282
|
+
if (hasThinking) {
|
|
283
|
+
assistantIndices.push(i);
|
|
284
|
+
}
|
|
285
|
+
}
|
|
286
|
+
}
|
|
287
|
+
|
|
288
|
+
// If we found assistant messages with thinking blocks
|
|
289
|
+
if (assistantIndices.length > 0) {
|
|
290
|
+
// Get the first assistant message with thinking
|
|
291
|
+
const assistantWithThinkingIndex = assistantIndices[0];
|
|
292
|
+
const assistantWithThinking = context[assistantWithThinkingIndex];
|
|
293
|
+
|
|
294
|
+
// Remove it from its current position
|
|
295
|
+
context.splice(assistantWithThinkingIndex, 1);
|
|
296
|
+
|
|
297
|
+
// Check if there's a system message in the context
|
|
298
|
+
const systemIndex = context.findIndex(msg => msg.getType() === 'system');
|
|
299
|
+
const hasSystem = systemIndex !== -1;
|
|
300
|
+
|
|
301
|
+
if (hasSystem) {
|
|
302
|
+
// We want the system message to be first after reversal
|
|
303
|
+
// This means we need to put it at the end position before reversal
|
|
304
|
+
// And the assistant message should be second after reversal
|
|
305
|
+
// This means we need to put it at the end - 1 position before reversal
|
|
306
|
+
|
|
307
|
+
// First, ensure the system message is at the end (will be first after reversal)
|
|
308
|
+
const systemMsg = context[systemIndex];
|
|
309
|
+
context.splice(systemIndex, 1);
|
|
310
|
+
context.push(systemMsg);
|
|
311
|
+
|
|
312
|
+
// Then, put the assistant message right before the system message (will be second after reversal)
|
|
313
|
+
context.splice(context.length - 1, 0, assistantWithThinking);
|
|
314
|
+
} else {
|
|
315
|
+
// No system message, so we want assistant to be first after reversal
|
|
316
|
+
// This means we need to put it at the end position before reversal
|
|
317
|
+
context.push(assistantWithThinking);
|
|
318
|
+
}
|
|
319
|
+
}
|
|
320
|
+
}
|
|
194
321
|
}
|
|
195
322
|
}
|
|
196
323
|
|
|
197
|
-
if (instructions && _messages.length > 0) {
|
|
198
|
-
context.push(_messages[0] as BaseMessage);
|
|
199
|
-
messages.shift();
|
|
200
|
-
}
|
|
201
|
-
|
|
202
324
|
const prunedMemory = messages;
|
|
203
325
|
summaryIndex = prunedMemory.length - 1;
|
|
204
326
|
remainingContextTokens -= currentTokenCount;
|
|
@@ -390,4 +390,310 @@ describe('Prune Messages with Thinking Mode Tests', () => {
|
|
|
390
390
|
// The function should not throw an error even though no thinking blocks are found
|
|
391
391
|
expect(() => pruneMessages({ messages })).not.toThrow();
|
|
392
392
|
});
|
|
393
|
+
|
|
394
|
+
it('should preserve AI <--> tool message correspondences when pruning', () => {
|
|
395
|
+
// Create a token counter
|
|
396
|
+
const tokenCounter = createTestTokenCounter();
|
|
397
|
+
|
|
398
|
+
// Create messages with tool calls
|
|
399
|
+
const assistantMessageWithToolCall = new AIMessage({
|
|
400
|
+
content: [
|
|
401
|
+
{
|
|
402
|
+
type: "text",
|
|
403
|
+
text: "Let me check that file:",
|
|
404
|
+
},
|
|
405
|
+
{
|
|
406
|
+
type: "tool_use",
|
|
407
|
+
id: "tool123",
|
|
408
|
+
name: "text_editor_mcp_textEditor",
|
|
409
|
+
input: "{\"command\": \"view\", \"path\": \"/path/to/file.txt\"}",
|
|
410
|
+
},
|
|
411
|
+
],
|
|
412
|
+
});
|
|
413
|
+
|
|
414
|
+
const toolResponseMessage = new ToolMessage({
|
|
415
|
+
content: [
|
|
416
|
+
{
|
|
417
|
+
type: "text",
|
|
418
|
+
text: "{\"success\":true,\"message\":\"File content\"}",
|
|
419
|
+
},
|
|
420
|
+
],
|
|
421
|
+
tool_call_id: "tool123",
|
|
422
|
+
name: "text_editor_mcp_textEditor",
|
|
423
|
+
});
|
|
424
|
+
|
|
425
|
+
const assistantMessageWithThinking = new AIMessage({
|
|
426
|
+
content: [
|
|
427
|
+
{
|
|
428
|
+
type: "thinking",
|
|
429
|
+
thinking: "This is a thinking block",
|
|
430
|
+
},
|
|
431
|
+
{
|
|
432
|
+
type: "text",
|
|
433
|
+
text: "Response with thinking",
|
|
434
|
+
},
|
|
435
|
+
],
|
|
436
|
+
});
|
|
437
|
+
|
|
438
|
+
const messages = [
|
|
439
|
+
new SystemMessage("System instruction"),
|
|
440
|
+
new HumanMessage("Hello"),
|
|
441
|
+
assistantMessageWithToolCall,
|
|
442
|
+
toolResponseMessage,
|
|
443
|
+
new HumanMessage("Next message"),
|
|
444
|
+
assistantMessageWithThinking,
|
|
445
|
+
];
|
|
446
|
+
|
|
447
|
+
// Calculate token counts for each message
|
|
448
|
+
const indexTokenCountMap: Record<string, number> = {};
|
|
449
|
+
for (let i = 0; i < messages.length; i++) {
|
|
450
|
+
indexTokenCountMap[i] = tokenCounter(messages[i]);
|
|
451
|
+
}
|
|
452
|
+
|
|
453
|
+
// Create pruneMessages function with thinking mode enabled and a low token limit
|
|
454
|
+
const pruneMessages = createPruneMessages({
|
|
455
|
+
maxTokens: 100, // Set a low token limit to force pruning
|
|
456
|
+
startIndex: 0,
|
|
457
|
+
tokenCounter,
|
|
458
|
+
indexTokenCountMap: { ...indexTokenCountMap },
|
|
459
|
+
thinkingEnabled: true,
|
|
460
|
+
});
|
|
461
|
+
|
|
462
|
+
// Prune messages
|
|
463
|
+
const result = pruneMessages({ messages });
|
|
464
|
+
|
|
465
|
+
// Find assistant message with tool call and its corresponding tool message in the pruned context
|
|
466
|
+
const assistantIndex = result.context.findIndex(msg =>
|
|
467
|
+
msg.getType() === 'ai' &&
|
|
468
|
+
Array.isArray(msg.content) &&
|
|
469
|
+
msg.content.some(item => item && typeof item === 'object' && item.type === 'tool_use' && item.id === 'tool123')
|
|
470
|
+
);
|
|
471
|
+
|
|
472
|
+
// If the assistant message with tool call is in the context, its corresponding tool message should also be there
|
|
473
|
+
if (assistantIndex !== -1) {
|
|
474
|
+
const toolIndex = result.context.findIndex(msg =>
|
|
475
|
+
msg.getType() === 'tool' &&
|
|
476
|
+
'tool_call_id' in msg &&
|
|
477
|
+
msg.tool_call_id === 'tool123'
|
|
478
|
+
);
|
|
479
|
+
|
|
480
|
+
expect(toolIndex).not.toBe(-1);
|
|
481
|
+
}
|
|
482
|
+
|
|
483
|
+
// If the tool message is in the context, its corresponding assistant message should also be there
|
|
484
|
+
const toolIndex = result.context.findIndex(msg =>
|
|
485
|
+
msg.getType() === 'tool' &&
|
|
486
|
+
'tool_call_id' in msg &&
|
|
487
|
+
msg.tool_call_id === 'tool123'
|
|
488
|
+
);
|
|
489
|
+
|
|
490
|
+
if (toolIndex !== -1) {
|
|
491
|
+
const assistantWithToolIndex = result.context.findIndex(msg =>
|
|
492
|
+
msg.getType() === 'ai' &&
|
|
493
|
+
Array.isArray(msg.content) &&
|
|
494
|
+
msg.content.some(item => item && typeof item === 'object' && item.type === 'tool_use' && item.id === 'tool123')
|
|
495
|
+
);
|
|
496
|
+
|
|
497
|
+
expect(assistantWithToolIndex).not.toBe(-1);
|
|
498
|
+
}
|
|
499
|
+
});
|
|
500
|
+
|
|
501
|
+
it('should ensure an assistant message with thinking appears at the beginning of the context when the latest message is an assistant message', () => {
|
|
502
|
+
// Create a token counter
|
|
503
|
+
const tokenCounter = createTestTokenCounter();
|
|
504
|
+
|
|
505
|
+
// Create messages with the latest message being an assistant message with thinking
|
|
506
|
+
const assistantMessageWithThinking = new AIMessage({
|
|
507
|
+
content: [
|
|
508
|
+
{
|
|
509
|
+
type: "thinking",
|
|
510
|
+
thinking: "This is a thinking block",
|
|
511
|
+
},
|
|
512
|
+
{
|
|
513
|
+
type: "text",
|
|
514
|
+
text: "Response with thinking",
|
|
515
|
+
},
|
|
516
|
+
],
|
|
517
|
+
});
|
|
518
|
+
|
|
519
|
+
// Create an assistant message with tool use
|
|
520
|
+
const assistantMessageWithToolUse = new AIMessage({
|
|
521
|
+
content: [
|
|
522
|
+
{
|
|
523
|
+
type: "text",
|
|
524
|
+
text: "Let me check that file:",
|
|
525
|
+
},
|
|
526
|
+
{
|
|
527
|
+
type: "tool_use",
|
|
528
|
+
id: "tool123",
|
|
529
|
+
name: "text_editor_mcp_textEditor",
|
|
530
|
+
input: "{\"command\": \"view\", \"path\": \"/path/to/file.txt\"}",
|
|
531
|
+
},
|
|
532
|
+
],
|
|
533
|
+
});
|
|
534
|
+
|
|
535
|
+
// Create a tool response message
|
|
536
|
+
const toolResponseMessage = new ToolMessage({
|
|
537
|
+
content: [
|
|
538
|
+
{
|
|
539
|
+
type: "text",
|
|
540
|
+
text: "{\"success\":true,\"message\":\"File content\"}",
|
|
541
|
+
},
|
|
542
|
+
],
|
|
543
|
+
tool_call_id: "tool123",
|
|
544
|
+
name: "text_editor_mcp_textEditor",
|
|
545
|
+
});
|
|
546
|
+
|
|
547
|
+
// Test case without system message
|
|
548
|
+
const messagesWithoutSystem = [
|
|
549
|
+
new HumanMessage("Hello"),
|
|
550
|
+
assistantMessageWithToolUse,
|
|
551
|
+
toolResponseMessage,
|
|
552
|
+
new HumanMessage("Next message"),
|
|
553
|
+
assistantMessageWithThinking, // Latest message is an assistant message with thinking
|
|
554
|
+
];
|
|
555
|
+
|
|
556
|
+
// Calculate token counts for each message
|
|
557
|
+
const indexTokenCountMapWithoutSystem: Record<string, number> = {};
|
|
558
|
+
for (let i = 0; i < messagesWithoutSystem.length; i++) {
|
|
559
|
+
indexTokenCountMapWithoutSystem[i] = tokenCounter(messagesWithoutSystem[i]);
|
|
560
|
+
}
|
|
561
|
+
|
|
562
|
+
// Create pruneMessages function with thinking mode enabled
|
|
563
|
+
const pruneMessagesWithoutSystem = createPruneMessages({
|
|
564
|
+
maxTokens: 100, // Set a token limit to force some pruning
|
|
565
|
+
startIndex: 0,
|
|
566
|
+
tokenCounter,
|
|
567
|
+
indexTokenCountMap: { ...indexTokenCountMapWithoutSystem },
|
|
568
|
+
thinkingEnabled: true,
|
|
569
|
+
});
|
|
570
|
+
|
|
571
|
+
// Prune messages
|
|
572
|
+
const resultWithoutSystem = pruneMessagesWithoutSystem({ messages: messagesWithoutSystem });
|
|
573
|
+
|
|
574
|
+
// Verify that the first message in the pruned context is an assistant message with thinking when no system message
|
|
575
|
+
expect(resultWithoutSystem.context.length).toBeGreaterThan(0);
|
|
576
|
+
expect(resultWithoutSystem.context[0].getType()).toBe('ai');
|
|
577
|
+
|
|
578
|
+
// Verify that the first message has a thinking block
|
|
579
|
+
const firstMsgContent = resultWithoutSystem.context[0].content as any[];
|
|
580
|
+
expect(Array.isArray(firstMsgContent)).toBe(true);
|
|
581
|
+
const hasThinkingBlock = firstMsgContent.some(item =>
|
|
582
|
+
item && typeof item === 'object' && item.type === 'thinking');
|
|
583
|
+
expect(hasThinkingBlock).toBe(true);
|
|
584
|
+
|
|
585
|
+
// Test case with system message
|
|
586
|
+
const messagesWithSystem = [
|
|
587
|
+
new SystemMessage("System instruction"),
|
|
588
|
+
new HumanMessage("Hello"),
|
|
589
|
+
assistantMessageWithToolUse,
|
|
590
|
+
toolResponseMessage,
|
|
591
|
+
new HumanMessage("Next message"),
|
|
592
|
+
assistantMessageWithThinking, // Latest message is an assistant message with thinking
|
|
593
|
+
];
|
|
594
|
+
|
|
595
|
+
// Calculate token counts for each message
|
|
596
|
+
const indexTokenCountMapWithSystem: Record<string, number> = {};
|
|
597
|
+
for (let i = 0; i < messagesWithSystem.length; i++) {
|
|
598
|
+
indexTokenCountMapWithSystem[i] = tokenCounter(messagesWithSystem[i]);
|
|
599
|
+
}
|
|
600
|
+
|
|
601
|
+
// Create pruneMessages function with thinking mode enabled
|
|
602
|
+
const pruneMessagesWithSystem = createPruneMessages({
|
|
603
|
+
maxTokens: 120, // Set a token limit to force some pruning but keep system message
|
|
604
|
+
startIndex: 0,
|
|
605
|
+
tokenCounter,
|
|
606
|
+
indexTokenCountMap: { ...indexTokenCountMapWithSystem },
|
|
607
|
+
thinkingEnabled: true,
|
|
608
|
+
});
|
|
609
|
+
|
|
610
|
+
// Prune messages
|
|
611
|
+
const resultWithSystem = pruneMessagesWithSystem({ messages: messagesWithSystem });
|
|
612
|
+
|
|
613
|
+
// Verify that the system message remains first, followed by an assistant message with thinking
|
|
614
|
+
expect(resultWithSystem.context.length).toBeGreaterThan(1);
|
|
615
|
+
expect(resultWithSystem.context[0].getType()).toBe('system');
|
|
616
|
+
expect(resultWithSystem.context[1].getType()).toBe('ai');
|
|
617
|
+
|
|
618
|
+
// Verify that the second message has a thinking block
|
|
619
|
+
const secondMsgContent = resultWithSystem.context[1].content as any[];
|
|
620
|
+
expect(Array.isArray(secondMsgContent)).toBe(true);
|
|
621
|
+
const hasThinkingBlockInSecond = secondMsgContent.some(item =>
|
|
622
|
+
item && typeof item === 'object' && item.type === 'thinking');
|
|
623
|
+
expect(hasThinkingBlockInSecond).toBe(true);
|
|
624
|
+
});
|
|
625
|
+
|
|
626
|
+
it('should look for thinking blocks starting from the most recent messages', () => {
|
|
627
|
+
// Create a token counter
|
|
628
|
+
const tokenCounter = createTestTokenCounter();
|
|
629
|
+
|
|
630
|
+
// Create messages with multiple thinking blocks
|
|
631
|
+
const olderAssistantMessageWithThinking = new AIMessage({
|
|
632
|
+
content: [
|
|
633
|
+
{
|
|
634
|
+
type: "thinking",
|
|
635
|
+
thinking: "This is an older thinking block",
|
|
636
|
+
},
|
|
637
|
+
{
|
|
638
|
+
type: "text",
|
|
639
|
+
text: "Older response with thinking",
|
|
640
|
+
},
|
|
641
|
+
],
|
|
642
|
+
});
|
|
643
|
+
|
|
644
|
+
const newerAssistantMessageWithThinking = new AIMessage({
|
|
645
|
+
content: [
|
|
646
|
+
{
|
|
647
|
+
type: "thinking",
|
|
648
|
+
thinking: "This is a newer thinking block",
|
|
649
|
+
},
|
|
650
|
+
{
|
|
651
|
+
type: "text",
|
|
652
|
+
text: "Newer response with thinking",
|
|
653
|
+
},
|
|
654
|
+
],
|
|
655
|
+
});
|
|
656
|
+
|
|
657
|
+
const messages = [
|
|
658
|
+
new SystemMessage("System instruction"),
|
|
659
|
+
new HumanMessage("Hello"),
|
|
660
|
+
olderAssistantMessageWithThinking,
|
|
661
|
+
new HumanMessage("Next message"),
|
|
662
|
+
newerAssistantMessageWithThinking,
|
|
663
|
+
];
|
|
664
|
+
|
|
665
|
+
// Calculate token counts for each message
|
|
666
|
+
const indexTokenCountMap: Record<string, number> = {};
|
|
667
|
+
for (let i = 0; i < messages.length; i++) {
|
|
668
|
+
indexTokenCountMap[i] = tokenCounter(messages[i]);
|
|
669
|
+
}
|
|
670
|
+
|
|
671
|
+
// Create pruneMessages function with thinking mode enabled
|
|
672
|
+
// Set a token limit that will force pruning of the older assistant message
|
|
673
|
+
const pruneMessages = createPruneMessages({
|
|
674
|
+
maxTokens: 80,
|
|
675
|
+
startIndex: 0,
|
|
676
|
+
tokenCounter,
|
|
677
|
+
indexTokenCountMap: { ...indexTokenCountMap },
|
|
678
|
+
thinkingEnabled: true,
|
|
679
|
+
});
|
|
680
|
+
|
|
681
|
+
// Prune messages
|
|
682
|
+
const result = pruneMessages({ messages });
|
|
683
|
+
|
|
684
|
+
// Find the first assistant message in the pruned context
|
|
685
|
+
const firstAssistantIndex = result.context.findIndex(msg => msg.getType() === 'ai');
|
|
686
|
+
expect(firstAssistantIndex).not.toBe(-1);
|
|
687
|
+
|
|
688
|
+
const firstAssistantMsg = result.context[firstAssistantIndex];
|
|
689
|
+
expect(Array.isArray(firstAssistantMsg.content)).toBe(true);
|
|
690
|
+
|
|
691
|
+
// Verify that the first assistant message has a thinking block
|
|
692
|
+
const thinkingBlock = (firstAssistantMsg.content as any[]).find(item =>
|
|
693
|
+
item && typeof item === 'object' && item.type === 'thinking');
|
|
694
|
+
expect(thinkingBlock).toBeDefined();
|
|
695
|
+
|
|
696
|
+
// Verify that it's the newer thinking block
|
|
697
|
+
expect(thinkingBlock.thinking).toContain("newer thinking block");
|
|
698
|
+
});
|
|
393
699
|
});
|