diff --git a/src/app/api/v1/messages/count_tokens/route.ts b/src/app/api/v1/messages/count_tokens/route.ts index bc532d20e3..051c2f4860 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 { estimateTokens } from "@/shared/utils/costEstimator"; +import { countTextTokens } 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"; @@ -101,31 +101,31 @@ export async function POST(request) { function buildEstimatedCountResponse(body) { const messages = Array.isArray(body?.messages) ? body.messages : []; - let totalChars = 0; + let inputTokens = 0; for (const msg of messages) { if (typeof msg?.content === "string") { - totalChars += msg.content.length; + inputTokens += countTextTokens(msg.content); continue; } if (Array.isArray(msg?.content)) { for (const part of msg.content) { if (part?.type === "text" && typeof part.text === "string") { - totalChars += part.text.length; + inputTokens += countTextTokens(part.text); } } } } if (typeof body?.system === "string") { - totalChars += body.system.length; + inputTokens += countTextTokens(body.system); } return new Response( JSON.stringify({ - input_tokens: totalChars > 0 ? Math.ceil(totalChars / 4) : estimateTokens(""), - source: "estimated", + input_tokens: inputTokens, + source: "local", }), { headers: { "Content-Type": "application/json", ...CORS_HEADERS }, diff --git a/src/shared/utils/tiktokenCounter.ts b/src/shared/utils/tiktokenCounter.ts new file mode 100644 index 0000000000..d5f53877fa --- /dev/null +++ b/src/shared/utils/tiktokenCounter.ts @@ -0,0 +1,21 @@ +import { getEncoding, type Tiktoken } from "js-tiktoken"; + +let encoder: Tiktoken | null = null; + +function getEncoder(): Tiktoken { + if (!encoder) encoder = getEncoding("cl100k_base"); + return encoder; +} + +/** + * Exact token count for a string using cl100k_base (offline, no upstream call). + * Defensive: never throws in a counting path — falls back to a char heuristic. + */ +export function countTextTokens(text: string): number { + if (!text || typeof text !== "string") return 0; + try { + return getEncoder().encode(text).length; + } catch { + return Math.ceil(text.length / 4); + } +} diff --git a/tests/unit/messages-count-tokens-route.test.ts b/tests/unit/messages-count-tokens-route.test.ts index 062b8a654f..ffcdcdd461 100644 --- a/tests/unit/messages-count-tokens-route.test.ts +++ b/tests/unit/messages-count-tokens-route.test.ts @@ -97,8 +97,20 @@ test("messages/count_tokens falls back to estimate when model is missing", async assert.equal(response.status, 200); const body = (await response.json()) as any; - assert.equal(body.input_tokens, 3); - assert.equal(body.source, "estimated"); + assert.equal(body.input_tokens, 4); // tiktoken: "abcd"=1 + "12345678"=3 + assert.equal(body.source, "local"); +}); + +test("count_tokens fallback uses exact tiktoken count with source=local", async () => { + const req = new Request("http://localhost/v1/messages/count_tokens", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ messages: [{ role: "user", content: "hello world" }] }), + }); + const res = await POST(req); + const json = (await res.json()) as any; + assert.equal(json.source, "local"); + assert.equal(json.input_tokens, 2); // exact cl100k_base count, not Math.ceil(11/4)=3 }); test("messages/count_tokens falls back to estimate when real upstream count fails", async () => { @@ -121,8 +133,8 @@ test("messages/count_tokens falls back to estimate when real upstream count fail assert.equal(response.status, 200); const body = (await response.json()) as any; - assert.equal(body.input_tokens, 1); - assert.equal(body.source, "estimated"); + assert.equal(body.input_tokens, 1); // tiktoken: "abcd"=1 + assert.equal(body.source, "local"); } finally { globalThis.fetch = originalFetch; } diff --git a/tests/unit/tiktoken-counter.test.ts b/tests/unit/tiktoken-counter.test.ts new file mode 100644 index 0000000000..d5d65da562 --- /dev/null +++ b/tests/unit/tiktoken-counter.test.ts @@ -0,0 +1,19 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { countTextTokens } 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("countTextTokens handles empty and non-string safely", () => { + assert.equal(countTextTokens(""), 0); + assert.equal(countTextTokens(undefined as unknown as string), 0); +}); + +test("countTextTokens is additive-ish and monotonic for longer text", () => { + const short = countTextTokens("the quick brown fox"); + const long = countTextTokens("the quick brown fox jumps over the lazy dog"); + assert.ok(long > short); + assert.ok(short > 0); +});