pydantic-ai-slim 0.6.2__tar.gz → 0.7.1__tar.gz

{pydantic_ai_slim-0.6.2 → pydantic_ai_slim-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.6.2
+Version: 0.7.1
 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.6.2
+Requires-Dist: pydantic-graph==0.7.1
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
@@ -51,7 +51,7 @@ Requires-Dist: cohere>=5.16.0; (platform_system != 'Emscripten') and extra == 'c
 Provides-Extra: duckduckgo
 Requires-Dist: ddgs>=9.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.6.2; extra == 'evals'
+Requires-Dist: pydantic-evals==0.7.1; extra == 'evals'
 Provides-Extra: google
 Requires-Dist: google-genai>=1.28.0; extra == 'google'
 Provides-Extra: groq
@@ -59,17 +59,19 @@ Requires-Dist: groq>=0.25.0; extra == 'groq'
 Provides-Extra: huggingface
 Requires-Dist: huggingface-hub[inference]>=0.33.5; extra == 'huggingface'
 Provides-Extra: logfire
-Requires-Dist: logfire>=3.11.0; extra == 'logfire'
+Requires-Dist: logfire>=3.14.1; extra == 'logfire'
 Provides-Extra: mcp
 Requires-Dist: mcp>=1.10.0; (python_version >= '3.10') and extra == 'mcp'
 Provides-Extra: mistral
 Requires-Dist: mistralai>=1.9.2; extra == 'mistral'
 Provides-Extra: openai
-Requires-Dist: openai>=1.92.0; extra == 'openai'
+Requires-Dist: openai>=1.99.9; extra == 'openai'
 Provides-Extra: retries
 Requires-Dist: tenacity>=8.2.3; extra == 'retries'
 Provides-Extra: tavily
 Requires-Dist: tavily-python>=0.5.0; extra == 'tavily'
+Provides-Extra: temporal
+Requires-Dist: temporalio>=1.15.0; extra == 'temporal'
 Provides-Extra: vertexai
 Requires-Dist: google-auth>=2.36.0; extra == 'vertexai'
 Requires-Dist: requests>=2.32.2; extra == 'vertexai'

{pydantic_ai_slim-0.6.2 → pydantic_ai_slim-0.7.1}/pydantic_ai/_a2a.py

@@ -27,7 +27,7 @@ from pydantic_ai.messages import (
     VideoUrl,
 )
 
-from .agent import Agent, AgentDepsT, OutputDataT
+from .agent import AbstractAgent, AgentDepsT, OutputDataT
 
 # AgentWorker output type needs to be invariant for use in both parameter and return positions
 WorkerOutputT = TypeVar('WorkerOutputT')
@@ -59,7 +59,9 @@ except ImportError as _import_error:
 
 
 @asynccontextmanager
-async def worker_lifespan(app: FastA2A, worker: Worker, agent: Agent[AgentDepsT, OutputDataT]) -> AsyncIterator[None]:
+async def worker_lifespan(
+    app: FastA2A, worker: Worker, agent: AbstractAgent[AgentDepsT, OutputDataT]
+) -> AsyncIterator[None]:
     """Custom lifespan that runs the worker during application startup.
 
     This ensures the worker is started and ready to process tasks as soon as the application starts.
@@ -70,7 +72,7 @@ async def worker_lifespan(app: FastA2A, worker: Worker, agent: Agent[AgentDepsT,
 
 
 def agent_to_a2a(
-    agent: Agent[AgentDepsT, OutputDataT],
+    agent: AbstractAgent[AgentDepsT, OutputDataT],
     *,
     storage: Storage | None = None,
     broker: Broker | None = None,
@@ -116,7 +118,7 @@ def agent_to_a2a(
 class AgentWorker(Worker[list[ModelMessage]], Generic[WorkerOutputT, AgentDepsT]):
     """A worker that uses an agent to execute tasks."""
 
-    agent: Agent[AgentDepsT, WorkerOutputT]
+    agent: AbstractAgent[AgentDepsT, WorkerOutputT]
 
     async def run_task(self, params: TaskSendParams) -> None:
         task = await self.storage.load_task(params['id'])

{pydantic_ai_slim-0.6.2 → pydantic_ai_slim-0.7.1}/pydantic_ai/_agent_graph.py

@@ -306,10 +306,18 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         self,
         ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, T]],
     ) -> AsyncIterator[result.AgentStream[DepsT, T]]:
-        async with self._stream(ctx) as streamed_response:
+        assert not self._did_stream, 'stream() should only be called once per node'
+
+        model_settings, model_request_parameters, message_history, run_context = await self._prepare_request(ctx)
+        async with ctx.deps.model.request_stream(
+            message_history, model_settings, model_request_parameters, run_context
+        ) as streamed_response:
+            self._did_stream = True
+            ctx.state.usage.requests += 1
             agent_stream = result.AgentStream[DepsT, T](
                 streamed_response,
                 ctx.deps.output_schema,
+                model_request_parameters,
                 ctx.deps.output_validators,
                 build_run_context(ctx),
                 ctx.deps.usage_limits,
@@ -321,28 +329,6 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
             async for _ in agent_stream:
                 pass
 
-    @asynccontextmanager
-    async def _stream(
-        self,
-        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, T]],
-    ) -> AsyncIterator[models.StreamedResponse]:
-        assert not self._did_stream, 'stream() should only be called once per node'
-
-        model_settings, model_request_parameters = await self._prepare_request(ctx)
-        model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
-        message_history = await _process_message_history(
-            ctx.state.message_history, ctx.deps.history_processors, build_run_context(ctx)
-        )
-        async with ctx.deps.model.request_stream(
-            message_history, model_settings, model_request_parameters
-        ) as streamed_response:
-            self._did_stream = True
-            ctx.state.usage.requests += 1
-            yield streamed_response
-            # In case the user didn't manually consume the full stream, ensure it is fully consumed here,
-            # otherwise usage won't be properly counted:
-            async for _ in streamed_response:
-                pass
         model_response = streamed_response.get()
 
         self._finish_handling(ctx, model_response)
@@ -354,11 +340,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         if self._result is not None:
             return self._result  # pragma: no cover
 
-        model_settings, model_request_parameters = await self._prepare_request(ctx)
-        model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
-        message_history = await _process_message_history(
-            ctx.state.message_history, ctx.deps.history_processors, build_run_context(ctx)
-        )
+        model_settings, model_request_parameters, message_history, _ = await self._prepare_request(ctx)
         model_response = await ctx.deps.model.request(message_history, model_settings, model_request_parameters)
         ctx.state.usage.incr(_usage.Usage())
 
@@ -366,19 +348,33 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
 
     async def _prepare_request(
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
-    ) -> tuple[ModelSettings | None, models.ModelRequestParameters]:
+    ) -> tuple[ModelSettings | None, models.ModelRequestParameters, list[_messages.ModelMessage], RunContext[DepsT]]:
         ctx.state.message_history.append(self.request)
 
-        # Check usage
-        if ctx.deps.usage_limits:  # pragma: no branch
-            ctx.deps.usage_limits.check_before_request(ctx.state.usage)
-
-        # Increment run_step
         ctx.state.run_step += 1
 
+        run_context = build_run_context(ctx)
+
         model_settings = merge_model_settings(ctx.deps.model_settings, None)
+
         model_request_parameters = await _prepare_request_parameters(ctx)
-        return model_settings, model_request_parameters
+        model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
+
+        message_history = await _process_message_history(ctx.state, ctx.deps.history_processors, run_context)
+
+        usage = ctx.state.usage
+        if ctx.deps.usage_limits.count_tokens_before_request:
+            # Copy to avoid modifying the original usage object with the counted usage
+            usage = dataclasses.replace(usage)
+
+            counted_usage = await ctx.deps.model.count_tokens(
+                message_history, ctx.deps.model_settings, model_request_parameters
+            )
+            usage.incr(counted_usage)
+
+        ctx.deps.usage_limits.check_before_request(usage)
+
+        return model_settings, model_request_parameters, message_history, run_context
 
     def _finish_handling(
         self,
@@ -802,7 +798,7 @@ _messages_ctx_var: ContextVar[_RunMessages] = ContextVar('var')
 
 @contextmanager
 def capture_run_messages() -> Iterator[list[_messages.ModelMessage]]:
-    """Context manager to access the messages used in a [`run`][pydantic_ai.Agent.run], [`run_sync`][pydantic_ai.Agent.run_sync], or [`run_stream`][pydantic_ai.Agent.run_stream] call.
+    """Context manager to access the messages used in a [`run`][pydantic_ai.agent.AbstractAgent.run], [`run_sync`][pydantic_ai.agent.AbstractAgent.run_sync], or [`run_stream`][pydantic_ai.agent.AbstractAgent.run_stream] call.
 
     Useful when a run may raise an exception, see [model errors](../agents.md#model-errors) for more information.
 
@@ -868,11 +864,12 @@ def build_agent_graph(
 
 
 async def _process_message_history(
-    messages: list[_messages.ModelMessage],
+    state: GraphAgentState,
     processors: Sequence[HistoryProcessor[DepsT]],
     run_context: RunContext[DepsT],
 ) -> list[_messages.ModelMessage]:
     """Process message history through a sequence of processors."""
+    messages = state.message_history
     for processor in processors:
         takes_ctx = is_takes_ctx(processor)
 
@@ -889,4 +886,7 @@ async def _process_message_history(
             else:
                 sync_processor = cast(_HistoryProcessorSync, processor)
                 messages = await run_in_executor(sync_processor, messages)
+
+    # Replaces the message history in the state with the processed messages
+    state.message_history = messages
     return messages

{pydantic_ai_slim-0.6.2 → pydantic_ai_slim-0.7.1}/pydantic_ai/_cli.py

@@ -16,7 +16,7 @@ from typing_inspection.introspection import get_literal_values
 
 from . import __version__
 from ._run_context import AgentDepsT
-from .agent import Agent
+from .agent import AbstractAgent, Agent
 from .exceptions import UserError
 from .messages import ModelMessage
 from .models import KnownModelName, infer_model
@@ -220,7 +220,7 @@ Special prompts:
 
 async def run_chat(
     stream: bool,
-    agent: Agent[AgentDepsT, OutputDataT],
+    agent: AbstractAgent[AgentDepsT, OutputDataT],
     console: Console,
     code_theme: str,
     prog_name: str,
@@ -263,7 +263,7 @@ async def run_chat(
 
 
 async def ask_agent(
-    agent: Agent[AgentDepsT, OutputDataT],
+    agent: AbstractAgent[AgentDepsT, OutputDataT],
     prompt: str,
     stream: bool,
     console: Console,

{pydantic_ai_slim-0.6.2 → pydantic_ai_slim-0.7.1}/pydantic_ai/_output.py

@@ -977,6 +977,14 @@ class OutputToolset(AbstractToolset[AgentDepsT]):
         self.max_retries = max_retries
         self.output_validators = output_validators or []
 
+    @property
+    def id(self) -> str | None:
+        return '<output>'  # pragma: no cover
+
+    @property
+    def label(self) -> str:
+        return "the agent's output tools"
+
     async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
         return {
             tool_def.name: ToolsetTool(

{pydantic_ai_slim-0.6.2 → pydantic_ai_slim-0.7.1}/pydantic_ai/_tool_manager.py

@@ -41,6 +41,9 @@ 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."""
+        if ctx.run_step == self.ctx.run_step:
+            return self
+
         retries = {
             failed_tool_name: self.ctx.retries.get(failed_tool_name, 0) + 1 for failed_tool_name in self.failed_tools
         }

{pydantic_ai_slim-0.6.2 → pydantic_ai_slim-0.7.1}/pydantic_ai/ag_ui.py

@@ -25,7 +25,7 @@ from typing import (
 from pydantic import BaseModel, ValidationError
 
 from ._agent_graph import CallToolsNode, ModelRequestNode
-from .agent import Agent, AgentRun
+from .agent import AbstractAgent, AgentRun
 from .exceptions import UserError
 from .messages import (
     AgentStreamEvent,
@@ -72,6 +72,7 @@ try:
         ThinkingTextMessageContentEvent,
         ThinkingTextMessageEndEvent,
         ThinkingTextMessageStartEvent,
+        Tool as AGUITool,
         ToolCallArgsEvent,
         ToolCallEndEvent,
         ToolCallResultEvent,
@@ -99,6 +100,7 @@ except ImportError as e:  # pragma: no cover
         'you can use the `ag-ui` optional group — `pip install "pydantic-ai-slim[ag-ui]"`'
     ) from e
 
+
 __all__ = [
     'SSE_CONTENT_TYPE',
     'StateDeps',
@@ -117,7 +119,7 @@ class AGUIApp(Generic[AgentDepsT, OutputDataT], Starlette):
 
     def __init__(
         self,
-        agent: Agent[AgentDepsT, OutputDataT],
+        agent: AbstractAgent[AgentDepsT, OutputDataT],
         *,
         # Agent.iter parameters.
         output_type: OutputSpec[Any] | None = None,
@@ -206,7 +208,7 @@ class AGUIApp(Generic[AgentDepsT, OutputDataT], Starlette):
 
 
 async def handle_ag_ui_request(
-    agent: Agent[AgentDepsT, Any],
+    agent: AbstractAgent[AgentDepsT, Any],
     request: Request,
     *,
     output_type: OutputSpec[Any] | None = None,
@@ -266,7 +268,7 @@ async def handle_ag_ui_request(
 
 
 async def run_ag_ui(
-    agent: Agent[AgentDepsT, Any],
+    agent: AbstractAgent[AgentDepsT, Any],
     run_input: RunAgentInput,
     accept: str = SSE_CONTENT_TYPE,
     *,
@@ -304,16 +306,7 @@ async def run_ag_ui(
         # AG-UI tools can't be prefixed as that would result in a mismatch between the tool names in the
         # Pydantic AI events and actual AG-UI tool names, preventing the tool from being called. If any
         # conflicts arise, the AG-UI tool should be renamed or a `PrefixedToolset` used for local toolsets.
-        toolset = DeferredToolset[AgentDepsT](
-            [
-                ToolDefinition(
-                    name=tool.name,
-                    description=tool.description,
-                    parameters_json_schema=tool.parameters,
-                )
-                for tool in run_input.tools
-            ]
-        )
+        toolset = _AGUIFrontendToolset[AgentDepsT](run_input.tools)
         toolsets = [*toolsets, toolset] if toolsets else [toolset]
 
     try:
@@ -686,3 +679,21 @@ class _ToolCallNotFoundError(_RunError, ValueError):
             message=f'Tool call with ID {tool_call_id} not found in the history.',
             code='tool_call_not_found',
         )
+
+
+class _AGUIFrontendToolset(DeferredToolset[AgentDepsT]):
+    def __init__(self, tools: list[AGUITool]):
+        super().__init__(
+            [
+                ToolDefinition(
+                    name=tool.name,
+                    description=tool.description,
+                    parameters_json_schema=tool.parameters,
+                )
+                for tool in tools
+            ]
+        )
+
+    @property
+    def label(self) -> str:
+        return 'the AG-UI frontend tools'  # pragma: no cover