From 992fe98386d1bc22f007cd713367154c20c7c5b8 Mon Sep 17 00:00:00 2001 From: Jan Leon Date: Wed, 22 Jul 2026 07:35:01 +0200 Subject: [PATCH] feat(compression): select model-aware tokenizers (#8009) * Add model-aware tokenizer selection * fix: recognize cx Codex model prefix --- open-sse/services/compression/hardBudget.ts | 43 ++++++++------- open-sse/services/compression/stats.ts | 24 +++++--- src/app/api/v1/messages/count_tokens/route.ts | 37 +++++++------ src/shared/utils/tiktokenCounter.ts | 55 ++++++++++++++++--- tests/unit/tiktoken-counter.test.ts | 25 ++++++++- 5 files changed, 134 insertions(+), 50 deletions(-) diff --git a/open-sse/services/compression/hardBudget.ts b/open-sse/services/compression/hardBudget.ts index 33f99a6e02..80706d2781 100644 --- a/open-sse/services/compression/hardBudget.ts +++ b/open-sse/services/compression/hardBudget.ts @@ -10,7 +10,11 @@ import type { CompressionResult } from "./types.ts"; import { scoreToken } from "./ultraHeuristic.ts"; -import { countTextTokens } from "../../../src/shared/utils/tiktokenCounter.ts"; +import { + countTextTokens, + tokenizerContextFromBody, + type TokenizerContext, +} from "../../../src/shared/utils/tiktokenCounter.ts"; import { createCompressionStats } from "./stats.ts"; interface HardBudgetOptions { @@ -69,11 +73,11 @@ interface TaggedUnit { preserve: boolean; } -function tagUnits(units: string[]): TaggedUnit[] { +function tagUnits(units: string[], tokenizerContext: TokenizerContext): TaggedUnit[] { return units.map((u, i) => ({ i, u, - tokens: countTextTokens(u), + tokens: countTextTokens(u, tokenizerContext), score: scoreUnit(u), preserve: mustPreserve(u), })); @@ -84,9 +88,7 @@ function dropToTarget(tagged: TaggedUnit[], targetTokens: number): Set { let tokCount = tagged.reduce((s, x) => s + x.tokens, 0); // Sort droppable candidates by score ascending (lowest first = drop first) - const candidates = tagged - .filter((x) => !x.preserve) - .sort((a, b) => a.score - b.score); + const candidates = tagged.filter((x) => !x.preserve).sort((a, b) => a.score - b.score); for (const candidate of candidates) { if (tokCount <= targetTokens) break; @@ -104,14 +106,18 @@ function rebuildText(tagged: TaggedUnit[], dropped: Set): string { .join("\n"); } -function compressText(text: string, targetTokens: number): string { - const currentTokens = countTextTokens(text); +function compressText( + text: string, + targetTokens: number, + tokenizerContext: TokenizerContext +): string { + const currentTokens = countTextTokens(text, tokenizerContext); if (currentTokens <= targetTokens) return text; const units = splitUnits(text); if (units.length <= 1) return text; - const tagged = tagUnits(units); + const tagged = tagUnits(units, tokenizerContext); const dropped = dropToTarget(tagged, targetTokens); if (dropped.size === 0) return text; @@ -137,17 +143,17 @@ export function applyHardBudget( const messages = extractMessages(body); if (messages.length === 0) return { body, compressed: false, stats: null }; + const tokenizerContext = tokenizerContextFromBody(body); + // Measure total tokens across all messages const totalText = messages .map((m) => (typeof m.content === "string" ? m.content : JSON.stringify(m.content))) .join(" "); - const totalTokens = countTextTokens(totalText); + const totalTokens = countTextTokens(totalText, tokenizerContext); // targetTokens wins when both are set const effectiveTarget = - targetTokens != null - ? targetTokens - : Math.floor(totalTokens * (targetRatio as number)); + targetTokens != null ? targetTokens : Math.floor(totalTokens * (targetRatio as number)); if (totalTokens <= effectiveTarget) { return { body, compressed: false, stats: null }; @@ -158,16 +164,14 @@ export function applyHardBudget( // come back N× over budget). const newMessages = messages.map((m) => { if (typeof m.content !== "string") return m; - const msgTokens = countTextTokens(m.content); + const msgTokens = countTextTokens(m.content, tokenizerContext); const perMsgTarget = totalTokens > 0 ? Math.floor(effectiveTarget * (msgTokens / totalTokens)) : effectiveTarget; - const out = compressText(m.content, perMsgTarget); + const out = compressText(m.content, perMsgTarget, tokenizerContext); return out === m.content ? m : { ...m, content: out }; }); - const changed = newMessages.some( - (m, i) => JSON.stringify(m) !== JSON.stringify(messages[i]) - ); + const changed = newMessages.some((m, i) => JSON.stringify(m) !== JSON.stringify(messages[i])); // Measure the result to detect when preserve-guarded content makes the target // unreachable, so callers are not silently left over budget. @@ -175,7 +179,8 @@ export function applyHardBudget( const resultTokens = countTextTokens( usedMessages .map((m) => (typeof m.content === "string" ? m.content : JSON.stringify(m.content))) - .join(" ") + .join(" "), + tokenizerContext ); const overBudget = resultTokens > effectiveTarget; diff --git a/open-sse/services/compression/stats.ts b/open-sse/services/compression/stats.ts index f2298c7f8a..4dfdfbeb39 100644 --- a/open-sse/services/compression/stats.ts +++ b/open-sse/services/compression/stats.ts @@ -7,6 +7,11 @@ import { DEFAULT_RTK_CONFIG, DEFAULT_COMPRESSION_LANGUAGE_CONFIG, } from "./types.ts"; +import { + countTextTokens, + isCodexTokenizerContext, + tokenizerContextFromBody, +} from "../../../src/shared/utils/tiktokenCounter.ts"; import { anthropicImageTokens, ANTHROPIC_IMAGE_BLOCK_OVERHEAD_TOKENS } from "omniglyph"; const CHARS_PER_TOKEN = 4; @@ -29,9 +34,7 @@ function isAnthropicPngImageBlock(value: unknown): value is AnthropicImageBlock const source = block.source as Record | undefined; if (!source || typeof source !== "object") return false; return ( - source.type === "base64" && - source.media_type === "image/png" && - typeof source.data === "string" + source.type === "base64" && source.media_type === "image/png" && typeof source.data === "string" ); } @@ -119,17 +122,24 @@ function blankImageBlocksAndSumImageTokens(body: Record): { export function estimateCompressionTokens(text: string | object | null | undefined): number { if (!text) return 0; if (typeof text === "string") { - return Math.ceil(text.length / CHARS_PER_TOKEN); + return charTokensOf(text); } try { + const tokenizerContext = tokenizerContextFromBody(text); + const useExactTokenizer = isCodexTokenizerContext(tokenizerContext); const { clone, imageTokens } = blankImageBlocksAndSumImageTokens( text as Record ); if (imageTokens === 0) { - // No recognized image blocks — byte-identical to the legacy behavior. - return Math.ceil(JSON.stringify(text).length / CHARS_PER_TOKEN); + // Keep the legacy character estimate for generic payloads. Codex payloads use + // the model-appropriate tokenizer so their compression stats match hard budgets. + return useExactTokenizer + ? countTextTokens(JSON.stringify(text), tokenizerContext) + : charTokensOf(text); } - return Math.ceil(JSON.stringify(clone).length / CHARS_PER_TOKEN) + imageTokens; + return useExactTokenizer + ? countTextTokens(JSON.stringify(clone), tokenizerContext) + imageTokens + : charTokensOf(clone) + imageTokens; } catch { // Non-serializable/unexpected shape → fall back to the legacy char-count, // never throw out of an estimator. diff --git a/src/app/api/v1/messages/count_tokens/route.ts b/src/app/api/v1/messages/count_tokens/route.ts index 7ae269e95d..af18840799 100644 --- a/src/app/api/v1/messages/count_tokens/route.ts +++ b/src/app/api/v1/messages/count_tokens/route.ts @@ -1,7 +1,7 @@ import { CORS_HEADERS } from "@/shared/utils/cors"; import { v1CountTokensSchema } from "@/shared/validation/schemas"; import { isValidationFailure, validateBody } from "@/shared/validation/helpers"; -import { countTextTokens } from "@/shared/utils/tiktokenCounter"; +import { countTextTokens, type TokenizerContext } from "@/shared/utils/tiktokenCounter"; import { getExecutor } from "@omniroute/open-sse/executors/index.ts"; import { runWithProxyContext } from "@omniroute/open-sse/utils/proxyFetch.ts"; import { getModelInfo } from "@/sse/services/model"; @@ -40,7 +40,10 @@ export async function POST(request) { } const body = validation.data; - const estimated = buildEstimatedCountResponse(body); + const tokenizerContext: TokenizerContext = { + model: typeof body.model === "string" ? body.model : undefined, + }; + const estimated = buildEstimatedCountResponse(body, tokenizerContext); const requestedModel = typeof body.model === "string" ? body.model : ""; if (!requestedModel) { return estimated; @@ -117,22 +120,24 @@ function safeStringify(value) { // content, and `thinking` blocks — counting only `text` (as before) reported // near-zero for those messages and silently broke Claude Code's auto-compaction // (#2337). Image / redacted_thinking blocks are not text-estimable and count 0. -function estimateContentBlockTokens(part) { +function estimateContentBlockTokens(part, tokenizerContext: TokenizerContext) { if (!part || typeof part !== "object") return 0; let tokens = 0; switch (part.type) { case "text": - if (typeof part.text === "string") tokens += countTextTokens(part.text); + if (typeof part.text === "string") tokens += countTextTokens(part.text, tokenizerContext); break; case "tool_use": - if (typeof part.name === "string") tokens += countTextTokens(part.name); - if (part.input !== undefined) tokens += countTextTokens(safeStringify(part.input)); + if (typeof part.name === "string") tokens += countTextTokens(part.name, tokenizerContext); + if (part.input !== undefined) + tokens += countTextTokens(safeStringify(part.input), tokenizerContext); break; case "tool_result": - tokens += estimateToolResultTokens(part.content); + tokens += estimateToolResultTokens(part.content, tokenizerContext); break; case "thinking": - if (typeof part.thinking === "string") tokens += countTextTokens(part.thinking); + if (typeof part.thinking === "string") + tokens += countTextTokens(part.thinking, tokenizerContext); break; default: break; @@ -142,13 +147,13 @@ function estimateContentBlockTokens(part) { // A `tool_result` content can be a plain string or an array of nested blocks // (text / image). Count string content and nested text blocks. -function estimateToolResultTokens(content) { - if (typeof content === "string") return countTextTokens(content); +function estimateToolResultTokens(content, tokenizerContext: TokenizerContext) { + if (typeof content === "string") return countTextTokens(content, tokenizerContext); if (Array.isArray(content)) { let tokens = 0; for (const block of content) { if (block?.type === "text" && typeof block.text === "string") { - tokens += countTextTokens(block.text); + tokens += countTextTokens(block.text, tokenizerContext); } } return tokens; @@ -156,29 +161,29 @@ function estimateToolResultTokens(content) { return 0; } -function buildEstimatedCountResponse(body) { +function buildEstimatedCountResponse(body, tokenizerContext: TokenizerContext = {}) { const messages = Array.isArray(body?.messages) ? body.messages : []; let inputTokens = 0; for (const msg of messages) { if (typeof msg?.content === "string") { - inputTokens += countTextTokens(msg.content); + inputTokens += countTextTokens(msg.content, tokenizerContext); continue; } if (Array.isArray(msg?.content)) { for (const part of msg.content) { - inputTokens += estimateContentBlockTokens(part); + inputTokens += estimateContentBlockTokens(part, tokenizerContext); } } } if (typeof body?.system === "string") { - inputTokens += countTextTokens(body.system); + inputTokens += countTextTokens(body.system, tokenizerContext); } else if (Array.isArray(body?.system)) { for (const block of body.system) { if (block?.type === "text" && typeof block.text === "string") { - inputTokens += countTextTokens(block.text); + inputTokens += countTextTokens(block.text, tokenizerContext); } } } diff --git a/src/shared/utils/tiktokenCounter.ts b/src/shared/utils/tiktokenCounter.ts index d5f53877fa..f36eb999a3 100644 --- a/src/shared/utils/tiktokenCounter.ts +++ b/src/shared/utils/tiktokenCounter.ts @@ -1,20 +1,61 @@ import { getEncoding, type Tiktoken } from "js-tiktoken"; -let encoder: Tiktoken | null = null; +export type TokenizerEncoding = "cl100k_base" | "o200k_base"; -function getEncoder(): Tiktoken { - if (!encoder) encoder = getEncoding("cl100k_base"); - return encoder; +export interface TokenizerContext { + provider?: string | null; + model?: string | null; +} + +export function tokenizerContextFromBody(body: unknown): TokenizerContext { + if (!body || typeof body !== "object" || Array.isArray(body)) return {}; + const record = body as Record; + return { + provider: typeof record.provider === "string" ? record.provider : undefined, + model: typeof record.model === "string" ? record.model : undefined, + }; +} + +const encoders = new Map(); + +function normalize(value: unknown): string { + return typeof value === "string" ? value.trim().toLowerCase() : ""; +} + +export function isCodexTokenizerContext(context?: TokenizerContext): boolean { + const provider = normalize(context?.provider); + const model = normalize(context?.model); + return ( + provider === "codex" || + provider === "cx" || + model.startsWith("codex/") || + model.startsWith("cx/") || + model.includes("codex") + ); +} + +export function resolveTokenizerEncoding(context?: TokenizerContext): TokenizerEncoding { + return isCodexTokenizerContext(context) ? "o200k_base" : "cl100k_base"; +} + +function getEncoder(encoding: TokenizerEncoding): Tiktoken { + const cached = encoders.get(encoding); + if (cached) return cached; + const created = getEncoding(encoding); + encoders.set(encoding, created); + return created; } /** - * Exact token count for a string using cl100k_base (offline, no upstream call). + * Exact token count for a string using the selected offline tokenizer. + * Existing callers retain cl100k_base; Codex callers may pass provider/model context + * to use o200k_base. * Defensive: never throws in a counting path — falls back to a char heuristic. */ -export function countTextTokens(text: string): number { +export function countTextTokens(text: string, context?: TokenizerContext): number { if (!text || typeof text !== "string") return 0; try { - return getEncoder().encode(text).length; + return getEncoder(resolveTokenizerEncoding(context)).encode(text).length; } catch { return Math.ceil(text.length / 4); } diff --git a/tests/unit/tiktoken-counter.test.ts b/tests/unit/tiktoken-counter.test.ts index d5d65da562..110618d1b0 100644 --- a/tests/unit/tiktoken-counter.test.ts +++ b/tests/unit/tiktoken-counter.test.ts @@ -1,11 +1,34 @@ import { test } from "node:test"; import assert from "node:assert/strict"; -import { countTextTokens } from "../../src/shared/utils/tiktokenCounter.ts"; +import { + countTextTokens, + isCodexTokenizerContext, + resolveTokenizerEncoding, +} from "../../src/shared/utils/tiktokenCounter.ts"; test("countTextTokens returns exact tiktoken count for a known string", () => { assert.equal(countTextTokens("hello world"), 2); // cl100k_base }); +test("Codex context selects o200k_base without changing the default", () => { + assert.equal(resolveTokenizerEncoding(), "cl100k_base"); + assert.equal(resolveTokenizerEncoding({ provider: "codex" }), "o200k_base"); + assert.equal(resolveTokenizerEncoding({ provider: "cx" }), "o200k_base"); + assert.equal(resolveTokenizerEncoding({ model: "codex/gpt-5.6-sol" }), "o200k_base"); + assert.equal(resolveTokenizerEncoding({ model: "cx/gpt-5.6-sol" }), "o200k_base"); + assert.equal(resolveTokenizerEncoding({ provider: "openai", model: "gpt-5.6" }), "cl100k_base"); + assert.equal(isCodexTokenizerContext({ provider: "codex" }), true); + assert.equal(isCodexTokenizerContext({ provider: "openai" }), false); +}); + +test("Codex token counting uses the o200k encoder", () => { + const text = "antidisestablishmentarianism 中文ภาษาไทย"; + assert.notEqual( + countTextTokens(text, { provider: "codex" }), + countTextTokens(text, { provider: "openai" }) + ); +}); + test("countTextTokens handles empty and non-string safely", () => { assert.equal(countTextTokens(""), 0); assert.equal(countTextTokens(undefined as unknown as string), 0);