pydantic-ai-slim 0.0.6a4__py3-none-any.whl → 0.0.7__py3-none-any.whl

pydantic_ai/agent.py

@@ -11,15 +11,15 @@ from typing_extensions import assert_never
 
 from . import (
     _result,
-    _retriever as _r,
     _system_prompt,
+    _tool as _r,
     _utils,
     exceptions,
     messages as _messages,
     models,
     result,
 )
-from .dependencies import AgentDeps, CallContext, RetrieverContextFunc, RetrieverParams, RetrieverPlainFunc
+from .dependencies import AgentDeps, RunContext, ToolContextFunc, ToolParams, ToolPlainFunc
 from .result import ResultData
 
 __all__ = ('Agent',)
@@ -58,7 +58,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     _result_validators: list[_result.ResultValidator[AgentDeps, ResultData]] = field(repr=False)
     _allow_text_result: bool = field(repr=False)
     _system_prompts: tuple[str, ...] = field(repr=False)
-    _retrievers: dict[str, _r.Retriever[AgentDeps, Any]] = field(repr=False)
+    _function_tools: dict[str, _r.Tool[AgentDeps, Any]] = field(repr=False)
     _default_retries: int = field(repr=False)
     _system_prompt_functions: list[_system_prompt.SystemPromptRunner[AgentDeps]] = field(repr=False)
     _deps_type: type[AgentDeps] = field(repr=False)
@@ -75,8 +75,8 @@ class Agent(Generic[AgentDeps, ResultData]):
     def __init__(
         self,
         model: models.Model | models.KnownModelName | None = None,
-        result_type: type[ResultData] = str,
         *,
+        result_type: type[ResultData] = str,
         system_prompt: str | Sequence[str] = (),
         deps_type: type[AgentDeps] = NoneType,
         retries: int = 1,
@@ -105,7 +105,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                 it's evaluated to create a [`Model`][pydantic_ai.models.Model] instance immediately,
                 which checks for the necessary environment variables. Set this to `false`
                 to defer the evaluation until the first run. Useful if you want to
-                [override the model][pydantic_ai.Agent.override_model] for testing.
+                [override the model][pydantic_ai.Agent.override] for testing.
         """
         if model is None or defer_model_check:
             self.model = model
@@ -119,7 +119,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         self._allow_text_result = self._result_schema is None or self._result_schema.allow_text_result
 
         self._system_prompts = (system_prompt,) if isinstance(system_prompt, str) else tuple(system_prompt)
-        self._retrievers: dict[str, _r.Retriever[AgentDeps, Any]] = {}
+        self._function_tools: dict[str, _r.Tool[AgentDeps, Any]] = {}
         self._deps_type = deps_type
         self._default_retries = retries
         self._system_prompt_functions = []
@@ -150,14 +150,6 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         deps = self._get_deps(deps)
 
-        new_message_index, messages = await self._prepare_messages(deps, user_prompt, message_history)
-        self.last_run_messages = messages
-
-        for retriever in self._retrievers.values():
-            retriever.reset()
-
-        cost = result.Cost()
-
         with _logfire.span(
             'agent run {prompt=}',
             prompt=user_prompt,
@@ -165,6 +157,14 @@ class Agent(Generic[AgentDeps, ResultData]):
             custom_model=custom_model,
             model_name=model_used.name(),
         ) as run_span:
+            new_message_index, messages = await self._prepare_messages(deps, user_prompt, message_history)
+            self.last_run_messages = messages
+
+            for tool in self._function_tools.values():
+                tool.reset()
+
+            cost = result.Cost()
+
             run_step = 0
             while True:
                 run_step += 1
@@ -243,14 +243,6 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         deps = self._get_deps(deps)
 
-        new_message_index, messages = await self._prepare_messages(deps, user_prompt, message_history)
-        self.last_run_messages = messages
-
-        for retriever in self._retrievers.values():
-            retriever.reset()
-
-        cost = result.Cost()
-
         with _logfire.span(
             'agent run stream {prompt=}',
             prompt=user_prompt,
@@ -258,6 +250,14 @@ class Agent(Generic[AgentDeps, ResultData]):
             custom_model=custom_model,
             model_name=model_used.name(),
         ) as run_span:
+            new_message_index, messages = await self._prepare_messages(deps, user_prompt, message_history)
+            self.last_run_messages = messages
+
+            for tool in self._function_tools.values():
+                tool.reset()
+
+            cost = result.Cost()
+
             run_step = 0
             while True:
                 run_step += 1
@@ -296,42 +296,51 @@ class Agent(Generic[AgentDeps, ResultData]):
                                 cost += model_response.cost()
 
     @contextmanager
-    def override_deps(self, overriding_deps: AgentDeps) -> Iterator[None]:
-        """Context manager to temporarily override agent dependencies, this is particularly useful when testing.
+    def override(
+        self,
+        *,
+        deps: AgentDeps | _utils.Unset = _utils.UNSET,
+        model: models.Model | models.KnownModelName | _utils.Unset = _utils.UNSET,
+    ) -> Iterator[None]:
+        """Context manager to temporarily override agent dependencies and model.
+
+        This is particularly useful when testing.
 
         Args:
-            overriding_deps: The dependencies to use instead of the dependencies passed to the agent run.
+            deps: The dependencies to use instead of the dependencies passed to the agent run.
+            model: The model to use instead of the model passed to the agent run.
         """
-        override_deps_before = self._override_deps
-        self._override_deps = _utils.Some(overriding_deps)
-        try:
-            yield
-        finally:
-            self._override_deps = override_deps_before
+        if _utils.is_set(deps):
+            override_deps_before = self._override_deps
+            self._override_deps = _utils.Some(deps)
+        else:
+            override_deps_before = _utils.UNSET
 
-    @contextmanager
-    def override_model(self, overriding_model: models.Model | models.KnownModelName) -> Iterator[None]:
-        """Context manager to temporarily override the model used by the agent.
+        # noinspection PyTypeChecker
+        if _utils.is_set(model):
+            override_model_before = self._override_model
+            # noinspection PyTypeChecker
+            self._override_model = _utils.Some(models.infer_model(model))  # pyright: ignore[reportArgumentType]
+        else:
+            override_model_before = _utils.UNSET
 
-        Args:
-            overriding_model: The model to use instead of the model passed to the agent run.
-        """
-        override_model_before = self._override_model
-        self._override_model = _utils.Some(models.infer_model(overriding_model))
         try:
             yield
         finally:
-            self._override_model = override_model_before
+            if _utils.is_set(override_deps_before):
+                self._override_deps = override_deps_before
+            if _utils.is_set(override_model_before):
+                self._override_model = override_model_before
 
     @overload
     def system_prompt(
-        self, func: Callable[[CallContext[AgentDeps]], str], /
-    ) -> Callable[[CallContext[AgentDeps]], str]: ...
+        self, func: Callable[[RunContext[AgentDeps]], str], /
+    ) -> Callable[[RunContext[AgentDeps]], str]: ...
 
     @overload
     def system_prompt(
-        self, func: Callable[[CallContext[AgentDeps]], Awaitable[str]], /
-    ) -> Callable[[CallContext[AgentDeps]], Awaitable[str]]: ...
+        self, func: Callable[[RunContext[AgentDeps]], Awaitable[str]], /
+    ) -> Callable[[RunContext[AgentDeps]], Awaitable[str]]: ...
 
     @overload
     def system_prompt(self, func: Callable[[], str], /) -> Callable[[], str]: ...
@@ -344,7 +353,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     ) -> _system_prompt.SystemPromptFunc[AgentDeps]:
         """Decorator to register a system prompt function.
 
-        Optionally takes [`CallContext`][pydantic_ai.dependencies.CallContext] as it's only argument.
+        Optionally takes [`RunContext`][pydantic_ai.dependencies.RunContext] as it's only argument.
         Can decorate a sync or async functions.
 
         Overloads for every possible signature of `system_prompt` are included so the decorator doesn't obscure
@@ -352,7 +361,7 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         Example:
         ```py
-        from pydantic_ai import Agent, CallContext
+        from pydantic_ai import Agent, RunContext
 
         agent = Agent('test', deps_type=str)
 
@@ -361,12 +370,12 @@ class Agent(Generic[AgentDeps, ResultData]):
             return 'foobar'
 
         @agent.system_prompt
-        async def async_system_prompt(ctx: CallContext[str]) -> str:
+        async def async_system_prompt(ctx: RunContext[str]) -> str:
             return f'{ctx.deps} is the best'
 
         result = agent.run_sync('foobar', deps='spam')
         print(result.data)
-        #> success (no retriever calls)
+        #> success (no tool calls)
         ```
         """
         self._system_prompt_functions.append(_system_prompt.SystemPromptRunner(func))
@@ -374,13 +383,13 @@ class Agent(Generic[AgentDeps, ResultData]):
 
     @overload
     def result_validator(
-        self, func: Callable[[CallContext[AgentDeps], ResultData], ResultData], /
-    ) -> Callable[[CallContext[AgentDeps], ResultData], ResultData]: ...
+        self, func: Callable[[RunContext[AgentDeps], ResultData], ResultData], /
+    ) -> Callable[[RunContext[AgentDeps], ResultData], ResultData]: ...
 
     @overload
     def result_validator(
-        self, func: Callable[[CallContext[AgentDeps], ResultData], Awaitable[ResultData]], /
-    ) -> Callable[[CallContext[AgentDeps], ResultData], Awaitable[ResultData]]: ...
+        self, func: Callable[[RunContext[AgentDeps], ResultData], Awaitable[ResultData]], /
+    ) -> Callable[[RunContext[AgentDeps], ResultData], Awaitable[ResultData]]: ...
 
     @overload
     def result_validator(self, func: Callable[[ResultData], ResultData], /) -> Callable[[ResultData], ResultData]: ...
@@ -395,7 +404,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     ) -> _result.ResultValidatorFunc[AgentDeps, ResultData]:
         """Decorator to register a result validator function.
 
-        Optionally takes [`CallContext`][pydantic_ai.dependencies.CallContext] as it's first argument.
+        Optionally takes [`RunContext`][pydantic_ai.dependencies.RunContext] as it's first argument.
         Can decorate a sync or async functions.
 
         Overloads for every possible signature of `result_validator` are included so the decorator doesn't obscure
@@ -403,7 +412,7 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         Example:
         ```py
-        from pydantic_ai import Agent, CallContext, ModelRetry
+        from pydantic_ai import Agent, ModelRetry, RunContext
 
         agent = Agent('test', deps_type=str)
 
@@ -414,61 +423,57 @@ class Agent(Generic[AgentDeps, ResultData]):
             return data
 
         @agent.result_validator
-        async def result_validator_deps(ctx: CallContext[str], data: str) -> str:
+        async def result_validator_deps(ctx: RunContext[str], data: str) -> str:
             if ctx.deps in data:
                 raise ModelRetry('wrong response')
             return data
 
         result = agent.run_sync('foobar', deps='spam')
         print(result.data)
-        #> success (no retriever calls)
+        #> success (no tool calls)
         ```
         """
         self._result_validators.append(_result.ResultValidator(func))
         return func
 
     @overload
-    def retriever(
-        self, func: RetrieverContextFunc[AgentDeps, RetrieverParams], /
-    ) -> RetrieverContextFunc[AgentDeps, RetrieverParams]: ...
+    def tool(self, func: ToolContextFunc[AgentDeps, ToolParams], /) -> ToolContextFunc[AgentDeps, ToolParams]: ...
 
     @overload
-    def retriever(
+    def tool(
         self, /, *, retries: int | None = None
-    ) -> Callable[
-        [RetrieverContextFunc[AgentDeps, RetrieverParams]], RetrieverContextFunc[AgentDeps, RetrieverParams]
-    ]: ...
+    ) -> Callable[[ToolContextFunc[AgentDeps, ToolParams]], ToolContextFunc[AgentDeps, ToolParams]]: ...
 
-    def retriever(
+    def tool(
         self,
-        func: RetrieverContextFunc[AgentDeps, RetrieverParams] | None = None,
+        func: ToolContextFunc[AgentDeps, ToolParams] | None = None,
         /,
         *,
         retries: int | None = None,
     ) -> Any:
-        """Decorator to register a retriever function which takes
-        [`CallContext`][pydantic_ai.dependencies.CallContext] as its first argument.
+        """Decorator to register a tool function which takes
+        [`RunContext`][pydantic_ai.dependencies.RunContext] as its first argument.
 
         Can decorate a sync or async functions.
 
         The docstring is inspected to extract both the tool description and description of each parameter,
-        [learn more](../agents.md#retrievers-tools-and-schema).
+        [learn more](../agents.md#function-tools-and-schema).
 
-        We can't add overloads for every possible signature of retriever, since the return type is a recursive union
-        so the signature of functions decorated with `@agent.retriever` is obscured.
+        We can't add overloads for every possible signature of tool, since the return type is a recursive union
+        so the signature of functions decorated with `@agent.tool` is obscured.
 
         Example:
         ```py
-        from pydantic_ai import Agent, CallContext
+        from pydantic_ai import Agent, RunContext
 
         agent = Agent('test', deps_type=int)
 
-        @agent.retriever
-        def foobar(ctx: CallContext[int], x: int) -> int:
+        @agent.tool
+        def foobar(ctx: RunContext[int], x: int) -> int:
             return ctx.deps + x
 
-        @agent.retriever(retries=2)
-        async def spam(ctx: CallContext[str], y: float) -> float:
+        @agent.tool(retries=2)
+        async def spam(ctx: RunContext[str], y: float) -> float:
             return ctx.deps + y
 
         result = agent.run_sync('foobar', deps=1)
@@ -477,58 +482,56 @@ class Agent(Generic[AgentDeps, ResultData]):
         ```
 
         Args:
-            func: The retriever function to register.
-            retries: The number of retries to allow for this retriever, defaults to the agent's default retries,
+            func: The tool function to register.
+            retries: The number of retries to allow for this tool, defaults to the agent's default retries,
                 which defaults to 1.
         """  # noqa: D205
         if func is None:
 
-            def retriever_decorator(
-                func_: RetrieverContextFunc[AgentDeps, RetrieverParams],
-            ) -> RetrieverContextFunc[AgentDeps, RetrieverParams]:
+            def tool_decorator(
+                func_: ToolContextFunc[AgentDeps, ToolParams],
+            ) -> ToolContextFunc[AgentDeps, ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_retriever(_utils.Either(left=func_), retries)
+                self._register_tool(_utils.Either(left=func_), retries)
                 return func_
 
-            return retriever_decorator
+            return tool_decorator
         else:
             # noinspection PyTypeChecker
-            self._register_retriever(_utils.Either(left=func), retries)
+            self._register_tool(_utils.Either(left=func), retries)
             return func
 
     @overload
-    def retriever_plain(self, func: RetrieverPlainFunc[RetrieverParams], /) -> RetrieverPlainFunc[RetrieverParams]: ...
+    def tool_plain(self, func: ToolPlainFunc[ToolParams], /) -> ToolPlainFunc[ToolParams]: ...
 
     @overload
-    def retriever_plain(
+    def tool_plain(
         self, /, *, retries: int | None = None
-    ) -> Callable[[RetrieverPlainFunc[RetrieverParams]], RetrieverPlainFunc[RetrieverParams]]: ...
+    ) -> Callable[[ToolPlainFunc[ToolParams]], ToolPlainFunc[ToolParams]]: ...
 
-    def retriever_plain(
-        self, func: RetrieverPlainFunc[RetrieverParams] | None = None, /, *, retries: int | None = None
-    ) -> Any:
-        """Decorator to register a retriever function which DOES NOT take `CallContext` as an argument.
+    def tool_plain(self, func: ToolPlainFunc[ToolParams] | None = None, /, *, retries: int | None = None) -> Any:
+        """Decorator to register a tool function which DOES NOT take `RunContext` as an argument.
 
         Can decorate a sync or async functions.
 
         The docstring is inspected to extract both the tool description and description of each parameter,
-        [learn more](../agents.md#retrievers-tools-and-schema).
+        [learn more](../agents.md#function-tools-and-schema).
 
-        We can't add overloads for every possible signature of retriever, since the return type is a recursive union
-        so the signature of functions decorated with `@agent.retriever` is obscured.
+        We can't add overloads for every possible signature of tool, since the return type is a recursive union
+        so the signature of functions decorated with `@agent.tool` is obscured.
 
         Example:
         ```py
-        from pydantic_ai import Agent, CallContext
+        from pydantic_ai import Agent, RunContext
 
         agent = Agent('test')
 
-        @agent.retriever
-        def foobar(ctx: CallContext[int]) -> int:
+        @agent.tool
+        def foobar(ctx: RunContext[int]) -> int:
             return 123
 
-        @agent.retriever(retries=2)
-        async def spam(ctx: CallContext[str]) -> float:
+        @agent.tool(retries=2)
+        async def spam(ctx: RunContext[str]) -> float:
             return 3.14
 
         result = agent.run_sync('foobar', deps=1)
@@ -537,38 +540,36 @@ class Agent(Generic[AgentDeps, ResultData]):
         ```
 
         Args:
-            func: The retriever function to register.
-            retries: The number of retries to allow for this retriever, defaults to the agent's default retries,
+            func: The tool function to register.
+            retries: The number of retries to allow for this tool, defaults to the agent's default retries,
                 which defaults to 1.
         """
         if func is None:
 
-            def retriever_decorator(
-                func_: RetrieverPlainFunc[RetrieverParams],
-            ) -> RetrieverPlainFunc[RetrieverParams]:
+            def tool_decorator(
+                func_: ToolPlainFunc[ToolParams],
+            ) -> ToolPlainFunc[ToolParams]:
                 # noinspection PyTypeChecker
-                self._register_retriever(_utils.Either(right=func_), retries)
+                self._register_tool(_utils.Either(right=func_), retries)
                 return func_
 
-            return retriever_decorator
+            return tool_decorator
         else:
-            self._register_retriever(_utils.Either(right=func), retries)
+            self._register_tool(_utils.Either(right=func), retries)
             return func
 
-    def _register_retriever(
-        self, func: _r.RetrieverEitherFunc[AgentDeps, RetrieverParams], retries: int | None
-    ) -> None:
-        """Private utility to register a retriever function."""
+    def _register_tool(self, func: _r.ToolEitherFunc[AgentDeps, ToolParams], retries: int | None) -> None:
+        """Private utility to register a tool function."""
         retries_ = retries if retries is not None else self._default_retries
-        retriever = _r.Retriever[AgentDeps, RetrieverParams](func, retries_)
+        tool = _r.Tool[AgentDeps, ToolParams](func, retries_)
 
-        if self._result_schema and retriever.name in self._result_schema.tools:
-            raise ValueError(f'Retriever name conflicts with result schema name: {retriever.name!r}')
+        if self._result_schema and tool.name in self._result_schema.tools:
+            raise ValueError(f'Tool name conflicts with result schema name: {tool.name!r}')
 
-        if retriever.name in self._retrievers:
-            raise ValueError(f'Retriever name conflicts with existing retriever: {retriever.name!r}')
+        if tool.name in self._function_tools:
+            raise ValueError(f'Tool name conflicts with existing tool: {tool.name!r}')
 
-        self._retrievers[retriever.name] = retriever
+        self._function_tools[tool.name] = tool
 
     async def _get_agent_model(
         self, model: models.Model | models.KnownModelName | None
@@ -583,11 +584,11 @@ class Agent(Generic[AgentDeps, ResultData]):
         """
         model_: models.Model
         if some_model := self._override_model:
-            # we don't want `override_model()` to cover up errors from the model not being defined, hence this check
+            # we don't want `override()` to cover up errors from the model not being defined, hence this check
             if model is None and self.model is None:
                 raise exceptions.UserError(
                     '`model` must be set either when creating the agent or when calling it. '
-                    '(Even when `override_model()` is customizing the model that will actually be called)'
+                    '(Even when `override(model=...)` is customizing the model that will actually be called)'
                 )
             model_ = some_model.value
             custom_model = None
@@ -601,7 +602,7 @@ class Agent(Generic[AgentDeps, ResultData]):
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
         result_tools = list(self._result_schema.tools.values()) if self._result_schema else None
-        agent_model = await model_.agent_model(self._retrievers, self._allow_text_result, result_tools)
+        agent_model = await model_.agent_model(self._function_tools, self._allow_text_result, result_tools)
         return model_, custom_model, agent_model
 
     async def _prepare_messages(
@@ -663,12 +664,12 @@ class Agent(Generic[AgentDeps, ResultData]):
             if not model_response.calls:
                 raise exceptions.UnexpectedModelBehavior('Received empty tool call message')
 
-            # otherwise we run all retriever functions in parallel
+            # otherwise we run all tool functions in parallel
             messages: list[_messages.Message] = []
             tasks: list[asyncio.Task[_messages.Message]] = []
             for call in model_response.calls:
-                if retriever := self._retrievers.get(call.tool_name):
-                    tasks.append(asyncio.create_task(retriever.run(deps, call), name=call.tool_name))
+                if tool := self._function_tools.get(call.tool_name):
+                    tasks.append(asyncio.create_task(tool.run(deps, call), name=call.tool_name))
                 else:
                     messages.append(self._unknown_tool(call.tool_name))
 
@@ -719,7 +720,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                 if self._result_schema.find_tool(structured_msg):
                     return _MarkFinalResult(model_response)
 
-            # the model is calling a retriever function, consume the response to get the next message
+            # the model is calling a tool function, consume the response to get the next message
             async for _ in model_response:
                 pass
             structured_msg = model_response.get()
@@ -727,11 +728,11 @@ class Agent(Generic[AgentDeps, ResultData]):
                 raise exceptions.UnexpectedModelBehavior('Received empty tool call message')
             messages: list[_messages.Message] = [structured_msg]
 
-            # we now run all retriever functions in parallel
+            # we now run all tool functions in parallel
             tasks: list[asyncio.Task[_messages.Message]] = []
             for call in structured_msg.calls:
-                if retriever := self._retrievers.get(call.tool_name):
-                    tasks.append(asyncio.create_task(retriever.run(deps, call), name=call.tool_name))
+                if tool := self._function_tools.get(call.tool_name):
+                    tasks.append(asyncio.create_task(tool.run(deps, call), name=call.tool_name))
                 else:
                     messages.append(self._unknown_tool(call.tool_name))
 
@@ -763,7 +764,7 @@ class Agent(Generic[AgentDeps, ResultData]):
 
     def _unknown_tool(self, tool_name: str) -> _messages.RetryPrompt:
         self._incr_result_retry()
-        names = list(self._retrievers.keys())
+        names = list(self._function_tools.keys())
         if self._result_schema:
             names.extend(self._result_schema.tool_names())
         if names:
@@ -775,7 +776,7 @@ class Agent(Generic[AgentDeps, ResultData]):
     def _get_deps(self, deps: AgentDeps) -> AgentDeps:
         """Get deps for a run.
 
-        If we've overridden deps via `_override_deps_stack`, use that, otherwise use the deps passed to the call.
+        If we've overridden deps via `_override_deps`, use that, otherwise use the deps passed to the call.
 
         We could do runtime type checking of deps against `self._deps_type`, but that's a slippery slope.
         """

pydantic_ai/dependencies.py

@@ -13,13 +13,13 @@ else:
 
 __all__ = (
     'AgentDeps',
-    'CallContext',
+    'RunContext',
     'ResultValidatorFunc',
     'SystemPromptFunc',
-    'RetrieverReturnValue',
-    'RetrieverContextFunc',
-    'RetrieverPlainFunc',
-    'RetrieverParams',
+    'ToolReturnValue',
+    'ToolContextFunc',
+    'ToolPlainFunc',
+    'ToolParams',
     'JsonData',
 )
 
@@ -28,7 +28,7 @@ AgentDeps = TypeVar('AgentDeps')
 
 
 @dataclass
-class CallContext(Generic[AgentDeps]):
+class RunContext(Generic[AgentDeps]):
     """Information about the current call."""
 
     deps: AgentDeps
@@ -39,23 +39,23 @@ class CallContext(Generic[AgentDeps]):
     """Name of the tool being called."""
 
 
-RetrieverParams = ParamSpec('RetrieverParams')
+ToolParams = ParamSpec('ToolParams')
 """Retrieval function param spec."""
 
 SystemPromptFunc = Union[
-    Callable[[CallContext[AgentDeps]], str],
-    Callable[[CallContext[AgentDeps]], Awaitable[str]],
+    Callable[[RunContext[AgentDeps]], str],
+    Callable[[RunContext[AgentDeps]], Awaitable[str]],
     Callable[[], str],
     Callable[[], Awaitable[str]],
 ]
-"""A function that may or maybe not take `CallContext` as an argument, and may or may not be async.
+"""A function that may or maybe not take `RunContext` as an argument, and may or may not be async.
 
 Usage `SystemPromptFunc[AgentDeps]`.
 """
 
 ResultValidatorFunc = Union[
-    Callable[[CallContext[AgentDeps], ResultData], ResultData],
-    Callable[[CallContext[AgentDeps], ResultData], Awaitable[ResultData]],
+    Callable[[RunContext[AgentDeps], ResultData], ResultData],
+    Callable[[RunContext[AgentDeps], ResultData], Awaitable[ResultData]],
     Callable[[ResultData], ResultData],
     Callable[[ResultData], Awaitable[ResultData]],
 ]
@@ -69,15 +69,15 @@ Usage `ResultValidator[AgentDeps, ResultData]`.
 JsonData: TypeAlias = 'None | str | int | float | Sequence[JsonData] | Mapping[str, JsonData]'
 """Type representing any JSON data."""
 
-RetrieverReturnValue = Union[JsonData, Awaitable[JsonData]]
-"""Return value of a retriever function."""
-RetrieverContextFunc = Callable[Concatenate[CallContext[AgentDeps], RetrieverParams], RetrieverReturnValue]
-"""A retriever function that takes `CallContext` as the first argument.
+ToolReturnValue = Union[JsonData, Awaitable[JsonData]]
+"""Return value of a tool function."""
+ToolContextFunc = Callable[Concatenate[RunContext[AgentDeps], ToolParams], ToolReturnValue]
+"""A tool function that takes `RunContext` as the first argument.
 
-Usage `RetrieverContextFunc[AgentDeps, RetrieverParams]`.
+Usage `ToolContextFunc[AgentDeps, ToolParams]`.
 """
-RetrieverPlainFunc = Callable[RetrieverParams, RetrieverReturnValue]
-"""A retriever function that does not take `CallContext` as the first argument.
+ToolPlainFunc = Callable[ToolParams, ToolReturnValue]
+"""A tool function that does not take `RunContext` as the first argument.
 
-Usage `RetrieverPlainFunc[RetrieverParams]`.
+Usage `ToolPlainFunc[ToolParams]`.
 """

pydantic_ai/exceptions.py

@@ -6,7 +6,7 @@ __all__ = 'ModelRetry', 'UserError', 'UnexpectedModelBehavior'
 
 
 class ModelRetry(Exception):
-    """Exception raised when a retriever function should be retried.
+    """Exception raised when a tool function should be retried.
 
     The agent will return the message to the model and ask it to try calling the function/tool again.
     """