ai
Version:
AI SDK by Vercel - build apps like ChatGPT, Claude, Gemini, and more with a single interface for any model using the Vercel AI Gateway or go direct to OpenAI, Anthropic, Google, or any other model provider.
271 lines (245 loc) • 9.08 kB
text/typescript
import type {
AssistantModelMessage,
ModelMessage,
ToolModelMessage,
} from '@ai-sdk/provider-utils';
/**
* Prunes model messages from a list of model messages.
*
* @param messages - The list of model messages to prune.
* @param reasoning - How to remove reasoning content from assistant messages. Default is `'none'`.
* @param toolCalls - How to prune tool call/results/approval content. Default is `[]`.
* @param emptyMessages - Whether to keep or remove messages whose content is empty after pruning. Default is `'remove'`.
*
* @returns The pruned list of model messages.
*/
export function pruneMessages({
messages,
reasoning = 'none',
toolCalls = [],
emptyMessages = 'remove',
}: {
messages: ModelMessage[];
reasoning?: 'all' | 'before-last-message' | 'none';
toolCalls?:
| 'all'
| 'before-last-message'
| `before-last-${number}-messages`
| 'none'
| Array<{
type: 'all' | 'before-last-message' | `before-last-${number}-messages`;
tools?: string[];
}>;
emptyMessages?: 'keep' | 'remove';
}): ModelMessage[] {
// filter reasoning parts:
if (reasoning === 'all' || reasoning === 'before-last-message') {
messages = messages.map((message, messageIndex) => {
if (
message.role !== 'assistant' ||
typeof message.content === 'string' ||
(reasoning === 'before-last-message' &&
messageIndex === messages.length - 1)
) {
return message;
}
return {
...message,
content: message.content.filter(
part => part.type !== 'reasoning' && part.type !== 'reasoning-file',
),
};
});
}
// filter tool calls, results, errors, and approvals:
if (toolCalls === 'none') {
toolCalls = [];
} else if (toolCalls === 'all') {
toolCalls = [{ type: 'all' }];
} else if (toolCalls === 'before-last-message') {
toolCalls = [{ type: 'before-last-message' }];
} else if (typeof toolCalls === 'string') {
toolCalls = [{ type: toolCalls }];
}
for (const toolCall of toolCalls) {
// determine how many trailing messages to keep:
const keepLastMessagesCount =
toolCall.type === 'all'
? undefined
: toolCall.type === 'before-last-message'
? 1
: Number(
toolCall.type
.slice('before-last-'.length)
.slice(0, -'-messages'.length),
);
// scan kept messages to identify tool calls and approvals that need to be kept:
const keptToolCallIds: Set<string> = new Set();
const keptApprovalIds: Set<string> = new Set();
if (keepLastMessagesCount != null && keepLastMessagesCount !== 0) {
for (const message of messages.slice(-keepLastMessagesCount)) {
if (
(message.role === 'assistant' || message.role === 'tool') &&
typeof message.content !== 'string'
) {
for (const part of message.content) {
if (part.type === 'tool-call' || part.type === 'tool-result') {
keptToolCallIds.add(part.toolCallId);
} else if (
part.type === 'tool-approval-request' ||
part.type === 'tool-approval-response'
) {
keptApprovalIds.add(part.approvalId);
}
}
}
}
}
// Build global maps from tool call id and approval id to tool name.
// These must be global (not per-message) because a `tool-approval-response`
// lives in a separate `tool` message from its `tool-approval-request`
// (assistant message), so the tool name of a response can only be resolved
// by looking across messages. Resolving names per-message left responses
// unresolved, which caused them to be kept while their request was pruned,
// producing orphaned approval responses.
const toolCallIdToToolName = new Map<string, string>();
for (const message of messages) {
if (
(message.role === 'assistant' || message.role === 'tool') &&
typeof message.content !== 'string'
) {
for (const part of message.content) {
if (part.type === 'tool-call' || part.type === 'tool-result') {
toolCallIdToToolName.set(part.toolCallId, part.toolName);
}
}
}
}
const approvalIdToToolCallId = new Map<string, string>();
const approvalIdToToolName = new Map<string, string>();
for (const message of messages) {
if (
(message.role === 'assistant' || message.role === 'tool') &&
typeof message.content !== 'string'
) {
for (const part of message.content) {
if (part.type === 'tool-approval-request') {
approvalIdToToolCallId.set(part.approvalId, part.toolCallId);
const toolName = toolCallIdToToolName.get(part.toolCallId);
if (toolName != null) {
approvalIdToToolName.set(part.approvalId, toolName);
}
}
}
}
}
// Approval requests depend on their originating tool call. When an
// approval response is in the retained messages, trace through its request
// so pruning does not leave a pending approval without the call to execute.
for (const approvalId of keptApprovalIds) {
const toolCallId = approvalIdToToolCallId.get(approvalId);
if (toolCallId != null) {
keptToolCallIds.add(toolCallId);
}
}
messages = messages.map((message, messageIndex) => {
if (
(message.role !== 'assistant' && message.role !== 'tool') ||
typeof message.content === 'string' ||
(keepLastMessagesCount &&
messageIndex >= messages.length - keepLastMessagesCount)
) {
return message;
}
const keepParts = message.content.map(part => {
// keep non-tool parts:
if (
part.type !== 'tool-call' &&
part.type !== 'tool-result' &&
part.type !== 'tool-approval-request' &&
part.type !== 'tool-approval-response'
) {
return true;
}
// keep parts that are associated with a tool call or approval that needs to be kept:
if (
((part.type === 'tool-call' || part.type === 'tool-result') &&
keptToolCallIds.has(part.toolCallId)) ||
((part.type === 'tool-approval-request' ||
part.type === 'tool-approval-response') &&
keptApprovalIds.has(part.approvalId))
) {
return true;
}
// keep parts that are not associated with a tool that should be removed:
const partToolName =
part.type === 'tool-call' || part.type === 'tool-result'
? part.toolName
: approvalIdToToolName.get(part.approvalId);
return (
toolCall.tools != null &&
partToolName != null &&
!toolCall.tools.includes(partToolName)
);
});
// Consecutive reasoning parts can belong to the same provider item.
// Remove the whole group when its following content was pruned.
if (message.role === 'assistant') {
for (
let reasoningIndex = 0;
reasoningIndex < message.content.length;
reasoningIndex++
) {
const part = message.content[reasoningIndex];
if (part.type !== 'reasoning' && part.type !== 'reasoning-file') {
continue;
}
let reasoningEndIndex = reasoningIndex + 1;
while (reasoningEndIndex < message.content.length) {
const nextPart = message.content[reasoningEndIndex];
if (
nextPart.type !== 'reasoning' &&
nextPart.type !== 'reasoning-file'
) {
break;
}
reasoningEndIndex++;
}
let followingPartIndex = reasoningEndIndex;
let hasFollowingPart = false;
let hasKeptFollowingPart = false;
while (followingPartIndex < message.content.length) {
const followingPart = message.content[followingPartIndex];
if (
followingPart.type === 'reasoning' ||
followingPart.type === 'reasoning-file'
) {
break;
}
hasFollowingPart = true;
hasKeptFollowingPart ||= keepParts[followingPartIndex];
followingPartIndex++;
}
if (hasFollowingPart && !hasKeptFollowingPart) {
for (
let index = reasoningIndex;
index < reasoningEndIndex;
index++
) {
keepParts[index] = false;
}
}
reasoningIndex = followingPartIndex - 1;
}
}
return {
...message,
content: message.content.filter((_, index) => keepParts[index]),
} as AssistantModelMessage | ToolModelMessage;
});
}
if (emptyMessages === 'remove') {
messages = messages.filter(message => message.content.length > 0);
}
return messages;
}