This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch durable-tool-approval in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 07267aaffff29b219791dd9ef96f75bd02f22403 Author: Kaxil Naik <[email protected]> AuthorDate: Thu Nov 12 10:31:15 2026 +0000 Keep generated durable toolset ids out of pydantic-ai's reserved format pydantic-ai keeps toolset ids in angle brackets for roles it fills itself, such as <agent>, so the id given to a FunctionToolset without one is now __airflow_function_toolset_<n> instead of <airflow-function_toolset-<n>>. The helpers that give the ids are only used by AgentOperator, so they are private now: _get_callable_toolset_id, _get_chosen_toolset_id (was chosen_toolset_id), _give_durable_ids and _GivenIds. --- .../providers/common/ai/durable/capability.py | 30 ++++++++--------- .../airflow/providers/common/ai/operators/agent.py | 18 +++++----- .../unit/common/ai/durable/test_capability.py | 38 +++++++++++----------- 3 files changed, 43 insertions(+), 43 deletions(-) diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py index 0433ef47f11..6dd5bc249e8 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py +++ b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py @@ -105,7 +105,7 @@ else: _DAG_MODULE_PREFIX = re.compile(r"^unusual_prefix_[0-9a-f]{40}_") -def get_callable_toolset_id(function: object) -> str | None: +def _get_callable_toolset_id(function: object) -> str | None: """ Return the id of a toolset that ``function`` builds per run: its module and qualified name. @@ -120,7 +120,7 @@ def get_callable_toolset_id(function: object) -> str | None: return f"{_DAG_MODULE_PREFIX.sub('', module)}.{qualname}" -def chosen_toolset_id(leaf: AbstractToolset[Any]) -> str | None: +def _get_chosen_toolset_id(leaf: AbstractToolset[Any]) -> str | None: """ Return the id pydantic-ai names a leaf toolset's durable steps after, or ``None`` when it has none. @@ -131,24 +131,24 @@ def chosen_toolset_id(leaf: AbstractToolset[Any]) -> str | None: return getattr(leaf, "_explicit_id", leaf.id) -# The id ``give_durable_ids`` gives the n-th ``FunctionToolset`` without one. The same agent gets the -# same ids on every attempt, so its requests to the model match the recorded ones. Angle brackets -# mark an id the framework chose, as in pydantic-ai's own ``<agent>``; there is no colon, which -# pydantic-ai rejects in a toolset id because it builds instruction ids from it. -_ASSIGNED_ID_PREFIX = "<airflow-function_toolset-" +# The id ``_give_durable_ids`` gives the n-th ``FunctionToolset`` without one. The same agent gets the +# same ids on every attempt, so its requests to the model match the recorded ones. The dunder prefix +# keeps it apart from ids people choose. pydantic-ai keeps ids in angle brackets for roles it fills +# itself (``<agent>``), and rejects a colon in a toolset id because it builds instruction ids from it. +_ASSIGNED_ID_PREFIX = "__airflow_function_toolset_" def _assigned_toolset_id(number: int) -> str: - return f"{_ASSIGNED_ID_PREFIX}{number}>" + return f"{_ASSIGNED_ID_PREFIX}{number}" -# Each leaf ``give_durable_ids`` copied, keyed by ``id()``, with its copy. The leaf is kept so +# Each leaf ``_give_durable_ids`` copied, keyed by ``id()``, with its copy. The leaf is kept so # that its ``id()`` cannot be reused by another object while the map is in use. -GivenIds: TypeAlias = "dict[int, tuple[AbstractToolset[Any], AbstractToolset[Any]]]" +_GivenIds: TypeAlias = "dict[int, tuple[AbstractToolset[Any], AbstractToolset[Any]]]" -def give_durable_ids( - toolset: AbstractToolset[Any] | ToolsetFunc[Any], given: GivenIds +def _give_durable_ids( + toolset: AbstractToolset[Any] | ToolsetFunc[Any], given: _GivenIds ) -> AbstractToolset[Any]: """ Return ``toolset`` with an id on each leaf pydantic-ai runs as durable steps, where one can be given. @@ -157,12 +157,12 @@ def give_durable_ids( it has an id. An id-less ``FunctionToolset`` gets one numbered in the order the toolsets are given, which its step names leave out (a function tool's step is named after the tool); a function that builds a toolset per run, or an id-less ``DynamicToolset``, gets - :func:`get_callable_toolset_id`. A leaf that already has an id, and an ``MCPToolset``, are + ``_get_callable_toolset_id``. A leaf that already has an id, and an ``MCPToolset``, are left as they are. Leaves are copied, not changed, and ``given`` maps each leaf to its copy so a toolset listed twice stays one toolset to pydantic-ai. """ if not isinstance(toolset, AbstractToolset): - return DynamicToolset(toolset, id=get_callable_toolset_id(toolset)) + return DynamicToolset(toolset, id=_get_callable_toolset_id(toolset)) if not any(_lacks_id(leaf) for leaf in iter_toolsets(toolset)): # Rebuilding wrappers via visit_and_replace breaks wrapper subclasses with their own __init__. return toolset @@ -176,7 +176,7 @@ def give_durable_ids( copied = DynamicToolset( leaf.toolset_func, per_run_step=leaf.per_run_step, - id=get_callable_toolset_id(leaf.toolset_func), + id=_get_callable_toolset_id(leaf.toolset_func), ) else: numbered = sum(isinstance(other, FunctionToolset) for _, other in given.values()) diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index 583c179e283..61e44e4fc1d 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -39,8 +39,8 @@ from pydantic_ai.usage import RunUsage from airflow.providers.common.ai.durable.capability import ( DURABLE_UNIT_TOOLSETS, AirflowDurability, - chosen_toolset_id, - give_durable_ids, + _get_chosen_toolset_id, + _give_durable_ids, ) from airflow.providers.common.ai.durable.journal import DurableJournal, build_task_storage, journal_scope from airflow.providers.common.ai.exceptions import ( @@ -116,7 +116,7 @@ if TYPE_CHECKING: from pydantic_ai.usage import UsageLimits from airflow.providers.common.ai.durable.base import DurableStorageProtocol - from airflow.providers.common.ai.durable.capability import GivenIds + from airflow.providers.common.ai.durable.capability import _GivenIds from airflow.providers.common.compat.sdk import TaskInstanceKey from airflow.sdk import Context from airflow.sdk.execution_time.context import TaskStateStoreAccessor @@ -648,15 +648,15 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa pydantic-ai names the durable steps of a ``FunctionToolset``, ``MCPToolset`` or ``DynamicToolset`` after the toolset's ``id``, and refuses to run one without it. The toolsets ``_build_agent`` hands the agent are given ids here as it gives them - (``give_durable_ids``), and a leaf still without one is refused: an ``MCPToolset``, whose + (``_give_durable_ids``), and a leaf still without one is refused: an ``MCPToolset``, whose server URL or command can carry credentials, or a function without a qualified name. Checked here so the Dag fails to load rather than at the first run. """ - given: GivenIds = {} + given: _GivenIds = {} leaves = [ leaf for toolset in self._running_toolsets() - for leaf in iter_toolsets(give_durable_ids(toolset, given)) + for leaf in iter_toolsets(_give_durable_ids(toolset, given)) if isinstance(leaf, DURABLE_UNIT_TOOLSETS) ] # One instance listed twice is one toolset to pydantic-ai; two instances sharing an id @@ -665,7 +665,7 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa first_by_id: dict[str, AbstractToolset[Any]] = {} duplicates: list[str] = [] for leaf in leaves: - if (toolset_id := chosen_toolset_id(leaf)) is None and isinstance(leaf, DynamicToolset): + if (toolset_id := _get_chosen_toolset_id(leaf)) is None and isinstance(leaf, DynamicToolset): raise ValueError( "durable=True names the steps of a toolset built per run after the function that " f"builds it, but {leaf.toolset_func!r} has no __qualname__. Pass it as " @@ -836,10 +836,10 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa # masked here as well because toolset wrappers a capability adds (get_wrapper_toolset, # as ToolLoggingCapability does) run below the capability hooks and would otherwise see # the unmasked result. The durable journal masks what it records on its own. - given_ids: GivenIds = {} + given_ids: _GivenIds = {} def prepare(toolset: AbstractToolset[Any] | ToolsetFunc[Any]) -> AbstractToolset[Any]: - return ensure_masked(give_durable_ids(toolset, given_ids) if self.durable else toolset) + return ensure_masked(_give_durable_ids(toolset, given_ids) if self.durable else toolset) if self.toolsets: toolsets: list[AbstractToolset] = [prepare(ts) for ts in self.toolsets] diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py index 482ca00604b..7d643479abf 100644 --- a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py +++ b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py @@ -53,10 +53,10 @@ from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits from airflow.providers.common.ai.durable import AirflowDurability from airflow.providers.common.ai.durable.base import RUNS_KEY from airflow.providers.common.ai.durable.capability import ( - GivenIds, + _get_callable_toolset_id, + _give_durable_ids, + _GivenIds, _JournalCodec, - get_callable_toolset_id, - give_durable_ids, ) from airflow.providers.common.ai.durable.journal import DurableJournal, Unrecordable, journal_scope from airflow.providers.common.ai.exceptions import DurableJournalError @@ -1115,7 +1115,7 @@ class _ToolsBuilder: class TestToolsetIds: - """Toolsets without an ``id`` replay under the ids ``give_durable_ids`` gives them, as in AgentOperator.""" + """Toolsets without an ``id`` replay under the ids ``_give_durable_ids`` gives them, as in AgentOperator.""" @staticmethod def tool_step_names(storage) -> list[str]: @@ -1146,7 +1146,7 @@ class TestToolsetIds: return Agent( FunctionModel(tool_then_answer(calls, "query", fail_final=fail_final)), name="analyst", - toolsets=[give_durable_ids(toolset, {})], + toolsets=[_give_durable_ids(toolset, {})], capabilities=[AirflowDurability()], ) @@ -1177,7 +1177,7 @@ class TestToolsetIds: agent = Agent( FunctionModel(tool_then_answer(calls, "query", fail_final=[True])), name="analyst", - toolsets=[give_durable_ids(_query_tools(calls, *tools), {})], + toolsets=[_give_durable_ids(_query_tools(calls, *tools), {})], capabilities=[AirflowDurability()], ) with pytest.raises(RuntimeError, match="worker died"): @@ -1195,7 +1195,7 @@ class TestToolsetIds: agent = Agent( FunctionModel(tool_then_answer(calls, "query", fail_final=[True])), name="analyst", - toolsets=[give_durable_ids(toolset, {})], + toolsets=[_give_durable_ids(toolset, {})], capabilities=[AirflowDurability()], ) @@ -1206,27 +1206,27 @@ class TestToolsetIds: def test_gives_ids_to_copies_and_one_copy_to_a_toolset_listed_twice(self): """ - ``give_durable_ids`` writes ``FunctionToolset``'s private ``_id``, which only its constructor sets. + ``_give_durable_ids`` writes ``FunctionToolset``'s private ``_id``, which only its constructor sets. This test and the replay tests above fail if pydantic-ai stops reading the id from there. """ shared = FunctionToolset() - given: GivenIds = {} + given: _GivenIds = {} - prefixed = give_durable_ids(shared.prefixed("a"), given) - plain = give_durable_ids(shared, given) + prefixed = _give_durable_ids(shared.prefixed("a"), given) + plain = _give_durable_ids(shared, given) assert isinstance(prefixed, WrapperToolset) assert prefixed.wrapped is plain - assert plain.id == "<airflow-function_toolset-0>" + assert plain.id == "__airflow_function_toolset_0" assert shared.id is None def test_numbers_the_function_toolsets_it_gives_ids_in_order(self): - given: GivenIds = {} + given: _GivenIds = {} - ids = [give_durable_ids(FunctionToolset(), given).id for _ in range(2)] + ids = [_give_durable_ids(FunctionToolset(), given).id for _ in range(2)] - assert ids == ["<airflow-function_toolset-0>", "<airflow-function_toolset-1>"] + assert ids == ["__airflow_function_toolset_0", "__airflow_function_toolset_1"] @pytest.mark.parametrize( "toolset", @@ -1237,10 +1237,10 @@ class TestToolsetIds: ], ) def test_leaves_a_toolset_it_cannot_or_need_not_give_an_id_as_it_is(self, toolset): - assert give_durable_ids(toolset, {}) is toolset + assert _give_durable_ids(toolset, {}) is toolset def test_a_dynamic_toolset_without_an_id_gets_its_functions_name(self): - given = give_durable_ids(DynamicToolset(_build_tools, per_run_step=False), {}) + given = _give_durable_ids(DynamicToolset(_build_tools, per_run_step=False), {}) assert isinstance(given, DynamicToolset) assert (given.id, given.per_run_step) == (f"{__name__}._build_tools", False) @@ -1253,7 +1253,7 @@ class TestGetCallableToolsetId: build_tools.__module__ = f"unusual_prefix_{'0123456789abcdef' * 2}01234567_orders_dag" - assert get_callable_toolset_id(build_tools) == f"orders_dag.{build_tools.__qualname__}" + assert _get_callable_toolset_id(build_tools) == f"orders_dag.{build_tools.__qualname__}" @pytest.mark.parametrize( "function", @@ -1263,7 +1263,7 @@ class TestGetCallableToolsetId: ], ) def test_none_for_a_callable_without_a_qualified_name(self, function): - assert get_callable_toolset_id(function) is None + assert _get_callable_toolset_id(function) is None class Answer(BaseModel):
