pydantic-ai-slim 0.7.5__tar.gz → 0.8.0__tar.gz

{pydantic_ai_slim-0.7.5 → pydantic_ai_slim-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.7.5
+Version: 0.8.0
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Project-URL: Homepage, https://github.com/pydantic/pydantic-ai/tree/main/pydantic_ai_slim
 Project-URL: Source, https://github.com/pydantic/pydantic-ai/tree/main/pydantic_ai_slim
@@ -35,7 +35,7 @@ Requires-Dist: genai-prices>=0.0.22
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.7.5
+Requires-Dist: pydantic-graph==0.8.0
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -57,7 +57,7 @@ Requires-Dist: cohere>=5.16.0; (platform_system != 'Emscripten') and extra == 'c
 Provides-Extra: duckduckgo
 Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.7.5; extra == 'evals'
+Requires-Dist: pydantic-evals==0.8.0; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.31.0; extra == 'google'
 Provides-Extra: groq
@@ -67,7 +67,7 @@ Requires-Dist: huggingface-hub[inference]>=0.33.5; extra == 'huggingface'
 Provides-Extra: logfire
 Requires-Dist: logfire>=3.14.1; extra == 'logfire'
 Provides-Extra: mcp
-Requires-Dist: mcp>=1.10.0; (python_version >= '3.10') and extra == 'mcp'
+Requires-Dist: mcp>=1.12.3; (python_version >= '3.10') and extra == 'mcp'
 Provides-Extra: mistral
 Requires-Dist: mistralai>=1.9.2; extra == 'mistral'
 Provides-Extra: openai

{pydantic_ai_slim-0.7.5 → pydantic_ai_slim-0.8.0}/pydantic_ai/_cli.py

@@ -228,6 +228,7 @@ async def run_chat(
     prog_name: str,
     config_dir: Path | None = None,
     deps: AgentDepsT = None,
+    message_history: list[ModelMessage] | None = None,
 ) -> int:
     prompt_history_path = (config_dir or PYDANTIC_AI_HOME) / PROMPT_HISTORY_FILENAME
     prompt_history_path.parent.mkdir(parents=True, exist_ok=True)
@@ -235,7 +236,7 @@ async def run_chat(
     session: PromptSession[Any] = PromptSession(history=FileHistory(str(prompt_history_path)))
 
     multiline = False
-    messages: list[ModelMessage] = []
+    messages: list[ModelMessage] = message_history[:] if message_history else []
 
     while True:
         try:

{pydantic_ai_slim-0.7.5 → pydantic_ai_slim-0.8.0}/pydantic_ai/ag_ui.py

@@ -28,11 +28,11 @@ from ._agent_graph import CallToolsNode, ModelRequestNode
 from .agent import AbstractAgent, AgentRun
 from .exceptions import UserError
 from .messages import (
-    AgentStreamEvent,
     FunctionToolResultEvent,
     ModelMessage,
     ModelRequest,
     ModelResponse,
+    ModelResponseStreamEvent,
     PartDeltaEvent,
     PartStartEvent,
     SystemPromptPart,
@@ -403,7 +403,7 @@ async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEve
 
 async def _handle_model_request_event(
     stream_ctx: _RequestStreamContext,
-    agent_event: AgentStreamEvent,
+    agent_event: ModelResponseStreamEvent,
 ) -> AsyncIterator[BaseEvent]:
     """Handle an agent event and yield AG-UI protocol events.
 

{pydantic_ai_slim-0.7.5 → pydantic_ai_slim-0.8.0}/pydantic_ai/agent/__init__.py

@@ -26,7 +26,14 @@ from .. import (
     models,
     usage as _usage,
 )
-from .._agent_graph import HistoryProcessor
+from .._agent_graph import (
+    CallToolsNode,
+    EndStrategy,
+    HistoryProcessor,
+    ModelRequestNode,
+    UserPromptNode,
+    capture_run_messages,
+)
 from .._output import OutputToolset
 from .._tool_manager import ToolManager
 from ..builtin_tools import AbstractBuiltinTool
@@ -60,13 +67,6 @@ from ..toolsets.prepared import PreparedToolset
 from .abstract import AbstractAgent, EventStreamHandler, RunOutputDataT
 from .wrapper import WrapperAgent
 
-# Re-exporting like this improves auto-import behavior in PyCharm
-capture_run_messages = _agent_graph.capture_run_messages
-EndStrategy = _agent_graph.EndStrategy
-CallToolsNode = _agent_graph.CallToolsNode
-ModelRequestNode = _agent_graph.ModelRequestNode
-UserPromptNode = _agent_graph.UserPromptNode
-
 if TYPE_CHECKING:
     from ..mcp import MCPServer
 
@@ -677,16 +677,23 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
     def _run_span_end_attributes(
         self, state: _agent_graph.GraphAgentState, usage: _usage.RunUsage, settings: InstrumentationSettings
     ):
+        if settings.version == 1:
+            attr_name = 'all_messages_events'
+            value = [
+                InstrumentedModel.event_to_dict(e) for e in settings.messages_to_otel_events(state.message_history)
+            ]
+        else:
+            attr_name = 'pydantic_ai.all_messages'
+            value = settings.messages_to_otel_messages(state.message_history)
+
         return {
             **usage.opentelemetry_attributes(),
-            'all_messages_events': json.dumps(
-                [InstrumentedModel.event_to_dict(e) for e in settings.messages_to_otel_events(state.message_history)]
-            ),
+            attr_name: json.dumps(value),
             'logfire.json_schema': json.dumps(
                 {
                     'type': 'object',
                     'properties': {
-                        'all_messages_events': {'type': 'array'},
+                        attr_name: {'type': 'array'},
                         'final_result': {'type': 'object'},
                     },
                 }

{pydantic_ai_slim-0.7.5 → pydantic_ai_slim-0.8.0}/pydantic_ai/agent/abstract.py

@@ -5,7 +5,7 @@ from abc import ABC, abstractmethod
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterator, Mapping, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from types import FrameType
-from typing import TYPE_CHECKING, Any, Callable, Generic, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, Generic, cast, overload
 
 from typing_extensions import Self, TypeAlias, TypeIs, TypeVar
 
@@ -34,13 +34,6 @@ from ..tools import (
 from ..toolsets import AbstractToolset
 from ..usage import RunUsage, UsageLimits
 
-# Re-exporting like this improves auto-import behavior in PyCharm
-capture_run_messages = _agent_graph.capture_run_messages
-EndStrategy = _agent_graph.EndStrategy
-CallToolsNode = _agent_graph.CallToolsNode
-ModelRequestNode = _agent_graph.ModelRequestNode
-UserPromptNode = _agent_graph.UserPromptNode
-
 if TYPE_CHECKING:
     from fasta2a.applications import FastA2A
     from fasta2a.broker import Broker
@@ -60,11 +53,7 @@ RunOutputDataT = TypeVar('RunOutputDataT')
 """Type variable for the result data of a run where `output_type` was customized on the run call."""
 
 EventStreamHandler: TypeAlias = Callable[
-    [
-        RunContext[AgentDepsT],
-        AsyncIterable[Union[_messages.AgentStreamEvent, _messages.HandleResponseEvent]],
-    ],
-    Awaitable[None],
+    [RunContext[AgentDepsT], AsyncIterable[_messages.AgentStreamEvent]], Awaitable[None]
 ]
 """A function that receives agent [`RunContext`][pydantic_ai.tools.RunContext] and an async iterable of events from the model's streaming response and the agent's execution of tools."""
 
@@ -452,7 +441,9 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
                     async with node.stream(graph_ctx) as stream:
                         final_result_event = None
 
-                        async def stream_to_final(stream: AgentStream) -> AsyncIterator[_messages.AgentStreamEvent]:
+                        async def stream_to_final(
+                            stream: AgentStream,
+                        ) -> AsyncIterator[_messages.ModelResponseStreamEvent]:
                             nonlocal final_result_event
                             async for event in stream:
                                 yield event
@@ -899,12 +890,18 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
             lifespan=lifespan,
         )
 
-    async def to_cli(self: Self, deps: AgentDepsT = None, prog_name: str = 'pydantic-ai') -> None:
+    async def to_cli(
+        self: Self,
+        deps: AgentDepsT = None,
+        prog_name: str = 'pydantic-ai',
+        message_history: list[_messages.ModelMessage] | None = None,
+    ) -> None:
         """Run the agent in a CLI chat interface.
 
         Args:
             deps: The dependencies to pass to the agent.
             prog_name: The name of the program to use for the CLI. Defaults to 'pydantic-ai'.
+            message_history: History of the conversation so far.
 
         Example:
         ```python {title="agent_to_cli.py" test="skip"}
@@ -920,14 +917,28 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
 
         from pydantic_ai._cli import run_chat
 
-        await run_chat(stream=True, agent=self, deps=deps, console=Console(), code_theme='monokai', prog_name=prog_name)
+        await run_chat(
+            stream=True,
+            agent=self,
+            deps=deps,
+            console=Console(),
+            code_theme='monokai',
+            prog_name=prog_name,
+            message_history=message_history,
+        )
 
-    def to_cli_sync(self: Self, deps: AgentDepsT = None, prog_name: str = 'pydantic-ai') -> None:
+    def to_cli_sync(
+        self: Self,
+        deps: AgentDepsT = None,
+        prog_name: str = 'pydantic-ai',
+        message_history: list[_messages.ModelMessage] | None = None,
+    ) -> None:
         """Run the agent in a CLI chat interface with the non-async interface.
 
         Args:
             deps: The dependencies to pass to the agent.
             prog_name: The name of the program to use for the CLI. Defaults to 'pydantic-ai'.
+            message_history: History of the conversation so far.
 
         ```python {title="agent_to_cli_sync.py" test="skip"}
         from pydantic_ai import Agent
@@ -937,4 +948,6 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         agent.to_cli_sync(prog_name='assistant')
         ```
         """
-        return get_event_loop().run_until_complete(self.to_cli(deps=deps, prog_name=prog_name))
+        return get_event_loop().run_until_complete(
+            self.to_cli(deps=deps, prog_name=prog_name, message_history=message_history)
+        )

{pydantic_ai_slim-0.7.5 → pydantic_ai_slim-0.8.0}/pydantic_ai/direct.py

@@ -275,7 +275,9 @@ class StreamedResponseSync:
     """
 
     _async_stream_cm: AbstractAsyncContextManager[StreamedResponse]
-    _queue: queue.Queue[messages.AgentStreamEvent | Exception | None] = field(default_factory=queue.Queue, init=False)
+    _queue: queue.Queue[messages.ModelResponseStreamEvent | Exception | None] = field(
+        default_factory=queue.Queue, init=False
+    )
     _thread: threading.Thread | None = field(default=None, init=False)
     _stream_response: StreamedResponse | None = field(default=None, init=False)
     _exception: Exception | None = field(default=None, init=False)
@@ -295,8 +297,8 @@ class StreamedResponseSync:
     ) -> None:
         self._cleanup()
 
-    def __iter__(self) -> Iterator[messages.AgentStreamEvent]:
-        """Stream the response as an iterable of [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s."""
+    def __iter__(self) -> Iterator[messages.ModelResponseStreamEvent]:
+        """Stream the response as an iterable of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s."""
         self._check_context_manager_usage()
 
         while True:

{pydantic_ai_slim-0.7.5 → pydantic_ai_slim-0.8.0}/pydantic_ai/durable_exec/temporal/_function_toolset.py

@@ -51,7 +51,10 @@ class TemporalFunctionToolset(TemporalWrapperToolset[AgentDepsT]):
                     'Removing or renaming tools during an agent run is not supported with Temporal.'
                 ) from e
 
-            return await self.wrapped.call_tool(name, params.tool_args, ctx, tool)
+            # The tool args will already have been validated into their proper types in the `ToolManager`,
+            # but `execute_activity` would have turned them into simple Python types again, so we need to re-validate them.
+            args_dict = tool.args_validator.validate_python(params.tool_args)
+            return await self.wrapped.call_tool(name, args_dict, ctx, tool)
 
         # Set type hint explicitly so that Temporal can take care of serialization and deserialization
         call_tool_activity.__annotations__['deps'] = deps_type
@@ -85,7 +88,11 @@ class TemporalFunctionToolset(TemporalWrapperToolset[AgentDepsT]):
         return await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
             activity=self.call_tool_activity,
             args=[
-                _CallToolParams(name=name, tool_args=tool_args, serialized_run_context=serialized_run_context),
+                _CallToolParams(
+                    name=name,
+                    tool_args=tool_args,
+                    serialized_run_context=serialized_run_context,
+                ),
                 ctx.deps,
             ],
             **tool_activity_config,

{pydantic_ai_slim-0.7.5 → pydantic_ai_slim-0.8.0}/pydantic_ai/mcp.py

@@ -18,14 +18,13 @@ import pydantic_core
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from typing_extensions import Self, assert_never, deprecated
 
-from pydantic_ai._run_context import RunContext
-from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.tools import RunContext, ToolDefinition
 
 from .toolsets.abstract import AbstractToolset, ToolsetTool
 
 try:
     from mcp import types as mcp_types
-    from mcp.client.session import ClientSession, LoggingFnT
+    from mcp.client.session import ClientSession, ElicitationFnT, LoggingFnT
     from mcp.client.sse import sse_client
     from mcp.client.stdio import StdioServerParameters, stdio_client
     from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
@@ -57,14 +56,49 @@ class MCPServer(AbstractToolset[Any], ABC):
     """
 
     tool_prefix: str | None
+    """A prefix to add to all tools that are registered with the server.
+
+    If not empty, will include a trailing underscore(`_`).
+
+    e.g. if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
+    """
+
     log_level: mcp_types.LoggingLevel | None
+    """The log level to set when connecting to the server, if any.
+
+    See <https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/logging#logging> for more details.
+
+    If `None`, no log level will be set.
+    """
+
     log_handler: LoggingFnT | None
+    """A handler for logging messages from the server."""
+
     timeout: float
+    """The timeout in seconds to wait for the client to initialize."""
+
     read_timeout: float
+    """Maximum time in seconds to wait for new messages before timing out.
+
+    This timeout applies to the long-lived connection after it's established.
+    If no new messages are received within this time, the connection will be considered stale
+    and may be closed. Defaults to 5 minutes (300 seconds).
+    """
+
     process_tool_call: ProcessToolCallback | None
+    """Hook to customize tool calling and optionally pass extra metadata."""
+
     allow_sampling: bool
+    """Whether to allow MCP sampling through this client."""
+
     sampling_model: models.Model | None
+    """The model to use for sampling."""
+
     max_retries: int
+    """The maximum number of times to retry a tool call."""
+
+    elicitation_callback: ElicitationFnT | None = None
+    """Callback function to handle elicitation requests from the server."""
 
     _id: str | None
 
@@ -87,6 +121,7 @@ class MCPServer(AbstractToolset[Any], ABC):
         allow_sampling: bool = True,
         sampling_model: models.Model | None = None,
         max_retries: int = 1,
+        elicitation_callback: ElicitationFnT | None = None,
         *,
         id: str | None = None,
     ):
@@ -99,6 +134,7 @@ class MCPServer(AbstractToolset[Any], ABC):
         self.allow_sampling = allow_sampling
         self.sampling_model = sampling_model
         self.max_retries = max_retries
+        self.elicitation_callback = elicitation_callback
 
         self._id = id or tool_prefix
 
@@ -247,6 +283,7 @@ class MCPServer(AbstractToolset[Any], ABC):
                         read_stream=self._read_stream,
                         write_stream=self._write_stream,
                         sampling_callback=self._sampling_callback if self.allow_sampling else None,
+                        elicitation_callback=self.elicitation_callback,
                         logging_callback=self.log_handler,
                         read_timeout_seconds=timedelta(seconds=self.read_timeout),
                     )
@@ -404,46 +441,15 @@ class MCPServerStdio(MCPServer):
 
     # last fields are re-defined from the parent class so they appear as fields
     tool_prefix: str | None
-    """A prefix to add to all tools that are registered with the server.
-
-    If not empty, will include a trailing underscore(`_`).
-
-    e.g. if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
-    """
-
     log_level: mcp_types.LoggingLevel | None
-    """The log level to set when connecting to the server, if any.
-
-    See <https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/logging#logging> for more details.
-
-    If `None`, no log level will be set.
-    """
-
     log_handler: LoggingFnT | None
-    """A handler for logging messages from the server."""
-
     timeout: float
-    """The timeout in seconds to wait for the client to initialize."""
-
     read_timeout: float
-    """Maximum time in seconds to wait for new messages before timing out.
-
-    This timeout applies to the long-lived connection after it's established.
-    If no new messages are received within this time, the connection will be considered stale
-    and may be closed. Defaults to 5 minutes (300 seconds).
-    """
-
     process_tool_call: ProcessToolCallback | None
-    """Hook to customize tool calling and optionally pass extra metadata."""
-
     allow_sampling: bool
-    """Whether to allow MCP sampling through this client."""
-
     sampling_model: models.Model | None
-    """The model to use for sampling."""
-
     max_retries: int
-    """The maximum number of times to retry a tool call."""
+    elicitation_callback: ElicitationFnT | None = None
 
     def __init__(
         self,
@@ -460,6 +466,7 @@ class MCPServerStdio(MCPServer):
         allow_sampling: bool = True,
         sampling_model: models.Model | None = None,
         max_retries: int = 1,
+        elicitation_callback: ElicitationFnT | None = None,
         *,
         id: str | None = None,
     ):
@@ -479,6 +486,7 @@ class MCPServerStdio(MCPServer):
             allow_sampling: Whether to allow MCP sampling through this client.
             sampling_model: The model to use for sampling.
             max_retries: The maximum number of times to retry a tool call.
+            elicitation_callback: Callback function to handle elicitation requests from the server.
             id: An optional unique ID for the MCP server. An MCP server needs to have an ID in order to be used in a durable execution environment like Temporal, in which case the ID will be used to identify the server's activities within the workflow.
         """
         self.command = command
@@ -496,6 +504,7 @@ class MCPServerStdio(MCPServer):
             allow_sampling,
             sampling_model,
             max_retries,
+            elicitation_callback,
             id=id,
         )
 
@@ -560,50 +569,15 @@ class _MCPServerHTTP(MCPServer):
 
     # last fields are re-defined from the parent class so they appear as fields
     tool_prefix: str | None
-    """A prefix to add to all tools that are registered with the server.
-
-    If not empty, will include a trailing underscore (`_`).
-
-    For example, if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
-    """
-
     log_level: mcp_types.LoggingLevel | None
-    """The log level to set when connecting to the server, if any.
-
-    See <https://modelcontextprotocol.io/introduction#logging> for more details.
-
-    If `None`, no log level will be set.
-    """
-
     log_handler: LoggingFnT | None
-    """A handler for logging messages from the server."""
-
     timeout: float
-    """Initial connection timeout in seconds for establishing the connection.
-
-    This timeout applies to the initial connection setup and handshake.
-    If the connection cannot be established within this time, the operation will fail.
-    """
-
     read_timeout: float
-    """Maximum time in seconds to wait for new messages before timing out.
-
-    This timeout applies to the long-lived connection after it's established.
-    If no new messages are received within this time, the connection will be considered stale
-    and may be closed. Defaults to 5 minutes (300 seconds).
-    """
-
     process_tool_call: ProcessToolCallback | None
-    """Hook to customize tool calling and optionally pass extra metadata."""
-
     allow_sampling: bool
-    """Whether to allow MCP sampling through this client."""
-
     sampling_model: models.Model | None
-    """The model to use for sampling."""
-
     max_retries: int
-    """The maximum number of times to retry a tool call."""
+    elicitation_callback: ElicitationFnT | None = None
 
     def __init__(
         self,
@@ -621,6 +595,7 @@ class _MCPServerHTTP(MCPServer):
         allow_sampling: bool = True,
         sampling_model: models.Model | None = None,
         max_retries: int = 1,
+        elicitation_callback: ElicitationFnT | None = None,
         **_deprecated_kwargs: Any,
     ):
         """Build a new MCP server.
@@ -639,6 +614,7 @@ class _MCPServerHTTP(MCPServer):
             allow_sampling: Whether to allow MCP sampling through this client.
             sampling_model: The model to use for sampling.
             max_retries: The maximum number of times to retry a tool call.
+            elicitation_callback: Callback function to handle elicitation requests from the server.
         """
         if 'sse_read_timeout' in _deprecated_kwargs:
             if read_timeout is not None:
@@ -668,6 +644,7 @@ class _MCPServerHTTP(MCPServer):
             allow_sampling,
             sampling_model,
             max_retries,
+            elicitation_callback,
             id=id,
         )
 

{pydantic_ai_slim-0.7.5 → pydantic_ai_slim-0.8.0}/pydantic_ai/messages.py

@@ -110,7 +110,9 @@ class FileUrl(ABC):
     - `GoogleModel`: `VideoUrl.vendor_metadata` is used as `video_metadata`: https://ai.google.dev/gemini-api/docs/video-understanding#customize-video-processing
     """
 
-    _media_type: str | None = field(init=False, repr=False, compare=False)
+    _media_type: Annotated[str | None, pydantic.Field(alias='media_type', default=None, exclude=True)] = field(
+        compare=False, default=None
+    )
 
     def __init__(
         self,
@@ -124,6 +126,7 @@ class FileUrl(ABC):
         self.force_download = force_download
         self._media_type = media_type
 
+    @pydantic.computed_field
     @property
     def media_type(self) -> str:
         """Return the media type of the file, based on the URL or the provided `media_type`."""
@@ -160,8 +163,16 @@ class VideoUrl(FileUrl):
         vendor_metadata: dict[str, Any] | None = None,
         media_type: str | None = None,
         kind: Literal['video-url'] = 'video-url',
+        *,
+        # Required for inline-snapshot which expects all dataclass `__init__` methods to take all field names as kwargs.
+        _media_type: str | None = None,
     ) -> None:
-        super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type)
+        super().__init__(
+            url=url,
+            force_download=force_download,
+            vendor_metadata=vendor_metadata,
+            media_type=media_type or _media_type,
+        )
         self.kind = kind
 
     def _infer_media_type(self) -> VideoMediaType:
@@ -223,8 +234,16 @@ class AudioUrl(FileUrl):
         vendor_metadata: dict[str, Any] | None = None,
         media_type: str | None = None,
         kind: Literal['audio-url'] = 'audio-url',
+        *,
+        # Required for inline-snapshot which expects all dataclass `__init__` methods to take all field names as kwargs.
+        _media_type: str | None = None,
     ) -> None:
-        super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type)
+        super().__init__(
+            url=url,
+            force_download=force_download,
+            vendor_metadata=vendor_metadata,
+            media_type=media_type or _media_type,
+        )
         self.kind = kind
 
     def _infer_media_type(self) -> AudioMediaType:
@@ -273,8 +292,16 @@ class ImageUrl(FileUrl):
         vendor_metadata: dict[str, Any] | None = None,
         media_type: str | None = None,
         kind: Literal['image-url'] = 'image-url',
+        *,
+        # Required for inline-snapshot which expects all dataclass `__init__` methods to take all field names as kwargs.
+        _media_type: str | None = None,
     ) -> None:
-        super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type)
+        super().__init__(
+            url=url,
+            force_download=force_download,
+            vendor_metadata=vendor_metadata,
+            media_type=media_type or _media_type,
+        )
         self.kind = kind
 
     def _infer_media_type(self) -> ImageMediaType:
@@ -318,8 +345,16 @@ class DocumentUrl(FileUrl):
         vendor_metadata: dict[str, Any] | None = None,
         media_type: str | None = None,
         kind: Literal['document-url'] = 'document-url',
+        *,
+        # Required for inline-snapshot which expects all dataclass `__init__` methods to take all field names as kwargs.
+        _media_type: str | None = None,
     ) -> None:
-        super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type)
+        super().__init__(
+            url=url,
+            force_download=force_download,
+            vendor_metadata=vendor_metadata,
+            media_type=media_type or _media_type,
+        )
         self.kind = kind
 
     def _infer_media_type(self) -> str:
@@ -1263,13 +1298,10 @@ class FinalResultEvent:
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
-ModelResponseStreamEvent = Annotated[Union[PartStartEvent, PartDeltaEvent], pydantic.Discriminator('event_kind')]
-"""An event in the model response stream, either starting a new part or applying a delta to an existing one."""
-
-AgentStreamEvent = Annotated[
+ModelResponseStreamEvent = Annotated[
     Union[PartStartEvent, PartDeltaEvent, FinalResultEvent], pydantic.Discriminator('event_kind')
 ]
-"""An event in the agent stream."""
+"""An event in the model response stream, starting a new part, applying a delta to an existing one, or indicating the final result."""
 
 
 @dataclass(repr=False)
@@ -1338,3 +1370,7 @@ HandleResponseEvent = Annotated[
     Union[FunctionToolCallEvent, FunctionToolResultEvent, BuiltinToolCallEvent, BuiltinToolResultEvent],
     pydantic.Discriminator('event_kind'),
 ]
+"""An event yielded when handling a model response, indicating tool calls and results."""
+
+AgentStreamEvent = Annotated[Union[ModelResponseStreamEvent, HandleResponseEvent], pydantic.Discriminator('event_kind')]
+"""An event in the agent stream: model response stream events and response-handling events."""