pydantic-ai-slim 1.0.14__py3-none-any.whl → 1.0.15__py3-none-any.whl

pydantic_ai/__init__.py

@@ -9,7 +9,14 @@ from .agent import (
     UserPromptNode,
     capture_run_messages,
 )
-from .builtin_tools import CodeExecutionTool, UrlContextTool, WebSearchTool, WebSearchUserLocation
+from .builtin_tools import (
+    CodeExecutionTool,
+    ImageGenerationTool,
+    MemoryTool,
+    UrlContextTool,
+    WebSearchTool,
+    WebSearchUserLocation,
+)
 from .exceptions import (
     AgentRunError,
     ApprovalRequired,
@@ -30,11 +37,13 @@ from .messages import (
     BaseToolCallPart,
     BaseToolReturnPart,
     BinaryContent,
+    BinaryImage,
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentFormat,
     DocumentMediaType,
     DocumentUrl,
+    FilePart,
     FileUrl,
     FinalResultEvent,
     FinishReason,
@@ -79,6 +88,7 @@ from .profiles import (
     ModelProfile,
     ModelProfileSpec,
 )
+from .run import AgentRun, AgentRunResult, AgentRunResultEvent
 from .settings import ModelSettings
 from .tools import DeferredToolRequests, DeferredToolResults, RunContext, Tool, ToolApproved, ToolDefinition, ToolDenied
 from .toolsets import (
@@ -131,6 +141,7 @@ __all__ = (
     'DocumentMediaType',
     'DocumentUrl',
     'FileUrl',
+    'FilePart',
     'FinalResultEvent',
     'FinishReason',
     'FunctionToolCallEvent',
@@ -139,6 +150,7 @@ __all__ = (
     'ImageFormat',
     'ImageMediaType',
     'ImageUrl',
+    'BinaryImage',
     'ModelMessage',
     'ModelMessagesTypeAdapter',
     'ModelRequest',
@@ -197,6 +209,8 @@ __all__ = (
     'WebSearchUserLocation',
     'UrlContextTool',
     'CodeExecutionTool',
+    'ImageGenerationTool',
+    'MemoryTool',
     # output
     'ToolOutput',
     'NativeOutput',
@@ -211,5 +225,9 @@ __all__ = (
     'RunUsage',
     'RequestUsage',
     'UsageLimits',
+    # run
+    'AgentRun',
+    'AgentRunResult',
+    'AgentRunResultEvent',
 )
 __version__ = _metadata_version('pydantic_ai_slim')

pydantic_ai/_agent_graph.py

@@ -87,10 +87,10 @@ Can optionally accept a `RunContext` as a parameter.
 class GraphAgentState:
     """State kept across the execution of the agent graph."""
 
-    message_history: list[_messages.ModelMessage]
-    usage: _usage.RunUsage
-    retries: int
-    run_step: int
+    message_history: list[_messages.ModelMessage] = dataclasses.field(default_factory=list)
+    usage: _usage.RunUsage = dataclasses.field(default_factory=_usage.RunUsage)
+    retries: int = 0
+    run_step: int = 0
 
     def increment_retries(self, max_result_retries: int, error: BaseException | None = None) -> None:
         self.retries += 1
@@ -222,7 +222,7 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
                 if self.user_prompt is None:
                     # Skip ModelRequestNode and go directly to CallToolsNode
                     return CallToolsNode[DepsT, NodeRunEndT](last_message)
-                elif any(isinstance(part, _messages.ToolCallPart) for part in last_message.parts):
+                elif last_message.tool_calls:
                     raise exceptions.UserError(
                         'Cannot provide a new user prompt when the message history contains unprocessed tool calls.'
                     )
@@ -230,7 +230,6 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
         # Build the run context after `ctx.deps.prompt` has been updated
         run_context = build_run_context(ctx)
 
-        parts: list[_messages.ModelRequestPart] = []
         if messages:
             await self._reevaluate_dynamic_prompts(messages, run_context)
 
@@ -272,7 +271,7 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
             raise exceptions.UserError(
                 'Tool call results were provided, but the message history does not contain a `ModelResponse`.'
             )
-        if not any(isinstance(part, _messages.ToolCallPart) for part in last_model_response.parts):
+        if not last_model_response.tool_calls:
             raise exceptions.UserError(
                 'Tool call results were provided, but the message history does not contain any unprocessed tool calls.'
             )
@@ -356,9 +355,6 @@ async def _prepare_request_parameters(
     if isinstance(output_schema, _output.NativeOutputSchema):
         output_object = output_schema.object_def
 
-    # ToolOrTextOutputSchema, NativeOutputSchema, and PromptedOutputSchema all inherit from TextOutputSchema
-    allow_text_output = isinstance(output_schema, _output.TextOutputSchema)
-
     function_tools: list[ToolDefinition] = []
     output_tools: list[ToolDefinition] = []
     for tool_def in ctx.deps.tool_manager.tool_defs:
@@ -373,7 +369,8 @@ async def _prepare_request_parameters(
         output_mode=output_schema.mode,
         output_tools=output_tools,
         output_object=output_object,
-        allow_text_output=allow_text_output,
+        allow_text_output=output_schema.allows_text,
+        allow_image_output=output_schema.allows_image,
     )
 
 
@@ -543,27 +540,58 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
         if self._events_iterator is None:
             # Ensure that the stream is only run once
 
+            output_schema = ctx.deps.output_schema
+
             async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:  # noqa: C901
+                if not self.model_response.parts:
+                    # we got an empty response.
+                    # this sometimes happens with anthropic (and perhaps other models)
+                    # when the model has already returned text along side tool calls
+                    if text_processor := output_schema.text_processor:
+                        # in this scenario, if text responses are allowed, we return text from the most recent model
+                        # response, if any
+                        for message in reversed(ctx.state.message_history):
+                            if isinstance(message, _messages.ModelResponse):
+                                text = ''
+                                for part in message.parts:
+                                    if isinstance(part, _messages.TextPart):
+                                        text += part.content
+                                    elif isinstance(part, _messages.BuiltinToolCallPart):
+                                        # Text parts before a built-in tool call are essentially thoughts,
+                                        # not part of the final result output, so we reset the accumulated text
+                                        text = ''  # pragma: no cover
+                                if text:
+                                    self._next_node = await self._handle_text_response(ctx, text, text_processor)
+                                    return
+
+                    # Go back to the model request node with an empty request, which means we'll essentially
+                    # resubmit the most recent request that resulted in an empty response,
+                    # as the empty response and request will not create any items in the API payload,
+                    # in the hope the model will return a non-empty response this time.
+                    ctx.state.increment_retries(ctx.deps.max_result_retries)
+                    self._next_node = ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[]))
+                    return
+
                 text = ''
                 tool_calls: list[_messages.ToolCallPart] = []
-                invisible_parts: bool = False
+                files: list[_messages.BinaryContent] = []
 
                 for part in self.model_response.parts:
                     if isinstance(part, _messages.TextPart):
                         text += part.content
                     elif isinstance(part, _messages.ToolCallPart):
                         tool_calls.append(part)
+                    elif isinstance(part, _messages.FilePart):
+                        files.append(part.content)
                     elif isinstance(part, _messages.BuiltinToolCallPart):
                         # Text parts before a built-in tool call are essentially thoughts,
                         # not part of the final result output, so we reset the accumulated text
                         text = ''
-                        invisible_parts = True
                         yield _messages.BuiltinToolCallEvent(part)  # pyright: ignore[reportDeprecated]
                     elif isinstance(part, _messages.BuiltinToolReturnPart):
-                        invisible_parts = True
                         yield _messages.BuiltinToolResultEvent(part)  # pyright: ignore[reportDeprecated]
                     elif isinstance(part, _messages.ThinkingPart):
-                        invisible_parts = True
+                        pass
                     else:
                         assert_never(part)
 
@@ -572,47 +600,35 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                 # This accounts for cases like anthropic returns that might contain a text response
                 # and a tool call response, where the text response just indicates the tool call will happen.
                 try:
+                    alternatives: list[str] = []
                     if tool_calls:
                         async for event in self._handle_tool_calls(ctx, tool_calls):
                             yield event
-                    elif text:
-                        # No events are emitted during the handling of text responses, so we don't need to yield anything
-                        self._next_node = await self._handle_text_response(ctx, text)
-                    elif invisible_parts:
-                        # handle responses with only thinking or built-in tool parts.
-                        # this can happen with models that support thinking mode when they don't provide
-                        # actionable output alongside their thinking content. so we tell the model to try again.
-                        m = _messages.RetryPromptPart(
-                            content='Responses without text or tool calls are not permitted.',
-                        )
-                        raise ToolRetryError(m)
+                        return
+                    elif output_schema.toolset:
+                        alternatives.append('include your response in a tool call')
                     else:
-                        # we got an empty response with no tool calls, text, thinking, or built-in tool calls.
-                        # this sometimes happens with anthropic (and perhaps other models)
-                        # when the model has already returned text along side tool calls
-                        # in this scenario, if text responses are allowed, we return text from the most recent model
-                        # response, if any
-                        if isinstance(ctx.deps.output_schema, _output.TextOutputSchema):
-                            for message in reversed(ctx.state.message_history):
-                                if isinstance(message, _messages.ModelResponse):
-                                    text = ''
-                                    for part in message.parts:
-                                        if isinstance(part, _messages.TextPart):
-                                            text += part.content
-                                        elif isinstance(part, _messages.BuiltinToolCallPart):
-                                            # Text parts before a built-in tool call are essentially thoughts,
-                                            # not part of the final result output, so we reset the accumulated text
-                                            text = ''  # pragma: no cover
-                                    if text:
-                                        self._next_node = await self._handle_text_response(ctx, text)
-                                        return
-
-                        # Go back to the model request node with an empty request, which means we'll essentially
-                        # resubmit the most recent request that resulted in an empty response,
-                        # as the empty response and request will not create any items in the API payload,
-                        # in the hope the model will return a non-empty response this time.
-                        ctx.state.increment_retries(ctx.deps.max_result_retries)
-                        self._next_node = ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[]))
+                        alternatives.append('call a tool')
+
+                    if output_schema.allows_image:
+                        if image := next((file for file in files if isinstance(file, _messages.BinaryImage)), None):
+                            self._next_node = await self._handle_image_response(ctx, image)
+                            return
+                        alternatives.append('return an image')
+
+                    if text_processor := output_schema.text_processor:
+                        if text:
+                            self._next_node = await self._handle_text_response(ctx, text, text_processor)
+                            return
+                        alternatives.insert(0, 'return text')
+
+                    # handle responses with only parts that don't constitute output.
+                    # This can happen with models that support thinking mode when they don't provide
+                    # actionable output alongside their thinking content. so we tell the model to try again.
+                    m = _messages.RetryPromptPart(
+                        content=f'Please {" or ".join(alternatives)}.',
+                    )
+                    raise ToolRetryError(m)
                 except ToolRetryError as e:
                     ctx.state.increment_retries(ctx.deps.max_result_retries, e)
                     self._next_node = ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[e.tool_retry]))
@@ -655,6 +671,28 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
                 _messages.ModelRequest(parts=output_parts, instructions=instructions)
             )
 
+    async def _handle_text_response(
+        self,
+        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
+        text: str,
+        text_processor: _output.BaseOutputProcessor[NodeRunEndT],
+    ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
+        run_context = build_run_context(ctx)
+
+        result_data = await text_processor.process(text, run_context)
+
+        for validator in ctx.deps.output_validators:
+            result_data = await validator.validate(result_data, run_context)
+        return self._handle_final_result(ctx, result.FinalResult(result_data), [])
+
+    async def _handle_image_response(
+        self,
+        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
+        image: _messages.BinaryImage,
+    ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
+        result_data = cast(NodeRunEndT, image)
+        return self._handle_final_result(ctx, result.FinalResult(result_data), [])
+
     def _handle_final_result(
         self,
         ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
@@ -669,26 +707,6 @@ class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
 
         return End(final_result)
 
-    async def _handle_text_response(
-        self,
-        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
-        text: str,
-    ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
-        output_schema = ctx.deps.output_schema
-        run_context = build_run_context(ctx)
-
-        if isinstance(output_schema, _output.TextOutputSchema):
-            result_data = await output_schema.process(text, run_context)
-        else:
-            m = _messages.RetryPromptPart(
-                content='Plain text responses are not permitted, please include your response in a tool call',
-            )
-            raise ToolRetryError(m)
-
-        for validator in ctx.deps.output_validators:
-            result_data = await validator.validate(result_data, run_context)
-        return self._handle_final_result(ctx, result.FinalResult(result_data), [])
-
     __repr__ = dataclasses_no_defaults_repr
 
 
@@ -821,6 +839,7 @@ async def process_tool_calls(  # noqa: C901
             tool_calls=calls_to_run,
             tool_call_results=calls_to_run_results,
             tracer=ctx.deps.tracer,
+            usage=ctx.state.usage,
             usage_limits=ctx.deps.usage_limits,
             output_parts=output_parts,
             output_deferred_calls=deferred_calls,
@@ -867,7 +886,8 @@ async def _call_tools(
     tool_calls: list[_messages.ToolCallPart],
     tool_call_results: dict[str, DeferredToolResult],
     tracer: Tracer,
-    usage_limits: _usage.UsageLimits | None,
+    usage: _usage.RunUsage,
+    usage_limits: _usage.UsageLimits,
     output_parts: list[_messages.ModelRequestPart],
     output_deferred_calls: dict[Literal['external', 'unapproved'], list[_messages.ToolCallPart]],
 ) -> AsyncIterator[_messages.HandleResponseEvent]:
@@ -875,6 +895,11 @@ async def _call_tools(
     user_parts_by_index: dict[int, _messages.UserPromptPart] = {}
     deferred_calls_by_index: dict[int, Literal['external', 'unapproved']] = {}
 
+    if usage_limits.tool_calls_limit is not None:
+        projected_usage = deepcopy(usage)
+        projected_usage.tool_calls += len(tool_calls)
+        usage_limits.check_before_tool_call(projected_usage)
+
     for call in tool_calls:
         yield _messages.FunctionToolCallEvent(call)
 
@@ -888,13 +913,19 @@ async def _call_tools(
 
         async def handle_call_or_result(
             coro_or_task: Awaitable[
-                tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, _messages.UserPromptPart | None]
+                tuple[
+                    _messages.ToolReturnPart | _messages.RetryPromptPart, str | Sequence[_messages.UserContent] | None
+                ]
             ]
-            | Task[tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, _messages.UserPromptPart | None]],
+            | Task[
+                tuple[
+                    _messages.ToolReturnPart | _messages.RetryPromptPart, str | Sequence[_messages.UserContent] | None
+                ]
+            ],
             index: int,
         ) -> _messages.HandleResponseEvent | None:
             try:
-                tool_part, tool_user_part = (
+                tool_part, tool_user_content = (
                     (await coro_or_task) if inspect.isawaitable(coro_or_task) else coro_or_task.result()
                 )
             except exceptions.CallDeferred:
@@ -903,15 +934,15 @@ async def _call_tools(
                 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
+                if tool_user_content:
+                    user_parts_by_index[index] = _messages.UserPromptPart(content=tool_user_content)
 
-                return _messages.FunctionToolResultEvent(tool_part)
+                return _messages.FunctionToolResultEvent(tool_part, content=tool_user_content)
 
         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, tool_call_results.get(call.tool_call_id), usage_limits),
+                    _call_tool(tool_manager, call, tool_call_results.get(call.tool_call_id)),
                     index,
                 ):
                     yield event
@@ -919,7 +950,7 @@ async def _call_tools(
         else:
             tasks = [
                 asyncio.create_task(
-                    _call_tool(tool_manager, call, tool_call_results.get(call.tool_call_id), usage_limits),
+                    _call_tool(tool_manager, call, tool_call_results.get(call.tool_call_id)),
                     name=call.tool_name,
                 )
                 for call in tool_calls
@@ -946,15 +977,14 @@ async def _call_tool(
     tool_manager: ToolManager[DepsT],
     tool_call: _messages.ToolCallPart,
     tool_call_result: DeferredToolResult | None,
-    usage_limits: _usage.UsageLimits | None,
-) -> tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, _messages.UserPromptPart | None]:
+) -> tuple[_messages.ToolReturnPart | _messages.RetryPromptPart, str | Sequence[_messages.UserContent] | None]:
     try:
         if tool_call_result is None:
-            tool_result = await tool_manager.handle_call(tool_call, usage_limits=usage_limits)
+            tool_result = await tool_manager.handle_call(tool_call)
         elif isinstance(tool_call_result, ToolApproved):
             if tool_call_result.override_args is not None:
                 tool_call = dataclasses.replace(tool_call, args=tool_call_result.override_args)
-            tool_result = await tool_manager.handle_call(tool_call, usage_limits=usage_limits)
+            tool_result = await tool_manager.handle_call(tool_call)
         elif isinstance(tool_call_result, ToolDenied):
             return _messages.ToolReturnPart(
                 tool_name=tool_call.tool_name,
@@ -1024,14 +1054,7 @@ async def _call_tool(
         metadata=tool_return.metadata,
     )
 
-    user_part: _messages.UserPromptPart | None = None
-    if tool_return.content:
-        user_part = _messages.UserPromptPart(
-            content=tool_return.content,
-            part_kind='user-prompt',
-        )
-
-    return return_part, user_part
+    return return_part, tool_return.content or None
 
 
 @dataclasses.dataclass

pydantic_ai/_cli.py

@@ -18,7 +18,7 @@ from . import __version__
 from ._run_context import AgentDepsT
 from .agent import AbstractAgent, Agent
 from .exceptions import UserError
-from .messages import ModelMessage, TextPart
+from .messages import ModelMessage, ModelResponse
 from .models import KnownModelName, infer_model
 from .output import OutputDataT
 
@@ -351,14 +351,11 @@ def handle_slash_command(
         console.print('[dim]Exiting…[/dim]')
         return 0, multiline
     elif ident_prompt == '/cp':
-        try:
-            parts = messages[-1].parts
-        except IndexError:
+        if not messages or not isinstance(messages[-1], ModelResponse):
             console.print('[dim]No output available to copy.[/dim]')
         else:
-            text_to_copy = ''.join(part.content for part in parts if isinstance(part, TextPart))
-            text_to_copy = text_to_copy.strip()
-            if text_to_copy:
+            text_to_copy = messages[-1].text
+            if text_to_copy and (text_to_copy := text_to_copy.strip()):
                 pyperclip.copy(text_to_copy)
                 console.print('[dim]Copied last output to clipboard.[/dim]')
             else: