pydantic-ai-slim 0.8.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

pydantic_ai/models/anthropic.py

@@ -6,7 +6,7 @@ from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from typing import Any, Literal, Union, cast, overload
+from typing import Any, Literal, cast, overload
 
 from typing_extensions import assert_never
 
@@ -99,7 +99,7 @@ except ImportError as _import_error:
 LatestAnthropicModelNames = ModelParam
 """Latest Anthropic models."""
 
-AnthropicModelName = Union[str, LatestAnthropicModelNames]
+AnthropicModelName = str | LatestAnthropicModelNames
 """Possible Anthropic model names.
 
 Since Anthropic supports a variety of date-stamped models, we explicitly list the latest models but
@@ -290,7 +290,7 @@ class AnthropicModel(Model):
         for item in response.content:
             if isinstance(item, BetaTextBlock):
                 items.append(TextPart(content=item.text))
-            elif isinstance(item, (BetaWebSearchToolResultBlock, BetaCodeExecutionToolResultBlock)):
+            elif isinstance(item, BetaWebSearchToolResultBlock | BetaCodeExecutionToolResultBlock):
                 items.append(
                     BuiltinToolReturnPart(
                         provider_name='anthropic',
@@ -327,7 +327,7 @@ class AnthropicModel(Model):
                 )
 
         return ModelResponse(
-            items,
+            parts=items,
             usage=_map_usage(response),
             model_name=response.model,
             provider_response_id=response.id,
@@ -536,7 +536,7 @@ class AnthropicModel(Model):
         }
 
 
-def _map_usage(message: BetaMessage | BetaRawMessageStreamEvent) -> usage.RequestUsage:
+def _map_usage(message: BetaMessage | BetaRawMessageStartEvent | BetaRawMessageDeltaEvent) -> usage.RequestUsage:
     if isinstance(message, BetaMessage):
         response_usage = message.usage
     elif isinstance(message, BetaRawMessageStartEvent):
@@ -544,12 +544,7 @@ def _map_usage(message: BetaMessage | BetaRawMessageStreamEvent) -> usage.Reques
     elif isinstance(message, BetaRawMessageDeltaEvent):
         response_usage = message.usage
     else:
-        # No usage information provided in:
-        # - RawMessageStopEvent
-        # - RawContentBlockStartEvent
-        # - RawContentBlockDeltaEvent
-        # - RawContentBlockStopEvent
-        return usage.RequestUsage()
+        assert_never(message)
 
     # Store all integer-typed usage values in the details, except 'output_tokens' which is represented exactly by
     # `response_tokens`
@@ -586,10 +581,8 @@ class AnthropicStreamedResponse(StreamedResponse):
         current_block: BetaContentBlock | None = None
 
         async for event in self._response:
-            self._usage += _map_usage(event)
-
             if isinstance(event, BetaRawMessageStartEvent):
-                pass
+                self._usage = _map_usage(event)
 
             elif isinstance(event, BetaRawContentBlockStartEvent):
                 current_block = event.content_block
@@ -652,9 +645,9 @@ class AnthropicStreamedResponse(StreamedResponse):
                     pass
 
             elif isinstance(event, BetaRawMessageDeltaEvent):
-                pass
+                self._usage = _map_usage(event)
 
-            elif isinstance(event, (BetaRawContentBlockStopEvent, BetaRawMessageStopEvent)):  # pragma: no branch
+            elif isinstance(event, BetaRawContentBlockStopEvent | BetaRawMessageStopEvent):  # pragma: no branch
                 current_block = None
 
     @property

pydantic_ai/models/bedrock.py

@@ -2,13 +2,12 @@ from __future__ import annotations
 
 import functools
 import typing
-import warnings
 from collections.abc import AsyncIterator, Iterable, Iterator, Mapping
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
 from itertools import count
-from typing import TYPE_CHECKING, Any, Generic, Literal, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload
 
 import anyio
 import anyio.to_thread
@@ -125,7 +124,7 @@ LatestBedrockModelNames = Literal[
 ]
 """Latest Bedrock models."""
 
-BedrockModelName = Union[str, LatestBedrockModelNames]
+BedrockModelName = 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.
@@ -303,7 +302,7 @@ class BedrockConverseModel(Model):
         )
         response_id = response.get('ResponseMetadata', {}).get('RequestId', None)
         return ModelResponse(
-            items,
+            parts=items,
             usage=u,
             model_name=self.model_name,
             provider_response_id=response_id,
@@ -490,7 +489,7 @@ class BedrockConverseModel(Model):
                         else:
                             # NOTE: We don't pass the thinking part to Bedrock for models other than Claude since it raises an error.
                             pass
-                    elif isinstance(item, (BuiltinToolCallPart, BuiltinToolReturnPart)):
+                    elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):
                         pass
                     else:
                         assert isinstance(item, ToolCallPart)
@@ -546,7 +545,7 @@ class BedrockConverseModel(Model):
                         content.append({'video': {'format': format, 'source': {'bytes': item.data}}})
                     else:
                         raise NotImplementedError('Binary content is not supported yet.')
-                elif isinstance(item, (ImageUrl, DocumentUrl, VideoUrl)):
+                elif isinstance(item, ImageUrl | DocumentUrl | VideoUrl):
                     downloaded_item = await download_item(item, data_format='bytes', type_format='extension')
                     format = downloaded_item['data_type']
                     if item.kind == 'image-url':
@@ -601,7 +600,7 @@ class BedrockStreamedResponse(StreamedResponse):
     _provider_name: str
     _timestamp: datetime = field(default_factory=_utils.now_utc)
 
-    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
+    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
@@ -610,60 +609,55 @@ class BedrockStreamedResponse(StreamedResponse):
         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']:  # pragma: no branch
-                    self._usage += self._map_usage(chunk['metadata'])
-                continue
-            if 'contentBlockStart' in chunk:
-                index = chunk['contentBlockStart']['contentBlockIndex']
-                start = chunk['contentBlockStart']['start']
-                if 'toolUse' in start:  # pragma: no branch
-                    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:  # pragma: no branch
-                        yield maybe_event
-            if 'contentBlockDelta' in chunk:
-                index = chunk['contentBlockDelta']['contentBlockIndex']
-                delta = chunk['contentBlockDelta']['delta']
-                if 'reasoningContent' in delta:
-                    if text := delta['reasoningContent'].get('text'):
+            match chunk:
+                case {'messageStart': _}:
+                    continue
+                case {'messageStop': _}:
+                    continue
+                case {'metadata': metadata}:
+                    if 'usage' in metadata:  # pragma: no branch
+                        self._usage += self._map_usage(metadata)
+                    continue
+                case {'contentBlockStart': content_block_start}:
+                    index = content_block_start['contentBlockIndex']
+                    start = content_block_start['start']
+                    if 'toolUse' in start:  # pragma: no branch
+                        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:  # pragma: no branch
+                            yield maybe_event
+                case {'contentBlockDelta': content_block_delta}:
+                    index = content_block_delta['contentBlockIndex']
+                    delta = content_block_delta['delta']
+                    if 'reasoningContent' in delta:
                         yield self._parts_manager.handle_thinking_delta(
                             vendor_part_id=index,
-                            content=text,
+                            content=delta['reasoningContent'].get('text'),
                             signature=delta['reasoningContent'].get('signature'),
                         )
-                    else:  # pragma: no cover
-                        warnings.warn(
-                            f'Only text reasoning content is supported yet, but you got {delta["reasoningContent"]}. '
-                            'Please report this to the maintainers.',
-                            UserWarning,
+                    if 'text' in delta:
+                        maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=index, content=delta['text'])
+                        if maybe_event is not None:  # pragma: no branch
+                            yield maybe_event
+                    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 'text' in delta:
-                    maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=index, content=delta['text'])
-                    if maybe_event is not None:  # pragma: no branch
-                        yield maybe_event
-                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:  # pragma: no branch
-                        yield maybe_event
+                        if maybe_event:  # pragma: no branch
+                            yield maybe_event
+                case _:
+                    pass  # pyright wants match statements to be exhaustive
 
     @property
     def model_name(self) -> str:

pydantic_ai/models/cohere.py

@@ -2,7 +2,7 @@ from __future__ import annotations as _annotations
 
 from collections.abc import Iterable
 from dataclasses import dataclass, field
-from typing import Literal, Union, cast
+from typing import Literal, cast
 
 from typing_extensions import assert_never
 
@@ -72,7 +72,7 @@ LatestCohereModelNames = Literal[
 ]
 """Latest Cohere models."""
 
-CohereModelName = Union[str, LatestCohereModelNames]
+CohereModelName = str | LatestCohereModelNames
 """Possible Cohere model names.
 
 Since Cohere supports a variety of date-stamped models, we explicitly list the latest models but
@@ -228,7 +228,7 @@ class CohereModel(Model):
                         pass
                     elif isinstance(item, ToolCallPart):
                         tool_calls.append(self._map_tool_call(item))
-                    elif isinstance(item, (BuiltinToolCallPart, BuiltinToolReturnPart)):  # pragma: no cover
+                    elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
                         # This is currently never returned from cohere
                         pass
                     else:

pydantic_ai/models/fallback.py

@@ -1,9 +1,9 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 from contextlib import AsyncExitStack, asynccontextmanager, suppress
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 from opentelemetry.trace import get_current_span
 

pydantic_ai/models/function.py

@@ -2,14 +2,14 @@ from __future__ import annotations as _annotations
 
 import inspect
 import re
-from collections.abc import AsyncIterator, Awaitable, Iterable, Sequence
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
 from contextlib import asynccontextmanager
-from dataclasses import dataclass, field
+from dataclasses import KW_ONLY, dataclass, field
 from datetime import datetime
 from itertools import chain
-from typing import Any, Callable, Union
+from typing import Any, TypeAlias
 
-from typing_extensions import TypeAlias, assert_never, overload
+from typing_extensions import assert_never, overload
 
 from .. import _utils, usage
 from .._run_context import RunContext
@@ -44,8 +44,8 @@ class FunctionModel(Model):
     Apart from `__init__`, all methods are private or match those of the base class.
     """
 
-    function: FunctionDef | None = None
-    stream_function: StreamFunctionDef | None = None
+    function: FunctionDef | None
+    stream_function: StreamFunctionDef | None
 
     _model_name: str = field(repr=False)
     _system: str = field(default='function', repr=False)
@@ -120,10 +120,10 @@ class FunctionModel(Model):
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
         agent_info = AgentInfo(
-            model_request_parameters.function_tools,
-            model_request_parameters.allow_text_output,
-            model_request_parameters.output_tools,
-            model_settings,
+            function_tools=model_request_parameters.function_tools,
+            allow_text_output=model_request_parameters.allow_text_output,
+            output_tools=model_request_parameters.output_tools,
+            model_settings=model_settings,
         )
 
         assert self.function is not None, 'FunctionModel must receive a `function` to support non-streamed requests'
@@ -149,10 +149,10 @@ class FunctionModel(Model):
         run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         agent_info = AgentInfo(
-            model_request_parameters.function_tools,
-            model_request_parameters.allow_text_output,
-            model_request_parameters.output_tools,
-            model_settings,
+            function_tools=model_request_parameters.function_tools,
+            allow_text_output=model_request_parameters.allow_text_output,
+            output_tools=model_request_parameters.output_tools,
+            model_settings=model_settings,
         )
 
         assert self.stream_function is not None, (
@@ -182,7 +182,7 @@ class FunctionModel(Model):
         return self._system
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, kw_only=True)
 class AgentInfo:
     """Information about an agent.
 
@@ -212,13 +212,17 @@ class DeltaToolCall:
 
     name: str | None = None
     """Incremental change to the name of the tool."""
+
     json_args: str | None = None
     """Incremental change to the arguments as JSON"""
+
+    _: KW_ONLY
+
     tool_call_id: str | None = None
     """Incremental change to the tool call ID."""
 
 
-@dataclass
+@dataclass(kw_only=True)
 class DeltaThinkingPart:
     """Incremental change to a thinking part.
 
@@ -237,18 +241,16 @@ DeltaToolCalls: TypeAlias = dict[int, DeltaToolCall]
 DeltaThinkingCalls: TypeAlias = dict[int, DeltaThinkingPart]
 """A mapping of thinking call IDs to incremental changes."""
 
-# TODO: Change the signature to Callable[[list[ModelMessage], ModelSettings, ModelRequestParameters], ...]
-FunctionDef: TypeAlias = Callable[[list[ModelMessage], AgentInfo], Union[ModelResponse, Awaitable[ModelResponse]]]
+FunctionDef: TypeAlias = Callable[[list[ModelMessage], AgentInfo], ModelResponse | Awaitable[ModelResponse]]
 """A function used to generate a non-streamed response."""
 
-# TODO: Change signature as indicated above
 StreamFunctionDef: TypeAlias = Callable[
-    [list[ModelMessage], AgentInfo], AsyncIterator[Union[str, DeltaToolCalls, DeltaThinkingCalls]]
+    [list[ModelMessage], AgentInfo], AsyncIterator[str | DeltaToolCalls | DeltaThinkingCalls]
 ]
 """A function used to generate a streamed response.
 
-While this is defined as having return type of `AsyncIterator[Union[str, DeltaToolCalls, DeltaThinkingCalls]]`, it should
-really be considered as `Union[AsyncIterator[str], AsyncIterator[DeltaToolCalls], AsyncIterator[DeltaThinkingCalls]]`,
+While this is defined as having return type of `AsyncIterator[str | DeltaToolCalls | DeltaThinkingCalls]`, it should
+really be considered as `AsyncIterator[str] | AsyncIterator[DeltaToolCalls] | AsyncIterator[DeltaThinkingCalls]`,
 
 E.g. you need to yield all text, all `DeltaToolCalls`, or all `DeltaThinkingCalls`, not mix them.
 """
@@ -326,7 +328,7 @@ def _estimate_usage(messages: Iterable[ModelMessage]) -> usage.RequestUsage:
     for message in messages:
         if isinstance(message, ModelRequest):
             for part in message.parts:
-                if isinstance(part, (SystemPromptPart, UserPromptPart)):
+                if isinstance(part, SystemPromptPart | UserPromptPart):
                     request_tokens += _estimate_string_tokens(part.content)
                 elif isinstance(part, ToolReturnPart):
                     request_tokens += _estimate_string_tokens(part.model_response_str())

pydantic_ai/models/gemini.py

@@ -5,7 +5,7 @@ from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Annotated, Any, Literal, Protocol, Union, cast
+from typing import Annotated, Any, Literal, Protocol, cast
 from uuid import uuid4
 
 import httpx
@@ -51,7 +51,7 @@ LatestGeminiModelNames = Literal[
 ]
 """Latest Gemini models."""
 
-GeminiModelName = Union[str, LatestGeminiModelNames]
+GeminiModelName = str | LatestGeminiModelNames
 """Possible Gemini model names.
 
 Since Gemini supports a variety of date-stamped models, we explicitly list the latest models but
@@ -211,7 +211,9 @@ class GeminiModel(Model):
         generation_config = _settings_to_generation_config(model_settings)
         if model_request_parameters.output_mode == 'native':
             if tools:
-                raise UserError('Gemini does not support structured output and tools at the same time.')
+                raise UserError(
+                    'Gemini does not support `NativeOutput` and tools at the same time. Use `output_type=ToolOutput(...)` instead.'
+                )
 
             generation_config['response_mime_type'] = 'application/json'
 
@@ -615,7 +617,7 @@ def _content_model_response(m: ModelResponse) -> _GeminiContent:
         elif isinstance(item, TextPart):
             if item.content:
                 parts.append(_GeminiTextPart(text=item.content))
-        elif isinstance(item, (BuiltinToolCallPart, BuiltinToolReturnPart)):  # pragma: no cover
+        elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
             # This is currently never returned from gemini
             pass
         else:
@@ -735,16 +737,13 @@ def _part_discriminator(v: Any) -> str:
 
 # See <https://ai.google.dev/api/caching#Part>
 # we don't currently support other part types
-# TODO discriminator
 _GeminiPartUnion = Annotated[
-    Union[
-        Annotated[_GeminiTextPart, pydantic.Tag('text')],
-        Annotated[_GeminiFunctionCallPart, pydantic.Tag('function_call')],
-        Annotated[_GeminiFunctionResponsePart, pydantic.Tag('function_response')],
-        Annotated[_GeminiInlineDataPart, pydantic.Tag('inline_data')],
-        Annotated[_GeminiFileDataPart, pydantic.Tag('file_data')],
-        Annotated[_GeminiThoughtPart, pydantic.Tag('thought')],
-    ],
+    Annotated[_GeminiTextPart, pydantic.Tag('text')]
+    | Annotated[_GeminiFunctionCallPart, pydantic.Tag('function_call')]
+    | Annotated[_GeminiFunctionResponsePart, pydantic.Tag('function_response')]
+    | Annotated[_GeminiInlineDataPart, pydantic.Tag('inline_data')]
+    | Annotated[_GeminiFileDataPart, pydantic.Tag('file_data')]
+    | Annotated[_GeminiThoughtPart, pydantic.Tag('thought')],
     pydantic.Discriminator(_part_discriminator),
 ]
 

pydantic_ai/models/google.py

@@ -5,7 +5,7 @@ from collections.abc import AsyncIterator, Awaitable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any, Literal, Union, cast, overload
+from typing import Any, Literal, cast, overload
 from uuid import uuid4
 
 from typing_extensions import assert_never
@@ -91,7 +91,7 @@ LatestGoogleModelNames = Literal[
 ]
 """Latest Gemini models."""
 
-GoogleModelName = Union[str, LatestGoogleModelNames]
+GoogleModelName = str | LatestGoogleModelNames
 """Possible Gemini model names.
 
 Since Gemini supports a variety of date-stamped models, we explicitly list the latest models but
@@ -264,6 +264,14 @@ class GoogleModel(Model):
         yield await self._process_streamed_response(response, model_request_parameters)  # type: ignore
 
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolDict] | None:
+        if model_request_parameters.builtin_tools:
+            if model_request_parameters.output_tools:
+                raise UserError(
+                    'Gemini does not support output tools and built-in tools at the same time. Use `output_type=PromptedOutput(...)` instead.'
+                )
+            if model_request_parameters.function_tools:
+                raise UserError('Gemini does not support user tools and built-in tools at the same time.')
+
         tools: list[ToolDict] = [
             ToolDict(function_declarations=[_function_declaration_from_tool(t)])
             for t in model_request_parameters.tool_defs.values()
@@ -334,7 +342,9 @@ class GoogleModel(Model):
         response_schema = None
         if model_request_parameters.output_mode == 'native':
             if tools:
-                raise UserError('Gemini does not support structured output and tools at the same time.')
+                raise UserError(
+                    'Gemini does not support `NativeOutput` and tools at the same time. Use `output_type=ToolOutput(...)` instead.'
+                )
             response_mime_type = 'application/json'
             output_object = model_request_parameters.output_object
             assert output_object is not None
@@ -349,7 +359,7 @@ class GoogleModel(Model):
             'headers': {'Content-Type': 'application/json', 'User-Agent': get_user_agent()}
         }
         if timeout := model_settings.get('timeout'):
-            if isinstance(timeout, (int, float)):
+            if isinstance(timeout, int | float):
                 http_options['timeout'] = int(1000 * timeout)
             else:
                 raise UserError('Google does not support setting ModelSettings.timeout to a httpx.Timeout')
@@ -559,6 +569,10 @@ class GeminiStreamedResponse(StreamedResponse):
                     )
                     if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
+                elif part.executable_code is not None:
+                    pass
+                elif part.code_execution_result is not None:
+                    pass
                 else:
                     assert part.function_response is not None, f'Unexpected part: {part}'  # pragma: no cover