Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions providers/common/ai/docs/operators/llm.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
----------

Expand All @@ -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
-------
Expand Down
6 changes: 4 additions & 2 deletions providers/common/ai/docs/operators/llm_branch.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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 <howto/operator:llm>`.
``approval_timeout``, ``approval_notifiers``, and the rest of the approval
behaviour are inherited from :ref:`LLMOperator <howto/operator:llm>`.

How It Works
------------
Expand Down Expand Up @@ -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``.

Expand Down
4 changes: 2 additions & 2 deletions providers/common/ai/docs/operators/llm_file_analysis.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
-----------------
Expand Down
7 changes: 5 additions & 2 deletions providers/common/ai/docs/operators/llm_schema_compare.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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 <howto/operator:llm>`.
``approval_timeout``, ``allow_modifications``, ``approval_notifiers``, and the
rest of the approval behaviour are inherited from
:ref:`LLMOperator <howto/operator:llm>`.

Conditional ETL Based on Schema Compatibility
----------------------------------------------
Expand Down Expand Up @@ -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
-------
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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
Expand All @@ -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``)
"""

Expand Down Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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`` --
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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()
Expand Down Expand Up @@ -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):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand Down