guan404ming commented on code in PR #74493:
URL: https://github.com/apache/airflow/pull/74493#discussion_r4396689392
##########
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:
Set `_can_skip_downstream = True` on the class and
`self._can_skip_downstream = self._may_review and not fail_on_reject` in
`__init__`, added a test that serializes a mapped task in both forms, and the
`skip()` comment now covers the mapped case.
##########
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:
Added both `decision_policy` cases; the `fail_on_reject=False` one fails if
`__init__` goes back to `self.require_approval`.
##########
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:
Added a paragraph after it, worded like `llm_branch.rst`, on
`fail_on_reject=False`, `ignore_downstream_trigger_rules`, and a `"fail"`
timeout still failing the task.
##########
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:
`_skip_downstream_on_reject` now takes the `HITLRejectException` and logs
its message, and the branch timeout test asserts the `approval_timeout` text.
##########
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:
Both now say "every downstream task except teardowns".
##########
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:
Added `fail_on_reject` and `ignore_downstream_trigger_rules` to the three
inherited lists, but kept `Any`, since mypy rejects `str | None` as an override
of `LLMApprovalMixin.execute_complete`, which returns `str`.
##########
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:
Dropped it and parametrized the `test_llm.py` test over `LLMOperator`,
`LLMFileAnalysisOperator` and `LLMSchemaCompareOperator`.
--
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]
