pydantic-ai-slim 0.7.3__py3-none-any.whl → 0.7.5__py3-none-any.whl

pydantic_ai/models/google.py

@@ -13,7 +13,7 @@ from typing_extensions import assert_never
 from .. import UnexpectedModelBehavior, _utils, usage
 from .._output import OutputObjectDefinition
 from .._run_context import RunContext
-from ..builtin_tools import CodeExecutionTool, WebSearchTool
+from ..builtin_tools import CodeExecutionTool, UrlContextTool, WebSearchTool
 from ..exceptions import UserError
 from ..messages import (
     BinaryContent,
@@ -72,6 +72,7 @@ try:
         ToolConfigDict,
         ToolDict,
         ToolListUnionDict,
+        UrlContextDict,
     )
 
     from ..providers.google import GoogleProvider
@@ -218,7 +219,7 @@ class GoogleModel(Model):
         )
         if self._provider.name != 'google-gla':
             # The fields are not supported by the Gemini API per https://github.com/googleapis/python-genai/blob/7e4ec284dc6e521949626f3ed54028163ef9121d/google/genai/models.py#L1195-L1214
-            config.update(
+            config.update(  # pragma: lax no cover
                 system_instruction=generation_config.get('system_instruction'),
                 tools=cast(list[ToolDict], generation_config.get('tools')),
                 # Annoyingly, GenerationConfigDict has fewer fields than GenerateContentConfigDict, and no extra fields are allowed.
@@ -270,6 +271,8 @@ class GoogleModel(Model):
         for tool in model_request_parameters.builtin_tools:
             if isinstance(tool, WebSearchTool):
                 tools.append(ToolDict(google_search=GoogleSearchDict()))
+            elif isinstance(tool, UrlContextTool):
+                tools.append(ToolDict(url_context=UrlContextDict()))
             elif isinstance(tool, CodeExecutionTool):  # pragma: no branch
                 tools.append(ToolDict(code_execution=ToolCodeExecutionDict()))
             else:  # pragma: no cover
@@ -374,23 +377,25 @@ class GoogleModel(Model):
     def _process_response(self, response: GenerateContentResponse) -> ModelResponse:
         if not response.candidates or len(response.candidates) != 1:
             raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response')  # pragma: no cover
-        if response.candidates[0].content is None or response.candidates[0].content.parts is None:
-            if response.candidates[0].finish_reason == 'SAFETY':
+        candidate = response.candidates[0]
+        if candidate.content is None or candidate.content.parts is None:
+            if candidate.finish_reason == 'SAFETY':
                 raise UnexpectedModelBehavior('Safety settings triggered', str(response))
             else:
                 raise UnexpectedModelBehavior(
                     'Content field missing from Gemini response', str(response)
                 )  # pragma: no cover
-        parts = response.candidates[0].content.parts or []
+        parts = candidate.content.parts or []
         vendor_id = response.response_id or None
         vendor_details: dict[str, Any] | None = None
-        finish_reason = response.candidates[0].finish_reason
+        finish_reason = candidate.finish_reason
         if finish_reason:  # pragma: no branch
             vendor_details = {'finish_reason': finish_reason.value}
         usage = _metadata_as_usage(response)
         return _process_response_from_parts(
             parts,
             response.model_version or self._model_name,
+            self._provider.name,
             usage,
             vendor_id=vendor_id,
             vendor_details=vendor_details,
@@ -410,6 +415,7 @@ class GoogleModel(Model):
             _model_name=self._model_name,
             _response=peekable_response,
             _timestamp=first_chunk.create_time or _utils.now_utc(),
+            _provider_name=self._provider.name,
         )
 
     async def _map_messages(self, messages: list[ModelMessage]) -> tuple[ContentDict | None, list[ContentUnionDict]]:
@@ -519,6 +525,7 @@ class GeminiStreamedResponse(StreamedResponse):
     _model_name: GoogleModelName
     _response: AsyncIterator[GenerateContentResponse]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
@@ -526,10 +533,16 @@ class GeminiStreamedResponse(StreamedResponse):
 
             assert chunk.candidates is not None
             candidate = chunk.candidates[0]
-            if candidate.content is None:
-                raise UnexpectedModelBehavior('Streamed response has no content field')  # pragma: no cover
-            assert candidate.content.parts is not None
-            for part in candidate.content.parts:
+            if candidate.content is None or candidate.content.parts is None:
+                if candidate.finish_reason == 'STOP':  # pragma: no cover
+                    # Normal completion - skip this chunk
+                    continue
+                elif candidate.finish_reason == 'SAFETY':  # pragma: no cover
+                    raise UnexpectedModelBehavior('Safety settings triggered', str(chunk))
+                else:  # pragma: no cover
+                    raise UnexpectedModelBehavior('Content field missing from streaming Gemini response', str(chunk))
+            parts = candidate.content.parts or []
+            for part in parts:
                 if part.text is not None:
                     if part.thought:
                         yield self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=part.text)
@@ -554,6 +567,11 @@ class GeminiStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""
@@ -589,6 +607,7 @@ def _content_model_response(m: ModelResponse) -> ContentDict:
 def _process_response_from_parts(
     parts: list[Part],
     model_name: GoogleModelName,
+    provider_name: str,
     usage: usage.RequestUsage,
     vendor_id: str | None,
     vendor_details: dict[str, Any] | None = None,
@@ -626,7 +645,12 @@ def _process_response_from_parts(
                 f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
             )
     return ModelResponse(
-        parts=items, model_name=model_name, usage=usage, provider_request_id=vendor_id, provider_details=vendor_details
+        parts=items,
+        model_name=model_name,
+        usage=usage,
+        provider_request_id=vendor_id,
+        provider_details=vendor_details,
+        provider_name=provider_name,
     )
 
 
@@ -654,7 +678,7 @@ def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
     if cached_content_token_count := metadata.cached_content_token_count:
         details['cached_content_tokens'] = cached_content_token_count
 
-    if thoughts_token_count := metadata.thoughts_token_count:
+    if thoughts_token_count := (metadata.thoughts_token_count or 0):
         details['thoughts_tokens'] = thoughts_token_count
 
     if tool_use_prompt_token_count := metadata.tool_use_prompt_token_count:
@@ -687,7 +711,7 @@ def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
 
     return usage.RequestUsage(
         input_tokens=metadata.prompt_token_count or 0,
-        output_tokens=metadata.candidates_token_count or 0,
+        output_tokens=(metadata.candidates_token_count or 0) + thoughts_token_count,
         cache_read_tokens=cached_content_token_count or 0,
         input_audio_tokens=input_audio_tokens,
         output_audio_tokens=output_audio_tokens,

pydantic_ai/models/groq.py

@@ -290,6 +290,7 @@ class GroqModel(Model):
             model_name=response.model,
             timestamp=timestamp,
             provider_request_id=response.id,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -309,6 +310,7 @@ class GroqModel(Model):
             _model_name=self._model_name,
             _model_profile=self.profile,
             _timestamp=number_to_datetime(first_chunk.created),
+            _provider_name=self._provider.name,
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[chat.ChatCompletionToolParam]:
@@ -444,6 +446,7 @@ class GroqStreamedResponse(StreamedResponse):
     _model_profile: ModelProfile
     _response: AsyncIterable[chat.ChatCompletionChunk]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
@@ -482,6 +485,11 @@ class GroqStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai/models/huggingface.py

@@ -272,6 +272,7 @@ class HuggingFaceModel(Model):
             model_name=response.model,
             timestamp=timestamp,
             provider_request_id=response.id,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -291,6 +292,7 @@ class HuggingFaceModel(Model):
             _model_profile=self.profile,
             _response=peekable_response,
             _timestamp=datetime.fromtimestamp(first_chunk.created, tz=timezone.utc),
+            _provider_name=self._provider.name,
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ChatCompletionInputTool]:
@@ -437,6 +439,7 @@ class HuggingFaceStreamedResponse(StreamedResponse):
     _model_profile: ModelProfile
     _response: AsyncIterable[ChatCompletionStreamOutput]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
@@ -474,6 +477,11 @@ class HuggingFaceStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai/models/instrumented.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+import itertools
 import json
 from collections.abc import AsyncIterator, Iterator, Mapping
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
-from typing import Any, Callable, Literal
+from typing import Any, Callable, Literal, cast
 from urllib.parse import urlparse
 
 from opentelemetry._events import (
@@ -18,8 +19,14 @@ from opentelemetry.trace import Span, Tracer, TracerProvider, get_tracer_provide
 from opentelemetry.util.types import AttributeValue
 from pydantic import TypeAdapter
 
+from .. import _otel_messages
 from .._run_context import RunContext
-from ..messages import ModelMessage, ModelRequest, ModelResponse
+from ..messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    SystemPromptPart,
+)
 from ..settings import ModelSettings
 from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse
 from .wrapper import WrapperModel
@@ -80,6 +87,8 @@ class InstrumentationSettings:
     event_logger: EventLogger = field(repr=False)
     event_mode: Literal['attributes', 'logs'] = 'attributes'
     include_binary_content: bool = True
+    include_content: bool = True
+    version: Literal[1, 2] = 1
 
     def __init__(
         self,
@@ -90,6 +99,7 @@ class InstrumentationSettings:
         event_logger_provider: EventLoggerProvider | None = None,
         include_binary_content: bool = True,
         include_content: bool = True,
+        version: Literal[1, 2] = 1,
     ):
         """Create instrumentation options.
 
@@ -109,6 +119,10 @@ class InstrumentationSettings:
             include_binary_content: Whether to include binary content in the instrumentation events.
             include_content: Whether to include prompts, completions, and tool call arguments and responses
                 in the instrumentation events.
+            version: Version of the data format.
+                Version 1 is based on the legacy event-based OpenTelemetry GenAI spec.
+                Version 2 stores messages in the attributes `gen_ai.input.messages` and `gen_ai.output.messages`.
+                Version 2 is still WIP and experimental, but will become the default in Pydantic AI v1.
         """
         from pydantic_ai import __version__
 
@@ -122,6 +136,7 @@ class InstrumentationSettings:
         self.event_mode = event_mode
         self.include_binary_content = include_binary_content
         self.include_content = include_content
+        self.version = version
 
         # As specified in the OpenTelemetry GenAI metrics spec:
         # https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/#metric-gen_aiclienttokenusage
@@ -179,6 +194,90 @@ class InstrumentationSettings:
             event.body = InstrumentedModel.serialize_any(event.body)
         return events
 
+    def messages_to_otel_messages(self, messages: list[ModelMessage]) -> list[_otel_messages.ChatMessage]:
+        result: list[_otel_messages.ChatMessage] = []
+        for message in messages:
+            if isinstance(message, ModelRequest):
+                for is_system, group in itertools.groupby(message.parts, key=lambda p: isinstance(p, SystemPromptPart)):
+                    message_parts: list[_otel_messages.MessagePart] = []
+                    for part in group:
+                        if hasattr(part, 'otel_message_parts'):
+                            message_parts.extend(part.otel_message_parts(self))
+                    result.append(
+                        _otel_messages.ChatMessage(role='system' if is_system else 'user', parts=message_parts)
+                    )
+            elif isinstance(message, ModelResponse):  # pragma: no branch
+                result.append(_otel_messages.ChatMessage(role='assistant', parts=message.otel_message_parts(self)))
+        return result
+
+    def handle_messages(self, input_messages: list[ModelMessage], response: ModelResponse, system: str, span: Span):
+        if self.version == 1:
+            events = self.messages_to_otel_events(input_messages)
+            for event in self.messages_to_otel_events([response]):
+                events.append(
+                    Event(
+                        'gen_ai.choice',
+                        body={
+                            'index': 0,
+                            'message': event.body,
+                        },
+                    )
+                )
+            for event in events:
+                event.attributes = {
+                    GEN_AI_SYSTEM_ATTRIBUTE: system,
+                    **(event.attributes or {}),
+                }
+            self._emit_events(span, events)
+        else:
+            output_messages = self.messages_to_otel_messages([response])
+            assert len(output_messages) == 1
+            output_message = cast(_otel_messages.OutputMessage, output_messages[0])
+            if response.provider_details and 'finish_reason' in response.provider_details:
+                output_message['finish_reason'] = response.provider_details['finish_reason']
+            instructions = InstrumentedModel._get_instructions(input_messages)  # pyright: ignore [reportPrivateUsage]
+            attributes = {
+                'gen_ai.input.messages': json.dumps(self.messages_to_otel_messages(input_messages)),
+                'gen_ai.output.messages': json.dumps([output_message]),
+                'logfire.json_schema': json.dumps(
+                    {
+                        'type': 'object',
+                        'properties': {
+                            'gen_ai.input.messages': {'type': 'array'},
+                            'gen_ai.output.messages': {'type': 'array'},
+                            **({'gen_ai.system_instructions': {'type': 'array'}} if instructions else {}),
+                            'model_request_parameters': {'type': 'object'},
+                        },
+                    }
+                ),
+            }
+            if instructions is not None:
+                attributes['gen_ai.system_instructions'] = json.dumps(
+                    [_otel_messages.TextPart(type='text', content=instructions)]
+                )
+            span.set_attributes(attributes)
+
+    def _emit_events(self, span: Span, events: list[Event]) -> None:
+        if self.event_mode == 'logs':
+            for event in events:
+                self.event_logger.emit(event)
+        else:
+            attr_name = 'events'
+            span.set_attributes(
+                {
+                    attr_name: json.dumps([InstrumentedModel.event_to_dict(event) for event in events]),
+                    'logfire.json_schema': json.dumps(
+                        {
+                            'type': 'object',
+                            'properties': {
+                                attr_name: {'type': 'array'},
+                                'model_request_parameters': {'type': 'object'},
+                            },
+                        }
+                    ),
+                }
+            )
+
 
 GEN_AI_SYSTEM_ATTRIBUTE = 'gen_ai.system'
 GEN_AI_REQUEST_MODEL_ATTRIBUTE = 'gen_ai.request.model'
@@ -269,7 +368,7 @@ class InstrumentedModel(WrapperModel):
                     # FallbackModel updates these span attributes.
                     attributes.update(getattr(span, 'attributes', {}))
                     request_model = attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE]
-                    system = attributes[GEN_AI_SYSTEM_ATTRIBUTE]
+                    system = cast(str, attributes[GEN_AI_SYSTEM_ATTRIBUTE])
 
                     response_model = response.model_name or request_model
 
@@ -297,18 +396,7 @@ class InstrumentedModel(WrapperModel):
                     if not span.is_recording():
                         return
 
-                    events = self.instrumentation_settings.messages_to_otel_events(messages)
-                    for event in self.instrumentation_settings.messages_to_otel_events([response]):
-                        events.append(
-                            Event(
-                                'gen_ai.choice',
-                                body={
-                                    # TODO finish_reason
-                                    'index': 0,
-                                    'message': event.body,
-                                },
-                            )
-                        )
+                    self.instrumentation_settings.handle_messages(messages, response, system, span)
                     span.set_attributes(
                         {
                             **response.usage.opentelemetry_attributes(),
@@ -316,12 +404,6 @@ class InstrumentedModel(WrapperModel):
                         }
                     )
                     span.update_name(f'{operation} {request_model}')
-                    for event in events:
-                        event.attributes = {
-                            GEN_AI_SYSTEM_ATTRIBUTE: system,
-                            **(event.attributes or {}),
-                        }
-                    self._emit_events(span, events)
 
                 yield finish
         finally:
@@ -330,27 +412,6 @@ class InstrumentedModel(WrapperModel):
                 # to prevent them from being redundantly recorded in the span itself by logfire.
                 record_metrics()
 
-    def _emit_events(self, span: Span, events: list[Event]) -> None:
-        if self.instrumentation_settings.event_mode == 'logs':
-            for event in events:
-                self.instrumentation_settings.event_logger.emit(event)
-        else:
-            attr_name = 'events'
-            span.set_attributes(
-                {
-                    attr_name: json.dumps([self.event_to_dict(event) for event in events]),
-                    'logfire.json_schema': json.dumps(
-                        {
-                            'type': 'object',
-                            'properties': {
-                                attr_name: {'type': 'array'},
-                                'model_request_parameters': {'type': 'object'},
-                            },
-                        }
-                    ),
-                }
-            )
-
     @staticmethod
     def model_attributes(model: Model):
         attributes: dict[str, AttributeValue] = {

pydantic_ai/models/mistral.py

@@ -353,6 +353,7 @@ class MistralModel(Model):
             model_name=response.model,
             timestamp=timestamp,
             provider_request_id=response.id,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -378,6 +379,7 @@ class MistralModel(Model):
             _response=peekable_response,
             _model_name=self._model_name,
             _timestamp=timestamp,
+            _provider_name=self._provider.name,
         )
 
     @staticmethod
@@ -584,6 +586,7 @@ class MistralStreamedResponse(StreamedResponse):
     _model_name: MistralModelName
     _response: AsyncIterable[MistralCompletionEvent]
     _timestamp: datetime
+    _provider_name: str
 
     _delta_content: str = field(default='', init=False)
 
@@ -631,6 +634,11 @@ class MistralStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai/models/openai.py

@@ -500,6 +500,7 @@ class OpenAIModel(Model):
             timestamp=timestamp,
             provider_details=vendor_details,
             provider_request_id=response.id,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -519,6 +520,7 @@ class OpenAIModel(Model):
             _model_profile=self.profile,
             _response=peekable_response,
             _timestamp=number_to_datetime(first_chunk.created),
+            _provider_name=self._provider.name,
         )
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[chat.ChatCompletionToolParam]:
@@ -571,6 +573,8 @@ class OpenAIModel(Model):
                     # Note: model responses from this model should only have one text item, so the following
                     # shouldn't merge multiple texts into one unless you switch models between runs:
                     message_param['content'] = '\n\n'.join(texts)
+                else:
+                    message_param['content'] = None
                 if tool_calls:
                     message_param['tool_calls'] = tool_calls
                 openai_messages.append(message_param)
@@ -803,6 +807,7 @@ class OpenAIResponsesModel(Model):
             model_name=response.model,
             provider_request_id=response.id,
             timestamp=timestamp,
+            provider_name=self._provider.name,
         )
 
     async def _process_streamed_response(
@@ -822,6 +827,7 @@ class OpenAIResponsesModel(Model):
             _model_name=self._model_name,
             _response=peekable_response,
             _timestamp=number_to_datetime(first_chunk.response.created_at),
+            _provider_name=self._provider.name,
         )
 
     @overload
@@ -1137,6 +1143,7 @@ class OpenAIStreamedResponse(StreamedResponse):
     _model_profile: ModelProfile
     _response: AsyncIterable[ChatCompletionChunk]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
@@ -1180,6 +1187,11 @@ class OpenAIStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""
@@ -1193,6 +1205,7 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
     _model_name: OpenAIModelName
     _response: AsyncIterable[responses.ResponseStreamEvent]
     _timestamp: datetime
+    _provider_name: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         async for chunk in self._response:
@@ -1313,6 +1326,11 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai/models/test.py

@@ -131,6 +131,7 @@ class TestModel(Model):
             _model_name=self._model_name,
             _structured_response=model_response,
             _messages=messages,
+            _provider_name=self._system,
         )
 
     @property
@@ -263,6 +264,7 @@ class TestStreamedResponse(StreamedResponse):
     _model_name: str
     _structured_response: ModelResponse
     _messages: InitVar[Iterable[ModelMessage]]
+    _provider_name: str
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
 
     def __post_init__(self, _messages: Iterable[ModelMessage]):
@@ -305,6 +307,11 @@ class TestStreamedResponse(StreamedResponse):
         """Get the model name of the response."""
         return self._model_name
 
+    @property
+    def provider_name(self) -> str:
+        """Get the provider name."""
+        return self._provider_name
+
     @property
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai/providers/anthropic.py

@@ -1,9 +1,10 @@
 from __future__ import annotations as _annotations
 
 import os
-from typing import overload
+from typing import Union, overload
 
 import httpx
+from typing_extensions import TypeAlias
 
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.models import cached_async_http_client
@@ -12,15 +13,18 @@ from pydantic_ai.profiles.anthropic import anthropic_model_profile
 from pydantic_ai.providers import Provider
 
 try:
-    from anthropic import AsyncAnthropic
-except ImportError as _import_error:  # pragma: no cover
+    from anthropic import AsyncAnthropic, AsyncAnthropicBedrock
+except ImportError as _import_error:
     raise ImportError(
         'Please install the `anthropic` package to use the Anthropic provider, '
         'you can use the `anthropic` optional group — `pip install "pydantic-ai-slim[anthropic]"`'
     ) from _import_error
 
 
-class AnthropicProvider(Provider[AsyncAnthropic]):
+AsyncAnthropicClient: TypeAlias = Union[AsyncAnthropic, AsyncAnthropicBedrock]
+
+
+class AnthropicProvider(Provider[AsyncAnthropicClient]):
     """Provider for Anthropic API."""
 
     @property
@@ -32,14 +36,14 @@ class AnthropicProvider(Provider[AsyncAnthropic]):
         return str(self._client.base_url)
 
     @property
-    def client(self) -> AsyncAnthropic:
+    def client(self) -> AsyncAnthropicClient:
         return self._client
 
     def model_profile(self, model_name: str) -> ModelProfile | None:
         return anthropic_model_profile(model_name)
 
     @overload
-    def __init__(self, *, anthropic_client: AsyncAnthropic | None = None) -> None: ...
+    def __init__(self, *, anthropic_client: AsyncAnthropicClient | None = None) -> None: ...
 
     @overload
     def __init__(self, *, api_key: str | None = None, http_client: httpx.AsyncClient | None = None) -> None: ...
@@ -48,7 +52,7 @@ class AnthropicProvider(Provider[AsyncAnthropic]):
         self,
         *,
         api_key: str | None = None,
-        anthropic_client: AsyncAnthropic | None = None,
+        anthropic_client: AsyncAnthropicClient | None = None,
         http_client: httpx.AsyncClient | None = None,
     ) -> None:
         """Create a new Anthropic provider.
@@ -71,7 +75,6 @@ class AnthropicProvider(Provider[AsyncAnthropic]):
                     'Set the `ANTHROPIC_API_KEY` environment variable or pass it via `AnthropicProvider(api_key=...)`'
                     'to use the Anthropic provider.'
                 )
-
             if http_client is not None:
                 self._client = AsyncAnthropic(api_key=api_key, http_client=http_client)
             else: