[anthropic] add count_tokens

Author: nathan-gage

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],