pydantic-ai-slim 0.4.8__tar.gz → 0.4.9__tar.gz

{pydantic_ai_slim-0.4.8 → pydantic_ai_slim-0.4.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.4.8
+Version: 0.4.9
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 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
@@ -30,7 +30,7 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.4.8
+Requires-Dist: pydantic-graph==0.4.9
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -48,11 +48,10 @@ Requires-Dist: prompt-toolkit>=3; extra == 'cli'
 Requires-Dist: rich>=13; extra == 'cli'
 Provides-Extra: cohere
 Requires-Dist: cohere>=5.16.0; (platform_system != 'Emscripten') and extra == 'cohere'
-Requires-Dist: tokenizers<=0.21.2; extra == 'cohere'
 Provides-Extra: duckduckgo
 Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.4.8; extra == 'evals'
+Requires-Dist: pydantic-evals==0.4.9; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.24.0; extra == 'google'
 Provides-Extra: groq

{pydantic_ai_slim-0.4.8 → pydantic_ai_slim-0.4.9}/pydantic_ai/_tool_manager.py

@@ -2,18 +2,17 @@ from __future__ import annotations
 
 import json
 from collections.abc import Iterable
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, field, replace
 from typing import Any, Generic
 
 from pydantic import ValidationError
 from typing_extensions import assert_never
 
-from pydantic_ai.output import DeferredToolCalls
-
 from . import messages as _messages
 from ._run_context import AgentDepsT, RunContext
 from .exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior
 from .messages import ToolCallPart
+from .output import DeferredToolCalls
 from .tools import ToolDefinition
 from .toolsets.abstract import AbstractToolset, ToolsetTool
 
@@ -28,6 +27,8 @@ class ToolManager(Generic[AgentDepsT]):
     """The toolset that provides the tools for this run step."""
     tools: dict[str, ToolsetTool[AgentDepsT]]
     """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]:
@@ -40,7 +41,10 @@ class ToolManager(Generic[AgentDepsT]):
 
     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."""
-        return await self.__class__.build(self.toolset, replace(ctx, retries=self.ctx.retries))
+        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))
 
     @property
     def tool_defs(self) -> list[ToolDefinition]:
@@ -54,20 +58,25 @@ class ToolManager(Generic[AgentDepsT]):
         except KeyError:
             return None
 
-    async def handle_call(self, call: ToolCallPart, allow_partial: bool = False) -> Any:
+    async def handle_call(
+        self, call: ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
+    ) -> Any:
         """Handle a tool call by validating the arguments, calling the tool, and handling retries.
 
         Args:
             call: The tool call part to handle.
             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 (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)
+            return await self._call_tool(call, allow_partial, wrap_validation_errors)
         else:
-            return await self._call_tool_traced(call, allow_partial)
+            return await self._call_tool_traced(call, allow_partial, wrap_validation_errors)
 
-    async def _call_tool(self, call: ToolCallPart, allow_partial: bool = False) -> Any:
+    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:
@@ -92,7 +101,7 @@ class ToolManager(Generic[AgentDepsT]):
             else:
                 args_dict = validator.validate_python(call.args or {}, allow_partial=pyd_allow_partial)
 
-            output = await self.toolset.call_tool(name, args_dict, ctx, tool)
+            return await self.toolset.call_tool(name, args_dict, ctx, tool)
         except (ValidationError, ModelRetry) as e:
             max_retries = tool.max_retries if tool is not None else 1
             current_retry = self.ctx.retries.get(name, 0)
@@ -100,30 +109,33 @@ class ToolManager(Generic[AgentDepsT]):
             if current_retry == max_retries:
                 raise UnexpectedModelBehavior(f'Tool {name!r} exceeded max retries count of {max_retries}') from e
             else:
-                if isinstance(e, ValidationError):
-                    m = _messages.RetryPromptPart(
-                        tool_name=name,
-                        content=e.errors(include_url=False, include_context=False),
-                        tool_call_id=call.tool_call_id,
-                    )
-                    e = ToolRetryError(m)
-                elif isinstance(e, ModelRetry):
-                    m = _messages.RetryPromptPart(
-                        tool_name=name,
-                        content=e.message,
-                        tool_call_id=call.tool_call_id,
-                    )
-                    e = ToolRetryError(m)
-                else:
-                    assert_never(e)
+                if wrap_validation_errors:
+                    if isinstance(e, ValidationError):
+                        m = _messages.RetryPromptPart(
+                            tool_name=name,
+                            content=e.errors(include_url=False, include_context=False),
+                            tool_call_id=call.tool_call_id,
+                        )
+                        e = ToolRetryError(m)
+                    elif isinstance(e, ModelRetry):
+                        m = _messages.RetryPromptPart(
+                            tool_name=name,
+                            content=e.message,
+                            tool_call_id=call.tool_call_id,
+                        )
+                        e = ToolRetryError(m)
+                    else:
+                        assert_never(e)
+
+                if not allow_partial:
+                    # If we're validating partial arguments, we don't want to count this as a failed tool as it may still succeed once the full arguments are received.
+                    self.failed_tools.add(name)
 
-                self.ctx.retries[name] = current_retry + 1
                 raise e
-        else:
-            self.ctx.retries.pop(name, None)
-            return output
 
-    async def _call_tool_traced(self, call: ToolCallPart, allow_partial: bool = False) -> Any:
+    async def _call_tool_traced(
+        self, call: ToolCallPart, allow_partial: bool = False, wrap_validation_errors: bool = True
+    ) -> 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,
@@ -152,7 +164,7 @@ class ToolManager(Generic[AgentDepsT]):
         }
         with self.ctx.tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
             try:
-                tool_result = await self._call_tool(call, allow_partial)
+                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():

{pydantic_ai_slim-0.4.8 → pydantic_ai_slim-0.4.9}/pydantic_ai/ag_ui.py

@@ -9,11 +9,13 @@ from __future__ import annotations
 import json
 import uuid
 from collections.abc import Iterable, Mapping, Sequence
-from dataclasses import dataclass, field
+from dataclasses import Field, dataclass, field, replace
 from http import HTTPStatus
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
+    ClassVar,
     Final,
     Generic,
     Protocol,
@@ -21,6 +23,11 @@ from typing import (
     runtime_checkable,
 )
 
+from pydantic_ai.exceptions import UserError
+
+if TYPE_CHECKING:
+    pass
+
 try:
     from ag_ui.core import (
         AssistantMessage,
@@ -288,8 +295,24 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
             if not run_input.messages:
                 raise _NoMessagesError
 
+            raw_state: dict[str, Any] = run_input.state or {}
             if isinstance(deps, StateHandler):
-                deps.state = run_input.state
+                if isinstance(deps.state, BaseModel):
+                    try:
+                        state = type(deps.state).model_validate(raw_state)
+                    except ValidationError as e:  # pragma: no cover
+                        raise _InvalidStateError from e
+                else:
+                    state = raw_state
+
+                deps = replace(deps, state=state)
+            elif raw_state:
+                raise UserError(
+                    f'AG-UI state is provided but `deps` of type `{type(deps).__name__}` does not implement the `StateHandler` protocol: it needs to be a dataclass with a non-optional `state` field.'
+                )
+            else:
+                # `deps` not being a `StateHandler` is OK if there is no state.
+                pass
 
             messages = _messages_from_ag_ui(run_input.messages)
 
@@ -311,7 +334,7 @@ class _Adapter(Generic[AgentDepsT, OutputDataT]):
             yield encoder.encode(
                 RunErrorEvent(message=e.message, code=e.code),
             )
-        except Exception as e:  # pragma: no cover
+        except Exception as e:
             yield encoder.encode(
                 RunErrorEvent(message=str(e)),
             )
@@ -531,7 +554,11 @@ def _messages_from_ag_ui(messages: list[Message]) -> list[ModelMessage]:
 
 @runtime_checkable
 class StateHandler(Protocol):
-    """Protocol for state handlers in agent runs."""
+    """Protocol for state handlers in agent runs. Requires the class to be a dataclass with a `state` field."""
+
+    # Has to be a dataclass so we can use `replace` to update the state.
+    # From https://github.com/python/typeshed/blob/9ab7fde0a0cd24ed7a72837fcb21093b811b80d8/stdlib/_typeshed/__init__.pyi#L352
+    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
 
     @property
     def state(self) -> State:
@@ -558,6 +585,7 @@ StateT = TypeVar('StateT', bound=BaseModel)
 """Type variable for the state type, which must be a subclass of `BaseModel`."""
 
 
+@dataclass
 class StateDeps(Generic[StateT]):
     """Provides AG-UI state management.
 
@@ -570,42 +598,7 @@ class StateDeps(Generic[StateT]):
     Implements the `StateHandler` protocol.
     """
 
-    def __init__(self, default: StateT) -> None:
-        """Initialize the state with the provided state type."""
-        self._state = default
-
-    @property
-    def state(self) -> StateT:
-        """Get the current state of the agent run.
-
-        Returns:
-            The current run state.
-        """
-        return self._state
-
-    @state.setter
-    def state(self, state: State) -> None:
-        """Set the state of the agent run.
-
-        This method is called to update the state of the agent run with the
-        provided state.
-
-        Implements the `StateHandler` protocol.
-
-        Args:
-            state: The run state, which must be `None` or model validate for the state type.
-
-        Raises:
-            InvalidStateError: If `state` does not validate.
-        """
-        if state is None:
-            # If state is None, we keep the current state, which will be the default state.
-            return
-
-        try:
-            self._state = type(self._state).model_validate(state)
-        except ValidationError as e:  # pragma: no cover
-            raise _InvalidStateError from e
+    state: StateT
 
 
 @dataclass(repr=False)

{pydantic_ai_slim-0.4.8 → pydantic_ai_slim-0.4.9}/pydantic_ai/profiles/google.py

@@ -49,7 +49,6 @@ class GoogleJsonSchemaTransformer(JsonSchemaTransformer):
             )
 
         schema.pop('title', None)
-        schema.pop('default', None)
         schema.pop('$schema', None)
         if (const := schema.pop('const', None)) is not None:
             # Gemini doesn't support const, but it does support enum with a single value

{pydantic_ai_slim-0.4.8 → pydantic_ai_slim-0.4.9}/pydantic_ai/result.py

@@ -67,7 +67,7 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
                 except ValidationError:
                     pass
         if self._final_result_event is not None:  # pragma: no branch
-            yield await self._validate_response(self._raw_stream_response.get(), allow_partial=False)
+            yield await self._validate_response(self._raw_stream_response.get())
 
     async def stream_responses(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[_messages.ModelResponse]:
         """Asynchronously stream the (unvalidated) model responses for the agent."""
@@ -128,7 +128,7 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
         async for _ in self:
             pass
 
-        return await self._validate_response(self._raw_stream_response.get(), allow_partial=False)
+        return await self._validate_response(self._raw_stream_response.get())
 
     async def _validate_response(self, message: _messages.ModelResponse, *, allow_partial: bool = False) -> OutputDataT:
         """Validate a structured result message."""
@@ -150,7 +150,9 @@ class AgentStream(Generic[AgentDepsT, OutputDataT]):
                 raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
                     f'Invalid response, unable to find tool call for {output_tool_name!r}'
                 )
-            return await self._tool_manager.handle_call(tool_call, allow_partial=allow_partial)
+            return await self._tool_manager.handle_call(
+                tool_call, allow_partial=allow_partial, wrap_validation_errors=False
+            )
         elif deferred_tool_calls := self._tool_manager.get_deferred_tool_calls(message.parts):
             if not self._output_schema.allows_deferred_tool_calls:
                 raise exceptions.UserError(

{pydantic_ai_slim-0.4.8 → pydantic_ai_slim-0.4.9}/pyproject.toml

@@ -63,12 +63,7 @@ dependencies = [
 logfire = ["logfire>=3.11.0"]
 # Models
 openai = ["openai>=1.92.0"]
-cohere = [
-    "cohere>=5.16.0; platform_system != 'Emscripten'",
-    # Remove once all wheels for 0.21.4+ are built successfully
-    # https://github.com/huggingface/tokenizers/actions/runs/16570140346/job/46860152621
-    "tokenizers<=0.21.2",
-]
+cohere = ["cohere>=5.16.0; platform_system != 'Emscripten'"]
 vertexai = ["google-auth>=2.36.0", "requests>=2.32.2"]
 google = ["google-genai>=1.24.0"]
 anthropic = ["anthropic>=0.52.0"]