You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2026/10/01 10:49:17 UTC
(airflow) branch main updated: Add a vendor-neutral managed-agent hook contract to Common AI (#73532)
This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e9c552cc827 Add a vendor-neutral managed-agent hook contract to Common AI (#73532)
e9c552cc827 is described below
commit e9c552cc82722679fa6c76545d0482d1cc62e401
Author: Kaxil Naik <ka...@apache.org>
AuthorDate: Thu Oct 1 11:49:07 2026 +0100
Add a vendor-neutral managed-agent hook contract to Common AI (#73532)
---
dev/breeze/tests/test_selective_checks.py | 16 +-
providers/amazon/docs/index.rst | 2 +
providers/amazon/docs/operators/bedrock.rst | 40 ++
providers/amazon/pyproject.toml | 4 +
.../airflow/providers/amazon/aws/hooks/bedrock.py | 229 +++++++++++-
.../tests/unit/amazon/aws/hooks/test_bedrock.py | 280 ++++++++++++++
providers/common/ai/docs/toolsets/index.rst | 35 +-
.../common/ai/docs/toolsets/managed_agent.rst | 337 +++++++++--------
.../providers/common/ai/durable/caching_toolset.py | 13 +-
.../src/airflow/providers/common/ai/exceptions.py | 15 +-
.../providers/common/ai/managed_agents/__init__.py | 49 +++
.../providers/common/ai/managed_agents/base.py | 251 +++++++++++++
.../providers/common/ai/managed_agents/failover.py | 145 ++++++++
.../providers/common/ai/toolsets/__init__.py | 4 +-
.../providers/common/ai/toolsets/managed_agent.py | 354 +++++++++---------
.../unit/common/ai/durable/test_caching_toolset.py | 55 +++
.../unit/common/ai/managed_agents/__init__.py | 16 +
.../unit/common/ai/managed_agents/test_base.py | 137 +++++++
.../unit/common/ai/managed_agents/test_failover.py | 205 +++++++++++
.../unit/common/ai/toolsets/test_managed_agent.py | 403 ++++++++-------------
providers/google/docs/index.rst | 2 +
.../google/docs/operators/cloud/vertex_ai.rst | 24 ++
providers/google/pyproject.toml | 4 +
.../google/cloud/hooks/vertex_ai/agent_engine.py | 159 +++++++-
.../cloud/hooks/vertex_ai/test_agent_engine.py | 145 ++++++++
.../observability/metrics/metrics_template.yaml | 20 +-
uv.lock | 16 +-
27 files changed, 2327 insertions(+), 633 deletions(-)
diff --git a/dev/breeze/tests/test_selective_checks.py b/dev/breeze/tests/test_selective_checks.py
index 533cd41a6a8..3c81c0ee901 100644
--- a/dev/breeze/tests/test_selective_checks.py
+++ b/dev/breeze/tests/test_selective_checks.py
@@ -1130,7 +1130,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
("providers/amazon/src/airflow/providers/amazon/provider.yaml",),
{
"selected-providers-list-as-string": "amazon apache.hive cncf.kubernetes "
- "common.compat common.messaging common.sql databricks duckdb exasol ftp google http imap "
+ "common.ai common.compat common.messaging common.sql databricks duckdb exasol ftp google http imap "
"microsoft.azure "
"mongo mysql openlineage postgres salesforce ssh teradata",
"all-python-versions": f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']",
@@ -1157,7 +1157,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
{
"description": "amazon...google",
"test_types": "Providers[amazon] Providers[apache.hive,cncf.kubernetes,"
- "common.compat,common.messaging,common.sql,databricks,duckdb,exasol,ftp,http,imap,"
+ "common.ai,common.compat,common.messaging,common.sql,databricks,duckdb,exasol,ftp,http,imap,"
"microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] "
"Providers[google]",
}
@@ -1203,7 +1203,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
("providers/amazon/src/airflow/providers/amazon/file.py",),
{
"selected-providers-list-as-string": "amazon apache.hive cncf.kubernetes "
- "common.compat common.messaging common.sql databricks duckdb exasol ftp google http imap "
+ "common.ai common.compat common.messaging common.sql databricks duckdb exasol ftp google http imap "
"microsoft.azure "
"mongo mysql openlineage postgres salesforce ssh teradata",
"all-python-versions": f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']",
@@ -1227,7 +1227,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
{
"description": "amazon...google",
"test_types": "Providers[amazon] Providers[apache.hive,cncf.kubernetes,"
- "common.compat,common.messaging,common.sql,databricks,duckdb,exasol,ftp,http,imap,"
+ "common.ai,common.compat,common.messaging,common.sql,databricks,duckdb,exasol,ftp,http,imap,"
"microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] "
"Providers[google]",
}
@@ -2797,7 +2797,7 @@ def test_expected_output_push(
),
{
"selected-providers-list-as-string": "amazon apache.beam apache.cassandra apache.kafka "
- "cncf.kubernetes common.compat common.messaging common.sql databricks "
+ "cncf.kubernetes common.ai common.compat common.messaging common.sql databricks "
"facebook google hashicorp http microsoft.azure microsoft.mssql mongo mysql "
"openlineage oracle postgres presto salesforce samba sftp ssh standard trino",
"all-python-versions": f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']",
@@ -2809,7 +2809,7 @@ def test_expected_output_push(
"skip-providers-tests": "false",
"docs-build": "true",
"docs-list-as-string": "apache-airflow helm-chart amazon apache.beam apache.cassandra "
- "apache.kafka cncf.kubernetes common.compat common.messaging common.sql databricks facebook google hashicorp http microsoft.azure "
+ "apache.kafka cncf.kubernetes common.ai common.compat common.messaging common.sql databricks facebook google hashicorp http microsoft.azure "
"microsoft.mssql mongo mysql openlineage oracle postgres "
"presto salesforce samba sftp ssh standard trino",
"skip-prek-hooks": (
@@ -2832,7 +2832,7 @@ def test_expected_output_push(
{
"description": "amazon...standard",
"test_types": "Providers[amazon] Providers[apache.beam,apache.cassandra,"
- "apache.kafka,cncf.kubernetes,common.compat,common.messaging,common.sql,databricks,facebook,"
+ "apache.kafka,cncf.kubernetes,common.ai,common.compat,common.messaging,common.sql,databricks,facebook,"
"hashicorp,http,microsoft.azure,microsoft.mssql,mongo,mysql,"
"openlineage,oracle,postgres,presto,salesforce,samba,sftp,ssh,trino] "
"Providers[google] "
@@ -3115,7 +3115,7 @@ def test_upgrade_to_newer_dependencies(
("providers/google/docs/some_file.rst",),
{
"docs-list-as-string": "amazon apache.beam apache.cassandra apache.kafka "
- "cncf.kubernetes common.compat common.messaging common.sql databricks facebook google hashicorp http "
+ "cncf.kubernetes common.ai common.compat common.messaging common.sql databricks facebook google hashicorp http "
"microsoft.azure microsoft.mssql mongo mysql openlineage oracle "
"postgres presto salesforce samba sftp ssh standard trino",
},
diff --git a/providers/amazon/docs/index.rst b/providers/amazon/docs/index.rst
index 970d7f5cd32..25da5ef57d5 100644
--- a/providers/amazon/docs/index.rst
+++ b/providers/amazon/docs/index.rst
@@ -149,6 +149,7 @@ Dependent package
======================================================================================================================== ====================
`apache-airflow-providers-apache-hive <https://airflow.apache.org/docs/apache-airflow-providers-apache-hive>`_ ``apache.hive``
`apache-airflow-providers-cncf-kubernetes <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes>`_ ``cncf.kubernetes``
+`apache-airflow-providers-common-ai <https://airflow.apache.org/docs/apache-airflow-providers-common-ai>`_ ``common.ai``
`apache-airflow-providers-common-messaging <https://airflow.apache.org/docs/apache-airflow-providers-common-messaging>`_ ``common.messaging``
`apache-airflow-providers-duckdb <https://airflow.apache.org/docs/apache-airflow-providers-duckdb>`_ ``duckdb``
`apache-airflow-providers-exasol <https://airflow.apache.org/docs/apache-airflow-providers-exasol>`_ ``exasol``
@@ -194,6 +195,7 @@ Extra Dependencies
``salesforce`` ``apache-airflow-providers-salesforce``
``ssh`` ``apache-airflow-providers-ssh``
``standard`` ``apache-airflow-providers-standard``
+``common.ai`` ``apache-airflow-providers-common-ai>=0.10.0``
``common.messaging`` ``apache-airflow-providers-common-messaging>=2.0.0``
``duckdb`` ``apache-airflow-providers-duckdb>=0.1.0``
``sqlalchemy`` ``sqlalchemy>=1.4.54,!=2.1.0``
diff --git a/providers/amazon/docs/operators/bedrock.rst b/providers/amazon/docs/operators/bedrock.rst
index c9669ca7dfe..74af2f116a7 100644
--- a/providers/amazon/docs/operators/bedrock.rst
+++ b/providers/amazon/docs/operators/bedrock.rst
@@ -110,6 +110,46 @@ through ``botocore_config``.
:start-after: [START howto_operator_bedrock_invoke_agent_runtime]
:end-before: [END howto_operator_bedrock_invoke_agent_runtime]
+.. _howto/hook:BedrockAgentCoreHook:managed-agent:
+
+Consult an Amazon Bedrock AgentCore Runtime from a Common AI agent
+==================================================================
+
+To let an agent running under Common AI's ``AgentOperator`` consult an AgentCore Runtime as one
+of its tools, use
+:class:`~airflow.providers.amazon.aws.hooks.bedrock.BedrockAgentCoreHook`.
+It implements the Common AI managed-agent contract, so ``hook.agent(runtime_arn)`` can be passed
+to a ``ManagedAgentToolset`` or combined with agents on other clouds in a
+``FailoverManagedAgentClient``. The hook needs the ``common.ai`` extra of this provider, which
+installs ``apache-airflow-providers-common-ai`` and therefore requires Airflow 3.
+
+The agent is the runtime ARN and the session is AgentCore's ``runtimeSessionId``, which the
+service requires to be 33 to 256 characters long. A prompt is
+sent as ``{"prompt": ...}`` in an ``application/json`` payload; the container behind the
+runtime defines the response shape, so pass ``text_key`` in ``vendor_options`` when you know which
+field holds the answer. The hook disables botocore retries unless the connection or caller configured them,
+because an invocation's effects are unknown, and raises a terminal error for non-JSON responses
+or bodies over the size limit rather than returning something misleading.
+
+.. code-block:: python
+
+ from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentCoreHook
+ from airflow.providers.common.ai.toolsets import ManagedAgentToolset
+
+ claims = BedrockAgentCoreHook(aws_conn_id="aws_default", region_name="us-east-1").agent(
+ "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/claims"
+ )
+ toolset = ManagedAgentToolset(
+ claims,
+ tool_name="ask_claims_agent",
+ description="Reviews an insurance claim and returns a coverage determination.",
+ vendor_options={"text_key": "determination"},
+ )
+
+For long-running invocations run as a pipeline step in their own right, use
+:class:`~airflow.providers.amazon.aws.operators.bedrock.BedrockInvokeAgentRuntimeOperator`
+instead; a toolset call blocks a worker thread for the duration of the call.
+
.. _howto/operator:BedrockDeleteAgentRuntimeOperator:
Delete an Amazon Bedrock AgentCore Runtime
diff --git a/providers/amazon/pyproject.toml b/providers/amazon/pyproject.toml
index 7a8b3e9c5d6..310290f1633 100644
--- a/providers/amazon/pyproject.toml
+++ b/providers/amazon/pyproject.toml
@@ -149,6 +149,9 @@ dependencies = [
"standard" = [
"apache-airflow-providers-standard"
]
+"common.ai" = [
+ "apache-airflow-providers-common-ai>=0.10.0", # use next version
+]
"common.messaging" = [
"apache-airflow-providers-common-messaging>=2.0.0"
]
@@ -168,6 +171,7 @@ dev = [
"apache-airflow-devel-common",
"apache-airflow-providers-apache-hive",
"apache-airflow-providers-cncf-kubernetes",
+ "apache-airflow-providers-common-ai",
"apache-airflow-providers-common-compat",
"apache-airflow-providers-common-messaging",
"apache-airflow-providers-common-sql",
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/bedrock.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/bedrock.py
index 9042f0754b7..c68050e10d0 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/bedrock.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/bedrock.py
@@ -16,7 +16,74 @@
# under the License.
from __future__ import annotations
+import json
+import threading
+from contextlib import closing
+from typing import TYPE_CHECKING, Any
+
+from botocore.config import Config
+from botocore.exceptions import ClientError
+
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+if TYPE_CHECKING:
+ from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError
+ from airflow.providers.common.ai.managed_agents.base import (
+ BaseManagedAgentHook,
+ ManagedAgentCapabilities,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ ManagedAgentResponse,
+ )
+else:
+ try:
+ from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError
+ from airflow.providers.common.ai.managed_agents.base import (
+ BaseManagedAgentHook,
+ ManagedAgentCapabilities,
+ ManagedAgentRef,
+ ManagedAgentResponse,
+ )
+ except ImportError:
+ # The Common AI provider is optional. This module still imports without it, and every
+ # managed-agent entry point on BedrockAgentCoreHook then says what is missing.
+ def _needs_common_ai(*args: Any, **kwargs: Any) -> Any:
+ raise AirflowOptionalProviderFeatureException(
+ "Consulting an AgentCore runtime as a managed agent needs the 'common.ai' extra of the "
+ "amazon provider: pip install 'apache-airflow-providers-amazon[common.ai]'."
+ )
+
+ class BaseManagedAgentHook:
+ """Stand-in for the Common AI contract base; ``agent()`` names the missing extra."""
+
+ agent = _needs_common_ai
+
+ ManagedAgentCapabilities = ManagedAgentRef = ManagedAgentResponse = _needs_common_ai
+ ManagedAgentInvocationError = _needs_common_ai
+
+# Fields the managed-agent contract already covers. Letting ``vendor_options`` carry them would
+# let a caller re-target the call, which is a change of authority, not an option.
+_RESERVED_OPTIONS = frozenset(
+ {"agentRuntimeArn", "payload", "contentType", "accept", "runtimeSessionId", "accountId", "mcpSessionId"}
+)
+# Error codes InvokeAgentRuntime can return that no retry or rephrase will fix. The container's own
+# errors come back inside a 200 body, so AgentCore has no error that means "rephrase the prompt".
+_TERMINAL_ERROR_CODES = frozenset(
+ {
+ "ValidationException",
+ "ResourceNotFoundException",
+ "AccessDeniedException",
+ "ServiceQuotaExceededException",
+ }
+)
+# Do not retry an invocation whose effects are unknown; Airflow's task-level retry is the right layer.
+_NO_RETRIES = Config(retries={"total_max_attempts": 1})
+# The service model's bounds for runtimeSessionId. botocore rejects a shorter value client-side with a
+# ParamValidationError, which is not a ClientError and would otherwise escape the error classes.
+_SESSION_ID_LENGTH = range(33, 257)
+_MAX_RESPONSE_BYTES = 1024 * 1024
+_TEXT_KEYS = ("output", "result", "text", "response")
class BedrockHook(AwsBaseHook):
@@ -128,21 +195,177 @@ class BedrockAgentCoreControlHook(AwsBaseHook):
super().__init__(*args, **kwargs)
-class BedrockAgentCoreHook(AwsBaseHook):
+class BedrockAgentCoreHook(AwsBaseHook, BaseManagedAgentHook):
"""
Interact with the Amazon Bedrock AgentCore runtime plane API.
Provide thin wrapper around :external+boto3:py:class:`boto3.client("bedrock-agentcore") <BedrockAgentCore.Client>`.
- Additional arguments (such as ``aws_conn_id``) may be specified and
- are passed down to the underlying AwsBaseHook.
+ Additional arguments (such as ``aws_conn_id`` and ``config``) may be specified and
+ are passed down to the underlying AwsBaseHook; the connection's ``config_kwargs`` apply
+ as they do for every other AWS hook.
+
+ With the ``common.ai`` extra installed, the hook also implements the Common AI
+ managed-agent contract, so ``hook.agent(runtime_arn)`` can be handed to a
+ ``ManagedAgentToolset``. The agent is the runtime ARN; the session is AgentCore's
+ ``runtimeSessionId``, which the service requires to be 33 to 256 characters long. A request
+ carrying a ``prompt`` is sent as ``{"prompt": ...}`` and a request carrying ``messages`` as
+ ``{"messages": [...]}``, both as ``application/json``. The container behind the runtime
+ defines its own response shape, so the answer text is taken from the first of ``output``,
+ ``result``, ``text`` or ``response`` that holds a string, or from the ``text_key`` vendor
+ option when the container's contract is known; otherwise the whole JSON body is returned as
+ text. The decoded body is always available on ``ManagedAgentResponse.raw``.
+
+ A remote invocation may have unknown effects, so the contract methods disable botocore's
+ retries unless the connection or the caller configured them, and let failures propagate to
+ Airflow's task-level retry instead. ``ManagedAgentRequest.timeout`` is honored as the botocore
+ connect and read timeout of the call; a client is built per distinct timeout and reused.
+ AgentCore has no error that means "rephrase the prompt" (a container's own errors arrive
+ inside a successful body), so the hook raises terminal errors or lets transient ones
+ propagate, never :class:`~airflow.providers.common.ai.exceptions.ManagedAgentRejected`.
+
+ .. code-block:: python
+
+ from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentCoreHook
+ from airflow.providers.common.ai.toolsets import ManagedAgentToolset
+
+ claims = BedrockAgentCoreHook(aws_conn_id="aws_prod", region_name="us-east-1").agent(
+ "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/claims"
+ )
+ toolset = ManagedAgentToolset(
+ claims, tool_name="ask_claims_agent", description="...", vendor_options={"text_key": "answer"}
+ )
+
+ Two ``vendor_options`` are read by the hook itself rather than forwarded to
+ ``InvokeAgentRuntime``: ``text_key`` names the response field that holds the answer text when
+ the container's contract is known (a body without a string there is an error), and
+ ``max_response_bytes`` bounds the body read into worker memory (default 1 MiB). Every other
+ option is passed to the API call as is.
.. seealso::
- :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+ - :class:`airflow.providers.common.ai.managed_agents.base.BaseManagedAgentHook`
"""
client_type = "bedrock-agentcore"
+ agent_platform = "aws.bedrock_agentcore"
def __init__(self, *args, **kwargs) -> None:
kwargs["client_type"] = self.client_type
super().__init__(*args, **kwargs)
+ self._clients: dict[float | None, Any] = {}
+ self._clients_lock = threading.Lock()
+
+ def resolve_agent(self, agent: str) -> ManagedAgentRef:
+ if not agent.startswith("arn:") or ":runtime/" not in agent:
+ raise ValueError(f"An AgentCore agent is a runtime ARN, got {agent!r}.")
+ return ManagedAgentRef(platform=self.agent_platform, name=agent)
+
+ def get_agent_capabilities(self, agent: str) -> ManagedAgentCapabilities:
+ return ManagedAgentCapabilities(sessions=True, structured_output=True, trace=True)
+
+ def invoke_agent(self, agent: str, request: ManagedAgentRequest) -> ManagedAgentResponse:
+ self.resolve_agent(agent)
+ if reserved := _RESERVED_OPTIONS.intersection(request.vendor_options):
+ raise ValueError(f"vendor_options cannot override contract fields: {sorted(reserved)}")
+ if request.session_id is not None and len(request.session_id) not in _SESSION_ID_LENGTH:
+ raise ValueError(
+ f"AgentCore requires a session id of {_SESSION_ID_LENGTH.start} to {_SESSION_ID_LENGTH[-1]} "
+ f"characters; got {len(request.session_id)}."
+ )
+ payload = (
+ {"prompt": request.prompt} if request.prompt is not None else {"messages": request.as_messages()}
+ )
+ kwargs: dict[str, Any] = dict(request.vendor_options)
+ text_key = kwargs.pop("text_key", None)
+ if text_key is not None and not isinstance(text_key, str):
+ raise ValueError(f"vendor_options['text_key'] must be a string, got {text_key!r}.")
+ max_response_bytes = kwargs.pop("max_response_bytes", _MAX_RESPONSE_BYTES)
+ if not isinstance(max_response_bytes, int) or max_response_bytes <= 0:
+ raise ValueError(
+ f"vendor_options['max_response_bytes'] must be a positive integer, got {max_response_bytes!r}."
+ )
+ if request.session_id is not None:
+ kwargs["runtimeSessionId"] = request.session_id
+ try:
+ response = self._get_client(request.timeout).invoke_agent_runtime(
+ agentRuntimeArn=agent,
+ payload=json.dumps(payload).encode(),
+ contentType="application/json",
+ accept="application/json",
+ **kwargs,
+ )
+ body = self._read_json_body(agent, response, max_response_bytes)
+ except ClientError as exc:
+ if exc.response.get("Error", {}).get("Code") in _TERMINAL_ERROR_CODES:
+ raise ManagedAgentInvocationError(f"{self._describe_call(agent)}: {exc}") from exc
+ raise # throttling, conflicts, server errors: Airflow's task retry is the right layer
+ return ManagedAgentResponse(
+ text=self._extract_text(agent, body, text_key),
+ raw={**response, "response": body},
+ structured=None if isinstance(body, str) else body,
+ session_id=response.get("runtimeSessionId"),
+ trace_ref=response.get("ResponseMetadata", {}).get("RequestId"),
+ )
+
+ def _get_client(self, timeout: float | None) -> Any:
+ """
+ One boto3 client per distinct request timeout, built on first use and reused.
+
+ The timeout is a client setting, so it cannot ride on the hook's shared client, and a
+ toolset uses one timeout, so this is one client per hook in practice. Clients are
+ thread-safe, which the toolset relies on when a model issues two calls in one turn.
+ """
+ with self._clients_lock:
+ client = self._clients.get(timeout)
+ if client is None:
+ client = self._clients[timeout] = self.get_client_type(config=self._call_config(timeout))
+ return client
+
+ def _call_config(self, timeout: float | None) -> Config:
+ """Return the connection's or caller's botocore config, with retries off unless they set them."""
+ base = self.config
+ if base.retries is None:
+ base = base.merge(_NO_RETRIES)
+ if timeout is None:
+ return base
+ return base.merge(Config(connect_timeout=timeout, read_timeout=timeout))
+
+ def _describe_call(self, agent: str) -> str:
+ return f"AgentCore agent {agent} via connection {self.aws_conn_id!r}"
+
+ def _read_json_body(self, agent: str, response: dict[str, Any], max_response_bytes: int) -> Any:
+ with closing(response["response"]) as stream:
+ content_type = response.get("contentType", "").split(";", 1)[0].strip().lower()
+ if content_type != "application/json":
+ raise ManagedAgentInvocationError(
+ f"{self._describe_call(agent)} returned {content_type or 'no Content-Type'}; "
+ "this hook handles application/json only."
+ )
+ data = stream.read(max_response_bytes + 1)
+ if len(data) > max_response_bytes:
+ raise ManagedAgentInvocationError(
+ f"{self._describe_call(agent)} returned more than max_response_bytes={max_response_bytes}."
+ )
+ try:
+ return json.loads(data)
+ except ValueError as exc:
+ raise ManagedAgentInvocationError(
+ f"{self._describe_call(agent)} returned a body that is not JSON: {exc}"
+ ) from exc
+
+ def _extract_text(self, agent: str, body: Any, text_key: str | None) -> str:
+ if text_key is not None:
+ value = body.get(text_key) if isinstance(body, dict) else None
+ if not isinstance(value, str):
+ raise ManagedAgentInvocationError(
+ f"{self._describe_call(agent)} returned no string at text_key={text_key!r}."
+ )
+ return value
+ if isinstance(body, str):
+ return body
+ if isinstance(body, dict):
+ for key in _TEXT_KEYS:
+ if isinstance(body.get(key), str):
+ return body[key]
+ return json.dumps(body, default=str)
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_bedrock.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_bedrock.py
index b8577a99290..8c541c8502d 100644
--- a/providers/amazon/tests/unit/amazon/aws/hooks/test_bedrock.py
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_bedrock.py
@@ -16,9 +16,16 @@
# under the License.
from __future__ import annotations
+import io
+import json
+import subprocess
+import sys
from unittest import mock
import pytest
+from botocore.config import Config
+from botocore.exceptions import ClientError
+from botocore.session import Session
from airflow.providers.amazon.aws.hooks.bedrock import (
BedrockAgentCoreControlHook,
@@ -29,6 +36,21 @@ from airflow.providers.amazon.aws.hooks.bedrock import (
BedrockRuntimeHook,
)
+try:
+ from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError
+ from airflow.providers.common.ai.managed_agents import (
+ ManagedAgentCapabilities,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ )
+except ImportError:
+ HAS_COMMON_AI = False
+else:
+ HAS_COMMON_AI = True
+
+# The managed-agent methods of BedrockAgentCoreHook need the optional common.ai extra.
+needs_common_ai = pytest.mark.skipif(not HAS_COMMON_AI, reason="requires apache-airflow-providers-common-ai")
+
class TestBedrockHooks:
@pytest.mark.parametrize(
@@ -69,3 +91,261 @@ class TestBedrockHookGetGuardrailIdByName:
hook = BedrockHook()
assert hook.get_guardrail_id_by_name("nonexistent") is None
+
+
+ARN = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/investigator"
+SESSION = "caller-session-0123456789abcdef0123456789abcdef" # AgentCore needs 33 to 256 characters
+
+
+@pytest.fixture
+def invoke_agent_runtime():
+ """The real hook and botocore client, stubbed at the API call."""
+ with mock.patch.dict("os.environ", {"AWS_EC2_METADATA_DISABLED": "true"}):
+ client = Session().create_client(
+ "bedrock-agentcore",
+ region_name="us-east-1",
+ aws_access_key_id="test",
+ aws_secret_access_key="test",
+ )
+ with mock.patch.object(BedrockAgentCoreHook, "get_client_type", autospec=True, return_value=client):
+ with mock.patch.object(client, "invoke_agent_runtime", autospec=True) as call:
+ yield call
+ client.close()
+
+
+def respond(call, body, **extra):
+ stream = io.BytesIO(json.dumps(body).encode())
+ call.return_value = {
+ "response": stream,
+ "contentType": "application/json",
+ "runtimeSessionId": "session-from-aws",
+ "ResponseMetadata": {"RequestId": "req-1"},
+ **extra,
+ }
+ return stream
+
+
+def create_bedrock_agent_core_hook(**kwargs) -> BedrockAgentCoreHook:
+ return BedrockAgentCoreHook(aws_conn_id="aws_default", region_name="us-east-1", **kwargs)
+
+
+@needs_common_ai
+class TestBedrockAgentCoreHookAgent:
+ def test_resolves_a_runtime_arn(self):
+ assert create_bedrock_agent_core_hook().resolve_agent(ARN) == ManagedAgentRef(
+ platform="aws.bedrock_agentcore", name=ARN
+ )
+
+ @pytest.mark.parametrize("agent", ["investigator", "arn:aws:bedrock:us-east-1:1:agent/x"])
+ def test_rejects_anything_but_a_runtime_arn(self, agent):
+ with pytest.raises(ValueError, match="runtime ARN"):
+ create_bedrock_agent_core_hook().resolve_agent(agent)
+
+ def test_capabilities(self):
+ assert create_bedrock_agent_core_hook().get_agent_capabilities(ARN) == ManagedAgentCapabilities(
+ sessions=True, structured_output=True, trace=True
+ )
+
+ def test_botocore_does_not_retry_an_invocation_with_unknown_effects(self):
+ assert create_bedrock_agent_core_hook()._call_config(None).retries == {"total_max_attempts": 1}
+
+ @pytest.mark.parametrize(
+ "config", [{"read_timeout": 900}, Config(read_timeout=900)], ids=["dict", "Config"]
+ )
+ def test_caller_config_is_kept_and_only_retries_defaulted(self, config):
+ call_config = create_bedrock_agent_core_hook(config=config)._call_config(None)
+ assert call_config.read_timeout == 900
+ assert call_config.retries == {"total_max_attempts": 1}
+
+ def test_caller_retries_win_over_the_default(self):
+ call_config = create_bedrock_agent_core_hook(
+ config=Config(retries={"total_max_attempts": 3})
+ )._call_config(None)
+ assert call_config.retries == {"total_max_attempts": 3}
+
+ def test_connection_config_kwargs_are_honored_like_every_other_aws_hook(self, monkeypatch):
+ # Handing the base an explicit config would make the connection's config_kwargs vanish.
+ extra = {"config_kwargs": {"read_timeout": 1234, "proxies": {"https": "http://proxy:3128"}}}
+ monkeypatch.setenv("AIRFLOW_CONN_AWS_PROBE", json.dumps({"conn_type": "aws", "extra": extra}))
+ call_config = BedrockAgentCoreHook(aws_conn_id="aws_probe", region_name="us-east-1")._call_config(
+ None
+ )
+ assert call_config.read_timeout == 1234
+ assert call_config.proxies == {"https": "http://proxy:3128"}
+ assert call_config.retries == {"total_max_attempts": 1}
+
+
+@needs_common_ai
+class TestBedrockAgentCoreHookInvokeAgent:
+ def test_prompt_session_and_passthrough_reach_the_api(self, invoke_agent_runtime):
+ stream = respond(invoke_agent_runtime, {"result": "Looks fine", "confidence": 0.9})
+ response = (
+ create_bedrock_agent_core_hook()
+ .agent(ARN)
+ .invoke(
+ ManagedAgentRequest(prompt="Check it", session_id=SESSION, vendor_options={"qualifier": "v2"})
+ )
+ )
+ kwargs = invoke_agent_runtime.call_args.kwargs
+ assert kwargs["agentRuntimeArn"] == ARN
+ assert json.loads(kwargs["payload"]) == {"prompt": "Check it"}
+ assert kwargs["contentType"] == kwargs["accept"] == "application/json"
+ assert kwargs["runtimeSessionId"] == SESSION
+ assert kwargs["qualifier"] == "v2"
+ assert response.text == "Looks fine"
+ assert response.structured == {"result": "Looks fine", "confidence": 0.9}
+ assert response.raw["response"] == {"result": "Looks fine", "confidence": 0.9}
+ assert response.session_id == "session-from-aws"
+ assert response.trace_ref == "req-1"
+ assert stream.closed
+
+ def test_messages_are_sent_as_messages(self, invoke_agent_runtime):
+ respond(invoke_agent_runtime, {"output": "ok"})
+ messages = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]
+ create_bedrock_agent_core_hook().agent(ARN).invoke(ManagedAgentRequest(messages=messages))
+ assert json.loads(invoke_agent_runtime.call_args.kwargs["payload"]) == {"messages": messages}
+
+ @pytest.mark.parametrize("session_id", ["short", "x" * 257])
+ def test_a_session_id_outside_the_services_bounds_is_refused_before_the_call(
+ self, invoke_agent_runtime, session_id
+ ):
+ with pytest.raises(ValueError, match="33 to 256"):
+ create_bedrock_agent_core_hook().agent(ARN).invoke(
+ ManagedAgentRequest(prompt="x", session_id=session_id)
+ )
+ invoke_agent_runtime.assert_not_called()
+
+ def test_text_key_overrides_the_default_lookup(self, invoke_agent_runtime):
+ respond(invoke_agent_runtime, {"output": "wrong", "answer": "right"})
+ request = ManagedAgentRequest(prompt="x", vendor_options={"text_key": "answer"})
+ assert create_bedrock_agent_core_hook().agent(ARN).invoke(request).text == "right"
+
+ @pytest.mark.parametrize("body", [{"output": "something"}, "just a string"], ids=["dict", "str"])
+ def test_missing_text_key_is_terminal_not_silent(self, invoke_agent_runtime, body):
+ respond(invoke_agent_runtime, body)
+ with pytest.raises(ManagedAgentInvocationError, match="text_key='answer'"):
+ create_bedrock_agent_core_hook().agent(ARN).invoke(
+ ManagedAgentRequest(prompt="x", vendor_options={"text_key": "answer"})
+ )
+
+ def test_a_body_that_is_not_json_is_terminal(self, invoke_agent_runtime):
+ respond(invoke_agent_runtime, "ignored")
+ invoke_agent_runtime.return_value["response"] = io.BytesIO(b"{not json")
+ with pytest.raises(ManagedAgentInvocationError, match="not JSON"):
+ create_bedrock_agent_core_hook().agent(ARN).invoke(ManagedAgentRequest(prompt="x"))
+
+ def test_terminal_errors_name_the_agent_and_the_connection(self, invoke_agent_runtime):
+ invoke_agent_runtime.side_effect = ClientError(
+ {"Error": {"Code": "AccessDeniedException", "Message": "m"}}, "InvokeAgentRuntime"
+ )
+ with pytest.raises(ManagedAgentInvocationError, match=f"{ARN} via connection 'aws_default'"):
+ create_bedrock_agent_core_hook().agent(ARN).invoke(ManagedAgentRequest(prompt="x"))
+
+ def test_unknown_shape_is_dumped_not_guessed(self, invoke_agent_runtime):
+ respond(invoke_agent_runtime, {"findings": ["a"], "score": 1})
+ response = create_bedrock_agent_core_hook().agent(ARN).invoke(ManagedAgentRequest(prompt="x"))
+ assert json.loads(response.text) == {"findings": ["a"], "score": 1}
+
+ @pytest.mark.parametrize("option", ["agentRuntimeArn", "accountId", "mcpSessionId"])
+ def test_vendor_options_cannot_retarget_the_call(self, invoke_agent_runtime, option):
+ with pytest.raises(ValueError, match=option):
+ create_bedrock_agent_core_hook().agent(ARN).invoke(
+ ManagedAgentRequest(prompt="x", vendor_options={option: "other"})
+ )
+ invoke_agent_runtime.assert_not_called()
+
+ def test_request_timeout_becomes_the_call_clients_timeouts(self, invoke_agent_runtime):
+ respond(invoke_agent_runtime, {"output": "ok"})
+ create_bedrock_agent_core_hook(config={"read_timeout": 900}).agent(ARN).invoke(
+ ManagedAgentRequest(prompt="x", timeout=15)
+ )
+ config = BedrockAgentCoreHook.get_client_type.call_args.kwargs["config"]
+ assert (config.read_timeout, config.connect_timeout) == (15, 15)
+ assert config.retries == {"total_max_attempts": 1}
+
+ def test_no_request_timeout_keeps_the_hooks_config(self, invoke_agent_runtime):
+ respond(invoke_agent_runtime, {"output": "ok"})
+ create_bedrock_agent_core_hook(config={"read_timeout": 900}).agent(ARN).invoke(
+ ManagedAgentRequest(prompt="x")
+ )
+ assert BedrockAgentCoreHook.get_client_type.call_args.kwargs["config"].read_timeout == 900
+
+ def test_one_client_per_timeout_reused_across_calls(self, invoke_agent_runtime):
+ bound = create_bedrock_agent_core_hook().agent(ARN)
+ for prompt, timeout in (("x", 15), ("y", 15), ("z", 30)):
+ respond(invoke_agent_runtime, {"output": "ok"})
+ bound.invoke(ManagedAgentRequest(prompt=prompt, timeout=timeout))
+ timeouts = [
+ c.kwargs["config"].read_timeout for c in BedrockAgentCoreHook.get_client_type.call_args_list
+ ]
+ assert timeouts == [15, 30]
+
+ @pytest.mark.parametrize(
+ ("code", "terminal"),
+ [
+ ("ValidationException", True),
+ ("ResourceNotFoundException", True),
+ ("AccessDeniedException", True),
+ ("ServiceQuotaExceededException", True),
+ ("ThrottlingException", False),
+ ("RuntimeClientError", False),
+ ],
+ )
+ def test_client_errors_are_terminal_or_propagate(self, invoke_agent_runtime, code, terminal):
+ invoke_agent_runtime.side_effect = ClientError(
+ {"Error": {"Code": code, "Message": "m"}}, "InvokeAgentRuntime"
+ )
+ # Resolved here rather than in the parametrize list, which is evaluated even where common.ai
+ # is absent and the class is skipped.
+ with pytest.raises(ManagedAgentInvocationError if terminal else ClientError):
+ create_bedrock_agent_core_hook().agent(ARN).invoke(ManagedAgentRequest(prompt="x"))
+
+ def test_non_json_content_type_is_terminal(self, invoke_agent_runtime):
+ stream = respond(invoke_agent_runtime, "ignored", contentType="text/event-stream")
+ with pytest.raises(ManagedAgentInvocationError, match="application/json only"):
+ create_bedrock_agent_core_hook().agent(ARN).invoke(ManagedAgentRequest(prompt="x"))
+ assert stream.closed
+
+ def test_oversized_body_is_terminal_and_closed(self, invoke_agent_runtime):
+ stream = respond(invoke_agent_runtime, "x" * 64)
+ with pytest.raises(ManagedAgentInvocationError, match="max_response_bytes=32"):
+ create_bedrock_agent_core_hook().agent(ARN).invoke(
+ ManagedAgentRequest(prompt="x", vendor_options={"max_response_bytes": 32})
+ )
+ assert stream.closed
+
+ @pytest.mark.parametrize(
+ "options",
+ [
+ {"text_key": 7},
+ {"max_response_bytes": 0},
+ {"max_response_bytes": "big"},
+ {"max_response_bytes": 1.5},
+ ],
+ ids=["text_key-not-str", "zero-bytes", "bytes-not-int", "bytes-float"],
+ )
+ def test_malformed_hook_options_are_rejected_before_the_call(self, invoke_agent_runtime, options):
+ with pytest.raises(ValueError, match="vendor_options"):
+ create_bedrock_agent_core_hook().agent(ARN).invoke(
+ ManagedAgentRequest(prompt="x", vendor_options=options)
+ )
+ invoke_agent_runtime.assert_not_called()
+
+
+def test_the_hook_imports_without_common_ai_and_agent_names_the_missing_extra():
+ """The contract base is optional: without it the module loads and ``agent()`` says what to install."""
+ script = (
+ "import sys\n"
+ "for m in ('airflow.providers.common.ai.exceptions', 'airflow.providers.common.ai.managed_agents.base'):\n"
+ " sys.modules[m] = None\n"
+ "from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentCoreHook\n"
+ "from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException\n"
+ "hook = BedrockAgentCoreHook(aws_conn_id='aws_default', region_name='us-east-1')\n"
+ "try:\n"
+ " hook.agent('arn:aws:bedrock-agentcore:us-east-1:1:runtime/x')\n"
+ "except AirflowOptionalProviderFeatureException as e:\n"
+ " assert 'amazon[common.ai]' in str(e), e\n"
+ "else:\n"
+ " raise SystemExit('agent() did not raise')\n"
+ )
+ subprocess.run([sys.executable, "-c", script], check=True)
diff --git a/providers/common/ai/docs/toolsets/index.rst b/providers/common/ai/docs/toolsets/index.rst
index fa9db38d70a..3d61337e8a3 100644
--- a/providers/common/ai/docs/toolsets/index.rst
+++ b/providers/common/ai/docs/toolsets/index.rst
@@ -79,7 +79,7 @@ Start with what you have
* - Work that means running code the model wrote, not calling a tool you chose
- ``SandboxToolset``
* - Reasoning that should happen on the vendor's own infrastructure
- - A subclass of ``BaseManagedAgentToolset`` that you write
+ - ``ManagedAgentToolset`` over a vendor hook that implements the managed-agent contract
The hook, SQL, object storage, DataFusion, MCP, Agent Skills and managed-agent guides each have a
*When to choose it* section giving the case for choosing it, what it cannot do, an
@@ -110,7 +110,7 @@ Airflow's 350+ provider hooks already have typed methods, rich docstrings,
and managed credentials. Toolsets expose them as pydantic-ai tools so that
LLM agents can call them during multi-turn reasoning.
-Seven 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
@@ -126,13 +126,13 @@ package root:
- :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`: give
the agent a shell and a filesystem inside an isolated sandbox, off the
Airflow worker. Guide: :doc:`../sandbox/index`.
-- :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.
- Guide: :doc:`managed_agent`.
-- :class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset`:
- composes several interchangeable managed agents behind a single tool.
- See :ref:`managed-agent-toolsets`.
+- :class:`~airflow.providers.common.ai.toolsets.managed_agent.ManagedAgentToolset`:
+ expose a **vendor-managed agent**, one whose reasoning loop runs on a cloud
+ provider's infrastructure, as a single tool. Takes any
+ :class:`~airflow.providers.common.ai.managed_agents.base.ManagedAgentClient`,
+ which the Amazon and Google providers produce from their hooks and which
+ :class:`~airflow.providers.common.ai.managed_agents.failover.FailoverManagedAgentClient`
+ composes across clouds. Guide: :doc:`managed_agent`.
Three more toolsets (:doc:`datafusion`, :doc:`logging`, :doc:`skills`) are not
re-exported from the package root, so import each of them from its own submodule::
@@ -191,16 +191,17 @@ making on purpose rather than inheriting.
for ``sbx``; ambient Modal token on the worker for Modal
- A microVM on the worker host (``sbx``), or Modal's infrastructure off
the worker (Modal)
- * - ``BaseManagedAgentToolset``
- - Undefined by the base class; the subclass decides
+ * - ``ManagedAgentToolset``
+ - The vendor hook's own connection
- The vendor's infrastructure
Two rows are worth pausing on. ``SandboxToolset`` deliberately takes no Airflow
credential; that is the whole point of it, and it substitutes a backend-level
-boundary, on the host or at the vendor, for the connection-level one. ``BaseManagedAgentToolset`` does not
-substitute anything; it simply leaves the question to whoever writes the
-subclass. Neither is a defect, but in both cases the access decision has moved
-somewhere Airflow cannot see it, and somebody has to make that decision again in
+boundary, on the host or at the vendor, for the connection-level one. ``ManagedAgentToolset``
+authenticates through the vendor hook's connection, but what the remote agent may
+touch is governed on the vendor's side. Neither is a defect, but in both cases the
+access decision has moved somewhere Airflow cannot see it, and somebody has to make
+that decision again in
the new place. The :ref:`defense-layer table <toolset-defense-layers>` is the
right companion when you do.
@@ -215,7 +216,7 @@ each build their own tool definitions and set ``sequential=True`` on them, which
pydantic-ai treats as a barrier: the tool runs alone, tools the model emitted
before it finish first, and tools emitted after it start only once it returns.
A slow call on any of those four therefore holds up the rest of that step, not
-just its own toolset. ``BaseManagedAgentToolset`` sets ``sequential=False``
+just its own toolset. ``ManagedAgentToolset`` sets ``sequential=False``
deliberately, because the wait it introduces is remote. ``MCPToolset`` and
``AgentSkillsToolset`` define no tools of their own: they pass through whatever
the upstream toolset declares, so the setting is not theirs to make. Do not read
@@ -322,7 +323,7 @@ This provider does not treat either as the default. The rule of thumb in
:doc:`../index` is the dividing line: if Airflow should *run* the AI step, and the
model should stay swappable, use ``common.ai``; if the Dag *submits work to* a
vendor-managed service and waits for the result, use that vendor's provider,
-and ``BaseManagedAgentToolset`` exists for the case where you want the second
+and ``ManagedAgentToolset`` exists for the case where you want the second
behaviour from inside an agent that is otherwise doing the first.
See also
diff --git a/providers/common/ai/docs/toolsets/managed_agent.rst b/providers/common/ai/docs/toolsets/managed_agent.rst
index 4ab923fa5f4..bc5d6464610 100644
--- a/providers/common/ai/docs/toolsets/managed_agent.rst
+++ b/providers/common/ai/docs/toolsets/managed_agent.rst
@@ -17,8 +17,8 @@
.. _managed-agent-toolsets:
-Vendor-managed agents: ``BaseManagedAgentToolset``
-==================================================
+Vendor-managed agents: ``ManagedAgentToolset``
+===============================================
.. note::
@@ -31,24 +31,45 @@ 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. A provider package would
-ship its own subclass, so credentials keep flowing through that provider's
-existing hook and no new connection types are needed.
+Two pieces make that consultation vendor-neutral. The **contract**, in
+:mod:`airflow.providers.common.ai.managed_agents`, is what a vendor hook
+implements: a request goes in, an answer comes out, and the agent is an argument
+rather than a hook of its own, the way a statement is an argument to
+``DbApiHook.run``. The **toolset**,
+:class:`~airflow.providers.common.ai.toolsets.managed_agent.ManagedAgentToolset`,
+is what a Dag passes to ``AgentOperator``: it presents any client of that
+contract to the calling model as one tool with a bare-prompt schema.
+
+.. code-block:: python
-A subclass implements two members:
+ from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentCoreHook
+ from airflow.providers.common.ai.operators.agent import AgentOperator
+ from airflow.providers.common.ai.toolsets import ManagedAgentToolset
-``agent_ref``
- Normalized identity of the remote agent (``platform`` and ``name``), logged
- on every call so a run can be audited for which agents it consulted.
+ claims = BedrockAgentCoreHook(aws_conn_id="aws_prod", region_name="us-east-1").agent(
+ "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/claims"
+ )
-``invoke_sync(prompt)``
- Send the prompt, return the agent's answer. Return the *answer*, not the
- transport envelope.
+ AgentOperator(
+ task_id="triage",
+ llm_conn_id="anthropic_default",
+ prompt="Review claim 4411 and decide whether to pay it.",
+ toolsets=[
+ ManagedAgentToolset(
+ claims,
+ tool_name="ask_claims_agent",
+ description="Reviews an insurance claim and returns a coverage determination.",
+ )
+ ],
+ )
-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.
+Credentials keep flowing through the vendor's own hook and connection; no new
+connection types are involved. The vendor hooks that adopt the contract today are
+:class:`~airflow.providers.amazon.aws.hooks.bedrock.BedrockAgentCoreHook`
+(install ``apache-airflow-providers-amazon[common.ai]``) and
+:class:`~airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.AgentEngineHook`
+(install ``apache-airflow-providers-google[common.ai]``). Each provider documents
+how it maps the contract onto its service.
``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
@@ -64,46 +85,70 @@ derives one from a method name when there is no docstring.
name and the description are the whole of what the model knows about the
agent.
-Sync or async?
---------------
+What the model receives
+-----------------------
-An agent run is driven by a single event loop, shared by the model requests and
-every toolset's tool calls. A call that blocks that loop stalls all of them for
-as long as the remote agent reasons -- seconds to minutes, not milliseconds --
-and it does so silently, with no error to point at.
+The model receives ``ManagedAgentResponse.text``: the answer, unwrapped from the
+vendor's transport envelope by the hook. The envelope itself is on
+``ManagedAgentResponse.raw`` for Python callers of the client and never reaches
+the model, so a vendor's citations, tool traces or metadata are neither lost nor
+pasted into a model's context. A Python caller who wants them calls the client
+directly::
-Implement ``invoke_sync`` when the vendor call is blocking. The base class runs
-it in a worker thread, keeping the loop free::
+ response = claims.invoke(ManagedAgentRequest(prompt="Summarize claim 4411"))
+ response.text # what a model would have seen
+ response.raw # the vendor envelope
- class MyManagedAgentToolset(BaseManagedAgentToolset):
- @property
- def agent_ref(self) -> dict[str, str]:
- return {"platform": "example.cloud", "name": self._agent_id}
+``ManagedAgentRequest.vendor_options`` carries anything the contract does not
+type through to the vendor call; ``ManagedAgentToolset(vendor_options=...)``
+sends the same options on every request, for per-agent settings such as Agent
+Engine's ``class_method`` or AgentCore's ``text_key``. Hooks reject options that would re-target the call,
+such as another agent identity, account or connection, and refuse a
+``session_id`` when the agent keeps no conversation state.
- def invoke_sync(self, prompt: str) -> str:
- return self._hook.conn.invoke(agent=self._agent_id, prompt=prompt)
+Error handling
+--------------
-Override ``invoke(prompt)`` -- the async coroutine -- when the call is already
-asynchronous. Implement one hook or the other; a subclass that implements
-neither is rejected when it is constructed.
+Failures sort into three buckets, and conflating them is the most common way a
+vendor adoption goes wrong:
-.. warning::
+.. list-table::
+ :header-rows: 1
+ :widths: 22 30 48
- A thread cannot be cancelled. A caller that stops waiting for
- ``invoke_sync`` does not stop the call -- it keeps a worker thread and its
- socket until the call returns. Set a timeout on the underlying request so
- that is bounded.
+ * - Raise
+ - When
+ - Who recovers
+ * - :class:`~airflow.providers.common.ai.exceptions.ManagedAgentRejected`
+ - The agent rejected the request in a way rephrasing could fix.
+ - The calling model. The toolset turns it into a pydantic-ai ``ModelRetry``,
+ bounded by ``max_retries`` (one rephrase by default; ``0`` makes it fatal).
+ * - :class:`~airflow.providers.common.ai.exceptions.ManagedAgentInvocationError`
+ - Terminal: bad credentials, missing agent, malformed request, exhausted quota.
+ - No rephrase; a failover group moves to its next member. Airflow's own task
+ retries still apply, so pair the task with a retry rule that stops on it if
+ retrying would only repeat the failure.
+ * - *let it propagate*
+ - Transient: 429, 5xx, connection reset, read timeout.
+ - Airflow's task-level retry. A rephrase does nothing for a 503.
+
+A vendor hook raises :class:`~airflow.providers.common.ai.exceptions.ManagedAgentRejected`
+rather than ``ModelRetry`` so that the contract module imports nothing from
+pydantic-ai; the toolset is the one place a rejection becomes something the
+calling model can act on. Neither platform adopted here has a rephrase-class
+error. AgentCore reports a container's own complaints inside a successful body,
+and Agent Engine's ``INVALID_ARGUMENT`` means an author-side mistake such as a
+wrong input key, so both hooks raise terminal errors or let transient ones
+propagate.
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:
+A toolset call runs in the worker and cannot defer to the Triggerer: it blocks
+for the duration of the call, in a worker thread so the agent's event loop keeps
+running. Some managed-agent platforms are built around a long-running job or a
+session rather than one request and one answer, and a toolset serves those
+poorly.
.. list-table::
:header-rows: 1
@@ -113,109 +158,76 @@ That draws a boundary worth respecting:
- 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
+ - :class:`~airflow.providers.common.ai.toolsets.managed_agent.ManagedAgentToolset`
* - 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`
+ :class:`~airflow.sdk.bases.resumablejobmixin.ResumableJobMixin`; for
+ Agent Engine query jobs, ``RunQueryJobOperator``
``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.
+the retry boundary is the task, not the tool call. Set ``timeout`` on the toolset
+so a call that stops being waited for is bounded: a thread cannot be cancelled,
+and the hooks apply the request's timeout to the vendor call itself.
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.
+``replayable`` is ``False`` by default, and ``AgentOperator(durable=True)``
+honors it: a managed-agent call is re-invoked on retry rather than served from
+the step cache, because the agent may have acted on systems Airflow cannot
+observe and replaying a cached answer could skip a side effect. Set
+``replayable=True`` only for an agent that is read-only.
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, so the policy
-stays deterministic Python rather than a prompt instruction a model may ignore.
-Groups nest.
+:class:`~airflow.providers.common.ai.managed_agents.failover.FailoverManagedAgentClient`
+composes several clients into one, trying them in order until one answers. It is
+itself a client, so a toolset over it presents a single tool and the calling
+model has no say in which provider serves the request, so 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
+ from airflow.providers.common.ai.managed_agents import FailoverManagedAgentClient
+ from airflow.providers.common.ai.toolsets import ManagedAgentToolset
- resilient = FailoverManagedAgentToolset(
+ resilient = ManagedAgentToolset(
+ FailoverManagedAgentClient([bedrock_claims, vertex_claims]), # same agent, two clouds
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.
+Members must be **substitutable**: the same agent deployed twice, not two
+specialists with different data. Two containerized agents built from one image
+qualify; an agent bound to one platform's own objects (a Cortex Agent over
+Snowflake semantic models) does not, because there is nothing equivalent to fail
+over *to*. The group cannot check this.
-**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.
+What it can check is **conversation state**. A failover starts a fresh conversation
+on the standby, which is correct for a one-shot consultation and wrong for a
+multi-turn one. The group therefore never reports ``capabilities.sessions`` and
+refuses a request that carries a ``session_id``, whatever its members support;
+a conversation belongs to one member, addressed directly.
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.
+- :class:`~airflow.providers.common.ai.exceptions.ManagedAgentInvocationError`
+ and transient failures move to the next member.
+- :class:`~airflow.providers.common.ai.exceptions.ManagedAgentRejected` is
+ re-raised immediately and never triggers failover. A prompt the primary could
+ not parse will not parse on the standby either.
- 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
+the cloud SDKs' exception trees: ``requests``, ``botocore`` and the Google SDK
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.
-
.. note::
For a **standalone** agent call, prefer plain Airflow task-level failover:
@@ -237,20 +249,47 @@ healthy one:
* - Metric
- Tags
* - ``managed_agent.failover``
- - ``from_platform``, ``to_platform`` (one per failover transition)
+ - ``from_platform``, ``to_platform``: one per failover transition, emitted by
+ the group
* - ``managed_agent.served``
- - ``platform`` and ``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
+ - ``tool``, ``platform``: one per answer, emitted by the toolset; a toolset
+ over a group reports ``platform=failover``, and either counter reports
+ ``unknown`` when an agent's identity could not be resolved
+
+``managed_agent.failover`` rising while ``managed_agent.served`` stays flat means
+the primary is down, and that is a dashboard question rather than a log grep. The
+counters do not say which member answered: that is in the task log, where each
+failover warning names the members involved and their positions, but not in XCom.
+Both counters 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.
+Implementing the contract for a new vendor
+------------------------------------------
+
+A vendor hook adopts
+:class:`~airflow.providers.common.ai.managed_agents.base.BaseManagedAgentHook`
+as a mixin beside its own base and implements three methods: ``resolve_agent``
+(normalize the agent identifier into a platform-qualified reference, without a
+network call), ``get_agent_capabilities`` (what the pair can do, so consumers can
+refuse rather than degrade) and ``invoke_agent`` (send a request, return an
+answer, sort failures into the three buckets). ``hook.agent(...)`` then returns a
+bound client the toolset accepts.
+
+Because ``common.ai`` requires Airflow 3 and most vendor providers still support
+Airflow 2, the provider declares ``common.ai`` as an optional extra and the hook
+module imports the contract under a guard: when the extra is missing, a stand-in
+base keeps the hook importable and ``hook.agent(...)`` raises
+``AirflowOptionalProviderFeatureException`` naming the extra to install. The Amazon
+provider already does this for optional base classes such as
+``ResumableJobMixin``; the Amazon and Google adoptions in this release are the
+templates.
+
+The only toolset most code needs is ``ManagedAgentToolset``. Subclass
+:class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset`
+directly only for an agent that has no hook at all; it takes an ``agent_ref`` and
+an ``invoke_sync`` (or an async ``invoke``) and supplies the same tool surface.
+A thread cannot be cancelled, so an ``invoke_sync`` must set a timeout on its own
+request.
When to choose it
-----------------
@@ -258,39 +297,31 @@ When to choose it
**Choose it when** the reasoning itself belongs on the vendor's infrastructure:
the agent is already deployed there, grounded in data that never leaves, and
Airflow's job is to submit one request and read one answer.
-:ref:`managed-agent-toolsets` covers the shape.
-
-Read the first bullet before planning around this route.
**What it cannot do**
-- It is an extension point, not a toolset you can use as-is. This provider ships
- the base class only:
- :class:`~airflow.providers.common.ai.toolsets.managed_agent.BaseManagedAgentToolset`
- declares ``agent_ref`` abstract and raises ``TypeError`` at construction unless
- a subclass implements ``invoke_sync`` or ``invoke``. No vendor subclass exists
- in this repository: the Snowflake Cortex, Bedrock AgentCore, Azure AI Foundry
- and Vertex AI Agent Engine names in its docstring describe the shape it expects,
- not implementations that ship. Using this route means writing that subclass.
- It has no allow-list to offer. The whole toolset is one tool: a prompt goes in,
an answer comes out. Whatever governs what the remote agent may touch lives on
the vendor's side, which is the trade you are making.
-- It does not define where the credential comes from. The base class leaves
- that decision to the subclass; :ref:`managed-agent-toolsets` frames the
- intended shape as routing authentication through the provider's own hook.
+- It cannot defer. The call blocks a worker thread for as long as the remote agent
+ reasons. Long-running submitted work belongs in the vendor's own operator.
- Durable replay is off by default. ``replayable`` is ``False`` because a managed
agent may act on systems Airflow cannot observe, so replaying from the cache
could skip a side effect. Read-only agents can opt in.
-- Its failover variant relies on two preconditions the code cannot check.
- :class:`~airflow.providers.common.ai.toolsets.managed_agent.FailoverManagedAgentToolset`
- requires members that are genuinely interchangeable (the same agent deployed
- twice, not two specialists over different data) and one-shot exchanges, because
- each member is invoked with a bare prompt and no thread reference, so a failover
- silently starts a fresh conversation rather than resuming the old one.
-
-**A real example.** There is none. No example Dag, no system test, and the only
-code in the documentation is an illustrative subclass on this page. Budget
-for writing and testing the subclass yourself.
-
-**Credentials and where it runs.** Both are the subclass's decision. Reasoning
-runs on the vendor's infrastructure; the worker sends a request and waits.
+- Its failover variant,
+ :class:`~airflow.providers.common.ai.managed_agents.failover.FailoverManagedAgentClient`,
+ requires members that are genuinely interchangeable: the same agent deployed
+ twice, not two specialists over different data. It refuses a request that
+ carries a ``session_id``, because a failover starts a fresh conversation on the
+ standby.
+
+**Which vendors.** Amazon Bedrock AgentCore through
+:class:`~airflow.providers.amazon.aws.hooks.bedrock.BedrockAgentCoreHook`
+and Vertex AI Agent Engine through
+:class:`~airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.AgentEngineHook`,
+each behind that provider's ``common.ai`` extra. Other vendors adopt the same
+contract; `Implementing the contract for a new vendor`_ describes how.
+
+**Credentials and where it runs.** Credentials flow through the vendor hook's own
+connection. Reasoning runs on the vendor's infrastructure; the worker sends a
+request and waits.
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/caching_toolset.py b/providers/common/ai/src/airflow/providers/common/ai/durable/caching_toolset.py
index 83a5598ae0f..fa4a61b9af7 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/caching_toolset.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/caching_toolset.py
@@ -75,9 +75,20 @@ class CachingToolset(WrapperToolset[Any]):
# Grab step index BEFORE any await -- ensures deterministic ordering
# even when multiple tool calls run concurrently via asyncio.gather.
step = self.counter.next_step()
+
+ # The toolset a tool came from may declare that a completed call must not be served from
+ # cache, because the call acted on a system Airflow cannot observe (a managed agent, for
+ # instance). ``tool.toolset`` survives every pydantic-ai wrapper, so the check is per tool
+ # and one such toolset inside a combined one does not stop its siblings from replaying.
+ # The step still counts so later steps keep their keys.
+ if not getattr(_innermost(tool.toolset), "replayable", True):
+ log.debug("Durable: toolset is not replayable; running the tool", step=step, tool=name)
+ if self.replay_usage is not None:
+ self.replay_usage.record_live_tool_call(step)
+ return await self.wrapped.call_tool(name, tool_args, ctx, tool)
+
key = build_tool_step_key(step)
fingerprint = fingerprint_tool_call(name, tool_args, ctx.tool_call_id)
-
found, cached, cached_fingerprint = self.storage.load_tool_result(key)
if found:
if cached_fingerprint == fingerprint:
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 9b1a3e1d929..2f9c61f9f4b 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
@@ -85,8 +85,19 @@ class ManagedAgentInvocationError(RuntimeError):
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.
+ task-level retry handles them, and requests the agent rejected in a way the
+ calling model could fix by rephrasing should raise
+ :class:`ManagedAgentRejected` instead.
+ """
+
+
+class ManagedAgentRejected(Exception):
+ """
+ Raised when a managed agent rejected a request in a way rephrasing could fix.
+
+ Vendor hooks raise this instead of pydantic-ai's ``ModelRetry`` so that the contract
+ module stays free of pydantic-ai; the managed-agent toolset translates it to
+ ``ModelRetry`` at the boundary, and the calling model sees the message and tries again.
"""
diff --git a/providers/common/ai/src/airflow/providers/common/ai/managed_agents/__init__.py b/providers/common/ai/src/airflow/providers/common/ai/managed_agents/__init__.py
new file mode 100644
index 00000000000..90b0e8143a2
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/managed_agents/__init__.py
@@ -0,0 +1,49 @@
+# 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.
+"""
+The vendor-neutral contract for managed agents, and the consumers written once over it.
+
+Vendor providers import :mod:`~airflow.providers.common.ai.managed_agents.base`; it pulls
+in nothing from pydantic-ai. The toolset that presents a managed agent to a calling model lives
+in :mod:`airflow.providers.common.ai.toolsets.managed_agent`.
+"""
+
+from __future__ import annotations
+
+from airflow.providers.common.ai.managed_agents.base import (
+ BaseManagedAgentHook,
+ BoundManagedAgent,
+ ManagedAgentCapabilities,
+ ManagedAgentClient,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ ManagedAgentResponse,
+ ManagedAgentUsage,
+)
+from airflow.providers.common.ai.managed_agents.failover import FailoverManagedAgentClient
+
+__all__ = [
+ "BaseManagedAgentHook",
+ "BoundManagedAgent",
+ "FailoverManagedAgentClient",
+ "ManagedAgentCapabilities",
+ "ManagedAgentClient",
+ "ManagedAgentRef",
+ "ManagedAgentRequest",
+ "ManagedAgentResponse",
+ "ManagedAgentUsage",
+]
diff --git a/providers/common/ai/src/airflow/providers/common/ai/managed_agents/base.py b/providers/common/ai/src/airflow/providers/common/ai/managed_agents/base.py
new file mode 100644
index 00000000000..d7712cc1746
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/managed_agents/base.py
@@ -0,0 +1,251 @@
+# 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.
+"""
+The contract a provider hook implements to expose a vendor-managed agent.
+
+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 a request and reads an answer. This module defines the shape of that
+exchange once, so that every consumer in ``common.ai`` (the toolset, the failover group)
+is written against one interface rather than one per cloud.
+
+The design follows ``DbApiHook`` in ``common.sql``: a small base mixed into each vendor's
+own hook, with the agent as an *argument* to every method, because a hook is scoped to a
+connection and one connection reaches many agents. Vendor providers adopt it the way they
+adopt ``BaseMessageQueueProvider`` from ``common.messaging``: behind an optional extra, in a
+module whose import of this contract is guarded, so a provider that floors Airflow 2 never
+has to raise its floor.
+
+This module imports nothing from pydantic-ai on purpose. A vendor hook's guarded import of
+the contract must stay cheap.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from functools import cached_property
+from typing import Any, ClassVar, Protocol, runtime_checkable
+
+
+@dataclass(frozen=True)
+class ManagedAgentRef:
+ """
+ Normalized identity of a remote agent.
+
+ :param platform: A stable, dotted platform id such as ``aws.bedrock_agentcore`` or
+ ``gcp.vertex_agent_engine``. Used as a metric tag, so keep the set small.
+ :param name: The vendor's canonical identifier for the agent: an ARN, a full resource
+ name, ``DATABASE.SCHEMA.NAME``.
+ :param version: The resolved version or revision when the platform exposes one. Recorded
+ so a behaviour change can be attributed to a deployment rather than to Airflow.
+ """
+
+ platform: str
+ name: str
+ version: str | None = None
+
+
+@dataclass(frozen=True)
+class ManagedAgentCapabilities:
+ """
+ What a ``(hook, agent)`` pair can do.
+
+ Consumers check these and refuse rather than degrade: a bound agent rejects a request that
+ carries a ``session_id`` when ``sessions`` is False, and a failover group never offers
+ sessions at all, because failing over discards the conversation the primary was holding.
+ """
+
+ sessions: bool = False
+ """Whether ``ManagedAgentRequest.session_id`` continues a conversation."""
+ structured_output: bool = False
+ """Whether ``ManagedAgentResponse.structured`` can carry a typed value."""
+ usage: bool = False
+ """Whether ``ManagedAgentResponse.usage`` is populated."""
+ trace: bool = False
+ """Whether ``ManagedAgentResponse.trace_ref`` is populated."""
+
+
+@dataclass(frozen=True)
+class ManagedAgentRequest:
+ """
+ One request to a managed agent.
+
+ Exactly one of ``prompt`` and ``messages`` must be set. Everything the contract does not
+ type travels in ``vendor_options``. A hook may read some of them itself and passes the rest
+ through to the vendor call.
+ Hooks reject options that would re-target the call (the agent identity, the connection),
+ because a model-facing caller must not be able to change what it is talking to.
+ """
+
+ prompt: str | None = None
+ messages: Sequence[dict[str, Any]] | None = None
+ session_id: str | None = None
+ timeout: float | None = None
+ """Seconds to wait for the vendor call. The hook must enforce it on the request itself."""
+ vendor_options: dict[str, Any] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ if (self.prompt is None) == (self.messages is None):
+ raise ValueError("ManagedAgentRequest needs exactly one of prompt or messages.")
+
+ def as_messages(self) -> list[dict[str, Any]]:
+ """Return the request as a message list, for vendors that only accept messages."""
+ if self.messages is not None:
+ return list(self.messages)
+ return [{"role": "user", "content": [{"type": "text", "text": self.prompt}]}]
+
+
+@dataclass(frozen=True)
+class ManagedAgentUsage:
+ """Usage the vendor reported for one invocation. Every field is optional because vendors differ."""
+
+ input_tokens: int | None = None
+ output_tokens: int | None = None
+
+
+@dataclass(frozen=True)
+class ManagedAgentResponse:
+ """
+ One answer from a managed agent.
+
+ ``text`` is what a calling model should read; the hook unwraps the vendor envelope to
+ produce it. ``raw`` is that envelope, always populated and never handed to a model, so a
+ Python caller loses nothing.
+ """
+
+ text: str
+ raw: Any
+ structured: Any | None = None
+ session_id: str | None = None
+ usage: ManagedAgentUsage | None = None
+ trace_ref: str | None = None
+ """A vendor request, invocation or trace id, for joining Airflow's record to the vendor's."""
+
+
+class BaseManagedAgentHook(ABC):
+ """
+ Mixin a vendor hook adopts to expose its managed agents through the common contract.
+
+ Mixed in beside the vendor's own base and never replacing it::
+
+ class BedrockAgentCoreHook(AwsBaseHook, BaseManagedAgentHook): ...
+
+ It therefore has no ``__init__`` and makes no assumption about ``get_conn``. The agent is
+ an argument to every method, the way a statement is an argument to ``DbApiHook.run``.
+
+ Method names are chosen to collide with nothing on the shipped vendor hooks. That matters
+ more than it looks: ``SnowflakeCortexAgentHook`` already defines ``run_agent`` and
+ ``describe_agent``, and an abstract method that a vendor base happens to define is
+ silently satisfied with the wrong signature.
+
+ Implementations sort failures into three classes, and conflating them is the most common
+ way an adoption goes wrong:
+
+ * :class:`~airflow.providers.common.ai.exceptions.ManagedAgentRejected` -- the agent
+ rejected the request in a way rephrasing could fix. The toolset turns it into a
+ pydantic-ai ``ModelRetry`` so the calling model tries again.
+ * :class:`~airflow.providers.common.ai.exceptions.ManagedAgentInvocationError` --
+ terminal: bad credentials, missing agent, revoked quota. Nothing on the agent side
+ recovers it; whether the task retries is the task's retry policy.
+ * Anything transient (429, 5xx, connection reset, read timeout) -- propagate unchanged.
+ Airflow's task-level retry is the right layer; a rephrase does nothing for a 503.
+ """
+
+ agent_platform: ClassVar[str]
+ """The ``platform`` every :class:`ManagedAgentRef` from this hook carries."""
+
+ @abstractmethod
+ def resolve_agent(self, agent: str) -> ManagedAgentRef:
+ """Normalize ``agent`` into a platform-qualified reference. Must not make a network call."""
+
+ @abstractmethod
+ def get_agent_capabilities(self, agent: str) -> ManagedAgentCapabilities:
+ """Report what ``agent`` on this connection can do. Must not make a network call."""
+
+ @abstractmethod
+ def invoke_agent(self, agent: str, request: ManagedAgentRequest) -> ManagedAgentResponse:
+ """Send ``request`` to ``agent`` and return its answer. Blocking."""
+
+ def agent(self, agent: str) -> BoundManagedAgent:
+ """Bind one agent on this connection. This is what the ``common.ai`` toolsets consume."""
+ return BoundManagedAgent(hook=self, agent=agent)
+
+
+@runtime_checkable
+class ManagedAgentClient(Protocol):
+ """
+ What ``common.ai``'s consumers are typed against.
+
+ A :class:`BoundManagedAgent` satisfies it. So does
+ :class:`~airflow.providers.common.ai.managed_agents.failover.FailoverManagedAgentClient`,
+ and so can anything that needs no Airflow connection at all.
+ """
+
+ @property
+ def ref(self) -> ManagedAgentRef: ...
+
+ @property
+ def capabilities(self) -> ManagedAgentCapabilities: ...
+
+ def invoke(self, request: ManagedAgentRequest) -> ManagedAgentResponse: ...
+
+
+@dataclass(frozen=True)
+class BoundManagedAgent:
+ """
+ A ``(hook, agent)`` pair. Forwards to the hook and resolves identity lazily.
+
+ This is also where the contract's "refuse rather than degrade" rule is enforced for every
+ adopter: a request that asks for something the pair's capabilities do not include is
+ rejected before the hook is called.
+ """
+
+ hook: BaseManagedAgentHook
+ agent: str
+
+ # Resolved on first use, never at construction, so a bad agent id fails the first call rather
+ # than Dag parsing; cached after that because identity is immutable for a frozen binding.
+ @cached_property
+ def ref(self) -> ManagedAgentRef:
+ return self.hook.resolve_agent(self.agent)
+
+ @cached_property
+ def capabilities(self) -> ManagedAgentCapabilities:
+ return self.hook.get_agent_capabilities(self.agent)
+
+ def invoke(self, request: ManagedAgentRequest) -> ManagedAgentResponse:
+ if request.session_id is not None and not self.capabilities.sessions:
+ raise ValueError(
+ f"{self.agent} does not keep conversation state, so session_id cannot be honored. "
+ "Refusing rather than silently starting a fresh conversation."
+ )
+ return self.hook.invoke_agent(self.agent, request)
+
+
+def describe(client: ManagedAgentClient) -> str:
+ """
+ Render a client's identity for a log line without letting identity resolution fail the caller.
+
+ A standby whose connection is misconfigured must not fail a call the primary served.
+ """
+ try:
+ ref = client.ref
+ except Exception as exc:
+ return f"<unresolved: {exc}>"
+ return f"{ref.name} on {ref.platform}"
diff --git a/providers/common/ai/src/airflow/providers/common/ai/managed_agents/failover.py b/providers/common/ai/src/airflow/providers/common/ai/managed_agents/failover.py
new file mode 100644
index 00000000000..63e04e845b5
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/managed_agents/failover.py
@@ -0,0 +1,145 @@
+# 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 collections.abc import Sequence
+
+from airflow.providers.common.ai.exceptions import ManagedAgentRejected
+from airflow.providers.common.ai.managed_agents.base import (
+ ManagedAgentCapabilities,
+ ManagedAgentClient,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ ManagedAgentResponse,
+ describe,
+)
+from airflow.providers.common.compat.sdk import Stats
+
+log = logging.getLogger(__name__)
+
+
+class FailoverManagedAgentClient:
+ """
+ Active/passive failover across interchangeable managed agents.
+
+ Members are tried in order and the first answer wins. The group is itself a
+ :class:`~airflow.providers.common.ai.managed_agents.base.ManagedAgentClient`, so
+ a toolset built over it presents a single tool and the calling model 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.
+
+ Members must be *substitutable*: the same agent deployed twice, not two specialists over
+ different data. Two containerized agents built from one image qualify; an agent bound to
+ one platform's own objects does not, because there is nothing equivalent to fail over to.
+ The group cannot check this.
+
+ What it can check is conversation state. A failover starts a fresh conversation on the
+ standby, which is correct for a one-shot consultation and wrong for a multi-turn one, so
+ the group never reports ``capabilities.sessions`` and refuses a request that carries a
+ ``session_id``, whatever its members support.
+
+ :param members: Interchangeable clients, 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.
+ Narrow it when the members' exception types are known.
+ :class:`~airflow.providers.common.ai.exceptions.ManagedAgentRejected` never
+ triggers failover, whatever this is set to: the standby would reject the same prompt.
+ """
+
+ def __init__(
+ self,
+ members: Sequence[ManagedAgentClient],
+ *,
+ failover_on: tuple[type[BaseException], ...] = (Exception,),
+ ) -> None:
+ if len(members) < 2:
+ raise ValueError(
+ f"A failover group needs at least two members; got {len(members)}. "
+ "Use the member directly instead."
+ )
+ # Copied, not aliased: a caller holding the original list could otherwise empty it.
+ self._members = tuple(members)
+ self._failover_on = failover_on
+
+ @property
+ def members(self) -> tuple[ManagedAgentClient, ...]:
+ return self._members
+
+ @property
+ def ref(self) -> ManagedAgentRef:
+ return ManagedAgentRef(platform="failover", name=" -> ".join(describe(m) for m in self._members))
+
+ @property
+ def capabilities(self) -> ManagedAgentCapabilities:
+ """The intersection of the members', except sessions, which a group never offers."""
+ caps = [m.capabilities for m in self._members]
+ return ManagedAgentCapabilities(
+ sessions=False,
+ structured_output=all(c.structured_output for c in caps),
+ usage=all(c.usage for c in caps),
+ trace=all(c.trace for c in caps),
+ )
+
+ def invoke(self, request: ManagedAgentRequest) -> ManagedAgentResponse:
+ if request.session_id is not None:
+ raise ValueError(
+ "A failover group cannot continue a conversation: the standby would start a fresh one. "
+ "Send session-bound requests to one member directly."
+ )
+ for position, member in enumerate(self._members[:-1]):
+ try:
+ response = member.invoke(request)
+ except ManagedAgentRejected:
+ raise
+ except self._failover_on:
+ standby = self._members[position + 1]
+ log.warning(
+ "Managed agent %s (member %d) failed; failing over to %s (member %d)",
+ describe(member),
+ position,
+ describe(standby),
+ position + 1,
+ exc_info=True,
+ )
+ # A failover is a success-shaped event: 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": _get_platform(member), "to_platform": _get_platform(standby)},
+ )
+ continue
+ return self._served(position, member, response)
+ # The last member gets no failover: whatever it raises is the group's answer.
+ last = len(self._members) - 1
+ return self._served(last, self._members[last], self._members[last].invoke(request))
+
+ @staticmethod
+ def _served(
+ position: int, member: ManagedAgentClient, response: ManagedAgentResponse
+ ) -> ManagedAgentResponse:
+ if position > 0:
+ log.info("Managed agent request served by standby %s (member %d)", describe(member), position)
+ return response
+
+
+def _get_platform(client: ManagedAgentClient) -> str:
+ try:
+ return client.ref.platform
+ except Exception:
+ return "unknown"
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 df797bfe623..771fd0c5a84 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
@@ -21,16 +21,16 @@ from __future__ import annotations
from airflow.providers.common.ai.toolsets.hook import HookToolset
from airflow.providers.common.ai.toolsets.managed_agent import (
BaseManagedAgentToolset,
- FailoverManagedAgentToolset,
+ ManagedAgentToolset,
)
from airflow.providers.common.ai.toolsets.object_storage import ObjectStorageToolset
from airflow.providers.common.ai.toolsets.sandbox import SandboxToolset
__all__ = [
"BaseManagedAgentToolset",
- "FailoverManagedAgentToolset",
"HookToolset",
"MCPToolset",
+ "ManagedAgentToolset",
"ObjectStorageToolset",
"SQLToolset",
"SandboxToolset",
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 d48729b0b92..dc11bf6007c 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
@@ -25,6 +25,12 @@ from pydantic_ai.exceptions import ModelRetry
from pydantic_ai.tools import ToolDefinition
from pydantic_ai.toolsets.abstract import ToolsetTool
+from airflow.providers.common.ai.exceptions import ManagedAgentRejected
+from airflow.providers.common.ai.managed_agents.base import (
+ ManagedAgentClient,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+)
from airflow.providers.common.ai.utils.tool_definition import (
build_args_validator,
return_schema_kwargs,
@@ -34,8 +40,6 @@ from airflow.providers.common.ai.utils.toolset_base import AirflowToolset, valid
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__)
@@ -50,6 +54,7 @@ _PROMPT_SCHEMA: dict[str, Any] = {
},
"required": ["prompt"],
}
+_PROMPT_VALIDATOR = build_args_validator(_PROMPT_SCHEMA)
class BaseManagedAgentToolset(AirflowToolset):
@@ -61,41 +66,40 @@ class BaseManagedAgentToolset(AirflowToolset):
Experimental: this can change or be removed in a minor release of this provider.
See :ref:`howto/stability`.
- 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, and the identifier it emits
- when calling the tool. A verb phrase naming the specialist reads best,
- e.g. ``ask_bookings_analyst``.
- :param description: What this agent knows and when to consult it. Optional --
- it falls back to ``tool_name`` rendered as prose, matching how
- ``HookToolset`` handles a method with no docstring. Worth writing anyway:
- it is what tells the model to consult the agent rather than answer from
- its own knowledge, and it is the only place to state a scope limit the
- name cannot carry ("cannot see revenue figures"). Since the argument
- schema is always a bare prompt, the name and this string are the whole
- of what the model knows about the agent.
- :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.
- :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.
+ A managed agent runs its own reasoning loop on the vendor's infrastructure. 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.
+
+ Most code should not subclass this. Use :class:`ManagedAgentToolset` over a
+ :class:`~airflow.providers.common.ai.managed_agents.base.ManagedAgentClient`, which
+ every provider hook that adopts
+ :class:`~airflow.providers.common.ai.managed_agents.base.BaseManagedAgentHook`
+ produces via ``hook.agent(...)``. Subclass this directly only for an agent that has no
+ hook at all. Subclasses implement :attr:`agent_ref` and :meth:`invoke_sync` (or override
+ the async :meth:`invoke`); tool naming, argument validation, result serialization,
+ logging and metrics are handled here so every implementation presents the same surface
+ to the model.
+
+ :param tool_name: Name the calling model sees, and the identifier it emits when calling
+ the tool. A verb phrase naming the specialist reads best, e.g. ``ask_bookings_analyst``.
+ :param description: What this agent knows and when to consult it. Optional -- it falls
+ back to ``tool_name`` rendered as prose, matching how ``HookToolset`` handles a method
+ with no docstring. Worth writing anyway: it is what tells the model to consult the
+ agent rather than answer from its own knowledge, and it is the only place to state a
+ scope limit the name cannot carry ("cannot see revenue figures"). Since the argument
+ schema is always a bare prompt, the name and this string are the whole of what the
+ model knows about the agent.
+ :param timeout: Seconds to wait for a single invocation. ``None`` defers to the platform
+ default. Exposed as :attr:`timeout` so an implementation can honor it.
+ :param max_retries: How many times the calling model may rephrase after the remote agent
+ rejects a request. ``0`` turns the first rejection into a hard error.
"""
- #: 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.
+ #: Whether ``durable=True`` may replay a completed invocation from its cache instead of
+ #: re-invoking. 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 may
+ #: opt in.
replayable: bool = False
def __init__(
@@ -124,54 +128,41 @@ class BaseManagedAgentToolset(AirflowToolset):
self._timeout = timeout
self._max_retries = max_retries
+ @property
+ def timeout(self) -> float | None:
+ """Seconds to wait for one invocation, or ``None`` for the platform default."""
+ return self._timeout
+
@property
@abstractmethod
- def agent_ref(self) -> dict[str, str]:
+ def agent_ref(self) -> ManagedAgentRef:
"""
- Normalised identity of the remote agent.
+ Normalized identity of the remote agent.
- Must contain ``platform`` and ``name``, e.g.
- ``{"platform": "snowflake.cortex", "name": "ANALYTICS.REVENUE.BOOKINGS_ANALYST"}``.
- 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.
+ Logged after every successful call, so the resolved remote identity behind a task
+ appears in that task's log even though the Dag only names a connection. Resolution
+ is never on the call's critical path: a failure here is logged, not raised.
"""
async def invoke(self, prompt: str) -> Any:
"""
Send ``prompt`` to the remote agent and return the agent's answer.
- Override this when the vendor call is already asynchronous. When it
- blocks, implement :meth:`invoke_sync` instead and let the default
- implementation here run it in a worker thread, which keeps it off the
- event loop that the whole agent run shares.
-
- 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.
+ Override this when the vendor call is already asynchronous. When it blocks, implement
+ :meth:`invoke_sync` instead and let the default implementation here run it in a worker
+ thread, which keeps it off the event loop that the whole agent run shares.
+
+ Return the answer, not the transport envelope. Failures sort into three classes:
+ ``ModelRetry`` (the model can fix it by rephrasing; a hook-backed client raises
+ :class:`~airflow.providers.common.ai.exceptions.ManagedAgentRejected` instead and
+ :class:`ManagedAgentToolset` translates it),
+ :class:`~airflow.providers.common.ai.exceptions.ManagedAgentInvocationError`
+ (terminal), and anything transient, which should propagate unchanged so Airflow's
+ task-level retry handles it.
+
+ **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``. A tool call has no post-task cleanup hook to fall back on.
:param prompt: The question or instruction to send to the remote agent.
"""
@@ -181,15 +172,8 @@ class BaseManagedAgentToolset(AirflowToolset):
"""
Blocking variant of :meth:`invoke`, run in a worker thread.
- This is the hook to implement when the vendor call blocks. Make it
- normally -- the base class keeps it off the event loop, so a call that
- takes minutes does not stall the calling agent's other tool calls.
-
- The contract is :meth:`invoke`'s: return the answer rather than the
- transport envelope, sort failures into the same three buckets, and
- release anything allocated on every path. A thread cannot be cancelled,
- so set a timeout on the underlying request: a caller that stops waiting
- does not stop this call.
+ A thread cannot be cancelled, so set a timeout on the underlying request: a caller
+ that stops waiting does not stop this call.
:param prompt: The question or instruction to send to the remote agent.
"""
@@ -226,7 +210,7 @@ class BaseManagedAgentToolset(AirflowToolset):
# becomes a hard error -- so raise it only when the remote agent's
# rejections are genuinely worth re-prompting.
max_retries=self._get_tool_max_retries(ctx),
- args_validator=build_args_validator(_PROMPT_SCHEMA),
+ args_validator=_PROMPT_VALIDATOR,
)
}
@@ -238,132 +222,120 @@ class BaseManagedAgentToolset(AirflowToolset):
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"])
+ # Identity is resolved after the call, never before it: a toolset whose identity comes
+ # from a misconfigured connection must not fail a call that would have succeeded, and
+ # a failover group's identity joins every member's, standbys included.
+ ref = self._resolve_agent_ref()
+ log.info(
+ "Consulted managed agent %s",
+ f"{ref.name} on {ref.platform}"
+ if ref is not None
+ else f"<unresolved identity> for tool {self._tool_name}",
+ )
+ # Emitted once per answer so managed-agent call volume is observable next to the
+ # ``managed_agent.failover`` counter. Tagged by platform to bound cardinality.
+ Stats.incr(
+ "managed_agent.served",
+ tags={"tool": self._tool_name, "platform": ref.platform if ref is not None else "unknown"},
+ )
return serialize_for_llm(result)
+ def _resolve_agent_ref(self) -> ManagedAgentRef | None:
+ """Resolve identity for a log line or a metric tag without letting resolution fail the call."""
+ try:
+ return self.agent_ref
+ except Exception:
+ log.warning("Managed agent identity could not be resolved", exc_info=True)
+ return None
-class FailoverManagedAgentToolset(BaseManagedAgentToolset):
+
+class ManagedAgentToolset(BaseManagedAgentToolset):
"""
- Present several interchangeable managed agents to the model as one tool.
+ Expose any managed-agent client as one tool.
.. note::
Experimental: this can change or be removed in a minor release of this provider.
See :ref:`howto/stability`.
- 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.
-
- :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. 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.
+ This is the toolset to use. It accepts any
+ :class:`~airflow.providers.common.ai.managed_agents.base.ManagedAgentClient`. The client is usually a
+ :class:`~airflow.providers.common.ai.managed_agents.base.BoundManagedAgent` from a
+ vendor hook's ``agent()`` method, or a
+ :class:`~airflow.providers.common.ai.managed_agents.failover.FailoverManagedAgentClient`
+ over several of them::
+
+ from airflow.providers.amazon.aws.hooks.bedrock import (
+ BedrockAgentCoreHook,
+ )
+ from airflow.providers.common.ai.toolsets import ManagedAgentToolset
+
+ claims = BedrockAgentCoreHook(aws_conn_id="aws_prod").agent(RUNTIME_ARN)
+ toolset = ManagedAgentToolset(
+ claims,
+ tool_name="ask_claims_agent",
+ description="Reviews an insurance claim and returns a coverage determination.",
+ )
+
+ The model receives ``response.text``. The vendor envelope in ``response.raw`` is for
+ Python callers of the client and never reaches the model.
+
+ :param client: The agent to consult.
+ :param tool_name: See :class:`BaseManagedAgentToolset`.
+ :param description: See :class:`BaseManagedAgentToolset`.
+ :param timeout: Passed to the client on every request as ``ManagedAgentRequest.timeout``.
+ ``None`` means whatever the vendor client defaults to, which for Agent Engine is no
+ deadline at all.
+ :param max_retries: See :class:`BaseManagedAgentToolset`.
+ :param replayable: Whether the durable cache may replay a completed call. Only set it
+ for an agent that is read-only.
+ :param vendor_options: Sent with every request as ``ManagedAgentRequest.vendor_options``,
+ for per-agent settings the vendor hook accepts there (Agent Engine's ``class_method``,
+ for instance). The hook decides which keys are allowed.
"""
def __init__(
self,
+ client: ManagedAgentClient,
*,
- members: Sequence[BaseManagedAgentToolset],
- failover_on: tuple[type[BaseException], ...] = (Exception,),
- **kwargs,
+ tool_name: str,
+ description: str | None = None,
+ timeout: float | None = None,
+ max_retries: int = 1,
+ replayable: bool = False,
+ vendor_options: dict[str, Any] | None = None,
) -> 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."
+ super().__init__(
+ tool_name=tool_name, description=description, timeout=timeout, max_retries=max_retries
+ )
+ # Not isinstance(client, ManagedAgentClient): on Python 3.10 and 3.11 a runtime_checkable
+ # Protocol check reads every member, which would resolve the agent's identity here.
+ if not callable(getattr(client, "invoke", None)) or not hasattr(type(client), "ref"):
+ raise TypeError(
+ f"{type(client).__name__} is not a ManagedAgentClient. Pass hook.agent(...) rather than "
+ "the hook, or an object with `ref`, `capabilities` and `invoke`."
)
- # 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.
- self.replayable = all(m.replayable for m in members)
+ self._client = client
+ self._vendor_options = dict(vendor_options or {})
+ self.replayable = replayable
@property
- def agent_ref(self) -> dict[str, str]:
- return {
- "platform": "failover",
- "name": " -> ".join(m.agent_ref.get("name", "?") for m in self._members),
- }
+ def client(self) -> ManagedAgentClient:
+ return self._client
- 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={
- "tool": self._tool_name,
- "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={
- "tool": self._tool_name,
- "platform": ref.get("platform", "unknown"),
- "role": "standby" if served_by_standby else "primary",
- "position": str(position),
- },
- )
- return result
+ @property
+ def agent_ref(self) -> ManagedAgentRef:
+ return self._client.ref
+
+ def invoke_sync(self, prompt: str) -> str:
+ request = ManagedAgentRequest(
+ prompt=prompt, timeout=self.timeout, vendor_options=dict(self._vendor_options)
+ )
+ try:
+ response = self._client.invoke(request)
+ except ManagedAgentRejected as exc:
+ # The contract keeps pydantic-ai out of the vendor hooks; this is the one place
+ # a rejection becomes something the calling model can act on.
+ raise ModelRetry(str(exc)) from exc
+ return response.text
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_caching_toolset.py b/providers/common/ai/tests/unit/common/ai/durable/test_caching_toolset.py
index d0266e8996c..307b7285df9 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_caching_toolset.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_caching_toolset.py
@@ -22,6 +22,9 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from pydantic_ai.messages import ModelResponse, TextPart
from pydantic_ai.models import ModelRequestParameters
+from pydantic_ai.toolsets.abstract import ToolsetTool
+from pydantic_ai.toolsets.combined import CombinedToolset
+from pydantic_ai.toolsets.function import FunctionToolset
from airflow.providers.common.ai.durable.base import DURABLE_KEY_PREFIX as P, DurableStorageProtocol
from airflow.providers.common.ai.durable.caching_model import CachingModel
@@ -211,3 +214,55 @@ class TestSharedCounter:
assert model_keys == [f"{P}model_step_0", f"{P}model_step_2"]
assert tool_keys == [f"{P}tool_step_1"]
assert counter.total_steps == 3
+
+
+class TestCachingToolsetReplayable:
+ @pytest.mark.asyncio
+ async def test_a_non_replayable_toolset_is_never_served_from_cache(
+ self, mock_toolset, mock_storage, counter
+ ):
+ # A managed agent may have acted on a system Airflow cannot observe, so its toolset
+ # declares replayable=False and a cached answer must not stand in for a fresh call.
+ mock_toolset.replayable = False
+ mock_storage.load_tool_result.return_value = (
+ True,
+ "stale cached result",
+ fingerprint_tool_call("t", {}, "call_1"),
+ )
+ caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
+
+ tool = MagicMock(spec=ToolsetTool)
+ tool.toolset = mock_toolset
+
+ result = await caching.call_tool("t", {}, ctx_for(), tool=tool)
+
+ assert result == "fresh result"
+ mock_toolset.call_tool.assert_awaited_once()
+ mock_storage.load_tool_result.assert_not_called()
+ mock_storage.save_tool_result.assert_not_called()
+ assert counter.replayed_tool == 0
+ # The step is still consumed so later steps keep their keys.
+ assert counter.next_step() == 1
+
+ @pytest.mark.asyncio
+ async def test_the_flag_is_read_per_tool_through_prefixed_and_combined_wrappers(
+ self, mock_storage, counter
+ ):
+ # ``.prefixed()`` and ``CombinedToolset`` do not carry the attribute; the cache reads it off the
+ # toolset each tool came from, so one non-replayable member does not stop its siblings replaying.
+ class NonReplayable(FunctionToolset):
+ replayable = False
+
+ managed, plain = NonReplayable(), FunctionToolset()
+ wrapped = CombinedToolset([plain, managed.prefixed("claims")])
+ mock_storage.load_tool_result.return_value = (True, "stale", fingerprint_tool_call("t", {}, "call_1"))
+ caching = CachingToolset(wrapped=wrapped, storage=mock_storage, counter=counter)
+ from_managed, from_plain = MagicMock(spec=ToolsetTool), MagicMock(spec=ToolsetTool)
+ from_managed.toolset, from_plain.toolset = managed, plain
+
+ with patch.object(CombinedToolset, "call_tool", autospec=True, return_value="fresh") as call_tool:
+ assert await caching.call_tool("t", {}, ctx_for(), tool=from_managed) == "fresh"
+ assert await caching.call_tool("t", {}, ctx_for(), tool=from_plain) == "stale"
+
+ call_tool.assert_awaited_once()
+ mock_storage.load_tool_result.assert_called_once()
diff --git a/providers/common/ai/tests/unit/common/ai/managed_agents/__init__.py b/providers/common/ai/tests/unit/common/ai/managed_agents/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/managed_agents/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git a/providers/common/ai/tests/unit/common/ai/managed_agents/test_base.py b/providers/common/ai/tests/unit/common/ai/managed_agents/test_base.py
new file mode 100644
index 00000000000..a6f4ebff1c2
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/managed_agents/test_base.py
@@ -0,0 +1,137 @@
+# 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 subprocess
+import sys
+
+import pytest
+
+from airflow.providers.common.ai.managed_agents.base import (
+ BaseManagedAgentHook,
+ ManagedAgentCapabilities,
+ ManagedAgentClient,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ ManagedAgentResponse,
+ describe,
+)
+
+
+class RecordingHook(BaseManagedAgentHook):
+ """A hook with no vendor behind it, enough to exercise the contract's own logic."""
+
+ agent_platform = "test.platform"
+
+ def __init__(self):
+ self.calls: list[tuple[str, ManagedAgentRequest]] = []
+
+ def resolve_agent(self, agent: str) -> ManagedAgentRef:
+ if not agent:
+ raise ValueError("empty agent")
+ return ManagedAgentRef(platform=self.agent_platform, name=agent.upper())
+
+ def get_agent_capabilities(self, agent: str) -> ManagedAgentCapabilities:
+ return ManagedAgentCapabilities(sessions=True)
+
+ def invoke_agent(self, agent: str, request: ManagedAgentRequest) -> ManagedAgentResponse:
+ self.calls.append((agent, request))
+ return ManagedAgentResponse(text=f"{agent}: {request.prompt}", raw={"agent": agent})
+
+
+class TestManagedAgentRequest:
+ @pytest.mark.parametrize(
+ "kwargs",
+ [{}, {"prompt": "a", "messages": []}],
+ ids=["neither", "both"],
+ )
+ def test_needs_exactly_one_of_prompt_or_messages(self, kwargs):
+ with pytest.raises(ValueError, match="exactly one of prompt or messages"):
+ ManagedAgentRequest(**kwargs)
+
+ def test_a_prompt_renders_as_a_single_user_message(self):
+ assert ManagedAgentRequest(prompt="Why?").as_messages() == [
+ {"role": "user", "content": [{"type": "text", "text": "Why?"}]}
+ ]
+
+ def test_messages_are_copied_not_aliased(self):
+ messages = [{"role": "user", "content": "x"}]
+ rendered = ManagedAgentRequest(messages=messages).as_messages()
+ rendered.append({"role": "user", "content": "y"})
+ assert len(messages) == 1
+
+
+class TestBaseManagedAgentHook:
+ def test_has_no_init_so_it_can_be_mixed_in_beside_a_vendor_base(self):
+ # A vendor hook's own __init__ must win. The contract adds behaviour, not state.
+ assert "__init__" not in vars(BaseManagedAgentHook)
+
+ def test_bound_agent_forwards_and_resolves_lazily(self):
+ hook = RecordingHook()
+ bound = hook.agent("analyst")
+ assert isinstance(bound, ManagedAgentClient)
+ assert bound.ref == ManagedAgentRef(platform="test.platform", name="ANALYST")
+ assert bound.capabilities.sessions is True
+ response = bound.invoke(ManagedAgentRequest(prompt="q"))
+ assert response.text == "analyst: q"
+ assert hook.calls == [("analyst", ManagedAgentRequest(prompt="q"))]
+
+ def test_a_bad_agent_fails_at_resolution_not_construction(self):
+ bound = RecordingHook().agent("")
+ with pytest.raises(ValueError, match="empty agent"):
+ bound.ref
+
+ def test_a_session_is_refused_for_a_hook_without_sessions_before_the_hook_is_called(self):
+ class Stateless(RecordingHook):
+ def get_agent_capabilities(self, agent: str) -> ManagedAgentCapabilities:
+ return ManagedAgentCapabilities()
+
+ hook = Stateless()
+ with pytest.raises(ValueError, match="does not keep conversation state"):
+ hook.agent("analyst").invoke(ManagedAgentRequest(prompt="q", session_id="t:1"))
+ assert hook.calls == []
+
+ def test_a_session_reaches_a_hook_that_supports_them(self):
+ hook = RecordingHook() # declares sessions=True
+ hook.agent("analyst").invoke(ManagedAgentRequest(prompt="q", session_id="t:1"))
+ assert hook.calls[0][1].session_id == "t:1"
+
+
+class TestDescribe:
+ def test_renders_name_and_platform(self):
+ assert describe(RecordingHook().agent("analyst")) == "ANALYST on test.platform"
+
+ def test_never_raises(self):
+ assert describe(RecordingHook().agent("")) == "<unresolved: empty agent>"
+
+
+def test_contract_module_imports_without_pydantic_ai():
+ """
+ Vendor providers import the contract behind a guarded import. That import must stay cheap:
+ a provider that floors Airflow 2 pays for it on every hook import once the extra is installed.
+ """
+ code = (
+ "import sys\n"
+ "import airflow.providers.common.ai.managed_agents.base\n"
+ "import airflow.providers.common.ai.exceptions\n"
+ "loaded = sorted(m for m in sys.modules if m == 'pydantic_ai' or m.startswith('pydantic_ai.'))\n"
+ "print('PYDANTIC_AI_MODULES=' + repr(loaded))\n"
+ )
+ result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True)
+ # Airflow's logging may write to stdout on import; only the marker line is the assertion.
+ marker = [line for line in result.stdout.splitlines() if line.startswith("PYDANTIC_AI_MODULES=")]
+ assert marker == ["PYDANTIC_AI_MODULES=[]"], result.stdout
diff --git a/providers/common/ai/tests/unit/common/ai/managed_agents/test_failover.py b/providers/common/ai/tests/unit/common/ai/managed_agents/test_failover.py
new file mode 100644
index 00000000000..8a7ba89a650
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/managed_agents/test_failover.py
@@ -0,0 +1,205 @@
+# 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 unittest import mock
+
+import pytest
+
+from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError, ManagedAgentRejected
+from airflow.providers.common.ai.managed_agents import (
+ FailoverManagedAgentClient,
+ ManagedAgentCapabilities,
+ ManagedAgentClient,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ ManagedAgentResponse,
+)
+
+
+class FakeClient:
+ def __init__(
+ self,
+ name: str,
+ *,
+ platform: str = "fake.cloud",
+ raises: Exception | None = None,
+ ref_raises: Exception | None = None,
+ **capabilities: bool,
+ ):
+ self._name = name
+ self._platform = platform
+ self._raises = raises
+ self._ref_raises = ref_raises
+ self.capabilities = ManagedAgentCapabilities(**capabilities)
+ self.calls = 0
+
+ @property
+ def ref(self) -> ManagedAgentRef:
+ if self._ref_raises is not None:
+ raise self._ref_raises
+ return ManagedAgentRef(platform=self._platform, name=self._name)
+
+ def invoke(self, request: ManagedAgentRequest) -> ManagedAgentResponse:
+ self.calls += 1
+ if self._raises is not None:
+ raise self._raises
+ return ManagedAgentResponse(text=f"{self._name} answered", raw={})
+
+
+REQUEST = ManagedAgentRequest(prompt="q")
+
+
+class TestConstruction:
+ @pytest.mark.parametrize("count", [0, 1], ids=["none", "one"])
+ def test_needs_at_least_two_members(self, count):
+ with pytest.raises(ValueError, match="at least two members"):
+ FailoverManagedAgentClient([FakeClient("a")] * count)
+
+ def test_members_are_copied_so_the_caller_cannot_empty_the_group(self):
+ members = [FakeClient("a"), FakeClient("b")]
+ group = FailoverManagedAgentClient(members)
+ members.clear()
+ assert len(group.members) == 2
+
+ def test_session_capable_members_are_accepted_but_the_group_never_offers_sessions(self):
+ group = FailoverManagedAgentClient([FakeClient("a", sessions=True), FakeClient("b", sessions=True)])
+ assert group.capabilities.sessions is False
+
+ def test_is_itself_a_client_so_groups_nest(self):
+ inner = FailoverManagedAgentClient([FakeClient("a"), FakeClient("b")])
+ outer = FailoverManagedAgentClient([inner, FakeClient("c")])
+ assert isinstance(inner, ManagedAgentClient)
+ assert outer.invoke(REQUEST).text == "a answered"
+
+
+class TestIdentityAndCapabilities:
+ def test_ref_names_every_member_in_order(self):
+ group = FailoverManagedAgentClient([FakeClient("a"), FakeClient("b", platform="other.cloud")])
+ assert group.ref == ManagedAgentRef(platform="failover", name="a on fake.cloud -> b on other.cloud")
+
+ def test_ref_survives_a_member_whose_identity_cannot_be_resolved(self):
+ group = FailoverManagedAgentClient(
+ [FakeClient("a"), FakeClient("b", ref_raises=RuntimeError("Connection 'b' not found"))]
+ )
+ assert group.ref.name == "a on fake.cloud -> <unresolved: Connection 'b' not found>"
+
+ def test_capabilities_are_the_intersection(self):
+ group = FailoverManagedAgentClient(
+ [FakeClient("a", usage=True, trace=True), FakeClient("b", usage=True)],
+ )
+ assert group.capabilities == ManagedAgentCapabilities(usage=True)
+
+ def test_a_session_bound_request_is_refused_before_any_member_is_tried(self):
+ primary, standby = FakeClient("a", sessions=True), FakeClient("b", sessions=True)
+ with pytest.raises(ValueError, match="cannot continue a conversation"):
+ FailoverManagedAgentClient([primary, standby]).invoke(
+ ManagedAgentRequest(prompt="q", session_id="t:1")
+ )
+ assert (primary.calls, standby.calls) == (0, 0)
+
+
+class TestInvoke:
+ def test_primary_answer_wins_and_standby_is_untouched(self):
+ primary, standby = FakeClient("a"), FakeClient("b")
+ assert FailoverManagedAgentClient([primary, standby]).invoke(REQUEST).text == "a answered"
+ assert (primary.calls, standby.calls) == (1, 0)
+
+ @pytest.mark.parametrize(
+ "error",
+ [ManagedAgentInvocationError("down"), RuntimeError("503"), TimeoutError("read timed out")],
+ ids=["terminal", "transient", "timeout"],
+ )
+ def test_fails_over_when_primary_fails(self, error):
+ standby = FakeClient("b")
+ group = FailoverManagedAgentClient([FakeClient("a", raises=error), standby])
+ assert group.invoke(REQUEST).text == "b answered"
+ assert standby.calls == 1
+
+ def test_a_broken_primary_identity_does_not_prevent_failover(self):
+ # Identity is resolved after the call, so a member whose connection is misconfigured
+ # cannot fail a tool call another member can serve.
+ primary = FakeClient("a", raises=RuntimeError("503"), ref_raises=RuntimeError("no connection"))
+ standby = FakeClient("b", ref_raises=RuntimeError("no connection either"))
+ assert FailoverManagedAgentClient([primary, standby]).invoke(REQUEST).text == "b answered"
+
+ def test_rejection_does_not_burn_the_standby(self):
+ standby = FakeClient("b")
+ group = FailoverManagedAgentClient([FakeClient("a", raises=ManagedAgentRejected("vague")), standby])
+ with pytest.raises(ManagedAgentRejected, match="vague"):
+ group.invoke(REQUEST)
+ assert standby.calls == 0
+
+ def test_last_members_error_propagates_when_all_fail(self):
+ group = FailoverManagedAgentClient(
+ [FakeClient("a", raises=RuntimeError("a down")), FakeClient("b", raises=RuntimeError("b down"))]
+ )
+ with pytest.raises(RuntimeError, match="b down"):
+ group.invoke(REQUEST)
+
+ def test_narrowed_failover_on_lets_other_errors_through(self):
+ standby = FakeClient("b")
+ group = FailoverManagedAgentClient(
+ [FakeClient("a", raises=RuntimeError("not covered")), standby], failover_on=(TimeoutError,)
+ )
+ with pytest.raises(RuntimeError, match="not covered"):
+ group.invoke(REQUEST)
+ assert standby.calls == 0
+
+ def test_warns_on_failover_and_names_both_members(self, caplog):
+ group = FailoverManagedAgentClient([FakeClient("a", raises=RuntimeError("503")), FakeClient("b")])
+ with caplog.at_level("INFO"):
+ group.invoke(REQUEST)
+ assert "a on fake.cloud (member 0) failed; failing over to b on fake.cloud (member 1)" in caplog.text
+ assert "served by standby b on fake.cloud (member 1)" in caplog.text
+
+
+class TestMetrics:
+ """A failover is a success-shaped event, so the counter is the only signal that a primary is down."""
+
+ @mock.patch("airflow.providers.common.ai.managed_agents.failover.Stats.incr", autospec=True)
+ def test_primary_success_emits_nothing_here(self, mock_incr):
+ # The toolset counts answers; the group counts only transitions.
+ FailoverManagedAgentClient([FakeClient("a"), FakeClient("b")]).invoke(REQUEST)
+ mock_incr.assert_not_called()
+
+ @mock.patch("airflow.providers.common.ai.managed_agents.failover.Stats.incr", autospec=True)
+ def test_failover_is_counted_by_platform_pair(self, mock_incr):
+ group = FailoverManagedAgentClient(
+ [FakeClient("a", raises=RuntimeError("503")), FakeClient("b", platform="other.cloud")]
+ )
+ group.invoke(REQUEST)
+ mock_incr.assert_called_once_with(
+ "managed_agent.failover", tags={"from_platform": "fake.cloud", "to_platform": "other.cloud"}
+ )
+
+ @mock.patch("airflow.providers.common.ai.managed_agents.failover.Stats.incr", autospec=True)
+ def test_unresolvable_identity_tags_as_unknown(self, mock_incr):
+ group = FailoverManagedAgentClient(
+ [FakeClient("a", raises=RuntimeError("503"), ref_raises=RuntimeError("x")), FakeClient("b")]
+ )
+ group.invoke(REQUEST)
+ assert mock_incr.call_args.kwargs["tags"]["from_platform"] == "unknown"
+
+ @mock.patch("airflow.providers.common.ai.managed_agents.failover.Stats.incr", autospec=True)
+ def test_rejection_is_not_a_failover(self, mock_incr):
+ group = FailoverManagedAgentClient(
+ [FakeClient("a", raises=ManagedAgentRejected("x")), FakeClient("b")]
+ )
+ with pytest.raises(ManagedAgentRejected):
+ group.invoke(REQUEST)
+ mock_incr.assert_not_called()
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 2b01077722d..9bb03c1c27b 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
@@ -22,18 +22,28 @@ from typing import Any
from unittest import mock
import pytest
+from pydantic_ai import Agent
+from pydantic_ai._run_context import RunContext
from pydantic_ai.exceptions import ModelRetry
+from pydantic_ai.messages import ModelResponse, RetryPromptPart, TextPart, ToolCallPart, ToolReturnPart
+from pydantic_ai.models.function import FunctionModel
from pydantic_core import ValidationError
-from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError
-from airflow.providers.common.ai.toolsets.managed_agent import (
- BaseManagedAgentToolset,
- FailoverManagedAgentToolset,
+from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError, ManagedAgentRejected
+from airflow.providers.common.ai.managed_agents import (
+ FailoverManagedAgentClient,
+ ManagedAgentCapabilities,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ ManagedAgentResponse,
)
+from airflow.providers.common.ai.toolsets.managed_agent import BaseManagedAgentToolset, ManagedAgentToolset
+
+FAKE_REF = ManagedAgentRef(platform="fake.cloud", name="specialist-1")
class FakeManagedAgentToolset(BaseManagedAgentToolset):
- """Minimal implementation standing in for a provider's concrete toolset."""
+ """Minimal direct subclass standing in for an agent that has no hook."""
def __init__(self, *, result: Any = "the answer", raises: Exception | None = None, **kwargs):
kwargs.setdefault("tool_name", "ask_specialist")
@@ -44,8 +54,8 @@ class FakeManagedAgentToolset(BaseManagedAgentToolset):
self.prompts: list[str] = []
@property
- def agent_ref(self) -> dict[str, str]:
- return {"platform": "fake.cloud", "name": "specialist-1"}
+ def agent_ref(self) -> ManagedAgentRef:
+ return FAKE_REF
async def invoke(self, prompt: str) -> Any:
self.prompts.append(prompt)
@@ -54,20 +64,56 @@ class FakeManagedAgentToolset(BaseManagedAgentToolset):
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]
+class FakeClient:
+ """A ManagedAgentClient with no hook behind it."""
+
+ def __init__(
+ self,
+ name: str = "specialist-1",
+ *,
+ answer: str = "the answer",
+ raises: Exception | None = None,
+ raise_times: int = 1,
+ ref_raises: Exception | None = None,
+ sessions: bool = False,
+ ):
+ self._name = name
+ self._answer = answer
+ self._raises = raises
+ self._raise_times = raise_times
+ self._ref_raises = ref_raises
+ self.capabilities = ManagedAgentCapabilities(sessions=sessions)
+ self.requests: list[ManagedAgentRequest] = []
+
+ @property
+ def ref(self) -> ManagedAgentRef:
+ if self._ref_raises is not None:
+ raise self._ref_raises
+ return ManagedAgentRef(platform="fake.cloud", name=self._name)
+
+ def invoke(self, request: ManagedAgentRequest) -> ManagedAgentResponse:
+ self.requests.append(request)
+ if self._raises is not None and len(self.requests) <= self._raise_times:
+ raise self._raises
+ return ManagedAgentResponse(text=f"{self._answer}: {request.prompt}", raw={"envelope": True})
+
+def served(mock_incr: mock.Mock) -> list[mock._Call]:
+ """The ``managed_agent.served`` increments; ``AirflowToolset`` counts every call under its own metric too."""
+ return [c for c in mock_incr.call_args_list if c.args[0] == "managed_agent.served"]
+
+
+async def call(toolset: BaseManagedAgentToolset, prompt: str = "what is the number?") -> Any:
+ ctx = mock.MagicMock(spec=RunContext)
+ tools = await toolset.get_tools(ctx)
+ return await toolset.call_tool(toolset._tool_name, {"prompt": prompt}, ctx, tools[toolset._tool_name])
+
+
+class TestBaseManagedAgentToolsetConstruction:
@pytest.mark.parametrize(
- "description",
- [None, "", " ", "\n"],
- ids=["none", "empty", "whitespace", "newline"],
+ "description", [None, "", " ", "\n"], ids=["none", "empty", "whitespace", "newline"]
)
def test_absent_description_falls_back_to_the_tool_name(self, description):
- # Matches HookToolset, which derives a description from the method name
- # when there is no docstring. The name is the required identifier; the
- # description is guidance the author may omit.
toolset = FakeManagedAgentToolset(tool_name="ask_bookings_analyst", description=description)
assert toolset._description == "Ask bookings analyst"
@@ -79,30 +125,13 @@ class TestBaseManagedAgentToolsetConstruction:
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
-
@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
@@ -113,8 +142,8 @@ class TestBaseManagedAgentToolsetConstruction:
def test_subclass_implementing_neither_invoke_hook_is_rejected(self):
class Neither(BaseManagedAgentToolset):
@property
- def agent_ref(self) -> dict[str, str]:
- return {"platform": "fake.cloud", "name": "neither"}
+ def agent_ref(self) -> ManagedAgentRef:
+ return FAKE_REF
with pytest.raises(TypeError, match="must implement invoke_sync"):
Neither(tool_name="ask_nothing")
@@ -125,8 +154,8 @@ class TestSyncInvocation:
class _Blocking(BaseManagedAgentToolset):
@property
- def agent_ref(self) -> dict[str, str]:
- return {"platform": "fake.cloud", "name": "blocking"}
+ def agent_ref(self) -> ManagedAgentRef:
+ return FAKE_REF
def invoke_sync(self, prompt: str) -> Any:
time.sleep(0.2)
@@ -134,9 +163,6 @@ class TestSyncInvocation:
@pytest.mark.asyncio
async def test_invoke_sync_is_offloaded_so_the_loop_keeps_running(self):
- # A synchronous client would freeze every other tool call for the
- # duration of the remote call if it ran on the loop, so the base class
- # runs invoke_sync in a worker thread.
ticks = 0
async def ticker():
@@ -148,36 +174,28 @@ class TestSyncInvocation:
task = asyncio.create_task(ticker())
result = await self._Blocking(tool_name="ask_slow").invoke("q")
task.cancel()
-
assert result == "slept for q"
assert ticks >= 3, "the event loop was blocked while invoke_sync ran"
@pytest.mark.asyncio
async def test_invoke_sync_result_reaches_the_model_through_call_tool(self):
- toolset = self._Blocking(tool_name="ask_slow")
- tools = await toolset.get_tools(ctx=None)
- result = await toolset.call_tool("ask_slow", {"prompt": "q"}, None, tools["ask_slow"])
- assert result == "slept for q"
+ assert await call(self._Blocking(tool_name="ask_slow"), "q") == "slept for q"
@pytest.mark.asyncio
async def test_an_async_override_is_used_as_is(self):
- # FakeManagedAgentToolset overrides invoke() directly, the path a vendor
- # with a natively async client takes.
assert await FakeManagedAgentToolset(result="async answer").invoke("q") == "async answer"
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)
+ tools = await FakeManagedAgentToolset(tool_name="ask_bookings").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"]
@@ -185,50 +203,54 @@ class TestGetTools:
@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
+ tool_def = (await FakeManagedAgentToolset().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
+ validator = (await FakeManagedAgentToolset().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?")
+ await 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]"),
- ],
+ [("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
+ async def test_result_is_serialized_for_the_model(self, result, expected):
+ assert await call(FakeManagedAgentToolset(result=result)) == expected
@pytest.mark.asyncio
async def test_logs_the_agent_it_consulted(self, caplog):
- await self._call(FakeManagedAgentToolset())
+ with caplog.at_level("INFO"):
+ await call(FakeManagedAgentToolset())
assert "specialist-1" in caplog.text
assert "fake.cloud" in caplog.text
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats.incr", autospec=True)
+ async def test_every_answer_is_counted_by_tool_and_platform(self, mock_incr):
+ await call(FakeManagedAgentToolset(tool_name="ask_bookings"))
+ assert served(mock_incr) == [
+ mock.call("managed_agent.served", tags={"tool": "ask_bookings", "platform": "fake.cloud"})
+ ]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats.incr", autospec=True)
+ async def test_a_failed_call_is_not_counted_as_served(self, mock_incr):
+ with pytest.raises(RuntimeError):
+ await call(FakeManagedAgentToolset(raises=RuntimeError("503")))
+ assert served(mock_incr) == []
+
@pytest.mark.asyncio
@pytest.mark.parametrize(
"error",
@@ -240,198 +262,93 @@ class TestCallTool:
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).
+ # 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))
+ await 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_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")
- assert await self._call(self._group(primary, standby)) == "from primary"
- assert standby.prompts == []
+class TestManagedAgentToolset:
+ """The client-backed toolset, which is the one Dags are expected to use."""
@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?"]
+ async def test_model_gets_text_never_the_envelope(self):
+ client = FakeClient()
+ assert await call(ManagedAgentToolset(client, tool_name="ask"), "q") == "the answer: q"
@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 == []
+ async def test_timeout_travels_on_the_request(self):
+ client = FakeClient()
+ await call(ManagedAgentToolset(client, tool_name="ask", timeout=42.0), "q")
+ assert client.requests == [ManagedAgentRequest(prompt="q", timeout=42.0)]
- @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))
+ def test_something_that_is_not_a_client_is_refused_at_construction(self):
+ # Passing the hook instead of hook.agent(...) is the likely mistake; it must not wait for the first call.
+ with pytest.raises(TypeError, match="Pass hook.agent"):
+ ManagedAgentToolset(object(), tool_name="ask") # type: ignore[arg-type]
@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 == []
+ async def test_vendor_options_travel_on_every_request(self):
+ client = FakeClient()
+ await call(ManagedAgentToolset(client, tool_name="ask", vendor_options={"class_method": "plan"}), "q")
+ assert client.requests[0].vendor_options == {"class_method": "plan"}
@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
+ async def test_rejection_becomes_a_model_retry_at_the_boundary(self):
+ client = FakeClient(raises=ManagedAgentRejected("too vague"))
+ with pytest.raises(ModelRetry, match="too vague"):
+ await call(ManagedAgentToolset(client, tool_name="ask"), "q")
@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"],
+ "error", [ManagedAgentInvocationError("gone"), RuntimeError("503")], ids=["terminal", "transient"]
)
- 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.",
- )
+ async def test_other_errors_propagate_unchanged(self, error):
+ with pytest.raises(type(error), match=str(error)):
+ await call(ManagedAgentToolset(FakeClient(raises=error), tool_name="ask"), "q")
@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={
- "tool": "ask_resilient",
- "platform": "fake.cloud",
- "role": "primary",
- "position": "0",
- },
- )
+ async def test_a_broken_identity_does_not_fail_the_call(self, caplog):
+ # Identity is for the log line and the metric tag; the answer must not depend on it.
+ client = FakeClient(ref_raises=RuntimeError("Connection 'standby' not found"))
+ with caplog.at_level("WARNING"):
+ assert await call(ManagedAgentToolset(client, tool_name="ask"), "q") == "the answer: q"
+ assert "identity could not be resolved" in caplog.text
+ assert "Connection 'standby' not found" in caplog.text
@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={
- "tool": "ask_resilient",
- "from_platform": "fake.cloud",
- "to_platform": "fake.cloud",
- },
- ),
- 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",
- },
- ),
+ @mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats.incr", autospec=True)
+ async def test_a_failover_group_is_counted_under_the_failover_platform(self, mock_incr):
+ group = FailoverManagedAgentClient([FakeClient("a"), FakeClient("b")])
+ await call(ManagedAgentToolset(group, tool_name="ask"), "q")
+ assert served(mock_incr) == [
+ mock.call("managed_agent.served", tags={"tool": "ask", "platform": "failover"})
]
- @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()
+ def test_rephrase_loop_end_to_end_with_a_scripted_model(self):
+ """The rejected first prompt is re-asked once, then the answer reaches the model."""
+ client = FakeClient(raises=ManagedAgentRejected("Which quarter?"), raise_times=1)
+ seen: list[Any] = []
+
+ def scripted(messages, info):
+ if len(messages) == 1:
+ return ModelResponse(
+ parts=[ToolCallPart(info.function_tools[0].name, {"prompt": "Revenue?"})]
+ )
+ parts = messages[-1].parts
+ if any(isinstance(p, RetryPromptPart) for p in parts):
+ seen.extend(p for p in parts if isinstance(p, RetryPromptPart))
+ return ModelResponse(
+ parts=[ToolCallPart(info.function_tools[0].name, {"prompt": "Q3 revenue?"})]
+ )
+ seen.extend(p for p in parts if isinstance(p, ToolReturnPart))
+ return ModelResponse(parts=[TextPart("Done")])
+
+ toolset = ManagedAgentToolset(client, tool_name="ask_analyst", max_retries=1)
+ result = Agent(FunctionModel(scripted), toolsets=[toolset]).run_sync("go")
+
+ assert result.output == "Done"
+ assert [type(p).__name__ for p in seen] == ["RetryPromptPart", "ToolReturnPart"]
+ assert "Which quarter?" in seen[0].content
+ assert seen[1].content == "the answer: Q3 revenue?"
+ assert [r.prompt for r in client.requests] == ["Revenue?", "Q3 revenue?"]
diff --git a/providers/google/docs/index.rst b/providers/google/docs/index.rst
index 305af286442..2bfceb0ed7f 100644
--- a/providers/google/docs/index.rst
+++ b/providers/google/docs/index.rst
@@ -213,6 +213,7 @@ Dependent package
`apache-airflow-providers-apache-beam <https://airflow.apache.org/docs/apache-airflow-providers-apache-beam>`_ ``apache.beam``
`apache-airflow-providers-apache-cassandra <https://airflow.apache.org/docs/apache-airflow-providers-apache-cassandra>`_ ``apache.cassandra``
`apache-airflow-providers-cncf-kubernetes <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes>`_ ``cncf.kubernetes``
+`apache-airflow-providers-common-ai <https://airflow.apache.org/docs/apache-airflow-providers-common-ai>`_ ``common.ai``
`apache-airflow-providers-common-messaging <https://airflow.apache.org/docs/apache-airflow-providers-common-messaging>`_ ``common.messaging``
`apache-airflow-providers-facebook <https://airflow.apache.org/docs/apache-airflow-providers-facebook>`_ ``facebook``
`apache-airflow-providers-http <https://airflow.apache.org/docs/apache-airflow-providers-http>`_ ``http``
@@ -266,6 +267,7 @@ Extra Dependencies
``trino`` ``apache-airflow-providers-trino``
``http`` ``apache-airflow-providers-http``
``standard`` ``apache-airflow-providers-standard``
+``common.ai`` ``apache-airflow-providers-common-ai>=0.10.0``
``common.messaging`` ``apache-airflow-providers-common-messaging>=2.0.0``
==================== ====================================================
diff --git a/providers/google/docs/operators/cloud/vertex_ai.rst b/providers/google/docs/operators/cloud/vertex_ai.rst
index dd69e361de2..6cfd165fe53 100644
--- a/providers/google/docs/operators/cloud/vertex_ai.rst
+++ b/providers/google/docs/operators/cloud/vertex_ai.rst
@@ -83,6 +83,30 @@ The same operation can be performed in the deferrable mode.
:start-after: [START how_to_cloud_vertex_ai_run_query_job_operator_deferrable]
:end-before: [END how_to_cloud_vertex_ai_run_query_job_operator_deferrable]
+To let an agent running under Common AI's ``AgentOperator`` consult an Agent Engine as one of
+its tools, use
+:class:`~airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.AgentEngineHook`.
+It implements the Common AI managed-agent contract over ``query_reasoning_engine``, so
+``hook.agent(resource_name)`` can be passed to a ``ManagedAgentToolset`` or combined with agents
+on other clouds in a ``FailoverManagedAgentClient``. The hook needs the ``common.ai`` extra of
+this provider, which installs ``apache-airflow-providers-common-ai`` and therefore requires
+Airflow 3. The agent is the engine's full resource name,
+``projects/P/locations/L/reasoningEngines/ID``, so one hook reaches engines in several projects
+and regions. A prompt is sent under ``input`` to the engine's ``query`` method; ``input_key`` and
+``class_method`` in ``vendor_options`` change that per request. A mapping output with a string ``output`` field is the answer text, and
+any other output is returned as JSON. Query *jobs* remain the domain of ``RunQueryJobOperator``,
+which can defer.
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.toolsets import ManagedAgentToolset
+ from airflow.providers.google.cloud.hooks.vertex_ai.agent_engine import AgentEngineHook
+
+ analyst = AgentEngineHook(gcp_conn_id="google_cloud_default").agent(
+ "projects/my-project/locations/us-central1/reasoningEngines/1234567890"
+ )
+ toolset = ManagedAgentToolset(analyst, tool_name="ask_analyst", description="Answers revenue questions.")
+
To update an Agent Engine you can use
:class:`~airflow.providers.google.cloud.operators.vertex_ai.agent_engine.UpdateAgentEngineOperator`.
diff --git a/providers/google/pyproject.toml b/providers/google/pyproject.toml
index 25b59542989..451fdbd2f83 100644
--- a/providers/google/pyproject.toml
+++ b/providers/google/pyproject.toml
@@ -222,6 +222,9 @@ dependencies = [
"standard" = [
"apache-airflow-providers-standard"
]
+"common.ai" = [
+ "apache-airflow-providers-common-ai>=0.10.0", # use next version
+]
"common.messaging" = [
"apache-airflow-providers-common-messaging>=2.0.0"
]
@@ -235,6 +238,7 @@ dev = [
"apache-airflow-providers-apache-beam",
"apache-airflow-providers-apache-cassandra",
"apache-airflow-providers-cncf-kubernetes",
+ "apache-airflow-providers-common-ai",
"apache-airflow-providers-common-compat",
"apache-airflow-providers-common-messaging",
"apache-airflow-providers-common-sql",
diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py
index ab6c130df7d..64a15f1be60 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py
@@ -19,15 +19,19 @@
from __future__ import annotations
+import json
import time
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any
import google.auth.transport.requests
from asgiref.sync import sync_to_async
+from google.api_core.exceptions import InvalidArgument, NotFound, PermissionDenied, Unauthenticated
from google.cloud.aiplatform_v1 import ReasoningEngineExecutionServiceClient
+from google.cloud.aiplatform_v1.types import QueryReasoningEngineResponse
from vertexai import Client
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
from airflow.providers.google.common.consts import CLIENT_INFO
from airflow.providers.google.common.hooks.base_google import (
PROVIDE_PROJECT_ID,
@@ -37,9 +41,42 @@ from airflow.providers.google.common.hooks.base_google import (
if TYPE_CHECKING:
from google.api_core.retry import Retry
- from google.cloud.aiplatform_v1.types import QueryReasoningEngineResponse
from vertexai._genai import types
+ from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError
+ from airflow.providers.common.ai.managed_agents.base import (
+ BaseManagedAgentHook,
+ ManagedAgentCapabilities,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ ManagedAgentResponse,
+ )
+else:
+ try:
+ from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError
+ from airflow.providers.common.ai.managed_agents.base import (
+ BaseManagedAgentHook,
+ ManagedAgentCapabilities,
+ ManagedAgentRef,
+ ManagedAgentResponse,
+ )
+ except ImportError:
+ # The Common AI provider is optional. This module still imports without it, and every
+ # managed-agent entry point on AgentEngineHook then says what is missing.
+ def _needs_common_ai(*args: Any, **kwargs: Any) -> Any:
+ raise AirflowOptionalProviderFeatureException(
+ "Consulting an Agent Engine as a managed agent needs the 'common.ai' extra of the "
+ "google provider: pip install 'apache-airflow-providers-google[common.ai]'."
+ )
+
+ class BaseManagedAgentHook:
+ """Stand-in for the Common AI contract base; ``agent()`` names the missing extra."""
+
+ agent = _needs_common_ai
+
+ ManagedAgentCapabilities = ManagedAgentRef = ManagedAgentResponse = _needs_common_ai
+ ManagedAgentInvocationError = _needs_common_ai
+
VERTEX_AI_AGENT_ENGINE_API_VERSION = "v1beta1"
VERTEX_AI_AGENT_ENGINE_OPERATION_URL = (
@@ -66,15 +103,63 @@ def serialize_value(value: Any) -> Any:
return value
-class AgentEngineHook(GoogleBaseHook):
+# ``vendor_options`` of a managed-agent request. ``class_method`` and ``input_key`` shape the
+# request; ``retry`` and ``metadata`` reach ``query_reasoning_engine``, which accepts exactly those.
+_AGENT_OPTIONS = frozenset({"class_method", "input_key", "retry", "metadata"})
+
+
+def _agent_output_text(output: Any) -> str:
+ if output is None:
+ return ""
+ if isinstance(output, str):
+ return output
+ if isinstance(output, dict) and isinstance(output.get("output"), str):
+ return output["output"]
+ return json.dumps(output, default=str)
+
+
+class AgentEngineHook(GoogleBaseHook, BaseManagedAgentHook):
"""
Hook for Google Cloud Vertex AI Agent Engine APIs.
Wraps the ``agent_engines`` module of the Vertex AI SDK client and the
Reasoning Engine Execution Service GAPIC client:
https://docs.cloud.google.com/python/docs/reference/agentplatform/latest/vertexai._genai.agent_engines.AgentEngines
+
+ With the ``common.ai`` extra installed, the hook also implements the Common AI
+ managed-agent contract over the synchronous query path, so ``hook.agent(resource_name)``
+ can be handed to a ``ManagedAgentToolset``. The agent is the engine's full resource name,
+ ``projects/P/locations/L/reasoningEngines/ID``, so a single hook on one connection reaches
+ engines in several projects and regions. A request carrying a ``prompt`` is sent as
+ ``{input_key: prompt}`` to ``class_method`` (``query`` and ``input`` by default; both can be
+ set per request in ``vendor_options``, alongside ``retry`` and ``metadata``); a request
+ carrying ``messages`` is sent as ``{"messages": [...]}``. When the engine returns a mapping
+ with a string ``output``, that is the answer text; any other output is returned as JSON text
+ and kept on ``ManagedAgentResponse.structured``. The wire type is a protobuf ``Value``, so
+ integers come back as floats and an absent, null or empty output all read as an empty string.
+
+ Agent Engine reports an author-side mistake, such as an unknown ``class_method`` or an
+ input under the wrong key, as ``INVALID_ARGUMENT``. That is a configuration error, not
+ something a model rephrase could fix, so it is terminal here; the hook never raises
+ :class:`~airflow.providers.common.ai.exceptions.ManagedAgentRejected`. The query path keeps
+ no conversation state, so a request with a ``session_id`` is refused rather than silently
+ sent as a fresh call. Query *jobs* remain the domain of
+ :class:`~airflow.providers.google.cloud.operators.vertex_ai.agent_engine.RunQueryJobOperator`,
+ which can defer.
+
+ .. code-block:: python
+
+ from airflow.providers.common.ai.toolsets import ManagedAgentToolset
+ from airflow.providers.google.cloud.hooks.vertex_ai.agent_engine import AgentEngineHook
+
+ analyst = AgentEngineHook(gcp_conn_id="google_cloud_default").agent(
+ "projects/my-project/locations/us-central1/reasoningEngines/1234567890"
+ )
+ toolset = ManagedAgentToolset(analyst, tool_name="ask_analyst", description="...")
"""
+ agent_platform = "gcp.vertex_agent_engine"
+
def __init__(
self,
gcp_conn_id: str = "google_cloud_default",
@@ -86,6 +171,76 @@ class AgentEngineHook(GoogleBaseHook):
impersonation_chain=impersonation_chain,
**kwargs,
)
+ # One execution client per location for managed-agent calls: a GAPIC client opens its gRPC
+ # channel in __init__, and a model may consult the same agent several times in one run.
+ self._agent_clients: dict[str, ReasoningEngineExecutionServiceClient] = {}
+
+ @staticmethod
+ def _parse_agent(agent: str) -> tuple[str, str, str]:
+ parts = ReasoningEngineExecutionServiceClient.parse_reasoning_engine_path(agent)
+ # The SDK parser is non-greedy, so a child resource such as an operation name still matches.
+ if not parts or "/" in parts["reasoning_engine"]:
+ raise ValueError(
+ f"An Agent Engine agent is its full resource name projects/P/locations/L/reasoningEngines/ID, "
+ f"got {agent!r}."
+ )
+ return parts["project"], parts["location"], parts["reasoning_engine"]
+
+ def resolve_agent(self, agent: str) -> ManagedAgentRef:
+ self._parse_agent(agent)
+ return ManagedAgentRef(platform=self.agent_platform, name=agent)
+
+ def get_agent_capabilities(self, agent: str) -> ManagedAgentCapabilities:
+ return ManagedAgentCapabilities(structured_output=True)
+
+ def invoke_agent(self, agent: str, request: ManagedAgentRequest) -> ManagedAgentResponse:
+ _, location, _ = self._parse_agent(agent)
+ if request.session_id is not None:
+ raise ValueError(
+ "Agent Engine's query path keeps no conversation state; session_id is not supported."
+ )
+ unknown = set(request.vendor_options) - _AGENT_OPTIONS
+ if unknown:
+ raise ValueError(
+ f"vendor_options {sorted(unknown)} are not accepted; an Agent Engine request takes "
+ f"{sorted(_AGENT_OPTIONS)}."
+ )
+ options = dict(request.vendor_options)
+ class_method = options.pop("class_method", "query")
+ input_key = options.pop("input_key", "input")
+ if not isinstance(class_method, str) or not isinstance(input_key, str):
+ raise ValueError(
+ "vendor_options['class_method'] and vendor_options['input_key'] must be strings."
+ )
+ input_data = (
+ {input_key: request.prompt} if request.prompt is not None else {"messages": request.as_messages()}
+ )
+ client = self._agent_clients.get(location)
+ if client is None:
+ client = self._agent_clients[location] = self.get_reasoning_engine_execution_service_client(
+ location
+ )
+ try:
+ response = client.query_reasoning_engine(
+ request={"name": agent, "class_method": class_method, "input": input_data},
+ retry=options.get("retry"),
+ timeout=request.timeout,
+ metadata=options.get("metadata", ()),
+ )
+ except InvalidArgument as exc:
+ raise ManagedAgentInvocationError(
+ f"Agent Engine {agent} rejected the request (class_method={class_method!r}, "
+ f"input_key={input_key!r}): {exc}"
+ ) from exc
+ except (NotFound, PermissionDenied, Unauthenticated) as exc:
+ raise ManagedAgentInvocationError(f"Agent Engine {agent}: {exc}") from exc
+ raw = QueryReasoningEngineResponse.to_dict(response)
+ output = raw.get("output")
+ return ManagedAgentResponse(
+ text=_agent_output_text(output),
+ raw=raw,
+ structured=None if isinstance(output, str) else output,
+ )
def get_agent_engine_client(self, project_id: str, location: str):
"""Return the Vertex AI Agent Engine client."""
diff --git a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_agent_engine.py b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_agent_engine.py
index 25f7ff46b8f..6012acecfce 100644
--- a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_agent_engine.py
+++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_agent_engine.py
@@ -17,15 +17,34 @@
# under the License.
from __future__ import annotations
+import subprocess
+import sys
from unittest import mock
import pytest
+from google.api_core.exceptions import InvalidArgument, NotFound, PermissionDenied, Unauthenticated
+from google.cloud.aiplatform_v1 import ReasoningEngineExecutionServiceClient
+from google.cloud.aiplatform_v1.types import QueryReasoningEngineResponse
from airflow.providers.google.cloud.hooks.vertex_ai.agent_engine import AgentEngineAsyncHook, AgentEngineHook
from airflow.providers.google.common.consts import CLIENT_INFO
from unit.google.cloud.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id
+try:
+ from airflow.providers.common.ai.exceptions import ManagedAgentInvocationError
+ from airflow.providers.common.ai.managed_agents import (
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ )
+except ImportError:
+ HAS_COMMON_AI = False
+else:
+ HAS_COMMON_AI = True
+
+# The managed-agent methods of AgentEngineHook need the optional common.ai extra.
+needs_common_ai = pytest.mark.skipif(not HAS_COMMON_AI, reason="requires apache-airflow-providers-common-ai")
+
BASE_STRING = "airflow.providers.google.common.hooks.base_google.{}"
AGENT_ENGINE_STRING = "airflow.providers.google.cloud.hooks.vertex_ai.agent_engine.{}"
@@ -470,3 +489,129 @@ class TestAgentEngineAsyncHook:
config=CHECK_QUERY_CONFIG,
)
assert result == mock.sentinel.query_job
+
+
+@pytest.fixture
+def execution_client():
+ client = mock.create_autospec(ReasoningEngineExecutionServiceClient, instance=True)
+ client.reasoning_engine_path.side_effect = ReasoningEngineExecutionServiceClient.reasoning_engine_path
+ client.query_reasoning_engine.return_value = QueryReasoningEngineResponse(
+ output={"output": "42", "steps": 3}
+ )
+ with mock.patch.object(
+ AgentEngineHook, "get_reasoning_engine_execution_service_client", autospec=True, return_value=client
+ ):
+ yield client
+
+
+@pytest.fixture
+def hook():
+ with mock.patch(BASE_STRING.format("GoogleBaseHook.__init__"), new=mock_base_gcp_hook_default_project_id):
+ yield AgentEngineHook(gcp_conn_id="google_cloud_default")
+
+
+@needs_common_ai
+class TestAgentEngineHookAgent:
+ def test_resolves_a_full_resource_name(self, hook):
+ assert hook.resolve_agent(AGENT_ENGINE_NAME) == ManagedAgentRef(
+ platform="gcp.vertex_agent_engine", name=AGENT_ENGINE_NAME
+ )
+
+ @pytest.mark.parametrize(
+ "agent", ["123", "projects/p/locations/l/agents/123", AGENT_ENGINE_NAME + "/operations/1"]
+ )
+ def test_rejects_anything_but_a_full_resource_name(self, hook, agent):
+ with pytest.raises(ValueError, match="reasoningEngines"):
+ hook.resolve_agent(agent)
+
+
+@needs_common_ai
+class TestAgentEngineHookInvokeAgent:
+ def test_prompt_goes_to_the_query_method_and_output_is_unwrapped(self, hook, execution_client):
+ response = hook.agent(AGENT_ENGINE_NAME).invoke(ManagedAgentRequest(prompt="Sum?", timeout=12.5))
+ kwargs = execution_client.query_reasoning_engine.call_args.kwargs
+ assert kwargs["request"] == {
+ "name": AGENT_ENGINE_NAME,
+ "class_method": "query",
+ "input": {"input": "Sum?"},
+ }
+ assert kwargs["timeout"] == 12.5
+ assert response.text == "42"
+ assert response.structured == {"output": "42", "steps": 3.0}
+ assert response.raw == {"output": {"output": "42", "steps": 3.0}}
+
+ def test_class_method_and_input_key_are_set_per_request(self, hook, execution_client):
+ options = {"class_method": "ask", "input_key": "question"}
+ hook.agent(AGENT_ENGINE_NAME).invoke(ManagedAgentRequest(prompt="Sum?", vendor_options=options))
+ request = execution_client.query_reasoning_engine.call_args.kwargs["request"]
+ assert request["class_method"] == "ask"
+ assert request["input"] == {"question": "Sum?"}
+
+ def test_messages_are_sent_as_messages(self, hook, execution_client):
+ messages = [{"role": "user", "content": "hi"}]
+ hook.agent(AGENT_ENGINE_NAME).invoke(ManagedAgentRequest(messages=messages))
+ assert execution_client.query_reasoning_engine.call_args.kwargs["request"]["input"] == {
+ "messages": messages
+ }
+
+ def test_a_plain_string_output_is_the_text(self, hook, execution_client):
+ execution_client.query_reasoning_engine.return_value = QueryReasoningEngineResponse(output="plain")
+ response = hook.agent(AGENT_ENGINE_NAME).invoke(ManagedAgentRequest(prompt="x"))
+ assert response.text == "plain"
+ assert response.structured is None
+
+ @pytest.mark.parametrize("option", ["reasoning_engine_id", "project_id", "output_gcs_uri"])
+ def test_only_the_query_methods_own_options_are_accepted(self, hook, execution_client, option):
+ with pytest.raises(ValueError, match=option):
+ hook.agent(AGENT_ENGINE_NAME).invoke(
+ ManagedAgentRequest(prompt="x", vendor_options={option: "9"})
+ )
+ execution_client.query_reasoning_engine.assert_not_called()
+
+ def test_a_session_is_refused_rather_than_silently_dropped(self, hook, execution_client):
+ # The bound client refuses it first; the hook refuses it too for callers on the raw path.
+ with pytest.raises(ValueError, match="query path keeps no conversation state"):
+ hook.invoke_agent(AGENT_ENGINE_NAME, ManagedAgentRequest(prompt="x", session_id="t"))
+ execution_client.query_reasoning_engine.assert_not_called()
+
+ @pytest.mark.parametrize(
+ ("error", "expected"),
+ [
+ (InvalidArgument("bad input"), "class_method='query'"),
+ (NotFound("no such engine"), "no such engine"),
+ (PermissionDenied("denied"), "denied"),
+ (Unauthenticated("expired"), "expired"),
+ ],
+ )
+ def test_api_errors_are_terminal_and_name_the_engine(self, hook, execution_client, error, expected):
+ # INVALID_ARGUMENT is an author-side mistake on Agent Engine, not something a rephrase fixes.
+ execution_client.query_reasoning_engine.side_effect = error
+ with pytest.raises(ManagedAgentInvocationError, match=expected) as excinfo:
+ hook.agent(AGENT_ENGINE_NAME).invoke(ManagedAgentRequest(prompt="x"))
+ assert AGENT_ENGINE_NAME in str(excinfo.value)
+
+ def test_a_bad_agent_name_fails_before_any_client_is_built(self, hook):
+ with mock.patch.object(
+ AgentEngineHook, "get_reasoning_engine_execution_service_client", autospec=True
+ ) as get:
+ with pytest.raises(ValueError, match="reasoningEngines"):
+ hook.agent("123").invoke(ManagedAgentRequest(prompt="x"))
+ get.assert_not_called()
+
+
+def test_the_hook_imports_without_common_ai_and_agent_names_the_missing_extra():
+ """The contract base is optional: without it the module loads and ``agent()`` says what to install."""
+ script = (
+ "import sys\n"
+ "for m in ('airflow.providers.common.ai.exceptions', 'airflow.providers.common.ai.managed_agents.base'):\n"
+ " sys.modules[m] = None\n"
+ "from airflow.providers.google.cloud.hooks.vertex_ai.agent_engine import AgentEngineHook\n"
+ "from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException\n"
+ "try:\n"
+ " AgentEngineHook.agent(object(), 'projects/p/locations/l/reasoningEngines/1')\n"
+ "except AirflowOptionalProviderFeatureException as e:\n"
+ " assert 'google[common.ai]' in str(e), e\n"
+ "else:\n"
+ " raise SystemExit('agent() did not raise')\n"
+ )
+ subprocess.run([sys.executable, "-c", script], check=True)
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 6e1182bf281..448f9aa702a 100644
--- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
+++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
@@ -400,21 +400,23 @@ metrics:
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."
+ description: "Number of failover transitions inside a ``FailoverManagedAgentClient`` group,
+ emitted once each time a member fails and the next one is tried. Metric with from_platform
+ and to_platform tagging; a platform reads ``unknown`` when that member's identity could not
+ be resolved."
type: "counter"
legacy_name: "-"
- name_variables: ["tool", "from_platform", "to_platform"]
+ name_variables: ["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."
+ description: "Number of managed agent invocations answered through a managed-agent toolset,
+ emitted once per answer. Metric with tool and platform tagging; a toolset over a failover
+ group reports the platform as ``failover``, and ``unknown`` means the agent's identity could
+ not be resolved. Read next to ``managed_agent.failover``, which records each transition by
+ platform."
type: "counter"
legacy_name: "-"
- name_variables: ["tool", "platform", "role", "position"]
+ name_variables: ["tool", "platform"]
- name: "common_ai.tool_calls"
description: "Number of calls to a toolset of the common.ai provider. Metric with toolset
diff --git a/uv.lock b/uv.lock
index ecf33af6816..eaa61ed383d 100644
--- a/uv.lock
+++ b/uv.lock
@@ -3159,6 +3159,9 @@ apache-hive = [
cncf-kubernetes = [
{ name = "apache-airflow-providers-cncf-kubernetes" },
]
+common-ai = [
+ { name = "apache-airflow-providers-common-ai" },
+]
common-messaging = [
{ name = "apache-airflow-providers-common-messaging" },
]
@@ -3224,6 +3227,7 @@ dev = [
{ name = "apache-airflow-providers-amazon", extra = ["aiobotocore", "sqlalchemy"] },
{ name = "apache-airflow-providers-apache-hive" },
{ name = "apache-airflow-providers-cncf-kubernetes" },
+ { name = "apache-airflow-providers-common-ai" },
{ name = "apache-airflow-providers-common-compat" },
{ name = "apache-airflow-providers-common-messaging" },
{ name = "apache-airflow-providers-common-sql", extra = ["pandas", "polars"] },
@@ -3262,6 +3266,7 @@ requires-dist = [
{ name = "apache-airflow", editable = "." },
{ name = "apache-airflow-providers-apache-hive", marker = "extra == 'apache-hive'", editable = "providers/apache/hive" },
{ name = "apache-airflow-providers-cncf-kubernetes", marker = "extra == 'cncf-kubernetes'", editable = "providers/cncf/kubernetes" },
+ { name = "apache-airflow-providers-common-ai", marker = "extra == 'common-ai'", editable = "providers/common/ai" },
{ name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" },
{ name = "apache-airflow-providers-common-messaging", marker = "extra == 'common-messaging'", editable = "providers/common/messaging" },
{ name = "apache-airflow-providers-common-sql", editable = "providers/common/sql" },
@@ -3301,7 +3306,7 @@ requires-dist = [
{ name = "watchtower", specifier = ">=3.3.1,<4" },
{ name = "xmlsec", marker = "python_full_version < '3.13' and extra == 'python3-saml'", specifier = ">=1.3.14" },
]
-provides-extras = ["aiobotocore", "cncf-kubernetes", "s3fs", "msk", "python3-saml", "apache-hive", "exasol", "fab", "ftp", "google", "imap", "microsoft-azure", "mongo", "pandas", "openlineage", "salesforce", "ssh", "standard", "common-messaging", "duckdb", "sqlalchemy"]
+provides-extras = ["aiobotocore", "cncf-kubernetes", "s3fs", "msk", "python3-saml", "apache-hive", "exasol", "fab", "ftp", "google", "imap", "microsoft-azure", "mongo", "pandas", "openlineage", "salesforce", "ssh", "standard", "common-ai", "common-messaging", "duckdb", "sqlalchemy"]
[package.metadata.requires-dev]
dev = [
@@ -3311,6 +3316,7 @@ dev = [
{ name = "apache-airflow-providers-amazon", extras = ["sqlalchemy"], editable = "providers/amazon" },
{ name = "apache-airflow-providers-apache-hive", editable = "providers/apache/hive" },
{ name = "apache-airflow-providers-cncf-kubernetes", editable = "providers/cncf/kubernetes" },
+ { name = "apache-airflow-providers-common-ai", editable = "providers/common/ai" },
{ name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" },
{ name = "apache-airflow-providers-common-messaging", editable = "providers/common/messaging" },
{ name = "apache-airflow-providers-common-sql", editable = "providers/common/sql" },
@@ -5906,6 +5912,9 @@ apache-cassandra = [
cncf-kubernetes = [
{ name = "apache-airflow-providers-cncf-kubernetes" },
]
+common-ai = [
+ { name = "apache-airflow-providers-common-ai" },
+]
common-messaging = [
{ name = "apache-airflow-providers-common-messaging" },
]
@@ -5970,6 +5979,7 @@ dev = [
{ name = "apache-airflow-providers-apache-cassandra" },
{ name = "apache-airflow-providers-apache-kafka" },
{ name = "apache-airflow-providers-cncf-kubernetes" },
+ { name = "apache-airflow-providers-common-ai" },
{ name = "apache-airflow-providers-common-compat" },
{ name = "apache-airflow-providers-common-messaging" },
{ name = "apache-airflow-providers-common-sql", extra = ["pandas", "polars"] },
@@ -6001,6 +6011,7 @@ requires-dist = [
{ name = "apache-airflow-providers-apache-beam", marker = "extra == 'apache-beam'", editable = "providers/apache/beam" },
{ name = "apache-airflow-providers-apache-cassandra", marker = "extra == 'apache-cassandra'", editable = "providers/apache/cassandra" },
{ name = "apache-airflow-providers-cncf-kubernetes", marker = "extra == 'cncf-kubernetes'", editable = "providers/cncf/kubernetes" },
+ { name = "apache-airflow-providers-common-ai", marker = "extra == 'common-ai'", editable = "providers/common/ai" },
{ name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" },
{ name = "apache-airflow-providers-common-messaging", marker = "extra == 'common-messaging'", editable = "providers/common/messaging" },
{ name = "apache-airflow-providers-common-sql", editable = "providers/common/sql" },
@@ -6097,7 +6108,7 @@ requires-dist = [
{ name = "tenacity", specifier = ">=8.3.0" },
{ name = "types-protobuf", specifier = ">=5.27.0,!=5.29.1.20250402" },
]
-provides-extras = ["apache-beam", "cncf-kubernetes", "fab", "leveldb", "oracle", "facebook", "amazon", "apache-cassandra", "microsoft-azure", "microsoft-mssql", "mongo", "mysql", "openlineage", "postgres", "presto", "salesforce", "sftp", "ssh", "trino", "http", "standard", "common-messaging"]
+provides-extras = ["apache-beam", "cncf-kubernetes", "fab", "leveldb", "oracle", "facebook", "amazon", "apache-cassandra", "microsoft-azure", "microsoft-mssql", "mongo", "mysql", "openlineage", "postgres", "presto", "salesforce", "sftp", "ssh", "trino", "http", "standard", "common-ai", "common-messaging"]
[package.metadata.requires-dev]
dev = [
@@ -6108,6 +6119,7 @@ dev = [
{ name = "apache-airflow-providers-apache-cassandra", editable = "providers/apache/cassandra" },
{ name = "apache-airflow-providers-apache-kafka", editable = "providers/apache/kafka" },
{ name = "apache-airflow-providers-cncf-kubernetes", editable = "providers/cncf/kubernetes" },
+ { name = "apache-airflow-providers-common-ai", editable = "providers/common/ai" },
{ name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" },
{ name = "apache-airflow-providers-common-messaging", editable = "providers/common/messaging" },
{ name = "apache-airflow-providers-common-sql", editable = "providers/common/sql" },