pydantic-ai-slim 0.0.34__tar.gz → 0.0.36__tar.gz

{pydantic_ai_slim-0.0.34 → pydantic_ai_slim-0.0.36}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.34
+Version: 0.0.36
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -29,11 +29,13 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.0.34
+Requires-Dist: pydantic-graph==0.0.36
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.49.0; extra == 'anthropic'
+Provides-Extra: bedrock
+Requires-Dist: boto3>=1.34.116; extra == 'bedrock'
 Provides-Extra: cli
 Requires-Dist: argcomplete>=3.5.0; extra == 'cli'
 Requires-Dist: prompt-toolkit>=3; extra == 'cli'

{pydantic_ai_slim-0.0.34 → pydantic_ai_slim-0.0.36}/pydantic_ai/_pydantic.py

@@ -6,7 +6,7 @@ This module has to use numerous internal Pydantic APIs and is therefore brittle
 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.fields import FieldInfo
 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-0.0.34 → pydantic_ai_slim-0.0.36}/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-0.0.34 → pydantic_ai_slim-0.0.36}/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-0.0.34 → pydantic_ai_slim-0.0.36}/pydantic_ai/models/__init__.py

@@ -34,6 +34,49 @@ KnownModelName = Literal[
     '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}')
 

pydantic_ai_slim-0.0.36/pydantic_ai/models/bedrock.py

@@ -0,0 +1,451 @@
+from __future__ import annotations
+
+import functools
+import typing
+from collections.abc import AsyncIterator, Iterable
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import TYPE_CHECKING, Generic, Literal, Union, cast, overload
+
+import anyio
+import anyio.to_thread
+from typing_extensions import ParamSpec, assert_never
+
+from pydantic_ai import _utils, result
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    ModelResponseStreamEvent,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
+from pydantic_ai.providers import Provider, infer_provider
+from pydantic_ai.settings import ModelSettings
+from pydantic_ai.tools import ToolDefinition
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+    from botocore.eventstream import EventStream
+    from mypy_boto3_bedrock_runtime import BedrockRuntimeClient
+    from mypy_boto3_bedrock_runtime.type_defs import (
+        ContentBlockOutputTypeDef,
+        ConverseResponseTypeDef,
+        ConverseStreamMetadataEventTypeDef,
+        ConverseStreamOutputTypeDef,
+        InferenceConfigurationTypeDef,
+        MessageUnionTypeDef,
+        ToolChoiceTypeDef,
+        ToolTypeDef,
+    )
+
+
+LatestBedrockModelNames = Literal[
+    'amazon.titan-tg1-large',
+    'amazon.titan-text-lite-v1',
+    'amazon.titan-text-express-v1',
+    'us.amazon.nova-pro-v1:0',
+    'us.amazon.nova-lite-v1:0',
+    'us.amazon.nova-micro-v1:0',
+    'anthropic.claude-3-5-sonnet-20241022-v2:0',
+    'us.anthropic.claude-3-5-sonnet-20241022-v2:0',
+    'anthropic.claude-3-5-haiku-20241022-v1:0',
+    'us.anthropic.claude-3-5-haiku-20241022-v1:0',
+    'anthropic.claude-instant-v1',
+    'anthropic.claude-v2:1',
+    'anthropic.claude-v2',
+    'anthropic.claude-3-sonnet-20240229-v1:0',
+    'us.anthropic.claude-3-sonnet-20240229-v1:0',
+    'anthropic.claude-3-haiku-20240307-v1:0',
+    'us.anthropic.claude-3-haiku-20240307-v1:0',
+    'anthropic.claude-3-opus-20240229-v1:0',
+    'us.anthropic.claude-3-opus-20240229-v1:0',
+    'anthropic.claude-3-5-sonnet-20240620-v1:0',
+    'us.anthropic.claude-3-5-sonnet-20240620-v1:0',
+    'anthropic.claude-3-7-sonnet-20250219-v1:0',
+    'us.anthropic.claude-3-7-sonnet-20250219-v1:0',
+    'cohere.command-text-v14',
+    'cohere.command-r-v1:0',
+    'cohere.command-r-plus-v1:0',
+    'cohere.command-light-text-v14',
+    'meta.llama3-8b-instruct-v1:0',
+    'meta.llama3-70b-instruct-v1:0',
+    'meta.llama3-1-8b-instruct-v1:0',
+    'us.meta.llama3-1-8b-instruct-v1:0',
+    'meta.llama3-1-70b-instruct-v1:0',
+    'us.meta.llama3-1-70b-instruct-v1:0',
+    'meta.llama3-1-405b-instruct-v1:0',
+    'us.meta.llama3-2-11b-instruct-v1:0',
+    'us.meta.llama3-2-90b-instruct-v1:0',
+    'us.meta.llama3-2-1b-instruct-v1:0',
+    'us.meta.llama3-2-3b-instruct-v1:0',
+    'us.meta.llama3-3-70b-instruct-v1:0',
+    'mistral.mistral-7b-instruct-v0:2',
+    'mistral.mixtral-8x7b-instruct-v0:1',
+    'mistral.mistral-large-2402-v1:0',
+    'mistral.mistral-large-2407-v1:0',
+]
+"""Latest Bedrock models."""
+
+BedrockModelName = Union[str, LatestBedrockModelNames]
+"""Possible Bedrock model names.
+
+Since Bedrock supports a variety of date-stamped models, we explicitly list the latest models but allow any name in the type hints.
+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')
+
+
+@dataclass(init=False)
+class BedrockConverseModel(Model):
+    """A model that uses the Bedrock Converse API."""
+
+    client: BedrockRuntimeClient
+
+    _model_name: BedrockModelName = field(repr=False)
+    _system: str | None = field(default='bedrock', repr=False)
+
+    @property
+    def model_name(self) -> str:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider, ex: openai."""
+        return self._system
+
+    def __init__(
+        self,
+        model_name: BedrockModelName,
+        *,
+        provider: Literal['bedrock'] | Provider[BaseClient] = 'bedrock',
+    ):
+        """Initialize a Bedrock model.
+
+        Args:
+            model_name: The name of the model to use.
+            model_name: The name of the Bedrock model to use. List of model names available
+                [here](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html).
+            provider: The provider to use. Defaults to `'bedrock'`.
+        """
+        self._model_name = model_name
+
+        if isinstance(provider, str):
+            self.client = infer_provider(provider).client
+        else:
+            self.client = cast('BedrockRuntimeClient', provider.client)
+
+    def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolTypeDef]:
+        tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
+        if model_request_parameters.result_tools:
+            tools += [self._map_tool_definition(r) for r in model_request_parameters.result_tools]
+        return tools
+
+    @staticmethod
+    def _map_tool_definition(f: ToolDefinition) -> ToolTypeDef:
+        return {
+            'toolSpec': {
+                'name': f.name,
+                'description': f.description,
+                'inputSchema': {'json': f.parameters_json_schema},
+            }
+        }
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> tuple[ModelResponse, result.Usage]:
+        response = await self._messages_create(messages, False, model_settings, model_request_parameters)
+        return await self._process_response(response)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        response = await self._messages_create(messages, True, model_settings, model_request_parameters)
+        yield BedrockStreamedResponse(_model_name=self.model_name, _event_stream=response)
+
+    async def _process_response(self, response: ConverseResponseTypeDef) -> tuple[ModelResponse, result.Usage]:
+        items: list[ModelResponsePart] = []
+        if message := response['output'].get('message'):
+            for item in message['content']:
+                if text := item.get('text'):
+                    items.append(TextPart(content=text))
+                else:
+                    tool_use = item.get('toolUse')
+                    assert tool_use is not None, f'Found a content that is not a text or tool use: {item}'
+                    items.append(
+                        ToolCallPart(
+                            tool_name=tool_use['name'],
+                            args=tool_use['input'],
+                            tool_call_id=tool_use['toolUseId'],
+                        ),
+                    )
+        usage = result.Usage(
+            request_tokens=response['usage']['inputTokens'],
+            response_tokens=response['usage']['outputTokens'],
+            total_tokens=response['usage']['totalTokens'],
+        )
+        return ModelResponse(items, model_name=self.model_name), usage
+
+    @overload
+    async def _messages_create(
+        self,
+        messages: list[ModelMessage],
+        stream: Literal[True],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> EventStream[ConverseStreamOutputTypeDef]:
+        pass
+
+    @overload
+    async def _messages_create(
+        self,
+        messages: list[ModelMessage],
+        stream: Literal[False],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ConverseResponseTypeDef:
+        pass
+
+    async def _messages_create(
+        self,
+        messages: list[ModelMessage],
+        stream: bool,
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ConverseResponseTypeDef | EventStream[ConverseStreamOutputTypeDef]:
+        tools = self._get_tools(model_request_parameters)
+        support_tools_choice = self.model_name.startswith(('anthropic', 'us.anthropic'))
+        if not tools or not support_tools_choice:
+            tool_choice: ToolChoiceTypeDef = {}
+        elif not model_request_parameters.allow_text_result:
+            tool_choice = {'any': {}}
+        else:
+            tool_choice = {'auto': {}}
+
+        system_prompt, bedrock_messages = self._map_message(messages)
+        inference_config = self._map_inference_config(model_settings)
+
+        params = {
+            'modelId': self.model_name,
+            'messages': bedrock_messages,
+            'system': [{'text': system_prompt}],
+            'inferenceConfig': inference_config,
+            **(
+                {'toolConfig': {'tools': tools, **({'toolChoice': tool_choice} if tool_choice else {})}}
+                if tools
+                else {}
+            ),
+        }
+
+        if stream:
+            model_response = await anyio.to_thread.run_sync(functools.partial(self.client.converse_stream, **params))
+            model_response = model_response['stream']
+        else:
+            model_response = await anyio.to_thread.run_sync(functools.partial(self.client.converse, **params))
+        return model_response
+
+    @staticmethod
+    def _map_inference_config(
+        model_settings: ModelSettings | None,
+    ) -> InferenceConfigurationTypeDef:
+        model_settings = model_settings or {}
+        inference_config: InferenceConfigurationTypeDef = {}
+
+        if max_tokens := model_settings.get('max_tokens'):
+            inference_config['maxTokens'] = max_tokens
+        if temperature := model_settings.get('temperature'):
+            inference_config['temperature'] = temperature
+        if top_p := model_settings.get('top_p'):
+            inference_config['topP'] = top_p
+        # TODO(Marcelo): This is not included in model_settings yet.
+        # if stop_sequences := model_settings.get('stop_sequences'):
+        #     inference_config['stopSequences'] = stop_sequences
+
+        return inference_config
+
+    def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[MessageUnionTypeDef]]:
+        """Just maps a `pydantic_ai.Message` to the Bedrock `MessageUnionTypeDef`."""
+        system_prompt: str = ''
+        bedrock_messages: list[MessageUnionTypeDef] = []
+        for m in messages:
+            if isinstance(m, ModelRequest):
+                for part in m.parts:
+                    if isinstance(part, SystemPromptPart):
+                        system_prompt += part.content
+                    elif isinstance(part, UserPromptPart):
+                        if isinstance(part.content, str):
+                            bedrock_messages.append({'role': 'user', 'content': [{'text': part.content}]})
+                        else:
+                            raise NotImplementedError('User prompt can only be a string for now.')
+                    elif isinstance(part, ToolReturnPart):
+                        assert part.tool_call_id is not None
+                        bedrock_messages.append(
+                            {
+                                'role': 'user',
+                                'content': [
+                                    {
+                                        'toolResult': {
+                                            'toolUseId': part.tool_call_id,
+                                            'content': [{'text': part.model_response_str()}],
+                                            'status': 'success',
+                                        }
+                                    }
+                                ],
+                            }
+                        )
+                    elif isinstance(part, RetryPromptPart):
+                        # TODO(Marcelo): We need to add a test here.
+                        if part.tool_name is None:  # pragma: no cover
+                            bedrock_messages.append({'role': 'user', 'content': [{'text': part.model_response()}]})
+                        else:
+                            assert part.tool_call_id is not None
+                            bedrock_messages.append(
+                                {
+                                    'role': 'user',
+                                    'content': [
+                                        {
+                                            'toolResult': {
+                                                'toolUseId': part.tool_call_id,
+                                                'content': [{'text': part.model_response()}],
+                                                'status': 'error',
+                                            }
+                                        }
+                                    ],
+                                }
+                            )
+            elif isinstance(m, ModelResponse):
+                content: list[ContentBlockOutputTypeDef] = []
+                for item in m.parts:
+                    if isinstance(item, TextPart):
+                        content.append({'text': item.content})
+                    else:
+                        assert isinstance(item, ToolCallPart)
+                        content.append(self._map_tool_call(item))  # FIXME: MISSING key
+                bedrock_messages.append({'role': 'assistant', 'content': content})
+            else:
+                assert_never(m)
+        return system_prompt, bedrock_messages
+
+    @staticmethod
+    def _map_tool_call(t: ToolCallPart) -> ContentBlockOutputTypeDef:
+        assert t.tool_call_id is not None
+        return {
+            'toolUse': {
+                'toolUseId': t.tool_call_id,
+                'name': t.tool_name,
+                'input': t.args_as_dict(),
+            }
+        }
+
+
+@dataclass
+class BedrockStreamedResponse(StreamedResponse):
+    """Implementation of `StreamedResponse` for Bedrock models."""
+
+    _model_name: BedrockModelName
+    _event_stream: EventStream[ConverseStreamOutputTypeDef]
+    _timestamp: datetime = field(default_factory=_utils.now_utc)
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        """Return an async iterator of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s.
+
+        This method should be implemented by subclasses to translate the vendor-specific stream of events into
+        pydantic_ai-format events.
+        """
+        chunk: ConverseStreamOutputTypeDef
+        tool_id: str | None = None
+        async for chunk in _AsyncIteratorWrapper(self._event_stream):
+            # TODO(Marcelo): Switch this to `match` when we drop Python 3.9 support.
+            if 'messageStart' in chunk:
+                continue
+            if 'messageStop' in chunk:
+                continue
+            if 'metadata' in chunk:
+                if 'usage' in chunk['metadata']:
+                    self._usage += self._map_usage(chunk['metadata'])
+                continue
+            if 'contentBlockStart' in chunk:
+                index = chunk['contentBlockStart']['contentBlockIndex']
+                start = chunk['contentBlockStart']['start']
+                if 'toolUse' in start:
+                    tool_use_start = start['toolUse']
+                    tool_id = tool_use_start['toolUseId']
+                    tool_name = tool_use_start['name']
+                    maybe_event = self._parts_manager.handle_tool_call_delta(
+                        vendor_part_id=index,
+                        tool_name=tool_name,
+                        args=None,
+                        tool_call_id=tool_id,
+                    )
+                    if maybe_event:
+                        yield maybe_event
+            if 'contentBlockDelta' in chunk:
+                index = chunk['contentBlockDelta']['contentBlockIndex']
+                delta = chunk['contentBlockDelta']['delta']
+                if 'text' in delta:
+                    yield self._parts_manager.handle_text_delta(vendor_part_id=index, content=delta['text'])
+                if 'toolUse' in delta:
+                    tool_use = delta['toolUse']
+                    maybe_event = self._parts_manager.handle_tool_call_delta(
+                        vendor_part_id=index,
+                        tool_name=tool_use.get('name'),
+                        args=tool_use.get('input'),
+                        tool_call_id=tool_id,
+                    )
+                    if maybe_event:
+                        yield maybe_event
+
+    @property
+    def timestamp(self) -> datetime:
+        return self._timestamp
+
+    @property
+    def model_name(self) -> str:
+        """Get the model name of the response."""
+        return self._model_name
+
+    def _map_usage(self, metadata: ConverseStreamMetadataEventTypeDef) -> result.Usage:
+        return result.Usage(
+            request_tokens=metadata['usage']['inputTokens'],
+            response_tokens=metadata['usage']['outputTokens'],
+            total_tokens=metadata['usage']['totalTokens'],
+        )
+
+
+class _AsyncIteratorWrapper(Generic[T]):
+    """Wrap a synchronous iterator in an async iterator."""
+
+    def __init__(self, sync_iterator: Iterable[T]):
+        self.sync_iterator = iter(sync_iterator)
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> T:
+        try:
+            # Run the synchronous next() call in a thread pool
+            item = await anyio.to_thread.run_sync(next, self.sync_iterator)
+            return item
+        except RuntimeError as e:
+            if type(e.__cause__) is StopIteration:
+                raise StopAsyncIteration
+            else:
+                raise e

{pydantic_ai_slim-0.0.34 → pydantic_ai_slim-0.0.36}/pydantic_ai/models/vertexai.py

@@ -224,15 +224,13 @@ class BearerTokenAuth:
 
 
 VertexAiRegion = Literal[
-    'us-central1',
-    'us-east1',
-    'us-east4',
-    'us-south1',
-    'us-west1',
-    'us-west2',
-    'us-west3',
-    'us-west4',
-    'us-east5',
+    'asia-east1',
+    'asia-east2',
+    'asia-northeast1',
+    'asia-northeast3',
+    'asia-south1',
+    'asia-southeast1',
+    'australia-southeast1',
     'europe-central2',
     'europe-north1',
     'europe-southwest1',
@@ -243,27 +241,20 @@ VertexAiRegion = Literal[
     'europe-west6',
     'europe-west8',
     'europe-west9',
-    'europe-west12',
-    'africa-south1',
-    'asia-east1',
-    'asia-east2',
-    'asia-northeast1',
-    'asia-northeast2',
-    'asia-northeast3',
-    'asia-south1',
-    'asia-southeast1',
-    'asia-southeast2',
-    'australia-southeast1',
-    'australia-southeast2',
     'me-central1',
     'me-central2',
     'me-west1',
     'northamerica-northeast1',
-    'northamerica-northeast2',
     'southamerica-east1',
-    'southamerica-west1',
+    'us-central1',
+    'us-east1',
+    'us-east4',
+    'us-east5',
+    'us-south1',
+    'us-west1',
+    'us-west4',
 ]
 """Regions available for Vertex AI.
 
-More details [here](https://cloud.google.com/vertex-ai/docs/reference/rest#rest_endpoints).
+More details [here](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#genai-locations).
 """

{pydantic_ai_slim-0.0.34 → pydantic_ai_slim-0.0.36}/pydantic_ai/providers/__init__.py

@@ -60,5 +60,10 @@ def infer_provider(provider: str) -> Provider[Any]:
         from .google_gla import GoogleGLAProvider
 
         return GoogleGLAProvider()
+    # NOTE: We don't test because there are many ways the `boto3.client` can retrieve the credentials.
+    elif provider == 'bedrock':  # pragma: no cover
+        from .bedrock import BedrockProvider
+
+        return BedrockProvider()
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')

pydantic_ai_slim-0.0.36/pydantic_ai/providers/bedrock.py

@@ -0,0 +1,76 @@
+from __future__ import annotations as _annotations
+
+from typing import overload
+
+from pydantic_ai.providers import Provider
+
+try:
+    import boto3
+    from botocore.client import BaseClient
+    from botocore.exceptions import NoRegionError
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install `boto3` to use the Bedrock provider, '
+        "you can use the `bedrock` optional group — `pip install 'pydantic-ai-slim[bedrock]'`"
+    ) from _import_error
+
+
+class BedrockProvider(Provider[BaseClient]):
+    """Provider for AWS Bedrock."""
+
+    @property
+    def name(self) -> str:
+        return 'bedrock'
+
+    @property
+    def base_url(self) -> str:
+        return self._client.meta.endpoint_url
+
+    @property
+    def client(self) -> BaseClient:
+        return self._client
+
+    @overload
+    def __init__(self, *, bedrock_client: BaseClient) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        *,
+        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,
+    ) -> None: ...
+
+    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,
+    ) -> 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.
+        """
+        if bedrock_client is not None:
+            self._client = bedrock_client
+        else:
+            try:
+                self._client = boto3.client(  # type: ignore[reportUnknownMemberType]
+                    'bedrock-runtime',
+                    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,
+                )
+            except NoRegionError as exc:  # pragma: no cover
+                raise ValueError('You must provide a `region_name` or a boto3 client for Bedrock Runtime.') from exc

{pydantic_ai_slim-0.0.34 → pydantic_ai_slim-0.0.36}/pydantic_ai/providers/google_vertex.py

@@ -155,15 +155,13 @@ async def _creds_from_file(service_account_file: str | Path) -> ServiceAccountCr
 
 
 VertexAiRegion = Literal[
-    'us-central1',
-    'us-east1',
-    'us-east4',
-    'us-south1',
-    'us-west1',
-    'us-west2',
-    'us-west3',
-    'us-west4',
-    'us-east5',
+    'asia-east1',
+    'asia-east2',
+    'asia-northeast1',
+    'asia-northeast3',
+    'asia-south1',
+    'asia-southeast1',
+    'australia-southeast1',
     'europe-central2',
     'europe-north1',
     'europe-southwest1',
@@ -174,27 +172,20 @@ VertexAiRegion = Literal[
     'europe-west6',
     'europe-west8',
     'europe-west9',
-    'europe-west12',
-    'africa-south1',
-    'asia-east1',
-    'asia-east2',
-    'asia-northeast1',
-    'asia-northeast2',
-    'asia-northeast3',
-    'asia-south1',
-    'asia-southeast1',
-    'asia-southeast2',
-    'australia-southeast1',
-    'australia-southeast2',
     'me-central1',
     'me-central2',
     'me-west1',
     'northamerica-northeast1',
-    'northamerica-northeast2',
     'southamerica-east1',
-    'southamerica-west1',
+    'us-central1',
+    'us-east1',
+    'us-east4',
+    'us-east5',
+    'us-south1',
+    'us-west1',
+    'us-west4',
 ]
 """Regions available for Vertex AI.
 
-More details [here](https://cloud.google.com/vertex-ai/docs/reference/rest#rest_endpoints).
+More details [here](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#genai-locations).
 """

{pydantic_ai_slim-0.0.34 → pydantic_ai_slim-0.0.36}/pydantic_ai/settings.py

@@ -27,6 +27,7 @@ class ModelSettings(TypedDict, total=False):
     * Groq
     * Cohere
     * Mistral
+    * Bedrock
     """
 
     temperature: float
@@ -45,6 +46,7 @@ class ModelSettings(TypedDict, total=False):
     * Groq
     * Cohere
     * Mistral
+    * Bedrock
     """
 
     top_p: float
@@ -62,6 +64,7 @@ class ModelSettings(TypedDict, total=False):
     * Groq
     * Cohere
     * Mistral
+    * Bedrock
     """
 
     timeout: float | Timeout

{pydantic_ai_slim-0.0.34 → pydantic_ai_slim-0.0.36}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai-slim"
-version = "0.0.34"
+version = "0.0.36"
 description = "Agent Framework / shim to use Pydantic with LLMs, slim package"
 authors = [{ name = "Samuel Colvin", email = "samuel@pydantic.dev" }]
 license = "MIT"
@@ -36,7 +36,7 @@ dependencies = [
     "griffe>=1.3.2",
     "httpx>=0.27",
     "pydantic>=2.10",
-    "pydantic-graph==0.0.34",
+    "pydantic-graph==0.0.36",
     "exceptiongroup; python_version < '3.11'",
     "opentelemetry-api>=1.28.0",
     "typing-inspection>=0.4.0",
@@ -52,6 +52,7 @@ vertexai = ["google-auth>=2.36.0", "requests>=2.32.3"]
 anthropic = ["anthropic>=0.49.0"]
 groq = ["groq>=0.12.0"]
 mistral = ["mistralai>=1.2.5"]
+bedrock = ["boto3>=1.34.116"]
 # Tools
 duckduckgo = ["duckduckgo-search>=7.0.0"]
 tavily = ["tavily-python>=0.5.0"]
@@ -71,6 +72,7 @@ dev = [
     "pytest-pretty>=1.2.0",
     "pytest-recording>=0.13.2",
     "diff-cover>=9.2.0",
+    "boto3-stubs[bedrock-runtime]",
 ]
 
 [project.scripts]