From a8db1562838fcc7212b61721b106d93ea7ecda41 Mon Sep 17 00:00:00 2001 From: "Guan-Ming (Wesley) Chiu" <105915352+guan404ming@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:10:26 +0800 Subject: [PATCH] Skip downstream tasks on LLMBranchOperator reject --- .../common/ai/docs/operators/llm_branch.rst | 9 ++++--- .../common/ai/operators/llm_branch.py | 20 +++++++++++++-- .../common/ai/operators/test_llm_branch.py | 25 +++++++++++++++++++ 3 files changed, 49 insertions(+), 5 deletions(-) diff --git a/providers/common/ai/docs/operators/llm_branch.rst b/providers/common/ai/docs/operators/llm_branch.rst index 21558b2277f02..93988ef6e5c59 100644 --- a/providers/common/ai/docs/operators/llm_branch.rst +++ b/providers/common/ai/docs/operators/llm_branch.rst @@ -92,9 +92,10 @@ against the downstream task IDs before branching: :start-after: [START howto_operator_llm_branch_approval] :end-before: [END howto_operator_llm_branch_approval] -Rejecting the review, or letting ``approval_timeout`` expire, **fails** the -task (``HITLRejectException`` / ``HITLTimeoutError``), so downstream tasks -end up ``upstream_failed`` rather than skipped. +Rejecting the review **skips every downstream task**, matching +:class:`~airflow.providers.standard.operators.hitl.ApprovalOperator`. Set +``fail_on_reject=True`` to fail the task instead (generally discouraged). +Letting ``approval_timeout`` expire fails the task (``HITLTimeoutError``). ``require_approval=True`` requires a string prompt: a decorated callable returning a ``Sequence[UserContent]`` raises ``TypeError`` before the LLM @@ -133,6 +134,8 @@ Parameters means wait indefinitely. Default ``None``. - ``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 + skipping every downstream task. Generally discouraged. Default ``False``. Logging ------- 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 88c299432594b..cdebfbc98547d 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 @@ -25,6 +25,7 @@ from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.utils.logging import log_run_summary +from airflow.providers.standard.exceptions import HITLRejectException from airflow.providers.standard.operators.branch import BranchMixIn if TYPE_CHECKING: @@ -46,6 +47,10 @@ class LLMBranchOperator(LLMOperator, BranchMixIn): :param system_prompt: System-level instructions for the LLM agent. :param allow_multiple_branches: When ``False`` (default) the LLM returns a single task ID. When ``True`` the LLM may return one or more task IDs. + :param fail_on_reject: If ``True``, a rejected review fails the task + instead of skipping every downstream task. Generally discouraged, + as for :class:`~airflow.providers.standard.operators.hitl.ApprovalOperator`. + Default ``False``. :param agent_params: Additional keyword arguments passed to the pydantic-ai ``Agent`` constructor (e.g. ``retries``, ``model_settings``, ``tools``). @@ -53,7 +58,10 @@ class LLMBranchOperator(LLMOperator, BranchMixIn): :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` (``require_approval``, ``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. The review form + unselected downstream tasks once a reviewer approves. Rejecting the + review skips every downstream task, matching + :class:`~airflow.providers.standard.operators.hitl.ApprovalOperator`; + set ``fail_on_reject=True`` to fail the task instead. The review form lists the valid downstream task IDs; with ``allow_modifications=True`` the editable choice is rendered as a dropdown of those IDs (single-branch mode) or a multi-select of them (``allow_multiple_branches=True``), and @@ -69,11 +77,13 @@ def __init__( self, *, allow_multiple_branches: bool = False, + fail_on_reject: bool = False, **kwargs: Any, ) -> None: kwargs.pop("output_type", None) super().__init__(**kwargs) self.allow_multiple_branches = allow_multiple_branches + self.fail_on_reject = fail_on_reject def execute(self, context: Context) -> str | Iterable[str] | None: if self.require_approval: @@ -133,7 +143,13 @@ def execute(self, context: Context) -> str | Iterable[str] | None: def execute_complete(self, context: Context, generated_output: str, event: dict[str, Any]) -> Any: """Resume after human review, validating the reviewed choice before branching.""" - output = super().execute_complete(context, generated_output, event) + try: + output = super().execute_complete(context, generated_output, event) + except HITLRejectException: + if self.fail_on_reject: + raise + self.log.info("Rejected. Skipping all downstream tasks...") + return self.do_branch(context, None) branches = self._parse_reviewed_branches(output) selected = {branches} if isinstance(branches, str) else set(branches) invalid = selected - self.downstream_task_ids diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py index 7f2d11bf26336..ac99bae74f9bd 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py @@ -26,6 +26,7 @@ from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator from airflow.providers.common.compat.sdk import Param, ParamValidationError, TaskDeferred +from airflow.providers.standard.exceptions import HITLRejectException from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS @@ -397,6 +398,30 @@ def test_execute_complete_approved_multiple_branches(self, mock_do_branch): assert result == ["task_a", "task_c"] mock_do_branch.assert_called_once_with(ctx, ["task_a", "task_c"]) + @patch.object(LLMBranchOperator, "do_branch") + def test_execute_complete_reject_skips_all_downstream(self, mock_do_branch): + mock_do_branch.return_value = None + op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c") + op.downstream_task_ids = {"task_a", "task_b"} + event = {"chosen_options": ["Reject"], "responded_by_user": "admin"} + ctx = _make_context() + + result = op.execute_complete(ctx, generated_output="task_a", event=event) + + assert result is None + mock_do_branch.assert_called_once_with(ctx, None) + + @patch.object(LLMBranchOperator, "do_branch") + def test_execute_complete_reject_fails_with_fail_on_reject(self, mock_do_branch): + op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c", fail_on_reject=True) + op.downstream_task_ids = {"task_a", "task_b"} + event = {"chosen_options": ["Reject"], "responded_by_user": "admin"} + + with pytest.raises(HITLRejectException, match="rejected"): + op.execute_complete(_make_context(), generated_output="task_a", event=event) + + mock_do_branch.assert_not_called() + @patch.object(LLMBranchOperator, "do_branch") def test_execute_complete_with_modified_branch(self, mock_do_branch): """A reviewer-modified branch is used when it is a valid downstream task."""