pydantic-ai-slim 0.7.1__tar.gz → 0.7.3__tar.gz

{pydantic_ai_slim-0.7.1 → pydantic_ai_slim-0.7.3}/PKG-INFO

@@ -1,7 +1,11 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.7.1
+Version: 0.7.3
 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
+Project-URL: Documentation, https://ai.pydantic.dev/install/#slim-install
+Project-URL: Changelog, https://github.com/pydantic/pydantic-ai/releases
 Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>, Douwe Maan <douwe@pydantic.dev>
 License-Expression: MIT
 License-File: LICENSE
@@ -27,10 +31,11 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.9
 Requires-Dist: eval-type-backport>=0.2.0
 Requires-Dist: exceptiongroup; python_version < '3.11'
+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.1
+Requires-Dist: pydantic-graph==0.7.3
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -45,13 +50,14 @@ Requires-Dist: boto3>=1.39.0; extra == 'bedrock'
 Provides-Extra: cli
 Requires-Dist: argcomplete>=3.5.0; extra == 'cli'
 Requires-Dist: prompt-toolkit>=3; extra == 'cli'
+Requires-Dist: pyperclip>=1.9.0; extra == 'cli'
 Requires-Dist: rich>=13; extra == 'cli'
 Provides-Extra: cohere
 Requires-Dist: cohere>=5.16.0; (platform_system != 'Emscripten') and extra == 'cohere'
 Provides-Extra: duckduckgo
 Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.7.1; extra == 'evals'
+Requires-Dist: pydantic-evals==0.7.3; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.28.0; extra == 'google'
 Provides-Extra: groq
@@ -71,7 +77,7 @@ Requires-Dist: tenacity>=8.2.3; extra == 'retries'
 Provides-Extra: tavily
 Requires-Dist: tavily-python>=0.5.0; extra == 'tavily'
 Provides-Extra: temporal
-Requires-Dist: temporalio>=1.15.0; extra == 'temporal'
+Requires-Dist: temporalio==1.15.0; extra == 'temporal'
 Provides-Extra: vertexai
 Requires-Dist: google-auth>=2.36.0; extra == 'vertexai'
 Requires-Dist: requests>=2.32.2; extra == 'vertexai'

{pydantic_ai_slim-0.7.1 → pydantic_ai_slim-0.7.3}/pydantic_ai/_agent_graph.py

@@ -23,7 +23,7 @@ from pydantic_graph.nodes import End, NodeRunEndT
 from . import _output, _system_prompt, exceptions, messages as _messages, models, result, usage as _usage
 from .exceptions import ToolRetryError
 from .output import OutputDataT, OutputSpec
-from .settings import ModelSettings, merge_model_settings
+from .settings import ModelSettings
 from .tools import RunContext, ToolDefinition, ToolKind
 
 if TYPE_CHECKING:
@@ -76,7 +76,7 @@ class GraphAgentState:
     """State kept across the execution of the agent graph."""
 
     message_history: list[_messages.ModelMessage]
-    usage: _usage.Usage
+    usage: _usage.RunUsage
     retries: int
     run_step: int
 
@@ -158,28 +158,7 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
 
     async def run(
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
-    ) -> ModelRequestNode[DepsT, NodeRunEndT]:
-        return ModelRequestNode[DepsT, NodeRunEndT](request=await self._get_first_message(ctx))
-
-    async def _get_first_message(
-        self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
-    ) -> _messages.ModelRequest:
-        run_context = build_run_context(ctx)
-        history, next_message = await self._prepare_messages(
-            self.user_prompt, ctx.state.message_history, ctx.deps.get_instructions, run_context
-        )
-        ctx.state.message_history = history
-        run_context.messages = history
-
-        return next_message
-
-    async def _prepare_messages(
-        self,
-        user_prompt: str | Sequence[_messages.UserContent] | None,
-        message_history: list[_messages.ModelMessage] | None,
-        get_instructions: Callable[[RunContext[DepsT]], Awaitable[str | None]],
-        run_context: RunContext[DepsT],
-    ) -> tuple[list[_messages.ModelMessage], _messages.ModelRequest]:
+    ) -> Union[ModelRequestNode[DepsT, NodeRunEndT], CallToolsNode[DepsT, NodeRunEndT]]:  # noqa UP007
         try:
             ctx_messages = get_captured_run_messages()
         except LookupError:
@@ -191,29 +170,48 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
                 messages = ctx_messages.messages
                 ctx_messages.used = True
 
+        # Add message history to the `capture_run_messages` list, which will be empty at this point
+        messages.extend(ctx.state.message_history)
+        # Use the `capture_run_messages` list as the message history so that new messages are added to it
+        ctx.state.message_history = messages
+
+        run_context = build_run_context(ctx)
+
         parts: list[_messages.ModelRequestPart] = []
-        instructions = await get_instructions(run_context)
-        if message_history:
-            # Shallow copy messages
-            messages.extend(message_history)
+        if messages:
             # Reevaluate any dynamic system prompt parts
             await self._reevaluate_dynamic_prompts(messages, run_context)
         else:
             parts.extend(await self._sys_parts(run_context))
 
-        if user_prompt is not None:
-            parts.append(_messages.UserPromptPart(user_prompt))
-        elif (
-            len(parts) == 0
-            and message_history
-            and (last_message := message_history[-1])
-            and isinstance(last_message, _messages.ModelRequest)
-        ):
-            # Drop last message that came from history and reuse its parts
-            messages.pop()
-            parts.extend(last_message.parts)
+        if messages and (last_message := messages[-1]):
+            if isinstance(last_message, _messages.ModelRequest) and self.user_prompt is None:
+                # Drop last message from history and reuse its parts
+                messages.pop()
+                parts.extend(last_message.parts)
+            elif isinstance(last_message, _messages.ModelResponse):
+                if self.user_prompt is None:
+                    # `CallToolsNode` requires the tool manager to be prepared for the run step
+                    # This will raise errors for any tool name conflicts
+                    ctx.deps.tool_manager = await ctx.deps.tool_manager.for_run_step(run_context)
+
+                    # Skip ModelRequestNode and go directly to CallToolsNode
+                    return CallToolsNode[DepsT, NodeRunEndT](model_response=last_message)
+                elif any(isinstance(part, _messages.ToolCallPart) for part in last_message.parts):
+                    raise exceptions.UserError(
+                        'Cannot provide a new user prompt when the message history ends with '
+                        'a model response containing unprocessed tool calls. Either process the '
+                        'tool calls first (by calling `iter` with `user_prompt=None`) or append a '
+                        '`ModelRequest` with `ToolResultPart`s.'
+                    )
+
+        if self.user_prompt is not None:
+            parts.append(_messages.UserPromptPart(self.user_prompt))
+
+        instructions = await ctx.deps.get_instructions(run_context)
+        next_message = _messages.ModelRequest(parts, instructions=instructions)
 
-        return messages, _messages.ModelRequest(parts, instructions=instructions)
+        return ModelRequestNode[DepsT, NodeRunEndT](request=next_message)
 
     async def _reevaluate_dynamic_prompts(
         self, messages: list[_messages.ModelMessage], run_context: RunContext[DepsT]
@@ -250,9 +248,6 @@ async def _prepare_request_parameters(
     ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
 ) -> models.ModelRequestParameters:
     """Build tools and create an agent model."""
-    run_context = build_run_context(ctx)
-    ctx.deps.tool_manager = await ctx.deps.tool_manager.for_run_step(run_context)
-
     output_schema = ctx.deps.output_schema
     output_object = None
     if isinstance(output_schema, _output.NativeOutputSchema):
@@ -342,7 +337,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
 
         model_settings, model_request_parameters, message_history, _ = await self._prepare_request(ctx)
         model_response = await ctx.deps.model.request(message_history, model_settings, model_request_parameters)
-        ctx.state.usage.incr(_usage.Usage())
+        ctx.state.usage.requests += 1
 
         return self._finish_handling(ctx, model_response)
 
@@ -355,21 +350,21 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
 
         run_context = build_run_context(ctx)
 
-        model_settings = merge_model_settings(ctx.deps.model_settings, None)
+        # This will raise errors for any tool name conflicts
+        ctx.deps.tool_manager = await ctx.deps.tool_manager.for_run_step(run_context)
+
+        message_history = await _process_message_history(ctx.state, ctx.deps.history_processors, run_context)
 
         model_request_parameters = await _prepare_request_parameters(ctx)
         model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
 
-        message_history = await _process_message_history(ctx.state, ctx.deps.history_processors, run_context)
-
+        model_settings = ctx.deps.model_settings
         usage = ctx.state.usage
         if ctx.deps.usage_limits.count_tokens_before_request:
             # Copy to avoid modifying the original usage object with the counted usage
             usage = dataclasses.replace(usage)
 
-            counted_usage = await ctx.deps.model.count_tokens(
-                message_history, ctx.deps.model_settings, model_request_parameters
-            )
+            counted_usage = await ctx.deps.model.count_tokens(message_history, model_settings, model_request_parameters)
             usage.incr(counted_usage)
 
         ctx.deps.usage_limits.check_before_request(usage)
@@ -432,9 +427,11 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
         if self._events_iterator is None:
             # Ensure that the stream is only run once
 
-            async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:
+            async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:  # noqa: C901
                 texts: list[str] = []
                 tool_calls: list[_messages.ToolCallPart] = []
+                thinking_parts: list[_messages.ThinkingPart] = []
+
                 for part in self.model_response.parts:
                     if isinstance(part, _messages.TextPart):
                         # ignore empty content for text parts, see #437
@@ -447,11 +444,7 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                     elif isinstance(part, _messages.BuiltinToolReturnPart):
                         yield _messages.BuiltinToolResultEvent(part)
                     elif isinstance(part, _messages.ThinkingPart):
-                        # We don't need to do anything with thinking parts in this tool-calling node.
-                        # We need to handle text parts in case there are no tool calls and/or the desired output comes
-                        # from the text, but thinking parts should not directly influence the execution of tools or
-                        # determination of the next node of graph execution here.
-                        pass
+                        thinking_parts.append(part)
                     else:
                         assert_never(part)
 
@@ -465,8 +458,18 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                 elif texts:
                     # No events are emitted during the handling of text responses, so we don't need to yield anything
                     self._next_node = await self._handle_text_response(ctx, texts)
+                elif thinking_parts:
+                    # handle thinking-only responses (responses that contain only ThinkingPart instances)
+                    # this can happen with models that support thinking mode when they don't provide
+                    # actionable output alongside their thinking content.
+                    self._next_node = ModelRequestNode[DepsT, NodeRunEndT](
+                        _messages.ModelRequest(
+                            parts=[_messages.RetryPromptPart('Responses without text or tool calls are not permitted.')]
+                        )
+                    )
                 else:
-                    # we've got an empty response, this sometimes happens with anthropic (and perhaps other models)
+                    # we got an empty response with no tool calls, text, or thinking
+                    # this sometimes happens with anthropic (and perhaps other models)
                     # when the model has already returned text along side tool calls
                     # in this scenario, if text responses are allowed, we return text from the most recent model
                     # response, if any

{pydantic_ai_slim-0.7.1 → pydantic_ai_slim-0.7.3}/pydantic_ai/_cli.py

@@ -18,12 +18,13 @@ from . import __version__
 from ._run_context import AgentDepsT
 from .agent import AbstractAgent, Agent
 from .exceptions import UserError
-from .messages import ModelMessage
+from .messages import ModelMessage, TextPart
 from .models import KnownModelName, infer_model
 from .output import OutputDataT
 
 try:
     import argcomplete
+    import pyperclip
     from prompt_toolkit import PromptSession
     from prompt_toolkit.auto_suggest import AutoSuggestFromHistory, Suggestion
     from prompt_toolkit.buffer import Buffer
@@ -38,7 +39,7 @@ try:
     from rich.text import Text
 except ImportError as _import_error:
     raise ImportError(
-        'Please install `rich`, `prompt-toolkit` and `argcomplete` to use the Pydantic AI CLI, '
+        'Please install `rich`, `prompt-toolkit`, `pyperclip` and `argcomplete` to use the Pydantic AI CLI, '
         'you can use the `cli` optional group — `pip install "pydantic-ai-slim[cli]"`'
     ) from _import_error
 
@@ -114,6 +115,7 @@ Special prompts:
 * `/exit` - exit the interactive mode (ctrl-c and ctrl-d also work)
 * `/markdown` - show the last markdown output of the last question
 * `/multiline` - toggle multiline mode
+* `/cp` - copy the last response to clipboard
 """,
         formatter_class=argparse.RawTextHelpFormatter,
     )
@@ -237,7 +239,7 @@ async def run_chat(
 
     while True:
         try:
-            auto_suggest = CustomAutoSuggest(['/markdown', '/multiline', '/exit'])
+            auto_suggest = CustomAutoSuggest(['/markdown', '/multiline', '/exit', '/cp'])
             text = await session.prompt_async(f'{prog_name} ➤ ', auto_suggest=auto_suggest, multiline=multiline)
         except (KeyboardInterrupt, EOFError):  # pragma: no cover
             return 0
@@ -347,6 +349,19 @@ def handle_slash_command(
     elif ident_prompt == '/exit':
         console.print('[dim]Exiting…[/dim]')
         return 0, multiline
+    elif ident_prompt == '/cp':
+        try:
+            parts = messages[-1].parts
+        except IndexError:
+            console.print('[dim]No output available to copy.[/dim]')
+        else:
+            text_to_copy = '\n\n'.join(part.content for part in parts if isinstance(part, TextPart))
+            text_to_copy = text_to_copy.strip()
+            if text_to_copy:
+                pyperclip.copy(text_to_copy)
+                console.print('[dim]Copied last output to clipboard.[/dim]')
+            else:
+                console.print('[dim]No text content to copy.[/dim]')
     else:
         console.print(f'[red]Unknown command[/red] [magenta]`{ident_prompt}`[/magenta]')
     return None, multiline

{pydantic_ai_slim-0.7.1 → pydantic_ai_slim-0.7.3}/pydantic_ai/_parts_manager.py

@@ -72,6 +72,7 @@ class ModelResponsePartsManager:
         vendor_part_id: VendorId | None,
         content: str,
         thinking_tags: tuple[str, str] | None = None,
+        ignore_leading_whitespace: bool = False,
     ) -> ModelResponseStreamEvent | None:
         """Handle incoming text content, creating or updating a TextPart in the manager as appropriate.
 
@@ -85,6 +86,7 @@ class ModelResponsePartsManager:
                 a TextPart.
             content: The text content to append to the appropriate TextPart.
             thinking_tags: If provided, will handle content between the thinking tags as thinking parts.
+            ignore_leading_whitespace: If True, will ignore leading whitespace in the content.
 
         Returns:
             - A `PartStartEvent` if a new part was created.
@@ -128,10 +130,9 @@ class ModelResponsePartsManager:
             return self.handle_thinking_delta(vendor_part_id=vendor_part_id, content='')
 
         if existing_text_part_and_index is None:
-            # If the first text delta is all whitespace, don't emit a new part yet.
-            # This is a workaround for models that emit `<think>\n</think>\n\n` ahead of tool calls (e.g. Ollama + Qwen3),
-            # which we don't want to end up treating as a final result.
-            if content.isspace():
+            # This is a workaround for models that emit `<think>\n</think>\n\n` or an empty text part ahead of tool calls (e.g. Ollama + Qwen3),
+            # which we don't want to end up treating as a final result when using `run_stream` with `str` a valid `output_type`.
+            if ignore_leading_whitespace and (len(content) == 0 or content.isspace()):
                 return None
 
             # There is no existing text part that should be updated, so create a new one

{pydantic_ai_slim-0.7.1 → pydantic_ai_slim-0.7.3}/pydantic_ai/_run_context.py

@@ -12,7 +12,7 @@ from . import _utils, messages as _messages
 
 if TYPE_CHECKING:
     from .models import Model
-    from .result import Usage
+    from .result import RunUsage
 
 AgentDepsT = TypeVar('AgentDepsT', default=None, contravariant=True)
 """Type variable for agent dependencies."""
@@ -26,7 +26,7 @@ class RunContext(Generic[AgentDepsT]):
     """Dependencies for the agent."""
     model: Model
     """The model used in this run."""
-    usage: Usage
+    usage: RunUsage
     """LLM usage associated with the run."""
     prompt: str | Sequence[_messages.UserContent] | None = None
     """The original user prompt passed to the run."""

{pydantic_ai_slim-0.7.1 → pydantic_ai_slim-0.7.3}/pydantic_ai/_tool_manager.py

@@ -5,6 +5,7 @@ from collections.abc import Iterable
 from dataclasses import dataclass, field, replace
 from typing import Any, Generic
 
+from opentelemetry.trace import Tracer
 from pydantic import ValidationError
 from typing_extensions import assert_never
 
@@ -21,41 +22,46 @@ from .toolsets.abstract import AbstractToolset, ToolsetTool
 class ToolManager(Generic[AgentDepsT]):
     """Manages tools for an agent run step. It caches the agent run's toolset's tool definitions and handles calling tools and retries."""
 
-    ctx: RunContext[AgentDepsT]
-    """The agent run context for a specific run step."""
     toolset: AbstractToolset[AgentDepsT]
     """The toolset that provides the tools for this run step."""
-    tools: dict[str, ToolsetTool[AgentDepsT]]
+    ctx: RunContext[AgentDepsT] | None = None
+    """The agent run context for a specific run step."""
+    tools: dict[str, ToolsetTool[AgentDepsT]] | None = None
     """The cached tools for this run step."""
     failed_tools: set[str] = field(default_factory=set)
     """Names of tools that failed in this run step."""
 
-    @classmethod
-    async def build(cls, toolset: AbstractToolset[AgentDepsT], ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDepsT]:
-        """Build a new tool manager for a specific run step."""
-        return cls(
-            ctx=ctx,
-            toolset=toolset,
-            tools=await toolset.get_tools(ctx),
-        )
-
     async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDepsT]:
         """Build a new tool manager for the next run step, carrying over the retries from the current run step."""
-        if ctx.run_step == self.ctx.run_step:
-            return self
-
-        retries = {
-            failed_tool_name: self.ctx.retries.get(failed_tool_name, 0) + 1 for failed_tool_name in self.failed_tools
-        }
-        return await self.__class__.build(self.toolset, replace(ctx, retries=retries))
+        if self.ctx is not None:
+            if ctx.run_step == self.ctx.run_step:
+                return self
+
+            retries = {
+                failed_tool_name: self.ctx.retries.get(failed_tool_name, 0) + 1
+                for failed_tool_name in self.failed_tools
+            }
+            ctx = replace(ctx, retries=retries)
+
+        return self.__class__(
+            toolset=self.toolset,
+            ctx=ctx,
+            tools=await self.toolset.get_tools(ctx),
+        )
 
     @property
     def tool_defs(self) -> list[ToolDefinition]:
         """The tool definitions for the tools in this tool manager."""
+        if self.tools is None:
+            raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
+
         return [tool.tool_def for tool in self.tools.values()]
 
     def get_tool_def(self, name: str) -> ToolDefinition | None:
         """Get the tool definition for a given tool name, or `None` if the tool is unknown."""
+        if self.tools is None:
+            raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
+
         try:
             return self.tools[name].tool_def
         except KeyError:
@@ -71,15 +77,25 @@ class ToolManager(Generic[AgentDepsT]):
             allow_partial: Whether to allow partial validation of the tool arguments.
             wrap_validation_errors: Whether to wrap validation errors in a retry prompt part.
         """
+        if self.tools is None or self.ctx is None:
+            raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
+
         if (tool := self.tools.get(call.tool_name)) and tool.tool_def.kind == 'output':
             # Output tool calls are not traced
             return await self._call_tool(call, allow_partial, wrap_validation_errors)
         else:
-            return await self._call_tool_traced(call, allow_partial, wrap_validation_errors)
+            return await self._call_tool_traced(
+                call,
+                allow_partial,
+                wrap_validation_errors,
+                self.ctx.tracer,
+                self.ctx.trace_include_content,
+            )
+
+    async def _call_tool(self, call: ToolCallPart, allow_partial: bool, wrap_validation_errors: bool) -> Any:
+        if self.tools is None or self.ctx is None:
+            raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
 
-    async def _call_tool(
-        self, call: ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
-    ) -> Any:
         name = call.tool_name
         tool = self.tools.get(name)
         try:
@@ -137,14 +153,19 @@ class ToolManager(Generic[AgentDepsT]):
                 raise e
 
     async def _call_tool_traced(
-        self, call: ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
+        self,
+        call: ToolCallPart,
+        allow_partial: bool,
+        wrap_validation_errors: bool,
+        tracer: Tracer,
+        include_content: bool = False,
     ) -> Any:
         """See <https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span>."""
         span_attributes = {
             'gen_ai.tool.name': call.tool_name,
             # NOTE: this means `gen_ai.tool.call.id` will be included even if it was generated by pydantic-ai
             'gen_ai.tool.call.id': call.tool_call_id,
-            **({'tool_arguments': call.args_as_json_str()} if self.ctx.trace_include_content else {}),
+            **({'tool_arguments': call.args_as_json_str()} if include_content else {}),
             'logfire.msg': f'running tool: {call.tool_name}',
             # add the JSON schema so these attributes are formatted nicely in Logfire
             'logfire.json_schema': json.dumps(
@@ -156,7 +177,7 @@ class ToolManager(Generic[AgentDepsT]):
                                 'tool_arguments': {'type': 'object'},
                                 'tool_response': {'type': 'object'},
                             }
-                            if self.ctx.trace_include_content
+                            if include_content
                             else {}
                         ),
                         'gen_ai.tool.name': {},
@@ -165,16 +186,16 @@ class ToolManager(Generic[AgentDepsT]):
                 }
             ),
         }
-        with self.ctx.tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
+        with tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
             try:
                 tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
             except ToolRetryError as e:
                 part = e.tool_retry
-                if self.ctx.trace_include_content and span.is_recording():
+                if include_content and span.is_recording():
                     span.set_attribute('tool_response', part.model_response())
                 raise e
 
-            if self.ctx.trace_include_content and span.is_recording():
+            if include_content and span.is_recording():
                 span.set_attribute(
                     'tool_response',
                     tool_result

{pydantic_ai_slim-0.7.1 → pydantic_ai_slim-0.7.3}/pydantic_ai/ag_ui.py

@@ -51,7 +51,7 @@ from .settings import ModelSettings
 from .tools import AgentDepsT, ToolDefinition
 from .toolsets import AbstractToolset
 from .toolsets.deferred import DeferredToolset
-from .usage import Usage, UsageLimits
+from .usage import RunUsage, UsageLimits
 
 try:
     from ag_ui.core import (
@@ -127,7 +127,7 @@ class AGUIApp(Generic[AgentDepsT, OutputDataT], Starlette):
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
         usage_limits: UsageLimits | None = None,
-        usage: Usage | None = None,
+        usage: RunUsage | None = None,
         infer_name: bool = True,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
         # Starlette parameters.
@@ -216,7 +216,7 @@ async def handle_ag_ui_request(
     deps: AgentDepsT = None,
     model_settings: ModelSettings | None = None,
     usage_limits: UsageLimits | None = None,
-    usage: Usage | None = None,
+    usage: RunUsage | None = None,
     infer_name: bool = True,
     toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
 ) -> Response:
@@ -277,7 +277,7 @@ async def run_ag_ui(
     deps: AgentDepsT = None,
     model_settings: ModelSettings | None = None,
     usage_limits: UsageLimits | None = None,
-    usage: Usage | None = None,
+    usage: RunUsage | None = None,
     infer_name: bool = True,
     toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
 ) -> AsyncIterator[str]: