diff --git a/src/services/tokenEstimation.test.ts b/src/services/tokenEstimation.test.ts new file mode 100644 index 000000000..3a2bfb2c4 --- /dev/null +++ b/src/services/tokenEstimation.test.ts @@ -0,0 +1,80 @@ +import type { Anthropic } from '@anthropic-ai/sdk' +import { expect, mock, test } from 'bun:test' +import { jsonStringify } from '../utils/slowOperations.js' +import { __test, roughTokenCountEstimation } from './tokenEstimation.js' + +function createTextTool(): Anthropic.Beta.Messages.BetaToolUnion { + return { + name: 'lookup_docs', + description: 'Look up project documentation.', + input_schema: { + type: 'object', + properties: { + query: { type: 'string' }, + }, + required: ['query'], + }, + } +} + +test('countMessagesTokensWithClient falls back when shim client lacks countTokens', async () => { + const content = 'hello from an openai-compatible provider' + + const result = await __test.countMessagesTokensWithClient({ + messagesClient: {}, + model: 'gpt-4o', + messages: [{ role: 'user', content }], + tools: [], + filteredBetas: [], + containsThinking: false, + }) + + expect(result).toBe(roughTokenCountEstimation(content)) +}) + +test('countMessagesTokensWithClient includes tool overhead in fallback estimates', async () => { + const content = 'count this request with tool definitions' + const tools = [createTextTool()] + + const result = await __test.countMessagesTokensWithClient({ + messagesClient: {}, + model: 'gpt-4o', + messages: [{ role: 'user', content }], + tools, + filteredBetas: [], + containsThinking: false, + }) + + expect(result).toBe( + roughTokenCountEstimation(content) + + 500 + + roughTokenCountEstimation(jsonStringify(tools)), + ) +}) + +test('countMessagesTokensWithClient uses countTokens when the client supports it', async () => { + const countTokens = mock(async (_params: unknown) => ({ input_tokens: 42 })) + const messages: Anthropic.Beta.Messages.BetaMessageParam[] = [ + { role: 'user', content: 'use exact count when available' }, + ] + + const result = await __test.countMessagesTokensWithClient({ + messagesClient: { + countTokens: + countTokens as unknown as Anthropic['beta']['messages']['countTokens'], + }, + model: 'gpt-4o', + messages, + tools: [], + filteredBetas: [], + containsThinking: false, + }) + + expect(countTokens).toHaveBeenCalledTimes(1) + expect(countTokens.mock.calls[0]?.[0]).toEqual({ + model: 'gpt-4o', + messages, + tools: [], + }) + expect(result).toBe(42) +}) diff --git a/src/services/tokenEstimation.ts b/src/services/tokenEstimation.ts index 832bb63b5..e5da94640 100644 --- a/src/services/tokenEstimation.ts +++ b/src/services/tokenEstimation.ts @@ -31,6 +31,12 @@ import { withTokenCountVCR } from './vcr.js' // API constraint: max_tokens must be greater than thinking.budget_tokens const TOKEN_COUNT_THINKING_BUDGET = 1024 const TOKEN_COUNT_MAX_TOKENS = 2048 +// Keep this local to avoid importing analyzeContext.ts, which already depends on tokenEstimation. +const ROUGH_TOOL_TOKEN_COUNT_OVERHEAD = 500 + +type CountTokensMessagesClient = { + countTokens?: Anthropic['beta']['messages']['countTokens'] +} /** * Check if messages contain thinking blocks @@ -163,36 +169,27 @@ export async function countMessagesTokensWithAPI( model, source: 'count_tokens', }) + const messagesClient = ( + anthropic as { + beta?: { + messages?: CountTokensMessagesClient + } + } + ).beta?.messages const filteredBetas = getAPIProvider() === 'vertex' ? betas.filter(b => VERTEX_COUNT_TOKENS_ALLOWED_BETAS.has(b)) : betas - const response = await anthropic.beta.messages.countTokens({ - model: normalizeModelStringForAPI(model), - messages: - // When we pass tools and no messages, we need to pass a dummy message - // to get an accurate tool token count. - messages.length > 0 ? messages : [{ role: 'user', content: 'foo' }], + return countMessagesTokensWithClient({ + messagesClient, + model, + messages, tools, - ...(filteredBetas.length > 0 && { betas: filteredBetas }), - // Enable thinking if messages contain thinking blocks - ...(containsThinking && { - thinking: { - type: 'enabled', - budget_tokens: TOKEN_COUNT_THINKING_BUDGET, - }, - }), + filteredBetas, + containsThinking, }) - - if (typeof response.input_tokens !== 'number') { - // Vertex client throws - // Bedrock client succeeds with { Output: { __type: 'com.amazon.coral.service#UnknownOperationException' }, Version: '1.0' } - return null - } - - return response.input_tokens } catch (error) { logError(error) return null @@ -200,6 +197,79 @@ export async function countMessagesTokensWithAPI( }) } +async function countMessagesTokensWithClient({ + messagesClient, + model, + messages, + tools, + filteredBetas, + containsThinking, +}: { + messagesClient: CountTokensMessagesClient | undefined + model: string + messages: Anthropic.Beta.Messages.BetaMessageParam[] + tools: Anthropic.Beta.Messages.BetaToolUnion[] + filteredBetas: string[] + containsThinking: boolean +}): Promise { + if (typeof messagesClient?.countTokens !== 'function') { + return roughTokenCountEstimationForCountTokensFallback(messages, tools) + } + + const response = await messagesClient.countTokens({ + model: normalizeModelStringForAPI(model), + messages: + // When we pass tools and no messages, we need to pass a dummy message + // to get an accurate tool token count. + messages.length > 0 ? messages : [{ role: 'user', content: 'foo' }], + tools, + ...(filteredBetas.length > 0 && { betas: filteredBetas }), + // Enable thinking if messages contain thinking blocks + ...(containsThinking && { + thinking: { + type: 'enabled', + budget_tokens: TOKEN_COUNT_THINKING_BUDGET, + }, + }), + }) + + if (typeof response.input_tokens !== 'number') { + // Vertex client throws + // Bedrock client succeeds with { Output: { __type: 'com.amazon.coral.service#UnknownOperationException' }, Version: '1.0' } + return null + } + + return response.input_tokens +} + +function roughTokenCountEstimationForCountTokensFallback( + messages: Anthropic.Beta.Messages.BetaMessageParam[], + tools: Anthropic.Beta.Messages.BetaToolUnion[], +): number { + let totalTokens = 0 + + for (const message of messages) { + totalTokens += roughTokenCountEstimationForContent( + message.content as + | string + | Array + | Array + | undefined, + ) + } + + if (tools.length > 0) { + totalTokens += + ROUGH_TOOL_TOKEN_COUNT_OVERHEAD + + roughTokenCountEstimation(jsonStringify(tools)) + } + + return totalTokens +} + +// Test-only surface for fallback dispatch without process-wide module mocks. +export const __test = { countMessagesTokensWithClient } + export function roughTokenCountEstimation( content: string, bytesPerToken: number = 4,