108 lines
3.2 KiB
TypeScript
108 lines
3.2 KiB
TypeScript
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
|
import { describe, expect, it } from "vitest";
|
|
|
|
import {
|
|
estimateMessagesTokens,
|
|
pruneHistoryForContextShare,
|
|
splitMessagesByTokenShare,
|
|
} from "./compaction.js";
|
|
|
|
function makeMessage(id: number, size: number): AgentMessage {
|
|
return {
|
|
role: "user",
|
|
content: "x".repeat(size),
|
|
timestamp: id,
|
|
};
|
|
}
|
|
|
|
describe("splitMessagesByTokenShare", () => {
|
|
it("splits messages into two non-empty parts", () => {
|
|
const messages: AgentMessage[] = [
|
|
makeMessage(1, 4000),
|
|
makeMessage(2, 4000),
|
|
makeMessage(3, 4000),
|
|
makeMessage(4, 4000),
|
|
];
|
|
|
|
const parts = splitMessagesByTokenShare(messages, 2);
|
|
expect(parts.length).toBeGreaterThanOrEqual(2);
|
|
expect(parts[0]?.length).toBeGreaterThan(0);
|
|
expect(parts[1]?.length).toBeGreaterThan(0);
|
|
expect(parts.flat().length).toBe(messages.length);
|
|
});
|
|
|
|
it("preserves message order across parts", () => {
|
|
const messages: AgentMessage[] = [
|
|
makeMessage(1, 4000),
|
|
makeMessage(2, 4000),
|
|
makeMessage(3, 4000),
|
|
makeMessage(4, 4000),
|
|
makeMessage(5, 4000),
|
|
makeMessage(6, 4000),
|
|
];
|
|
|
|
const parts = splitMessagesByTokenShare(messages, 3);
|
|
expect(parts.flat().map((msg) => msg.timestamp)).toEqual(messages.map((msg) => msg.timestamp));
|
|
});
|
|
});
|
|
|
|
describe("pruneHistoryForContextShare", () => {
|
|
it("drops older chunks until the history budget is met", () => {
|
|
const messages: AgentMessage[] = [
|
|
makeMessage(1, 4000),
|
|
makeMessage(2, 4000),
|
|
makeMessage(3, 4000),
|
|
makeMessage(4, 4000),
|
|
];
|
|
const maxContextTokens = 2000; // budget is 1000 tokens (50%)
|
|
const pruned = pruneHistoryForContextShare({
|
|
messages,
|
|
maxContextTokens,
|
|
maxHistoryShare: 0.5,
|
|
parts: 2,
|
|
});
|
|
|
|
expect(pruned.droppedChunks).toBeGreaterThan(0);
|
|
expect(pruned.keptTokens).toBeLessThanOrEqual(Math.floor(maxContextTokens * 0.5));
|
|
expect(pruned.messages.length).toBeGreaterThan(0);
|
|
});
|
|
|
|
it("keeps the newest messages when pruning", () => {
|
|
const messages: AgentMessage[] = [
|
|
makeMessage(1, 4000),
|
|
makeMessage(2, 4000),
|
|
makeMessage(3, 4000),
|
|
makeMessage(4, 4000),
|
|
makeMessage(5, 4000),
|
|
makeMessage(6, 4000),
|
|
];
|
|
const totalTokens = estimateMessagesTokens(messages);
|
|
const maxContextTokens = Math.max(1, Math.floor(totalTokens * 0.5)); // budget = 25%
|
|
const pruned = pruneHistoryForContextShare({
|
|
messages,
|
|
maxContextTokens,
|
|
maxHistoryShare: 0.5,
|
|
parts: 2,
|
|
});
|
|
|
|
const keptIds = pruned.messages.map((msg) => msg.timestamp);
|
|
const expectedSuffix = messages.slice(-keptIds.length).map((msg) => msg.timestamp);
|
|
expect(keptIds).toEqual(expectedSuffix);
|
|
});
|
|
|
|
it("keeps history when already within budget", () => {
|
|
const messages: AgentMessage[] = [makeMessage(1, 1000)];
|
|
const maxContextTokens = 2000;
|
|
const pruned = pruneHistoryForContextShare({
|
|
messages,
|
|
maxContextTokens,
|
|
maxHistoryShare: 0.5,
|
|
parts: 2,
|
|
});
|
|
|
|
expect(pruned.droppedChunks).toBe(0);
|
|
expect(pruned.messages.length).toBe(messages.length);
|
|
expect(pruned.keptTokens).toBe(estimateMessagesTokens(messages));
|
|
});
|
|
});
|