diff --git a/providers/common/ai/docs/operators/llm.rst b/providers/common/ai/docs/operators/llm.rst index 426a2573206fc..893cef05fa04a 100644 --- a/providers/common/ai/docs/operators/llm.rst +++ b/providers/common/ai/docs/operators/llm.rst @@ -198,7 +198,11 @@ Set ``require_approval=True`` to pause the task after the LLM generates its output and wait for a human reviewer to approve or reject it via the Airflow HITL interface. Optionally allow the reviewer to edit the output before approving with ``allow_modifications=True``, and set a deadline with -``approval_timeout``: +``approval_timeout``. + +When ``approval_timeout`` expires without a review, the task fails by default. +Set ``on_approval_timeout="approve"`` or ``"reject"`` to answer the review with +that option instead, so an unattended pipeline keeps moving: .. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py :language: python @@ -223,6 +227,9 @@ Parameters for human review. Default ``False``. - ``approval_timeout``: Maximum time to wait for a review (``timedelta``). ``None`` means wait indefinitely. Default ``None``. +- ``on_approval_timeout``: Outcome when ``approval_timeout`` expires without a + review: ``"fail"`` (default), ``"approve"``, or ``"reject"``. Requires + ``approval_timeout``. - ``allow_modifications``: If ``True``, the reviewer can edit the output before approving. Default ``False``. diff --git a/providers/common/ai/docs/operators/llm_branch.rst b/providers/common/ai/docs/operators/llm_branch.rst index 1ef8c557abcf2..5d4945c730d61 100644 --- a/providers/common/ai/docs/operators/llm_branch.rst +++ b/providers/common/ai/docs/operators/llm_branch.rst @@ -99,15 +99,17 @@ teardown carve-out applies only to rejection: approving branches as usual, so a teardown that is not among the chosen branch(es) is skipped like any other unselected downstream task. Set ``fail_on_reject=True`` to fail the task on rejection instead (generally discouraged). Letting -``approval_timeout`` expire fails the task (``HITLTimeoutError``). +``approval_timeout`` expire fails the task (``HITLTimeoutError``) unless +``on_approval_timeout`` answers the review for you; a timeout-driven +rejection then skips downstream like any other rejection. ``require_approval=True`` requires a string prompt: a decorated callable 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``, ``on_approval_timeout``, and the rest of the approval +behaviour are inherited from :ref:`LLMOperator `. How It Works ------------ @@ -137,6 +139,9 @@ Parameters branch(es) and waits for human review before branching. Default ``False``. - ``approval_timeout``: Maximum time to wait for a review (``timedelta``). ``None`` means wait indefinitely. Default ``None``. +- ``on_approval_timeout``: Outcome when ``approval_timeout`` expires without a + review: ``"fail"`` (default), ``"approve"``, or ``"reject"``. Requires + ``approval_timeout``. - ``allow_modifications``: If ``True``, the reviewer can change the chosen branch(es) before approving. Default ``False``. - ``fail_on_reject``: If ``True``, a rejected review fails the task instead of diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst b/providers/common/ai/docs/operators/llm_file_analysis.rst index ac688ebfaf415..dc65e31bde069 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``, ``on_approval_timeout``, and +``allow_modifications`` -- 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..e452bcf075608 100644 --- a/providers/common/ai/docs/operators/llm_schema_compare.rst +++ b/providers/common/ai/docs/operators/llm_schema_compare.rst @@ -121,7 +121,7 @@ Set ``require_approval=True`` to pause the task after the comparison and wait for a human reviewer to approve the result before it is returned. The review body shows the compatibility verdict, a mismatch severity summary, and the full result JSON. Rejecting the review, or letting ``approval_timeout`` -expire, fails the task: +expire with the default ``on_approval_timeout="fail"``, fails the task: .. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py :language: python @@ -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``, ``on_approval_timeout``, ``allow_modifications``, and +the rest of the approval behaviour are inherited from +:ref:`LLMOperator `. Conditional ETL Based on Schema Compatibility ---------------------------------------------- @@ -191,6 +192,9 @@ Parameters waits for human review before returning the result. Default ``False``. - ``approval_timeout``: Maximum time to wait for a review (``timedelta``). ``None`` means wait indefinitely. Default ``None``. +- ``on_approval_timeout``: Outcome when ``approval_timeout`` expires without a + review: ``"fail"`` (default), ``"approve"``, or ``"reject"``. Requires + ``approval_timeout``. - ``allow_modifications``: If ``True``, the reviewer can edit the result JSON before approving. Default ``False``. diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py index 545a138e9df9f..54a4542acd7f7 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py +++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py @@ -159,6 +159,7 @@ def example_llm_operator_approval(): system_prompt="You are a financial analyst. Be concise and accurate.", require_approval=True, approval_timeout=timedelta(hours=24), + on_approval_timeout="reject", allow_modifications=True, ) 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..04efbc9d8e388 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 @@ -20,7 +20,7 @@ import json import logging from datetime import timedelta -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, ClassVar, Protocol from pydantic import BaseModel, TypeAdapter @@ -42,6 +42,7 @@ class DeferForApprovalProtocol(Protocol): approval_timeout: timedelta | None allow_modifications: bool + on_approval_timeout: str prompt: str task_id: str defer: Any @@ -62,16 +63,23 @@ class LLMApprovalMixin: before approving. The (possibly modified) output is then returned as the task result. + ``on_approval_timeout`` decides what happens when ``approval_timeout`` + expires without a response: ``"fail"`` raises ``HITLTimeoutError``, while + ``"approve"`` and ``"reject"`` answer the review with that option so the + task resumes as if a reviewer had chosen it. + Operators that use this mixin must set the following attributes: - ``require_approval`` (``bool``) - ``allow_modifications`` (``bool``) - ``approval_timeout`` (``timedelta | None``) + - ``on_approval_timeout`` (``str``) - ``prompt`` (``str``) """ APPROVE = "Approve" REJECT = "Reject" + TIMEOUT_DEFAULTS: ClassVar[dict[str, list[str]]] = {"approve": [APPROVE], "reject": [REJECT]} def validate_approval_prompt(self: DeferForApprovalProtocol) -> None: """Fail fast when the prompt cannot be rendered as text in the approval review body.""" @@ -99,7 +107,8 @@ def defer_for_approval( On Airflow 3.3+ the task parks in the ``awaiting_input`` state (no trigger or triggerer involved); on older versions it defers to :class:`HITLTrigger`. Either way it resumes in - ``execute_complete`` once a response (or timeout default) arrives. + ``execute_complete`` once a response (or timeout default) arrives. ``on_approval_timeout`` + supplies that timeout default; ``"fail"`` supplies none, so the review times out as an error. :param context: Airflow task context. :param output: The generated output to present for review. @@ -130,6 +139,7 @@ def defer_for_approval( output = TypeAdapter(type(output)).dump_json(output).decode() ti_id = context["task_instance"].id + timeout_defaults = LLMApprovalMixin.TIMEOUT_DEFAULTS.get(self.on_approval_timeout) if subject is None: subject = f"Review output for task `{self.task_id}`" @@ -160,7 +170,7 @@ def defer_for_approval( options=[LLMApprovalMixin.APPROVE, LLMApprovalMixin.REJECT], subject=subject, body=body, - defaults=None, + defaults=timeout_defaults, multiple=False, params=hitl_params, ) @@ -179,7 +189,7 @@ def defer_for_approval( trigger=HITLTrigger( ti_id=ti_id, options=[LLMApprovalMixin.APPROVE, LLMApprovalMixin.REJECT], - defaults=None, + defaults=timeout_defaults, params=hitl_params, multiple=False, timeout_datetime=utcnow() + self.approval_timeout if self.approval_timeout else None, @@ -199,8 +209,9 @@ def execute_complete(self, context: Context, generated_output: str, event: dict[ :param context: Airflow task context. :param generated_output: The output that was deferred for review. :param event: Trigger event payload containing ``chosen_options``, - ``params_input``, and ``responded_by_user``. - :raises HITLRejectException: If the reviewer rejected the output. + ``params_input``, ``responded_by_user``, and ``timedout``. + :raises HITLRejectException: If the reviewer, or the + ``on_approval_timeout="reject"`` default, rejected the output. :raises HITLTriggerEventError: If the trigger reported an error. :raises HITLTimeoutError: If the approval timed out. """ @@ -219,6 +230,8 @@ def execute_complete(self, context: Context, generated_output: str, event: dict[ responded_by_user = event.get("responded_by_user") chosen = event["chosen_options"] if self.APPROVE not in chosen: + if event.get("timedout"): + raise HITLRejectException("Output was rejected by the approval timeout default.") raise HITLRejectException(f"Output was rejected by the reviewer {responded_by_user}.") output = generated_output 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..00b86ea57e9da 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 @@ -21,7 +21,7 @@ from collections.abc import Sequence from datetime import timedelta from functools import cached_property -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar, Literal from pydantic import BaseModel @@ -84,7 +84,12 @@ class LLMOperator(BaseOperator, LLMApprovalMixin): output and waits for a human reviewer to approve or reject via the HITL interface. Default ``False``. :param approval_timeout: Maximum time to wait for a review. When - exceeded, the task fails with ``TimeoutError``. + exceeded, ``on_approval_timeout`` decides the outcome. + :param on_approval_timeout: What to do when ``approval_timeout`` expires + without a review. ``"fail"`` (default) fails the task with + ``HITLTimeoutError``; ``"approve"`` and ``"reject"`` answer the review + with that option, so the task resumes as if a reviewer had chosen it. + Requires ``approval_timeout`` to be set. :param allow_modifications: If ``True``, the reviewer can edit the output before approving. The modified value is returned as the task result. Default ``False``. @@ -118,6 +123,7 @@ def __init__( usage_limits: UsageLimits | None = None, require_approval: bool = False, approval_timeout: timedelta | None = None, + on_approval_timeout: Literal["fail", "approve", "reject"] = "fail", allow_modifications: bool = False, serialize_output: bool = False, **kwargs: Any, @@ -135,8 +141,18 @@ def __init__( self._serialize_model_output = serialize_output or not _CORE_WALKER self.agent_params = agent_params or {} self.usage_limits = usage_limits + if on_approval_timeout not in ("fail", *LLMApprovalMixin.TIMEOUT_DEFAULTS): + raise ValueError( + f"on_approval_timeout must be 'fail', 'approve', or 'reject', got {on_approval_timeout!r}." + ) + if on_approval_timeout != "fail" and approval_timeout is None: + raise ValueError( + f"on_approval_timeout={on_approval_timeout!r} has no effect without approval_timeout. " + "Set approval_timeout, or leave on_approval_timeout as 'fail'." + ) self.require_approval = require_approval self.approval_timeout = approval_timeout + self.on_approval_timeout = on_approval_timeout self.allow_modifications = allow_modifications @cached_property 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..d69539fb31e4e 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``, ``on_approval_timeout``, + ``allow_modifications``). 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..501b6754ce095 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``, ``on_approval_timeout``, + ``allow_modifications``). 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..dbbb98adc1479 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``, ``on_approval_timeout``, + ``allow_modifications``). 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..66399b07b1b18 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``, ``on_approval_timeout``, + ``allow_modifications``). 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..302571309233a 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 @@ -52,11 +52,13 @@ def __init__( prompt: str = "Summarize this", task_id: str = "test_task", approval_timeout: timedelta | None = None, + on_approval_timeout: str = "fail", allow_modifications: bool = False, ): self.prompt = prompt self.task_id = task_id self.approval_timeout = approval_timeout + self.on_approval_timeout = on_approval_timeout self.allow_modifications = allow_modifications self.defer = MagicMock() @@ -206,6 +208,22 @@ def test_timeout_sets_timeout_datetime(self, mock_upsert, mock_trigger_cls, mock defer_kwargs = op.defer.call_args[1] assert defer_kwargs["timeout"] == timeout + @pytest.mark.parametrize( + ("on_approval_timeout", "expected_defaults"), + [("fail", None), ("approve", ["Approve"]), ("reject", ["Reject"])], + ) + @patch(HITL_TRIGGER_PATH, autospec=True) + @patch(UPSERT_HITL_PATH) + def test_on_approval_timeout_sets_hitl_defaults( + self, mock_upsert, mock_trigger_cls, context, on_approval_timeout, expected_defaults + ): + op = FakeOperator(approval_timeout=timedelta(hours=1), on_approval_timeout=on_approval_timeout) + + op.defer_for_approval(context, "output") + + assert mock_upsert.call_args[1]["defaults"] == expected_defaults + assert mock_trigger_cls.call_args[1]["defaults"] == expected_defaults + @patch(HITL_TRIGGER_PATH, autospec=True) @patch(UPSERT_HITL_PATH) def test_no_timeout_passes_none(self, mock_upsert, mock_trigger_cls, approval_op, context): @@ -420,6 +438,12 @@ def test_event_missing_responded_by_user(self, approval_op): assert result == "output" + def test_timed_out_rejection_names_the_timeout_default(self, approval_op): + event = {"chosen_options": ["Reject"], "responded_by_user": None, "timedout": True} + + with pytest.raises(HITLRejectException, match="Output was rejected by the approval timeout default."): + approval_op.execute_complete({}, generated_output="output", event=event) + def test_rejection_message_includes_username(self, approval_op): event = {"chosen_options": ["Reject"], "responded_by_user": "alice"} @@ -453,6 +477,21 @@ def test_approval_timeout_carried_on_await(self, mock_upsert, context): assert exc_info.value.timeout == timeout + @pytest.mark.parametrize( + ("on_approval_timeout", "expected_defaults"), + [("fail", None), ("approve", ["Approve"]), ("reject", ["Reject"])], + ) + @patch(UPSERT_HITL_PATH) + def test_on_approval_timeout_sets_hitl_defaults_on_await( + self, mock_upsert, context, on_approval_timeout, expected_defaults + ): + op = FakeOperator(approval_timeout=timedelta(hours=1), on_approval_timeout=on_approval_timeout) + + with pytest.raises(TaskAwaitingInput): + op.defer_for_approval(context, "output") + + assert mock_upsert.call_args[1]["defaults"] == expected_defaults + @patch(UPSERT_HITL_PATH) def test_pydantic_output_stringified_on_await(self, mock_upsert, approval_op, context): class Answer(BaseModel): 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..7f1a708121d30 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 @@ -193,6 +193,21 @@ 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.on_approval_timeout == "fail" + + def test_unknown_on_approval_timeout_raises(self): + with pytest.raises(ValueError, match="on_approval_timeout must be"): + LLMOperator( + task_id="t", + prompt="p", + llm_conn_id="c", + approval_timeout=timedelta(hours=1), + on_approval_timeout="skip", + ) + + def test_on_approval_timeout_without_approval_timeout_raises(self): + with pytest.raises(ValueError, match="has no effect without approval_timeout"): + LLMOperator(task_id="t", prompt="p", llm_conn_id="c", on_approval_timeout="approve") @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail")