pydantic-ai-slim 1.0.2__tar.gz → 1.0.4__tar.gz

{pydantic_ai_slim-1.0.2 → pydantic_ai_slim-1.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 1.0.2
+Version: 1.0.4
 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
@@ -33,7 +33,7 @@ Requires-Dist: genai-prices>=0.0.23
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==1.0.2
+Requires-Dist: pydantic-graph==1.0.4
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -51,13 +51,13 @@ 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'
+Requires-Dist: cohere>=5.17.0; (platform_system != 'Emscripten') and extra == 'cohere'
 Provides-Extra: dbos
 Requires-Dist: dbos>=1.13.0; extra == 'dbos'
 Provides-Extra: duckduckgo
 Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==1.0.2; extra == 'evals'
+Requires-Dist: pydantic-evals==1.0.4; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.31.0; extra == 'google'
 Provides-Extra: groq
@@ -69,7 +69,7 @@ Requires-Dist: logfire[httpx]>=3.14.1; extra == 'logfire'
 Provides-Extra: mcp
 Requires-Dist: mcp>=1.12.3; extra == 'mcp'
 Provides-Extra: mistral
-Requires-Dist: mistralai>=1.9.2; extra == 'mistral'
+Requires-Dist: mistralai>=1.9.10; extra == 'mistral'
 Provides-Extra: openai
 Requires-Dist: openai>=1.99.9; extra == 'openai'
 Provides-Extra: retries

{pydantic_ai_slim-1.0.2 → pydantic_ai_slim-1.0.4}/pydantic_ai/_output.py

@@ -7,7 +7,7 @@ from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Generic, Literal, cast, overload
 
-from pydantic import TypeAdapter, ValidationError
+from pydantic import Json, TypeAdapter, ValidationError
 from pydantic_core import SchemaValidator, to_json
 from typing_extensions import Self, TypedDict, TypeVar, assert_never
 
@@ -624,21 +624,33 @@ class ObjectOutputProcessor(BaseOutputProcessor[OutputDataT]):
             json_schema = self._function_schema.json_schema
             json_schema['description'] = self._function_schema.description
         else:
-            type_adapter: TypeAdapter[Any]
+            json_schema_type_adapter: TypeAdapter[Any]
+            validation_type_adapter: TypeAdapter[Any]
             if _utils.is_model_like(output):
-                type_adapter = TypeAdapter(output)
+                json_schema_type_adapter = validation_type_adapter = TypeAdapter(output)
             else:
                 self.outer_typed_dict_key = 'response'
+                output_type: type[OutputDataT] = cast(type[OutputDataT], output)
+
                 response_data_typed_dict = TypedDict(  # noqa: UP013
                     'response_data_typed_dict',
-                    {'response': cast(type[OutputDataT], output)},  # pyright: ignore[reportInvalidTypeForm]
+                    {'response': output_type},  # pyright: ignore[reportInvalidTypeForm]
+                )
+                json_schema_type_adapter = TypeAdapter(response_data_typed_dict)
+
+                # More lenient validator: allow either the native type or a JSON string containing it
+                # i.e. `response: OutputDataT | Json[OutputDataT]`, as some models don't follow the schema correctly,
+                # e.g. `BedrockConverseModel('us.meta.llama3-2-11b-instruct-v1:0')`
+                response_validation_typed_dict = TypedDict(  # noqa: UP013
+                    'response_validation_typed_dict',
+                    {'response': output_type | Json[output_type]},  # pyright: ignore[reportInvalidTypeForm]
                 )
-                type_adapter = TypeAdapter(response_data_typed_dict)
+                validation_type_adapter = TypeAdapter(response_validation_typed_dict)
 
             # Really a PluggableSchemaValidator, but it's API-compatible
-            self.validator = cast(SchemaValidator, type_adapter.validator)
+            self.validator = cast(SchemaValidator, validation_type_adapter.validator)
             json_schema = _utils.check_object_json_schema(
-                type_adapter.json_schema(schema_generator=GenerateToolJsonSchema)
+                json_schema_type_adapter.json_schema(schema_generator=GenerateToolJsonSchema)
             )
 
             if self.outer_typed_dict_key:

{pydantic_ai_slim-1.0.2 → pydantic_ai_slim-1.0.4}/pydantic_ai/_parts_manager.py

@@ -156,6 +156,7 @@ class ModelResponsePartsManager:
         content: str | None = None,
         id: str | None = None,
         signature: str | None = None,
+        provider_name: str | None = None,
     ) -> ModelResponseStreamEvent:
         """Handle incoming thinking content, creating or updating a ThinkingPart in the manager as appropriate.
 
@@ -170,6 +171,7 @@ class ModelResponsePartsManager:
             content: The thinking content to append to the appropriate ThinkingPart.
             id: An optional id for the thinking part.
             signature: An optional signature for the thinking content.
+            provider_name: An optional provider name for the thinking part.
 
         Returns:
             A `PartStartEvent` if a new part was created, or a `PartDeltaEvent` if an existing part was updated.
@@ -196,27 +198,23 @@ class ModelResponsePartsManager:
                 existing_thinking_part_and_index = existing_part, part_index
 
         if existing_thinking_part_and_index is None:
-            if content is not None:
+            if content is not None or signature is not None:
                 # There is no existing thinking part that should be updated, so create a new one
                 new_part_index = len(self._parts)
-                part = ThinkingPart(content=content, id=id, signature=signature)
+                part = ThinkingPart(content=content or '', id=id, signature=signature, provider_name=provider_name)
                 if vendor_part_id is not None:  # pragma: no branch
                     self._vendor_id_to_part_index[vendor_part_id] = new_part_index
                 self._parts.append(part)
                 return PartStartEvent(index=new_part_index, part=part)
             else:
-                raise UnexpectedModelBehavior('Cannot create a ThinkingPart with no content')
+                raise UnexpectedModelBehavior('Cannot create a ThinkingPart with no content or signature')
         else:
-            if content is not None:
-                # Update the existing ThinkingPart with the new content delta
+            if content is not None or signature is not None:
+                # Update the existing ThinkingPart with the new content and/or signature delta
                 existing_thinking_part, part_index = existing_thinking_part_and_index
-                part_delta = ThinkingPartDelta(content_delta=content)
-                self._parts[part_index] = part_delta.apply(existing_thinking_part)
-                return PartDeltaEvent(index=part_index, delta=part_delta)
-            elif signature is not None:
-                # Update the existing ThinkingPart with the new signature delta
-                existing_thinking_part, part_index = existing_thinking_part_and_index
-                part_delta = ThinkingPartDelta(signature_delta=signature)
+                part_delta = ThinkingPartDelta(
+                    content_delta=content, signature_delta=signature, provider_name=provider_name
+                )
                 self._parts[part_index] = part_delta.apply(existing_thinking_part)
                 return PartDeltaEvent(index=part_index, delta=part_delta)
             else:

{pydantic_ai_slim-1.0.2 → pydantic_ai_slim-1.0.4}/pydantic_ai/_tool_manager.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
 from dataclasses import dataclass, field, replace
 from typing import Any, Generic
 
@@ -16,6 +19,8 @@ from .tools import ToolDefinition
 from .toolsets.abstract import AbstractToolset, ToolsetTool
 from .usage import UsageLimits
 
+_sequential_tool_calls_ctx_var: ContextVar[bool] = ContextVar('sequential_tool_calls', default=False)
+
 
 @dataclass
 class ToolManager(Generic[AgentDepsT]):
@@ -30,6 +35,16 @@ class ToolManager(Generic[AgentDepsT]):
     failed_tools: set[str] = field(default_factory=set)
     """Names of tools that failed in this run step."""
 
+    @classmethod
+    @contextmanager
+    def sequential_tool_calls(cls) -> Iterator[None]:
+        """Run tool calls sequentially during the context."""
+        token = _sequential_tool_calls_ctx_var.set(True)
+        try:
+            yield
+        finally:
+            _sequential_tool_calls_ctx_var.reset(token)
+
     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 self.ctx is not None:
@@ -58,7 +73,9 @@ class ToolManager(Generic[AgentDepsT]):
 
     def should_call_sequentially(self, calls: list[ToolCallPart]) -> bool:
         """Whether to require sequential tool calls for a list of tool calls."""
-        return any(tool_def.sequential for call in calls if (tool_def := self.get_tool_def(call.tool_name)))
+        return _sequential_tool_calls_ctx_var.get() or any(
+            tool_def.sequential for call in calls if (tool_def := self.get_tool_def(call.tool_name))
+        )
 
     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."""

{pydantic_ai_slim-1.0.2 → pydantic_ai_slim-1.0.4}/pydantic_ai/ag_ui.py

@@ -8,7 +8,7 @@ from __future__ import annotations
 
 import json
 import uuid
-from collections.abc import AsyncIterator, Callable, Iterable, Mapping, Sequence
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping, Sequence
 from dataclasses import Field, dataclass, replace
 from http import HTTPStatus
 from typing import (
@@ -17,14 +17,16 @@ from typing import (
     Final,
     Generic,
     Protocol,
+    TypeAlias,
     TypeVar,
     runtime_checkable,
 )
 
 from pydantic import BaseModel, ValidationError
 
+from . import _utils
 from ._agent_graph import CallToolsNode, ModelRequestNode
-from .agent import AbstractAgent, AgentRun
+from .agent import AbstractAgent, AgentRun, AgentRunResult
 from .exceptions import UserError
 from .messages import (
     FunctionToolResultEvent,
@@ -68,9 +70,8 @@ try:
         TextMessageContentEvent,
         TextMessageEndEvent,
         TextMessageStartEvent,
-        # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
-        # ThinkingEndEvent,
-        # ThinkingStartEvent,
+        ThinkingEndEvent,
+        ThinkingStartEvent,
         ThinkingTextMessageContentEvent,
         ThinkingTextMessageEndEvent,
         ThinkingTextMessageStartEvent,
@@ -108,6 +109,7 @@ __all__ = [
     'StateDeps',
     'StateHandler',
     'AGUIApp',
+    'OnCompleteFunc',
     'handle_ag_ui_request',
     'run_ag_ui',
 ]
@@ -115,6 +117,9 @@ __all__ = [
 SSE_CONTENT_TYPE: Final[str] = 'text/event-stream'
 """Content type header value for Server-Sent Events (SSE)."""
 
+OnCompleteFunc: TypeAlias = Callable[[AgentRunResult[Any]], None] | Callable[[AgentRunResult[Any]], Awaitable[None]]
+"""Callback function type that receives the `AgentRunResult` of the completed run. Can be sync or async."""
+
 
 class AGUIApp(Generic[AgentDepsT, OutputDataT], Starlette):
     """ASGI application for running Pydantic AI agents with AG-UI protocol support."""
@@ -221,6 +226,7 @@ async def handle_ag_ui_request(
     usage: RunUsage | None = None,
     infer_name: bool = True,
     toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    on_complete: OnCompleteFunc | None = None,
 ) -> Response:
     """Handle an AG-UI request by running the agent and returning a streaming response.
 
@@ -237,6 +243,8 @@ async def handle_ag_ui_request(
         usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
         infer_name: Whether to try to infer the agent name from the call frame if it's not set.
         toolsets: Optional additional toolsets for this run.
+        on_complete: Optional callback function called when the agent run completes successfully.
+            The callback receives the completed [`AgentRunResult`][pydantic_ai.agent.AgentRunResult] and can access `all_messages()` and other result data.
 
     Returns:
         A streaming Starlette response with AG-UI protocol events.
@@ -264,6 +272,7 @@ async def handle_ag_ui_request(
             usage=usage,
             infer_name=infer_name,
             toolsets=toolsets,
+            on_complete=on_complete,
         ),
         media_type=accept,
     )
@@ -282,6 +291,7 @@ async def run_ag_ui(
     usage: RunUsage | None = None,
     infer_name: bool = True,
     toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    on_complete: OnCompleteFunc | None = None,
 ) -> AsyncIterator[str]:
     """Run the agent with the AG-UI run input and stream AG-UI protocol events.
 
@@ -299,6 +309,8 @@ async def run_ag_ui(
         usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
         infer_name: Whether to try to infer the agent name from the call frame if it's not set.
         toolsets: Optional additional toolsets for this run.
+        on_complete: Optional callback function called when the agent run completes successfully.
+            The callback receives the completed [`AgentRunResult`][pydantic_ai.agent.AgentRunResult] and can access `all_messages()` and other result data.
 
     Yields:
         Streaming event chunks encoded as strings according to the accept header value.
@@ -357,6 +369,12 @@ async def run_ag_ui(
         ) as run:
             async for event in _agent_stream(run):
                 yield encoder.encode(event)
+
+        if on_complete is not None and run.result is not None:
+            if _utils.is_async_callable(on_complete):
+                await on_complete(run.result)
+            else:
+                await _utils.run_in_executor(on_complete, run.result)
     except _RunError as e:
         yield encoder.encode(
             RunErrorEvent(message=e.message, code=e.code),
@@ -396,10 +414,9 @@ async def _agent_stream(run: AgentRun[AgentDepsT, Any]) -> AsyncIterator[BaseEve
                     yield stream_ctx.part_end
                     stream_ctx.part_end = None
                 if stream_ctx.thinking:
-                    # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
-                    # yield ThinkingEndEvent(
-                    #     type=EventType.THINKING_END,
-                    # )
+                    yield ThinkingEndEvent(
+                        type=EventType.THINKING_END,
+                    )
                     stream_ctx.thinking = False
         elif isinstance(node, CallToolsNode):
             async with node.stream(run.ctx) as handle_stream:
@@ -431,10 +448,9 @@ async def _handle_model_request_event(  # noqa: C901
         part = agent_event.part
         if isinstance(part, ThinkingPart):  # pragma: no branch
             if not stream_ctx.thinking:
-                # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
-                # yield ThinkingStartEvent(
-                #     type=EventType.THINKING_START,
-                # )
+                yield ThinkingStartEvent(
+                    type=EventType.THINKING_START,
+                )
                 stream_ctx.thinking = True
 
             if part.content:
@@ -450,10 +466,9 @@ async def _handle_model_request_event(  # noqa: C901
                 )
         else:
             if stream_ctx.thinking:
-                # TODO: Enable once https://github.com/ag-ui-protocol/ag-ui/issues/289 is resolved.
-                # yield ThinkingEndEvent(
-                #     type=EventType.THINKING_END,
-                # )
+                yield ThinkingEndEvent(
+                    type=EventType.THINKING_END,
+                )
                 stream_ctx.thinking = False
 
             if isinstance(part, TextPart):

{pydantic_ai_slim-1.0.2 → pydantic_ai_slim-1.0.4}/pydantic_ai/agent/abstract.py

@@ -21,6 +21,7 @@ from .. import (
     result,
     usage as _usage,
 )
+from .._tool_manager import ToolManager
 from ..output import OutputDataT, OutputSpec
 from ..result import AgentStream, FinalResult, StreamedRunResult
 from ..run import AgentRun, AgentRunResult
@@ -714,6 +715,13 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
                             self.name = name
                             return
 
+    @staticmethod
+    @contextmanager
+    def sequential_tool_calls() -> Iterator[None]:
+        """Run tool calls sequentially during the context."""
+        with ToolManager.sequential_tool_calls():
+            yield
+
     @staticmethod
     def is_model_request_node(
         node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]],

{pydantic_ai_slim-1.0.2 → pydantic_ai_slim-1.0.4}/pydantic_ai/durable_exec/dbos/_agent.py

@@ -13,7 +13,6 @@ from pydantic_ai import (
     models,
     usage as _usage,
 )
-from pydantic_ai._run_context import AgentDepsT
 from pydantic_ai.agent import AbstractAgent, AgentRun, AgentRunResult, EventStreamHandler, RunOutputDataT, WrapperAgent
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.mcp import MCPServer
@@ -22,6 +21,7 @@ from pydantic_ai.output import OutputDataT, OutputSpec
 from pydantic_ai.result import StreamedRunResult
 from pydantic_ai.settings import ModelSettings
 from pydantic_ai.tools import (
+    AgentDepsT,
     DeferredToolResults,
     RunContext,
     Tool,
@@ -218,7 +218,10 @@ class DBOSAgent(WrapperAgent[AgentDepsT, OutputDataT], DBOSConfiguredInstance):
     @contextmanager
     def _dbos_overrides(self) -> Iterator[None]:
         # Override with DBOSModel and DBOSMCPServer in the toolsets.
-        with super().override(model=self._model, toolsets=self._toolsets, tools=[]):
+        with (
+            super().override(model=self._model, toolsets=self._toolsets, tools=[]),
+            self.sequential_tool_calls(),
+        ):
             yield
 
     @overload

{pydantic_ai_slim-1.0.2 → pydantic_ai_slim-1.0.4}/pydantic_ai/durable_exec/temporal/_agent.py

@@ -21,7 +21,6 @@ from pydantic_ai import (
     models,
     usage as _usage,
 )
-from pydantic_ai._run_context import AgentDepsT
 from pydantic_ai.agent import AbstractAgent, AgentRun, AgentRunResult, EventStreamHandler, RunOutputDataT, WrapperAgent
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.models import Model
@@ -29,6 +28,7 @@ from pydantic_ai.output import OutputDataT, OutputSpec
 from pydantic_ai.result import StreamedRunResult
 from pydantic_ai.settings import ModelSettings
 from pydantic_ai.tools import (
+    AgentDepsT,
     DeferredToolResults,
     RunContext,
     Tool,

{pydantic_ai_slim-1.0.2 → pydantic_ai_slim-1.0.4}/pydantic_ai/messages.py

@@ -895,7 +895,18 @@ class ThinkingPart:
     signature: str | None = None
     """The signature of the thinking.
 
-    The signature is only available on the Anthropic models.
+    Supported by:
+
+    * Anthropic (corresponds to the `signature` field)
+    * Bedrock (corresponds to the `signature` field)
+    * Google (corresponds to the `thought_signature` field)
+    * OpenAI (corresponds to the `encrypted_content` field)
+    """
+
+    provider_name: str | None = None
+    """The name of the provider that generated the response.
+
+    Signatures are only sent back to the same provider.
     """
 
     part_kind: Literal['thinking'] = 'thinking'
@@ -980,7 +991,10 @@ class BuiltinToolCallPart(BaseToolCallPart):
     _: KW_ONLY
 
     provider_name: str | None = None
-    """The name of the provider that generated the response."""
+    """The name of the provider that generated the response.
+
+    Built-in tool calls are only sent back to the same provider.
+    """
 
     part_kind: Literal['builtin-tool-call'] = 'builtin-tool-call'
     """Part type identifier, this is available on all parts as a discriminator."""
@@ -1198,6 +1212,12 @@ class ThinkingPartDelta:
     Note this is never treated as a delta — it can replace None.
     """
 
+    provider_name: str | None = None
+    """Optional provider name for the thinking part.
+
+    Signatures are only sent back to the same provider.
+    """
+
     part_delta_kind: Literal['thinking'] = 'thinking'
     """Part delta type identifier, used as a discriminator."""
 
@@ -1222,14 +1242,18 @@ class ThinkingPartDelta:
         if isinstance(part, ThinkingPart):
             new_content = part.content + self.content_delta if self.content_delta else part.content
             new_signature = self.signature_delta if self.signature_delta is not None else part.signature
-            return replace(part, content=new_content, signature=new_signature)
+            new_provider_name = self.provider_name if self.provider_name is not None else part.provider_name
+            return replace(part, content=new_content, signature=new_signature, provider_name=new_provider_name)
         elif isinstance(part, ThinkingPartDelta):
             if self.content_delta is None and self.signature_delta is None:
                 raise ValueError('Cannot apply ThinkingPartDelta with no content or signature')
-            if self.signature_delta is not None:
-                return replace(part, signature_delta=self.signature_delta)
             if self.content_delta is not None:
-                return replace(part, content_delta=self.content_delta)
+                part = replace(part, content_delta=(part.content_delta or '') + self.content_delta)
+            if self.signature_delta is not None:
+                part = replace(part, signature_delta=self.signature_delta)
+            if self.provider_name is not None:
+                part = replace(part, provider_name=self.provider_name)
+            return part
         raise ValueError(  # pragma: no cover
             f'Cannot apply ThinkingPartDeltas to non-ThinkingParts or non-ThinkingPartDeltas ({part=}, {self=})'
         )

{pydantic_ai_slim-1.0.2 → pydantic_ai_slim-1.0.4}/pydantic_ai/models/__init__.py

@@ -718,7 +718,11 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         )
         provider = 'google-vertex'
 
-    if provider == 'cohere':
+    if provider == 'gateway':
+        from ..providers.gateway import infer_model as infer_model_from_gateway
+
+        return infer_model_from_gateway(model_name)
+    elif provider == 'cohere':
         from .cohere import CohereModel
 
         return CohereModel(model_name, provider=provider)