From 32a04318dcfdae9aa46e0fc6e3651835d7f49397 Mon Sep 17 00:00:00 2001 From: "Guan-Ming (Wesley) Chiu" <105915352+guan404ming@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:55:50 +0800 Subject: [PATCH] Support notifiers in LLM approval reviews A pending review only showed up on the Required Actions page, so the people who had to act on it were never told. The task then sat waiting until somebody happened to look. --- providers/common/ai/docs/operators/llm.rst | 9 ++++++ .../common/ai/docs/operators/llm_branch.rst | 6 ++-- .../ai/docs/operators/llm_file_analysis.rst | 4 +-- .../ai/docs/operators/llm_schema_compare.rst | 7 +++-- .../providers/common/ai/mixins/approval.py | 13 +++++++++ .../providers/common/ai/operators/llm.py | 8 ++++- .../common/ai/operators/llm_branch.py | 3 +- .../common/ai/operators/llm_file_analysis.py | 3 +- .../common/ai/operators/llm_schema_compare.py | 3 +- .../providers/common/ai/operators/llm_sql.py | 3 +- .../unit/common/ai/mixins/test_approval.py | 29 +++++++++++++++++++ .../unit/common/ai/operators/test_llm.py | 8 ++++- 12 files changed, 84 insertions(+), 12 deletions(-) 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")