File size: 2,507 Bytes
68d7816 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | import type { Message, ToolDescription } from '#/llm/message';
import { emptyUsage, type TokenUsage } from '#/llm/usage';
import type { AssistantEntry, HistoryMessage } from './turn';
const MEDIA_TOKEN_ESTIMATE = 2000;
export interface ContextUsagePrefix {
systemPrompt?: string;
tools?: readonly ToolDescription[];
}
export function calculateContextTokens(usage: TokenUsage): number {
return usage.inputOther + usage.inputCacheRead + usage.inputCacheCreation + usage.output;
}
export function estimateTextTokens(text: string): number {
let asciiCount = 0;
let nonAsciiCount = 0;
for (const char of text) {
if ((char.codePointAt(0) as number) <= 127) {
asciiCount++;
} else {
nonAsciiCount++;
}
}
return Math.ceil(asciiCount / 4) + nonAsciiCount;
}
export function estimateMessageTokens(message: Message): number {
let total = estimateTextTokens(message.role);
for (const part of message.content) {
switch (part.type) {
case 'text':
total += estimateTextTokens(part.text);
break;
case 'think':
total += estimateTextTokens(part.think);
break;
case 'image_url':
case 'audio_url':
case 'video_url':
total += MEDIA_TOKEN_ESTIMATE;
break;
}
}
if (message.role === 'assistant') {
for (const call of message.toolCalls) {
total += estimateTextTokens(call.name);
total += estimateTextTokens(call.arguments ?? '');
}
}
return total;
}
export function estimateUsedContextTokens(
history: readonly HistoryMessage[],
prefix?: ContextUsagePrefix,
): number {
let lastUsageIndex = -1;
let usageTokens = 0;
for (let i = history.length - 1; i >= 0; i--) {
const entry = history[i] as HistoryMessage;
if (entry.message.role !== 'assistant') continue;
const tokens = calculateContextTokens((entry as AssistantEntry).meta?.usage ?? emptyUsage());
if (tokens > 0) {
lastUsageIndex = i;
usageTokens = tokens;
break;
}
}
let tokens = usageTokens;
for (let i = lastUsageIndex + 1; i < history.length; i++) {
tokens += estimateMessageTokens((history[i] as HistoryMessage).message);
}
if (lastUsageIndex === -1 && prefix !== undefined) {
if (prefix.systemPrompt !== undefined) {
tokens += estimateTextTokens(prefix.systemPrompt);
}
if (prefix.tools !== undefined && prefix.tools.length > 0) {
tokens += estimateTextTokens(JSON.stringify(prefix.tools));
}
}
return tokens;
}
|