Merge branch 'main' into gateway-string-docs

Author: DouweM

.github/workflows/ci.yml

@@ -202,7 +202,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        # TODO(Marcelo): Enable 3.11 again.
+        python-version: ["3.10", "3.12", "3.13"]
     env:
       CI: true
       COVERAGE_PROCESS_START: ./pyproject.toml

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -103,6 +103,13 @@
         'bedrock:us.anthropic.claude-opus-4-20250514-v1:0',
         'bedrock:anthropic.claude-sonnet-4-20250514-v1:0',
         'bedrock:us.anthropic.claude-sonnet-4-20250514-v1:0',
+        'bedrock:eu.anthropic.claude-sonnet-4-20250514-v1:0',
+        'bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0',
+        'bedrock:us.anthropic.claude-sonnet-4-5-20250929-v1:0',
+        'bedrock:eu.anthropic.claude-sonnet-4-5-20250929-v1:0',
+        'bedrock:anthropic.claude-haiku-4-5-20251001-v1:0',
+        'bedrock:us.anthropic.claude-haiku-4-5-20251001-v1:0',
+        'bedrock:eu.anthropic.claude-haiku-4-5-20251001-v1:0',
         'bedrock:cohere.command-text-v14',
         'bedrock:cohere.command-r-v1:0',
         'bedrock:cohere.command-r-plus-v1:0',
@@ -773,9 +780,10 @@ def infer_model(  # noqa: C901
 
     model_kind = provider_name
     if model_kind.startswith('gateway/'):
-        from ..providers.gateway import infer_gateway_model
+        from ..providers.gateway import normalize_gateway_provider
 
-        return infer_gateway_model(model_kind.removeprefix('gateway/'), model_name=model_name)
+        model_kind = provider_name.removeprefix('gateway/')
+        model_kind = normalize_gateway_provider(model_kind)
     if model_kind in (
         'openai',
         'azure',

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -22,7 +22,6 @@
     DocumentUrl,
     FinishReason,
     ImageUrl,
-    ModelHTTPError,
     ModelMessage,
     ModelProfileSpec,
     ModelRequest,
@@ -41,7 +40,7 @@
     usage,
 )
 from pydantic_ai._run_context import RunContext
-from pydantic_ai.exceptions import UserError
+from pydantic_ai.exceptions import ModelHTTPError, UserError
 from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse, download_item
 from pydantic_ai.providers import Provider, infer_provider
 from pydantic_ai.providers.bedrock import BedrockModelProfile
@@ -61,6 +60,7 @@
         ConverseStreamMetadataEventTypeDef,
         ConverseStreamOutputTypeDef,
         ConverseStreamResponseTypeDef,
+        CountTokensRequestTypeDef,
         DocumentBlockTypeDef,
         GuardrailConfigurationTypeDef,
         ImageBlockTypeDef,
@@ -77,7 +77,6 @@
         VideoBlockTypeDef,
     )
 
-
 LatestBedrockModelNames = Literal[
     'amazon.titan-tg1-large',
     'amazon.titan-text-lite-v1',
@@ -106,6 +105,13 @@
     'us.anthropic.claude-opus-4-20250514-v1:0',
     'anthropic.claude-sonnet-4-20250514-v1:0',
     'us.anthropic.claude-sonnet-4-20250514-v1:0',
+    'eu.anthropic.claude-sonnet-4-20250514-v1:0',
+    'anthropic.claude-sonnet-4-5-20250929-v1:0',
+    'us.anthropic.claude-sonnet-4-5-20250929-v1:0',
+    'eu.anthropic.claude-sonnet-4-5-20250929-v1:0',
+    'anthropic.claude-haiku-4-5-20251001-v1:0',
+    'us.anthropic.claude-haiku-4-5-20251001-v1:0',
+    'eu.anthropic.claude-haiku-4-5-20251001-v1:0',
     'cohere.command-text-v14',
     'cohere.command-r-v1:0',
     'cohere.command-r-plus-v1:0',
@@ -136,7 +142,6 @@
 See [the Bedrock docs](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) for a full list.
 """
 
-
 P = ParamSpec('P')
 T = typing.TypeVar('T')
 
@@ -149,6 +154,13 @@
     'tool_use': 'tool_call',
 }
 
+_AWS_BEDROCK_INFERENCE_GEO_PREFIXES: tuple[str, ...] = ('us.', 'eu.', 'apac.', 'jp.', 'au.', 'ca.')
+"""Geo prefixes for Bedrock inference profile IDs (e.g., 'eu.', 'us.').
+
+Used to strip the geo prefix so we can pass a pure foundation model ID/ARN to CountTokens,
+which does not accept profile IDs. Extend if new geos appear (e.g., 'global.', 'us-gov.').
+"""
+
 
 class BedrockModelSettings(ModelSettings, total=False):
     """Settings for Bedrock models.
@@ -228,7 +240,7 @@ def __init__(
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider('gateway/converse' if provider == 'gateway' else provider)
+            provider = infer_provider('gateway/bedrock' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = cast('BedrockRuntimeClient', provider.client)
 
@@ -275,6 +287,34 @@ async def request(
         model_response = await 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:
+        """Count the number of tokens, works with limited models.
+
+        Check the actual supported models on <https://docs.aws.amazon.com/bedrock/latest/userguide/count-tokens.html>
+        """
+        model_settings, model_request_parameters = self.prepare_request(model_settings, model_request_parameters)
+        system_prompt, bedrock_messages = await self._map_messages(messages, model_request_parameters)
+        params: CountTokensRequestTypeDef = {
+            'modelId': self._remove_inference_geo_prefix(self.model_name),
+            'input': {
+                'converse': {
+                    'messages': bedrock_messages,
+                    'system': system_prompt,
+                },
+            },
+        }
+        try:
+            response = await anyio.to_thread.run_sync(functools.partial(self.client.count_tokens, **params))
+        except ClientError as e:
+            status_code = e.response.get('ResponseMetadata', {}).get('HTTPStatusCode', 500)
+            raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.response) from e
+        return usage.RequestUsage(input_tokens=response['inputTokens'])
+
     @asynccontextmanager
     async def request_stream(
         self,
@@ -642,6 +682,14 @@ def _map_tool_call(t: ToolCallPart) -> ContentBlockOutputTypeDef:
             'toolUse': {'toolUseId': _utils.guard_tool_call_id(t=t), 'name': t.tool_name, 'input': t.args_as_dict()}
         }
 
+    @staticmethod
+    def _remove_inference_geo_prefix(model_name: BedrockModelName) -> BedrockModelName:
+        """Remove inference geographic prefix from model ID if present."""
+        for prefix in _AWS_BEDROCK_INFERENCE_GEO_PREFIXES:
+            if model_name.startswith(prefix):
+                return model_name.removeprefix(prefix)
+        return model_name
+
 
 @dataclass
 class BedrockStreamedResponse(StreamedResponse):

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -204,7 +204,7 @@ def __init__(
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider('gateway/gemini' if provider == 'gateway' else provider)
+            provider = infer_provider('gateway/google-vertex' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = provider.client
 

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -375,7 +375,7 @@ def __init__(
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider('gateway/chat' if provider == 'gateway' else provider)
+            provider = infer_provider('gateway/openai' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = provider.client
 
@@ -944,7 +944,7 @@ def __init__(
         self._model_name = model_name
 
         if isinstance(provider, str):
-            provider = infer_provider('gateway/responses' if provider == 'gateway' else provider)
+            provider = infer_provider('gateway/openai' if provider == 'gateway' else provider)
         self._provider = provider
         self.client = provider.client
 

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -158,8 +158,8 @@ def infer_provider(provider: str) -> Provider[Any]:
     if provider.startswith('gateway/'):
         from .gateway import gateway_provider
 
-        api_type = provider.removeprefix('gateway/')
-        return gateway_provider(api_type)
+        upstream_provider = provider.removeprefix('gateway/')
+        return gateway_provider(upstream_provider)
     elif provider in ('google-vertex', 'google-gla'):
         from .google import GoogleProvider