feature: token counting for different models

Author: caozhiyuan

src/lib/tokenizer.ts

@@ -37,10 +37,7 @@ const calculateToolCallsTokens = (
   let tokens = 0
   for (const toolCall of toolCalls) {
     tokens += constants.funcInit
-    tokens += encoder.encode(toolCall.id).length
-    tokens += encoder.encode(toolCall.type).length
-    tokens += encoder.encode(toolCall.function.name).length
-    tokens += encoder.encode(toolCall.function.arguments).length
+    tokens += encoder.encode(JSON.stringify(toolCall)).length
   }
   tokens += constants.funcEnd
   return tokens
@@ -57,9 +54,6 @@ const calculateContentPartsTokens = (
   for (const part of contentParts) {
     if (part.type === "image_url") {
       tokens += encoder.encode(part.image_url.url).length + 85
-      if (part.image_url.detail === "high") {
-        tokens += 85
-      }
     } else if (part.text) {
       tokens += encoder.encode(part.text).length
     }
@@ -346,6 +340,7 @@ export const getTokenCount = async (
     inputTokens += numTokensForTools(payload.tools, encoder, constants)
   }
   const outputTokens = calculateTokens(outputMessages, encoder, constants)
+
   return {
     input: inputTokens,
     output: outputTokens,

src/routes/messages/count-tokens-handler.ts

@@ -13,12 +13,12 @@ import { translateToOpenAI } from "./non-stream-translation"
  */
 export async function handleCountTokens(c: Context) {
   try {
+    const anthropicBeta = c.req.header("anthropic-beta")
+
     const anthropicPayload = await c.req.json<AnthropicMessagesPayload>()
 
-    // Convert to OpenAI format for token counting
     const openAIPayload = translateToOpenAI(anthropicPayload)
 
-    // Find the selected model
     const selectedModel = state.models?.data.find(
       (model) => model.id === anthropicPayload.model,
     )
@@ -30,17 +30,39 @@ export async function handleCountTokens(c: Context) {
       })
     }
 
-    // Calculate token count
     const tokenCount = await getTokenCount(openAIPayload, selectedModel)
-    consola.debug("Token count:", tokenCount)
 
-    // Return response in Anthropic API format
+    if (anthropicPayload.tools && anthropicPayload.tools.length > 0) {
+      let mcpToolExist = false
+      if (anthropicBeta?.startsWith("claude-code")) {
+        mcpToolExist = anthropicPayload.tools.some((tool) =>
+          tool.name.startsWith("mcp__"),
+        )
+      }
+      if (!mcpToolExist) {
+        if (anthropicPayload.model.startsWith("claude")) {
+          // https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/overview#pricing
+          tokenCount.input = tokenCount.input + 346
+        } else if (anthropicPayload.model.startsWith("grok")) {
+          tokenCount.input = tokenCount.input + 480
+        }
+      }
+    }
+
+    let finalTokenCount = tokenCount.input + tokenCount.output
+    if (anthropicPayload.model.startsWith("claude")) {
+      finalTokenCount = Math.round(finalTokenCount * 1.15)
+    } else if (anthropicPayload.model.startsWith("grok")) {
+      finalTokenCount = Math.round(finalTokenCount * 1.03)
+    }
+
+    consola.info("Token count:", finalTokenCount)
+
     return c.json({
-      input_tokens: tokenCount.input,
+      input_tokens: finalTokenCount,
     })
   } catch (error) {
     consola.error("Error counting tokens:", error)
-    // Return default value on error
     return c.json({
       input_tokens: 1,
     })