pydantic-ai-slim 1.0.1__tar.gz → 1.0.3__tar.gz

{pydantic_ai_slim-1.0.1 → pydantic_ai_slim-1.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 1.0.1
+Version: 1.0.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
@@ -29,11 +29,11 @@ Classifier: Topic :: Internet
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Requires-Dist: exceptiongroup; python_version < '3.11'
-Requires-Dist: genai-prices>=0.0.22
+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.1
+Requires-Dist: pydantic-graph==1.0.3
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -51,11 +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.1; extra == 'evals'
+Requires-Dist: pydantic-evals==1.0.3; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.31.0; extra == 'google'
 Provides-Extra: groq
@@ -67,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.1 → pydantic_ai_slim-1.0.3}/pydantic_ai/_agent_graph.py

@@ -2,6 +2,8 @@ from __future__ import annotations as _annotations
 
 import asyncio
 import dataclasses
+import inspect
+from asyncio import Task
 from collections import defaultdict, deque
 from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
@@ -740,7 +742,6 @@ async def process_function_tools(  # noqa: C901
     deferred_tool_results: dict[str, DeferredToolResult] = {}
     if build_run_context(ctx).tool_call_approved and ctx.deps.tool_call_results is not None:
         deferred_tool_results = ctx.deps.tool_call_results
-
         # Deferred tool calls are "run" as well, by reading their value from the tool call results
         calls_to_run.extend(tool_calls_by_kind['external'])
         calls_to_run.extend(tool_calls_by_kind['unapproved'])
@@ -819,7 +820,6 @@ async def _call_tools(
     for call in tool_calls:
         yield _messages.FunctionToolCallEvent(call)
 
-    # Run all tool tasks in parallel
     with tracer.start_as_current_span(
         'running tools',
         attributes={
@@ -827,39 +827,58 @@ async def _call_tools(
             'logfire.msg': f'running {len(tool_calls)} tool{"" if len(tool_calls) == 1 else "s"}',
         },
     ):
-        tasks = [
-            asyncio.create_task(
-                _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id), usage_limits),
-                name=call.tool_name,
-            )
-            for call in tool_calls
-        ]
-
-        pending = tasks
-        while pending:
-            done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
-            for task in done:
-                index = tasks.index(task)
-                try:
-                    tool_part, tool_user_part = task.result()
-                except exceptions.CallDeferred:
-                    deferred_calls_by_index[index] = 'external'
-                except exceptions.ApprovalRequired:
-                    deferred_calls_by_index[index] = 'unapproved'
-                else:
-                    yield _messages.FunctionToolResultEvent(tool_part)
 
-                    tool_parts_by_index[index] = tool_part
-                    if tool_user_part:
-                        user_parts_by_index[index] = tool_user_part
+        async def handle_call_or_result(
+            coro_or_task: Awaitable[
+                tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, _messages.UserPromptPart | None]
+            ]
+            | Task[tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, _messages.UserPromptPart | None]],
+            index: int,
+        ) -> _messages.HandleResponseEvent | None:
+            try:
+                tool_part, tool_user_part = (
+                    (await coro_or_task) if inspect.isawaitable(coro_or_task) else coro_or_task.result()
+                )
+            except exceptions.CallDeferred:
+                deferred_calls_by_index[index] = 'external'
+            except exceptions.ApprovalRequired:
+                deferred_calls_by_index[index] = 'unapproved'
+            else:
+                tool_parts_by_index[index] = tool_part
+                if tool_user_part:
+                    user_parts_by_index[index] = tool_user_part
+
+                return _messages.FunctionToolResultEvent(tool_part)
+
+        if tool_manager.should_call_sequentially(tool_calls):
+            for index, call in enumerate(tool_calls):
+                if event := await handle_call_or_result(
+                    _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id), usage_limits),
+                    index,
+                ):
+                    yield event
+
+        else:
+            tasks = [
+                asyncio.create_task(
+                    _call_tool(tool_manager, call, deferred_tool_results.get(call.tool_call_id), usage_limits),
+                    name=call.tool_name,
+                )
+                for call in tool_calls
+            ]
+
+            pending = tasks
+            while pending:
+                done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
+                for task in done:
+                    index = tasks.index(task)
+                    if event := await handle_call_or_result(coro_or_task=task, index=index):
+                        yield event
 
     # We append the results at the end, rather than as they are received, to retain a consistent ordering
     # This is mostly just to simplify testing
-    for k in sorted(tool_parts_by_index):
-        output_parts.append(tool_parts_by_index[k])
-
-    for k in sorted(user_parts_by_index):
-        output_parts.append(user_parts_by_index[k])
+    output_parts.extend([tool_parts_by_index[k] for k in sorted(tool_parts_by_index)])
+    output_parts.extend([user_parts_by_index[k] for k in sorted(user_parts_by_index)])
 
     for k in sorted(deferred_calls_by_index):
         output_deferred_calls[deferred_calls_by_index[k]].append(tool_calls[k])

{pydantic_ai_slim-1.0.1 → pydantic_ai_slim-1.0.3}/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.1 → pydantic_ai_slim-1.0.3}/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.
@@ -199,7 +201,7 @@ class ModelResponsePartsManager:
             if content 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, 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)
@@ -207,16 +209,12 @@ class ModelResponsePartsManager:
             else:
                 raise UnexpectedModelBehavior('Cannot create a ThinkingPart with no content')
         else:
-            if content is not None:
-                # Update the existing ThinkingPart with the new content 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
+            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(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.1 → pydantic_ai_slim-1.0.3}/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:
@@ -56,6 +71,12 @@ class ToolManager(Generic[AgentDepsT]):
 
         return [tool.tool_def for tool in self.tools.values()]
 
+    def should_call_sequentially(self, calls: list[ToolCallPart]) -> bool:
+        """Whether to require sequential tool calls for a list of tool calls."""
+        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."""
         if self.tools is None:

{pydantic_ai_slim-1.0.1 → pydantic_ai_slim-1.0.3}/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.1 → pydantic_ai_slim-1.0.3}/pydantic_ai/agent/__init__.py

@@ -1119,6 +1119,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
         strict: bool | None = None,
+        sequential: bool = False,
         requires_approval: bool = False,
     ) -> Any:
         """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
@@ -1164,6 +1165,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
             schema_generator: The JSON schema generator class to use for this tool. Defaults to `GenerateToolJsonSchema`.
             strict: Whether to enforce JSON schema compliance (only affects OpenAI).
                 See [`ToolDefinition`][pydantic_ai.tools.ToolDefinition] for more info.
+            sequential: Whether the function requires a sequential/serial execution environment. Defaults to False.
             requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False.
                 See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
         """
@@ -1180,6 +1182,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
                 require_parameter_descriptions,
                 schema_generator,
                 strict,
+                sequential,
                 requires_approval,
             )
             return func_

{pydantic_ai_slim-1.0.1 → pydantic_ai_slim-1.0.3}/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.3/pydantic_ai/durable_exec/dbos/__init__.py

@@ -0,0 +1,6 @@
+from ._agent import DBOSAgent
+from ._mcp_server import DBOSMCPServer
+from ._model import DBOSModel
+from ._utils import StepConfig
+
+__all__ = ['DBOSAgent', 'DBOSModel', 'DBOSMCPServer', 'StepConfig']