This is an automated email from the ASF dual-hosted git repository.
guan404ming pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 44ea07505c6 Skip downstream tasks on LLMBranchOperator reject (#71073)
44ea07505c6 is described below
commit 44ea07505c6e312a6322dbdb2b72cdb79b4484c4
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Thu Aug 27 22:05:07 2026 +0800
Skip downstream tasks on LLMBranchOperator reject (#71073)
* Skip downstream tasks on LLMBranchOperator reject
* Keep teardowns running and log reviewer on reject
* Scope teardown carve-out docs to the reject path
---
providers/common/ai/docs/operators/llm_branch.rst | 14 +++++++--
.../providers/common/ai/operators/llm_branch.py | 22 ++++++++++++--
.../unit/common/ai/operators/test_llm_branch.py | 34 ++++++++++++++++++++++
3 files changed, 65 insertions(+), 5 deletions(-)
diff --git a/providers/common/ai/docs/operators/llm_branch.rst
b/providers/common/ai/docs/operators/llm_branch.rst
index 21558b2277f..1ef8c557abc 100644
--- a/providers/common/ai/docs/operators/llm_branch.rst
+++ b/providers/common/ai/docs/operators/llm_branch.rst
@@ -92,14 +92,20 @@ 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 the direct downstream tasks except teardowns**,
+matching
+:class:`~airflow.providers.standard.operators.hitl.ApprovalOperator`. The
+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``).
``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 <howto/operator:llm>`.
@@ -133,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``.
+- ``fail_on_reject``: If ``True``, a rejected review fails the task instead of
+ skipping the downstream tasks. 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 88c29943259..5de42fb2db6 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 typing import TYPE_CHECKING, Any
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 the downstream tasks. 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 the direct downstream tasks except teardowns, 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 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
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,15 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
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 by %s. Skipping downstream tasks...",
event.get("responded_by_user"))
+ tasks = context["task"].get_direct_relatives(upstream=False)
+ self.skip(ti=context["ti"], tasks=(t for t in tasks if not
t.is_teardown))
+ return 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 7f2d11bf263..7ff0b9bf83c 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.mixins.approval import
LLMApprovalMixin
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,39 @@ class TestLLMBranchOperatorApproval:
assert result == ["task_a", "task_c"]
mock_do_branch.assert_called_once_with(ctx, ["task_a", "task_c"])
+ @patch.object(LLMBranchOperator, "skip")
+ @patch.object(LLMBranchOperator, "do_branch")
+ def test_execute_complete_reject_skips_downstream_except_teardowns(self,
mock_do_branch, mock_skip):
+ op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c")
+ op.downstream_task_ids = {"task_a", "cleanup"}
+ event = {"chosen_options": ["Reject"], "responded_by_user": "admin"}
+ task_a = MagicMock(is_teardown=False)
+ cleanup = MagicMock(is_teardown=True)
+ task = MagicMock()
+ task.get_direct_relatives.return_value = [task_a, cleanup]
+ ti = MagicMock()
+ ctx = MagicMock(**{"__getitem__": lambda self, key: {"task": task,
"ti": ti}[key]})
+
+ result = op.execute_complete(ctx, generated_output="task_a",
event=event)
+
+ assert result is None
+ task.get_direct_relatives.assert_called_once_with(upstream=False)
+ mock_skip.assert_called_once()
+ assert mock_skip.call_args.kwargs["ti"] is ti
+ assert list(mock_skip.call_args.kwargs["tasks"]) == [task_a]
+ mock_do_branch.assert_not_called()
+
+ @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."""
