Skip to content

feat(core): simplify truncation logic to only keep the newest message #18904

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
nicohrubec wants to merge 1 commit into from
Closed

feat(core): simplify truncation logic to only keep the newest message #18904

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-media-truncation.mjs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,10 @@ async function run() {
model: 'claude-3-haiku-20240307',
max_tokens: 1024,
messages: [
{
role: 'user',
content: 'what number is this?',
},
{
role: 'user',
content: [
Expand All @@ -66,10 +70,6 @@ async function run() {
},
],
},
{
role: 'user',
content: 'what number is this?',
},
],
temperature: 0.7,
});
Expand Down
5 changes: 1 addition & 4 deletions dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -677,6 +677,7 @@ describe('Anthropic integration', () => {
'sentry.origin': 'auto.ai.anthropic',
'gen_ai.system': 'anthropic',
'gen_ai.request.model': 'claude-3-haiku-20240307',
// Only the last message (with filtered media) should be kept
'gen_ai.request.messages': JSON.stringify([
{
role: 'user',
Expand All @@ -691,10 +692,6 @@ describe('Anthropic integration', () => {
},
],
},
{
role: 'user',
content: 'what number is this?',
},
]),
}),
description: 'messages claude-3-haiku-20240307',
Expand Down
59 changes: 17 additions & 42 deletions packages/core/src/tracing/ai/messageTruncation.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -374,19 +374,19 @@ function stripInlineMediaFromMessages(messages: unknown[]): unknown[] {
* Truncate an array of messages to fit within a byte limit.
*
* Strategy:
* - Keeps the newest messages (from the end of the array)
* - Uses O(n) algorithm: precompute sizes once, then find largest suffix under budget
* - If no complete messages fit, attempts to truncate the newest single message
* - Always keeps only the last (newest) message
* - Strips inline media from the message
* - Truncates the message content if it exceeds the byte limit
*
* @param messages - Array of messages to truncate
* @param maxBytes - Maximum total byte limit for all messages
* @returns Truncated array of messages
* @param maxBytes - Maximum total byte limit for the message
* @returns Array containing only the last message (possibly truncated)
*
* @example
* ```ts
* const messages = [msg1, msg2, msg3, msg4]; // newest is msg4
* const truncated = truncateMessagesByBytes(messages, 10000);
* // Returns [msg3, msg4] if they fit, or [msg4] if only it fits, etc.
* // Returns [msg4] (truncated if needed)
* ```
*/
function truncateMessagesByBytes(messages: unknown[], maxBytes: number): unknown[] {
Expand All @@ -395,46 +395,21 @@ function truncateMessagesByBytes(messages: unknown[], maxBytes: number): unknown
return messages;
}

// strip inline media first. This will often get us below the threshold,
// while preserving human-readable information about messages sent.
const stripped = stripInlineMediaFromMessages(messages);

// Fast path: if all messages fit, return as-is
const totalBytes = jsonBytes(stripped);
if (totalBytes <= maxBytes) {
return stripped;
}
// Always keep only the last message
const lastMessage = messages[messages.length - 1];

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

// Find the largest suffix (newest messages) that fits within the budget
let bytesUsed = 0;
let startIndex = stripped.length; // Index where the kept suffix starts

for (let i = stripped.length - 1; i >= 0; i--) {
const messageSize = messageSizes[i];

if (messageSize && bytesUsed + messageSize > maxBytes) {
// Adding this message would exceed the budget
break;
}

if (messageSize) {
bytesUsed += messageSize;
}
startIndex = i;
}

// If no complete messages fit, try truncating just the newest message
if (startIndex === stripped.length) {
// we're truncating down to one message, so all others dropped.
const newestMessage = stripped[stripped.length - 1];
return truncateSingleMessage(newestMessage, maxBytes);
// Check if it fits
const messageBytes = jsonBytes(strippedMessage);
if (messageBytes <= maxBytes) {
return stripped;
}

// Return the suffix that fits
return stripped.slice(startIndex);
// Truncate the single message if needed
return truncateSingleMessage(strippedMessage, maxBytes);
}

/**
Expand Down
96 changes: 25 additions & 71 deletions packages/core/test/lib/tracing/ai-message-truncation.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,33 +96,8 @@ describe('message truncation utilities', () => {

// original messages objects must not be mutated
expect(JSON.stringify(messages, null, 2)).toBe(messagesJson);
// only the last message should be kept (with media stripped)
expect(result).toStrictEqual([
{
role: 'user',
content: [
{
type: 'image',
source: {
type: 'base64',
media_type: 'image/png',
data: removed,
},
},
],
},
{
role: 'user',
content: {
image_url: removed,
},
},
{
role: 'agent',
type: 'image',
content: {
b64_json: removed,
},
},
{
role: 'system',
inlineData: {
Expand Down Expand Up @@ -177,39 +152,35 @@ describe('message truncation utilities', () => {
const giant = 'this is a long string '.repeat(1_000);
const big = 'this is a long string '.repeat(100);

it('drops older messages to fit in the limit', () => {
it('keeps only the last message without truncation when it fits the limit', () => {
// Multiple messages that together exceed 20KB, but last message is small
const messages = [
`0 ${giant}`,
{ type: 'text', content: `1 ${big}` },
{ type: 'text', content: `2 ${big}` },
{ type: 'text', content: `3 ${giant}` },
{ type: 'text', content: `4 ${big}` },
`5 ${big}`,
{ type: 'text', content: `6 ${big}` },
{ type: 'text', content: `7 ${big}` },
{ type: 'text', content: `8 ${big}` },
{ type: 'text', content: `9 ${big}` },
{ type: 'text', content: `10 ${big}` },
{ type: 'text', content: `11 ${big}` },
{ type: 'text', content: `12 ${big}` },
{ content: `1 ${humongous}` },
{ content: `2 ${humongous}` },
{ content: `3 ${big}` }, // last message - small enough to fit
];

const messagesJson = JSON.stringify(messages, null, 2);
const result = truncateGenAiMessages(messages);
// should not mutate original messages list
expect(JSON.stringify(messages, null, 2)).toBe(messagesJson);

// just retain the messages that fit in the budget
expect(result).toStrictEqual([
`5 ${big}`,
{ type: 'text', content: `6 ${big}` },
{ type: 'text', content: `7 ${big}` },
{ type: 'text', content: `8 ${big}` },
{ type: 'text', content: `9 ${big}` },
{ type: 'text', content: `10 ${big}` },
{ type: 'text', content: `11 ${big}` },
{ type: 'text', content: `12 ${big}` },
]);
// Should only keep the last message, unchanged
expect(result).toStrictEqual([{ content: `3 ${big}` }]);
});

it('keeps only the last message with truncation when it does not fit the limit', () => {
const messages = [{ content: `1 ${humongous}` }, { content: `2 ${humongous}` }, { content: `3 ${humongous}` }];
const result = truncateGenAiMessages(messages);
const truncLen = 20_000 - JSON.stringify({ content: '' }).length;
expect(result).toStrictEqual([{ content: `3 ${humongous}`.substring(0, truncLen) }]);
});

it('drops if last message cannot be safely truncated', () => {
const messages = [
{ content: `1 ${humongous}` },
{ content: `2 ${humongous}` },
{ what_even_is_this: `? ${humongous}` },
];
const result = truncateGenAiMessages(messages);
expect(result).toStrictEqual([]);
});

it('fully drops message if content cannot be made to fit', () => {
Expand Down Expand Up @@ -315,22 +286,5 @@ describe('message truncation utilities', () => {
},
]);
});

it('truncates first message if none fit', () => {
const messages = [{ content: `1 ${humongous}` }, { content: `2 ${humongous}` }, { content: `3 ${humongous}` }];
const result = truncateGenAiMessages(messages);
const truncLen = 20_000 - JSON.stringify({ content: '' }).length;
expect(result).toStrictEqual([{ content: `3 ${humongous}`.substring(0, truncLen) }]);
});

it('drops if first message cannot be safely truncated', () => {
const messages = [
{ content: `1 ${humongous}` },
{ content: `2 ${humongous}` },
{ what_even_is_this: `? ${humongous}` },
];
const result = truncateGenAiMessages(messages);
expect(result).toStrictEqual([]);
});
});
});
Loading