pydantic-ai-slim 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

pydantic_ai/_run_context.py

@@ -0,0 +1,56 @@
+from __future__ import annotations as _annotations
+
+import dataclasses
+from collections.abc import Sequence
+from dataclasses import field
+from typing import TYPE_CHECKING, Generic
+
+from typing_extensions import TypeVar
+
+from . import _utils, messages as _messages
+
+if TYPE_CHECKING:
+    from .models import Model
+    from .result import Usage
+
+AgentDepsT = TypeVar('AgentDepsT', default=None, contravariant=True)
+"""Type variable for agent dependencies."""
+
+
+@dataclasses.dataclass(repr=False)
+class RunContext(Generic[AgentDepsT]):
+    """Information about the current call."""
+
+    deps: AgentDepsT
+    """Dependencies for the agent."""
+    model: Model
+    """The model used in this run."""
+    usage: Usage
+    """LLM usage associated with the run."""
+    prompt: str | Sequence[_messages.UserContent] | None
+    """The original user prompt passed to the run."""
+    messages: list[_messages.ModelMessage] = field(default_factory=list)
+    """Messages exchanged in the conversation so far."""
+    tool_call_id: str | None = None
+    """The ID of the tool call."""
+    tool_name: str | None = None
+    """Name of the tool being called."""
+    retry: int = 0
+    """Number of retries so far."""
+    run_step: int = 0
+    """The current step in the run."""
+
+    def replace_with(
+        self,
+        retry: int | None = None,
+        tool_name: str | None | _utils.Unset = _utils.UNSET,
+    ) -> RunContext[AgentDepsT]:
+        # Create a new `RunContext` a new `retry` value and `tool_name`.
+        kwargs = {}
+        if retry is not None:
+            kwargs['retry'] = retry
+        if tool_name is not _utils.UNSET:  # pragma: no branch
+            kwargs['tool_name'] = tool_name
+        return dataclasses.replace(self, **kwargs)
+
+    __repr__ = _utils.dataclasses_no_defaults_repr

pydantic_ai/_system_prompt.py

@@ -6,7 +6,8 @@ from dataclasses import dataclass, field
 from typing import Any, Callable, Generic, cast
 
 from . import _utils
-from .tools import AgentDepsT, RunContext, SystemPromptFunc
+from ._run_context import AgentDepsT, RunContext
+from .tools import SystemPromptFunc
 
 
 @dataclass

pydantic_ai/_utils.py

@@ -3,6 +3,7 @@ from __future__ import annotations as _annotations
 import asyncio
 import functools
 import inspect
+import re
 import time
 import uuid
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterator
@@ -16,7 +17,17 @@ from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union, overlo
 from anyio.to_thread import run_sync
 from pydantic import BaseModel, TypeAdapter
 from pydantic.json_schema import JsonSchemaValue
-from typing_extensions import ParamSpec, TypeAlias, TypeGuard, TypeIs, is_typeddict
+from typing_extensions import (
+    ParamSpec,
+    TypeAlias,
+    TypeGuard,
+    TypeIs,
+    get_args,
+    get_origin,
+    is_typeddict,
+)
+from typing_inspection import typing_objects
+from typing_inspection.introspection import is_union_origin
 
 from pydantic_graph._utils import AbstractSpan
 
@@ -327,3 +338,102 @@ def is_async_callable(obj: Any) -> Any:
         obj = obj.func
 
     return inspect.iscoroutinefunction(obj) or (callable(obj) and inspect.iscoroutinefunction(obj.__call__))  # type: ignore
+
+
+def _update_mapped_json_schema_refs(s: dict[str, Any], name_mapping: dict[str, str]) -> None:
+    """Update $refs in a schema to use the new names from name_mapping."""
+    if '$ref' in s:
+        ref = s['$ref']
+        if ref.startswith('#/$defs/'):  # pragma: no branch
+            original_name = ref[8:]  # Remove '#/$defs/'
+            new_name = name_mapping.get(original_name, original_name)
+            s['$ref'] = f'#/$defs/{new_name}'
+
+    # Recursively update refs in properties
+    if 'properties' in s:
+        props: dict[str, dict[str, Any]] = s['properties']
+        for prop in props.values():
+            _update_mapped_json_schema_refs(prop, name_mapping)
+
+    # Handle arrays
+    if 'items' in s and isinstance(s['items'], dict):
+        items: dict[str, Any] = s['items']
+        _update_mapped_json_schema_refs(items, name_mapping)
+    if 'prefixItems' in s:
+        prefix_items: list[dict[str, Any]] = s['prefixItems']
+        for item in prefix_items:
+            _update_mapped_json_schema_refs(item, name_mapping)
+
+    # Handle unions
+    for union_type in ['anyOf', 'oneOf']:
+        if union_type in s:
+            union_items: list[dict[str, Any]] = s[union_type]
+            for item in union_items:
+                _update_mapped_json_schema_refs(item, name_mapping)
+
+
+def merge_json_schema_defs(schemas: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], dict[str, dict[str, Any]]]:
+    """Merges the `$defs` from different JSON schemas into a single deduplicated `$defs`, handling name collisions of `$defs` that are not the same, and rewrites `$ref`s to point to the new `$defs`.
+
+    Returns a tuple of the rewritten schemas and a dictionary of the new `$defs`.
+    """
+    all_defs: dict[str, dict[str, Any]] = {}
+    rewritten_schemas: list[dict[str, Any]] = []
+
+    for schema in schemas:
+        if '$defs' not in schema:
+            rewritten_schemas.append(schema)
+            continue
+
+        schema = schema.copy()
+        defs = schema.pop('$defs', None)
+        schema_name_mapping: dict[str, str] = {}
+
+        # Process definitions and build mapping
+        for name, def_schema in defs.items():
+            if name not in all_defs:
+                all_defs[name] = def_schema
+                schema_name_mapping[name] = name
+            elif def_schema != all_defs[name]:
+                new_name = name
+                if title := schema.get('title'):
+                    new_name = f'{title}_{name}'
+
+                i = 1
+                original_new_name = new_name
+                new_name = f'{new_name}_{i}'
+                while new_name in all_defs:
+                    i += 1
+                    new_name = f'{original_new_name}_{i}'
+
+                all_defs[new_name] = def_schema
+                schema_name_mapping[name] = new_name
+
+        _update_mapped_json_schema_refs(schema, schema_name_mapping)
+        rewritten_schemas.append(schema)
+
+    return rewritten_schemas, all_defs
+
+
+def strip_markdown_fences(text: str) -> str:
+    if text.startswith('{'):
+        return text
+
+    regex = r'```(?:\w+)?\n(\{.*\})\n```'
+    match = re.search(regex, text, re.DOTALL)
+    if match:
+        return match.group(1)
+
+    return text
+
+
+def get_union_args(tp: Any) -> tuple[Any, ...]:
+    """Extract the arguments of a Union type if `tp` is a union, otherwise return an empty tuple."""
+    if typing_objects.is_typealiastype(tp):
+        tp = tp.__value__
+
+    origin = get_origin(tp)
+    if is_union_origin(origin):
+        return get_args(tp)
+    else:
+        return ()

pydantic_ai/agent.py

@@ -14,6 +14,7 @@ from opentelemetry.trace import NoOpTracer, use_span
 from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import Literal, Never, Self, TypeIs, TypeVar, deprecated
 
+from pydantic_ai.profiles import ModelProfile
 from pydantic_graph import End, Graph, GraphRun, GraphRunContext
 from pydantic_graph._utils import get_event_loop
 
@@ -30,7 +31,8 @@ from . import (
 )
 from ._agent_graph import HistoryProcessor
 from .models.instrumented import InstrumentationSettings, InstrumentedModel, instrument_model
-from .result import FinalResult, OutputDataT, StreamedRunResult
+from .output import OutputDataT, OutputSpec
+from .result import FinalResult, StreamedRunResult
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
     AgentDepsT,
@@ -91,7 +93,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
     """Class for defining "agents" - a way to have a specific type of "conversation" with an LLM.
 
     Agents are generic in the dependency type they take [`AgentDepsT`][pydantic_ai.tools.AgentDepsT]
-    and the result data type they return, [`OutputDataT`][pydantic_ai.result.OutputDataT].
+    and the output type they return, [`OutputDataT`][pydantic_ai.output.OutputDataT].
 
     By default, if neither generic parameter is customised, agents have type `Agent[None, str]`.
 
@@ -128,7 +130,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
     be merged with this value, with the runtime argument taking priority.
     """
 
-    output_type: _output.OutputType[OutputDataT]
+    output_type: OutputSpec[OutputDataT]
     """
     The type of data output by agent runs, used to validate the data returned by the model, defaults to `str`.
     """
@@ -141,7 +143,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
     _deps_type: type[AgentDepsT] = dataclasses.field(repr=False)
     _deprecated_result_tool_name: str | None = dataclasses.field(repr=False)
     _deprecated_result_tool_description: str | None = dataclasses.field(repr=False)
-    _output_schema: _output.OutputSchema[OutputDataT] | None = dataclasses.field(repr=False)
+    _output_schema: _output.BaseOutputSchema[OutputDataT] = dataclasses.field(repr=False)
     _output_validators: list[_output.OutputValidator[AgentDepsT, OutputDataT]] = dataclasses.field(repr=False)
     _instructions: str | None = dataclasses.field(repr=False)
     _instructions_functions: list[_system_prompt.SystemPromptRunner[AgentDepsT]] = dataclasses.field(repr=False)
@@ -163,7 +165,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         model: models.Model | models.KnownModelName | str | None = None,
         *,
-        output_type: _output.OutputType[OutputDataT] = str,
+        output_type: OutputSpec[OutputDataT] = str,
         instructions: str
         | _system_prompt.SystemPromptFunc[AgentDepsT]
         | Sequence[str | _system_prompt.SystemPromptFunc[AgentDepsT]]
@@ -324,8 +326,14 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
             warnings.warn('`result_retries` is deprecated, use `max_result_retries` instead', DeprecationWarning)
             output_retries = result_retries
 
+        default_output_mode = (
+            self.model.profile.default_structured_output_mode if isinstance(self.model, models.Model) else None
+        )
         self._output_schema = _output.OutputSchema[OutputDataT].build(
-            output_type, self._deprecated_result_tool_name, self._deprecated_result_tool_description
+            output_type,
+            default_mode=default_output_mode,
+            name=self._deprecated_result_tool_name,
+            description=self._deprecated_result_tool_description,
         )
         self._output_validators = []
 
@@ -382,7 +390,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: _output.OutputType[RunOutputDataT],
+        output_type: OutputSpec[RunOutputDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -412,7 +420,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: _output.OutputType[RunOutputDataT] | None = None,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -500,7 +508,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None,
         *,
-        output_type: _output.OutputType[RunOutputDataT],
+        output_type: OutputSpec[RunOutputDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -532,7 +540,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: _output.OutputType[RunOutputDataT] | None = None,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -631,7 +639,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
 
         deps = self._get_deps(deps)
         new_message_index = len(message_history) if message_history else 0
-        output_schema = self._prepare_output_schema(output_type)
+        output_schema = self._prepare_output_schema(output_type, model_used.profile)
 
         output_type_ = output_type or self.output_type
 
@@ -674,14 +682,20 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         )
 
         async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
-            if self._instructions is None and not self._instructions_functions:
-                return None
+            parts = [
+                self._instructions,
+                *[await func.run(run_context) for func in self._instructions_functions],
+            ]
 
-            instructions = [self._instructions] if self._instructions else []
-            for instructions_runner in self._instructions_functions:
-                instructions.append(await instructions_runner.run(run_context))
-            concatenated_instructions = '\n'.join(instruction for instruction in instructions if instruction)
-            return concatenated_instructions.strip() if concatenated_instructions else None
+            model_profile = model_used.profile
+            if isinstance(output_schema, _output.PromptedOutputSchema):
+                instructions = output_schema.instructions(model_profile.prompted_output_template)
+                parts.append(instructions)
+
+            parts = [p for p in parts if p]
+            if not parts:
+                return None
+            return '\n\n'.join(parts).strip()
 
         # Copy the function tools so that retry state is agent-run-specific
         # Note that the retry count is reset to 0 when this happens due to the `default=0` and `init=False`.
@@ -779,7 +793,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: _output.OutputType[RunOutputDataT] | None = None,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -809,7 +823,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: _output.OutputType[RunOutputDataT] | None = None,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -892,7 +906,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent],
         *,
-        output_type: _output.OutputType[RunOutputDataT],
+        output_type: OutputSpec[RunOutputDataT],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -923,7 +937,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
-        output_type: _output.OutputType[RunOutputDataT] | None = None,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -1002,10 +1016,13 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
                             async for maybe_part_event in streamed_response:
                                 if isinstance(maybe_part_event, _messages.PartStartEvent):
                                     new_part = maybe_part_event.part
-                                    if isinstance(new_part, _messages.TextPart):
-                                        if _output.allow_text_output(output_schema):
-                                            return FinalResult(s, None, None)
-                                    elif isinstance(new_part, _messages.ToolCallPart) and output_schema:
+                                    if isinstance(new_part, _messages.TextPart) and isinstance(
+                                        output_schema, _output.TextOutputSchema
+                                    ):
+                                        return FinalResult(s, None, None)
+                                    elif isinstance(new_part, _messages.ToolCallPart) and isinstance(
+                                        output_schema, _output.ToolOutputSchema
+                                    ):  # pragma: no branch
                                         for call, _ in output_schema.find_tool([new_part]):
                                             return FinalResult(s, call.tool_name, call.tool_call_id)
                             return None
@@ -1561,8 +1578,8 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         if tool.name in self._function_tools:
             raise exceptions.UserError(f'Tool name conflicts with existing tool: {tool.name!r}')
 
-        if self._output_schema and tool.name in self._output_schema.tools:
-            raise exceptions.UserError(f'Tool name conflicts with result schema name: {tool.name!r}')
+        if tool.name in self._output_schema.tools:
+            raise exceptions.UserError(f'Tool name conflicts with output tool name: {tool.name!r}')
 
         self._function_tools[tool.name] = tool
 
@@ -1637,18 +1654,23 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         raise AttributeError('The `last_run_messages` attribute has been removed, use `capture_run_messages` instead.')
 
     def _prepare_output_schema(
-        self, output_type: _output.OutputType[RunOutputDataT] | None
-    ) -> _output.OutputSchema[RunOutputDataT] | None:
+        self, output_type: OutputSpec[RunOutputDataT] | None, model_profile: ModelProfile
+    ) -> _output.OutputSchema[RunOutputDataT]:
         if output_type is not None:
             if self._output_validators:
                 raise exceptions.UserError('Cannot set a custom run `output_type` when the agent has output validators')
-            return _output.OutputSchema[RunOutputDataT].build(
+            schema = _output.OutputSchema[RunOutputDataT].build(
                 output_type,
-                self._deprecated_result_tool_name,
-                self._deprecated_result_tool_description,
+                name=self._deprecated_result_tool_name,
+                description=self._deprecated_result_tool_description,
+                default_mode=model_profile.default_structured_output_mode,
             )
         else:
-            return self._output_schema  # pyright: ignore[reportReturnType]
+            schema = self._output_schema.with_default_mode(model_profile.default_structured_output_mode)
+
+        schema.raise_if_unsupported(model_profile)
+
+        return schema  # pyright: ignore[reportReturnType]
 
     @staticmethod
     def is_model_request_node(
@@ -1691,14 +1713,23 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         return isinstance(node, End)
 
     @asynccontextmanager
-    async def run_mcp_servers(self) -> AsyncIterator[None]:
+    async def run_mcp_servers(
+        self, model: models.Model | models.KnownModelName | str | None = None
+    ) -> AsyncIterator[None]:
         """Run [`MCPServerStdio`s][pydantic_ai.mcp.MCPServerStdio] so they can be used by the agent.
 
         Returns: a context manager to start and shutdown the servers.
         """
+        try:
+            sampling_model: models.Model | None = self._get_model(model)
+        except exceptions.UserError:  # pragma: no cover
+            sampling_model = None
+
         exit_stack = AsyncExitStack()
         try:
             for mcp_server in self._mcp_servers:
+                if sampling_model is not None:  # pragma: no branch
+                    mcp_server.sampling_model = sampling_model
                 await exit_stack.enter_async_context(mcp_server)
             yield
         finally: