Resolve comments from jed
diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index 83ab8f0..f7d1b9e 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst
@@ -24,7 +24,7 @@ and managed credentials. Toolsets expose them as pydantic-ai tools so that LLM agents can call them during multi-turn reasoning. -Four toolsets are exported directly from the ``airflow.providers.common.ai.toolsets`` +Six toolsets are exported directly from the ``airflow.providers.common.ai.toolsets`` package root: - :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` — generic @@ -37,17 +37,13 @@ - :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset` — give the agent a shell and a filesystem inside an isolated sandbox, off the Airflow worker. See :ref:`which boundary that is <sandbox-boundaries>`. +- :class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset` + — base class that provider packages subclass to expose a **vendor-managed + agent**, one whose reasoning loop runs on a cloud provider's infrastructure. +- :class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset` + — composes several interchangeable managed agents behind a single tool. + See :ref:`managed-agent-toolsets` below. -A fourth pair, -:class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset` -and -:class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset`, -covers **vendor-managed agents** -- agents whose reasoning loop runs on a cloud -provider's infrastructure. The first is a base class that provider packages -subclass; the second composes several interchangeable ones behind a single tool. -See :ref:`managed-agent-toolsets` below. - -All of them implement pydantic-ai's Three more toolsets are documented later on this page. They are not re-exported from the package root, so import each of them from its own submodule:: @@ -1358,7 +1354,8 @@ ``failover_on`` defaults to ``Exception`` because ``common.ai`` cannot enumerate the cloud SDKs' exception trees — ``requests``, ``botocore`` and the Azure SDK -share no common base. Narrow it when the members' exception types are known. +share no common base. It can be narrowed when the members' exception types are +known. ``replayable`` on a group is ``True`` only when every member is, because the durable cache cannot know which member produced the answer it holds.
diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py index d445a6b..5e965d0 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py
@@ -24,7 +24,6 @@ from pydantic_ai.tools import ToolDefinition from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool -from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError from airflow.providers.common.ai.utils.tool_definition import ( build_args_validator, return_schema_kwargs, @@ -33,6 +32,8 @@ from airflow.providers.common.compat.sdk import Stats if TYPE_CHECKING: + from collections.abc import Sequence + from pydantic_ai._run_context import RunContext log = logging.getLogger(__name__) @@ -79,6 +80,9 @@ the platform default, which subclasses supply -- a number chosen here would silently disagree with the vendor operator's documented timeout for the same service. + :param max_retries: How many times the calling model may rephrase after the + remote agent raises ``ModelRetry``. ``0`` turns the first ``ModelRetry`` + into a hard error, which disables that recovery path entirely. """ #: Whether a completed invocation may be replayed from the durable cache @@ -93,13 +97,17 @@ tool_name: str, description: str | None = None, timeout: float | None = None, + max_retries: int = 1, ) -> None: if not tool_name: raise ValueError("tool_name must be a non-empty string.") + if max_retries < 0: + raise ValueError(f"max_retries must not be negative, got {max_retries}.") self._tool_name = tool_name # Same fallback as HookToolset uses for a method with no docstring. self._description = (description or "").strip() or tool_name.replace("_", " ").capitalize() self._timeout = timeout + self._max_retries = max_retries @property @abstractmethod @@ -109,12 +117,9 @@ Must contain ``platform`` and ``name``, e.g. ``{"platform": "snowflake.cortex", "name": "ANALYTICS.REVENUE.BOOKINGS_ANALYST"}``. - Logged on every invocation, so the resolved remote identity behind a task - appears in that task's log even though the Dag only names a connection. - It is not pushed to XCom: ``FailoverManagedAgentToolset`` reports the - group rather than the responder, and the operational question -- how often - a standby is answering -- is carried by the ``managed_agent.served`` - counter instead. + Logged whenever this toolset's tool is called, so the resolved remote + identity behind a task appears in that task's log even though the Dag + only names a connection. """ @abstractmethod @@ -172,11 +177,13 @@ self._tool_name: ToolsetTool( toolset=self, tool_def=tool_def, - # One rephrase attempt, matching HookToolset. Deliberately not - # more: a managed agent invocation is expensive, and a prompt the - # remote agent could not parse rarely becomes parseable on a - # second try. - max_retries=1, + # How many times the calling model may rephrase after ``invoke`` + # raises ``ModelRetry``. One by default, matching HookToolset: a + # managed agent invocation is expensive, so the budget is small. + # Zero disables the ``ModelRetry`` path entirely -- the first one + # becomes a hard error -- so raise it only when the remote agent's + # rejections are genuinely worth re-prompting. + max_retries=self._max_retries, args_validator=build_args_validator(_PROMPT_SCHEMA), ) } @@ -224,26 +231,19 @@ it elsewhere. Since most platforms fall on the stateful side, treat one-shot as something a group is deliberately restricted to, not a safe default. - Prefer plain Airflow task-level failover for a standalone call: two tasks, - the second with ``trigger_rule=TriggerRule.ALL_FAILED``, keeps which - provider served the request visible in the grid at no code cost. This class - is for the case a task boundary cannot express -- a managed agent consulted - as a tool *inside* a longer agent run, where failing the task would discard - the calling agent's accumulated context and re-run every earlier tool call. - :param members: Interchangeable toolsets, tried in order. At least two. :param failover_on: Exception types that move to the next member. Defaults to ``Exception`` because ``common.ai`` cannot enumerate the cloud SDKs' exception trees (``requests``, ``botocore`` and the Azure SDK share no - common base), so the safe default is broad. Narrow it when the members' - exception types are known. ``ModelRetry`` is always re-raised and never - triggers failover, whatever this is set to. + common base), so the safe default is broad. It can be narrowed when the + members' exception types are known. ``ModelRetry`` is always re-raised + and never triggers failover, whatever this is set to. """ def __init__( self, *, - members: list[BaseManagedAgentToolset], + members: Sequence[BaseManagedAgentToolset], failover_on: tuple[type[BaseException], ...] = (Exception,), **kwargs, ) -> None: @@ -253,7 +253,10 @@ "A failover group needs at least two members; " f"got {len(members)}. Use the member toolset directly instead." ) - self._members = members + # Copied, not aliased: the loop in invoke() relies on the group being + # non-empty, and a caller holding the original list could otherwise empty + # it after construction. + self._members = tuple(members) self._failover_on = failover_on # Replay is only safe if every member is safe to replay: the cache cannot # know which member produced the answer it holds. @@ -295,6 +298,7 @@ Stats.incr( "managed_agent.failover", tags={ + "tool": self._tool_name, "from_platform": ref.get("platform", "unknown"), "to_platform": standby.get("platform", "unknown"), }, @@ -308,10 +312,10 @@ Stats.incr( "managed_agent.served", tags={ + "tool": self._tool_name, "platform": ref.get("platform", "unknown"), "role": "standby" if served_by_standby else "primary", + "position": str(position), }, ) return result - # Unreachable: the last member either returns or raises above. - raise ManagedAgentInvocationError("Failover group exhausted with no result.")
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py index eb5957a..7751a54 100644 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py
@@ -94,6 +94,20 @@ def test_not_replayable_by_default(self): assert FakeManagedAgentToolset().replayable is False + @pytest.mark.asyncio + async def test_max_retries_defaults_to_one_and_is_configurable(self): + # One rephrase by default; 0 disables the ModelRetry path entirely, so it + # is a deliberate choice rather than the default. + default = await FakeManagedAgentToolset().get_tools(ctx=None) + assert default["ask_specialist"].max_retries == 1 + + tuned = await FakeManagedAgentToolset(max_retries=3).get_tools(ctx=None) + assert tuned["ask_specialist"].max_retries == 3 + + def test_negative_max_retries_rejected(self): + with pytest.raises(ValueError, match="max_retries must not be negative"): + FakeManagedAgentToolset(max_retries=-1) + class TestGetTools: @pytest.mark.asyncio @@ -193,6 +207,21 @@ self._group(*members) @pytest.mark.asyncio + async def test_members_are_copied_so_the_caller_cannot_empty_the_group(self): + # invoke() relies on the group being non-empty; aliasing the caller's list + # would let it be emptied after construction and make invoke return None, + # which reaches the model as the string "null". + # Constructed directly, not via _group(), which copies the list itself. + members = [FakeManagedAgentToolset(result="from primary"), FakeManagedAgentToolset()] + group = FailoverManagedAgentToolset( + members=members, + tool_name="ask_resilient", + description="Answers questions, on whichever cloud is up.", + ) + members.clear() + assert await self._call(group) == "from primary" + + @pytest.mark.asyncio async def test_primary_answer_wins_and_standby_is_untouched(self): primary = FakeManagedAgentToolset(result="from primary") standby = FakeManagedAgentToolset(result="from standby") @@ -293,7 +322,13 @@ async def test_primary_success_counts_as_primary(self, mock_stats): await self._group(FakeManagedAgentToolset(), FakeManagedAgentToolset()).invoke("q") mock_stats.incr.assert_called_once_with( - "managed_agent.served", tags={"platform": "fake.cloud", "role": "primary"} + "managed_agent.served", + tags={ + "tool": "ask_resilient", + "platform": "fake.cloud", + "role": "primary", + "position": "0", + }, ) @pytest.mark.asyncio @@ -305,9 +340,22 @@ assert mock_stats.incr.call_args_list == [ mock.call( "managed_agent.failover", - tags={"from_platform": "fake.cloud", "to_platform": "fake.cloud"}, + tags={ + "tool": "ask_resilient", + "from_platform": "fake.cloud", + "to_platform": "fake.cloud", + }, ), - mock.call("managed_agent.served", tags={"platform": "fake.cloud", "role": "standby"}), + mock.call( + "managed_agent.served", + tags={ + "tool": "ask_resilient", + "platform": "fake.cloud", + "role": "standby", + # The member that answered, not just that a standby did. + "position": "1", + }, + ), ] @pytest.mark.asyncio
diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index edac39b..b4a39ea 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
@@ -398,6 +398,23 @@ legacy_name: "-" name_variables: [] + - name: "managed_agent.failover" + description: "Number of failover transitions inside a ``FailoverManagedAgentToolset`` group, + emitted once each time a member fails and the next one is tried. + Metric with tool, from_platform and to_platform tagging." + type: "counter" + legacy_name: "-" + name_variables: ["tool", "from_platform", "to_platform"] + + - name: "managed_agent.served" + description: "Number of managed agent invocations answered inside a + ``FailoverManagedAgentToolset`` group, emitted once per answer. Metric with tool, + platform, role (``primary`` or ``standby``) and position tagging, where position is + the zero-based index of the member that answered." + type: "counter" + legacy_name: "-" + name_variables: ["tool", "platform", "role", "position"] + # ========== # Gauges # ==========