kaxil commented on code in PR #74493:
URL: https://github.com/apache/airflow/pull/74493#discussion_r4394101963
##########
providers/common/ai/tests/unit/common/ai/operators/test_llm.py:
##########
@@ -1545,6 +1545,44 @@ def
test_approval_resume_survives_unparsable_usage_limits(self, mock_upsert):
assert resumed == "draft"
accessor.delete.assert_called_once_with(USAGE_BUDGET_KEY)
+ @patch.object(LLMOperator, "skip")
+ def
test_execute_complete_reject_skips_downstream_without_fail_on_reject(self,
mock_skip):
+ op = LLMOperator(
+ task_id="t", prompt="p", llm_conn_id="c", require_approval=True,
fail_on_reject=False
+ )
+ downstream = MagicMock(is_teardown=False)
+ task = MagicMock()
+ task.get_direct_relatives.return_value = [downstream,
MagicMock(is_teardown=True)]
+ ti = MagicMock()
+ ctx = MagicMock(**{"__getitem__": lambda self, key: {"task": task,
"ti": ti}[key]})
+
+ result = op.execute_complete(
+ ctx,
+ generated_output="output",
+ event={"chosen_options": ["Reject"], "responded_by_user": {"id":
"u1", "name": "admin"}},
+ )
+
+ assert result is None
+ assert mock_skip.call_args.kwargs["ti"] is ti
+ assert list(mock_skip.call_args.kwargs["tasks"]) == [downstream]
+
+ @pytest.mark.parametrize(
+ ("require_approval", "fail_on_reject", "expected"),
+ [(True, False, True), (True, True, False), (False, False, False)],
Review Comment:
In all three cases `require_approval` equals `_may_review`, so this still
passes if the check in `__init__` goes back to `self.require_approval`. Could
we add a case with `require_approval=False` and a reviewing `decision_policy`
(one with `min_confidence` set), once with `fail_on_reject=False` (expect
`True`) and once with `fail_on_reject=True` (expect `False`)?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -575,6 +600,19 @@ def _resolve_task_instance(self, context: Context) -> Any:
return None
return ti if callable(getattr(ti, "xcom_push", None)) else None
+ def _skip_downstream_on_reject(self, context: Context, event: dict[str,
Any]) -> None:
+ self.log.info("Rejected by %s. Skipping downstream tasks...",
self._describe_responder(event))
Review Comment:
On a timeout rejection this logs "Rejected by the approval timeout default",
which doesn't say which setting or deadline did it. The `HITLRejectException`
caught in `_complete_review` already says `approval_timeout expired with
on_approval_timeout='reject'`. Could we pass it in and log its message, so the
task log shows nobody clicked Reject? With this PR the line is reachable from
four more operators.
##########
providers/common/ai/docs/approval_gates.rst:
##########
@@ -191,9 +191,8 @@ Timeouts
``[scheduler] trigger_timeout_check_interval`` seconds (15 by default), so a
review can
stay open that long past its deadline. On Airflow 3.1 and 3.2 the triggerer
enforces the
deadline. ``on_approval_timeout="approve"`` returns the generated output when
the
-timeout expires. ``"reject"`` still fails this operator; only
-:class:`~airflow.providers.common.ai.operators.llm_branch.LLMBranchOperator`
turns a
-rejection into a downstream skip. The chosen option is also pre-highlighted as
the
+timeout expires. ``"reject"`` answers the review with a rejection, which fails
the task
+unless ``fail_on_reject=False``. The chosen option is also pre-highlighted as
the
Review Comment:
This is the only place the page mentions `fail_on_reject`. The rejection
paragraph further up ("A rejection, or a timeout with `"fail"` or `"reject"`,
fails the task like any other exception, so `retries` apply") still says it
unconditionally, but with `fail_on_reject=False` a rejection succeeds the task
and skips its direct downstream. Could that paragraph get a sentence on it? The
wording around `llm_branch.rst:148` would fit.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -309,6 +324,10 @@ def __init__(
self.approval_timeout = approval_timeout
self.on_approval_timeout = on_approval_timeout
self.allow_modifications = allow_modifications
+ self.fail_on_reject = fail_on_reject
+ self.ignore_downstream_trigger_rules = ignore_downstream_trigger_rules
+ if self._may_review and not fail_on_reject:
+ self._can_skip_downstream = True
Review Comment:
Mapped tasks don't see this, because both mapped paths take the flag from
the class, not the instance. A decorated `.expand()` reads
`operator_class._can_skip_downstream` (task-sdk `bases/decorator.py:728`),
which is still `False`, so a mapped `@task.llm(require_approval=True,
fail_on_reject=False)` serializes as not skip-capable. On a mapped TI `skip()`
only pushes the XCom and returns (it raises only at `map_index == -1`), so the
task succeeds with `None` and `NotPreviouslySkippedDep` never looks at this
parent. As far as I can tell from the dep, the downstream then runs on the
rejected item.
The other way round, `.partial().expand()` uses `issubclass(operator_class,
SkipMixin)` (`mappedoperator.py:260`), so every mapped LLM operator now turns
the skip check on, review or not.
Could we set `_can_skip_downstream = True` on the `LLMOperator` class and
assign `self._can_skip_downstream = self._may_review and not fail_on_reject`
here? Mapped tasks then over-approximate (an extra XCom read) rather than
letting a rejection through. A test that serializes a mapped `@task.llm` would
pin it, and the comment in `_skip_downstream_on_reject` saying `skip()` always
raises needs the mapped case.
##########
providers/common/ai/docs/operators/llm.rst:
##########
@@ -307,6 +307,12 @@ Parameters
``{"id": ..., "name": ...}`` dicts where ``id`` is the auth manager's user
id.
``None`` (default) lets any user with the permission respond. Fixed at first
run on Airflow 3.1 to 3.3; replaced on clear or retry from 3.4. Needs
Airflow 3.1+.
+- ``fail_on_reject``: If ``False``, a rejected review skips the direct
downstream
+ tasks except teardowns and the task succeeds, as ``ApprovalOperator`` does.
+ Default ``True``: a rejection fails the task.
+- ``ignore_downstream_trigger_rules``: If ``True``, a rejection that does not
fail
+ the task skips every downstream task rather than only the direct ones.
Review Comment:
Teardowns are filtered out on this path too, so "every downstream task
except teardowns"? The `LLMOperator` docstring has the same wording.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -184,10 +184,10 @@ def execute_complete(
generated_output: str,
event: dict[str, Any],
decision: dict[str, Any] | None = None,
- ) -> str:
+ ) -> Any:
Review Comment:
This can only return `str | None` now, so could the annotation say that
rather than `Any`? Also, the inherited approval parameter list in this class's
docstring (and in `LLMSchemaCompareOperator` and `LLMFileAnalysisOperator`)
doesn't mention `fail_on_reject` or `ignore_downstream_trigger_rules` yet. Only
the branch operator's list was updated.
##########
providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py:
##########
@@ -810,6 +810,22 @@ def test_execute_complete_approved_returns_dict(self):
assert resumed == result.model_dump()
+ @mock.patch.object(LLMSchemaCompareOperator, "skip")
Review Comment:
`LLMSchemaCompareOperator` doesn't override any of the resume methods, so
this runs the same code as the new `test_llm.py` test, and `mock_skip` is never
asserted. Maybe drop it, or parametrize one test over the subclasses so
`LLMFileAnalysisOperator` is covered as well?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
