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 e49f29f0ed90b9c570607d1b9724f089eb1940b4 Author: Vikram Koka <[email protected]> AuthorDate: Wed Aug 12 21:45:31 2026 -0700 Add BaseManagedAgentToolset for vendor-managed AI agents Cloud vendors now run agents on the user's behalf -- Snowflake Cortex Agents, Amazon Bedrock AgentCore, Azure AI Foundry hosted agents, Vertex AI Agent Engine. Their reasoning loops execute on the vendor's infrastructure, so they are not something AgentOperator can run, only something an Airflow task can consult. Each provider that has added support so far arrived at a different identity scheme, request shape, and invoke verb, so a Dag author who learns one learns nothing transferable about the next. What all of them do agree on is a single consultation: a question in, an answer out, once. That is the shape of a tool call, so the contract is scoped to exactly that and leaves each provider's own lifecycle operators alone. Provider packages subclass it, keeping credentials flowing through their existing hooks and requiring no new connection types. Errors are deliberately sorted into three buckets rather than two: what the calling model can fix by rephrasing, what is terminal, and what is transient and belongs to Airflow's task retry. Collapsing the third into either of the others is how these implementations go wrong. Durable replay is opt-in per implementation, since a managed agent may act on systems Airflow cannot observe and replaying a cached answer could skip a side effect. FailoverManagedAgentToolset composes interchangeable agents -- the same agent deployed on two clouds -- into one tool for active/passive failover. It is itself a BaseManagedAgentToolset, so the model sees a single tool and the policy stays deterministic rather than becoming a prompt instruction a model may ignore, and groups nest. Task-level failover remains the better choice for a standalone call; this covers the case a task boundary cannot express, where the agent is consulted mid-run and failing the task would discard the calling agent's accumulated context. Also promotes the private tool-result serialiser out of HookToolset into the shared tool_definition utils so both toolsets use one implementation. --- providers/common/ai/docs/toolsets.rst | 205 ++++++++++++- providers/common/ai/provider.yaml | 1 + .../src/airflow/providers/common/ai/exceptions.py | 11 + .../providers/common/ai/toolsets/__init__.py | 15 +- .../airflow/providers/common/ai/toolsets/hook.py | 23 +- .../providers/common/ai/toolsets/managed_agent.py | 304 +++++++++++++++++++ .../providers/common/ai/utils/tool_definition.py | 17 ++ .../ai/tests/unit/common/ai/toolsets/test_hook.py | 16 +- .../unit/common/ai/toolsets/test_managed_agent.py | 325 +++++++++++++++++++++ 9 files changed, 891 insertions(+), 26 deletions(-) diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index 34a8b43f417..47c72cf91e6 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -34,7 +34,16 @@ Three toolsets are included: `MCP servers <https://modelcontextprotocol.io/>`__ configured via Airflow connections. -All three implement pydantic-ai's +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 `AbstractToolset <https://ai.pydantic.dev/toolsets/>`__ interface and can be passed to any pydantic-ai ``Agent``, including via :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`. @@ -839,3 +848,197 @@ Before deploying an agent task to production: 8. **Prompt injection**: Be cautious when the prompt includes untrusted data (user input, external API responses, upstream XCom). Consider sanitizing inputs before passing them to the agent. + +.. _managed-agent-toolsets: + +Managed Agent Toolsets +---------------------- + +Cloud vendors now run agents on your behalf — Snowflake Cortex Agents, Amazon +Bedrock AgentCore runtimes, Azure AI Foundry hosted agents, Vertex AI Agent +Engine. Their reasoning loops execute on the vendor's infrastructure, so they +are not something ``AgentOperator`` runs; they are something an Airflow task +*consults*. + +:class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset` +is the contract for exposing one of those as a tool. Each provider package +ships its own subclass, so credentials keep flowing through that provider's +existing hook and no new connection types are needed. + +A subclass implements two members: + +``agent_ref`` + Normalised identity of the remote agent — ``platform`` and ``name`` — logged + on every call so a run can be audited for which agents it consulted. + +``invoke(prompt)`` + Send the prompt, return the agent's answer. Return the *answer*, not the + transport envelope. + +Tool naming, argument validation, result serialisation, and logging are handled +by the base class, so every provider's implementation presents the same surface +to the calling model. + +.. 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. + +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 +"""""""""""""" + +Failures sort into three buckets, and conflating them is the most common way an +implementation goes wrong: + +.. list-table:: + :header-rows: 1 + :widths: 22 30 48 + + * - Raise + - When + - Who recovers + * - ``ModelRetry`` + - The agent rejected the request in a way rephrasing could fix. + - The calling model, bounded by its ``usage_limits``. + * - ``ManagedAgentInvocationError`` + - Terminal: bad credentials, missing agent, revoked quota. + - Nobody — the task fails fast instead of burning retries. + * - *let it propagate* + - Transient: 429, 5xx, connection reset, read timeout. + - Airflow's task-level retry. A rephrase does nothing for a 503. + +Durable execution +""""""""""""""""" + +``replayable`` is ``False`` by default. A managed agent may act on systems +Airflow cannot observe, so replaying a cached answer on retry could skip a side +effect. Implementations whose agent is read-only should set it to ``True`` to +avoid paying for the same invocation twice. + +Deferral +"""""""" + +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 --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index b6133e2e2e4..5335de6179b 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -490,6 +490,7 @@ toolsets: - airflow.providers.common.ai.toolsets.mcp - airflow.providers.common.ai.toolsets.skills - airflow.providers.common.ai.toolsets.langchain_bridge + - airflow.providers.common.ai.toolsets.managed_agent retry-policies: - integration-name: Common AI diff --git a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py index d0412fb2a48..d02748d4a3a 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py +++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py @@ -37,3 +37,14 @@ class LLMFileAnalysisLimitExceededError(LLMFileAnalysisError): class LLMFileAnalysisMultimodalRequiredError(LLMFileAnalysisUnsupportedFormatError): """Raised when image/PDF inputs are used without ``multi_modal=True``.""" + + +class ManagedAgentInvocationError(RuntimeError): + """ + Raised when a managed agent cannot be reached and retrying will not help. + + Reserved for terminal conditions -- bad credentials, a missing agent, a + revoked quota. Transient failures should propagate unchanged so Airflow's + task-level retry handles them, and requests the model could fix by + rephrasing should raise ``pydantic_ai.exceptions.ModelRetry`` instead. + """ diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py index 6c30fa4a733..655d9b0823d 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py @@ -19,8 +19,19 @@ from __future__ import annotations from airflow.providers.common.ai.toolsets.hook import HookToolset - -__all__ = ["HookToolset", "MCPToolset", "SQLToolset", "airflow_toolset_to_langchain_tools"] +from airflow.providers.common.ai.toolsets.managed_agent import ( + BaseManagedAgentToolset, + FailoverManagedAgentToolset, +) + +__all__ = [ + "BaseManagedAgentToolset", + "FailoverManagedAgentToolset", + "HookToolset", + "MCPToolset", + "SQLToolset", + "airflow_toolset_to_langchain_tools", +] def __getattr__(name: str): diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py index 63037e1a4f8..13412b82c35 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py @@ -19,7 +19,6 @@ from __future__ import annotations import inspect -import json import re import types from typing import TYPE_CHECKING, Any, Union, get_args, get_origin, get_type_hints @@ -27,7 +26,11 @@ from typing import TYPE_CHECKING, Any, Union, get_args, get_origin, get_type_hin from pydantic_ai.tools import ToolDefinition from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool -from airflow.providers.common.ai.utils.tool_definition import build_args_validator, return_schema_kwargs +from airflow.providers.common.ai.utils.tool_definition import ( + build_args_validator, + return_schema_kwargs, + serialize_for_llm, +) if TYPE_CHECKING: from collections.abc import Callable @@ -109,7 +112,7 @@ class HookToolset(AbstractToolset[Any]): # sequential=True because hook methods perform synchronous I/O # (network calls, DB queries) and should not run concurrently. # return_schema is "string": call_tool serializes every result with - # _serialize_for_llm, so the tool always returns a (JSON-encoded) + # serialize_for_llm, so the tool always returns a (JSON-encoded) # string regardless of the method's own return annotation. This lets # code mode render `-> str` instead of `-> Any`. tool_def = ToolDefinition( @@ -137,7 +140,7 @@ class HookToolset(AbstractToolset[Any]): method_name = name.removeprefix(self._tool_name_prefix) if self._tool_name_prefix else name method: Callable[..., Any] = getattr(self._hook, method_name) result = method(**tool_args) - return _serialize_for_llm(result) + return serialize_for_llm(result) # --------------------------------------------------------------------------- @@ -260,15 +263,3 @@ def _parse_param_docs(docstring: str) -> dict[str, str]: params[m.group(1)] = " ".join(m.group(2).split()) return params - - -def _serialize_for_llm(value: Any) -> str: - """Convert a Python return value to a string suitable for an LLM.""" - if value is None: - return "null" - if isinstance(value, str): - return value - try: - return json.dumps(value, default=str) - except (TypeError, ValueError): - return str(value) 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 new file mode 100644 index 00000000000..5166132dbcc --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py @@ -0,0 +1,304 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import logging +from abc import abstractmethod +from typing import TYPE_CHECKING, Any + +from pydantic_ai.exceptions import ModelRetry +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, + serialize_for_llm, +) +from airflow.providers.common.compat.sdk import Stats + +if TYPE_CHECKING: + from pydantic_ai._run_context import RunContext + +log = logging.getLogger(__name__) + +_PROMPT_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "The question or instruction to send to this agent.", + } + }, + "required": ["prompt"], +} + + +class BaseManagedAgentToolset(AbstractToolset[Any]): + """ + Base class exposing a vendor-managed agent as a single pydantic-ai tool. + + A managed agent runs its own reasoning loop on the vendor's infrastructure + (Snowflake Cortex Agents, Amazon Bedrock AgentCore, Azure AI Foundry hosted + agents, Vertex AI Agent Engine). Airflow submits one request and reads one + answer, so the Airflow-side agent features -- toolsets, human-in-the-loop + review, durable step replay -- apply to the *calling* agent and never reach + inside the managed agent. + + Subclasses implement :meth:`agent_ref` and :meth:`invoke`. Tool naming, + argument validation, result serialisation and logging are handled here so + every provider's implementation presents the same surface to the model. + + :param tool_name: Name the calling model sees. A verb phrase naming the + specialist reads best, e.g. ``ask_bookings_analyst``. + :param description: What this agent knows and when to consult it. Required: + a remote agent's competence cannot be introspected, and this is the only + basis the calling model has for choosing between specialists. + :param timeout: Seconds to wait for a single invocation. ``None`` defers to + the platform default, which subclasses supply -- a number chosen here + would silently disagree with the vendor operator's documented timeout + for the same service. + """ + + #: Whether a completed invocation may be replayed from the durable cache + #: instead of re-invoked. Off by default because a managed agent may act on + #: systems Airflow cannot observe, so replaying a cached answer could skip a + #: side effect. Read-only agents should opt in. + replayable: bool = False + + def __init__( + self, + *, + tool_name: str, + description: str, + timeout: float | None = None, + ) -> None: + if not tool_name: + raise ValueError("tool_name must be a non-empty string.") + if not description or not description.strip(): + raise ValueError( + "description is required: the calling model uses it to decide which " + "specialist to consult, and it cannot be derived from the agent's identifier." + ) + self._tool_name = tool_name + self._description = description + self._timeout = timeout + + @property + @abstractmethod + def agent_ref(self) -> dict[str, str]: + """ + Normalised identity of the remote agent. + + Must contain ``platform`` and ``name``, e.g. + ``{"platform": "snowflake.cortex", "name": "ANALYTICS.REVENUE.BOOKINGS_ANALYST"}``. + Logged on every invocation so a run can be audited for which agents were + consulted. + """ + + @abstractmethod + async def invoke(self, prompt: str) -> Any: + """ + Send ``prompt`` to the remote agent and return the agent's answer. + + Return the answer, not the transport envelope -- whatever the calling + model should actually read. Unwrapping is the implementation's job. + + Failures sort into three buckets, and conflating them is the most common + way an implementation goes wrong: + + * ``pydantic_ai.exceptions.ModelRetry`` -- the remote agent rejected the + request in a way rephrasing could fix. The calling model sees the + message and tries again, bounded by its ``usage_limits``. + * :class:`~airflow.providers.common.ai.exceptions.ManagedAgentInvocationError` + -- terminal. Bad credentials, missing agent, revoked quota. Neither a + rephrase nor a task retry helps, so fail fast. + * Anything transient (429, 5xx, connection reset, read timeout) -- let it + 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. + """ + + @property + def id(self) -> str: + return f"managed-agent-{self._tool_name}" + + async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]: + tool_def = ToolDefinition( + name=self._tool_name, + description=self._description, + parameters_json_schema=_PROMPT_SCHEMA, + # Each invocation is an independent request to a remote service, so + # consulting several specialists concurrently is both safe and the point. + sequential=False, + **return_schema_kwargs({"type": "string"}), + ) + return { + self._tool_name: ToolsetTool( + toolset=self, + tool_def=tool_def, + max_retries=1, + args_validator=build_args_validator(_PROMPT_SCHEMA), + ) + } + + async def call_tool( + self, + name: str, + tool_args: dict[str, Any], + ctx: RunContext[Any], + tool: ToolsetTool[Any], + ) -> Any: + ref = self.agent_ref + 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 + if position: + 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 position else "primary", + }, + ) + 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/src/airflow/providers/common/ai/utils/tool_definition.py b/providers/common/ai/src/airflow/providers/common/ai/utils/tool_definition.py index 9f4fb04b87b..28639701d94 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/utils/tool_definition.py +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/tool_definition.py @@ -19,6 +19,7 @@ from __future__ import annotations import dataclasses +import json from typing import Any, Literal from pydantic_ai.tools import ToolDefinition @@ -45,6 +46,22 @@ def return_schema_kwargs(schema: dict[str, Any]) -> dict[str, Any]: return {} +def serialize_for_llm(value: Any) -> str: + """ + Convert a Python return value to a string suitable for an LLM. + + :param value: The tool's return value. + """ + if value is None: + return "null" + if isinstance(value, str): + return value + try: + return json.dumps(value, default=str) + except (TypeError, ValueError): + return str(value) + + def _fragment_to_core_schema(fragment: dict[str, Any]) -> core_schema.CoreSchema: any_of = fragment.get("anyOf") if isinstance(any_of, list): diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py index ae2d4c6f867..b8ce959db02 100644 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py @@ -27,9 +27,11 @@ from airflow.providers.common.ai.toolsets.hook import ( _build_json_schema_from_signature, _extract_description, _parse_param_docs, - _serialize_for_llm, ) -from airflow.providers.common.ai.utils.tool_definition import _SUPPORTS_RETURN_SCHEMA +from airflow.providers.common.ai.utils.tool_definition import ( + _SUPPORTS_RETURN_SCHEMA, + serialize_for_llm, +) class _FakeHook: @@ -320,20 +322,20 @@ class TestParseParamDocs: class TestSerializeForLlm: def test_string_passthrough(self): - assert _serialize_for_llm("hello") == "hello" + assert serialize_for_llm("hello") == "hello" def test_none_returns_null(self): - assert _serialize_for_llm(None) == "null" + assert serialize_for_llm(None) == "null" def test_dict_to_json(self): - result = _serialize_for_llm({"key": "value"}) + result = serialize_for_llm({"key": "value"}) assert result == '{"key": "value"}' def test_list_to_json(self): - result = _serialize_for_llm([1, 2, 3]) + result = serialize_for_llm([1, 2, 3]) assert result == "[1, 2, 3]" def test_non_serializable_falls_back_to_str(self): obj = object() - result = _serialize_for_llm(obj) + result = serialize_for_llm(obj) assert "object" in 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 new file mode 100644 index 00000000000..55a5b7094f3 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py @@ -0,0 +1,325 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import Any +from unittest import mock + +import pytest +from pydantic_ai.exceptions import ModelRetry +from pydantic_core import ValidationError + +from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError +from airflow.providers.common.ai.toolsets.managed_agent import ( + BaseManagedAgentToolset, + FailoverManagedAgentToolset, +) + + +class FakeManagedAgentToolset(BaseManagedAgentToolset): + """Minimal implementation standing in for a provider's concrete toolset.""" + + def __init__(self, *, result: Any = "the answer", raises: Exception | None = None, **kwargs): + kwargs.setdefault("tool_name", "ask_specialist") + kwargs.setdefault("description", "Answers questions about the thing.") + super().__init__(**kwargs) + self._result = result + self._raises = raises + self.prompts: list[str] = [] + + @property + def agent_ref(self) -> dict[str, str]: + return {"platform": "fake.cloud", "name": "specialist-1"} + + async def invoke(self, prompt: str) -> Any: + self.prompts.append(prompt) + if self._raises is not None: + raise self._raises + return self._result + + +class TestBaseManagedAgentToolsetConstruction: + def test_is_abstract(self): + with pytest.raises(TypeError, match="abstract"): + BaseManagedAgentToolset(tool_name="x", description="y") # type: ignore[abstract] + + @pytest.mark.parametrize( + "description", + ["", " ", "\n"], + ids=["empty", "whitespace", "newline"], + ) + def test_blank_description_rejected(self, description): + with pytest.raises(ValueError, match="description is required"): + FakeManagedAgentToolset(description=description) + + def test_empty_tool_name_rejected(self): + with pytest.raises(ValueError, match="tool_name must be a non-empty string"): + FakeManagedAgentToolset(tool_name="") + + def test_timeout_defaults_to_none_so_subclasses_supply_it(self): + assert FakeManagedAgentToolset()._timeout is None + + def test_subclass_may_default_its_own_timeout(self): + class WithPlatformDefault(FakeManagedAgentToolset): + def __init__(self, **kwargs): + kwargs.setdefault("timeout", 600.0) + super().__init__(**kwargs) + + assert WithPlatformDefault()._timeout == 600.0 + + def test_id_is_derived_from_tool_name(self): + assert FakeManagedAgentToolset(tool_name="ask_bookings").id == "managed-agent-ask_bookings" + + def test_not_replayable_by_default(self): + assert FakeManagedAgentToolset().replayable is False + + +class TestGetTools: + @pytest.mark.asyncio + async def test_exposes_exactly_one_tool_under_its_name(self): + toolset = FakeManagedAgentToolset(tool_name="ask_bookings") + tools = await toolset.get_tools(ctx=None) + assert list(tools) == ["ask_bookings"] + + @pytest.mark.asyncio + async def test_tool_definition_carries_name_description_and_prompt_schema(self): + toolset = FakeManagedAgentToolset(tool_name="ask_bookings", description="Knows bookings.") + tool_def = (await toolset.get_tools(ctx=None))["ask_bookings"].tool_def + + assert tool_def.name == "ask_bookings" + assert tool_def.description == "Knows bookings." + assert tool_def.parameters_json_schema["required"] == ["prompt"] + assert tool_def.parameters_json_schema["properties"]["prompt"]["type"] == "string" + + @pytest.mark.asyncio + async def test_not_sequential_so_specialists_can_be_consulted_concurrently(self): + toolset = FakeManagedAgentToolset() + tool_def = (await toolset.get_tools(ctx=None))["ask_specialist"].tool_def + assert tool_def.sequential is False + + @pytest.mark.asyncio + async def test_args_validator_rejects_a_missing_prompt(self): + toolset = FakeManagedAgentToolset() + validator = (await toolset.get_tools(ctx=None))["ask_specialist"].args_validator + with pytest.raises(ValidationError): + validator.validate_json("{}") + + +class TestCallTool: + async def _call(self, toolset, prompt="what is the number?"): + tools = await toolset.get_tools(ctx=None) + tool = tools[toolset._tool_name] + return await toolset.call_tool(toolset._tool_name, {"prompt": prompt}, None, tool) + + @pytest.mark.asyncio + async def test_passes_the_prompt_through_to_invoke(self): + toolset = FakeManagedAgentToolset() + await self._call(toolset, prompt="how many widgets?") + assert toolset.prompts == ["how many widgets?"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("result", "expected"), + [ + ("plain text", "plain text"), + (None, "null"), + ({"total": 42}, '{"total": 42}'), + ([1, 2], "[1, 2]"), + ], + ids=["str", "none", "dict", "list"], + ) + async def test_result_is_serialised_for_the_model(self, result, expected): + assert await self._call(FakeManagedAgentToolset(result=result)) == expected + + @pytest.mark.asyncio + async def test_logs_the_agent_it_consulted(self, caplog): + await self._call(FakeManagedAgentToolset()) + assert "specialist-1" in caplog.text + assert "fake.cloud" in caplog.text + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "error", + [ + ModelRetry("rephrase that"), + ManagedAgentInvocationError("bad credentials"), + RuntimeError("503 from upstream"), + ], + ids=["model_retry", "terminal", "transient"], + ) + async def test_invoke_errors_propagate_unchanged(self, error): + # The base class must not reclassify what invoke() raised: the three + # buckets are handled by different layers (model, task failure, task retry). + with pytest.raises(type(error), match=str(error)): + await self._call(FakeManagedAgentToolset(raises=error)) + + +class TestFailoverManagedAgentToolset: + @staticmethod + def _group(*members, **kwargs): + kwargs.setdefault("tool_name", "ask_resilient") + kwargs.setdefault("description", "Answers questions, on whichever cloud is up.") + return FailoverManagedAgentToolset(members=list(members), **kwargs) + + async def _call(self, group, prompt="what is the number?"): + tools = await group.get_tools(ctx=None) + return await group.call_tool(group._tool_name, {"prompt": prompt}, None, tools[group._tool_name]) + + @pytest.mark.parametrize("count", [0, 1], ids=["none", "one"]) + def test_needs_at_least_two_members(self, count): + members = [FakeManagedAgentToolset() for _ in range(count)] + with pytest.raises(ValueError, match="at least two members"): + self._group(*members) + + @pytest.mark.asyncio + async def test_primary_answer_wins_and_standby_is_untouched(self): + primary = FakeManagedAgentToolset(result="from primary") + standby = FakeManagedAgentToolset(result="from standby") + assert await self._call(self._group(primary, standby)) == "from primary" + assert standby.prompts == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "error", + [ + ManagedAgentInvocationError("region is down"), + RuntimeError("503 from upstream"), + TimeoutError("read timeout"), + ], + ids=["terminal", "transient", "timeout"], + ) + async def test_fails_over_when_primary_fails(self, error): + primary = FakeManagedAgentToolset(raises=error) + standby = FakeManagedAgentToolset(result="from standby") + assert await self._call(self._group(primary, standby)) == "from standby" + assert standby.prompts == ["what is the number?"] + + @pytest.mark.asyncio + async def test_model_retry_does_not_burn_the_standby(self): + # A prompt the primary could not parse will not parse on the standby + # either, so the model must get the chance to rephrase instead. + primary = FakeManagedAgentToolset(raises=ModelRetry("rephrase that")) + standby = FakeManagedAgentToolset(result="from standby") + with pytest.raises(ModelRetry, match="rephrase that"): + await self._call(self._group(primary, standby)) + assert standby.prompts == [] + + @pytest.mark.asyncio + async def test_last_members_error_propagates_when_all_fail(self): + primary = FakeManagedAgentToolset(raises=RuntimeError("primary down")) + standby = FakeManagedAgentToolset(raises=ManagedAgentInvocationError("standby down")) + with pytest.raises(ManagedAgentInvocationError, match="standby down"): + await self._call(self._group(primary, standby)) + + @pytest.mark.asyncio + async def test_narrowed_failover_on_lets_other_errors_through(self): + primary = FakeManagedAgentToolset(raises=RuntimeError("a bug, not an outage")) + standby = FakeManagedAgentToolset(result="from standby") + group = self._group(primary, standby, failover_on=(ManagedAgentInvocationError,)) + with pytest.raises(RuntimeError, match="a bug, not an outage"): + await self._call(group) + assert standby.prompts == [] + + @pytest.mark.asyncio + async def test_warns_on_failover_and_names_the_standby(self, caplog): + primary = FakeManagedAgentToolset(raises=ManagedAgentInvocationError("down")) + standby = FakeManagedAgentToolset(result="ok") + await self._call(self._group(primary, standby)) + assert "failing over" in caplog.text + assert "served by standby" in caplog.text + + @pytest.mark.asyncio + async def test_groups_nest(self): + inner = self._group( + FakeManagedAgentToolset(raises=ManagedAgentInvocationError("a down")), + FakeManagedAgentToolset(raises=ManagedAgentInvocationError("b down")), + ) + outer = self._group(inner, FakeManagedAgentToolset(result="from outer standby")) + assert await self._call(outer) == "from outer standby" + + @pytest.mark.parametrize( + ("primary_replayable", "standby_replayable", "expected"), + [(True, True, True), (True, False, False), (False, False, False)], + ids=["both", "one", "neither"], + ) + def test_replayable_only_when_every_member_is(self, primary_replayable, standby_replayable, expected): + # The durable cache cannot know which member produced the answer it + # holds, so a single non-replayable member makes the group unsafe. + primary, standby = FakeManagedAgentToolset(), FakeManagedAgentToolset() + primary.replayable, standby.replayable = primary_replayable, standby_replayable + assert self._group(primary, standby).replayable is expected + + def test_agent_ref_describes_the_group(self): + ref = self._group(FakeManagedAgentToolset(), FakeManagedAgentToolset()).agent_ref + assert ref["platform"] == "failover" + assert ref["name"] == "specialist-1 -> specialist-1" + + +class TestFailoverMetrics: + """A failover is a success-shaped event, so the counters are the only signal + distinguishing a healthy primary from one that has been down for a week.""" + + @staticmethod + def _group(*members): + return FailoverManagedAgentToolset( + members=list(members), + tool_name="ask_resilient", + description="Answers questions, on whichever cloud is up.", + ) + + @pytest.mark.asyncio + @mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats") + 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"} + ) + + @pytest.mark.asyncio + @mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats") + async def test_failover_emits_both_counters(self, mock_stats): + primary = FakeManagedAgentToolset(raises=ManagedAgentInvocationError("down")) + await self._group(primary, FakeManagedAgentToolset()).invoke("q") + + assert mock_stats.incr.call_args_list == [ + mock.call( + "managed_agent.failover", + tags={"from_platform": "fake.cloud", "to_platform": "fake.cloud"}, + ), + mock.call("managed_agent.served", tags={"platform": "fake.cloud", "role": "standby"}), + ] + + @pytest.mark.asyncio + @mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats") + async def test_total_outage_records_the_failover_but_no_answer(self, mock_stats): + group = self._group( + FakeManagedAgentToolset(raises=ManagedAgentInvocationError("a down")), + FakeManagedAgentToolset(raises=ManagedAgentInvocationError("b down")), + ) + with pytest.raises(ManagedAgentInvocationError): + await group.invoke("q") + + emitted = [c.args[0] for c in mock_stats.incr.call_args_list] + assert emitted == ["managed_agent.failover"], "no answer means no served counter" + + @pytest.mark.asyncio + @mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats") + async def test_model_retry_is_not_a_failover(self, mock_stats): + group = self._group(FakeManagedAgentToolset(raises=ModelRetry("rephrase")), FakeManagedAgentToolset()) + with pytest.raises(ModelRetry): + await group.invoke("q") + mock_stats.incr.assert_not_called()
