This is an automated email from the ASF dual-hosted git repository.
shahar1 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 c860159eee8 Add Azure AI Foundry Agents operators (Create, Update,
Delete, Run) (#68799)
c860159eee8 is described below
commit c860159eee8c9e86670e1418a90f530719419093
Author: Morgan <[email protected]>
AuthorDate: Sat Jul 11 04:18:01 2026 -0300
Add Azure AI Foundry Agents operators (Create, Update, Delete, Run) (#68799)
---
docs/spelling_wordlist.txt | 4 +
providers/microsoft/azure/README.rst | 2 +
providers/microsoft/azure/docs/index.rst | 2 +
.../microsoft/azure/docs/operators/ai_agents.rst | 128 ++++++
providers/microsoft/azure/provider.yaml | 52 +++
providers/microsoft/azure/pyproject.toml | 3 +
.../providers/microsoft/azure/_ai_agents.py | 72 ++++
.../providers/microsoft/azure/get_provider_info.py | 56 +++
.../providers/microsoft/azure/hooks/ai_agents.py | 454 ++++++++++++++++++++
.../microsoft/azure/operators/ai_agents.py | 398 +++++++++++++++++
.../microsoft/azure/triggers/ai_agents.py | 175 ++++++++
.../microsoft/azure/example_azure_ai_agents.py | 230 ++++++++++
.../unit/microsoft/azure/hooks/test_ai_agents.py | 475 +++++++++++++++++++++
.../microsoft/azure/operators/test_ai_agents.py | 463 ++++++++++++++++++++
.../tests/unit/microsoft/azure/test__ai_agents.py | 77 ++++
.../microsoft/azure/triggers/test_ai_agents.py | 197 +++++++++
uv.lock | 21 +
17 files changed, 2809 insertions(+)
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index 2b7c13ee660..728ab391dc7 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -24,7 +24,9 @@ AdsInsights
adsinsights
afterall
agentcore
+AgentDetails
agentic
+AgentVersionDetails
ai
aio
aiobotocore
@@ -449,6 +451,8 @@ deferrable
deidentify
DeidentifyTemplate
del
+DeleteAgentResponse
+DeleteAgentVersionResponse
delim
deliverability
deltalake
diff --git a/providers/microsoft/azure/README.rst
b/providers/microsoft/azure/README.rst
index 450b5b11bd1..63f0e97e181 100644
--- a/providers/microsoft/azure/README.rst
+++ b/providers/microsoft/azure/README.rst
@@ -56,7 +56,9 @@ PIP package Version required
``apache-airflow`` ``>=2.11.0``
``apache-airflow-providers-common-compat`` ``>=1.13.0``
``adlfs`` ``>=2023.10.0``
+``aiohttp`` ``>=3.14.0``
``azure-batch`` ``<15.0.0,>=8.0.0``
+``azure-ai-projects`` ``>=2.2.0``
``azure-cosmos`` ``>=4.6.0``
``azure-mgmt-cosmosdb`` ``>=3.0.0``
``azure-datalake-store`` ``>=0.0.45``
diff --git a/providers/microsoft/azure/docs/index.rst
b/providers/microsoft/azure/docs/index.rst
index fd9442fa473..0ad9ca9512f 100644
--- a/providers/microsoft/azure/docs/index.rst
+++ b/providers/microsoft/azure/docs/index.rst
@@ -110,7 +110,9 @@ PIP package Version required
``apache-airflow`` ``>=2.11.0``
``apache-airflow-providers-common-compat`` ``>=1.13.0``
``adlfs`` ``>=2023.10.0``
+``aiohttp`` ``>=3.14.0``
``azure-batch`` ``<15.0.0,>=8.0.0``
+``azure-ai-projects`` ``>=2.2.0``
``azure-cosmos`` ``>=4.6.0``
``azure-mgmt-cosmosdb`` ``>=3.0.0``
``azure-datalake-store`` ``>=0.0.45``
diff --git a/providers/microsoft/azure/docs/operators/ai_agents.rst
b/providers/microsoft/azure/docs/operators/ai_agents.rst
new file mode 100644
index 00000000000..a41cf6ebca4
--- /dev/null
+++ b/providers/microsoft/azure/docs/operators/ai_agents.rst
@@ -0,0 +1,128 @@
+ .. 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.
+
+
+Azure AI Foundry Hosted Agents Operators
+========================================
+
+Azure AI Foundry Hosted agents let you run a containerized agent in Microsoft
Foundry Agent Service.
+Build and push your agent image to Azure Container Registry first, then use
these operators to create
+agent versions, invoke the hosted endpoint, and clean up the deployment.
+
+Prerequisite Tasks
+^^^^^^^^^^^^^^^^^^
+
+.. include:: /operators/_partials/prerequisite_tasks.rst
+
+The operators use the ``azure_ai_agents_default`` connection by default.
Configure the Azure AI Foundry
+project endpoint in the connection host field or in the ``endpoint``
connection extra. The endpoint
+format is:
+
+.. code-block:: text
+
+ https://<aiservices-id>.services.ai.azure.com/api/projects/<project-name>
+
+The container image must be available in Azure Container Registry and must
implement one of the
+Hosted agent protocols exposed by Microsoft Foundry, such as ``responses`` or
``invocations``.
+
+.. _howto/operator:CreateAzureAIAgentOperator:
+
+CreateAzureAIAgentOperator
+--------------------------
+
+To create an Azure AI Hosted agent from a container image, use the
+:class:`~airflow.providers.microsoft.azure.operators.ai_agents.CreateAzureAIAgentOperator`.
+The operator returns the created Hosted agent version as a serializable
dictionary.
+Optional agent fields such as ``metadata``, ``description``, and
``blueprint_reference``
+can be passed through to the Foundry API.
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_ai_agents.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_azure_ai_agent_create]
+ :end-before: [END howto_operator_azure_ai_agent_create]
+
+.. _howto/operator:UpdateAzureAIAgentOperator:
+
+UpdateAzureAIAgentOperator
+--------------------------
+
+Azure AI Hosted agent updates are published as new immutable versions. To
create a new version, use the
+:class:`~airflow.providers.microsoft.azure.operators.ai_agents.UpdateAzureAIAgentOperator`.
+The operator accepts the same parameters as ``CreateAzureAIAgentOperator`` and
returns the newly
+created Hosted agent version as a serializable dictionary.
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_ai_agents.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_azure_ai_agent_update]
+ :end-before: [END howto_operator_azure_ai_agent_update]
+
+Set ``deferrable=True`` to release the worker slot while the new version
deploys; the wait then
+runs in the triggerer. The same parameter is available on
``CreateAzureAIAgentOperator``.
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_ai_agents.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_azure_ai_agent_update_deferrable]
+ :end-before: [END howto_operator_azure_ai_agent_update_deferrable]
+
+.. _howto/operator:RunAzureAIAgentOperator:
+
+RunAzureAIAgentOperator
+-----------------------
+
+To invoke an Azure AI Hosted agent, use the
+:class:`~airflow.providers.microsoft.azure.operators.ai_agents.RunAzureAIAgentOperator`.
+The ``responses`` protocol sends the payload to the agent's OpenAI-compatible
endpoint, and the
+``invocations`` protocol posts it to the agent's invocations endpoint.
+Pass ``agent_session_id`` to reuse an existing Hosted agent session with the
``invocations``
+protocol and ``user_isolation_key`` to scope endpoint resources to a specific
end user.
+The operator returns the JSON-compatible protocol response.
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_ai_agents.py
+ :language: python
+ :dedent: 8
+ :start-after: [START howto_operator_azure_ai_agent_run]
+ :end-before: [END howto_operator_azure_ai_agent_run]
+
+.. _howto/operator:DeleteAzureAIAgentOperator:
+
+DeleteAzureAIAgentOperator
+--------------------------
+
+To delete an Azure AI Hosted agent and all of its versions, use the
+:class:`~airflow.providers.microsoft.azure.operators.ai_agents.DeleteAzureAIAgentOperator`.
+Pass ``agent_version`` to delete only one version. Set ``force=True`` to also
delete active sessions
+associated with the agent or version.
+The operator returns the serialized deletion response.
+
+.. exampleinclude:: /../tests/system/microsoft/azure/example_azure_ai_agents.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_azure_ai_agent_delete]
+ :end-before: [END howto_operator_azure_ai_agent_delete]
+
+Reference
+---------
+
+For further information, look at:
+
+* `Deploy an Azure AI Hosted agent
<https://learn.microsoft.com/en-us/azure/foundry/agents/how-to/deploy-hosted-agent>`__
+* `Microsoft Foundry Agent Service documentation
<https://learn.microsoft.com/en-us/azure/foundry/agents/overview>`__
+* `Microsoft Foundry Agents REST API reference
<https://ai.azure.com/api-reference/agents/>`__
+* `Azure AI Projects Python SDK reference
<https://learn.microsoft.com/en-us/python/api/azure-ai-projects/azure.ai.projects.aiprojectclient?view=azure-python-preview>`__
diff --git a/providers/microsoft/azure/provider.yaml
b/providers/microsoft/azure/provider.yaml
index 617242aea77..fe6366023e7 100644
--- a/providers/microsoft/azure/provider.yaml
+++ b/providers/microsoft/azure/provider.yaml
@@ -176,6 +176,12 @@ integrations:
external-doc-url: https://azure.microsoft.com/
logo: /docs/integration-logos/Microsoft-Azure.png
tags: [azure]
+ - integration-name: Microsoft Azure AI Foundry Hosted Agents
+ external-doc-url: https://ai.azure.com/api-reference/agents/
+ how-to-guide:
+ - /docs/apache-airflow-providers-microsoft-azure/operators/ai_agents.rst
+ logo: /docs/integration-logos/Microsoft-Azure.png
+ tags: [azure]
- integration-name: Microsoft Azure Service Bus
external-doc-url: https://azure.microsoft.com/en-us/services/service-bus/
logo: /docs/integration-logos/Service-Bus.svg
@@ -217,6 +223,9 @@ operators:
- integration-name: Microsoft Azure Batch
python-modules:
- airflow.providers.microsoft.azure.operators.batch
+ - integration-name: Microsoft Azure AI Foundry Hosted Agents
+ python-modules:
+ - airflow.providers.microsoft.azure.operators.ai_agents
- integration-name: Microsoft Azure Container Instances
python-modules:
- airflow.providers.microsoft.azure.operators.container_instances
@@ -284,6 +293,9 @@ hooks:
- integration-name: Microsoft Azure Batch
python-modules:
- airflow.providers.microsoft.azure.hooks.batch
+ - integration-name: Microsoft Azure AI Foundry Hosted Agents
+ python-modules:
+ - airflow.providers.microsoft.azure.hooks.ai_agents
- integration-name: Microsoft Azure Data Lake Storage
python-modules:
- airflow.providers.microsoft.azure.hooks.data_lake
@@ -316,6 +328,9 @@ triggers:
- integration-name: Microsoft Azure Batch
python-modules:
- airflow.providers.microsoft.azure.triggers.batch
+ - integration-name: Microsoft Azure AI Foundry Hosted Agents
+ python-modules:
+ - airflow.providers.microsoft.azure.triggers.ai_agents
- integration-name: Microsoft Azure Compute
python-modules:
- airflow.providers.microsoft.azure.triggers.compute
@@ -402,6 +417,43 @@ connection-types:
label: Workload Identity Tenant ID
schema:
type: ["string", "null"]
+ - hook-class-name:
airflow.providers.microsoft.azure.hooks.ai_agents.AzureAIAgentsHook
+ hook-name: "Azure AI Foundry Hosted Agents"
+ connection-type: azure_ai_agents
+ ui-field-behaviour:
+ hidden-fields: ["schema", "port"]
+ relabeling:
+ host: Project Endpoint
+ login: Azure Client ID
+ password: Azure Secret
+ placeholders:
+ host:
https://<aiservices-id>.services.ai.azure.com/api/projects/<project-name>
+ login: client_id (token credentials auth)
+ password: secret (token credentials auth)
+ tenantId: tenantId (token credentials auth)
+ cloud_environment: AzurePublicCloud (default) | AzureUSGovernment |
AzureChinaCloud
+ endpoint: Overrides Project Endpoint from host
+ conn-fields:
+ tenantId:
+ label: Azure Tenant ID
+ schema:
+ type: ["string", "null"]
+ cloud_environment:
+ label: Azure Cloud Environment
+ schema:
+ type: ["string", "null"]
+ endpoint:
+ label: Project Endpoint
+ schema:
+ type: ["string", "null"]
+ managed_identity_client_id:
+ label: Managed Identity Client ID
+ schema:
+ type: ["string", "null"]
+ workload_identity_tenant_id:
+ label: Workload Identity Tenant ID
+ schema:
+ type: ["string", "null"]
- hook-class-name:
airflow.providers.microsoft.azure.hooks.compute.AzureComputeHook
hook-name: "Azure Compute"
connection-type: azure_compute
diff --git a/providers/microsoft/azure/pyproject.toml
b/providers/microsoft/azure/pyproject.toml
index 10a3c5030fd..9b2d117593b 100644
--- a/providers/microsoft/azure/pyproject.toml
+++ b/providers/microsoft/azure/pyproject.toml
@@ -62,9 +62,12 @@ dependencies = [
"apache-airflow>=2.11.0",
"apache-airflow-providers-common-compat>=1.13.0",
"adlfs>=2023.10.0",
+ "aiohttp>=3.14.0",
# azure-batch 15.x is a full rewrite of the Azure SDK (track 2) that
removes BatchServiceClient, batch_auth,
# and the other references in AzureBatchHook. Lifting the upper bound cap
needs a full hook rewrite.
"azure-batch>=8.0.0,<15.0.0",
+ # 2.2.0 adds force deletion support for Hosted agents and agent versions.
+ "azure-ai-projects>=2.2.0",
"azure-cosmos>=4.6.0",
"azure-mgmt-cosmosdb>=3.0.0",
"azure-datalake-store>=0.0.45",
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/_ai_agents.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/_ai_agents.py
new file mode 100644
index 00000000000..1936010f9b3
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/_ai_agents.py
@@ -0,0 +1,72 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+ from pydantic import JsonValue
+
+VERSION_INTERMEDIATE_STATUSES = {"creating", "deleting"}
+VERSION_SUCCESS_STATUSES = {"active"}
+VERSION_FAILURE_STATUSES = {"failed"}
+VERSION_DELETED_STATUS = "deleted"
+
+
+def _serialize_resource(resource: Any) -> JsonValue:
+ """Serialize an SDK model or response object into XCom-safe primitives."""
+ if resource is None or isinstance(resource, str | int | float | bool):
+ return resource
+ if isinstance(resource, list | tuple):
+ return [_serialize_resource(item) for item in resource]
+ if isinstance(resource, dict):
+ if not all(isinstance(key, str) for key in resource):
+ raise TypeError("Azure AI Hosted agent resources must use string
keys for XCom serialization.")
+ return {key: _serialize_resource(value) for key, value in
resource.items()}
+ if hasattr(resource, "as_dict"):
+ return _serialize_resource(resource.as_dict())
+ if hasattr(resource, "model_dump"):
+ return _serialize_resource(resource.model_dump())
+ raise TypeError(
+ f"Cannot serialize Azure AI Hosted agent resource of type
{type(resource).__name__} for XCom."
+ )
+
+
+def _get_resource_attr(resource: Any, attr: str) -> Any:
+ """Get an attribute from an SDK resource or mapping."""
+ if isinstance(resource, dict):
+ return resource.get(attr)
+ return getattr(resource, attr, None)
+
+
+def _get_version_status(version: Any) -> str:
+ """Return a normalized Hosted agent version status string."""
+ status = _get_resource_attr(version, "status")
+ if hasattr(status, "value"):
+ status = status.value
+ if status is None:
+ raise ValueError("Azure AI Hosted agent version did not include a
status.")
+ return str(status).lower()
+
+
+def _get_agent_version(version: Any) -> str:
+ """Return the version identifier from a Hosted agent version payload."""
+ agent_version = _get_resource_attr(version, "version")
+ if agent_version is None:
+ raise ValueError("Azure AI Hosted agent response did not include a
version.")
+ return str(agent_version)
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
index ee2e388980f..e6a6be7c87b 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
@@ -103,6 +103,13 @@ def get_provider_info():
"logo": "/docs/integration-logos/Microsoft-Azure.png",
"tags": ["azure"],
},
+ {
+ "integration-name": "Microsoft Azure AI Foundry Hosted Agents",
+ "external-doc-url":
"https://ai.azure.com/api-reference/agents/";,
+ "how-to-guide":
["/docs/apache-airflow-providers-microsoft-azure/operators/ai_agents.rst"],
+ "logo": "/docs/integration-logos/Microsoft-Azure.png",
+ "tags": ["azure"],
+ },
{
"integration-name": "Microsoft Azure Service Bus",
"external-doc-url":
"https://azure.microsoft.com/en-us/services/service-bus/";,
@@ -158,6 +165,10 @@ def get_provider_info():
"integration-name": "Microsoft Azure Batch",
"python-modules":
["airflow.providers.microsoft.azure.operators.batch"],
},
+ {
+ "integration-name": "Microsoft Azure AI Foundry Hosted Agents",
+ "python-modules":
["airflow.providers.microsoft.azure.operators.ai_agents"],
+ },
{
"integration-name": "Microsoft Azure Container Instances",
"python-modules":
["airflow.providers.microsoft.azure.operators.container_instances"],
@@ -246,6 +257,10 @@ def get_provider_info():
"integration-name": "Microsoft Azure Batch",
"python-modules":
["airflow.providers.microsoft.azure.hooks.batch"],
},
+ {
+ "integration-name": "Microsoft Azure AI Foundry Hosted Agents",
+ "python-modules":
["airflow.providers.microsoft.azure.hooks.ai_agents"],
+ },
{
"integration-name": "Microsoft Azure Data Lake Storage",
"python-modules":
["airflow.providers.microsoft.azure.hooks.data_lake"],
@@ -288,6 +303,10 @@ def get_provider_info():
"integration-name": "Microsoft Azure Batch",
"python-modules":
["airflow.providers.microsoft.azure.triggers.batch"],
},
+ {
+ "integration-name": "Microsoft Azure AI Foundry Hosted Agents",
+ "python-modules":
["airflow.providers.microsoft.azure.triggers.ai_agents"],
+ },
{
"integration-name": "Microsoft Azure Compute",
"python-modules":
["airflow.providers.microsoft.azure.triggers.compute"],
@@ -391,6 +410,43 @@ def get_provider_info():
},
},
},
+ {
+ "hook-class-name":
"airflow.providers.microsoft.azure.hooks.ai_agents.AzureAIAgentsHook",
+ "hook-name": "Azure AI Foundry Hosted Agents",
+ "connection-type": "azure_ai_agents",
+ "ui-field-behaviour": {
+ "hidden-fields": ["schema", "port"],
+ "relabeling": {
+ "host": "Project Endpoint",
+ "login": "Azure Client ID",
+ "password": "Azure Secret",
+ },
+ "placeholders": {
+ "host":
"https://<aiservices-id>.services.ai.azure.com/api/projects/<project-name>",
+ "login": "client_id (token credentials auth)",
+ "password": "secret (token credentials auth)",
+ "tenantId": "tenantId (token credentials auth)",
+ "cloud_environment": "AzurePublicCloud (default) |
AzureUSGovernment | AzureChinaCloud",
+ "endpoint": "Overrides Project Endpoint from host",
+ },
+ },
+ "conn-fields": {
+ "tenantId": {"label": "Azure Tenant ID", "schema":
{"type": ["string", "null"]}},
+ "cloud_environment": {
+ "label": "Azure Cloud Environment",
+ "schema": {"type": ["string", "null"]},
+ },
+ "endpoint": {"label": "Project Endpoint", "schema":
{"type": ["string", "null"]}},
+ "managed_identity_client_id": {
+ "label": "Managed Identity Client ID",
+ "schema": {"type": ["string", "null"]},
+ },
+ "workload_identity_tenant_id": {
+ "label": "Workload Identity Tenant ID",
+ "schema": {"type": ["string", "null"]},
+ },
+ },
+ },
{
"hook-class-name":
"airflow.providers.microsoft.azure.hooks.compute.AzureComputeHook",
"hook-name": "Azure Compute",
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/ai_agents.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/ai_agents.py
new file mode 100644
index 00000000000..623e629800f
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/ai_agents.py
@@ -0,0 +1,454 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, cast
+from urllib.parse import quote
+
+from azure.ai.projects import AIProjectClient
+from azure.ai.projects.aio import AIProjectClient as AsyncAIProjectClient
+from azure.core.exceptions import ResourceNotFoundError
+from azure.core.rest import HttpRequest
+from azure.identity import ClientSecretCredential
+from azure.identity.aio import ClientSecretCredential as
AsyncClientSecretCredential
+
+from airflow.providers.common.compat.connection import get_async_connection
+from airflow.providers.common.compat.sdk import BaseHook
+from airflow.providers.microsoft.azure._ai_agents import (
+ VERSION_DELETED_STATUS,
+ _get_version_status,
+ _serialize_resource,
+)
+from airflow.providers.microsoft.azure.hooks.base_azure import
_AZURE_CLOUD_ENVIRONMENTS
+from airflow.providers.microsoft.azure.utils import (
+ add_managed_identity_connection_widgets,
+ get_async_default_azure_credential,
+ get_field,
+ get_sync_default_azure_credential,
+)
+
+if TYPE_CHECKING:
+ from azure.ai.projects.models import (
+ AgentBlueprintReference,
+ AgentDefinition,
+ AgentDetails,
+ AgentVersionDetails,
+ DeleteAgentResponse,
+ DeleteAgentVersionResponse,
+ )
+ from azure.core.credentials import TokenCredential
+ from azure.core.credentials_async import AsyncTokenCredential
+ from pydantic import JsonValue
+
+ from airflow.sdk import Connection
+
+
+DEFAULT_REQUEST_TIMEOUT = 60.0
+
+
+class AzureAIAgentsHook(BaseHook):
+ """
+ Hook for Microsoft Foundry Hosted agents, backed by the
``azure-ai-projects`` SDK.
+
+ Wraps the Agents operations of :class:`azure.ai.projects.AIProjectClient`:
+
https://learn.microsoft.com/en-us/azure/foundry/agents/how-to/deploy-hosted-agent
+
+ :param azure_ai_agents_conn_id: The Azure AI Agents connection id.
+ Default is ``azure_ai_agents_default``.
+ :param endpoint: Optional Azure AI Foundry project endpoint. If not
provided, the hook uses the
+ connection host or the ``endpoint`` connection extra. Default is
``None``.
+ :param api_version: Foundry Agent Service API version. Default is ``v1``.
+ :param timeout: Optional connection/read timeout for service requests, in
seconds.
+ Default is ``60.0``.
+ """
+
+ conn_name_attr = "azure_ai_agents_conn_id"
+ default_conn_name = "azure_ai_agents_default"
+ conn_type = "azure_ai_agents"
+ hook_name = "Azure AI Foundry Hosted Agents"
+
+ def __init__(
+ self,
+ azure_ai_agents_conn_id: str = default_conn_name,
+ endpoint: str | None = None,
+ api_version: str = "v1",
+ timeout: float | None = DEFAULT_REQUEST_TIMEOUT,
+ ) -> None:
+ super().__init__()
+ self.conn_id = azure_ai_agents_conn_id
+ self.endpoint = endpoint
+ self.api_version = api_version
+ self.timeout = timeout
+ self._sync_client: AIProjectClient | None = None
+
+ @classmethod
+ @add_managed_identity_connection_widgets
+ def get_connection_form_widgets(cls) -> dict[str, Any]:
+ """Return connection widgets to add to connection form."""
+ from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
+ from flask_babel import lazy_gettext
+ from wtforms import StringField
+
+ return {
+ "tenantId": StringField(lazy_gettext("Azure Tenant ID"),
widget=BS3TextFieldWidget()),
+ "cloud_environment": StringField(
+ lazy_gettext("Azure Cloud Environment"),
widget=BS3TextFieldWidget()
+ ),
+ "endpoint": StringField(lazy_gettext("Project Endpoint"),
widget=BS3TextFieldWidget()),
+ }
+
+ @classmethod
+ def get_ui_field_behaviour(cls) -> dict[str, Any]:
+ """Return custom field behaviour."""
+ return {
+ "hidden_fields": ["schema", "port"],
+ "relabeling": {
+ "host": "Project Endpoint",
+ "login": "Azure Client ID",
+ "password": "Azure Secret",
+ },
+ "placeholders": {
+ "host":
"https://<aiservices-id>.services.ai.azure.com/api/projects/<project-name>",
+ "login": "client_id (token credentials auth)",
+ "password": "secret (token credentials auth)",
+ "tenantId": "tenantId (token credentials auth)",
+ "cloud_environment": "AzurePublicCloud (default) |
AzureUSGovernment | AzureChinaCloud",
+ "endpoint": "Overrides Project Endpoint from host",
+ },
+ }
+
+ def get_conn(self) -> AIProjectClient:
+ """Return the Azure AI Foundry project client."""
+ if self._sync_client is None:
+ conn = self.get_connection(self.conn_id)
+ # Hosted agents are a Foundry preview feature; allow_preview makes
the SDK send
+ # the required Foundry-Features request header.
+ self._sync_client = AIProjectClient(
+ endpoint=self._get_endpoint(conn),
+ credential=self._get_credential(conn),
+ api_version=self.api_version,
+ allow_preview=True,
+ connection_timeout=self.timeout,
+ read_timeout=self.timeout,
+ )
+ return self._sync_client
+
+ def _get_endpoint(self, conn: Connection) -> str:
+ endpoint = self.endpoint or conn.host or
self._get_field(conn.extra_dejson, "endpoint")
+ if not endpoint:
+ raise ValueError(
+ "Azure AI Foundry project endpoint must be provided by the
hook, connection host, "
+ "or connection extra."
+ )
+ return endpoint.rstrip("/")
+
+ def _get_authority(self, extras: dict[str, Any]) -> str:
+ cloud_env_name = self._get_field(extras, "cloud_environment") or
"AzurePublicCloud"
+ cloud_env = _AZURE_CLOUD_ENVIRONMENTS.get(
+ cloud_env_name, _AZURE_CLOUD_ENVIRONMENTS["AzurePublicCloud"]
+ )
+ return cloud_env["authority"]
+
+ @staticmethod
+ def _should_use_client_secret_credential(conn: Connection, tenant: str |
None) -> bool:
+ credential_fields = (conn.login, conn.password, tenant)
+ if any(credential_fields) and not all(credential_fields):
+ raise ValueError(
+ "Azure Client ID, Azure Secret, and Azure Tenant ID must all
be provided when "
+ "authenticating with a service principal."
+ )
+ return all(credential_fields)
+
+ def _get_credential(self, conn: Connection) -> TokenCredential:
+ extras = conn.extra_dejson
+ tenant = self._get_field(extras, "tenantId")
+
+ if self._should_use_client_secret_credential(conn, tenant):
+ self.log.info("Getting connection using specific credentials.")
+ return ClientSecretCredential(
+ client_id=cast("str", conn.login),
+ client_secret=cast("str", conn.password),
+ tenant_id=cast("str", tenant),
+ authority=self._get_authority(extras),
+ )
+
+ self.log.info("Using DefaultAzureCredential as credential.")
+ return get_sync_default_azure_credential(
+ managed_identity_client_id=self._get_field(extras,
"managed_identity_client_id"),
+ workload_identity_tenant_id=self._get_field(extras,
"workload_identity_tenant_id"),
+ )
+
+ def _get_field(self, extras: dict[str, Any], field_name: str) -> str |
None:
+ return cast(
+ "str | None",
+ get_field(
+ conn_id=self.conn_id,
+ conn_type=self.conn_type,
+ extras=extras,
+ field_name=field_name,
+ ),
+ )
+
+ def create_agent_version(
+ self,
+ agent_name: str,
+ definition: dict[str, Any] | AgentDefinition,
+ *,
+ metadata: dict[str, str] | None = None,
+ description: str | None = None,
+ blueprint_reference: AgentBlueprintReference | None = None,
+ ) -> AgentVersionDetails:
+ """
+ Create a Hosted agent version, creating the Hosted agent itself on
first use.
+
+ :param agent_name: Hosted agent name.
+ :param definition: Hosted agent container definition.
+ :param metadata: Optional metadata attached to the Hosted agent.
Default is ``None``.
+ :param description: Optional human-readable Hosted agent description.
Default is ``None``.
+ :param blueprint_reference: Optional managed identity blueprint
reference. Default is ``None``.
+ :return: Created Hosted agent version.
+ """
+ return self.get_conn().agents.create_version(
+ agent_name=agent_name,
+ definition=cast("AgentDefinition", definition),
+ metadata=metadata,
+ description=description,
+ blueprint_reference=blueprint_reference,
+ )
+
+ def get_agent_version(self, agent_name: str, agent_version: str) ->
AgentVersionDetails:
+ """
+ Get a Hosted agent version.
+
+ :param agent_name: Hosted agent name.
+ :param agent_version: Hosted agent version.
+ :return: Hosted agent version details.
+ """
+ return self.get_conn().agents.get_version(agent_name=agent_name,
agent_version=agent_version)
+
+ def get_agent(self, agent_name: str) -> AgentDetails:
+ """
+ Get a Hosted agent.
+
+ :param agent_name: Hosted agent name.
+ :return: Hosted agent details.
+ """
+ return self.get_conn().agents.get(agent_name=agent_name)
+
+ def delete_agent(self, agent_name: str, *, force: bool = False) ->
DeleteAgentResponse:
+ """
+ Delete a Hosted agent and all versions.
+
+ :param agent_name: Hosted agent name.
+ :param force: Whether to force deletion when the Hosted agent has
active sessions.
+ Default is ``False``.
+ :return: Deletion response.
+ """
+ return self.get_conn().agents.delete(agent_name=agent_name,
force=force)
+
+ def delete_agent_version(
+ self, agent_name: str, agent_version: str, *, force: bool = False
+ ) -> DeleteAgentVersionResponse:
+ """
+ Delete one Hosted agent version.
+
+ :param agent_name: Hosted agent name.
+ :param agent_version: Hosted agent version.
+ :param force: Whether to force deletion when the Hosted agent version
has active sessions.
+ Default is ``False``.
+ :return: Deletion response.
+ """
+ return self.get_conn().agents.delete_version(
+ agent_name=agent_name, agent_version=agent_version, force=force
+ )
+
+ def is_agent_version_deleted(self, agent_name: str, agent_version: str) ->
bool:
+ """
+ Return True if the Hosted agent version no longer exists or is deleted.
+
+ :param agent_name: Hosted agent name.
+ :param agent_version: Hosted agent version.
+ :return: ``True`` when the Hosted agent version no longer exists or
has deleted status.
+ """
+ try:
+ version = self.get_agent_version(agent_name=agent_name,
agent_version=agent_version)
+ except ResourceNotFoundError:
+ return True
+ return _get_version_status(version) == VERSION_DELETED_STATUS
+
+ def is_agent_deleted(self, agent_name: str) -> bool:
+ """
+ Return True if the Hosted agent no longer exists.
+
+ :param agent_name: Hosted agent name.
+ :return: ``True`` when the Hosted agent no longer exists.
+ """
+ try:
+ self.get_agent(agent_name=agent_name)
+ except ResourceNotFoundError:
+ return True
+ return False
+
+ def invoke_agent_responses(
+ self,
+ agent_name: str,
+ input_data: dict[str, Any],
+ *,
+ user_isolation_key: str | None = None,
+ ) -> dict[str, JsonValue]:
+ """
+ Invoke a Hosted agent through the OpenAI Responses protocol.
+
+ :param agent_name: Hosted agent name.
+ :param input_data: Request payload for the Responses protocol.
+ :param user_isolation_key: Optional user isolation key for endpoint
resources. Default is ``None``.
+ :return: JSON-compatible Responses protocol response.
+ """
+ request_kwargs: dict[str, Any] = {}
+ if user_isolation_key:
+ request_kwargs["extra_headers"] = {"x-ms-user-isolation-key":
user_isolation_key}
+ with self.get_conn().get_openai_client(agent_name=agent_name) as
openai_client:
+ response = openai_client.responses.create(**input_data,
**request_kwargs)
+ return cast("dict[str, JsonValue]", response.model_dump(mode="json"))
+
+ def invoke_agent_invocations(
+ self,
+ agent_name: str,
+ input_data: dict[str, Any],
+ *,
+ agent_session_id: str | None = None,
+ user_isolation_key: str | None = None,
+ ) -> JsonValue:
+ """
+ Invoke a Hosted agent through the Invocations protocol.
+
+ The SDK does not expose the Invocations protocol yet, so the request
is sent through
+ the project client pipeline, which handles authentication and preview
feature headers.
+
+ :param agent_name: Hosted agent name.
+ :param input_data: Request payload for the Invocations protocol.
+ :param agent_session_id: Optional Hosted agent session id. Default is
``None``.
+ :param user_isolation_key: Optional user isolation key for endpoint
resources. Default is ``None``.
+ :return: JSON-compatible Invocations protocol payload; its shape is
defined by the Hosted agent
+ container.
+ """
+ params = {"api-version": self.api_version}
+ if agent_session_id:
+ params["agent_session_id"] = agent_session_id
+ headers = {"x-ms-user-isolation-key": user_isolation_key} if
user_isolation_key else None
+ request = HttpRequest(
+ "POST",
+ f"/agents/{quote(agent_name,
safe='')}/endpoint/protocols/invocations",
+ params=params,
+ headers=headers,
+ json=input_data,
+ )
+ response = self.get_conn().send_request(request)
+ response.raise_for_status()
+ if not response.content:
+ return None
+ return _serialize_resource(response.json())
+
+
+class AzureAIAgentsAsyncHook(AzureAIAgentsHook):
+ """
+ Async hook for Microsoft Foundry Hosted agents.
+
+ :param azure_ai_agents_conn_id: The Azure AI Agents connection id.
+ Default is ``azure_ai_agents_default``.
+ :param endpoint: Optional Azure AI Foundry project endpoint override.
Default is ``None``.
+ :param api_version: Foundry Agent Service API version. Default is ``v1``.
+ :param timeout: Optional connection/read timeout for service requests, in
seconds.
+ Default is ``60.0``.
+ """
+
+ def __init__(
+ self,
+ azure_ai_agents_conn_id: str = AzureAIAgentsHook.default_conn_name,
+ endpoint: str | None = None,
+ api_version: str = "v1",
+ timeout: float | None = DEFAULT_REQUEST_TIMEOUT,
+ ) -> None:
+ self._async_client: AsyncAIProjectClient | None = None
+ self._async_credential: AsyncTokenCredential | None = None
+ super().__init__(
+ azure_ai_agents_conn_id=azure_ai_agents_conn_id,
+ endpoint=endpoint,
+ api_version=api_version,
+ timeout=timeout,
+ )
+
+ async def close(self) -> None:
+ """Close the async Azure AI Foundry project client and credential."""
+ try:
+ if self._async_client is not None:
+ await self._async_client.close()
+ finally:
+ self._async_client = None
+ if self._async_credential is not None:
+ try:
+ await self._async_credential.close()
+ finally:
+ self._async_credential = None
+
+ async def get_async_conn(self) -> AsyncAIProjectClient:
+ """Return the async Azure AI Foundry project client."""
+ if self._async_client is not None:
+ return self._async_client
+
+ conn = await get_async_connection(self.conn_id)
+ self._async_credential = self._get_async_credential(conn)
+ self._async_client = AsyncAIProjectClient(
+ endpoint=self._get_endpoint(conn),
+ credential=self._async_credential,
+ api_version=self.api_version,
+ allow_preview=True,
+ connection_timeout=self.timeout,
+ read_timeout=self.timeout,
+ )
+ return self._async_client
+
+ def _get_async_credential(self, conn: Connection) -> AsyncTokenCredential:
+ extras = conn.extra_dejson
+ tenant = self._get_field(extras, "tenantId")
+
+ if self._should_use_client_secret_credential(conn, tenant):
+ self.log.info("Getting connection using specific credentials.")
+ return AsyncClientSecretCredential(
+ client_id=cast("str", conn.login),
+ client_secret=cast("str", conn.password),
+ tenant_id=cast("str", tenant),
+ authority=self._get_authority(extras),
+ )
+
+ self.log.info("Using DefaultAzureCredential as credential.")
+ return get_async_default_azure_credential(
+ managed_identity_client_id=self._get_field(extras,
"managed_identity_client_id"),
+ workload_identity_tenant_id=self._get_field(extras,
"workload_identity_tenant_id"),
+ )
+
+ async def async_get_agent_version(self, agent_name: str, agent_version:
str) -> AgentVersionDetails:
+ """
+ Get a Hosted agent version asynchronously.
+
+ :param agent_name: Hosted agent name.
+ :param agent_version: Hosted agent version.
+ :return: Hosted agent version details.
+ """
+ client = await self.get_async_conn()
+ return await client.agents.get_version(agent_name=agent_name,
agent_version=agent_version)
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/ai_agents.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/ai_agents.py
new file mode 100644
index 00000000000..69040ee856e
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/ai_agents.py
@@ -0,0 +1,398 @@
+#
+# 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 time
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, cast
+
+from airflow.providers.common.compat.sdk import BaseOperator, conf
+from airflow.providers.microsoft.azure._ai_agents import (
+ VERSION_FAILURE_STATUSES,
+ VERSION_INTERMEDIATE_STATUSES,
+ VERSION_SUCCESS_STATUSES,
+ _get_agent_version,
+ _get_resource_attr,
+ _get_version_status,
+ _serialize_resource,
+)
+from airflow.providers.microsoft.azure.hooks.ai_agents import AzureAIAgentsHook
+from airflow.providers.microsoft.azure.triggers.ai_agents import
AzureAIAgentVersionTrigger
+
+if TYPE_CHECKING:
+ from azure.ai.projects.models import AgentBlueprintReference
+ from pydantic import JsonValue
+
+ from airflow.sdk import Context
+
+
+class CreateAzureAIAgentOperator(BaseOperator):
+ """
+ Create an Azure AI Hosted agent version from a container image definition.
+
+ Creating a version also creates the Hosted agent itself on first use.
+
+ :param agent_name: Hosted agent name.
+ :param definition: Hosted agent definition. Include
``container_configuration.image`` with
+ the Azure Container Registry image URI.
+ :param metadata: Optional metadata attached to the Hosted agent. Default
is ``None``.
+ :param description: Optional human-readable Hosted agent description.
Default is ``None``.
+ :param blueprint_reference: Optional managed identity blueprint reference.
Default is ``None``.
+ :param wait_for_completion: Whether to wait until the created version
reaches ``active``.
+ Default is ``True``.
+ :param poll_interval: Time in seconds between status checks. Default is
``30.0``.
+ :param timeout: Time in seconds to wait for the version to become active.
Default is ``3600``.
+ :param deferrable: Run in deferrable mode when ``wait_for_completion`` is
``True``.
+ Default is configured by ``operators.default_deferrable``.
+ :param azure_ai_agents_conn_id: Azure AI Agents connection id.
+ Default is ``azure_ai_agents_default``.
+ :param endpoint: Optional Azure AI Foundry project endpoint override.
Default is ``None``.
+ :param api_version: Foundry Agent Service API version. Default is ``v1``.
+ """
+
+ template_fields: Sequence[str] = (
+ "agent_name",
+ "definition",
+ "metadata",
+ "description",
+ "blueprint_reference",
+ "azure_ai_agents_conn_id",
+ "endpoint",
+ "api_version",
+ )
+ template_fields_renderers = {
+ "definition": "json",
+ "metadata": "json",
+ "blueprint_reference": "json",
+ }
+ ui_color = "#0678d4"
+
+ def __init__(
+ self,
+ *,
+ agent_name: str,
+ definition: dict[str, Any],
+ metadata: dict[str, str] | None = None,
+ description: str | None = None,
+ blueprint_reference: AgentBlueprintReference | None = None,
+ wait_for_completion: bool = True,
+ poll_interval: float = 30.0,
+ timeout: float = 60 * 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ azure_ai_agents_conn_id: str = AzureAIAgentsHook.default_conn_name,
+ endpoint: str | None = None,
+ api_version: str = "v1",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.agent_name = agent_name
+ self.definition = definition
+ self.metadata = metadata
+ self.description = description
+ self.blueprint_reference = blueprint_reference
+ self.wait_for_completion = wait_for_completion
+ self.poll_interval = poll_interval
+ self.timeout = timeout
+ self.deferrable = deferrable
+ self.azure_ai_agents_conn_id = azure_ai_agents_conn_id
+ self.endpoint = endpoint
+ self.api_version = api_version
+
+ @cached_property
+ def hook(self) -> AzureAIAgentsHook:
+ """Create and return an AzureAIAgentsHook."""
+ return AzureAIAgentsHook(
+ azure_ai_agents_conn_id=self.azure_ai_agents_conn_id,
+ endpoint=self.endpoint,
+ api_version=self.api_version,
+ )
+
+ def execute(self, context: Context) -> dict[str, JsonValue]:
+ """Create an Azure AI Hosted agent version and optionally wait for it
to become active."""
+ self.log.info("Creating Azure AI Hosted agent %s version.",
self.agent_name)
+ version = self.hook.create_agent_version(
+ agent_name=self.agent_name,
+ definition=self.definition,
+ metadata=self.metadata,
+ description=self.description,
+ blueprint_reference=self.blueprint_reference,
+ )
+ if not self.wait_for_completion:
+ return cast("dict[str, JsonValue]", _serialize_resource(version))
+ agent_version = _get_agent_version(version)
+ if self.deferrable:
+ self.defer(
+ timeout=self.execution_timeout,
+ trigger=AzureAIAgentVersionTrigger(
+ azure_ai_agents_conn_id=self.azure_ai_agents_conn_id,
+ endpoint=self.endpoint,
+ api_version=self.api_version,
+ agent_name=self.agent_name,
+ agent_version=agent_version,
+ timeout=self.timeout,
+ poll_interval=self.poll_interval,
+ ),
+ method_name="execute_complete",
+ )
+ return self._wait_for_version(agent_version=agent_version)
+
+ def _wait_for_version(self, *, agent_version: str) -> dict[str, JsonValue]:
+ end_time = time.monotonic() + self.timeout
+ while True:
+ version = self.hook.get_agent_version(agent_name=self.agent_name,
agent_version=agent_version)
+ status = _get_version_status(version)
+ if status in VERSION_SUCCESS_STATUSES:
+ return cast("dict[str, JsonValue]",
_serialize_resource(version))
+ if status in VERSION_FAILURE_STATUSES:
+ error = _get_resource_attr(version, "error") or "No error
details were returned"
+ raise RuntimeError(
+ f"Azure AI Hosted agent {self.agent_name} version
{agent_version} failed: {error}."
+ )
+ if status not in VERSION_INTERMEDIATE_STATUSES:
+ raise RuntimeError(
+ f"Azure AI Hosted agent {self.agent_name} version
{agent_version} "
+ f"reached unknown status {status}."
+ )
+ if time.monotonic() >= end_time:
+ raise TimeoutError(
+ f"Timeout waiting for Azure AI Hosted agent
{self.agent_name} version {agent_version}."
+ )
+
+ self.log.info(
+ "Azure AI Hosted agent %s version %s is in status %s. Sleeping
for %s seconds.",
+ self.agent_name,
+ agent_version,
+ status,
+ self.poll_interval,
+ )
+ time.sleep(self.poll_interval)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None)
-> dict[str, JsonValue]:
+ """Resume after the version trigger completes."""
+ if event is None:
+ raise RuntimeError("Trigger returned no event.")
+
+ status = event.get("status")
+ message = event.get("message", "No message returned from trigger.")
+ if status == "success":
+ if "version" not in event:
+ raise RuntimeError("Trigger success event did not include
version payload.")
+ self.log.info(message)
+ return event["version"]
+ if status == "timeout":
+ raise TimeoutError(message)
+ if status == "error":
+ raise RuntimeError(message)
+ raise ValueError(f"Unexpected trigger event status: {status!r}")
+
+
+class UpdateAzureAIAgentOperator(CreateAzureAIAgentOperator):
+ """
+ Create a new version of an existing Azure AI Hosted agent.
+
+ The Hosted agent API publishes updates as new immutable versions, so this
operator accepts
+ the same parameters as :class:`CreateAzureAIAgentOperator`.
+ """
+
+
+class RunAzureAIAgentOperator(BaseOperator):
+ """
+ Invoke an Azure AI Hosted agent.
+
+ :param agent_name: Hosted agent name.
+ :param input_data: Request payload for the selected protocol.
+ :param protocol: Hosted agent protocol to use. Supported values are
``responses`` and
+ ``invocations``. Default is ``responses``.
+ :param agent_session_id: Optional session id used by the ``invocations``
protocol.
+ Default is ``None``.
+ :param user_isolation_key: Optional per-user isolation key for
endpoint-scoped resources.
+ Default is ``None``.
+ :param azure_ai_agents_conn_id: Azure AI Agents connection id.
+ Default is ``azure_ai_agents_default``.
+ :param endpoint: Optional Azure AI Foundry project endpoint override.
Default is ``None``.
+ :param api_version: Foundry Agent Service API version. Default is ``v1``.
+ """
+
+ template_fields: Sequence[str] = (
+ "agent_name",
+ "input_data",
+ "protocol",
+ "agent_session_id",
+ "user_isolation_key",
+ "azure_ai_agents_conn_id",
+ "endpoint",
+ "api_version",
+ )
+ template_fields_renderers = {"input_data": "json"}
+ ui_color = "#0678d4"
+
+ def __init__(
+ self,
+ *,
+ agent_name: str,
+ input_data: dict[str, Any],
+ protocol: str = "responses",
+ agent_session_id: str | None = None,
+ user_isolation_key: str | None = None,
+ azure_ai_agents_conn_id: str = AzureAIAgentsHook.default_conn_name,
+ endpoint: str | None = None,
+ api_version: str = "v1",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.agent_name = agent_name
+ self.input_data = input_data
+ self.protocol = protocol
+ self.agent_session_id = agent_session_id
+ self.user_isolation_key = user_isolation_key
+ self.azure_ai_agents_conn_id = azure_ai_agents_conn_id
+ self.endpoint = endpoint
+ self.api_version = api_version
+
+ @cached_property
+ def hook(self) -> AzureAIAgentsHook:
+ """Create and return an AzureAIAgentsHook."""
+ return AzureAIAgentsHook(
+ azure_ai_agents_conn_id=self.azure_ai_agents_conn_id,
+ endpoint=self.endpoint,
+ api_version=self.api_version,
+ )
+
+ def execute(self, context: Context) -> JsonValue:
+ """Invoke an Azure AI Hosted agent and return the response payload."""
+ self.log.info("Invoking Azure AI Hosted agent %s with %s protocol.",
self.agent_name, self.protocol)
+ if self.protocol == "responses":
+ return self.hook.invoke_agent_responses(
+ agent_name=self.agent_name,
+ input_data=self.input_data,
+ user_isolation_key=self.user_isolation_key,
+ )
+ if self.protocol == "invocations":
+ return self.hook.invoke_agent_invocations(
+ agent_name=self.agent_name,
+ input_data=self.input_data,
+ agent_session_id=self.agent_session_id,
+ user_isolation_key=self.user_isolation_key,
+ )
+ raise ValueError("protocol must be either 'responses' or
'invocations'.")
+
+
+class DeleteAzureAIAgentOperator(BaseOperator):
+ """
+ Delete an Azure AI Hosted agent or one Hosted agent version.
+
+ :param agent_name: Hosted agent name.
+ :param agent_version: Optional Hosted agent version. When omitted, the
whole agent is deleted.
+ Default is ``None``.
+ :param force: Whether to cascade-delete active sessions when deleting the
agent or version.
+ Default is ``False``.
+ :param wait_for_completion: Whether to wait until the resource is deleted.
Default is ``True``.
+ :param poll_interval: Time in seconds between status checks. Default is
``30.0``.
+ :param timeout: Time in seconds to wait for deletion to complete. Default
is ``3600``.
+ :param azure_ai_agents_conn_id: Azure AI Agents connection id.
+ Default is ``azure_ai_agents_default``.
+ :param endpoint: Optional Azure AI Foundry project endpoint override.
Default is ``None``.
+ :param api_version: Foundry Agent Service API version. Default is ``v1``.
+ """
+
+ template_fields: Sequence[str] = (
+ "agent_name",
+ "agent_version",
+ "azure_ai_agents_conn_id",
+ "endpoint",
+ "api_version",
+ )
+ ui_color = "#0678d4"
+
+ def __init__(
+ self,
+ *,
+ agent_name: str,
+ agent_version: str | None = None,
+ force: bool = False,
+ wait_for_completion: bool = True,
+ poll_interval: float = 30.0,
+ timeout: float = 60 * 60,
+ azure_ai_agents_conn_id: str = AzureAIAgentsHook.default_conn_name,
+ endpoint: str | None = None,
+ api_version: str = "v1",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.agent_name = agent_name
+ self.agent_version = agent_version
+ self.force = force
+ self.wait_for_completion = wait_for_completion
+ self.poll_interval = poll_interval
+ self.timeout = timeout
+ self.azure_ai_agents_conn_id = azure_ai_agents_conn_id
+ self.endpoint = endpoint
+ self.api_version = api_version
+
+ @cached_property
+ def hook(self) -> AzureAIAgentsHook:
+ """Create and return an AzureAIAgentsHook."""
+ return AzureAIAgentsHook(
+ azure_ai_agents_conn_id=self.azure_ai_agents_conn_id,
+ endpoint=self.endpoint,
+ api_version=self.api_version,
+ )
+
+ def execute(self, context: Context) -> dict[str, JsonValue]:
+ """Delete an Azure AI Hosted agent or version and optionally wait for
deletion."""
+ delete_response: Any
+ if self.agent_version is None:
+ self.log.info("Deleting Azure AI Hosted agent %s.",
self.agent_name)
+ delete_response =
self.hook.delete_agent(agent_name=self.agent_name, force=self.force)
+ else:
+ self.log.info(
+ "Deleting Azure AI Hosted agent %s version %s.",
self.agent_name, self.agent_version
+ )
+ delete_response = self.hook.delete_agent_version(
+ agent_name=self.agent_name, agent_version=self.agent_version,
force=self.force
+ )
+
+ if self.wait_for_completion:
+ self._wait_for_deletion()
+ return cast("dict[str, JsonValue]",
_serialize_resource(delete_response))
+
+ def _wait_for_deletion(self) -> None:
+ end_time = time.monotonic() + self.timeout
+ resource_description = f"agent {self.agent_name}"
+ if self.agent_version is not None:
+ resource_description += f" version {self.agent_version}"
+ while True:
+ if self.agent_version is None:
+ is_deleted =
self.hook.is_agent_deleted(agent_name=self.agent_name)
+ else:
+ is_deleted = self.hook.is_agent_version_deleted(
+ agent_name=self.agent_name,
agent_version=self.agent_version
+ )
+ if is_deleted:
+ self.log.info("Azure AI Hosted %s was deleted.",
resource_description)
+ return
+ if time.monotonic() >= end_time:
+ raise TimeoutError(f"Timeout waiting for Azure AI Hosted
{resource_description} deletion.")
+
+ self.log.info(
+ "Azure AI Hosted %s is still retrievable. Sleeping for %s
seconds.",
+ resource_description,
+ self.poll_interval,
+ )
+ time.sleep(self.poll_interval)
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/ai_agents.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/ai_agents.py
new file mode 100644
index 00000000000..704c4af7da4
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/ai_agents.py
@@ -0,0 +1,175 @@
+#
+# 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 asyncio
+import time
+from collections.abc import AsyncIterator
+from typing import Any
+
+from airflow.providers.microsoft.azure._ai_agents import (
+ VERSION_FAILURE_STATUSES,
+ VERSION_INTERMEDIATE_STATUSES,
+ VERSION_SUCCESS_STATUSES,
+ _get_resource_attr,
+ _get_version_status,
+ _serialize_resource,
+)
+from airflow.providers.microsoft.azure.hooks.ai_agents import
AzureAIAgentsAsyncHook
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class AzureAIAgentVersionTrigger(BaseTrigger):
+ """
+ Trigger that polls an Azure AI Hosted agent version until it becomes
active.
+
+ :param azure_ai_agents_conn_id: Azure AI Agents connection id.
+ :param endpoint: Optional Azure AI Foundry project endpoint override.
Default is ``None``.
+ :param api_version: Foundry Agent Service API version.
+ :param agent_name: Hosted agent name.
+ :param agent_version: Hosted agent version to poll.
+ :param timeout: Time in seconds to wait for the version to become active.
+ :param poll_interval: Poll interval in seconds.
+ """
+
+ def __init__(
+ self,
+ *,
+ azure_ai_agents_conn_id: str,
+ endpoint: str | None = None,
+ api_version: str,
+ agent_name: str,
+ agent_version: str,
+ timeout: float,
+ poll_interval: float,
+ ) -> None:
+ super().__init__()
+ self.azure_ai_agents_conn_id = azure_ai_agents_conn_id
+ self.endpoint = endpoint
+ self.api_version = api_version
+ self.agent_name = agent_name
+ self.agent_version = agent_version
+ self.timeout = timeout
+ self.poll_interval = poll_interval
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize trigger arguments and classpath."""
+ return (
+ f"{self.__class__.__module__}.{self.__class__.__name__}",
+ {
+ "azure_ai_agents_conn_id": self.azure_ai_agents_conn_id,
+ "endpoint": self.endpoint,
+ "api_version": self.api_version,
+ "agent_name": self.agent_name,
+ "agent_version": self.agent_version,
+ "timeout": self.timeout,
+ "poll_interval": self.poll_interval,
+ },
+ )
+
+ def _build_trigger_event(self, version: Any) -> TriggerEvent | None:
+ """Build a terminal TriggerEvent for a Hosted agent version."""
+ status = _get_version_status(version)
+ if status in VERSION_INTERMEDIATE_STATUSES:
+ return None
+
+ serialized_version = _serialize_resource(version)
+ if status in VERSION_SUCCESS_STATUSES:
+ return TriggerEvent(
+ {
+ "status": "success",
+ "message": (
+ f"Azure AI Hosted agent {self.agent_name} version
{self.agent_version} is active."
+ ),
+ "version": serialized_version,
+ }
+ )
+ if status in VERSION_FAILURE_STATUSES:
+ error = _get_resource_attr(version, "error") or "No error details
were returned"
+ return TriggerEvent(
+ {
+ "status": "error",
+ "message": (
+ f"Azure AI Hosted agent {self.agent_name} version
{self.agent_version} failed: "
+ f"{error}."
+ ),
+ "version": serialized_version,
+ }
+ )
+ return TriggerEvent(
+ {
+ "status": "error",
+ "message": (
+ f"Azure AI Hosted agent {self.agent_name} version
{self.agent_version} "
+ f"reached unknown status {status}."
+ ),
+ "version": serialized_version,
+ }
+ )
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ """Poll the Hosted agent version status until terminal state or
timeout."""
+ hook = AzureAIAgentsAsyncHook(
+ azure_ai_agents_conn_id=self.azure_ai_agents_conn_id,
+ endpoint=self.endpoint,
+ api_version=self.api_version,
+ )
+
+ try:
+ end_time = time.monotonic() + self.timeout
+ while True:
+ version = await hook.async_get_agent_version(
+ agent_name=self.agent_name,
+ agent_version=self.agent_version,
+ )
+ event = self._build_trigger_event(version)
+ if event:
+ yield event
+ return
+ if time.monotonic() >= end_time:
+ yield TriggerEvent(
+ {
+ "status": "timeout",
+ "message": (
+ f"Timeout waiting for Azure AI Hosted agent
{self.agent_name} "
+ f"version {self.agent_version}."
+ ),
+ "version": {"name": self.agent_name, "version":
self.agent_version},
+ }
+ )
+ return
+
+ await asyncio.sleep(self.poll_interval)
+ except Exception as e:
+ self.log.exception(
+ "Exception while polling Azure AI Hosted agent %s version %s.",
+ self.agent_name,
+ self.agent_version,
+ )
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": (
+ f"Failed while polling Azure AI Hosted agent
{self.agent_name} "
+ f"version {self.agent_version}: {e}"
+ ),
+ "version": {"name": self.agent_name, "version":
self.agent_version},
+ }
+ )
+ finally:
+ await hook.close()
diff --git
a/providers/microsoft/azure/tests/system/microsoft/azure/example_azure_ai_agents.py
b/providers/microsoft/azure/tests/system/microsoft/azure/example_azure_ai_agents.py
new file mode 100644
index 00000000000..aa9e888e423
--- /dev/null
+++
b/providers/microsoft/azure/tests/system/microsoft/azure/example_azure_ai_agents.py
@@ -0,0 +1,230 @@
+#
+# 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.
+"""
+System test for Azure AI Foundry Hosted agent operators.
+
+Requires a real Azure AI Foundry project and a Hosted agent container image
pushed to
+Azure Container Registry:
+
+* ``AIRFLOW_CONN_AZURE_AI_AGENTS_DEFAULT``: Azure AI Agents connection URI.
+* ``AZURE_AI_AGENTS_ENDPOINT``: Azure AI Foundry project endpoint.
+* ``AZURE_AI_AGENTS_CONTAINER_IMAGE``: Hosted agent container image URI.
+* ``AZURE_AI_AGENTS_MODEL_DEPLOYMENT_NAME``: Model deployment available in the
project.
+* ``AZURE_AI_AGENTS_RUN_PROTOCOL``: Optional runtime protocol to invoke
(``responses`` or ``invocations``).
+* ``AZURE_AI_AGENTS_USE_MODEL``: Optional, defaults to ``false`` for a
deterministic smoke test.
+* ``AZURE_AI_AGENTS_USE_MOCKS``: Optional, defaults to ``true`` so the sample
agent does not need
+ external service credentials.
+"""
+
+from __future__ import annotations
+
+import atexit
+import json
+import os
+import tempfile
+from datetime import datetime
+from typing import Any
+
+from airflow.providers.microsoft.azure.operators.ai_agents import (
+ CreateAzureAIAgentOperator,
+ DeleteAzureAIAgentOperator,
+ RunAzureAIAgentOperator,
+ UpdateAzureAIAgentOperator,
+)
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import DAG, chain
+else:
+ from airflow.models.baseoperator import chain # type:
ignore[attr-defined,no-redef]
+ from airflow.models.dag import DAG # type:
ignore[attr-defined,no-redef,assignment]
+
+try:
+ from airflow.sdk import TriggerRule
+except ImportError:
+ from airflow.utils.trigger_rule import TriggerRule # type:
ignore[no-redef,attr-defined]
+
+DAG_ID = "example_azure_ai_agents"
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") or "default"
+
+
+def _get_env(name: str, default: str = "") -> str:
+ return os.environ.get(name) or os.environ.get(f"AIRFLOW_VAR_{name}",
default)
+
+
+AZURE_AI_AGENTS_CONN_ID = _get_env("AZURE_AI_AGENTS_CONN_ID",
"azure_ai_agents_default")
+AZURE_AI_AGENTS_CONN_URI = _get_env("AIRFLOW_CONN_AZURE_AI_AGENTS_DEFAULT")
+ENDPOINT = _get_env("AZURE_AI_AGENTS_ENDPOINT")
+MODEL_DEPLOYMENT_NAME = _get_env("AZURE_AI_AGENTS_MODEL_DEPLOYMENT_NAME",
"gpt-4o")
+CONTAINER_IMAGE = _get_env(
+ "AZURE_AI_AGENTS_CONTAINER_IMAGE",
"myregistry.azurecr.io/airflow-hosted-agent:latest"
+)
+AGENT_NAME = _get_env("AZURE_AI_AGENT_NAME", f"airflow-ai-agent-{ENV_ID}")
+RUN_AGENT_PROTOCOL = _get_env("AZURE_AI_AGENTS_RUN_PROTOCOL").lower()
+if RUN_AGENT_PROTOCOL not in {"", "responses", "invocations"}:
+ raise RuntimeError("AZURE_AI_AGENTS_RUN_PROTOCOL must be either
'responses' or 'invocations'.")
+HOSTED_AGENT_PROTOCOL = RUN_AGENT_PROTOCOL or "responses"
+run_agent: RunAzureAIAgentOperator | None
+
+HOSTED_AGENT_DEFINITION: dict[str, Any] = {
+ "kind": "hosted",
+ "container_configuration": {
+ "image": CONTAINER_IMAGE,
+ },
+ "cpu": "1",
+ "memory": "2Gi",
+ "protocol_versions": [
+ {"protocol": HOSTED_AGENT_PROTOCOL, "version": "1.0.0"},
+ ],
+ "environment_variables": {
+ "AZURE_AI_MODEL_DEPLOYMENT_NAME": MODEL_DEPLOYMENT_NAME,
+ "AIRFLOW_AGENT_FOUNDRY_PROJECT_ENDPOINT": ENDPOINT,
+ "AIRFLOW_AGENT_USE_MODEL": _get_env("AZURE_AI_AGENTS_USE_MODEL",
"false"),
+ "AIRFLOW_AGENT_USE_MOCKS": _get_env("AZURE_AI_AGENTS_USE_MOCKS",
"true"),
+ },
+}
+
+UPDATED_AGENT_DEFINITION: dict[str, Any] = {
+ **HOSTED_AGENT_DEFINITION,
+ "environment_variables": {
+ **HOSTED_AGENT_DEFINITION["environment_variables"],
+ "AIRFLOW_AGENT_BEHAVIOR": "troubleshoot-airflow-task-failures",
+ },
+}
+
+SMOKE_FAILURE_CONTEXT = json.dumps(
+ {
+ "dag_id": "azure_ai_agents_demo_failing_etl",
+ "run_id": "manual__azure_ai_agents_smoke",
+ "dag_file": "azure_ai_agents/demo_failing_etl.py",
+ "failed_task": {
+ "task_id": "transform",
+ "state": "failed",
+ "try_number": 1,
+ },
+ "log_excerpt": "KeyError: 'rowz'",
+ }
+)
+
+RUN_AGENT_PROMPT = (
+ "Analyze this Airflow failure context and return JSON with summary, "
+ "root_cause, and suggested_fix. Do not claim PR or Slack actions unless "
+ "the hosted agent tools actually performed them.\n\n"
+ f"{SMOKE_FAILURE_CONTEXT}"
+)
+
+RUN_AGENT_INPUT = {"input" if RUN_AGENT_PROTOCOL == "responses" else
"message": RUN_AGENT_PROMPT}
+
+
+def create_connection_file(conn_id_name: str, conn_uri: str) -> str | None:
+ if not conn_uri:
+ return None
+ with tempfile.NamedTemporaryFile("w", suffix=".env", delete=False) as
conn_file:
+ conn_file.write(f"{conn_id_name}={conn_uri}\n")
+ atexit.register(os.unlink, conn_file.name)
+ return conn_file.name
+
+
+CONN_FILE_PATH = create_connection_file(AZURE_AI_AGENTS_CONN_ID,
AZURE_AI_AGENTS_CONN_URI)
+
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ # [START howto_operator_azure_ai_agent_create]
+ create_agent = CreateAzureAIAgentOperator(
+ task_id="create_agent",
+ agent_name=AGENT_NAME,
+ definition=HOSTED_AGENT_DEFINITION,
+ poll_interval=10,
+ timeout=900,
+ deferrable=False,
+ azure_ai_agents_conn_id=AZURE_AI_AGENTS_CONN_ID,
+ endpoint=ENDPOINT,
+ )
+ # [END howto_operator_azure_ai_agent_create]
+
+ # [START howto_operator_azure_ai_agent_update]
+ update_agent = UpdateAzureAIAgentOperator(
+ task_id="update_agent",
+ agent_name=AGENT_NAME,
+ definition=UPDATED_AGENT_DEFINITION,
+ poll_interval=10,
+ timeout=900,
+ deferrable=False,
+ azure_ai_agents_conn_id=AZURE_AI_AGENTS_CONN_ID,
+ endpoint=ENDPOINT,
+ )
+ # [END howto_operator_azure_ai_agent_update]
+
+ # [START howto_operator_azure_ai_agent_update_deferrable]
+ update_agent_deferrable = UpdateAzureAIAgentOperator(
+ task_id="update_agent_deferrable",
+ agent_name=AGENT_NAME,
+ definition=UPDATED_AGENT_DEFINITION,
+ poll_interval=10,
+ timeout=900,
+ deferrable=True,
+ azure_ai_agents_conn_id=AZURE_AI_AGENTS_CONN_ID,
+ endpoint=ENDPOINT,
+ )
+ # [END howto_operator_azure_ai_agent_update_deferrable]
+
+ if RUN_AGENT_PROTOCOL:
+ # [START howto_operator_azure_ai_agent_run]
+ run_agent = RunAzureAIAgentOperator(
+ task_id="run_agent",
+ agent_name=AGENT_NAME,
+ protocol=RUN_AGENT_PROTOCOL,
+ input_data=RUN_AGENT_INPUT,
+ azure_ai_agents_conn_id=AZURE_AI_AGENTS_CONN_ID,
+ endpoint=ENDPOINT,
+ )
+ # [END howto_operator_azure_ai_agent_run]
+ else:
+ run_agent = None
+
+ # [START howto_operator_azure_ai_agent_delete]
+ delete_agent = DeleteAzureAIAgentOperator(
+ task_id="delete_agent",
+ agent_name=AGENT_NAME,
+ force=True,
+ poll_interval=10,
+ timeout=900,
+ trigger_rule=TriggerRule.ALL_DONE,
+ azure_ai_agents_conn_id=AZURE_AI_AGENTS_CONN_ID,
+ endpoint=ENDPOINT,
+ )
+ # [END howto_operator_azure_ai_agent_delete]
+
+ if run_agent is None:
+ chain(create_agent, update_agent, update_agent_deferrable,
delete_agent)
+ else:
+ chain(create_agent, update_agent, update_agent_deferrable, run_agent,
delete_agent)
+
+ from tests_common.test_utils.watcher import watcher
+
+ list(dag.tasks) >> watcher()
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+test_run = get_test_run(dag, conn_file_path=CONN_FILE_PATH)
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_ai_agents.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_ai_agents.py
new file mode 100644
index 00000000000..95d74dacae2
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_ai_agents.py
@@ -0,0 +1,475 @@
+#
+# 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 azure.core.exceptions import ResourceNotFoundError
+
+from airflow.models import Connection
+from airflow.providers.microsoft.azure.hooks.ai_agents import (
+ DEFAULT_REQUEST_TIMEOUT,
+ AzureAIAgentsAsyncHook,
+ AzureAIAgentsHook,
+)
+
+MODULE = "airflow.providers.microsoft.azure.hooks.ai_agents"
+CONN_ID = "azure_ai_agents_test"
+ENDPOINT = "https://test.services.ai.azure.com/api/projects/test-project";
+AGENT_NAME = "agent-123"
+DEFINITION = {
+ "kind": "hosted",
+ "container_configuration": {"image": "registry.azurecr.io/agent:v1"},
+ "cpu": "1",
+ "memory": "2Gi",
+ "protocol_versions": [{"protocol": "responses", "version": "1.0.0"}],
+}
+METADATA = {"team": "airflow"}
+DESCRIPTION = "Airflow hosted agent"
+BLUEPRINT_REFERENCE = {"type": "ManagedAgentIdentityBlueprint",
"blueprint_id": "blueprint-1"}
+
+
[email protected]
+def hook_with_mocked_client():
+ hook = AzureAIAgentsHook(azure_ai_agents_conn_id=CONN_ID)
+ hook._sync_client = mock.MagicMock()
+ return hook
+
+
+class TestAzureAIAgentsHook:
+ def test_connection_form_widgets(self):
+ pytest.importorskip("flask_appbuilder")
+ widgets = AzureAIAgentsHook.get_connection_form_widgets()
+
+ assert "tenantId" in widgets
+ assert "cloud_environment" in widgets
+ assert "endpoint" in widgets
+ assert "managed_identity_client_id" in widgets
+ assert "workload_identity_tenant_id" in widgets
+
+ def test_ui_field_behaviour(self):
+ behaviour = AzureAIAgentsHook.get_ui_field_behaviour()
+
+ assert behaviour["hidden_fields"] == ["schema", "port"]
+ assert behaviour["relabeling"]["host"] == "Project Endpoint"
+ assert behaviour["relabeling"]["login"] == "Azure Client ID"
+ assert behaviour["relabeling"]["password"] == "Azure Secret"
+
+ @mock.patch(f"{MODULE}.ClientSecretCredential", autospec=True)
+ def test_get_credential_uses_client_secret_credential(self,
mock_credential_cls, create_mock_connection):
+ create_mock_connection(
+ Connection(
+ conn_id=CONN_ID,
+ conn_type="azure_ai_agents",
+ host=ENDPOINT,
+ login="client-id",
+ password="client-secret",
+ extra={"tenantId": "tenant-id", "cloud_environment":
"AzureUSGovernment"},
+ )
+ )
+ hook = AzureAIAgentsHook(azure_ai_agents_conn_id=CONN_ID)
+
+ result = hook._get_credential(hook.get_connection(CONN_ID))
+
+ assert result == mock_credential_cls.return_value
+ mock_credential_cls.assert_called_once_with(
+ client_id="client-id",
+ client_secret="client-secret",
+ tenant_id="tenant-id",
+ authority="login.microsoftonline.us",
+ )
+
+ @pytest.mark.parametrize(
+ ("login", "password", "tenant"),
+ [
+ ("client-id", None, None),
+ (None, "client-secret", None),
+ (None, None, "tenant-id"),
+ ("client-id", "client-secret", None),
+ ("client-id", None, "tenant-id"),
+ (None, "client-secret", "tenant-id"),
+ ],
+ )
+ def test_get_credential_rejects_partial_service_principal(self, login,
password, tenant):
+ conn = Connection(
+ conn_id=CONN_ID,
+ conn_type="azure_ai_agents",
+ login=login,
+ password=password,
+ extra={"tenantId": tenant} if tenant else None,
+ )
+ hook = AzureAIAgentsHook(azure_ai_agents_conn_id=CONN_ID)
+
+ with pytest.raises(ValueError, match="must all be provided"):
+ hook._get_credential(conn)
+
+ @mock.patch(f"{MODULE}.get_sync_default_azure_credential", autospec=True)
+ @mock.patch(f"{MODULE}.AIProjectClient", autospec=True)
+ def test_client_uses_default_credential_and_endpoint_extra(
+ self, mock_client_cls, mock_default_credential, create_mock_connection
+ ):
+ create_mock_connection(
+ Connection(
+ conn_id=CONN_ID,
+ conn_type="azure_ai_agents",
+ extra={
+ "endpoint": ENDPOINT,
+ "managed_identity_client_id": "managed-identity-client-id",
+ "workload_identity_tenant_id":
"workload-identity-tenant-id",
+ },
+ )
+ )
+ hook = AzureAIAgentsHook(azure_ai_agents_conn_id=CONN_ID)
+
+ result = hook.get_conn()
+
+ assert result == mock_client_cls.return_value
+ mock_default_credential.assert_called_once_with(
+ managed_identity_client_id="managed-identity-client-id",
+ workload_identity_tenant_id="workload-identity-tenant-id",
+ )
+ mock_client_cls.assert_called_once_with(
+ endpoint=ENDPOINT,
+ credential=mock_default_credential.return_value,
+ api_version="v1",
+ allow_preview=True,
+ connection_timeout=DEFAULT_REQUEST_TIMEOUT,
+ read_timeout=DEFAULT_REQUEST_TIMEOUT,
+ )
+
+ @mock.patch(f"{MODULE}.get_sync_default_azure_credential", autospec=True)
+ @mock.patch(f"{MODULE}.AIProjectClient", autospec=True)
+ def test_client_uses_custom_api_version_and_timeout(
+ self, mock_client_cls, mock_default_credential, create_mock_connection
+ ):
+ create_mock_connection(Connection(conn_id=CONN_ID,
conn_type="azure_ai_agents", host=ENDPOINT))
+ hook = AzureAIAgentsHook(azure_ai_agents_conn_id=CONN_ID,
api_version="v2", timeout=10)
+
+ hook.get_conn()
+
+ mock_client_cls.assert_called_once_with(
+ endpoint=ENDPOINT,
+ credential=mock_default_credential.return_value,
+ api_version="v2",
+ allow_preview=True,
+ connection_timeout=10,
+ read_timeout=10,
+ )
+
+ @mock.patch(f"{MODULE}.get_sync_default_azure_credential", autospec=True)
+ @mock.patch(f"{MODULE}.AIProjectClient", autospec=True)
+ def test_client_is_cached(self, mock_client_cls, mock_default_credential,
create_mock_connection):
+ create_mock_connection(Connection(conn_id=CONN_ID,
conn_type="azure_ai_agents", host=ENDPOINT))
+ hook = AzureAIAgentsHook(azure_ai_agents_conn_id=CONN_ID)
+
+ assert hook.get_conn() is hook.get_conn()
+ mock_client_cls.assert_called_once()
+ mock_default_credential.assert_called_once()
+
+ def test_get_endpoint_hook_endpoint_overrides_connection_endpoint(self,
create_mock_connection):
+ create_mock_connection(Connection(conn_id=CONN_ID,
conn_type="azure_ai_agents", host=ENDPOINT))
+ endpoint_override =
"https://override.services.ai.azure.com/api/projects/project";
+ hook = AzureAIAgentsHook(azure_ai_agents_conn_id=CONN_ID,
endpoint=endpoint_override)
+
+ assert hook._get_endpoint(hook.get_connection(CONN_ID)) ==
endpoint_override
+
+ def test_get_endpoint_raises_when_endpoint_missing(self,
create_mock_connection):
+ create_mock_connection(Connection(conn_id=CONN_ID,
conn_type="azure_ai_agents"))
+ hook = AzureAIAgentsHook(azure_ai_agents_conn_id=CONN_ID)
+
+ with pytest.raises(ValueError, match="Azure AI Foundry project
endpoint must be provided"):
+ hook._get_endpoint(hook.get_connection(CONN_ID))
+
+ def test_create_agent_version(self, hook_with_mocked_client):
+ hook = hook_with_mocked_client
+
+ result = hook.create_agent_version(
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ metadata=METADATA,
+ description=DESCRIPTION,
+ blueprint_reference=BLUEPRINT_REFERENCE,
+ )
+
+ hook.get_conn().agents.create_version.assert_called_once_with(
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ metadata=METADATA,
+ description=DESCRIPTION,
+ blueprint_reference=BLUEPRINT_REFERENCE,
+ )
+ assert result == hook.get_conn().agents.create_version.return_value
+
+ def test_get_agent_version(self, hook_with_mocked_client):
+ hook = hook_with_mocked_client
+
+ result = hook.get_agent_version(agent_name=AGENT_NAME,
agent_version="1")
+
+
hook.get_conn().agents.get_version.assert_called_once_with(agent_name=AGENT_NAME,
agent_version="1")
+ assert result == hook.get_conn().agents.get_version.return_value
+
+ def test_get_agent(self, hook_with_mocked_client):
+ hook = hook_with_mocked_client
+
+ result = hook.get_agent(agent_name=AGENT_NAME)
+
+
hook.get_conn().agents.get.assert_called_once_with(agent_name=AGENT_NAME)
+ assert result == hook.get_conn().agents.get.return_value
+
+ @pytest.mark.parametrize("force", [False, True])
+ def test_delete_agent(self, hook_with_mocked_client, force):
+ hook = hook_with_mocked_client
+
+ result = hook.delete_agent(agent_name=AGENT_NAME, force=force)
+
+
hook.get_conn().agents.delete.assert_called_once_with(agent_name=AGENT_NAME,
force=force)
+ assert result == hook.get_conn().agents.delete.return_value
+
+ @pytest.mark.parametrize("force", [False, True])
+ def test_delete_agent_version(self, hook_with_mocked_client, force):
+ hook = hook_with_mocked_client
+
+ result = hook.delete_agent_version(agent_name=AGENT_NAME,
agent_version="2", force=force)
+
+ hook.get_conn().agents.delete_version.assert_called_once_with(
+ agent_name=AGENT_NAME, agent_version="2", force=force
+ )
+ assert result == hook.get_conn().agents.delete_version.return_value
+
+ def test_is_agent_version_deleted_when_resource_does_not_exist(self,
hook_with_mocked_client):
+ hook = hook_with_mocked_client
+ hook.get_conn().agents.get_version.side_effect =
ResourceNotFoundError("not found")
+
+ assert hook.is_agent_version_deleted(agent_name=AGENT_NAME,
agent_version="1") is True
+
+ def test_is_agent_version_deleted_when_status_is_deleted(self,
hook_with_mocked_client):
+ hook = hook_with_mocked_client
+ hook.get_conn().agents.get_version.return_value = {"status": "deleted"}
+
+ assert hook.is_agent_version_deleted(agent_name=AGENT_NAME,
agent_version="1") is True
+
+ def test_is_agent_version_deleted_when_resource_exists(self,
hook_with_mocked_client):
+ hook = hook_with_mocked_client
+ hook.get_conn().agents.get_version.return_value = {"status": "active"}
+
+ assert hook.is_agent_version_deleted(agent_name=AGENT_NAME,
agent_version="1") is False
+
+ def test_is_agent_deleted_when_resource_does_not_exist(self,
hook_with_mocked_client):
+ hook = hook_with_mocked_client
+ hook.get_conn().agents.get.side_effect = ResourceNotFoundError("not
found")
+
+ assert hook.is_agent_deleted(agent_name=AGENT_NAME) is True
+
+ def test_is_agent_deleted_when_resource_exists(self,
hook_with_mocked_client):
+ hook = hook_with_mocked_client
+ hook.get_conn().agents.get.return_value = {"name": AGENT_NAME}
+
+ assert hook.is_agent_deleted(agent_name=AGENT_NAME) is False
+
+ def test_invoke_agent_responses(self, hook_with_mocked_client):
+ hook = hook_with_mocked_client
+ openai_client_manager = hook.get_conn().get_openai_client.return_value
+ openai_client = openai_client_manager.__enter__.return_value
+ openai_client.responses.create.return_value.model_dump.return_value =
{"output_text": "hello"}
+
+ result = hook.invoke_agent_responses(
+ agent_name=AGENT_NAME,
+ input_data={"input": "hello"},
+ user_isolation_key="user-key",
+ )
+
+ assert result == {"output_text": "hello"}
+
hook.get_conn().get_openai_client.assert_called_once_with(agent_name=AGENT_NAME)
+ openai_client.responses.create.assert_called_once_with(
+ input="hello",
+ extra_headers={"x-ms-user-isolation-key": "user-key"},
+ )
+ openai_client_manager.__exit__.assert_called_once_with(None, None,
None)
+
+ def test_invoke_agent_responses_without_isolation_key(self,
hook_with_mocked_client):
+ hook = hook_with_mocked_client
+ openai_client_manager = hook.get_conn().get_openai_client.return_value
+ openai_client = openai_client_manager.__enter__.return_value
+
+ hook.invoke_agent_responses(agent_name=AGENT_NAME,
input_data={"input": "hello"})
+
+ openai_client.responses.create.assert_called_once_with(input="hello")
+ openai_client_manager.__exit__.assert_called_once_with(None, None,
None)
+
+ @mock.patch(f"{MODULE}.HttpRequest", autospec=True)
+ def test_invoke_agent_invocations(self, mock_request_cls,
hook_with_mocked_client):
+ hook = hook_with_mocked_client
+ response = hook.get_conn().send_request.return_value
+ response.content = b'{"result": "done"}'
+ response.json.return_value = {"result": "done"}
+
+ result = hook.invoke_agent_invocations(
+ agent_name=AGENT_NAME,
+ input_data={"message": "hello"},
+ agent_session_id="session-1",
+ user_isolation_key="user-key",
+ )
+
+ assert result == {"result": "done"}
+ mock_request_cls.assert_called_once_with(
+ "POST",
+ f"/agents/{AGENT_NAME}/endpoint/protocols/invocations",
+ params={"api-version": "v1", "agent_session_id": "session-1"},
+ headers={"x-ms-user-isolation-key": "user-key"},
+ json={"message": "hello"},
+ )
+
hook.get_conn().send_request.assert_called_once_with(mock_request_cls.return_value)
+ response.raise_for_status.assert_called_once_with()
+
+ @mock.patch(f"{MODULE}.HttpRequest", autospec=True)
+ def test_invoke_agent_invocations_quotes_resource_id(self,
mock_request_cls, hook_with_mocked_client):
+ hook = hook_with_mocked_client
+ hook.get_conn().send_request.return_value.content = b""
+
+ result = hook.invoke_agent_invocations(
+ agent_name="agent/name with spaces", input_data={"message":
"hello"}
+ )
+
+ assert result is None
+ assert (
+ mock_request_cls.call_args.args[1]
+ ==
"/agents/agent%2Fname%20with%20spaces/endpoint/protocols/invocations"
+ )
+ assert mock_request_cls.call_args.kwargs["headers"] is None
+
+ @mock.patch(f"{MODULE}.HttpRequest", autospec=True)
+ def test_invoke_agent_invocations_rejects_non_json_response(
+ self, mock_request_cls, hook_with_mocked_client
+ ):
+ hook = hook_with_mocked_client
+ response = hook.get_conn().send_request.return_value
+ response.content = b'{"result": "done"}'
+ response.json.return_value = {"result": object()}
+
+ with pytest.raises(TypeError, match="Cannot serialize.*object.*for
XCom"):
+ hook.invoke_agent_invocations(agent_name=AGENT_NAME,
input_data={"message": "hello"})
+
+ response.raise_for_status.assert_called_once_with()
+
+
+class TestAzureAIAgentsAsyncHook:
+ pytestmark = pytest.mark.asyncio
+
+ @mock.patch(f"{MODULE}.get_async_default_azure_credential", autospec=True)
+ @mock.patch(f"{MODULE}.AsyncAIProjectClient", autospec=True)
+ @mock.patch(f"{MODULE}.get_async_connection")
+ async def test_get_async_conn(self, mock_get_connection, mock_client_cls,
mock_default_credential):
+ mock_get_connection.return_value = Connection(
+ conn_id=CONN_ID, conn_type="azure_ai_agents", host=ENDPOINT
+ )
+ hook = AzureAIAgentsAsyncHook(azure_ai_agents_conn_id=CONN_ID)
+
+ result = await hook.get_async_conn()
+
+ assert result == mock_client_cls.return_value
+ mock_client_cls.assert_called_once_with(
+ endpoint=ENDPOINT,
+ credential=mock_default_credential.return_value,
+ api_version="v1",
+ allow_preview=True,
+ connection_timeout=DEFAULT_REQUEST_TIMEOUT,
+ read_timeout=DEFAULT_REQUEST_TIMEOUT,
+ )
+
+ assert await hook.get_async_conn() is result
+ mock_client_cls.assert_called_once()
+
+ @mock.patch(f"{MODULE}.AsyncClientSecretCredential", autospec=True)
+ @mock.patch(f"{MODULE}.AsyncAIProjectClient", autospec=True)
+ @mock.patch(f"{MODULE}.get_async_connection")
+ async def test_get_async_conn_uses_client_secret_credential(
+ self, mock_get_connection, mock_client_cls, mock_credential_cls
+ ):
+ mock_get_connection.return_value = Connection(
+ conn_id=CONN_ID,
+ conn_type="azure_ai_agents",
+ host=ENDPOINT,
+ login="client-id",
+ password="client-secret",
+ extra={"tenantId": "tenant-id"},
+ )
+ hook = AzureAIAgentsAsyncHook(azure_ai_agents_conn_id=CONN_ID)
+
+ await hook.get_async_conn()
+
+ mock_credential_cls.assert_called_once_with(
+ client_id="client-id",
+ client_secret="client-secret",
+ tenant_id="tenant-id",
+ authority="login.microsoftonline.com",
+ )
+
+ async def
test_get_async_credential_rejects_partial_service_principal(self):
+ conn = Connection(
+ conn_id=CONN_ID,
+ conn_type="azure_ai_agents",
+ login="client-id",
+ )
+ hook = AzureAIAgentsAsyncHook(azure_ai_agents_conn_id=CONN_ID)
+
+ with pytest.raises(ValueError, match="must all be provided"):
+ hook._get_async_credential(conn)
+
+ async def test_async_get_agent_version(self):
+ hook = AzureAIAgentsAsyncHook(azure_ai_agents_conn_id=CONN_ID)
+ client = mock.MagicMock()
+ client.agents.get_version = mock.AsyncMock(return_value={"version":
"1"})
+ hook.get_async_conn = mock.AsyncMock(return_value=client)
+
+ result = await hook.async_get_agent_version(agent_name=AGENT_NAME,
agent_version="1")
+
+ assert result == {"version": "1"}
+
client.agents.get_version.assert_awaited_once_with(agent_name=AGENT_NAME,
agent_version="1")
+
+ async def test_close(self):
+ hook = AzureAIAgentsAsyncHook(azure_ai_agents_conn_id=CONN_ID)
+ client = mock.MagicMock(spec=["close"])
+ client.close = mock.AsyncMock()
+ credential = mock.MagicMock(spec=["close"])
+ credential.close = mock.AsyncMock()
+ hook._async_client = client
+ hook._async_credential = credential
+
+ await hook.close()
+
+ client.close.assert_awaited_once_with()
+ credential.close.assert_awaited_once_with()
+ assert hook._async_client is None
+ assert hook._async_credential is None
+
+ async def test_close_closes_credential_when_client_close_fails(self):
+ hook = AzureAIAgentsAsyncHook(azure_ai_agents_conn_id=CONN_ID)
+ client = mock.MagicMock(spec=["close"])
+ client.close = mock.AsyncMock(side_effect=RuntimeError("client close
failed"))
+ credential = mock.MagicMock(spec=["close"])
+ credential.close = mock.AsyncMock()
+ hook._async_client = client
+ hook._async_credential = credential
+
+ with pytest.raises(RuntimeError, match="client close failed"):
+ await hook.close()
+
+ credential.close.assert_awaited_once_with()
+ assert hook._async_client is None
+ assert hook._async_credential is None
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_ai_agents.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_ai_agents.py
new file mode 100644
index 00000000000..3612acccb1d
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_ai_agents.py
@@ -0,0 +1,463 @@
+#
+# 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.compat.sdk import TaskDeferred
+from airflow.providers.microsoft.azure.hooks.ai_agents import AzureAIAgentsHook
+from airflow.providers.microsoft.azure.operators.ai_agents import (
+ CreateAzureAIAgentOperator,
+ DeleteAzureAIAgentOperator,
+ RunAzureAIAgentOperator,
+ UpdateAzureAIAgentOperator,
+)
+from airflow.providers.microsoft.azure.triggers.ai_agents import
AzureAIAgentVersionTrigger
+
+MODULE = "airflow.providers.microsoft.azure.operators.ai_agents"
+CONN_ID = "azure_ai_agents_test"
+ENDPOINT = "https://test.services.ai.azure.com/api/projects/test-project";
+AGENT_NAME = "agent-123"
+DEFINITION = {
+ "kind": "hosted",
+ "container_configuration": {"image": "registry.azurecr.io/agent:v1"},
+ "cpu": "1",
+ "memory": "2Gi",
+ "protocol_versions": [{"protocol": "responses", "version": "1.0.0"}],
+}
+METADATA = {"team": "airflow"}
+DESCRIPTION = "Airflow hosted agent"
+BLUEPRINT_REFERENCE = {"type": "ManagedAgentIdentityBlueprint",
"blueprint_id": "blueprint-1"}
+
+
+class TestCreateAzureAIAgentOperator:
+ @mock.patch.object(AzureAIAgentsHook, "get_agent_version", autospec=True)
+ @mock.patch.object(AzureAIAgentsHook, "create_agent_version",
autospec=True)
+ def test_execute_waits_until_version_is_active(self, mock_create_version,
mock_get_version):
+ mock_create_version.return_value = {"name": AGENT_NAME, "version":
"1", "status": "creating"}
+ mock_get_version.return_value = {"name": AGENT_NAME, "version": "1",
"status": "active"}
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ metadata=METADATA,
+ description=DESCRIPTION,
+ blueprint_reference=BLUEPRINT_REFERENCE,
+ poll_interval=0,
+ azure_ai_agents_conn_id=CONN_ID,
+ endpoint=ENDPOINT,
+ )
+
+ result = operator.execute(context={})
+
+ assert result == {"name": AGENT_NAME, "version": "1", "status":
"active"}
+ mock_create_version.assert_called_once_with(
+ mock.ANY,
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ metadata=METADATA,
+ description=DESCRIPTION,
+ blueprint_reference=BLUEPRINT_REFERENCE,
+ )
+ mock_get_version.assert_called_once_with(mock.ANY,
agent_name=AGENT_NAME, agent_version="1")
+ assert operator.hook.conn_id == CONN_ID
+ assert operator.hook.endpoint == ENDPOINT
+
+ @mock.patch.object(AzureAIAgentsHook, "create_agent_version",
autospec=True)
+ def test_execute_without_waiting(self, mock_create_version):
+ mock_create_version.return_value = {"name": AGENT_NAME, "version":
"1", "status": "creating"}
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ wait_for_completion=False,
+ )
+
+ result = operator.execute(context={})
+
+ assert result == {"name": AGENT_NAME, "version": "1", "status":
"creating"}
+
+ @mock.patch.object(AzureAIAgentsHook, "create_agent_version",
autospec=True)
+ def test_execute_without_waiting_does_not_require_version(self,
mock_create_version):
+ mock_create_version.return_value = {"name": AGENT_NAME, "status":
"creating"}
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ wait_for_completion=False,
+ )
+
+ result = operator.execute(context={})
+
+ assert result == {"name": AGENT_NAME, "status": "creating"}
+
+ @mock.patch.object(AzureAIAgentsHook, "create_agent_version",
autospec=True)
+ def test_execute_raises_when_wait_response_has_no_version(self,
mock_create_version):
+ mock_create_version.return_value = {"name": AGENT_NAME, "status":
"creating"}
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ )
+
+ with pytest.raises(ValueError, match="did not include a version"):
+ operator.execute(context={})
+
+ @mock.patch.object(AzureAIAgentsHook, "get_agent_version", autospec=True)
+ @mock.patch.object(AzureAIAgentsHook, "create_agent_version",
autospec=True)
+ def test_execute_raises_when_version_fails(self, mock_create_version,
mock_get_version):
+ mock_create_version.return_value = {"name": AGENT_NAME, "version":
"1", "status": "creating"}
+ mock_get_version.return_value = {"name": AGENT_NAME, "version": "1",
"status": "failed"}
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent", agent_name=AGENT_NAME,
definition=DEFINITION
+ )
+
+ with pytest.raises(
+ RuntimeError,
+ match=f"Azure AI Hosted agent {AGENT_NAME} version 1 failed: No
error details were returned.",
+ ):
+ operator.execute(context={})
+
+ @mock.patch.object(AzureAIAgentsHook, "get_agent_version", autospec=True)
+ @mock.patch.object(AzureAIAgentsHook, "create_agent_version",
autospec=True)
+ def test_execute_raises_when_version_reaches_unknown_status(self,
mock_create_version, mock_get_version):
+ mock_create_version.return_value = {"name": AGENT_NAME, "version":
"1", "status": "creating"}
+ mock_get_version.return_value = {"name": AGENT_NAME, "version": "1",
"status": "paused"}
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent", agent_name=AGENT_NAME,
definition=DEFINITION
+ )
+
+ with pytest.raises(RuntimeError, match="unknown status paused"):
+ operator.execute(context={})
+
+ @mock.patch(f"{MODULE}.time.sleep", autospec=True)
+ @mock.patch(f"{MODULE}.time.monotonic", autospec=True, side_effect=[0, 0])
+ @mock.patch.object(AzureAIAgentsHook, "get_agent_version", autospec=True)
+ @mock.patch.object(AzureAIAgentsHook, "create_agent_version",
autospec=True)
+ def test_execute_raises_when_version_times_out(
+ self, mock_create_version, mock_get_version, mock_monotonic, mock_sleep
+ ):
+ mock_create_version.return_value = {"name": AGENT_NAME, "version":
"1", "status": "creating"}
+ mock_get_version.return_value = {"name": AGENT_NAME, "version": "1",
"status": "creating"}
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ timeout=0,
+ poll_interval=0,
+ )
+
+ with pytest.raises(TimeoutError, match="Timeout waiting for Azure AI
Hosted agent"):
+ operator.execute(context={})
+
+ mock_sleep.assert_not_called()
+
+ @mock.patch.object(AzureAIAgentsHook, "create_agent_version",
autospec=True)
+ def test_execute_defers(self, mock_create_version):
+ mock_create_version.return_value = {"name": AGENT_NAME, "version":
"1", "status": "creating"}
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ deferrable=True,
+ timeout=10,
+ poll_interval=2,
+ azure_ai_agents_conn_id=CONN_ID,
+ endpoint=ENDPOINT,
+ api_version="v2",
+ )
+
+ with pytest.raises(TaskDeferred) as exc:
+ operator.execute(context={})
+
+ trigger = exc.value.trigger
+ assert isinstance(trigger, AzureAIAgentVersionTrigger)
+ assert trigger.azure_ai_agents_conn_id == CONN_ID
+ assert trigger.endpoint == ENDPOINT
+ assert trigger.api_version == "v2"
+ assert trigger.agent_name == AGENT_NAME
+ assert trigger.agent_version == "1"
+ assert trigger.timeout == 10
+ assert trigger.poll_interval == 2
+
+ def test_execute_complete_success(self):
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ )
+
+ result = operator.execute_complete(
+ context={},
+ event={
+ "status": "success",
+ "message": "active",
+ "version": {"name": AGENT_NAME, "version": "1", "status":
"active"},
+ },
+ )
+
+ assert result == {"name": AGENT_NAME, "version": "1", "status":
"active"}
+
+ def test_execute_complete_requires_version_payload(self):
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ )
+
+ with pytest.raises(RuntimeError, match="did not include version
payload"):
+ operator.execute_complete(context={}, event={"status": "success",
"message": "active"})
+
+
+class TestUpdateAzureAIAgentOperator:
+ @mock.patch.object(AzureAIAgentsHook, "get_agent_version", autospec=True)
+ @mock.patch.object(AzureAIAgentsHook, "create_agent_version",
autospec=True)
+ def test_execute_creates_new_version(self, mock_create_version,
mock_get_version):
+ mock_create_version.return_value = {"name": AGENT_NAME, "version":
"2", "status": "creating"}
+ mock_get_version.return_value = {"name": AGENT_NAME, "version": "2",
"status": "active"}
+ operator = UpdateAzureAIAgentOperator(
+ task_id="update_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ poll_interval=0,
+ )
+
+ result = operator.execute(context={})
+
+ assert result == {"name": AGENT_NAME, "version": "2", "status":
"active"}
+ mock_create_version.assert_called_once_with(
+ mock.ANY,
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ metadata=None,
+ description=None,
+ blueprint_reference=None,
+ )
+ mock_get_version.assert_called_once_with(mock.ANY,
agent_name=AGENT_NAME, agent_version="2")
+
+ @mock.patch.object(AzureAIAgentsHook, "create_agent_version",
autospec=True)
+ def test_execute_defers(self, mock_create_version):
+ mock_create_version.return_value = {"name": AGENT_NAME, "version":
"2", "status": "creating"}
+ operator = UpdateAzureAIAgentOperator(
+ task_id="update_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ deferrable=True,
+ timeout=10,
+ poll_interval=2,
+ )
+
+ with pytest.raises(TaskDeferred) as exc:
+ operator.execute(context={})
+
+ trigger = exc.value.trigger
+ assert isinstance(trigger, AzureAIAgentVersionTrigger)
+ assert trigger.agent_version == "2"
+
+
+class TestRunAzureAIAgentOperator:
+ @mock.patch.object(AzureAIAgentsHook, "invoke_agent_responses",
autospec=True)
+ def test_execute_responses_protocol(self, mock_invoke):
+ mock_invoke.return_value = {"output_text": "hello"}
+ operator = RunAzureAIAgentOperator(
+ task_id="run_agent",
+ agent_name=AGENT_NAME,
+ protocol="responses",
+ input_data={"input": "hello"},
+ user_isolation_key="user-key",
+ )
+
+ result = operator.execute(context={})
+
+ assert result == {"output_text": "hello"}
+ mock_invoke.assert_called_once_with(
+ mock.ANY,
+ agent_name=AGENT_NAME,
+ input_data={"input": "hello"},
+ user_isolation_key="user-key",
+ )
+
+ @mock.patch.object(AzureAIAgentsHook, "invoke_agent_invocations",
autospec=True)
+ def test_execute_invocations_protocol(self, mock_invoke):
+ mock_invoke.return_value = {"result": "done"}
+ operator = RunAzureAIAgentOperator(
+ task_id="run_agent",
+ agent_name=AGENT_NAME,
+ protocol="invocations",
+ input_data={"message": "hello"},
+ agent_session_id="session-1",
+ user_isolation_key="user-key",
+ )
+
+ result = operator.execute(context={})
+
+ assert result == {"result": "done"}
+ mock_invoke.assert_called_once_with(
+ mock.ANY,
+ agent_name=AGENT_NAME,
+ input_data={"message": "hello"},
+ agent_session_id="session-1",
+ user_isolation_key="user-key",
+ )
+
+ def test_execute_raises_for_unknown_protocol(self):
+ operator = RunAzureAIAgentOperator(
+ task_id="run_agent",
+ agent_name=AGENT_NAME,
+ protocol="unknown",
+ input_data={"message": "hello"},
+ )
+
+ with pytest.raises(ValueError, match="protocol must be either"):
+ operator.execute(context={})
+
+
+class TestDeleteAzureAIAgentOperator:
+ @mock.patch.object(AzureAIAgentsHook, "is_agent_deleted", autospec=True)
+ @mock.patch.object(AzureAIAgentsHook, "delete_agent", autospec=True)
+ def test_execute_deletes_agent(self, mock_delete_agent, mock_is_deleted):
+ mock_delete_agent.return_value = {"object": "agent.deleted", "name":
AGENT_NAME, "deleted": True}
+ mock_is_deleted.return_value = True
+ operator = DeleteAzureAIAgentOperator(
+ task_id="delete_agent",
+ agent_name=AGENT_NAME,
+ poll_interval=0,
+ )
+
+ with mock.patch.object(operator.log, "info") as mock_log_info:
+ result = operator.execute(context={})
+
+ assert result == {"object": "agent.deleted", "name": AGENT_NAME,
"deleted": True}
+ mock_delete_agent.assert_called_once_with(mock.ANY,
agent_name=AGENT_NAME, force=False)
+ mock_is_deleted.assert_called_once_with(mock.ANY,
agent_name=AGENT_NAME)
+ mock_log_info.assert_any_call("Azure AI Hosted %s was deleted.",
f"agent {AGENT_NAME}")
+
+ @mock.patch.object(AzureAIAgentsHook, "delete_agent", autospec=True)
+ def test_execute_deletes_agent_with_force(self, mock_delete_agent):
+ mock_delete_agent.return_value = {"object": "agent.deleted", "name":
AGENT_NAME, "deleted": True}
+ operator = DeleteAzureAIAgentOperator(
+ task_id="delete_agent",
+ agent_name=AGENT_NAME,
+ force=True,
+ wait_for_completion=False,
+ )
+
+ result = operator.execute(context={})
+
+ assert result == {"object": "agent.deleted", "name": AGENT_NAME,
"deleted": True}
+ mock_delete_agent.assert_called_once_with(mock.ANY,
agent_name=AGENT_NAME, force=True)
+
+ @mock.patch.object(AzureAIAgentsHook, "is_agent_version_deleted",
autospec=True)
+ @mock.patch.object(AzureAIAgentsHook, "delete_agent_version",
autospec=True)
+ def test_execute_deletes_agent_version(self, mock_delete_version,
mock_is_deleted):
+ mock_delete_version.return_value = {"object": "agent.version.deleted",
"deleted": True}
+ mock_is_deleted.return_value = True
+ operator = DeleteAzureAIAgentOperator(
+ task_id="delete_agent",
+ agent_name=AGENT_NAME,
+ agent_version="2",
+ poll_interval=0,
+ )
+
+ with mock.patch.object(operator.log, "info") as mock_log_info:
+ result = operator.execute(context={})
+
+ assert result == {"object": "agent.version.deleted", "deleted": True}
+ mock_delete_version.assert_called_once_with(
+ mock.ANY, agent_name=AGENT_NAME, agent_version="2", force=False
+ )
+ mock_is_deleted.assert_called_once_with(mock.ANY,
agent_name=AGENT_NAME, agent_version="2")
+ mock_log_info.assert_any_call("Azure AI Hosted %s was deleted.",
f"agent {AGENT_NAME} version 2")
+
+ @mock.patch.object(AzureAIAgentsHook, "delete_agent_version",
autospec=True)
+ def test_execute_deletes_agent_version_with_force(self,
mock_delete_version):
+ mock_delete_version.return_value = {"object": "agent.version.deleted",
"deleted": True}
+ operator = DeleteAzureAIAgentOperator(
+ task_id="delete_agent_version",
+ agent_name=AGENT_NAME,
+ agent_version="2",
+ force=True,
+ wait_for_completion=False,
+ )
+
+ result = operator.execute(context={})
+
+ assert result == {"object": "agent.version.deleted", "deleted": True}
+ mock_delete_version.assert_called_once_with(
+ mock.ANY, agent_name=AGENT_NAME, agent_version="2", force=True
+ )
+
+ @mock.patch(f"{MODULE}.time.sleep", autospec=True)
+ @mock.patch.object(AzureAIAgentsHook, "is_agent_deleted", autospec=True)
+ @mock.patch.object(AzureAIAgentsHook, "delete_agent", autospec=True)
+ def test_execute_polls_until_agent_is_deleted(self, mock_delete_agent,
mock_is_deleted, mock_sleep):
+ mock_delete_agent.return_value = {"object": "agent.deleted", "name":
AGENT_NAME, "deleted": True}
+ mock_is_deleted.side_effect = [False, False, True]
+ operator = DeleteAzureAIAgentOperator(
+ task_id="delete_agent",
+ agent_name=AGENT_NAME,
+ poll_interval=1,
+ )
+
+ operator.execute(context={})
+
+ assert mock_is_deleted.call_count == 3
+ assert mock_sleep.call_count == 2
+ mock_sleep.assert_called_with(1)
+
+ @mock.patch(f"{MODULE}.time.sleep", autospec=True)
+ @mock.patch(f"{MODULE}.time.monotonic", autospec=True, side_effect=[0, 0])
+ @mock.patch.object(AzureAIAgentsHook, "is_agent_deleted", autospec=True)
+ @mock.patch.object(AzureAIAgentsHook, "delete_agent", autospec=True)
+ def test_execute_raises_when_deletion_times_out(
+ self, mock_delete_agent, mock_is_deleted, mock_monotonic, mock_sleep
+ ):
+ mock_is_deleted.return_value = False
+ operator = DeleteAzureAIAgentOperator(
+ task_id="delete_agent",
+ agent_name=AGENT_NAME,
+ timeout=0,
+ poll_interval=0,
+ )
+
+ with pytest.raises(TimeoutError, match="Timeout waiting for Azure AI
Hosted agent"):
+ operator.execute(context={})
+
+ mock_sleep.assert_not_called()
+
+
[email protected](
+ ("event", "exception", "message"),
+ [
+ (None, RuntimeError, "Trigger returned no event."),
+ ({"status": "timeout", "message": "timed out"}, TimeoutError, "timed
out"),
+ ({"status": "error", "message": "failed"}, RuntimeError, "failed"),
+ ({"status": "unknown"}, ValueError, "Unexpected trigger event status"),
+ ],
+)
+def test_execute_complete_event_errors(event, exception, message):
+ operator = CreateAzureAIAgentOperator(
+ task_id="create_agent",
+ agent_name=AGENT_NAME,
+ definition=DEFINITION,
+ )
+
+ with pytest.raises(exception, match=message):
+ operator.execute_complete(context={}, event=event)
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/test__ai_agents.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/test__ai_agents.py
new file mode 100644
index 00000000000..d20a2859490
--- /dev/null
+++ b/providers/microsoft/azure/tests/unit/microsoft/azure/test__ai_agents.py
@@ -0,0 +1,77 @@
+#
+# 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 pytest
+
+from airflow.providers.microsoft.azure._ai_agents import (
+ _get_agent_version,
+ _get_version_status,
+ _serialize_resource,
+)
+
+AGENT_NAME = "agent-123"
+
+
+def test_get_version_status_raises_when_status_missing():
+ with pytest.raises(ValueError, match="did not include a status"):
+ _get_version_status({})
+
+
+def test_get_version_status_normalizes_enum_values():
+ class FakeEnum:
+ value = "Active"
+
+ assert _get_version_status({"status": FakeEnum()}) == "active"
+
+
+def test_get_agent_version_raises_when_version_missing():
+ with pytest.raises(ValueError, match="did not include a version"):
+ _get_agent_version({})
+
+
+def test_serialize_resource_uses_as_dict_serializer():
+ class SdkModel:
+ def as_dict(self):
+ return {"name": AGENT_NAME, "versions": [{"version": "1"}]}
+
+ assert _serialize_resource(SdkModel()) == {"name": AGENT_NAME, "versions":
[{"version": "1"}]}
+
+
+def test_serialize_resource_uses_model_dump_serializer():
+ class PydanticModel:
+ def model_dump(self):
+ return {"output_text": "hello"}
+
+ assert _serialize_resource(PydanticModel()) == {"output_text": "hello"}
+
+
+def test_serialize_resource_rejects_arbitrary_objects():
+ class ArbitraryObject:
+ value = "class-value"
+
+ def __init__(self):
+ self.value = "instance-value"
+
+ with pytest.raises(TypeError, match="Cannot
serialize.*ArbitraryObject.*for XCom"):
+ _serialize_resource(ArbitraryObject())
+
+
+def test_serialize_resource_rejects_non_string_mapping_keys():
+ with pytest.raises(TypeError, match="must use string keys"):
+ _serialize_resource({1: "value"})
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_ai_agents.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_ai_agents.py
new file mode 100644
index 00000000000..5db1e9c152b
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_ai_agents.py
@@ -0,0 +1,197 @@
+#
+# 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.microsoft.azure.hooks.ai_agents import
AzureAIAgentsAsyncHook
+from airflow.providers.microsoft.azure.triggers.ai_agents import
AzureAIAgentVersionTrigger
+
+MODULE = "airflow.providers.microsoft.azure.triggers.ai_agents"
+CONN_ID = "azure_ai_agents_test"
+ENDPOINT = "https://test.services.ai.azure.com/api/projects/test-project";
+AGENT_NAME = "agent-123"
+
+
+def build_trigger(**overrides):
+ kwargs = {
+ "azure_ai_agents_conn_id": CONN_ID,
+ "endpoint": ENDPOINT,
+ "api_version": "v1",
+ "agent_name": AGENT_NAME,
+ "agent_version": "1",
+ "timeout": 10,
+ "poll_interval": 2,
+ **overrides,
+ }
+ return AzureAIAgentVersionTrigger(**kwargs)
+
+
+async def get_trigger_event(trigger):
+ generator = trigger.run()
+ event = await anext(generator)
+ await generator.aclose()
+ return event
+
+
+class TestAzureAIAgentVersionTrigger:
+ def test_serialize(self):
+ trigger = build_trigger(api_version="v2")
+
+ classpath, kwargs = trigger.serialize()
+
+ assert classpath ==
"airflow.providers.microsoft.azure.triggers.ai_agents.AzureAIAgentVersionTrigger"
+ assert kwargs == {
+ "azure_ai_agents_conn_id": CONN_ID,
+ "endpoint": ENDPOINT,
+ "api_version": "v2",
+ "agent_name": AGENT_NAME,
+ "agent_version": "1",
+ "timeout": 10,
+ "poll_interval": 2,
+ }
+
+ def test_serialize_default_endpoint(self):
+ trigger = AzureAIAgentVersionTrigger(
+ azure_ai_agents_conn_id=CONN_ID,
+ api_version="v1",
+ agent_name=AGENT_NAME,
+ agent_version="1",
+ timeout=10,
+ poll_interval=2,
+ )
+
+ _, kwargs = trigger.serialize()
+
+ assert kwargs["endpoint"] is None
+
+ @mock.patch(f"{MODULE}._serialize_resource", autospec=True)
+ def test_build_trigger_event_does_not_serialize_intermediate_version(self,
mock_serialize):
+ trigger = build_trigger()
+
+ assert trigger._build_trigger_event({"status": "creating"}) is None
+ mock_serialize.assert_not_called()
+
+ @pytest.mark.asyncio
+ @mock.patch.object(AzureAIAgentsAsyncHook, "close", autospec=True)
+ @mock.patch.object(AzureAIAgentsAsyncHook, "async_get_agent_version",
autospec=True)
+ async def test_run_success(self, mock_get_version, mock_close):
+ mock_get_version.return_value = {"name": AGENT_NAME, "version": "1",
"status": "active"}
+ trigger = build_trigger()
+
+ event = await get_trigger_event(trigger)
+
+ assert event.payload == {
+ "status": "success",
+ "message": f"Azure AI Hosted agent {AGENT_NAME} version 1 is
active.",
+ "version": {"name": AGENT_NAME, "version": "1", "status":
"active"},
+ }
+ mock_close.assert_awaited_once()
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.asyncio.sleep", autospec=True)
+ @mock.patch.object(AzureAIAgentsAsyncHook, "async_get_agent_version",
autospec=True)
+ async def test_run_polls_until_success(self, mock_get_version, mock_sleep):
+ creating = {"name": AGENT_NAME, "version": "1", "status": "creating"}
+ active = {"name": AGENT_NAME, "version": "1", "status": "active"}
+ mock_get_version.side_effect = [creating, creating, active]
+ trigger = build_trigger(poll_interval=1)
+
+ event = await get_trigger_event(trigger)
+
+ assert event.payload["status"] == "success"
+ assert mock_get_version.call_count == 3
+ assert mock_sleep.call_count == 2
+ mock_sleep.assert_awaited_with(1)
+
+ @pytest.mark.asyncio
+ @mock.patch.object(AzureAIAgentsAsyncHook, "async_get_agent_version",
autospec=True)
+ async def test_run_failure(self, mock_get_version):
+ mock_get_version.return_value = {
+ "name": AGENT_NAME,
+ "version": "1",
+ "status": "failed",
+ "error": {"message": "boom"},
+ }
+ trigger = build_trigger()
+
+ event = await get_trigger_event(trigger)
+
+ assert event.payload["status"] == "error"
+ assert "boom" in event.payload["message"]
+ assert event.payload["version"]["status"] == "failed"
+
+ @pytest.mark.asyncio
+ @mock.patch.object(AzureAIAgentsAsyncHook, "async_get_agent_version",
autospec=True)
+ async def test_run_failure_without_error_details(self, mock_get_version):
+ mock_get_version.return_value = {"name": AGENT_NAME, "version": "1",
"status": "failed"}
+ trigger = build_trigger()
+
+ event = await get_trigger_event(trigger)
+
+ assert event.payload["status"] == "error"
+ assert event.payload["message"] == (
+ f"Azure AI Hosted agent {AGENT_NAME} version 1 failed: No error
details were returned."
+ )
+ assert event.payload["version"]["status"] == "failed"
+
+ @pytest.mark.asyncio
+ @mock.patch.object(AzureAIAgentsAsyncHook, "async_get_agent_version",
autospec=True)
+ async def test_run_unknown_status(self, mock_get_version):
+ mock_get_version.return_value = {"name": AGENT_NAME, "version": "1",
"status": "paused"}
+ trigger = build_trigger()
+
+ event = await get_trigger_event(trigger)
+
+ assert event.payload["status"] == "error"
+ assert "unknown status paused" in event.payload["message"]
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.asyncio.sleep", autospec=True)
+ @mock.patch(f"{MODULE}.time.monotonic", autospec=True, side_effect=[0, 0])
+ @mock.patch.object(AzureAIAgentsAsyncHook, "async_get_agent_version",
autospec=True)
+ async def test_run_timeout(self, mock_get_version, mock_monotonic,
mock_sleep):
+ mock_get_version.return_value = {"name": AGENT_NAME, "version": "1",
"status": "creating"}
+ trigger = build_trigger(timeout=0)
+
+ event = await get_trigger_event(trigger)
+
+ assert event.payload == {
+ "status": "timeout",
+ "message": f"Timeout waiting for Azure AI Hosted agent
{AGENT_NAME} version 1.",
+ "version": {"name": AGENT_NAME, "version": "1"},
+ }
+ mock_sleep.assert_not_called()
+
+ @pytest.mark.asyncio
+ @mock.patch.object(AzureAIAgentsAsyncHook, "close", autospec=True)
+ @mock.patch.object(AzureAIAgentsAsyncHook, "async_get_agent_version",
autospec=True)
+ async def test_run_exception(self, mock_get_version, mock_close):
+ mock_get_version.side_effect = RuntimeError("boom")
+ trigger = build_trigger()
+
+ event = await get_trigger_event(trigger)
+
+ assert event.payload == {
+ "status": "error",
+ "message": f"Failed while polling Azure AI Hosted agent
{AGENT_NAME} version 1: boom",
+ "version": {"name": AGENT_NAME, "version": "1"},
+ }
+ mock_close.assert_awaited_once()
diff --git a/uv.lock b/uv.lock
index 7f8c32196b5..2b08c898b6c 100644
--- a/uv.lock
+++ b/uv.lock
@@ -6225,8 +6225,10 @@ version = "13.5.1"
source = { editable = "providers/microsoft/azure" }
dependencies = [
{ name = "adlfs" },
+ { name = "aiohttp" },
{ name = "apache-airflow" },
{ name = "apache-airflow-providers-common-compat" },
+ { name = "azure-ai-projects" },
{ name = "azure-batch" },
{ name = "azure-cosmos" },
{ name = "azure-datalake-store" },
@@ -6299,6 +6301,7 @@ docs = [
[package.metadata]
requires-dist = [
{ name = "adlfs", specifier = ">=2023.10.0" },
+ { name = "aiohttp", specifier = ">=3.14.0" },
{ name = "apache-airflow", editable = "." },
{ name = "apache-airflow-providers-amazon", marker = "extra == 'amazon'",
editable = "providers/amazon" },
{ name = "apache-airflow-providers-common-compat", editable =
"providers/common/compat" },
@@ -6307,6 +6310,7 @@ requires-dist = [
{ name = "apache-airflow-providers-openlineage", marker = "extra ==
'openlineage'", editable = "providers/openlineage" },
{ name = "apache-airflow-providers-oracle", marker = "extra == 'oracle'",
editable = "providers/oracle" },
{ name = "apache-airflow-providers-sftp", marker = "extra == 'sftp'",
editable = "providers/sftp" },
+ { name = "azure-ai-projects", specifier = ">=2.2.0" },
{ name = "azure-batch", specifier = ">=8.0.0,<15.0.0" },
{ name = "azure-cosmos", specifier = ">=4.6.0" },
{ name = "azure-datalake-store", specifier = ">=0.0.45" },
@@ -9478,6 +9482,23 @@ wheels = [
{ url =
"https://files.pythonhosted.org/packages/ef/c3/f30a7a63e664acc7c2545ca0491b6ce8264536e0e5cad3965f1d1b91e960/aws_xray_sdk-2.15.0-py2.py3-none-any.whl";,
hash =
"sha256:422d62ad7d52e373eebb90b642eb1bb24657afe03b22a8df4a8b2e5108e278a3", size
= 103228, upload-time = "2025-10-29T21:00:24.12Z" },
]
+[[package]]
+name = "azure-ai-projects"
+version = "2.2.0"
+source = { registry = "https://pypi.org/simple"; }
+dependencies = [
+ { name = "azure-core" },
+ { name = "azure-identity" },
+ { name = "azure-storage-blob" },
+ { name = "isodate" },
+ { name = "openai" },
+ { name = "typing-extensions" },
+]
+sdist = { url =
"https://files.pythonhosted.org/packages/b9/d8/24342aea74fe75b0a8378b6eff665b9c1cb63f855c1a96f70a0095e474a2/azure_ai_projects-2.2.0.tar.gz";,
hash =
"sha256:58ee31bb031cfb004051145c545294bb0d32de679c670c312ef384845bd72cef", size
= 668496, upload-time = "2026-05-30T00:20:59.099Z" }
+wheels = [
+ { url =
"https://files.pythonhosted.org/packages/40/cf/90f27a2b48c9b748f84194b07e565f900e7f0ce0500da9b9f067dca599d3/azure_ai_projects-2.2.0-py3-none-any.whl";,
hash =
"sha256:8f89bdaca4df1bd479d3bd2bd0f19a0905d60be6d17b84a69e8fabd82eac5906", size
= 344307, upload-time = "2026-05-30T00:21:00.672Z" },
+]
+
[[package]]
name = "azure-batch"
version = "14.2.0"
