Skip to content

Commit e9c552c

Browse files
kaxilkaxil
authored
Add a vendor-neutral managed-agent hook contract to Common AI (#73532)
1 parent 4d304d0 commit e9c552c

27 files changed

Lines changed: 2327 additions & 633 deletions

File tree

‎dev/breeze/tests/test_selective_checks.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1130,7 +1130,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
11301130
("providers/amazon/src/airflow/providers/amazon/provider.yaml",),
11311131
{
11321132
"selected-providers-list-as-string": "amazon apache.hive cncf.kubernetes "
1133-
"common.compat common.messaging common.sql databricks duckdb exasol ftp google http imap "
1133+
"common.ai common.compat common.messaging common.sql databricks duckdb exasol ftp google http imap "
11341134
"microsoft.azure "
11351135
"mongo mysql openlineage postgres salesforce ssh teradata",
11361136
"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):
11571157
{
11581158
"description": "amazon...google",
11591159
"test_types": "Providers[amazon] Providers[apache.hive,cncf.kubernetes,"
1160-
"common.compat,common.messaging,common.sql,databricks,duckdb,exasol,ftp,http,imap,"
1160+
"common.ai,common.compat,common.messaging,common.sql,databricks,duckdb,exasol,ftp,http,imap,"
11611161
"microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] "
11621162
"Providers[google]",
11631163
}
@@ -1203,7 +1203,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str):
12031203
("providers/amazon/src/airflow/providers/amazon/file.py",),
12041204
{
12051205
"selected-providers-list-as-string": "amazon apache.hive cncf.kubernetes "
1206-
"common.compat common.messaging common.sql databricks duckdb exasol ftp google http imap "
1206+
"common.ai common.compat common.messaging common.sql databricks duckdb exasol ftp google http imap "
12071207
"microsoft.azure "
12081208
"mongo mysql openlineage postgres salesforce ssh teradata",
12091209
"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):
12271227
{
12281228
"description": "amazon...google",
12291229
"test_types": "Providers[amazon] Providers[apache.hive,cncf.kubernetes,"
1230-
"common.compat,common.messaging,common.sql,databricks,duckdb,exasol,ftp,http,imap,"
1230+
"common.ai,common.compat,common.messaging,common.sql,databricks,duckdb,exasol,ftp,http,imap,"
12311231
"microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] "
12321232
"Providers[google]",
12331233
}
@@ -2797,7 +2797,7 @@ def test_expected_output_push(
27972797
),
27982798
{
27992799
"selected-providers-list-as-string": "amazon apache.beam apache.cassandra apache.kafka "
2800-
"cncf.kubernetes common.compat common.messaging common.sql databricks "
2800+
"cncf.kubernetes common.ai common.compat common.messaging common.sql databricks "
28012801
"facebook google hashicorp http microsoft.azure microsoft.mssql mongo mysql "
28022802
"openlineage oracle postgres presto salesforce samba sftp ssh standard trino",
28032803
"all-python-versions": f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']",
@@ -2809,7 +2809,7 @@ def test_expected_output_push(
28092809
"skip-providers-tests": "false",
28102810
"docs-build": "true",
28112811
"docs-list-as-string": "apache-airflow helm-chart amazon apache.beam apache.cassandra "
2812-
"apache.kafka cncf.kubernetes common.compat common.messaging common.sql databricks facebook google hashicorp http microsoft.azure "
2812+
"apache.kafka cncf.kubernetes common.ai common.compat common.messaging common.sql databricks facebook google hashicorp http microsoft.azure "
28132813
"microsoft.mssql mongo mysql openlineage oracle postgres "
28142814
"presto salesforce samba sftp ssh standard trino",
28152815
"skip-prek-hooks": (
@@ -2832,7 +2832,7 @@ def test_expected_output_push(
28322832
{
28332833
"description": "amazon...standard",
28342834
"test_types": "Providers[amazon] Providers[apache.beam,apache.cassandra,"
2835-
"apache.kafka,cncf.kubernetes,common.compat,common.messaging,common.sql,databricks,facebook,"
2835+
"apache.kafka,cncf.kubernetes,common.ai,common.compat,common.messaging,common.sql,databricks,facebook,"
28362836
"hashicorp,http,microsoft.azure,microsoft.mssql,mongo,mysql,"
28372837
"openlineage,oracle,postgres,presto,salesforce,samba,sftp,ssh,trino] "
28382838
"Providers[google] "
@@ -3115,7 +3115,7 @@ def test_upgrade_to_newer_dependencies(
31153115
("providers/google/docs/some_file.rst",),
31163116
{
31173117
"docs-list-as-string": "amazon apache.beam apache.cassandra apache.kafka "
3118-
"cncf.kubernetes common.compat common.messaging common.sql databricks facebook google hashicorp http "
3118+
"cncf.kubernetes common.ai common.compat common.messaging common.sql databricks facebook google hashicorp http "
31193119
"microsoft.azure microsoft.mssql mongo mysql openlineage oracle "
31203120
"postgres presto salesforce samba sftp ssh standard trino",
31213121
},

‎providers/amazon/docs/index.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,7 @@ Dependent package
149149
======================================================================================================================== ====================
150150
`apache-airflow-providers-apache-hive <https://airflow.apache.org/docs/apache-airflow-providers-apache-hive>`_ ``apache.hive``
151151
`apache-airflow-providers-cncf-kubernetes <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes>`_ ``cncf.kubernetes``
152+
`apache-airflow-providers-common-ai <https://airflow.apache.org/docs/apache-airflow-providers-common-ai>`_ ``common.ai``
152153
`apache-airflow-providers-common-messaging <https://airflow.apache.org/docs/apache-airflow-providers-common-messaging>`_ ``common.messaging``
153154
`apache-airflow-providers-duckdb <https://airflow.apache.org/docs/apache-airflow-providers-duckdb>`_ ``duckdb``
154155
`apache-airflow-providers-exasol <https://airflow.apache.org/docs/apache-airflow-providers-exasol>`_ ``exasol``
@@ -194,6 +195,7 @@ Extra Dependencies
194195
``salesforce`` ``apache-airflow-providers-salesforce``
195196
``ssh`` ``apache-airflow-providers-ssh``
196197
``standard`` ``apache-airflow-providers-standard``
198+
``common.ai`` ``apache-airflow-providers-common-ai>=0.10.0``
197199
``common.messaging`` ``apache-airflow-providers-common-messaging>=2.0.0``
198200
``duckdb`` ``apache-airflow-providers-duckdb>=0.1.0``
199201
``sqlalchemy`` ``sqlalchemy>=1.4.54,!=2.1.0``

‎providers/amazon/docs/operators/bedrock.rst‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,46 @@ through ``botocore_config``.
110110
:start-after: [START howto_operator_bedrock_invoke_agent_runtime]
111111
:end-before: [END howto_operator_bedrock_invoke_agent_runtime]
112112

113+
.. _howto/hook:BedrockAgentCoreHook:managed-agent:
114+
115+
Consult an Amazon Bedrock AgentCore Runtime from a Common AI agent
116+
==================================================================
117+
118+
To let an agent running under Common AI's ``AgentOperator`` consult an AgentCore Runtime as one
119+
of its tools, use
120+
:class:`~airflow.providers.amazon.aws.hooks.bedrock.BedrockAgentCoreHook`.
121+
It implements the Common AI managed-agent contract, so ``hook.agent(runtime_arn)`` can be passed
122+
to a ``ManagedAgentToolset`` or combined with agents on other clouds in a
123+
``FailoverManagedAgentClient``. The hook needs the ``common.ai`` extra of this provider, which
124+
installs ``apache-airflow-providers-common-ai`` and therefore requires Airflow 3.
125+
126+
The agent is the runtime ARN and the session is AgentCore's ``runtimeSessionId``, which the
127+
service requires to be 33 to 256 characters long. A prompt is
128+
sent as ``{"prompt": ...}`` in an ``application/json`` payload; the container behind the
129+
runtime defines the response shape, so pass ``text_key`` in ``vendor_options`` when you know which
130+
field holds the answer. The hook disables botocore retries unless the connection or caller configured them,
131+
because an invocation's effects are unknown, and raises a terminal error for non-JSON responses
132+
or bodies over the size limit rather than returning something misleading.
133+
134+
.. code-block:: python
135+
136+
from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentCoreHook
137+
from airflow.providers.common.ai.toolsets import ManagedAgentToolset
138+
139+
claims = BedrockAgentCoreHook(aws_conn_id="aws_default", region_name="us-east-1").agent(
140+
"arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/claims"
141+
)
142+
toolset = ManagedAgentToolset(
143+
claims,
144+
tool_name="ask_claims_agent",
145+
description="Reviews an insurance claim and returns a coverage determination.",
146+
vendor_options={"text_key": "determination"},
147+
)
148+
149+
For long-running invocations run as a pipeline step in their own right, use
150+
:class:`~airflow.providers.amazon.aws.operators.bedrock.BedrockInvokeAgentRuntimeOperator`
151+
instead; a toolset call blocks a worker thread for the duration of the call.
152+
113153
.. _howto/operator:BedrockDeleteAgentRuntimeOperator:
114154

115155
Delete an Amazon Bedrock AgentCore Runtime

‎providers/amazon/pyproject.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,9 @@ dependencies = [
149149
"standard" = [
150150
"apache-airflow-providers-standard"
151151
]
152+
"common.ai" = [
153+
"apache-airflow-providers-common-ai>=0.10.0", # use next version
154+
]
152155
"common.messaging" = [
153156
"apache-airflow-providers-common-messaging>=2.0.0"
154157
]
@@ -168,6 +171,7 @@ dev = [
168171
"apache-airflow-devel-common",
169172
"apache-airflow-providers-apache-hive",
170173
"apache-airflow-providers-cncf-kubernetes",
174+
"apache-airflow-providers-common-ai",
171175
"apache-airflow-providers-common-compat",
172176
"apache-airflow-providers-common-messaging",
173177
"apache-airflow-providers-common-sql",

0 commit comments

Comments
 (0)