Add support for AWS Bedrock Converse API (#994)

Co-authored-by: Wh1isper <[email protected]>

Author: Kludex

docs/api/models/bedrock.md

@@ -0,0 +1,7 @@
+# `pydantic_ai.models.bedrock`
+
+## Setup
+
+For details on how to set up authentication with this model, see [model configuration for Bedrock](../../models.md#bedrock).
+
+::: pydantic_ai.models.bedrock

docs/api/providers.md

@@ -7,3 +7,5 @@
 ::: pydantic_ai.providers.openai
 
 ::: pydantic_ai.providers.deepseek
+
+::: pydantic_ai.providers.bedrock

docs/models.md

@@ -7,6 +7,7 @@ PydanticAI is Model-agnostic and has built in support for the following model pr
 * [Groq](#groq)
 * [Mistral](#mistral)
 * [Cohere](#cohere)
+* [Bedrock](#bedrock)
 
 See [OpenAI-compatible models](#openai-compatible-models) for more examples on how to use models such as [OpenRouter](#openrouter), and [Grok (xAI)](#grok-xai) that support the OpenAI SDK.
 
@@ -540,6 +541,93 @@ agent = Agent(model)
 ...
 ```
 
+## Bedrock
+
+### Install
+
+To use [`BedrockConverseModel`][pydantic_ai.models.bedrock.BedrockConverseModel], you need to either install [`pydantic-ai`](install.md), or install [`pydantic-ai-slim`](install.md#slim-install) with the `bedrock` optional group:
+
+```bash
+pip/uv-add 'pydantic-ai-slim[bedrock]'
+```
+
+### Configuration
+
+To use [AWS Bedrock](https://aws.amazon.com/bedrock/), you'll need an AWS account with Bedrock enabled and appropriate credentials. You can use either AWS credentials directly or a pre-configured boto3 client.
+
+[`BedrockModelName`][pydantic_ai.models.bedrock.BedrockModelName] contains a list of available Bedrock models, including models from Anthropic, Amazon, Cohere, Meta, and Mistral.
+
+### Environment variables
+
+You can set your AWS credentials as environment variables:
+
+```bash
+export AWS_ACCESS_KEY_ID='your-access-key'
+export AWS_SECRET_ACCESS_KEY='your-secret-key'
+export AWS_REGION='us-east-1'  # or your preferred region
+```
+
+You can then use [`BedrockConverseModel`][pydantic_ai.models.bedrock.BedrockConverseModel] by name:
+
+```python {title="bedrock_model_by_name.py", test="skip"}
+from pydantic_ai import Agent
+
+agent = Agent('bedrock:anthropic.claude-3-sonnet-20240229-v1:0')
+...
+```
+
+Or initialize the model directly with just the model name:
+
+```python {title="bedrock_model_init.py" test="skip"}
+from pydantic_ai import Agent
+from pydantic_ai.models.bedrock import BedrockConverseModel
+
+model = BedrockConverseModel('anthropic.claude-3-sonnet-20240229-v1:0')
+agent = Agent(model)
+...
+```
+
+### `provider` argument
+
+You can provide a custom [`BedrockProvider`][pydantic_ai.providers.bedrock.BedrockProvider] via the [`provider` argument][pydantic_ai.models.bedrock.BedrockConverseModel.__init__]. This is useful when you want to specify credentials directly or use a custom boto3 client:
+
+```python {title="bedrock_model_provider.py"}
+from pydantic_ai import Agent
+from pydantic_ai.models.bedrock import BedrockConverseModel
+from pydantic_ai.providers.bedrock import BedrockProvider
+
+# Using AWS credentials directly
+model = BedrockConverseModel(
+    'anthropic.claude-3-sonnet-20240229-v1:0',
+    provider=BedrockProvider(
+        region_name='us-east-1',
+        aws_access_key_id='your-access-key',
+        aws_secret_access_key='your-secret-key',
+    ),
+)
+agent = Agent(model)
+...
+```
+
+You can also pass a pre-configured boto3 client:
+
+```python {title="bedrock_model_boto3.py"}
+import boto3
+
+from pydantic_ai import Agent
+from pydantic_ai.models.bedrock import BedrockConverseModel
+from pydantic_ai.providers.bedrock import BedrockProvider
+
+# Using a pre-configured boto3 client
+bedrock_client = boto3.client('bedrock-runtime', region_name='us-east-1')
+model = BedrockConverseModel(
+    'anthropic.claude-3-sonnet-20240229-v1:0',
+    provider=BedrockProvider(bedrock_client=bedrock_client),
+)
+agent = Agent(model)
+...
+```
+
 ## OpenAI-compatible Models
 
 Many of the models are compatible with OpenAI API, and thus can be used with [`OpenAIModel`][pydantic_ai.models.openai.OpenAIModel] in PydanticAI.

mkdocs.yml

@@ -54,6 +54,7 @@ nav:
       - api/models/base.md
       - api/models/openai.md
       - api/models/anthropic.md
+      - api/models/bedrock.md
       - api/models/cohere.md
       - api/models/gemini.md
       - api/models/vertexai.md

pydantic_ai_slim/pydantic_ai/_pydantic.py

@@ -6,7 +6,7 @@
 from __future__ import annotations as _annotations
 
 from inspect import Parameter, signature
-from typing import TYPE_CHECKING, Any, Callable, TypedDict, cast
+from typing import TYPE_CHECKING, Any, Callable, cast
 
 from pydantic import ConfigDict
 from pydantic._internal import _decorators, _generate_schema, _typing_extra
@@ -15,7 +15,7 @@
 from pydantic.json_schema import GenerateJsonSchema
 from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
-from typing_extensions import get_origin
+from typing_extensions import TypedDict, get_origin
 
 from ._griffe import doc_descriptions
 from ._utils import check_object_json_schema, is_model_like

pydantic_ai_slim/pydantic_ai/common_tools/duckduckgo.py

@@ -1,10 +1,10 @@
 import functools
 from dataclasses import dataclass
-from typing import TypedDict
 
 import anyio
 import anyio.to_thread
 from pydantic import TypeAdapter
+from typing_extensions import TypedDict
 
 from pydantic_ai.tools import Tool
 

pydantic_ai_slim/pydantic_ai/common_tools/tavily.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass
-from typing import Literal, TypedDict
+from typing import Literal
 
 from pydantic import TypeAdapter
+from typing_extensions import TypedDict
 
 from pydantic_ai.tools import Tool
 

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -34,6 +34,49 @@
     'anthropic:claude-3-opus-latest',
     'claude-3-7-sonnet-latest',
     'claude-3-5-haiku-latest',
+    'bedrock:amazon.titan-tg1-large',
+    'bedrock:amazon.titan-text-lite-v1',
+    'bedrock:amazon.titan-text-express-v1',
+    'bedrock:us.amazon.nova-pro-v1:0',
+    'bedrock:us.amazon.nova-lite-v1:0',
+    'bedrock:us.amazon.nova-micro-v1:0',
+    'bedrock:anthropic.claude-3-5-sonnet-20241022-v2:0',
+    'bedrock:us.anthropic.claude-3-5-sonnet-20241022-v2:0',
+    'bedrock:anthropic.claude-3-5-haiku-20241022-v1:0',
+    'bedrock:us.anthropic.claude-3-5-haiku-20241022-v1:0',
+    'bedrock:anthropic.claude-instant-v1',
+    'bedrock:anthropic.claude-v2:1',
+    'bedrock:anthropic.claude-v2',
+    'bedrock:anthropic.claude-3-sonnet-20240229-v1:0',
+    'bedrock:us.anthropic.claude-3-sonnet-20240229-v1:0',
+    'bedrock:anthropic.claude-3-haiku-20240307-v1:0',
+    'bedrock:us.anthropic.claude-3-haiku-20240307-v1:0',
+    'bedrock:anthropic.claude-3-opus-20240229-v1:0',
+    'bedrock:us.anthropic.claude-3-opus-20240229-v1:0',
+    'bedrock:anthropic.claude-3-5-sonnet-20240620-v1:0',
+    'bedrock:us.anthropic.claude-3-5-sonnet-20240620-v1:0',
+    'bedrock:anthropic.claude-3-7-sonnet-20250219-v1:0',
+    'bedrock:us.anthropic.claude-3-7-sonnet-20250219-v1:0',
+    'bedrock:cohere.command-text-v14',
+    'bedrock:cohere.command-r-v1:0',
+    'bedrock:cohere.command-r-plus-v1:0',
+    'bedrock:cohere.command-light-text-v14',
+    'bedrock:meta.llama3-8b-instruct-v1:0',
+    'bedrock:meta.llama3-70b-instruct-v1:0',
+    'bedrock:meta.llama3-1-8b-instruct-v1:0',
+    'bedrock:us.meta.llama3-1-8b-instruct-v1:0',
+    'bedrock:meta.llama3-1-70b-instruct-v1:0',
+    'bedrock:us.meta.llama3-1-70b-instruct-v1:0',
+    'bedrock:meta.llama3-1-405b-instruct-v1:0',
+    'bedrock:us.meta.llama3-2-11b-instruct-v1:0',
+    'bedrock:us.meta.llama3-2-90b-instruct-v1:0',
+    'bedrock:us.meta.llama3-2-1b-instruct-v1:0',
+    'bedrock:us.meta.llama3-2-3b-instruct-v1:0',
+    'bedrock:us.meta.llama3-3-70b-instruct-v1:0',
+    'bedrock:mistral.mistral-7b-instruct-v0:2',
+    'bedrock:mistral.mixtral-8x7b-instruct-v0:1',
+    'bedrock:mistral.mistral-large-2402-v1:0',
+    'bedrock:mistral.mistral-large-2407-v1:0',
     'claude-3-5-sonnet-latest',
     'claude-3-opus-latest',
     'cohere:c4ai-aya-expanse-32b',
@@ -324,7 +367,7 @@ def infer_model(model: Model | KnownModelName) -> Model:
         return TestModel()
 
     try:
-        provider, model_name = model.split(':')
+        provider, model_name = model.split(':', maxsplit=1)
     except ValueError:
         model_name = model
         # TODO(Marcelo): We should deprecate this way.
@@ -368,6 +411,10 @@ def infer_model(model: Model | KnownModelName) -> Model:
 
         # TODO(Marcelo): Missing provider API.
         return AnthropicModel(model_name)
+    elif provider == 'bedrock':
+        from .bedrock import BedrockConverseModel
+
+        return BedrockConverseModel(model_name)
     else:
         raise UserError(f'Unknown model: {model}')