diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 72c256e9c4..4309f896b8 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -344,6 +344,7 @@ def __init__( self._event_stream_handler = event_stream_handler + self._override_name: ContextVar[_utils.Option[str]] = ContextVar('_override_name', default=None) self._override_deps: ContextVar[_utils.Option[AgentDepsT]] = ContextVar('_override_deps', default=None) self._override_model: ContextVar[_utils.Option[models.Model]] = ContextVar('_override_model', default=None) self._override_toolsets: ContextVar[_utils.Option[Sequence[AbstractToolset[AgentDepsT]]]] = ContextVar( @@ -384,7 +385,8 @@ def name(self) -> str | None: If `None`, we try to infer the agent name from the call frame when the agent is first run. """ - return self._name + name_ = self._override_name.get() + return name_.value if name_ else self._name @name.setter def name(self, value: str | None) -> None: @@ -712,6 +714,7 @@ def _run_span_end_attributes( def override( self, *, + name: str | _utils.Unset = _utils.UNSET, deps: AgentDepsT | _utils.Unset = _utils.UNSET, model: models.Model | models.KnownModelName | str | _utils.Unset = _utils.UNSET, toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET, @@ -724,12 +727,18 @@ def override( You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures). Args: + name: The name to use instead of the name passed to the agent constructor and 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. toolsets: The toolsets to use instead of the toolsets passed to the agent constructor and agent run. tools: The tools to use instead of the tools registered with the agent. instructions: The instructions to use instead of the instructions registered with the agent. """ + if _utils.is_set(name): + name_token = self._override_name.set(_utils.Some(name)) + else: + name_token = None + if _utils.is_set(deps): deps_token = self._override_deps.set(_utils.Some(deps)) else: @@ -759,6 +768,8 @@ def override( try: yield finally: + if name_token is not None: + self._override_name.reset(name_token) if deps_token is not None: self._override_deps.reset(deps_token) if model_token is not None: diff --git a/pydantic_ai_slim/pydantic_ai/agent/abstract.py b/pydantic_ai_slim/pydantic_ai/agent/abstract.py index c5d59f7561..e5ebfe0d49 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/abstract.py +++ b/pydantic_ai_slim/pydantic_ai/agent/abstract.py @@ -854,6 +854,7 @@ async def main(): def override( self, *, + name: str | _utils.Unset = _utils.UNSET, deps: AgentDepsT | _utils.Unset = _utils.UNSET, model: models.Model | models.KnownModelName | str | _utils.Unset = _utils.UNSET, toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET, @@ -866,6 +867,7 @@ def override( You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures). Args: + name: The name to use instead of the name passed to the agent constructor and 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. toolsets: The toolsets to use instead of the toolsets passed to the agent constructor and agent run. diff --git a/pydantic_ai_slim/pydantic_ai/agent/wrapper.py b/pydantic_ai_slim/pydantic_ai/agent/wrapper.py index ba735f0907..339ed5ebdd 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/wrapper.py +++ b/pydantic_ai_slim/pydantic_ai/agent/wrapper.py @@ -210,6 +210,7 @@ async def main(): def override( self, *, + name: str | _utils.Unset = _utils.UNSET, deps: AgentDepsT | _utils.Unset = _utils.UNSET, model: models.Model | models.KnownModelName | str | _utils.Unset = _utils.UNSET, toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET, @@ -222,6 +223,7 @@ def override( You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures). Args: + name: The name to use instead of the name passed to the agent constructor and 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. toolsets: The toolsets to use instead of the toolsets passed to the agent constructor and agent run. @@ -229,6 +231,7 @@ def override( instructions: The instructions to use instead of the instructions registered with the agent. """ with self.wrapped.override( + name=name, deps=deps, model=model, toolsets=toolsets, diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/dbos/_agent.py b/pydantic_ai_slim/pydantic_ai/durable_exec/dbos/_agent.py index eb609cfc4a..e6212a2ed6 100644 --- a/pydantic_ai_slim/pydantic_ai/durable_exec/dbos/_agent.py +++ b/pydantic_ai_slim/pydantic_ai/durable_exec/dbos/_agent.py @@ -804,6 +804,7 @@ async def main(): def override( self, *, + name: str | _utils.Unset = _utils.UNSET, deps: AgentDepsT | _utils.Unset = _utils.UNSET, model: models.Model | models.KnownModelName | str | _utils.Unset = _utils.UNSET, toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET, @@ -816,6 +817,7 @@ def override( You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures). Args: + name: The name to use instead of the name passed to the agent constructor and 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. toolsets: The toolsets to use instead of the toolsets passed to the agent constructor and agent run. @@ -828,6 +830,7 @@ def override( ) with super().override( + name=name, deps=deps, model=model, toolsets=toolsets, diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py index 5338f228a9..1da6e4d902 100644 --- a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py +++ b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_agent.py @@ -862,6 +862,7 @@ async def main(): def override( self, *, + name: str | _utils.Unset = _utils.UNSET, deps: AgentDepsT | _utils.Unset = _utils.UNSET, model: models.Model | models.KnownModelName | str | _utils.Unset = _utils.UNSET, toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET, @@ -874,6 +875,7 @@ def override( You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures). Args: + name: The name to use instead of the name passed to the agent constructor and 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. toolsets: The toolsets to use instead of the toolsets passed to the agent constructor and agent run. @@ -895,6 +897,7 @@ def override( ) with super().override( + name=name, deps=deps, model=model, toolsets=toolsets, diff --git a/tests/test_agent.py b/tests/test_agent.py index 00113bd4da..68ee1f3a61 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -2537,6 +2537,14 @@ async def test_agent_name_changes(): assert new_agent.name == 'my_agent' +def test_agent_name_override(): + agent = Agent('test', name='custom_name') + + with agent.override(name='overridden_name'): + agent.run_sync('Hello') + assert agent.name == 'overridden_name' + + def test_name_from_global(create_module: Callable[[str], Any]): module_code = """ from pydantic_ai import Agent