pydantic-ai-slim 0.0.55__py3-none-any.whl → 0.1.0__py3-none-any.whl

pydantic_ai/format_as_xml.py

@@ -1,116 +1,9 @@
-from __future__ import annotations as _annotations
+from typing_extensions import deprecated
 
-from collections.abc import Iterable, Iterator, Mapping
-from dataclasses import asdict, dataclass, is_dataclass
-from datetime import date
-from typing import Any
-from xml.etree import ElementTree
+from .format_prompt import format_as_xml as _format_as_xml
 
-from pydantic import BaseModel
 
-__all__ = ('format_as_xml',)
-
-
-def format_as_xml(
-    obj: Any,
-    root_tag: str = 'examples',
-    item_tag: str = 'example',
-    include_root_tag: bool = True,
-    none_str: str = 'null',
-    indent: str | None = '  ',
-) -> str:
-    """Format a Python object as XML.
-
-    This is useful since LLMs often find it easier to read semi-structured data (e.g. examples) as XML,
-    rather than JSON etc.
-
-    Supports: `str`, `bytes`, `bytearray`, `bool`, `int`, `float`, `date`, `datetime`, `Mapping`,
-    `Iterable`, `dataclass`, and `BaseModel`.
-
-    Args:
-        obj: Python Object to serialize to XML.
-        root_tag: Outer tag to wrap the XML in, use `None` to omit the outer tag.
-        item_tag: Tag to use for each item in an iterable (e.g. list), this is overridden by the class name
-            for dataclasses and Pydantic models.
-        include_root_tag: Whether to include the root tag in the output
-            (The root tag is always included if it includes a body - e.g. when the input is a simple value).
-        none_str: String to use for `None` values.
-        indent: Indentation string to use for pretty printing.
-
-    Returns:
-        XML representation of the object.
-
-    Example:
-    ```python {title="format_as_xml_example.py" lint="skip"}
-    from pydantic_ai.format_as_xml import format_as_xml
-
-    print(format_as_xml({'name': 'John', 'height': 6, 'weight': 200}, root_tag='user'))
-    '''
-    <user>
-      <name>John</name>
-      <height>6</height>
-      <weight>200</weight>
-    </user>
-    '''
-    ```
-    """
-    el = _ToXml(item_tag=item_tag, none_str=none_str).to_xml(obj, root_tag)
-    if not include_root_tag and el.text is None:
-        join = '' if indent is None else '\n'
-        return join.join(_rootless_xml_elements(el, indent))
-    else:
-        if indent is not None:
-            ElementTree.indent(el, space=indent)
-        return ElementTree.tostring(el, encoding='unicode')
-
-
-@dataclass
-class _ToXml:
-    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)
-        if value is None:
-            element.text = self.none_str
-        elif isinstance(value, str):
-            element.text = value
-        elif isinstance(value, (bytes, bytearray)):
-            element.text = value.decode(errors='ignore')
-        elif isinstance(value, (bool, int, float)):
-            element.text = str(value)
-        elif isinstance(value, date):
-            element.text = value.isoformat()
-        elif isinstance(value, Mapping):
-            self._mapping_to_xml(element, value)  # pyright: ignore[reportUnknownArgumentType]
-        elif is_dataclass(value) and not isinstance(value, type):
-            if tag is None:
-                element = ElementTree.Element(value.__class__.__name__)
-            dc_dict = asdict(value)
-            self._mapping_to_xml(element, dc_dict)
-        elif isinstance(value, BaseModel):
-            if tag is None:
-                element = ElementTree.Element(value.__class__.__name__)
-            self._mapping_to_xml(element, value.model_dump(mode='python'))
-        elif isinstance(value, Iterable):
-            for item in value:  # pyright: ignore[reportUnknownVariableType]
-                item_el = self.to_xml(item, None)
-                element.append(item_el)
-        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:
-        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))
-
-
-def _rootless_xml_elements(root: ElementTree.Element, indent: str | None) -> Iterator[str]:
-    for sub_element in root:
-        if indent is not None:
-            ElementTree.indent(sub_element, space=indent)
-        yield ElementTree.tostring(sub_element, encoding='unicode')
+@deprecated('`format_as_xml` has moved, import it via `from pydantic_ai import format_as_xml`')
+def format_as_xml(prompt: str) -> str:
+    """`format_as_xml` has moved, import it via `from pydantic_ai import format_as_xml` instead."""
+    return _format_as_xml(prompt)

pydantic_ai/format_prompt.py

@@ -0,0 +1,116 @@
+from __future__ import annotations as _annotations
+
+from collections.abc import Iterable, Iterator, Mapping
+from dataclasses import asdict, dataclass, is_dataclass
+from datetime import date
+from typing import Any
+from xml.etree import ElementTree
+
+from pydantic import BaseModel
+
+__all__ = ('format_as_xml',)
+
+
+def format_as_xml(
+    obj: Any,
+    root_tag: str = 'examples',
+    item_tag: str = 'example',
+    include_root_tag: bool = True,
+    none_str: str = 'null',
+    indent: str | None = '  ',
+) -> str:
+    """Format a Python object as XML.
+
+    This is useful since LLMs often find it easier to read semi-structured data (e.g. examples) as XML,
+    rather than JSON etc.
+
+    Supports: `str`, `bytes`, `bytearray`, `bool`, `int`, `float`, `date`, `datetime`, `Mapping`,
+    `Iterable`, `dataclass`, and `BaseModel`.
+
+    Args:
+        obj: Python Object to serialize to XML.
+        root_tag: Outer tag to wrap the XML in, use `None` to omit the outer tag.
+        item_tag: Tag to use for each item in an iterable (e.g. list), this is overridden by the class name
+            for dataclasses and Pydantic models.
+        include_root_tag: Whether to include the root tag in the output
+            (The root tag is always included if it includes a body - e.g. when the input is a simple value).
+        none_str: String to use for `None` values.
+        indent: Indentation string to use for pretty printing.
+
+    Returns:
+        XML representation of the object.
+
+    Example:
+    ```python {title="format_as_xml_example.py" lint="skip"}
+    from pydantic_ai import format_as_xml
+
+    print(format_as_xml({'name': 'John', 'height': 6, 'weight': 200}, root_tag='user'))
+    '''
+    <user>
+      <name>John</name>
+      <height>6</height>
+      <weight>200</weight>
+    </user>
+    '''
+    ```
+    """
+    el = _ToXml(item_tag=item_tag, none_str=none_str).to_xml(obj, root_tag)
+    if not include_root_tag and el.text is None:
+        join = '' if indent is None else '\n'
+        return join.join(_rootless_xml_elements(el, indent))
+    else:
+        if indent is not None:
+            ElementTree.indent(el, space=indent)
+        return ElementTree.tostring(el, encoding='unicode')
+
+
+@dataclass
+class _ToXml:
+    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)
+        if value is None:
+            element.text = self.none_str
+        elif isinstance(value, str):
+            element.text = value
+        elif isinstance(value, (bytes, bytearray)):
+            element.text = value.decode(errors='ignore')
+        elif isinstance(value, (bool, int, float)):
+            element.text = str(value)
+        elif isinstance(value, date):
+            element.text = value.isoformat()
+        elif isinstance(value, Mapping):
+            self._mapping_to_xml(element, value)  # pyright: ignore[reportUnknownArgumentType]
+        elif is_dataclass(value) and not isinstance(value, type):
+            if tag is None:
+                element = ElementTree.Element(value.__class__.__name__)
+            dc_dict = asdict(value)
+            self._mapping_to_xml(element, dc_dict)
+        elif isinstance(value, BaseModel):
+            if tag is None:
+                element = ElementTree.Element(value.__class__.__name__)
+            self._mapping_to_xml(element, value.model_dump(mode='python'))
+        elif isinstance(value, Iterable):
+            for item in value:  # pyright: ignore[reportUnknownVariableType]
+                item_el = self.to_xml(item, None)
+                element.append(item_el)
+        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:
+        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))
+
+
+def _rootless_xml_elements(root: ElementTree.Element, indent: str | None) -> Iterator[str]:
+    for sub_element in root:
+        if indent is not None:
+            ElementTree.indent(sub_element, space=indent)
+        yield ElementTree.tostring(sub_element, encoding='unicode')

pydantic_ai/messages.py

@@ -15,6 +15,34 @@ from typing_extensions import TypeAlias
 from ._utils import generate_tool_call_id as _generate_tool_call_id, now_utc as _now_utc
 from .exceptions import UnexpectedModelBehavior
 
+AudioMediaType: TypeAlias = Literal['audio/wav', 'audio/mpeg']
+ImageMediaType: TypeAlias = Literal['image/jpeg', 'image/png', 'image/gif', 'image/webp']
+DocumentMediaType: TypeAlias = Literal[
+    'application/pdf',
+    'text/plain',
+    'text/csv',
+    'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+    'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
+    'text/html',
+    'text/markdown',
+    'application/vnd.ms-excel',
+]
+VideoMediaType: TypeAlias = Literal[
+    'video/x-matroska',
+    'video/quicktime',
+    'video/mp4',
+    'video/webm',
+    'video/x-flv',
+    'video/mpeg',
+    'video/x-ms-wmv',
+    'video/3gpp',
+]
+
+AudioFormat: TypeAlias = Literal['wav', 'mp3']
+ImageFormat: TypeAlias = Literal['jpeg', 'png', 'gif', 'webp']
+DocumentFormat: TypeAlias = Literal['csv', 'doc', 'docx', 'html', 'md', 'pdf', 'txt', 'xls', 'xlsx']
+VideoFormat: TypeAlias = Literal['mkv', 'mov', 'mp4', 'webm', 'flv', 'mpeg', 'mpg', 'wmv', 'three_gp']
+
 
 @dataclass
 class SystemPromptPart:
@@ -42,6 +70,47 @@ class SystemPromptPart:
         return Event('gen_ai.system.message', body={'content': self.content, 'role': 'system'})
 
 
+@dataclass
+class VideoUrl:
+    """A URL to an video."""
+
+    url: str
+    """The URL of the video."""
+
+    kind: Literal['video-url'] = 'video-url'
+    """Type identifier, this is available on all parts as a discriminator."""
+
+    @property
+    def media_type(self) -> VideoMediaType:  # pragma: no cover
+        """Return the media type of the video, based on the url."""
+        if self.url.endswith('.mkv'):
+            return 'video/x-matroska'
+        elif self.url.endswith('.mov'):
+            return 'video/quicktime'
+        elif self.url.endswith('.mp4'):
+            return 'video/mp4'
+        elif self.url.endswith('.webm'):
+            return 'video/webm'
+        elif self.url.endswith('.flv'):
+            return 'video/x-flv'
+        elif self.url.endswith(('.mpeg', '.mpg')):
+            return 'video/mpeg'
+        elif self.url.endswith('.wmv'):
+            return 'video/x-ms-wmv'
+        elif self.url.endswith('.three_gp'):
+            return 'video/3gpp'
+        else:
+            raise ValueError(f'Unknown video file extension: {self.url}')
+
+    @property
+    def format(self) -> VideoFormat:
+        """The file format of the video.
+
+        The choice of supported formats were based on the Bedrock Converse API. Other APIs don't require to use a format.
+        """
+        return _video_format(self.media_type)
+
+
 @dataclass
 class AudioUrl:
     """A URL to an audio file."""
@@ -123,23 +192,6 @@ class DocumentUrl:
         return _document_format(self.media_type)
 
 
-AudioMediaType: TypeAlias = Literal['audio/wav', 'audio/mpeg']
-ImageMediaType: TypeAlias = Literal['image/jpeg', 'image/png', 'image/gif', 'image/webp']
-DocumentMediaType: TypeAlias = Literal[
-    'application/pdf',
-    'text/plain',
-    'text/csv',
-    'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
-    'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
-    'text/html',
-    'text/markdown',
-    'application/vnd.ms-excel',
-]
-AudioFormat: TypeAlias = Literal['wav', 'mp3']
-ImageFormat: TypeAlias = Literal['jpeg', 'png', 'gif', 'webp']
-DocumentFormat: TypeAlias = Literal['csv', 'doc', 'docx', 'html', 'md', 'pdf', 'txt', 'xls', 'xlsx']
-
-
 @dataclass
 class BinaryContent:
     """Binary content, e.g. an audio or image file."""
@@ -163,6 +215,11 @@ class BinaryContent:
         """Return `True` if the media type is an image type."""
         return self.media_type.startswith('image/')
 
+    @property
+    def is_video(self) -> bool:
+        """Return `True` if the media type is a video type."""
+        return self.media_type.startswith('video/')
+
     @property
     def is_document(self) -> bool:
         """Return `True` if the media type is a document type."""
@@ -189,10 +246,12 @@ class BinaryContent:
             return _image_format(self.media_type)
         elif self.is_document:
             return _document_format(self.media_type)
+        elif self.is_video:
+            return _video_format(self.media_type)
         raise ValueError(f'Unknown media type: {self.media_type}')
 
 
-UserContent: TypeAlias = 'str | ImageUrl | AudioUrl | DocumentUrl | BinaryContent'
+UserContent: TypeAlias = 'str | ImageUrl | AudioUrl | DocumentUrl | VideoUrl | BinaryContent'
 
 
 def _document_format(media_type: str) -> DocumentFormat:
@@ -229,6 +288,27 @@ def _image_format(media_type: str) -> ImageFormat:
         raise ValueError(f'Unknown image media type: {media_type}')
 
 
+def _video_format(media_type: str) -> VideoFormat:
+    if media_type == 'video/x-matroska':
+        return 'mkv'
+    elif media_type == 'video/quicktime':
+        return 'mov'
+    elif media_type == 'video/mp4':
+        return 'mp4'
+    elif media_type == 'video/webm':
+        return 'webm'
+    elif media_type == 'video/x-flv':
+        return 'flv'
+    elif media_type == 'video/mpeg':
+        return 'mpeg'
+    elif media_type == 'video/x-ms-wmv':
+        return 'wmv'
+    elif media_type == 'video/3gpp':
+        return 'three_gp'
+    else:  # pragma: no cover
+        raise ValueError(f'Unknown video media type: {media_type}')
+
+
 @dataclass
 class UserPromptPart:
     """A user prompt, generally written by the end user.
@@ -315,7 +395,7 @@ class RetryPromptPart:
     * the model returned plain text when a structured response was expected
     * Pydantic validation of a structured response failed, here content is derived from a Pydantic
       [`ValidationError`][pydantic_core.ValidationError]
-    * a result validator raised a [`ModelRetry`][pydantic_ai.exceptions.ModelRetry] exception
+    * an output validator raised a [`ModelRetry`][pydantic_ai.exceptions.ModelRetry] exception
     """
 
     content: list[pydantic_core.ErrorDetails] | str
@@ -377,6 +457,9 @@ class ModelRequest:
     parts: list[ModelRequestPart]
     """The parts of the user message."""
 
+    instructions: str | None = None
+    """The instructions for the model."""
+
     kind: Literal['request'] = 'request'
     """Message type identifier, this is available on all parts as a discriminator."""
 
@@ -691,10 +774,10 @@ class PartDeltaEvent:
 
 @dataclass
 class FinalResultEvent:
-    """An event indicating the response to the current model request matches the result schema."""
+    """An event indicating the response to the current model request matches the output schema and will produce a result."""
 
     tool_name: str | None
-    """The name of the result tool that was called. `None` if the result is from text content and not from a tool."""
+    """The name of the output tool that was called. `None` if the result is from text content and not from a tool."""
     tool_call_id: str | None
     """The tool call ID, if any, that this result is associated with."""
     event_kind: Literal['final_result'] = 'final_result'

pydantic_ai/models/__init__.py

@@ -19,7 +19,7 @@ from typing_extensions import Literal, TypeAliasType
 
 from .._parts_manager import ModelResponsePartsManager
 from ..exceptions import UserError
-from ..messages import ModelMessage, ModelResponse, ModelResponseStreamEvent
+from ..messages import ModelMessage, ModelRequest, ModelResponse, ModelResponseStreamEvent
 from ..settings import ModelSettings
 from ..usage import Usage
 
@@ -107,6 +107,7 @@ KnownModelName = TypeAliasType(
         'google-gla:gemini-2.0-flash-lite-preview-02-05',
         'google-gla:gemini-2.0-pro-exp-02-05',
         'google-gla:gemini-2.5-pro-exp-03-25',
+        'google-gla:gemini-2.5-pro-preview-03-25',
         'google-vertex:gemini-1.0-pro',
         'google-vertex:gemini-1.5-flash',
         'google-vertex:gemini-1.5-flash-8b',
@@ -118,6 +119,7 @@ KnownModelName = TypeAliasType(
         'google-vertex:gemini-2.0-flash-lite-preview-02-05',
         'google-vertex:gemini-2.0-pro-exp-02-05',
         'google-vertex:gemini-2.5-pro-exp-03-25',
+        'google-vertex:gemini-2.5-pro-preview-03-25',
         'gpt-3.5-turbo',
         'gpt-3.5-turbo-0125',
         'gpt-3.5-turbo-0301',
@@ -137,6 +139,12 @@ KnownModelName = TypeAliasType(
         'gpt-4-turbo-2024-04-09',
         'gpt-4-turbo-preview',
         'gpt-4-vision-preview',
+        'gpt-4.1',
+        'gpt-4.1-2025-04-14',
+        'gpt-4.1-mini',
+        'gpt-4.1-mini-2025-04-14',
+        'gpt-4.1-nano',
+        'gpt-4.1-nano-2025-04-14',
         'gpt-4o',
         'gpt-4o-2024-05-13',
         'gpt-4o-2024-08-06',
@@ -206,6 +214,12 @@ KnownModelName = TypeAliasType(
         'openai:gpt-4-turbo-2024-04-09',
         'openai:gpt-4-turbo-preview',
         'openai:gpt-4-vision-preview',
+        'openai:gpt-4.1',
+        'openai:gpt-4.1-2025-04-14',
+        'openai:gpt-4.1-mini',
+        'openai:gpt-4.1-mini-2025-04-14',
+        'openai:gpt-4.1-nano',
+        'openai:gpt-4.1-nano-2025-04-14',
         'openai:gpt-4o',
         'openai:gpt-4o-2024-05-13',
         'openai:gpt-4o-2024-08-06',
@@ -240,11 +254,11 @@ KnownModelName = TypeAliasType(
 
 @dataclass
 class ModelRequestParameters:
-    """Configuration for an agent's request to a model, specifically related to tools and result handling."""
+    """Configuration for an agent's request to a model, specifically related to tools and output handling."""
 
     function_tools: list[ToolDefinition]
-    allow_text_result: bool
-    result_tools: list[ToolDefinition]
+    allow_text_output: bool
+    output_tools: list[ToolDefinition]
 
 
 class Model(ABC):
@@ -306,6 +320,12 @@ class Model(ABC):
         """The base URL for the provider API, if available."""
         return None
 
+    def _get_instructions(self, messages: list[ModelMessage]) -> str | None:
+        """Get instructions from the first ModelRequest found when iterating messages in reverse."""
+        for message in reversed(messages):
+            if isinstance(message, ModelRequest):
+                return message.instructions
+
 
 @dataclass
 class StreamedResponse(ABC):

pydantic_ai/models/_json_schema.py

@@ -0,0 +1,156 @@
+import re
+from abc import ABC, abstractmethod
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from pydantic_ai.exceptions import UserError
+
+JsonSchema = dict[str, Any]
+
+
+@dataclass(init=False)
+class WalkJsonSchema(ABC):
+    """Walks a JSON schema, applying transformations to it at each level.
+
+    Note: We may eventually want to rework tools to build the JSON schema from the type directly, using a subclass of
+    pydantic.json_schema.GenerateJsonSchema, rather than making use of this machinery.
+    """
+
+    def __init__(
+        self, schema: JsonSchema, *, prefer_inlined_defs: bool = False, simplify_nullable_unions: bool = False
+    ):
+        self.schema = deepcopy(schema)
+        self.prefer_inlined_defs = prefer_inlined_defs
+        self.simplify_nullable_unions = simplify_nullable_unions
+
+        self.defs: dict[str, JsonSchema] = self.schema.pop('$defs', {})
+        self.refs_stack = tuple[str, ...]()
+        self.recursive_refs = set[str]()
+
+    @abstractmethod
+    def transform(self, schema: JsonSchema) -> JsonSchema:
+        """Make changes to the schema."""
+        return schema
+
+    def walk(self) -> JsonSchema:
+        handled = self._handle(deepcopy(self.schema))
+
+        if not self.prefer_inlined_defs and self.defs:
+            handled['$defs'] = {k: self._handle(v) for k, v in self.defs.items()}
+
+        elif self.recursive_refs:  # pragma: no cover
+            # If we are preferring inlined defs and there are recursive refs, we _have_ to use a $defs+$ref structure
+            # We try to use whatever the original root key was, but if it is already in use,
+            # we modify it to avoid collisions.
+            defs = {key: self.defs[key] for key in self.recursive_refs}
+            root_ref = self.schema.get('$ref')
+            root_key = None if root_ref is None else re.sub(r'^#/\$defs/', '', root_ref)
+            if root_key is None:
+                root_key = self.schema.get('title', 'root')
+                while root_key in defs:
+                    # Modify the root key until it is not already in use
+                    root_key = f'{root_key}_root'
+
+            defs[root_key] = handled
+            return {'$defs': defs, '$ref': f'#/$defs/{root_key}'}
+
+        return handled
+
+    def _handle(self, schema: JsonSchema) -> JsonSchema:
+        if self.prefer_inlined_defs:
+            while ref := schema.get('$ref'):
+                key = re.sub(r'^#/\$defs/', '', ref)
+                if key in self.refs_stack:
+                    self.recursive_refs.add(key)
+                    break  # recursive ref can't be unpacked
+                self.refs_stack += (key,)
+                def_schema = self.defs.get(key)
+                if def_schema is None:  # pragma: no cover
+                    raise UserError(f'Could not find $ref definition for {key}')
+                schema = def_schema
+
+        # Handle the schema based on its type / structure
+        type_ = schema.get('type')
+        if type_ == 'object':
+            schema = self._handle_object(schema)
+        elif type_ == 'array':
+            schema = self._handle_array(schema)
+        elif type_ is None:
+            schema = self._handle_union(schema, 'anyOf')
+            schema = self._handle_union(schema, 'oneOf')
+
+        # Apply the base transform
+        schema = self.transform(schema)
+
+        return schema
+
+    def _handle_object(self, schema: JsonSchema) -> JsonSchema:
+        if properties := schema.get('properties'):
+            handled_properties = {}
+            for key, value in properties.items():
+                handled_properties[key] = self._handle(value)
+            schema['properties'] = handled_properties
+
+        if (additional_properties := schema.get('additionalProperties')) is not None:
+            if isinstance(additional_properties, bool):
+                schema['additionalProperties'] = additional_properties
+            else:  # pragma: no cover
+                schema['additionalProperties'] = self._handle(additional_properties)
+
+        if (pattern_properties := schema.get('patternProperties')) is not None:
+            handled_pattern_properties = {}
+            for key, value in pattern_properties.items():
+                handled_pattern_properties[key] = self._handle(value)
+            schema['patternProperties'] = handled_pattern_properties
+
+        return schema
+
+    def _handle_array(self, schema: JsonSchema) -> JsonSchema:
+        if prefix_items := schema.get('prefixItems'):
+            schema['prefixItems'] = [self._handle(item) for item in prefix_items]
+
+        if items := schema.get('items'):
+            schema['items'] = self._handle(items)
+
+        return schema
+
+    def _handle_union(self, schema: JsonSchema, union_kind: Literal['anyOf', 'oneOf']) -> JsonSchema:
+        members = schema.get(union_kind)
+        if not members:
+            return schema
+
+        handled = [self._handle(member) for member in members]
+
+        # convert nullable unions to nullable types
+        if self.simplify_nullable_unions:
+            handled = self._simplify_nullable_union(handled)
+
+        if len(handled) == 1:
+            # In this case, no need to retain the union
+            return handled[0]
+
+        # If we have keys besides the union kind (such as title or discriminator), keep them without modifications
+        schema = schema.copy()
+        schema[union_kind] = handled
+        return schema
+
+    @staticmethod
+    def _simplify_nullable_union(cases: list[JsonSchema]) -> list[JsonSchema]:
+        # TODO: Should we move this to relevant subclasses? Or is it worth keeping here to make reuse easier?
+        if len(cases) == 2 and {'type': 'null'} in cases:
+            # Find the non-null schema
+            non_null_schema = next(
+                (item for item in cases if item != {'type': 'null'}),
+                None,
+            )
+            if non_null_schema:
+                # Create a new schema based on the non-null part, mark as nullable
+                new_schema = deepcopy(non_null_schema)
+                new_schema['nullable'] = True
+                return [new_schema]
+            else:  # pragma: no cover
+                # they are both null, so just return one of them
+                return [cases[0]]
+
+        return cases  # pragma: no cover