This is an automated email from the ASF dual-hosted git repository. vikramkoka pushed a commit to branch common_ai_managed_toolset in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 8c1db549538e0a4c491808714ae56b1e7068ce87 Merge: 5f8bc3854a4 7bbbd2fdacc Author: Vikram Koka <[email protected]> AuthorDate: Fri Aug 21 06:46:54 2026 -0700 Merge branch 'common_ai_managed_toolset' of https://github.com/apache/airflow into common_ai_managed_toolset providers/common/ai/docs/toolsets.rst | 5 +++-- .../ai/src/airflow/providers/common/ai/toolsets/managed_agent.py | 2 -- 2 files changed, 3 insertions(+), 4 deletions(-) diff --cc providers/common/ai/docs/toolsets.rst index 4e613d302f4,eb393713bd8..cd6d9b4ac83 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@@ -879,50 -876,12 +879,51 @@@ Tool naming, argument validation, resul by the base class, so every provider's implementation presents the same surface to the calling model. +``tool_name`` is the required identifier — it is what the model emits when it +calls the tool, and the Dag author chooses it. ``description`` is optional and +falls back to the tool name rendered as prose, the same way ``HookToolset`` +derives one from a method name when there is no docstring. + .. note:: - ``description`` is a required constructor argument. A remote agent's - competence cannot be introspected the way ``HookToolset`` reads a hook's - docstrings, and the description is the only basis the calling model has for - choosing between specialists. + Writing a description is still worth the line. It is what tells the model to + consult the agent rather than answer from its own knowledge, and it is the + only place to record a scope limit the name cannot carry — "cannot see + revenue figures". Because the argument schema is always a bare prompt, the + name and the description are the whole of what the model knows about the + agent. + +Toolset or operator? - """""""""""""""""""" ++"""""""""""""""""""""" + +Most managed-agent platforms do not offer a plain one-request-one-answer API. Some +require polling a job; others require creating a session and tearing it down around +each exchange. A toolset can do either, but only by blocking inside +``invoke()`` — it cannot defer to the Triggerer, and it has no post-task hook to +clean up with if the worker dies mid-call. + +That draws a boundary worth respecting: + +.. list-table:: + :header-rows: 1 + :widths: 45 55 + + * - Shape + - Surface to use + * - A short consultation *inside* an agent's reasoning, where failing the task + would discard the calling agent's accumulated context + - A managed agent toolset + * - Long-running submitted work as a pipeline step in its own right + - That provider's own operator, with deferral or + :class:`~airflow.sdk.bases.resumablejobmixin.ResumableJobMixin` + ++ +``ResumableJobMixin`` exists for exactly the second case: it persists the external +job ID to the task state store before polling, so a worker crash reconnects to the +running job instead of submitting a duplicate. A toolset cannot offer that, because +the retry boundary is the task, not the tool call — on retry the agent loop restarts +and re-issues the call. Durable execution covers the *completed* call (see +``replayable`` below); it does not cover a call that was still in flight. Error handling """""""""""""" @@@ -958,94 -917,9 +959,94 @@@ avoid paying for the same invocation tw Deferral """""""" -A toolset call runs in the worker and cannot defer to the triggerer. Managed -agents with no request/response mode — an async job that writes output to -object storage, or a stateful session — are reachable through a toolset only by -blocking for the duration. That is acceptable for a tool call inside an agent's -reasoning, but it is not a substitute for that provider's own deferrable -operator when the Dag simply needs to submit work and wait. +A toolset call runs in the worker and cannot defer to the Triggerer — it blocks +for the duration of the call. See `Toolset or operator?`_ above for when that is +acceptable and when the provider's own deferrable operator is the right surface +instead. + +Failover between interchangeable agents - """"""""""""""""""""""""""""""""""""""" ++""""""""""""""""""""""""""""""""""" + +:class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset` +composes several managed agents into one tool, trying them in order until one +answers. It is itself a ``BaseManagedAgentToolset``, so the calling model sees a +single tool and has no say in which provider serves the request — the policy +stays deterministic Python rather than a prompt instruction a model may ignore. +Groups nest. + +.. code-block:: python + + from airflow.providers.common.ai.toolsets import FailoverManagedAgentToolset + + resilient = FailoverManagedAgentToolset( + tool_name="ask_claims_agent", + description="Reviews an insurance claim and returns a coverage determination.", + members=[bedrock_claims_agent, foundry_claims_agent], # same image, two clouds + ) + +Members must satisfy two preconditions the class cannot check. + +**Substitutability.** The same agent deployed twice, not two specialists with +different data. Two containerised agents built from one image qualify; agents +bound to one platform's own objects — a Cortex Agent over Snowflake semantic +models — do not, because there is nothing equivalent to fail over *to*. + +**Statelessness per invocation.** Server-side conversation state is the norm +across managed-agent platforms, not the exception — optional on some (Cortex +``thread_id``), mandatory on others where a session is created and torn down +around each exchange. Each member is invoked with a bare prompt and no thread +reference, so a failover silently starts a fresh conversation on the standby: +correct for a one-shot consultation, wrong for a multi-turn one. Treat one-shot +as a restriction a group is deliberately held to, not a safe default. + +The three error buckets do real work here: + +- ``ManagedAgentInvocationError`` and transient failures move to the next member. +- ``ModelRetry`` is re-raised immediately and never triggers failover. A prompt + the primary could not parse will not parse on the standby either, so failing + over would spend the standby's budget reproducing the same error. +- The last member's exception propagates unchanged, so a total outage still fails + the task rather than returning something misleading. + +``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. + +``replayable`` on a group is ``True`` only when every member is, because the +durable cache cannot know which member produced the answer it holds. + +.. note:: + + For a **standalone** agent call, prefer plain Airflow task-level failover: + two tasks, the second with ``trigger_rule=TriggerRule.ALL_FAILED``. That keeps + which provider served the request visible in the grid at no code cost, and + makes failover rate a task metric. 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. + +Two counters make failover visible, because a failover is a *success-shaped* +event — without them a primary that has been down for a week looks identical to a +healthy one: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Metric + - Tags + * - ``managed_agent.failover`` + - ``from_platform``, ``to_platform`` — one per failover transition + * - ``managed_agent.served`` + - ``platform``, ``role`` (``primary`` / ``standby``) — one per answer + +The standby-served fraction is a ratio over ``managed_agent.served`` alone, so +"are we quietly running on the standby?" is a dashboard question rather than a log +grep. Both are tagged by platform rather than agent name to keep cardinality +bounded. + +One limitation remains: which member served a *particular* answer is in the task +log but not in XCom. ``agent_ref`` on a group describes the group, not the +responder, because the responder is not known until after the call. The counters +cover the operational question; per-answer provenance for an audit trail would +need ``AgentOperator`` to collect per-toolset metadata. diff --cc providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py index d445a6b0e92,03ff1c54f3c..ce8ef6431f1 --- 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 @@@ -138,17 -129,6 +138,16 @@@ class BaseManagedAgentToolset(AbstractT propagate unchanged. Airflow's task-level retry is the right layer; a rephrase does nothing for a 503. + **Release anything you allocate, on every path.** Platforms that require a + session bill for its lifetime, so an implementation that opens one here + must close it in a ``finally`` -- including when ``ModelRetry`` propagates, + which is a return path the calling model treats as recoverable and will + therefore hit repeatedly. A tool call has no post-task cleanup hook to + fall back on: if the worker dies mid-call the handle is lost, and nothing + will reap the remote session. Implementations whose sessions are long + enough for that to matter belong in that provider's own operator, where + deferral and :class:`~airflow.sdk.bases.resumablejobmixin.ResumableJobMixin` + can reconnect to the existing job instead of leaking it. - :param prompt: The question or instruction to send to the remote agent. """ @@@ -192,126 -166,3 +191,125 @@@ log.info("Consulting managed agent %s on %s", ref.get("name"), ref.get("platform")) result = await self.invoke(tool_args["prompt"]) return serialize_for_llm(result) + - +class FailoverManagedAgentToolset(BaseManagedAgentToolset): + """ + Present several interchangeable managed agents to the model as one tool. + + Active/passive failover for a managed agent: members are tried in order and + the first answer wins. Because this is itself a + :class:`BaseManagedAgentToolset`, the calling model sees a single tool and + has no say in which provider serves the request -- the policy stays + deterministic Python rather than a prompt instruction a model may ignore. + Groups nest, so a group can itself be a member of another group. + + Members must satisfy two preconditions that this class cannot check: + + *Substitutability.* The same agent deployed twice, not two specialists with + different data. Two containerised agents built from one image (Bedrock + AgentCore and Azure AI Foundry hosted agents, say) qualify; agents backed by + different corpora or bound to one platform's own objects -- a Cortex Agent + over Snowflake semantic models -- do not, because there is no equivalent to + fail over *to*. + + *Statelessness per invocation.* Server-side conversation state is the norm + rather than the exception across managed-agent platforms -- optional on some + (Cortex ``thread_id``), mandatory on others, where a session must be created + and torn down around every exchange. Each member here is invoked with a bare + prompt and no thread reference, so a failover silently starts a fresh + conversation on the standby. That is correct for a one-shot consultation and + wrong for a multi-turn one: failover discards the thread rather than resuming + 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. + """ + + def __init__( + self, + *, + members: list[BaseManagedAgentToolset], + failover_on: tuple[type[BaseException], ...] = (Exception,), + **kwargs, + ) -> None: + super().__init__(**kwargs) + if len(members) < 2: + raise ValueError( + "A failover group needs at least two members; " + f"got {len(members)}. Use the member toolset directly instead." + ) + self._members = 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. + self.replayable = all(m.replayable for m in members) + + @property + def agent_ref(self) -> dict[str, str]: + return { + "platform": "failover", + "name": " -> ".join(m.agent_ref.get("name", "?") for m in self._members), + } + + async def invoke(self, prompt: str) -> Any: + last = len(self._members) - 1 + for position, member in enumerate(self._members): + ref = member.agent_ref + try: + result = await member.invoke(prompt) + except ModelRetry: + # The model can fix this by rephrasing, and the standby would + # reject the same prompt identically. Failing over would spend + # the standby's budget to reproduce the same error. + raise + except self._failover_on: + if position == last: + raise + standby = self._members[position + 1].agent_ref + log.warning( + "Managed agent %s on %s failed; failing over to %s", + ref.get("name"), + ref.get("platform"), + standby.get("name"), + exc_info=True, + ) + # Metrics, not just logs: a failover is a success-shaped event, so + # without a counter a primary that has been down for a week looks + # identical to a healthy one. Tagged by platform rather than agent + # name to keep cardinality bounded. + Stats.incr( + "managed_agent.failover", + tags={ + "from_platform": ref.get("platform", "unknown"), + "to_platform": standby.get("platform", "unknown"), + }, + ) + continue + served_by_standby = position > 0 + if served_by_standby: + log.info("Managed agent request served by standby %s", ref.get("name")) + # Emitted on every answer so the standby-served fraction is a ratio of + # this counter, not something that has to be scanned out of XCom. + Stats.incr( + "managed_agent.served", + tags={ + "platform": ref.get("platform", "unknown"), + "role": "standby" if served_by_standby else "primary", + }, + ) + return result + # Unreachable: the last member either returns or raises above. + raise ManagedAgentInvocationError("Failover group exhausted with no result.")
