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 428e873c9fc Support full downstream skip on LLM branch reject (#72183)
428e873c9fc is described below
commit 428e873c9fc0ce164dadf42bd51f8223fc9c60a0
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Wed Sep 9 21:53:24 2026 +0800
Support full downstream skip on LLM branch reject (#72183)
* Support full downstream skip on LLM branch reject
Rejecting a branch review only skipped the direct downstream tasks, so a
task with a permissive trigger rule further down the chain still ran even
though the reviewer had turned the branch down.
* Cover reject skip on a real Dag and document approval precondition
Reviewers asked for the transitive-plus-teardown skip to be pinned
against a real Dag topology rather than mocked getters, and for the
docs to state that the reject flags only apply with require_approval.
Co-Authored-By: Claude Fable 5.1 <[email protected]>
Claude-Session: https://claude.ai/code/session_01Fh4nN2o42JdSjvhWyLEpW6
---------
Co-authored-by: Claude Fable 5.1 <[email protected]>
---
providers/common/ai/docs/operators/llm_branch.rst | 19 ++++++----
.../providers/common/ai/operators/llm_branch.py | 19 ++++++++--
.../unit/common/ai/operators/test_llm_branch.py | 42 ++++++++++++++++------
3 files changed, 60 insertions(+), 20 deletions(-)
diff --git a/providers/common/ai/docs/operators/llm_branch.rst
b/providers/common/ai/docs/operators/llm_branch.rst
index 1ef8c557abc..d1f11a7309c 100644
--- a/providers/common/ai/docs/operators/llm_branch.rst
+++ b/providers/common/ai/docs/operators/llm_branch.rst
@@ -98,16 +98,19 @@ matching
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``).
+task on rejection instead (generally discouraged), or
+``ignore_downstream_trigger_rules=True`` to skip every downstream task rather
+than only the direct ones, so a task whose trigger rule would still run it is
+skipped too. 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>`.
+Apart from ``fail_on_reject`` and ``ignore_downstream_trigger_rules``, which
+are specific to this operator, ``approval_timeout`` and the rest of the
+approval behaviour are inherited from :ref:`LLMOperator <howto/operator:llm>`.
How It Works
------------
@@ -140,7 +143,11 @@ Parameters
- ``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``.
+ skipping the downstream tasks. Generally discouraged. Only takes effect
+ with ``require_approval=True``. Default ``False``.
+- ``ignore_downstream_trigger_rules``: If ``True``, a rejected review skips
every
+ downstream task rather than only the direct ones. Only takes effect with
+ ``require_approval=True``. 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 5de42fb2db6..b3970596bf6 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
@@ -50,7 +50,11 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
: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``.
+ Only takes effect with ``require_approval=True``. Default ``False``.
+ :param ignore_downstream_trigger_rules: If ``True``, a rejected review
skips
+ every downstream task rather than only the direct ones, so a task whose
+ trigger rule would still run it is skipped too. Only takes effect with
+ ``require_approval=True``. Default ``False``.
:param agent_params: Additional keyword arguments passed to the pydantic-ai
``Agent`` constructor (e.g. ``retries``, ``model_settings``,
``tools``).
@@ -61,7 +65,9 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
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
+ set ``fail_on_reject=True`` to fail the task instead, or
+ ``ignore_downstream_trigger_rules=True`` to skip every downstream task
+ rather than only the direct ones. 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
@@ -78,12 +84,14 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
*,
allow_multiple_branches: bool = False,
fail_on_reject: bool = False,
+ ignore_downstream_trigger_rules: 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
+ self.ignore_downstream_trigger_rules = ignore_downstream_trigger_rules
def execute(self, context: Context) -> str | Iterable[str] | None:
if self.require_approval:
@@ -149,7 +157,12 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
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)
+ task = context["task"]
+ tasks = (
+ task.get_flat_relatives(upstream=False)
+ if self.ignore_downstream_trigger_rules
+ else 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)
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 7ff0b9bf83c..94f09c7046d 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
@@ -27,6 +27,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 airflow.providers.standard.operators.empty import EmptyOperator
from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_3_PLUS
@@ -398,26 +399,45 @@ class TestLLMBranchOperatorApproval:
assert result == ["task_a", "task_c"]
mock_do_branch.assert_called_once_with(ctx, ["task_a", "task_c"])
+ @pytest.mark.db_test
+ @pytest.mark.parametrize(
+ ("ignore_downstream_trigger_rules", "with_teardown", "expected"),
+ [
+ (False, True, {"op2"}),
+ (False, False, {"op2"}),
+ (True, True, {"op2", "op3"}),
+ (True, False, {"op2", "op3", "op4"}),
+ ],
+ )
@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"}
+ def test_execute_complete_reject_skips_downstream_except_teardowns(
+ self, mock_do_branch, mock_skip, dag_maker,
ignore_downstream_trigger_rules, with_teardown, expected
+ ):
+ with dag_maker(serialized=True):
+ op1 = LLMBranchOperator(
+ task_id="op1",
+ prompt="p",
+ llm_conn_id="c",
+ require_approval=True,
+
ignore_downstream_trigger_rules=ignore_downstream_trigger_rules,
+ )
+ op2 = EmptyOperator(task_id="op2")
+ op3 = EmptyOperator(task_id="op3")
+ op4 = EmptyOperator(task_id="op4")
+ if with_teardown:
+ op4.as_teardown()
+ op1 >> op2 >> op3 >> op4
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]})
+ ctx = MagicMock(**{"__getitem__": lambda self, key: {"task": op1,
"ti": ti}[key]})
- result = op.execute_complete(ctx, generated_output="task_a",
event=event)
+ result = op1.execute_complete(ctx, generated_output="op2", 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]
+ assert {t.task_id for t in mock_skip.call_args.kwargs["tasks"]} ==
expected
mock_do_branch.assert_not_called()
@patch.object(LLMBranchOperator, "do_branch")
