116 lines
3.5 KiB
TypeScript
116 lines
3.5 KiB
TypeScript
import type { AiAssistantMessage, AiAssistantMessageMetadata } from './types';
|
|
|
|
const APPROX_CHARS_PER_TOKEN = 4;
|
|
const MAX_TOOL_OUTPUT_CHARS = 2000;
|
|
|
|
export interface ContextUsage {
|
|
/** Tokens used by the latest model request, or an estimate before the first response */
|
|
usedTokens: number;
|
|
/** Maximum context window size for the model */
|
|
maxTokens: number;
|
|
/** Percentage of context used (0-100) */
|
|
percentage: number;
|
|
/** Whether the context is nearly full (>= threshold) */
|
|
isNearlyFull: boolean;
|
|
/** Whether the context is completely full (>= 100%) */
|
|
isFull: boolean;
|
|
}
|
|
|
|
/** Threshold percentage at which we consider context "nearly full" */
|
|
export const CONTEXT_NEARLY_FULL_THRESHOLD = 90;
|
|
|
|
/** Threshold percentage at which we consider context "full" and block input */
|
|
export const CONTEXT_FULL_THRESHOLD = 95;
|
|
|
|
function getSerializedLength(value: unknown): number {
|
|
if (value == null) {
|
|
return 0;
|
|
}
|
|
|
|
if (typeof value === 'string') {
|
|
return value.length;
|
|
}
|
|
|
|
try {
|
|
return JSON.stringify(value).length;
|
|
} catch {
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
function estimatePartChars(part: unknown): number {
|
|
if (!part || typeof part !== 'object') {
|
|
return getSerializedLength(part);
|
|
}
|
|
|
|
const record = part as Record<string, unknown>;
|
|
const type = typeof record.type === 'string' ? record.type : '';
|
|
|
|
if (type === 'step-start') {
|
|
return 0;
|
|
}
|
|
|
|
if (type.startsWith('tool-')) {
|
|
const outputChars =
|
|
record.state === 'output-available' ? Math.min(getSerializedLength(record.output), MAX_TOOL_OUTPUT_CHARS) : 0;
|
|
|
|
return type.length + getSerializedLength(record.input) + outputChars + getSerializedLength(record.errorText);
|
|
}
|
|
|
|
if (typeof record.text === 'string') {
|
|
return record.text.length;
|
|
}
|
|
|
|
return getSerializedLength(record);
|
|
}
|
|
|
|
function estimateContextTokens(messages: AiAssistantMessage[]): number {
|
|
const chars = messages.reduce((messageTotal, message) => {
|
|
const parts = Array.isArray(message.parts) ? message.parts : [];
|
|
const partsChars = parts.reduce((partTotal, part) => partTotal + estimatePartChars(part), 0);
|
|
|
|
return messageTotal + message.role.length + partsChars;
|
|
}, 0);
|
|
|
|
return Math.ceil(chars / APPROX_CHARS_PER_TOKEN);
|
|
}
|
|
|
|
/**
|
|
* Tracks the latest provider-reported request size as the context guard.
|
|
* Do not sum historical usage: every assistant turn can include a full prompt
|
|
* replay, so summing turns overstates current context. The character estimate
|
|
* is only a fallback before the first provider usage report exists.
|
|
*/
|
|
export function calculateContextUsage(messages: AiAssistantMessage[], contextWindow: number): ContextUsage {
|
|
const latestTokenUsage = [...messages]
|
|
.reverse()
|
|
.map((message) => (message.metadata as AiAssistantMessageMetadata | undefined)?.tokenUsage)
|
|
.find((tokenUsage) => tokenUsage !== undefined);
|
|
const usedTokens = latestTokenUsage?.totalTokens ?? estimateContextTokens(messages);
|
|
|
|
const percentage = contextWindow > 0 ? Math.min(100, Math.round((usedTokens / contextWindow) * 100)) : 0;
|
|
|
|
return {
|
|
usedTokens,
|
|
maxTokens: contextWindow,
|
|
percentage,
|
|
isNearlyFull: percentage >= CONTEXT_NEARLY_FULL_THRESHOLD,
|
|
isFull: percentage >= CONTEXT_FULL_THRESHOLD
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Formats token count for display (e.g., "125K" for 125000)
|
|
*/
|
|
export function formatTokenCount(tokens: number): string {
|
|
if (tokens >= 1_000_000) {
|
|
return `${(tokens / 1_000_000).toFixed(1)}M`;
|
|
}
|
|
|
|
if (tokens >= 1_000) {
|
|
return `${Math.round(tokens / 1_000)}K`;
|
|
}
|
|
|
|
return tokens.toString();
|
|
}
|