( messages: Anthropic.Beta.Messages.BetaMessageParam[], tools: Anthropic.Beta.Messages.BetaToolUnion[], )
| 138 | } |
| 139 | |
| 140 | export async function countMessagesTokensWithAPI( |
| 141 | messages: Anthropic.Beta.Messages.BetaMessageParam[], |
| 142 | tools: Anthropic.Beta.Messages.BetaToolUnion[], |
| 143 | ): Promise<number | null> { |
| 144 | return withTokenCountVCR(messages, tools, async () => { |
| 145 | try { |
| 146 | const model = getMainLoopModel() |
| 147 | const betas = getModelBetas(model) |
| 148 | const containsThinking = hasThinkingBlocks(messages) |
| 149 | |
| 150 | if (getAPIProvider() === 'bedrock') { |
| 151 | // @anthropic-sdk/bedrock-sdk doesn't support countTokens currently |
| 152 | return countTokensWithBedrock({ |
| 153 | model: normalizeModelStringForAPI(model), |
| 154 | messages, |
| 155 | tools, |
| 156 | betas, |
| 157 | containsThinking, |
| 158 | }) |
| 159 | } |
| 160 | |
| 161 | const anthropic = await getAnthropicClient({ |
| 162 | maxRetries: 1, |
| 163 | model, |
| 164 | source: 'count_tokens', |
| 165 | }) |
| 166 | |
| 167 | const filteredBetas = |
| 168 | getAPIProvider() === 'vertex' |
| 169 | ? betas.filter(b => VERTEX_COUNT_TOKENS_ALLOWED_BETAS.has(b)) |
| 170 | : betas |
| 171 | |
| 172 | const response = await anthropic.beta.messages.countTokens({ |
| 173 | model: normalizeModelStringForAPI(model), |
| 174 | messages: |
| 175 | // When we pass tools and no messages, we need to pass a dummy message |
| 176 | // to get an accurate tool token count. |
| 177 | messages.length > 0 ? messages : [{ role: 'user', content: 'foo' }], |
| 178 | tools, |
| 179 | ...(filteredBetas.length > 0 && { betas: filteredBetas }), |
| 180 | // Enable thinking if messages contain thinking blocks |
| 181 | ...(containsThinking && { |
| 182 | thinking: { |
| 183 | type: 'enabled', |
| 184 | budget_tokens: TOKEN_COUNT_THINKING_BUDGET, |
| 185 | }, |
| 186 | }), |
| 187 | }) |
| 188 | |
| 189 | if (typeof response.input_tokens !== 'number') { |
| 190 | // Vertex client throws |
| 191 | // Bedrock client succeeds with { Output: { __type: 'com.amazon.coral.service#UnknownOperationException' }, Version: '1.0' } |
| 192 | return null |
| 193 | } |
| 194 | |
| 195 | return response.input_tokens |
| 196 | } catch (error) { |
| 197 | logError(error) |
no test coverage detected