guan404ming commented on code in PR #72159:
URL: https://github.com/apache/airflow/pull/72159#discussion_r3999633775
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py:
##########
@@ -165,6 +177,12 @@ def defer_for_approval(
params=hitl_params,
)
+ for notifier in self.approval_notifiers:
+ try:
+ notifier({**context, "subject": subject, "body": body})
Review Comment:
Good catch. Now set as task attributes like HITLOperator; docs use `{{
task.subject }}` / `{{ task.body }}`.
##########
providers/common/ai/tests/unit/common/ai/mixins/test_approval.py:
##########
@@ -174,6 +178,42 @@ def test_array_schema_passes_list_param_value(
defer_kwargs = approval_op_with_modifications.defer.call_args[1]
assert defer_kwargs["kwargs"]["generated_output"] == '["task_a"]'
+ @patch(HITL_TRIGGER_PATH, autospec=True)
+ @patch(UPSERT_HITL_PATH)
+ def test_notifiers_fire_once_the_review_is_open(self, mock_upsert,
mock_trigger_cls, context):
+ notifier = MagicMock(spec=BaseNotifier)
+ order = MagicMock()
+ order.attach_mock(mock_upsert, "open_review")
+ order.attach_mock(notifier, "notify")
+ op = FakeOperator(approval_notifiers=[notifier])
+
+ op.defer_for_approval(context, "output")
+
+ notifier.assert_called_once_with(
Review Comment:
Added real BaseNotifier subclass with own `subject` field, plus `[failing,
healthy]` case.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -138,6 +142,9 @@ def __init__(
self.require_approval = require_approval
self.approval_timeout = approval_timeout
self.allow_modifications = allow_modifications
+ self.approval_notifiers = (
+ [approval_notifiers] if isinstance(approval_notifiers,
BaseNotifier) else approval_notifiers or []
Review Comment:
Agreed. Wrapped in `list(...)` and raise `TypeError` for non-BaseNotifier
elements at parse time.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py:
##########
@@ -83,7 +83,8 @@ class LLMSQLQueryOperator(LLMOperator):
Human-in-the-Loop approval parameters are inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
- (``require_approval``, ``approval_timeout``, ``allow_modifications``).
+ (``require_approval``, ``approval_timeout``, ``allow_modifications``,
+ ``approval_notifiers``).
Review Comment:
Added a sentence pointing to LLMOperator docs for the inherited approval
parameters.
--
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]
