72 lines
3.9 KiB
TypeScript
72 lines
3.9 KiB
TypeScript
import { parseProviderUsage } from "./provider-usage";
|
|
|
|
export type TokenLimitCategory = "NONE" | "CONTEXT_LIMIT" | "OUTPUT_LIMIT" | "MAX_OUTPUT_TOKENS" | "UNKNOWN";
|
|
|
|
export interface TokenLimitClassification {
|
|
category: TokenLimitCategory;
|
|
detected: boolean;
|
|
retryable: boolean;
|
|
reason: string | null;
|
|
actual: number | null;
|
|
limit: number | null;
|
|
}
|
|
|
|
export interface TokenLimitOptions {
|
|
contextTokens?: number;
|
|
outputTokens?: number;
|
|
maxOutputTokens?: number;
|
|
}
|
|
|
|
function textOf(input: unknown): string {
|
|
try {
|
|
const serialized = JSON.stringify(input);
|
|
return (serialized ?? String(input)).toLowerCase();
|
|
} catch { return String(input).toLowerCase(); }
|
|
}
|
|
|
|
export const classifyLimit = classifyTokenLimit;
|
|
|
|
function numberAt(input: unknown, names: Set<string>, seen = new Set<object>()): number | null {
|
|
if (typeof input !== "object" || input === null) return null;
|
|
if (seen.has(input)) return null;
|
|
seen.add(input);
|
|
if (Array.isArray(input)) {
|
|
for (const value of input) { const result = numberAt(value, names, seen); if (result !== null) return result; }
|
|
return null;
|
|
}
|
|
for (const [key, value] of Object.entries(input)) {
|
|
const normalized = key.replace(/[A-Z]/gu, (letter) => `_${letter.toLowerCase()}`).toLowerCase();
|
|
if (names.has(normalized) && typeof value === "number" && Number.isFinite(value)) return value;
|
|
const result = numberAt(value, names, seen);
|
|
if (result !== null) return result;
|
|
}
|
|
return null;
|
|
}
|
|
|
|
function result(category: TokenLimitCategory, reason: string | null, actual: number | null, limit: number | null): TokenLimitClassification {
|
|
return { category, detected: category !== "NONE", retryable: category === "CONTEXT_LIMIT" || category === "OUTPUT_LIMIT", reason, actual, limit };
|
|
}
|
|
|
|
/** Recognize provider errors, finish reasons and measured usage consistently. */
|
|
export function classifyTokenLimit(input: unknown, options: TokenLimitOptions = {}): TokenLimitClassification {
|
|
const text = textOf(input);
|
|
const usage = parseProviderUsage(input);
|
|
const contextActual = numberAt(input, new Set(["context_tokens", "contexttokens", "prompt_tokens", "prompttokens", "input_tokens", "inputtokens"]));
|
|
const outputActual = usage.outputTokens ?? numberAt(input, new Set(["output_tokens", "outputtokens", "completion_tokens", "completiontokens"]));
|
|
const contextMarker = /(context[_ -]?length|context[_ -]?window|maximum[_ -]?context|prompt[_ -]?too[_ -]?long|input[_ -]?too[_ -]?long|context[_ -]?limit|context[_ -]?exceed)/u.test(text);
|
|
const outputMarker = /(output[_ -]?limit|completion[_ -]?limit|output[_ -]?too[_ -]?long|output[_ -]?exceed|completion[_ -]?tokens?.{0,20}(limit|exceed|too))/u.test(text);
|
|
const maxMarker = /(max(?:imum)?[_ -]?output[_ -]?tokens|max[_ -]?completion[_ -]?tokens|max[_ -]?tokens|finish[_ -]?reason.{0,12}(length|max_tokens)|stop[_ -]?reason.{0,12}max_tokens|incomplete.{0,30}max_output_tokens)/u.test(text);
|
|
|
|
if (contextMarker || (options.contextTokens !== undefined && contextActual !== null && contextActual > options.contextTokens)) {
|
|
return result("CONTEXT_LIMIT", "context token limit exceeded", contextActual, options.contextTokens ?? null);
|
|
}
|
|
if (outputMarker || (options.outputTokens !== undefined && outputActual !== null && outputActual > options.outputTokens)) {
|
|
return result("OUTPUT_LIMIT", "output token limit exceeded", outputActual, options.outputTokens ?? null);
|
|
}
|
|
if (maxMarker || (options.maxOutputTokens !== undefined && outputActual !== null && outputActual >= options.maxOutputTokens && /(?:length|max[_ -]?tokens)/u.test(text))) {
|
|
return result("MAX_OUTPUT_TOKENS", "generation stopped at max_output_tokens", outputActual, options.maxOutputTokens ?? null);
|
|
}
|
|
if (text.includes("token") && /(limit|exceed|truncat|too long)/u.test(text)) return result("UNKNOWN", "provider reported an unclassified token limit", null, null);
|
|
return result("NONE", null, null, null);
|
|
}
|