pydantic-ai-slim 1.0.8__py3-none-any.whl → 1.0.10__py3-none-any.whl

pydantic_ai/durable_exec/temporal/_model.py

@@ -4,7 +4,7 @@ from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any
+from typing import Any, cast
 
 from pydantic import ConfigDict, with_config
 from temporalio import activity, workflow
@@ -30,7 +30,8 @@ from ._run_context import TemporalRunContext
 @with_config(ConfigDict(arbitrary_types_allowed=True))
 class _RequestParams:
     messages: list[ModelMessage]
-    model_settings: ModelSettings | None
+    # `model_settings` can't be a `ModelSettings` because Temporal would end up dropping fields only defined on its subclasses.
+    model_settings: dict[str, Any] | None
     model_request_parameters: ModelRequestParameters
     serialized_run_context: Any
 
@@ -82,7 +83,11 @@ class TemporalModel(WrapperModel):
 
         @activity.defn(name=f'{activity_name_prefix}__model_request')
         async def request_activity(params: _RequestParams) -> ModelResponse:
-            return await self.wrapped.request(params.messages, params.model_settings, params.model_request_parameters)
+            return await self.wrapped.request(
+                params.messages,
+                cast(ModelSettings | None, params.model_settings),
+                params.model_request_parameters,
+            )
 
         self.request_activity = request_activity
 
@@ -92,7 +97,10 @@ class TemporalModel(WrapperModel):
 
             run_context = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
             async with self.wrapped.request_stream(
-                params.messages, params.model_settings, params.model_request_parameters, run_context
+                params.messages,
+                cast(ModelSettings | None, params.model_settings),
+                params.model_request_parameters,
+                run_context,
             ) as streamed_response:
                 await self.event_stream_handler(run_context, streamed_response)
 
@@ -124,7 +132,7 @@ class TemporalModel(WrapperModel):
             activity=self.request_activity,
             arg=_RequestParams(
                 messages=messages,
-                model_settings=model_settings,
+                model_settings=cast(dict[str, Any] | None, model_settings),
                 model_request_parameters=model_request_parameters,
                 serialized_run_context=None,
             ),
@@ -161,7 +169,7 @@ class TemporalModel(WrapperModel):
             args=[
                 _RequestParams(
                     messages=messages,
-                    model_settings=model_settings,
+                    model_settings=cast(dict[str, Any] | None, model_settings),
                     model_request_parameters=model_request_parameters,
                     serialized_run_context=serialized_run_context,
                 ),

pydantic_ai/durable_exec/temporal/_run_context.py

@@ -9,7 +9,7 @@ from pydantic_ai.tools import AgentDepsT, RunContext
 class TemporalRunContext(RunContext[AgentDepsT]):
     """The [`RunContext`][pydantic_ai.tools.RunContext] subclass to use to serialize and deserialize the run context for use inside a Temporal activity.
 
-    By default, only the `deps`, `retries`, `tool_call_id`, `tool_name`, `tool_call_approved`, `retry` and `run_step` attributes will be available.
+    By default, only the `deps`, `retries`, `tool_call_id`, `tool_name`, `tool_call_approved`, `retry`, `max_retries` and `run_step` attributes will be available.
     To make another attribute available, create a `TemporalRunContext` subclass with a custom `serialize_run_context` class method that returns a dictionary that includes the attribute and pass it to [`TemporalAgent`][pydantic_ai.durable_exec.temporal.TemporalAgent].
     """
 
@@ -42,6 +42,7 @@ class TemporalRunContext(RunContext[AgentDepsT]):
             'tool_name': ctx.tool_name,
             'tool_call_approved': ctx.tool_call_approved,
             'retry': ctx.retry,
+            'max_retries': ctx.max_retries,
             'run_step': ctx.run_step,
         }
 

pydantic_ai/format_prompt.py

@@ -1,15 +1,17 @@
 from __future__ import annotations as _annotations
 
 from collections.abc import Iterable, Iterator, Mapping
-from dataclasses import asdict, dataclass, is_dataclass
+from dataclasses import asdict, dataclass, field, fields, is_dataclass
 from datetime import date
-from typing import Any
+from typing import Any, Literal
 from xml.etree import ElementTree
 
 from pydantic import BaseModel
 
 __all__ = ('format_as_xml',)
 
+from pydantic.fields import ComputedFieldInfo, FieldInfo
+
 
 def format_as_xml(
     obj: Any,
@@ -17,6 +19,7 @@ def format_as_xml(
     item_tag: str = 'item',
     none_str: str = 'null',
     indent: str | None = '  ',
+    include_field_info: Literal['once'] | bool = False,
 ) -> str:
     """Format a Python object as XML.
 
@@ -33,6 +36,10 @@ def format_as_xml(
             for dataclasses and Pydantic models.
         none_str: String to use for `None` values.
         indent: Indentation string to use for pretty printing.
+        include_field_info: Whether to include attributes like Pydantic `Field` attributes and dataclasses `field()`
+            `metadata` as XML attributes. In both cases the allowed `Field` attributes and `field()` metadata keys are
+            `title` and `description`. If a field is repeated in the data (e.g. in a list) by setting `once`
+            the attributes are included only in the first occurrence of an XML element relative to the same field.
 
     Returns:
         XML representation of the object.
@@ -51,7 +58,12 @@ def format_as_xml(
     '''
     ```
     """
-    el = _ToXml(item_tag=item_tag, none_str=none_str).to_xml(obj, root_tag)
+    el = _ToXml(
+        data=obj,
+        item_tag=item_tag,
+        none_str=none_str,
+        include_field_info=include_field_info,
+    ).to_xml(root_tag)
     if root_tag is None and el.text is None:
         join = '' if indent is None else '\n'
         return join.join(_rootless_xml_elements(el, indent))
@@ -63,11 +75,26 @@ def format_as_xml(
 
 @dataclass
 class _ToXml:
+    data: Any
     item_tag: str
     none_str: str
-
-    def to_xml(self, value: Any, tag: str | None) -> ElementTree.Element:
-        element = ElementTree.Element(self.item_tag if tag is None else tag)
+    include_field_info: Literal['once'] | bool
+    # a map of Pydantic and dataclasses Field paths to their metadata:
+    # a field unique string representation and its class
+    _fields_info: dict[str, tuple[str, FieldInfo | ComputedFieldInfo]] = field(default_factory=dict)
+    # keep track of fields we have extracted attributes from
+    _included_fields: set[str] = field(default_factory=set)
+    # keep track of class names for dataclasses and Pydantic models, that occur in lists
+    _element_names: dict[str, str] = field(default_factory=dict)
+    # flag for parsing dataclasses and Pydantic models once
+    _is_info_extracted: bool = False
+    _FIELD_ATTRIBUTES = ('title', 'description')
+
+    def to_xml(self, tag: str | None = None) -> ElementTree.Element:
+        return self._to_xml(value=self.data, path='', tag=tag)
+
+    def _to_xml(self, value: Any, path: str, tag: str | None = None) -> ElementTree.Element:
+        element = self._create_element(self.item_tag if tag is None else tag, path)
         if value is None:
             element.text = self.none_str
         elif isinstance(value, str):
@@ -79,31 +106,96 @@ class _ToXml:
         elif isinstance(value, date):
             element.text = value.isoformat()
         elif isinstance(value, Mapping):
-            self._mapping_to_xml(element, value)  # pyright: ignore[reportUnknownArgumentType]
+            if tag is None and path in self._element_names:
+                element.tag = self._element_names[path]
+            self._mapping_to_xml(element, value, path)  # pyright: ignore[reportUnknownArgumentType]
         elif is_dataclass(value) and not isinstance(value, type):
+            self._init_structure_info()
             if tag is None:
-                element = ElementTree.Element(value.__class__.__name__)
-            dc_dict = asdict(value)
-            self._mapping_to_xml(element, dc_dict)
+                element.tag = value.__class__.__name__
+            self._mapping_to_xml(element, asdict(value), path)
         elif isinstance(value, BaseModel):
+            self._init_structure_info()
             if tag is None:
-                element = ElementTree.Element(value.__class__.__name__)
-            self._mapping_to_xml(element, value.model_dump(mode='python'))
+                element.tag = value.__class__.__name__
+            # by dumping the model we loose all metadata in nested data structures,
+            # but we have collected it when called _init_structure_info
+            self._mapping_to_xml(element, value.model_dump(), path)
         elif isinstance(value, Iterable):
-            for item in value:  # pyright: ignore[reportUnknownVariableType]
-                item_el = self.to_xml(item, None)
-                element.append(item_el)
+            for n, item in enumerate(value):  # pyright: ignore[reportUnknownVariableType,reportUnknownArgumentType]
+                element.append(self._to_xml(value=item, path=f'{path}.[{n}]' if path else f'[{n}]'))
         else:
             raise TypeError(f'Unsupported type for XML formatting: {type(value)}')
         return element
 
-    def _mapping_to_xml(self, element: ElementTree.Element, mapping: Mapping[Any, Any]) -> None:
+    def _create_element(self, tag: str, path: str) -> ElementTree.Element:
+        element = ElementTree.Element(tag)
+        if path in self._fields_info:
+            field_repr, field_info = self._fields_info[path]
+            if self.include_field_info and self.include_field_info != 'once' or field_repr not in self._included_fields:
+                field_attributes = self._extract_attributes(field_info)
+                for k, v in field_attributes.items():
+                    element.set(k, v)
+                self._included_fields.add(field_repr)
+        return element
+
+    def _init_structure_info(self):
+        """Create maps with all data information (fields info and class names), if not already created."""
+        if not self._is_info_extracted:
+            self._parse_data_structures(self.data)
+            self._is_info_extracted = True
+
+    def _mapping_to_xml(
+        self,
+        element: ElementTree.Element,
+        mapping: Mapping[Any, Any],
+        path: str = '',
+    ) -> None:
         for key, value in mapping.items():
             if isinstance(key, int):
                 key = str(key)
             elif not isinstance(key, str):
                 raise TypeError(f'Unsupported key type for XML formatting: {type(key)}, only str and int are allowed')
-            element.append(self.to_xml(value, key))
+            element.append(self._to_xml(value=value, path=f'{path}.{key}' if path else key, tag=key))
+
+    def _parse_data_structures(
+        self,
+        value: Any,
+        path: str = '',
+    ):
+        """Parse data structures as dataclasses or Pydantic models to extract element names and attributes."""
+        if value is None or isinstance(value, (str | int | float | date | bytearray | bytes | bool)):
+            return
+        elif isinstance(value, Mapping):
+            for k, v in value.items():  # pyright: ignore[reportUnknownVariableType]
+                self._parse_data_structures(v, f'{path}.{k}' if path else f'{k}')
+        elif is_dataclass(value) and not isinstance(value, type):
+            self._element_names[path] = value.__class__.__name__
+            for field in fields(value):
+                new_path = f'{path}.{field.name}' if path else field.name
+                if self.include_field_info and field.metadata:
+                    attributes = {k: v for k, v in field.metadata.items() if k in self._FIELD_ATTRIBUTES}
+                    if attributes:
+                        field_repr = f'{value.__class__.__name__}.{field.name}'
+                        self._fields_info[new_path] = (field_repr, FieldInfo(**attributes))
+                self._parse_data_structures(getattr(value, field.name), new_path)
+        elif isinstance(value, BaseModel):
+            self._element_names[path] = value.__class__.__name__
+            for model_fields in (value.__class__.model_fields, value.__class__.model_computed_fields):
+                for field, info in model_fields.items():
+                    new_path = f'{path}.{field}' if path else field
+                    if self.include_field_info and (isinstance(info, ComputedFieldInfo) or not info.exclude):
+                        field_repr = f'{value.__class__.__name__}.{field}'
+                        self._fields_info[new_path] = (field_repr, info)
+                    self._parse_data_structures(getattr(value, field), new_path)
+        elif isinstance(value, Iterable):
+            for n, item in enumerate(value):  # pyright: ignore[reportUnknownVariableType,reportUnknownArgumentType]
+                new_path = f'{path}.[{n}]' if path else f'[{n}]'
+                self._parse_data_structures(item, new_path)
+
+    @classmethod
+    def _extract_attributes(cls, info: FieldInfo | ComputedFieldInfo) -> dict[str, str]:
+        return {attr: str(value) for attr in cls._FIELD_ATTRIBUTES if (value := getattr(info, attr, None)) is not None}
 
 
 def _rootless_xml_elements(root: ElementTree.Element, indent: str | None) -> Iterator[str]:

pydantic_ai/messages.py

@@ -668,8 +668,11 @@ class BaseToolReturnPart:
     content: Any
     """The return value."""
 
-    tool_call_id: str
-    """The tool call identifier, this is used by some models including OpenAI."""
+    tool_call_id: str = field(default_factory=_generate_tool_call_id)
+    """The tool call identifier, this is used by some models including OpenAI.
+
+    In case the tool call id is not provided by the model, Pydantic AI will generate a random one.
+    """
 
     _: KW_ONLY
 
@@ -708,14 +711,16 @@ class BaseToolReturnPart:
     def otel_message_parts(self, settings: InstrumentationSettings) -> list[_otel_messages.MessagePart]:
         from .models.instrumented import InstrumentedModel
 
-        return [
-            _otel_messages.ToolCallResponsePart(
-                type='tool_call_response',
-                id=self.tool_call_id,
-                name=self.tool_name,
-                **({'result': InstrumentedModel.serialize_any(self.content)} if settings.include_content else {}),
-            )
-        ]
+        part = _otel_messages.ToolCallResponsePart(
+            type='tool_call_response',
+            id=self.tool_call_id,
+            name=self.tool_name,
+        )
+
+        if settings.include_content and self.content is not None:
+            part['result'] = InstrumentedModel.serialize_any(self.content)
+
+        return [part]
 
     def has_content(self) -> bool:
         """Return `True` if the tool return has content."""
@@ -820,14 +825,16 @@ class RetryPromptPart:
         if self.tool_name is None:
             return [_otel_messages.TextPart(type='text', content=self.model_response())]
         else:
-            return [
-                _otel_messages.ToolCallResponsePart(
-                    type='tool_call_response',
-                    id=self.tool_call_id,
-                    name=self.tool_name,
-                    **({'result': self.model_response()} if settings.include_content else {}),
-                )
-            ]
+            part = _otel_messages.ToolCallResponsePart(
+                type='tool_call_response',
+                id=self.tool_call_id,
+                name=self.tool_name,
+            )
+
+            if settings.include_content:
+                part['result'] = self.model_response()
+
+            return [part]
 
     __repr__ = _utils.dataclasses_no_defaults_repr
 
@@ -1131,8 +1138,10 @@ class ModelResponse:
                         **({'content': part.content} if settings.include_content else {}),
                     )
                 )
-            elif isinstance(part, ToolCallPart):
+            elif isinstance(part, BaseToolCallPart):
                 call_part = _otel_messages.ToolCallPart(type='tool_call', id=part.tool_call_id, name=part.tool_name)
+                if isinstance(part, BuiltinToolCallPart):
+                    call_part['builtin'] = True
                 if settings.include_content and part.args is not None:
                     from .models.instrumented import InstrumentedModel
 
@@ -1142,6 +1151,19 @@ class ModelResponse:
                         call_part['arguments'] = {k: InstrumentedModel.serialize_any(v) for k, v in part.args.items()}
 
                 parts.append(call_part)
+            elif isinstance(part, BuiltinToolReturnPart):
+                return_part = _otel_messages.ToolCallResponsePart(
+                    type='tool_call_response',
+                    id=part.tool_call_id,
+                    name=part.tool_name,
+                    builtin=True,
+                )
+                if settings.include_content and part.content is not None:  # pragma: no branch
+                    from .models.instrumented import InstrumentedModel
+
+                    return_part['result'] = InstrumentedModel.serialize_any(part.content)
+
+                parts.append(return_part)
         return parts
 
     @property
@@ -1299,35 +1321,39 @@ class ToolCallPartDelta:
         return ToolCallPart(self.tool_name_delta, self.args_delta, self.tool_call_id or _generate_tool_call_id())
 
     @overload
-    def apply(self, part: ModelResponsePart) -> ToolCallPart: ...
+    def apply(self, part: ModelResponsePart) -> ToolCallPart | BuiltinToolCallPart: ...
 
     @overload
-    def apply(self, part: ModelResponsePart | ToolCallPartDelta) -> ToolCallPart | ToolCallPartDelta: ...
+    def apply(
+        self, part: ModelResponsePart | ToolCallPartDelta
+    ) -> ToolCallPart | BuiltinToolCallPart | ToolCallPartDelta: ...
 
-    def apply(self, part: ModelResponsePart | ToolCallPartDelta) -> ToolCallPart | ToolCallPartDelta:
+    def apply(
+        self, part: ModelResponsePart | ToolCallPartDelta
+    ) -> ToolCallPart | BuiltinToolCallPart | ToolCallPartDelta:
         """Apply this delta to a part or delta, returning a new part or delta with the changes applied.
 
         Args:
             part: The existing model response part or delta to update.
 
         Returns:
-            Either a new `ToolCallPart` or an updated `ToolCallPartDelta`.
+            Either a new `ToolCallPart` or `BuiltinToolCallPart`, or an updated `ToolCallPartDelta`.
 
         Raises:
-            ValueError: If `part` is neither a `ToolCallPart` nor a `ToolCallPartDelta`.
+            ValueError: If `part` is neither a `ToolCallPart`, `BuiltinToolCallPart`, nor a `ToolCallPartDelta`.
             UnexpectedModelBehavior: If applying JSON deltas to dict arguments or vice versa.
         """
-        if isinstance(part, ToolCallPart):
+        if isinstance(part, ToolCallPart | BuiltinToolCallPart):
             return self._apply_to_part(part)
 
         if isinstance(part, ToolCallPartDelta):
             return self._apply_to_delta(part)
 
         raise ValueError(  # pragma: no cover
-            f'Can only apply ToolCallPartDeltas to ToolCallParts or ToolCallPartDeltas, not {part}'
+            f'Can only apply ToolCallPartDeltas to ToolCallParts, BuiltinToolCallParts, or ToolCallPartDeltas, not {part}'
         )
 
-    def _apply_to_delta(self, delta: ToolCallPartDelta) -> ToolCallPart | ToolCallPartDelta:
+    def _apply_to_delta(self, delta: ToolCallPartDelta) -> ToolCallPart | BuiltinToolCallPart | ToolCallPartDelta:
         """Internal helper to apply this delta to another delta."""
         if self.tool_name_delta:
             # Append incremental text to the existing tool_name_delta
@@ -1358,8 +1384,8 @@ class ToolCallPartDelta:
 
         return delta
 
-    def _apply_to_part(self, part: ToolCallPart) -> ToolCallPart:
-        """Internal helper to apply this delta directly to a `ToolCallPart`."""
+    def _apply_to_part(self, part: ToolCallPart | BuiltinToolCallPart) -> ToolCallPart | BuiltinToolCallPart:
+        """Internal helper to apply this delta directly to a `ToolCallPart` or `BuiltinToolCallPart`."""
         if self.tool_name_delta:
             # Append incremental text to the existing tool_name
             tool_name = part.tool_name + self.tool_name_delta
@@ -1491,6 +1517,9 @@ class FunctionToolResultEvent:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
+@deprecated(
+    '`BuiltinToolCallEvent` is deprecated, look for `PartStartEvent` and `PartDeltaEvent` with `BuiltinToolCallPart` instead.'
+)
 @dataclass(repr=False)
 class BuiltinToolCallEvent:
     """An event indicating the start to a call to a built-in tool."""
@@ -1504,6 +1533,9 @@ class BuiltinToolCallEvent:
     """Event type identifier, used as a discriminator."""
 
 
+@deprecated(
+    '`BuiltinToolResultEvent` is deprecated, look for `PartStartEvent` and `PartDeltaEvent` with `BuiltinToolReturnPart` instead.'
+)
 @dataclass(repr=False)
 class BuiltinToolResultEvent:
     """An event indicating the result of a built-in tool call."""
@@ -1518,7 +1550,10 @@ class BuiltinToolResultEvent:
 
 
 HandleResponseEvent = Annotated[
-    FunctionToolCallEvent | FunctionToolResultEvent | BuiltinToolCallEvent | BuiltinToolResultEvent,
+    FunctionToolCallEvent
+    | FunctionToolResultEvent
+    | BuiltinToolCallEvent  # pyright: ignore[reportDeprecated]
+    | BuiltinToolResultEvent,  # pyright: ignore[reportDeprecated]
     pydantic.Discriminator('event_kind'),
 ]
 """An event yielded when handling a model response, indicating tool calls and results."""