diff --git a/providers/common/ai/docs/operators/llm.rst b/providers/common/ai/docs/operators/llm.rst index 426a2573206fc..efabdc3a106ad 100644 --- a/providers/common/ai/docs/operators/llm.rst +++ b/providers/common/ai/docs/operators/llm.rst @@ -205,6 +205,13 @@ approving with ``allow_modifications=True``, and set a deadline with :start-after: [START howto_operator_llm_approval] :end-before: [END howto_operator_llm_approval] +A pending review is only visible on the Required Actions page. Pass +``approval_notifiers`` to tell the reviewers about it through any Airflow +notifier (Slack, email, ...), the way +:class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with +``notifiers``. The notifiers run once the review is open, and a notifier that +raises fails the task before it starts waiting. + Parameters ---------- @@ -225,6 +232,8 @@ Parameters means wait indefinitely. Default ``None``. - ``allow_modifications``: If ``True``, the reviewer can edit the output before approving. Default ``False``. +- ``approval_notifiers``: Notifier, or list of notifiers, called once the review + is open. Default ``None``. Logging ------- diff --git a/providers/common/ai/docs/operators/llm_branch.rst b/providers/common/ai/docs/operators/llm_branch.rst index 1ef8c557abcf2..403a32a2d14f5 100644 --- a/providers/common/ai/docs/operators/llm_branch.rst +++ b/providers/common/ai/docs/operators/llm_branch.rst @@ -106,8 +106,8 @@ returning a ``Sequence[UserContent]`` raises ``TypeError`` before the LLM call. Apart from ``fail_on_reject``, which is specific to this operator, -``approval_timeout`` and the rest of the approval behaviour are inherited -from :ref:`LLMOperator `. +``approval_timeout``, ``approval_notifiers``, and the rest of the approval +behaviour are inherited from :ref:`LLMOperator `. How It Works ------------ @@ -139,6 +139,8 @@ Parameters means wait indefinitely. Default ``None``. - ``allow_modifications``: If ``True``, the reviewer can change the chosen branch(es) before approving. Default ``False``. +- ``approval_notifiers``: Notifier, or list of notifiers, called once the review + is open. Default ``None``. - ``fail_on_reject``: If ``True``, a rejected review fails the task instead of skipping the downstream tasks. Generally discouraged. Default ``False``. diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst b/providers/common/ai/docs/operators/llm_file_analysis.rst index ac688ebfaf415..c6e19a56157f7 100644 --- a/providers/common/ai/docs/operators/llm_file_analysis.rst +++ b/providers/common/ai/docs/operators/llm_file_analysis.rst @@ -159,8 +159,8 @@ Parameters downstream consumer needs the dict shape. This operator also inherits ``LLMOperator``'s HITL review parameters -- -``require_approval``, ``approval_timeout``, and ``allow_modifications`` -- see -:doc:`llm` for details. +``require_approval``, ``approval_timeout``, ``allow_modifications``, and +``approval_notifiers`` -- see :doc:`llm` for details. Supported Formats ----------------- diff --git a/providers/common/ai/docs/operators/llm_schema_compare.rst b/providers/common/ai/docs/operators/llm_schema_compare.rst index 768d96280be55..6df9b1ff522f8 100644 --- a/providers/common/ai/docs/operators/llm_schema_compare.rst +++ b/providers/common/ai/docs/operators/llm_schema_compare.rst @@ -132,8 +132,9 @@ expire, fails the task: returning a ``Sequence[UserContent]`` raises ``TypeError`` before the LLM call. -``approval_timeout``, ``allow_modifications``, and the rest of the approval -behaviour are inherited from :ref:`LLMOperator `. +``approval_timeout``, ``allow_modifications``, ``approval_notifiers``, and the +rest of the approval behaviour are inherited from +:ref:`LLMOperator `. Conditional ETL Based on Schema Compatibility ---------------------------------------------- @@ -193,6 +194,8 @@ Parameters means wait indefinitely. Default ``None``. - ``allow_modifications``: If ``True``, the reviewer can edit the result JSON before approving. Default ``False``. +- ``approval_notifiers``: Notifier, or list of notifiers, called once the review + is open. Default ``None``. Logging ------- diff --git a/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py index cccff32bfd150..45b00bd7628b9 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py +++ b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py @@ -34,6 +34,9 @@ log = logging.getLogger(__name__) if TYPE_CHECKING: + from collections.abc import Sequence + + from airflow.providers.common.compat.sdk import BaseNotifier from airflow.sdk import Context @@ -42,6 +45,7 @@ class DeferForApprovalProtocol(Protocol): approval_timeout: timedelta | None allow_modifications: bool + approval_notifiers: Sequence[BaseNotifier] prompt: str task_id: str defer: Any @@ -62,11 +66,17 @@ class LLMApprovalMixin: before approving. The (possibly modified) output is then returned as the task result. + ``approval_notifiers`` are called once the review is open, so a reviewer + learns about it without watching the Required Actions page, the way + :class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with + ``notifiers``. + Operators that use this mixin must set the following attributes: - ``require_approval`` (``bool``) - ``allow_modifications`` (``bool``) - ``approval_timeout`` (``timedelta | None``) + - ``approval_notifiers`` (``Sequence[BaseNotifier]``) - ``prompt`` (``str``) """ @@ -165,6 +175,9 @@ def defer_for_approval( params=hitl_params, ) + for notifier in self.approval_notifiers: + notifier(context) + if AIRFLOW_V_3_3_PLUS: # New core (3.3+): park the task in AWAITING_INPUT -- no trigger, no triggerer. The # task is resumed by the Core API response handler or the scheduler timeout sweep. diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py index 36cd616c1b386..b9ca61d268c4e 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py @@ -29,7 +29,7 @@ from airflow.providers.common.ai.mixins.approval import LLMApprovalMixin from airflow.providers.common.ai.utils.logging import log_run_summary from airflow.providers.common.ai.utils.output_type import rehydrate_pydantic_output -from airflow.providers.common.compat.sdk import BaseOperator +from airflow.providers.common.compat.sdk import BaseNotifier, BaseOperator try: # New enough cores register an operator's declared ``output_type`` classes for @@ -88,6 +88,8 @@ class LLMOperator(BaseOperator, LLMApprovalMixin): :param allow_modifications: If ``True``, the reviewer can edit the output before approving. The modified value is returned as the task result. Default ``False``. + :param approval_notifiers: Notifiers called once the review is open, so a + reviewer is told about it. Default ``None``. :param serialize_output: If ``True`` and ``output_type`` is a Pydantic ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via ``model_dump()`` before being pushed to XCom. Default ``False`` -- @@ -119,6 +121,7 @@ def __init__( require_approval: bool = False, approval_timeout: timedelta | None = None, allow_modifications: bool = False, + approval_notifiers: BaseNotifier | Sequence[BaseNotifier] | None = None, serialize_output: bool = False, **kwargs: Any, ) -> None: @@ -138,6 +141,9 @@ def __init__( self.require_approval = require_approval self.approval_timeout = approval_timeout self.allow_modifications = allow_modifications + self.approval_notifiers = ( + [approval_notifiers] if isinstance(approval_notifiers, BaseNotifier) else approval_notifiers or [] + ) @cached_property def llm_hook(self) -> PydanticAIHook: diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py index 5de42fb2db666..22693d93c6da7 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py @@ -56,7 +56,8 @@ class LLMBranchOperator(LLMOperator, BranchMixIn): Human-in-the-Loop approval parameters are inherited from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` - (``require_approval``, ``approval_timeout``, ``allow_modifications``). + (``require_approval``, ``approval_timeout``, ``allow_modifications``, + ``approval_notifiers``). The task pauses after the LLM chooses the branch(es) and only skips the unselected downstream tasks once a reviewer approves. Rejecting the review skips the direct downstream tasks except teardowns, matching diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py index 7b757ba22caf7..3ccfeed7667d3 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py @@ -71,7 +71,8 @@ class LLMFileAnalysisOperator(LLMOperator): Human-in-the-Loop approval parameters are inherited from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` - (``require_approval``, ``approval_timeout``, ``allow_modifications``). + (``require_approval``, ``approval_timeout``, ``allow_modifications``, + ``approval_notifiers``). The task pauses after the file analysis and only returns the result once a reviewer approves. """ diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py index 7992419296fa2..9a0957a51d6d6 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py @@ -110,7 +110,8 @@ class LLMSchemaCompareOperator(LLMOperator): Human-in-the-Loop approval parameters are inherited from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` - (``require_approval``, ``approval_timeout``, ``allow_modifications``). + (``require_approval``, ``approval_timeout``, ``allow_modifications``, + ``approval_notifiers``). The task pauses after the comparison and only returns the result once a reviewer approves. The review body shows the compatibility verdict, a mismatch severity summary, and the full result JSON. diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py index 262d52ff8d0f4..480c9b655a7d0 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py @@ -83,7 +83,8 @@ class LLMSQLQueryOperator(LLMOperator): Human-in-the-Loop approval parameters are inherited from :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` - (``require_approval``, ``approval_timeout``, ``allow_modifications``). + (``require_approval``, ``approval_timeout``, ``allow_modifications``, + ``approval_notifiers``). When ``allow_modifications=True`` and the reviewer edits the SQL, the modified query is re-validated against the same safety rules before being returned. diff --git a/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py b/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py index ace3ce7d3d98e..15147c66be75f 100644 --- a/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py +++ b/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py @@ -23,6 +23,7 @@ if not AIRFLOW_V_3_1_PLUS: pytest.skip("Human in the loop is only compatible with Airflow >= 3.1.0", allow_module_level=True) +from collections.abc import Sequence from datetime import timedelta from unittest.mock import MagicMock, patch from uuid import uuid4 @@ -32,6 +33,7 @@ from airflow.providers.common.ai.mixins.approval import ( LLMApprovalMixin, ) +from airflow.providers.common.compat.sdk import BaseNotifier from airflow.providers.standard.exceptions import HITLRejectException, HITLTriggerEventError if AIRFLOW_V_3_3_PLUS: @@ -53,11 +55,13 @@ def __init__( task_id: str = "test_task", approval_timeout: timedelta | None = None, allow_modifications: bool = False, + approval_notifiers: Sequence[BaseNotifier] = (), ): self.prompt = prompt self.task_id = task_id self.approval_timeout = approval_timeout self.allow_modifications = allow_modifications + self.approval_notifiers = approval_notifiers self.defer = MagicMock() self.log = MagicMock() @@ -174,6 +178,31 @@ def test_array_schema_passes_list_param_value( defer_kwargs = approval_op_with_modifications.defer.call_args[1] assert defer_kwargs["kwargs"]["generated_output"] == '["task_a"]' + @patch(HITL_TRIGGER_PATH, autospec=True) + @patch(UPSERT_HITL_PATH) + def test_notifiers_fire_once_the_review_is_open(self, mock_upsert, mock_trigger_cls, context): + notifier = MagicMock(spec=BaseNotifier) + order = MagicMock() + order.attach_mock(mock_upsert, "open_review") + order.attach_mock(notifier, "notify") + op = FakeOperator(approval_notifiers=[notifier]) + + op.defer_for_approval(context, "output") + + notifier.assert_called_once_with(context) + assert [call[0] for call in order.mock_calls] == ["open_review", "notify"] + + @patch(HITL_TRIGGER_PATH, autospec=True) + @patch(UPSERT_HITL_PATH) + def test_notifier_failure_stops_the_review(self, mock_upsert, mock_trigger_cls, context): + notifier = MagicMock(spec=BaseNotifier, side_effect=RuntimeError("smtp down")) + op = FakeOperator(approval_notifiers=[notifier]) + + with pytest.raises(RuntimeError, match="smtp down"): + op.defer_for_approval(context, "output") + + op.defer.assert_not_called() + @patch(HITL_TRIGGER_PATH, autospec=True) @patch(UPSERT_HITL_PATH) def test_no_modifications_params_empty(self, mock_upsert, mock_trigger_cls, approval_op, context): diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py index e2004b4031c32..d27a19b9acc6c 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py @@ -36,7 +36,7 @@ except ImportError: _CORE_WALKER = False -from airflow.providers.common.compat.sdk import TaskDeferred +from airflow.providers.common.compat.sdk import BaseNotifier, TaskDeferred if AIRFLOW_V_3_3_PLUS: # On 3.3+ cores require_approval pauses the task in AWAITING_INPUT; older cores defer @@ -193,6 +193,12 @@ def test_default_approval_flags(self): assert op.require_approval is False assert op.allow_modifications is False assert op.approval_timeout is None + assert op.approval_notifiers == [] + + def test_single_approval_notifier_normalized_to_list(self): + notifier = MagicMock(spec=BaseNotifier) + op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c", approval_notifiers=notifier) + assert op.approval_notifiers == [notifier] @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail")