Add ability to end agent runs as a result of a tool call

pydantic_ai_slim/pydantic_ai/_exceptions.py

@@ -0,0 +1,16 @@
+from __future__ import annotations as _annotations
+
+from typing import Any
+
+
+class StopAgentRun(Exception):
+    """Exception raised to stop the agent run and use the provided result as the output."""
+
+    result: Any
+    """The value to use as the result of the agent run."""
+    tool_name: str | None
+    """The name of the tool call, if available."""
+
+    def __init__(self, result: Any, tool_name: str | None) -> None:
+        self.result = result
+        self.tool_name = tool_name

pydantic_ai_slim/pydantic_ai/agent.py

@@ -13,6 +13,7 @@
 from typing_extensions import assert_never
 
 from . import (
+    _exceptions,
     _result,
     _system_prompt,
     _utils,
@@ -787,9 +788,13 @@ async def _handle_model_response(
                 else:
                     messages.append(self._unknown_tool(call.tool_name))
 
-            with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]):
-                task_results: Sequence[_messages.Message] = await asyncio.gather(*tasks)
-                messages.extend(task_results)
+            with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]) as span:
+                try:
+                    task_results: Sequence[_messages.Message] = await asyncio.gather(*tasks)
+                    messages.extend(task_results)
+                except _exceptions.StopAgentRun as e:
+                    span.set_attribute('stop_agent_run', e.tool_name)
+                    return _MarkFinalResult(data=e.result), []
             return None, messages
         else:
             assert_never(model_response)
@@ -854,9 +859,13 @@ async def _handle_streamed_model_response(
                 else:
                     messages.append(self._unknown_tool(call.tool_name))
 
-            with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]):
-                task_results: Sequence[_messages.Message] = await asyncio.gather(*tasks)
-                messages.extend(task_results)
+            with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]) as span:
+                try:
+                    task_results: Sequence[_messages.Message] = await asyncio.gather(*tasks)
+                    messages.extend(task_results)
+                except _exceptions.StopAgentRun as e:
+                    span.set_attribute('stop_agent_run', e.tool_name)
+                    return _MarkFinalResult(data=e.result), []
             return None, messages
 
     async def _validate_result(

pydantic_ai_slim/pydantic_ai/tools.py

@@ -3,19 +3,17 @@
 import inspect
 from collections.abc import Awaitable
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar, Union, cast
+from typing import Any, Callable, Generic, NoReturn, Union, cast
 
 from pydantic import ValidationError
 from pydantic_core import SchemaValidator
-from typing_extensions import Concatenate, ParamSpec, TypeAlias
+from typing_extensions import Concatenate, ParamSpec, TypeAlias, TypeVar
 
-from . import _pydantic, _utils, messages
+from . import _exceptions, _pydantic, _utils, messages
 from .exceptions import ModelRetry, UnexpectedModelBehavior
 
-if TYPE_CHECKING:
-    from .result import ResultData
-else:
-    ResultData = Any
+ResultData = TypeVar('ResultData', default=Any)
+"""Type variable for the result data of a run."""
 
 
 __all__ = (
@@ -38,7 +36,7 @@
 
 
 @dataclass
-class RunContext(Generic[AgentDeps]):
+class RunContext(Generic[AgentDeps, ResultData]):
     """Information about the current call."""
 
     deps: AgentDeps
@@ -48,6 +46,11 @@ class RunContext(Generic[AgentDeps]):
     tool_name: str | None = None
     """Name of the tool being called."""
 
+    def stop_run(self, result: ResultData) -> NoReturn:
+        """Stop the call to `agent.run` as soon as possible, using the provided value as the result."""
+        # NOTE: this means we ignore any other tools called concurrently
+        raise _exceptions.StopAgentRun(result, tool_name=self.tool_name)
+
 
 ToolParams = ParamSpec('ToolParams')
 """Retrieval function param spec."""

tests/test_tools.py

@@ -466,3 +466,12 @@ def ctx_tool(ctx: RunContext[int], x: int) -> int:
     """)
     result = mod.agent.run_sync('foobar', deps=5)
     assert result.data == snapshot('{"ctx_tool":5}')
+
+def test_ctx_stop_run():
+    """Ensure the ctx_stop_run_tool is used to complete the agent run."""
+    def ctx_stop_run_tool(ctx: RunContext[int]):
+        ctx.stop_run(ctx.deps * 'abc')
+
+    agent = Agent('test', result_type=str, tools=[Tool(ctx_stop_run_tool, takes_ctx=True)], deps_type=int)
+    result = agent.run_sync('foobar', deps=2)
+    assert result.data == snapshot('abcabc')