59 lines
1.9 KiB
TypeScript
59 lines
1.9 KiB
TypeScript
|
import assert from 'node:assert/strict';
|
||
|
|
import test from 'node:test';
|
||
|
|
import {
|
||
|
|
computeCompactionBuffer,
|
||
|
|
computeCompactionThreshold,
|
||
|
|
computeTotalInputTokens,
|
||
|
|
DEFAULT_MAX_OUTPUT_TOKENS,
|
||
|
|
shouldCompactByBudget,
|
||
|
|
} from './contextBudget.ts';
|
||
|
|
import type { ModelMessage } from 'ai';
|
||
|
|
|
||
|
|
test('computeCompactionThreshold reserves output and buffer', () => {
|
||
|
|
const threshold = computeCompactionThreshold({
|
||
|
|
contextWindow: 128_000,
|
||
|
|
maxOutputTokens: 4096,
|
||
|
|
});
|
||
|
|
const buffer = computeCompactionBuffer(128_000, 4096);
|
||
|
|
assert.equal(threshold, 128_000 - 4096 - buffer - 150);
|
||
|
|
assert.ok(threshold < 128_000 * 0.85);
|
||
|
|
});
|
||
|
|
|
||
|
|
test('shouldCompactByBudget triggers when total input exceeds threshold', () => {
|
||
|
|
const messages: ModelMessage[] = [{ role: 'user', content: 'x'.repeat(400_000) }];
|
||
|
|
assert.equal(shouldCompactByBudget({
|
||
|
|
messages,
|
||
|
|
contextWindow: 128_000,
|
||
|
|
maxOutputTokens: 4096,
|
||
|
|
providerId: 'openai',
|
||
|
|
}), true);
|
||
|
|
});
|
||
|
|
|
||
|
|
test('computeTotalInputTokens includes system and tool names', () => {
|
||
|
|
const messages: ModelMessage[] = [{ role: 'user', content: 'hello' }];
|
||
|
|
const withExtras = computeTotalInputTokens({
|
||
|
|
messages,
|
||
|
|
systemPrompt: 'system prompt',
|
||
|
|
toolNames: ['terminal_execute', 'sftp_read'],
|
||
|
|
providerId: 'anthropic',
|
||
|
|
});
|
||
|
|
const base = computeTotalInputTokens({ messages, providerId: 'anthropic' });
|
||
|
|
assert.ok(withExtras > base);
|
||
|
|
});
|
||
|
|
|
||
|
|
test('computeCompactionThreshold keeps a reasonable threshold for small context windows', () => {
|
||
|
|
const threshold8k = computeCompactionThreshold({
|
||
|
|
contextWindow: 8_192,
|
||
|
|
maxOutputTokens: DEFAULT_MAX_OUTPUT_TOKENS,
|
||
|
|
});
|
||
|
|
assert.ok(threshold8k > 1_000);
|
||
|
|
assert.ok(threshold8k < 8_192 * 0.85);
|
||
|
|
|
||
|
|
const threshold4k = computeCompactionThreshold({
|
||
|
|
contextWindow: 4_096,
|
||
|
|
maxOutputTokens: DEFAULT_MAX_OUTPUT_TOKENS,
|
||
|
|
});
|
||
|
|
assert.ok(threshold4k > 500);
|
||
|
|
assert.ok(threshold4k < 4_096 * 0.85);
|
||
|
|
});
|