Add AnthropicModel count_tokens (#3458)

nathan-gage · claude · web-flow · commit 80a1ce0cd6f8 · 2025-11-17T17:28:31.000-06:00
Co-authored-by: Claude &lt;noreply@anthropic.com&gt;
diff --git a/pydantic_ai_slim/pydantic_ai/models/anthropic.py b/pydantic_ai_slim/pydantic_ai/models/anthropic.py
@@ -55,7 +55,7 @@
 
 
 try:
-    from anthropic import NOT_GIVEN, APIStatusError, AsyncStream, omit as OMIT
+    from anthropic import NOT_GIVEN, APIStatusError, AsyncAnthropicBedrock, AsyncStream, omit as OMIT
     from anthropic.types.beta import (
         BetaBase64PDFBlockParam,
         BetaBase64PDFSourceParam,
@@ -76,6 +76,7 @@
         BetaMemoryTool20250818Param,
         BetaMessage,
         BetaMessageParam,
+        BetaMessageTokensCount,
         BetaMetadataParam,
         BetaPlainTextSourceParam,
         BetaRawContentBlockDeltaEvent,
@@ -239,6 +240,23 @@ async def request(
         model_response = self._process_response(response)
         return model_response
 
+    async def count_tokens(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> usage.RequestUsage:
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
+
+        response = await self._messages_count_tokens(
+            messages, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters
+        )
+
+        return usage.RequestUsage(input_tokens=response.input_tokens)
+
     @asynccontextmanager
     async def request_stream(
         self,
@@ -310,28 +328,12 @@ async def _messages_create(
         tools = self._get_tools(model_request_parameters, model_settings)
         tools, mcp_servers, beta_features = self._add_builtin_tools(tools, model_request_parameters)
 
-        tool_choice: BetaToolChoiceParam | None
-
-        if not tools:
-            tool_choice = None
-        else:
-            if not model_request_parameters.allow_text_output:
-                tool_choice = {'type': 'any'}
-            else:
-                tool_choice = {'type': 'auto'}
-
-            if (allow_parallel_tool_calls := model_settings.get('parallel_tool_calls')) is not None:
-                tool_choice['disable_parallel_tool_use'] = not allow_parallel_tool_calls
+        tool_choice = self._infer_tool_choice(tools, model_settings, model_request_parameters)
 
         system_prompt, anthropic_messages = await self._map_message(messages, model_request_parameters, model_settings)
 
         try:
-            extra_headers = model_settings.get('extra_headers', {})
-            extra_headers.setdefault('User-Agent', get_user_agent())
-            if beta_features:
-                if 'anthropic-beta' in extra_headers:
-                    beta_features.insert(0, extra_headers['anthropic-beta'])
-                extra_headers['anthropic-beta'] = ','.join(beta_features)
+            extra_headers = self._map_extra_headers(beta_features, model_settings)
 
             return await self.client.beta.messages.create(
                 max_tokens=model_settings.get('max_tokens', 4096),
@@ -356,6 +358,43 @@ async def _messages_create(
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
             raise  # pragma: lax no cover
 
+    async def _messages_count_tokens(
+        self,
+        messages: list[ModelMessage],
+        model_settings: AnthropicModelSettings,
+        model_request_parameters: ModelRequestParameters,
+    ) -> BetaMessageTokensCount:
+        if isinstance(self.client, AsyncAnthropicBedrock):
+            raise UserError('AsyncAnthropicBedrock client does not support `count_tokens` api.')
+
+        # standalone function to make it easier to override
+        tools = self._get_tools(model_request_parameters, model_settings)
+        tools, mcp_servers, beta_features = self._add_builtin_tools(tools, model_request_parameters)
+
+        tool_choice = self._infer_tool_choice(tools, model_settings, model_request_parameters)
+
+        system_prompt, anthropic_messages = await self._map_message(messages, model_request_parameters, model_settings)
+
+        try:
+            extra_headers = self._map_extra_headers(beta_features, model_settings)
+
+            return await self.client.beta.messages.count_tokens(
+                system=system_prompt or OMIT,
+                messages=anthropic_messages,
+                model=self._model_name,
+                tools=tools or OMIT,
+                tool_choice=tool_choice or OMIT,
+                mcp_servers=mcp_servers or OMIT,
+                thinking=model_settings.get('anthropic_thinking', OMIT),
+                timeout=model_settings.get('timeout', NOT_GIVEN),
+                extra_headers=extra_headers,
+                extra_body=model_settings.get('extra_body'),
+            )
+        except APIStatusError as e:
+            if (status_code := e.status_code) >= 400:
+                raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
+            raise  # pragma: lax no cover
+
     def _process_response(self, response: BetaMessage) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         items: list[ModelResponsePart] = []
@@ -492,6 +531,37 @@ def _add_builtin_tools(
                 )
         return tools, mcp_servers, beta_features
 
+    def _infer_tool_choice(
+        self,
+        tools: list[BetaToolUnionParam],
+        model_settings: AnthropicModelSettings,
+        model_request_parameters: ModelRequestParameters,
+    ) -> BetaToolChoiceParam | None:
+        if not tools:
+            return None
+        else:
+            tool_choice: BetaToolChoiceParam
+
+            if not model_request_parameters.allow_text_output:
+                tool_choice = {'type': 'any'}
+            else:
+                tool_choice = {'type': 'auto'}
+
+            if 'parallel_tool_calls' in model_settings:
+                tool_choice['disable_parallel_tool_use'] = not model_settings['parallel_tool_calls']
+
+            return tool_choice
+
+    def _map_extra_headers(self, beta_features: list[str], model_settings: AnthropicModelSettings) -> dict[str, str]:
+        """Apply beta_features to extra_headers in model_settings."""
+        extra_headers = model_settings.get('extra_headers', {})
+        extra_headers.setdefault('User-Agent', get_user_agent())
+        if beta_features:
+            if 'anthropic-beta' in extra_headers:
+                beta_features.insert(0, extra_headers['anthropic-beta'])
+            extra_headers['anthropic-beta'] = ','.join(beta_features)
+        return extra_headers
+
     async def _map_message(  # noqa: C901
         self,
         messages: list[ModelMessage],
diff --git a/tests/models/cassettes/test_anthropic/test_anthropic_count_tokens_error.yaml b/tests/models/cassettes/test_anthropic/test_anthropic_count_tokens_error.yaml
@@ -0,0 +1,46 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '105'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+    method: POST
+    parsed_body:
+      messages:
+      - content:
+        - text: hello
+          type: text
+        role: user
+      model: claude-does-not-exist
+    uri: https://api.anthropic.com/v1/messages/count_tokens?beta=true
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '136'
+      content-type:
+      - application/json
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      error:
+        message: 'model: claude-does-not-exist'
+        type: not_found_error
+      request_id: req_011CVEA3SF7rnb3DuBZytqQa
+      type: error
+    status:
+      code: 404
+      message: Not Found
+version: 1
diff --git a/tests/models/cassettes/test_anthropic/test_anthropic_model_usage_limit_exceeded.yaml b/tests/models/cassettes/test_anthropic/test_anthropic_model_usage_limit_exceeded.yaml
@@ -0,0 +1,40 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '139'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+    method: POST
+    parsed_body:
+      messages:
+      - content:
+        - text: The quick brown fox jumps over the lazydog.
+          type: text
+        role: user
+      model: claude-sonnet-4-5
+    uri: https://api.anthropic.com/v1/messages/count_tokens?beta=true
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '19'
+      content-type:
+      - application/json
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+    parsed_body:
+      input_tokens: 19
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/models/cassettes/test_anthropic/test_anthropic_model_usage_limit_not_exceeded.yaml b/tests/models/cassettes/test_anthropic/test_anthropic_model_usage_limit_not_exceeded.yaml
@@ -0,0 +1,108 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '139'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+    method: POST
+    parsed_body:
+      messages:
+      - content:
+        - text: The quick brown fox jumps over the lazydog.
+          type: text
+        role: user
+      model: claude-sonnet-4-5
+    uri: https://api.anthropic.com/v1/messages/count_tokens?beta=true
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '19'
+      content-type:
+      - application/json
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+    parsed_body:
+      input_tokens: 19
+    status:
+      code: 200
+      message: OK
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '172'
+      content-type:
+      - application/json
+      host:
+      - api.anthropic.com
+    method: POST
+    parsed_body:
+      max_tokens: 4096
+      messages:
+      - content:
+        - text: The quick brown fox jumps over the lazydog.
+          type: text
+        role: user
+      model: claude-sonnet-4-5
+      stream: false
+    uri: https://api.anthropic.com/v1/messages?beta=true
+  response:
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '729'
+      content-type:
+      - application/json
+      retry-after:
+      - '19'
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      transfer-encoding:
+      - chunked
+    parsed_body:
+      content:
+      - text: |-
+          I noticed a small typo in that famous pangram! It should be:
+
+          "The quick brown fox jumps over the **lazy dog**."
+
+          (There should be a space between "lazy" and "dog")
+
+          This sentence is often used for testing typewriters, fonts, and keyboards because it contains every letter of the English alphabet at least once.
+        type: text
+      id: msg_01QHpSAhCiB6L5pL23LjdRAy
+      model: claude-sonnet-4-5-20250929
+      role: assistant
+      stop_reason: end_turn
+      stop_sequence: null
+      type: message
+      usage:
+        cache_creation:
+          ephemeral_1h_input_tokens: 0
+          ephemeral_5m_input_tokens: 0
+        cache_creation_input_tokens: 0
+        cache_read_input_tokens: 0
+        input_tokens: 19
+        output_tokens: 77
+        service_tier: standard
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/models/test_anthropic.py b/tests/models/test_anthropic.py
