diff --git a/src/app/api/v1/messages/count_tokens/route.js b/src/app/api/v1/messages/count_tokens/route.js index c5a2918f..2e08cd09 100644 --- a/src/app/api/v1/messages/count_tokens/route.js +++ b/src/app/api/v1/messages/count_tokens/route.js @@ -11,6 +11,64 @@ export async function OPTIONS() { return new Response(null, { headers: CORS_HEADERS }); } +function countValueChars(value) { + if (value == null) return 0; + if (typeof value === "string") return value.length; + if (typeof value === "number" || typeof value === "boolean") { + return String(value).length; + } + if (Array.isArray(value)) { + return value.reduce((total, item) => total + countValueChars(item), 0); + } + if (typeof value === "object") { + return Object.entries(value).reduce((total, [key, item]) => { + return total + key.length + countValueChars(item); + }, 0); + } + return 0; +} + +function countContentBlockChars(block) { + if (block == null) return 0; + if (typeof block === "string") return block.length; + if (typeof block !== "object") return countValueChars(block); + + switch (block.type) { + case "text": + return countValueChars(block.text); + case "tool_use": + return countValueChars(block.name) + countValueChars(block.input); + case "tool_result": + return countValueChars(block.content); + case "thinking": + return countValueChars(block.thinking); + default: + return countValueChars(block); + } +} + +function countMessageChars(message) { + if (!message || typeof message !== "object") return 0; + const content = message.content; + + if (typeof content === "string") return content.length; + if (Array.isArray(content)) { + return content.reduce((total, block) => total + countContentBlockChars(block), 0); + } + return countValueChars(content); +} + +export function estimateAnthropicInputTokens(body = {}) { + const messages = Array.isArray(body.messages) ? body.messages : []; + let totalChars = countValueChars(body.system) + countValueChars(body.tools); + + for (const msg of messages) { + totalChars += countMessageChars(msg); + } + + return Math.ceil(totalChars / 4); +} + /** * POST /v1/messages/count_tokens - Mock token count response */ @@ -25,23 +83,7 @@ export async function POST(request) { }); } - // Estimate token count based on content length - const messages = body.messages || []; - let totalChars = 0; - for (const msg of messages) { - if (typeof msg.content === "string") { - totalChars += msg.content.length; - } else if (Array.isArray(msg.content)) { - for (const part of msg.content) { - if (part.type === "text" && part.text) { - totalChars += part.text.length; - } - } - } - } - - // Rough estimate: ~4 chars per token - const inputTokens = Math.ceil(totalChars / 4); + const inputTokens = estimateAnthropicInputTokens(body); return new Response(JSON.stringify({ input_tokens: inputTokens diff --git a/tests/unit/count-tokens.test.js b/tests/unit/count-tokens.test.js new file mode 100644 index 00000000..92f7f9c3 --- /dev/null +++ b/tests/unit/count-tokens.test.js @@ -0,0 +1,84 @@ +import { describe, expect, it } from "vitest"; + +import { POST } from "../../src/app/api/v1/messages/count_tokens/route.js"; + +async function countTokens(body) { + const response = await POST(new Request("https://9router.local/v1/messages/count_tokens", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + })); + + expect(response.status).toBe(200); + return response.json(); +} + +describe("Anthropic count_tokens estimator", () => { + it("preserves the existing plain text estimate", async () => { + const result = await countTokens({ + messages: [ + { + role: "user", + content: "hello world", + }, + ], + }); + + expect(result.input_tokens).toBe(3); + }); + + it("counts tool and thinking content blocks that carry context", async () => { + const result = await countTokens({ + messages: [ + { + role: "assistant", + content: [ + { + type: "tool_use", + id: "toolu_01", + name: "Read", + input: { file_path: "/tmp/example.txt" }, + }, + { + type: "thinking", + thinking: "Need to inspect the file before answering.", + }, + ], + }, + { + role: "user", + content: [ + { + type: "tool_result", + tool_use_id: "toolu_01", + content: "line1 line2 line3 some file content here", + }, + ], + }, + ], + }); + + expect(result.input_tokens).toBeGreaterThan(0); + }); + + it("counts system prompts and tool definitions", async () => { + const result = await countTokens({ + system: "You are a coding assistant.", + tools: [ + { + name: "Read", + description: "Read a file", + input_schema: { + type: "object", + properties: { + file_path: { type: "string" }, + }, + }, + }, + ], + messages: [], + }); + + expect(result.input_tokens).toBeGreaterThan(0); + }); +});