49 lines
2.1 KiB
TypeScript
49 lines
2.1 KiB
TypeScript
import type { ModelMessage } from "ai";
|
|
|
|
// Keep complete assistant/tool batches together. Never remove just one side of a tool call.
|
|
export function splitCompactableHistory(messages: ModelMessage[], keepUserTurns = 3) {
|
|
const userIndices = messages.flatMap((message, index) =>
|
|
message.role === "user" ? [index] : [],
|
|
);
|
|
let cut =
|
|
userIndices.length > keepUserTurns ? userIndices[userIndices.length - keepUserTurns] : 0;
|
|
if (!cut) {
|
|
const openCalls = new Set<string>();
|
|
for (let index = 0; index < messages.length - 6; index++) {
|
|
const message = messages[index];
|
|
if (message.role === "assistant" && Array.isArray(message.content)) {
|
|
for (const part of message.content)
|
|
if (part.type === "tool-call") openCalls.add(part.toolCallId);
|
|
}
|
|
if (message.role === "tool")
|
|
for (const part of message.content) openCalls.delete(part.toolCallId);
|
|
if (!openCalls.size && message.role !== "user") cut = index + 1;
|
|
}
|
|
}
|
|
return { older: messages.slice(0, cut), recent: messages.slice(cut) };
|
|
}
|
|
|
|
export function estimateContextTokens(messages: ModelMessage[]) {
|
|
// Conservatively count CJK characters and media rather than assuming four characters per token.
|
|
const text = JSON.stringify(messages);
|
|
const cjk = (text.match(/[^\x00-\x7f]/g) || []).length;
|
|
return Math.ceil((text.length - cjk) / 3) + cjk;
|
|
}
|
|
|
|
export async function compactAgentContext(params: {
|
|
messages: ModelMessage[];
|
|
previousSummary: string;
|
|
summarize: (text: string) => Promise<string>;
|
|
threshold?: number;
|
|
}) {
|
|
if (estimateContextTokens(params.messages) < (params.threshold ?? 18_000)) return params;
|
|
const { older, recent } = splitCompactableHistory(params.messages);
|
|
if (!older.length) throw new Error("当前任务上下文过长，已保留结果，请分成更小的任务继续。");
|
|
const summary = (
|
|
await params.summarize(JSON.stringify({ previousSummary: params.previousSummary, older }))
|
|
).trim();
|
|
if (!summary || summary.length > 12_000)
|
|
throw new Error("上下文整理未完成，已保留已完成的步骤。");
|
|
return { ...params, previousSummary: summary, messages: recent };
|
|
}
|