Skip to content

Commit ff334af

Browse files
nicohrubecnicohrubec
committed
feat(core): simplify truncation logic to only keep the newest message
1 parent 9115e19 commit ff334af

4 files changed

Lines changed: 47 additions & 121 deletions

File tree

‎dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-media-truncation.mjs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,10 @@ async function run() {
5353
model: 'claude-3-haiku-20240307',
5454
max_tokens: 1024,
5555
messages: [
56+
{
57+
role: 'user',
58+
content: 'what number is this?',
59+
},
5660
{
5761
role: 'user',
5862
content: [
@@ -66,10 +70,6 @@ async function run() {
6670
},
6771
],
6872
},
69-
{
70-
role: 'user',
71-
content: 'what number is this?',
72-
},
7373
],
7474
temperature: 0.7,
7575
});

‎dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -677,6 +677,7 @@ describe('Anthropic integration', () => {
677677
'sentry.origin': 'auto.ai.anthropic',
678678
'gen_ai.system': 'anthropic',
679679
'gen_ai.request.model': 'claude-3-haiku-20240307',
680+
// Only the last message (with filtered media) should be kept
680681
'gen_ai.request.messages': JSON.stringify([
681682
{
682683
role: 'user',
@@ -691,10 +692,6 @@ describe('Anthropic integration', () => {
691692
},
692693
],
693694
},
694-
{
695-
role: 'user',
696-
content: 'what number is this?',
697-
},
698695
]),
699696
}),
700697
description: 'messages claude-3-haiku-20240307',

‎packages/core/src/tracing/ai/messageTruncation.ts‎

Lines changed: 17 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -374,19 +374,19 @@ function stripInlineMediaFromMessages(messages: unknown[]): unknown[] {
374374
* Truncate an array of messages to fit within a byte limit.
375375
*
376376
* Strategy:
377-
* - Keeps the newest messages (from the end of the array)
378-
* - Uses O(n) algorithm: precompute sizes once, then find largest suffix under budget
379-
* - If no complete messages fit, attempts to truncate the newest single message
377+
* - Always keeps only the last (newest) message
378+
* - Strips inline media from the message
379+
* - Truncates the message content if it exceeds the byte limit
380380
*
381381
* @param messages - Array of messages to truncate
382-
* @param maxBytes - Maximum total byte limit for all messages
383-
* @returns Truncated array of messages
382+
* @param maxBytes - Maximum total byte limit for the message
383+
* @returns Array containing only the last message (possibly truncated)
384384
*
385385
* @example
386386
* ```ts
387387
* const messages = [msg1, msg2, msg3, msg4]; // newest is msg4
388388
* const truncated = truncateMessagesByBytes(messages, 10000);
389-
* // Returns [msg3, msg4] if they fit, or [msg4] if only it fits, etc.
389+
* // Returns [msg4] (truncated if needed)
390390
* ```
391391
*/
392392
function truncateMessagesByBytes(messages: unknown[], maxBytes: number): unknown[] {
@@ -395,46 +395,21 @@ function truncateMessagesByBytes(messages: unknown[], maxBytes: number): unknown
395395
return messages;
396396
}
397397

398-
// strip inline media first. This will often get us below the threshold,
399-
// while preserving human-readable information about messages sent.
400-
const stripped = stripInlineMediaFromMessages(messages);
401-
402-
// Fast path: if all messages fit, return as-is
403-
const totalBytes = jsonBytes(stripped);
404-
if (totalBytes <= maxBytes) {
405-
return stripped;
406-
}
398+
// Always keep only the last message
399+
const lastMessage = messages[messages.length - 1];
407400

408-
// Precompute each message's JSON size once for efficiency
409-
const messageSizes = stripped.map(jsonBytes);
401+
// Strip inline media from the single message
402+
const stripped = stripInlineMediaFromMessages([lastMessage]);
403+
const strippedMessage = stripped[0];
410404

411-
// Find the largest suffix (newest messages) that fits within the budget
412-
let bytesUsed = 0;
413-
let startIndex = stripped.length; // Index where the kept suffix starts
414-
415-
for (let i = stripped.length - 1; i >= 0; i--) {
416-
const messageSize = messageSizes[i];
417-
418-
if (messageSize && bytesUsed + messageSize > maxBytes) {
419-
// Adding this message would exceed the budget
420-
break;
421-
}
422-
423-
if (messageSize) {
424-
bytesUsed += messageSize;
425-
}
426-
startIndex = i;
427-
}
428-
429-
// If no complete messages fit, try truncating just the newest message
430-
if (startIndex === stripped.length) {
431-
// we're truncating down to one message, so all others dropped.
432-
const newestMessage = stripped[stripped.length - 1];
433-
return truncateSingleMessage(newestMessage, maxBytes);
405+
// Check if it fits
406+
const messageBytes = jsonBytes(strippedMessage);
407+
if (messageBytes <= maxBytes) {
408+
return stripped;
434409
}
435410

436-
// Return the suffix that fits
437-
return stripped.slice(startIndex);
411+
// Truncate the single message if needed
412+
return truncateSingleMessage(strippedMessage, maxBytes);
438413
}
439414

440415
/**

‎packages/core/test/lib/tracing/ai-message-truncation.test.ts‎

Lines changed: 25 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -96,33 +96,8 @@ describe('message truncation utilities', () => {
9696

9797
// original messages objects must not be mutated
9898
expect(JSON.stringify(messages, null, 2)).toBe(messagesJson);
99+
// only the last message should be kept (with media stripped)
99100
expect(result).toStrictEqual([
100-
{
101-
role: 'user',
102-
content: [
103-
{
104-
type: 'image',
105-
source: {
106-
type: 'base64',
107-
media_type: 'image/png',
108-
data: removed,
109-
},
110-
},
111-
],
112-
},
113-
{
114-
role: 'user',
115-
content: {
116-
image_url: removed,
117-
},
118-
},
119-
{
120-
role: 'agent',
121-
type: 'image',
122-
content: {
123-
b64_json: removed,
124-
},
125-
},
126101
{
127102
role: 'system',
128103
inlineData: {
@@ -177,39 +152,35 @@ describe('message truncation utilities', () => {
177152
const giant = 'this is a long string '.repeat(1_000);
178153
const big = 'this is a long string '.repeat(100);
179154

180-
it('drops older messages to fit in the limit', () => {
155+
it('keeps only the last message without truncation when it fits the limit', () => {
156+
// Multiple messages that together exceed 20KB, but last message is small
181157
const messages = [
182-
`0 ${giant}`,
183-
{ type: 'text', content: `1 ${big}` },
184-
{ type: 'text', content: `2 ${big}` },
185-
{ type: 'text', content: `3 ${giant}` },
186-
{ type: 'text', content: `4 ${big}` },
187-
`5 ${big}`,
188-
{ type: 'text', content: `6 ${big}` },
189-
{ type: 'text', content: `7 ${big}` },
190-
{ type: 'text', content: `8 ${big}` },
191-
{ type: 'text', content: `9 ${big}` },
192-
{ type: 'text', content: `10 ${big}` },
193-
{ type: 'text', content: `11 ${big}` },
194-
{ type: 'text', content: `12 ${big}` },
158+
{ content: `1 ${humongous}` },
159+
{ content: `2 ${humongous}` },
160+
{ content: `3 ${big}` }, // last message - small enough to fit
195161
];
196162

197-
const messagesJson = JSON.stringify(messages, null, 2);
198163
const result = truncateGenAiMessages(messages);
199-
// should not mutate original messages list
200-
expect(JSON.stringify(messages, null, 2)).toBe(messagesJson);
201164

202-
// just retain the messages that fit in the budget
203-
expect(result).toStrictEqual([
204-
`5 ${big}`,
205-
{ type: 'text', content: `6 ${big}` },
206-
{ type: 'text', content: `7 ${big}` },
207-
{ type: 'text', content: `8 ${big}` },
208-
{ type: 'text', content: `9 ${big}` },
209-
{ type: 'text', content: `10 ${big}` },
210-
{ type: 'text', content: `11 ${big}` },
211-
{ type: 'text', content: `12 ${big}` },
212-
]);
165+
// Should only keep the last message, unchanged
166+
expect(result).toStrictEqual([{ content: `3 ${big}` }]);
167+
});
168+
169+
it('keeps only the last message with truncation when it does not fit the limit', () => {
170+
const messages = [{ content: `1 ${humongous}` }, { content: `2 ${humongous}` }, { content: `3 ${humongous}` }];
171+
const result = truncateGenAiMessages(messages);
172+
const truncLen = 20_000 - JSON.stringify({ content: '' }).length;
173+
expect(result).toStrictEqual([{ content: `3 ${humongous}`.substring(0, truncLen) }]);
174+
});
175+
176+
it('drops if last message cannot be safely truncated', () => {
177+
const messages = [
178+
{ content: `1 ${humongous}` },
179+
{ content: `2 ${humongous}` },
180+
{ what_even_is_this: `? ${humongous}` },
181+
];
182+
const result = truncateGenAiMessages(messages);
183+
expect(result).toStrictEqual([]);
213184
});
214185

215186
it('fully drops message if content cannot be made to fit', () => {
@@ -315,22 +286,5 @@ describe('message truncation utilities', () => {
315286
},
316287
]);
317288
});
318-
319-
it('truncates first message if none fit', () => {
320-
const messages = [{ content: `1 ${humongous}` }, { content: `2 ${humongous}` }, { content: `3 ${humongous}` }];
321-
const result = truncateGenAiMessages(messages);
322-
const truncLen = 20_000 - JSON.stringify({ content: '' }).length;
323-
expect(result).toStrictEqual([{ content: `3 ${humongous}`.substring(0, truncLen) }]);
324-
});
325-
326-
it('drops if first message cannot be safely truncated', () => {
327-
const messages = [
328-
{ content: `1 ${humongous}` },
329-
{ content: `2 ${humongous}` },
330-
{ what_even_is_this: `? ${humongous}` },
331-
];
332-
const result = truncateGenAiMessages(messages);
333-
expect(result).toStrictEqual([]);
334-
});
335289
});
336290
});

0 commit comments

Comments
 (0)