pydantic-ai-slim 1.0.11__py3-none-any.whl → 1.0.13__py3-none-any.whl

pydantic_ai/__init__.py

@@ -22,10 +22,79 @@ from .exceptions import (
     UserError,
 )
 from .format_prompt import format_as_xml
-from .messages import AudioUrl, BinaryContent, DocumentUrl, ImageUrl, VideoUrl
+from .messages import (
+    AgentStreamEvent,
+    AudioFormat,
+    AudioMediaType,
+    AudioUrl,
+    BaseToolCallPart,
+    BaseToolReturnPart,
+    BinaryContent,
+    BuiltinToolCallPart,
+    BuiltinToolReturnPart,
+    DocumentFormat,
+    DocumentMediaType,
+    DocumentUrl,
+    FileUrl,
+    FinalResultEvent,
+    FinishReason,
+    FunctionToolCallEvent,
+    FunctionToolResultEvent,
+    HandleResponseEvent,
+    ImageFormat,
+    ImageMediaType,
+    ImageUrl,
+    ModelMessage,
+    ModelMessagesTypeAdapter,
+    ModelRequest,
+    ModelRequestPart,
+    ModelResponse,
+    ModelResponsePart,
+    ModelResponsePartDelta,
+    ModelResponseStreamEvent,
+    MultiModalContent,
+    PartDeltaEvent,
+    PartStartEvent,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    TextPartDelta,
+    ThinkingPart,
+    ThinkingPartDelta,
+    ToolCallPart,
+    ToolCallPartDelta,
+    ToolReturn,
+    ToolReturnPart,
+    UserContent,
+    UserPromptPart,
+    VideoFormat,
+    VideoMediaType,
+    VideoUrl,
+)
 from .output import NativeOutput, PromptedOutput, StructuredDict, TextOutput, ToolOutput
+from .profiles import (
+    DEFAULT_PROFILE,
+    InlineDefsJsonSchemaTransformer,
+    JsonSchemaTransformer,
+    ModelProfile,
+    ModelProfileSpec,
+)
 from .settings import ModelSettings
 from .tools import DeferredToolRequests, DeferredToolResults, RunContext, Tool, ToolApproved, ToolDefinition, ToolDenied
+from .toolsets import (
+    AbstractToolset,
+    ApprovalRequiredToolset,
+    CombinedToolset,
+    ExternalToolset,
+    FilteredToolset,
+    FunctionToolset,
+    PrefixedToolset,
+    PreparedToolset,
+    RenamedToolset,
+    ToolsetFunc,
+    ToolsetTool,
+    WrapperToolset,
+)
 from .usage import RequestUsage, RunUsage, UsageLimits
 
 __all__ = (
@@ -49,11 +118,59 @@ __all__ = (
     'UsageLimitExceeded',
     'UserError',
     # messages
-    'ImageUrl',
+    'AgentStreamEvent',
+    'AudioFormat',
+    'AudioMediaType',
     'AudioUrl',
-    'VideoUrl',
-    'DocumentUrl',
+    'BaseToolCallPart',
+    'BaseToolReturnPart',
     'BinaryContent',
+    'BuiltinToolCallPart',
+    'BuiltinToolReturnPart',
+    'DocumentFormat',
+    'DocumentMediaType',
+    'DocumentUrl',
+    'FileUrl',
+    'FinalResultEvent',
+    'FinishReason',
+    'FunctionToolCallEvent',
+    'FunctionToolResultEvent',
+    'HandleResponseEvent',
+    'ImageFormat',
+    'ImageMediaType',
+    'ImageUrl',
+    'ModelMessage',
+    'ModelMessagesTypeAdapter',
+    'ModelRequest',
+    'ModelRequestPart',
+    'ModelResponse',
+    'ModelResponsePart',
+    'ModelResponsePartDelta',
+    'ModelResponseStreamEvent',
+    'MultiModalContent',
+    'PartDeltaEvent',
+    'PartStartEvent',
+    'RetryPromptPart',
+    'SystemPromptPart',
+    'TextPart',
+    'TextPartDelta',
+    'ThinkingPart',
+    'ThinkingPartDelta',
+    'ToolCallPart',
+    'ToolCallPartDelta',
+    'ToolReturn',
+    'ToolReturnPart',
+    'UserContent',
+    'UserPromptPart',
+    'VideoFormat',
+    'VideoMediaType',
+    'VideoUrl',
+    # profiles
+    'ModelProfile',
+    'ModelProfileSpec',
+    'DEFAULT_PROFILE',
+    'InlineDefsJsonSchemaTransformer',
+    'JsonSchemaTransformer',
     # tools
     'Tool',
     'ToolDefinition',
@@ -62,6 +179,19 @@ __all__ = (
     'DeferredToolResults',
     'ToolApproved',
     'ToolDenied',
+    # toolsets
+    'AbstractToolset',
+    'ApprovalRequiredToolset',
+    'CombinedToolset',
+    'ExternalToolset',
+    'FilteredToolset',
+    'FunctionToolset',
+    'PrefixedToolset',
+    'PreparedToolset',
+    'RenamedToolset',
+    'ToolsetFunc',
+    'ToolsetTool',
+    'WrapperToolset',
     # builtin_tools
     'WebSearchTool',
     'WebSearchUserLocation',

pydantic_ai/_a2a.py

@@ -10,7 +10,7 @@ from typing import Any, Generic, TypeVar
 from pydantic import TypeAdapter
 from typing_extensions import assert_never
 
-from pydantic_ai.messages import (
+from pydantic_ai import (
     AudioUrl,
     BinaryContent,
     DocumentUrl,

pydantic_ai/_agent_graph.py

@@ -16,6 +16,7 @@ from opentelemetry.trace import Tracer
 from typing_extensions import TypeVar, assert_never
 
 from pydantic_ai._function_schema import _takes_ctx as is_takes_ctx  # type: ignore
+from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION
 from pydantic_ai._tool_manager import ToolManager
 from pydantic_ai._utils import dataclasses_no_defaults_repr, get_union_args, is_async_callable, run_in_executor
 from pydantic_ai.builtin_tools import AbstractBuiltinTool
@@ -704,6 +705,9 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
         tracer=ctx.deps.tracer,
         trace_include_content=ctx.deps.instrumentation_settings is not None
         and ctx.deps.instrumentation_settings.include_content,
+        instrumentation_version=ctx.deps.instrumentation_settings.version
+        if ctx.deps.instrumentation_settings
+        else DEFAULT_INSTRUMENTATION_VERSION,
         run_step=ctx.state.run_step,
         tool_call_approved=ctx.state.run_step == 0,
     )

pydantic_ai/_instrumentation.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+DEFAULT_INSTRUMENTATION_VERSION = 2
+"""Default instrumentation version for `InstrumentationSettings`."""
+
+
+@dataclass(frozen=True)
+class InstrumentationNames:
+    """Configuration for instrumentation span names and attributes based on version."""
+
+    # Agent run span configuration
+    agent_run_span_name: str
+    agent_name_attr: str
+
+    # Tool execution span configuration
+    tool_span_name: str
+    tool_arguments_attr: str
+    tool_result_attr: str
+
+    # Output Tool execution span configuration
+    output_tool_span_name: str
+
+    @classmethod
+    def for_version(cls, version: int) -> Self:
+        """Create instrumentation configuration for a specific version.
+
+        Args:
+            version: The instrumentation version (1, 2, or 3+)
+
+        Returns:
+            InstrumentationConfig instance with version-appropriate settings
+        """
+        if version <= 2:
+            return cls(
+                agent_run_span_name='agent run',
+                agent_name_attr='agent_name',
+                tool_span_name='running tool',
+                tool_arguments_attr='tool_arguments',
+                tool_result_attr='tool_response',
+                output_tool_span_name='running output function',
+            )
+        else:
+            return cls(
+                agent_run_span_name='invoke_agent',
+                agent_name_attr='gen_ai.agent.name',
+                tool_span_name='execute_tool',  # Will be formatted with tool name
+                tool_arguments_attr='gen_ai.tool.call.arguments',
+                tool_result_attr='gen_ai.tool.call.result',
+                output_tool_span_name='execute_tool',
+            )
+
+    def get_agent_run_span_name(self, agent_name: str) -> str:
+        """Get the formatted agent span name.
+
+        Args:
+            agent_name: Name of the agent being executed
+
+        Returns:
+            Formatted span name
+        """
+        if self.agent_run_span_name == 'invoke_agent':
+            return f'invoke_agent {agent_name}'
+        return self.agent_run_span_name
+
+    def get_tool_span_name(self, tool_name: str) -> str:
+        """Get the formatted tool span name.
+
+        Args:
+            tool_name: Name of the tool being executed
+
+        Returns:
+            Formatted span name
+        """
+        if self.tool_span_name == 'execute_tool':
+            return f'execute_tool {tool_name}'
+        return self.tool_span_name
+
+    def get_output_tool_span_name(self, tool_name: str) -> str:
+        """Get the formatted output tool span name.
+
+        Args:
+            tool_name: Name of the tool being executed
+
+        Returns:
+            Formatted span name
+        """
+        if self.output_tool_span_name == 'execute_tool':
+            return f'execute_tool {tool_name}'
+        return self.output_tool_span_name

pydantic_ai/{profiles/_json_schema.py → _json_schema.py}

@@ -6,7 +6,7 @@ from copy import deepcopy
 from dataclasses import dataclass
 from typing import Any, Literal
 
-from pydantic_ai.exceptions import UserError
+from .exceptions import UserError
 
 JsonSchema = dict[str, Any]
 
@@ -54,14 +54,14 @@ class JsonSchemaTransformer(ABC):
         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
+        elif self.recursive_refs:
             # 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:
+            if root_key is None:  # pragma: no cover
                 root_key = self.schema.get('title', 'root')
                 while root_key in defs:
                     # Modify the root key until it is not already in use
@@ -77,6 +77,8 @@ class JsonSchemaTransformer(ABC):
         if self.prefer_inlined_defs:
             while ref := schema.get('$ref'):
                 key = re.sub(r'^#/\$defs/', '', ref)
+                if key in self.recursive_refs:
+                    break
                 if key in self.refs_stack:
                     self.recursive_refs.add(key)
                     break  # recursive ref can't be unpacked

pydantic_ai/_output.py

@@ -11,6 +11,8 @@ from pydantic import Json, TypeAdapter, ValidationError
 from pydantic_core import SchemaValidator, to_json
 from typing_extensions import Self, TypedDict, TypeVar, assert_never
 
+from pydantic_ai._instrumentation import InstrumentationNames
+
 from . import _function_schema, _utils, messages as _messages
 from ._run_context import AgentDepsT, RunContext
 from .exceptions import ModelRetry, ToolRetryError, UserError
@@ -95,6 +97,7 @@ async def execute_traced_output_function(
         ToolRetryError: When wrap_validation_errors is True and a ModelRetry is caught
         ModelRetry: When wrap_validation_errors is False and a ModelRetry occurs
     """
+    instrumentation_names = InstrumentationNames.for_version(run_context.instrumentation_version)
     # Set up span attributes
     tool_name = run_context.tool_name or getattr(function_schema.function, '__name__', 'output_function')
     attributes = {
@@ -104,18 +107,29 @@ async def execute_traced_output_function(
     if run_context.tool_call_id:
         attributes['gen_ai.tool.call.id'] = run_context.tool_call_id
     if run_context.trace_include_content:
-        attributes['tool_arguments'] = to_json(args).decode()
-        attributes['logfire.json_schema'] = json.dumps(
-            {
-                'type': 'object',
-                'properties': {
-                    'tool_arguments': {'type': 'object'},
-                    'tool_response': {'type': 'object'},
-                },
-            }
-        )
+        attributes[instrumentation_names.tool_arguments_attr] = to_json(args).decode()
+
+    attributes['logfire.json_schema'] = json.dumps(
+        {
+            'type': 'object',
+            'properties': {
+                **(
+                    {
+                        instrumentation_names.tool_arguments_attr: {'type': 'object'},
+                        instrumentation_names.tool_result_attr: {'type': 'object'},
+                    }
+                    if run_context.trace_include_content
+                    else {}
+                ),
+                'gen_ai.tool.name': {},
+                **({'gen_ai.tool.call.id': {}} if run_context.tool_call_id else {}),
+            },
+        }
+    )
 
-    with run_context.tracer.start_as_current_span('running output function', attributes=attributes) as span:
+    with run_context.tracer.start_as_current_span(
+        instrumentation_names.get_output_tool_span_name(tool_name), attributes=attributes
+    ) as span:
         try:
             output = await function_schema.call(args, run_context)
         except ModelRetry as r:
@@ -135,7 +149,7 @@ async def execute_traced_output_function(
             from .models.instrumented import InstrumentedModel
 
             span.set_attribute(
-                'tool_response',
+                instrumentation_names.tool_result_attr,
                 output if isinstance(output, str) else json.dumps(InstrumentedModel.serialize_any(output)),
             )
 

pydantic_ai/_run_context.py

@@ -8,6 +8,8 @@ from typing import TYPE_CHECKING, Generic
 from opentelemetry.trace import NoOpTracer, Tracer
 from typing_extensions import TypeVar
 
+from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION
+
 from . import _utils, messages as _messages
 
 if TYPE_CHECKING:
@@ -36,6 +38,8 @@ class RunContext(Generic[AgentDepsT]):
     """The tracer to use for tracing the run."""
     trace_include_content: bool = False
     """Whether to include the content of the messages in the trace."""
+    instrumentation_version: int = DEFAULT_INSTRUMENTATION_VERSION
+    """Instrumentation settings version, if instrumentation is enabled."""
     retries: dict[str, int] = field(default_factory=dict)
     """Number of retries for each tool so far."""
     tool_call_id: str | None = None

pydantic_ai/_thinking_part.py

@@ -1,6 +1,6 @@
 from __future__ import annotations as _annotations
 
-from pydantic_ai.messages import TextPart, ThinkingPart
+from pydantic_ai import TextPart, ThinkingPart
 
 
 def split_content_into_text_and_thinking(content: str, thinking_tags: tuple[str, str]) -> list[ThinkingPart | TextPart]:

pydantic_ai/_tool_manager.py

@@ -12,6 +12,7 @@ from pydantic import ValidationError
 from typing_extensions import assert_never
 
 from . import messages as _messages
+from ._instrumentation import InstrumentationNames
 from ._run_context import AgentDepsT, RunContext
 from .exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior
 from .messages import ToolCallPart
@@ -115,6 +116,7 @@ class ToolManager(Generic[AgentDepsT]):
                 wrap_validation_errors,
                 self.ctx.tracer,
                 self.ctx.trace_include_content,
+                self.ctx.instrumentation_version,
                 usage_limits,
             )
 
@@ -203,15 +205,18 @@ class ToolManager(Generic[AgentDepsT]):
         allow_partial: bool,
         wrap_validation_errors: bool,
         tracer: Tracer,
-        include_content: bool = False,
+        include_content: bool,
+        instrumentation_version: int,
         usage_limits: UsageLimits | None = None,
     ) -> Any:
         """See <https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span>."""
+        instrumentation_names = InstrumentationNames.for_version(instrumentation_version)
+
         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 include_content else {}),
+            **({instrumentation_names.tool_arguments_attr: 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(
@@ -220,8 +225,8 @@ class ToolManager(Generic[AgentDepsT]):
                     'properties': {
                         **(
                             {
-                                'tool_arguments': {'type': 'object'},
-                                'tool_response': {'type': 'object'},
+                                instrumentation_names.tool_arguments_attr: {'type': 'object'},
+                                instrumentation_names.tool_result_attr: {'type': 'object'},
                             }
                             if include_content
                             else {}
@@ -232,18 +237,21 @@ class ToolManager(Generic[AgentDepsT]):
                 }
             ),
         }
-        with tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
+        with tracer.start_as_current_span(
+            instrumentation_names.get_tool_span_name(call.tool_name),
+            attributes=span_attributes,
+        ) as span:
             try:
                 tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors, usage_limits)
             except ToolRetryError as e:
                 part = e.tool_retry
                 if include_content and span.is_recording():
-                    span.set_attribute('tool_response', part.model_response())
+                    span.set_attribute(instrumentation_names.tool_result_attr, part.model_response())
                 raise e
 
             if include_content and span.is_recording():
                 span.set_attribute(
-                    'tool_response',
+                    instrumentation_names.tool_result_attr,
                     tool_result
                     if isinstance(tool_result, str)
                     else _messages.tool_return_ta.dump_json(tool_result).decode(),

pydantic_ai/_utils.py

@@ -6,7 +6,7 @@ import inspect
 import re
 import time
 import uuid
-from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterable, Iterator
 from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass, fields, is_dataclass
 from datetime import datetime, timezone
@@ -70,16 +70,33 @@ def check_object_json_schema(schema: JsonSchemaValue) -> ObjectJsonSchema:
 
     if schema.get('type') == 'object':
         return schema
-    elif schema.get('$ref') is not None:
-        maybe_result = schema.get('$defs', {}).get(schema['$ref'][8:])  # This removes the initial "#/$defs/".
-
-        if "'$ref': '#/$defs/" in str(maybe_result):
-            return schema  # We can't remove the $defs because the schema contains other references
-        return maybe_result
+    elif ref := schema.get('$ref'):
+        prefix = '#/$defs/'
+        # Return the referenced schema unless it contains additional nested references.
+        if (
+            ref.startswith(prefix)
+            and (resolved := schema.get('$defs', {}).get(ref[len(prefix) :]))
+            and resolved.get('type') == 'object'
+            and not _contains_ref(resolved)
+        ):
+            return resolved
+        return schema
     else:
         raise UserError('Schema must be an object')
 
 
+def _contains_ref(obj: JsonSchemaValue | list[JsonSchemaValue]) -> bool:
+    """Recursively check if an object contains any $ref keys."""
+    items: Iterable[JsonSchemaValue]
+    if isinstance(obj, dict):
+        if '$ref' in obj:
+            return True
+        items = obj.values()
+    else:
+        items = obj
+    return any(isinstance(item, dict | list) and _contains_ref(item) for item in items)  # pyright: ignore[reportUnknownArgumentType]
+
+
 T = TypeVar('T')