feat(gateway): support AWS Bedrock (#3203)

Co-authored-by: Douwe Maan <[email protected]>

Author: Kludex

pydantic_ai_slim/pydantic_ai/providers/bedrock.py

@@ -4,7 +4,7 @@
 import re
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Literal, overload
+from typing import Any, Literal, overload
 
 from pydantic_ai import ModelProfile
 from pydantic_ai.exceptions import UserError
@@ -21,6 +21,8 @@
     from botocore.client import BaseClient
     from botocore.config import Config
     from botocore.exceptions import NoRegionError
+    from botocore.session import Session
+    from botocore.tokens import FrozenAuthToken
 except ImportError as _import_error:
     raise ImportError(
         'Please install the `boto3` package to use the Bedrock provider, '
@@ -117,10 +119,23 @@ def __init__(self, *, bedrock_client: BaseClient) -> None: ...
     def __init__(
         self,
         *,
+        api_key: str,
+        base_url: str | None = None,
         region_name: str | None = None,
+        profile_name: str | None = None,
+        aws_read_timeout: float | None = None,
+        aws_connect_timeout: float | None = None,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        *,
         aws_access_key_id: str | None = None,
         aws_secret_access_key: str | None = None,
         aws_session_token: str | None = None,
+        base_url: str | None = None,
+        region_name: str | None = None,
         profile_name: str | None = None,
         aws_read_timeout: float | None = None,
         aws_connect_timeout: float | None = None,
@@ -130,42 +145,71 @@ def __init__(
         self,
         *,
         bedrock_client: BaseClient | None = None,
-        region_name: str | None = None,
         aws_access_key_id: str | None = None,
         aws_secret_access_key: str | None = None,
         aws_session_token: str | None = None,
+        base_url: str | None = None,
+        region_name: str | None = None,
         profile_name: str | None = None,
+        api_key: str | None = None,
         aws_read_timeout: float | None = None,
         aws_connect_timeout: float | None = None,
     ) -> None:
         """Initialize the Bedrock provider.
 
         Args:
             bedrock_client: A boto3 client for Bedrock Runtime. If provided, other arguments are ignored.
-            region_name: The AWS region name.
-            aws_access_key_id: The AWS access key ID.
-            aws_secret_access_key: The AWS secret access key.
-            aws_session_token: The AWS session token.
+            aws_access_key_id: The AWS access key ID. If not set, the `AWS_ACCESS_KEY_ID` environment variable will be used if available.
+            aws_secret_access_key: The AWS secret access key. If not set, the `AWS_SECRET_ACCESS_KEY` environment variable will be used if available.
+            aws_session_token: The AWS session token. If not set, the `AWS_SESSION_TOKEN` environment variable will be used if available.
+            api_key: The API key for Bedrock client. Can be used instead of `aws_access_key_id`, `aws_secret_access_key`, and `aws_session_token`. If not set, the `AWS_BEARER_TOKEN_BEDROCK` environment variable will be used if available.
+            base_url: The base URL for the Bedrock client.
+            region_name: The AWS region name. If not set, the `AWS_DEFAULT_REGION` environment variable will be used if available.
             profile_name: The AWS profile name.
             aws_read_timeout: The read timeout for Bedrock client.
             aws_connect_timeout: The connect timeout for Bedrock client.
         """
         if bedrock_client is not None:
             self._client = bedrock_client
         else:
+            read_timeout = aws_read_timeout or float(os.getenv('AWS_READ_TIMEOUT', 300))
+            connect_timeout = aws_connect_timeout or float(os.getenv('AWS_CONNECT_TIMEOUT', 60))
+            config: dict[str, Any] = {
+                'read_timeout': read_timeout,
+                'connect_timeout': connect_timeout,
+            }
             try:
-                read_timeout = aws_read_timeout or float(os.getenv('AWS_READ_TIMEOUT', 300))
-                connect_timeout = aws_connect_timeout or float(os.getenv('AWS_CONNECT_TIMEOUT', 60))
-                session = boto3.Session(
-                    aws_access_key_id=aws_access_key_id,
-                    aws_secret_access_key=aws_secret_access_key,
-                    aws_session_token=aws_session_token,
-                    region_name=region_name,
-                    profile_name=profile_name,
-                )
+                if api_key is not None:
+                    session = boto3.Session(
+                        botocore_session=_BearerTokenSession(api_key),
+                        region_name=region_name,
+                        profile_name=profile_name,
+                    )
+                    config['signature_version'] = 'bearer'
+                else:
+                    session = boto3.Session(
+                        aws_access_key_id=aws_access_key_id,
+                        aws_secret_access_key=aws_secret_access_key,
+                        aws_session_token=aws_session_token,
+                        region_name=region_name,
+                        profile_name=profile_name,
+                    )
                 self._client = session.client(  # type: ignore[reportUnknownMemberType]
                     'bedrock-runtime',
-                    config=Config(read_timeout=read_timeout, connect_timeout=connect_timeout),
+                    config=Config(**config),
+                    endpoint_url=base_url,
                 )
             except NoRegionError as exc:  # pragma: no cover
                 raise UserError('You must provide a `region_name` or a boto3 client for Bedrock Runtime.') from exc
+
+
+class _BearerTokenSession(Session):
+    def __init__(self, token: str):
+        super().__init__()
+        self.token = token
+
+    def get_auth_token(self, **_kwargs: Any) -> FrozenAuthToken:
+        return FrozenAuthToken(self.token)
+
+    def get_credentials(self) -> None:  # type: ignore[reportIncompatibleMethodOverride]
+        return None

pydantic_ai_slim/pydantic_ai/providers/gateway.py

@@ -11,6 +11,7 @@
 from pydantic_ai.models import Model, cached_async_http_client, get_user_agent
 
 if TYPE_CHECKING:
+    from botocore.client import BaseClient
     from google.genai import Client as GoogleClient
     from groq import AsyncGroq
     from openai import AsyncOpenAI
@@ -57,13 +58,25 @@ def gateway_provider(
 ) -> Provider[AsyncAnthropicClient]: ...
 
 
+@overload
 def gateway_provider(
-    upstream_provider: Literal['openai', 'openai-chat', 'openai-responses', 'groq', 'google-vertex', 'anthropic'] | str,
+    upstream_provider: Literal['bedrock'],
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+) -> Provider[BaseClient]: ...
+
+
+UpstreamProvider = Literal['openai', 'openai-chat', 'openai-responses', 'groq', 'google-vertex', 'anthropic', 'bedrock']
+
+
+def gateway_provider(
+    upstream_provider: UpstreamProvider | str,
     *,
     # Every provider
     api_key: str | None = None,
     base_url: str | None = None,
-    # OpenAI & Groq
+    # OpenAI, Groq & Anthropic
     http_client: httpx.AsyncClient | None = None,
 ) -> Provider[Any]:
     """Create a new Gateway provider.
@@ -73,7 +86,7 @@ def gateway_provider(
         api_key: The API key to use for authentication. If not provided, the `PYDANTIC_AI_GATEWAY_API_KEY`
             environment variable will be used if available.
         base_url: The base URL to use for the Gateway. If not provided, the `PYDANTIC_AI_GATEWAY_BASE_URL`
-            environment variable will be used if available. Otherwise, defaults to `http://localhost:8787/`.
+            environment variable will be used if available. Otherwise, defaults to `https://gateway.pydantic.dev/proxy`.
         http_client: The HTTP client to use for the Gateway.
     """
     api_key = api_key or os.getenv('PYDANTIC_AI_GATEWAY_API_KEY')
@@ -111,6 +124,14 @@ def gateway_provider(
                 http_client=http_client,
             )
         )
+    elif upstream_provider == 'bedrock':
+        from .bedrock import BedrockProvider
+
+        return BedrockProvider(
+            api_key=api_key,
+            base_url=_merge_url_path(base_url, 'bedrock'),
+            region_name='pydantic-ai-gateway',  # Fake region name to avoid NoRegionError
+        )
     elif upstream_provider == 'google-vertex':
         from google.genai import Client as GoogleClient
 

tests/providers/cassettes/test_gateway/test_gateway_provider_with_anthropic.yaml

@@ -40,7 +40,7 @@ interactions:
       content:
       - text: The capital of France is Paris.
         type: text
-      id: msg_015tco2dv5oh9rFq1PcZAduv
+      id: msg_0116L5r52AZ42YhvvdUuHEsk
       model: claude-3-5-sonnet-20241022
       role: assistant
       stop_reason: end_turn

tests/providers/cassettes/test_gateway/test_gateway_provider_with_bedrock.yaml

@@ -0,0 +1,50 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": [{"text": "What is the capital of France?"}]}], "system": [], "inferenceConfig":
+      {}}'
+    headers:
+      amz-sdk-invocation-id:
+      - !!binary |
+        MWYwNDlkMTQtMjVmMC00YTRhLWJhYmMtNTQ0MDdhMmRlNjgw
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+      content-length:
+      - '126'
+      content-type:
+      - !!binary |
+        YXBwbGljYXRpb24vanNvbg==
+    method: POST
+    uri: http://localhost:8787/bedrock/model/amazon.nova-micro-v1%3A0/converse
+  response:
+    headers:
+      content-length:
+      - '741'
+      content-type:
+      - application/json
+      pydantic-ai-gateway-price-estimate:
+      - 0.0000USD
+    parsed_body:
+      metrics:
+        latencyMs: 668
+      output:
+        message:
+          content:
+          - text: The capital of France is Paris. Paris is not only the capital city but also the most populous city in France,
+              and it is a major center for culture, commerce, fashion, and international diplomacy. The city is known for
+              its historical and architectural landmarks, including the Eiffel Tower, the Louvre Museum, Notre-Dame Cathedral,
+              and the Champs-Élysées. Paris plays a significant role in the global arts, fashion, research, technology, education,
+              and entertainment scenes.
+          role: assistant
+      stopReason: end_turn
+      usage:
+        inputTokens: 7
+        outputTokens: 96
+        pydantic_ai_gateway:
+          cost_estimate: 1.3685000000000002e-05
+        serverToolUsage: {}
+        totalTokens: 103
+    status:
+      code: 200
+      message: OK
+version: 1

tests/providers/test_gateway.py

@@ -14,6 +14,7 @@
 
 with try_import() as imports_successful:
     from pydantic_ai.models.anthropic import AnthropicModel
+    from pydantic_ai.models.bedrock import BedrockConverseModel
     from pydantic_ai.models.google import GoogleModel
     from pydantic_ai.models.groq import GroqModel
     from pydantic_ai.models.openai import OpenAIChatModel, OpenAIResponsesModel
@@ -150,3 +151,14 @@ async def test_gateway_provider_with_anthropic(allow_model_requests: None, gatew
 
     result = await agent.run('What is the capital of France?')
     assert result.output == snapshot('The capital of France is Paris.')
+
+
+async def test_gateway_provider_with_bedrock(allow_model_requests: None, gateway_api_key: str):
+    provider = gateway_provider('bedrock', api_key=gateway_api_key, base_url='http://localhost:8787')
+    model = BedrockConverseModel('amazon.nova-micro-v1:0', provider=provider)
+    agent = Agent(model)
+
+    result = await agent.run('What is the capital of France?')
+    assert result.output == snapshot(
+        'The capital of France is Paris. Paris is not only the capital city but also the most populous city in France, and it is a major center for culture, commerce, fashion, and international diplomacy. The city is known for its historical and architectural landmarks, including the Eiffel Tower, the Louvre Museum, Notre-Dame Cathedral, and the Champs-Élysées. Paris plays a significant role in the global arts, fashion, research, technology, education, and entertainment scenes.'
+    )