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 42741924a79 Support assigned reviewers in LLM approval reviews (#72157)
42741924a79 is described below

commit 42741924a79f7e19dc96672a2c4f7f0bfe93beb2
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Fri Sep 18 16:01:57 2026 +0800

    Support assigned reviewers in LLM approval reviews (#72157)
    
    * Support assigned reviewers in LLM approval reviews
    
    The review was always open to every user with the HITL permission, so an
    operator could not route an approval to the people who actually own the
    decision. The Task SDK already accepts assigned users for a review; the
    LLM operators simply never passed any.
    
    * Validate assigned reviewers at parse time and document their matching 
rules
    
    A malformed approval_assigned_users entry only failed inside 
defer_for_approval, after the model call was billed, and the version guard 
tested the raw argument while the mixin tested the normalized one. The old-core 
guard test never ran on any lane because its skipif contradicted the 
class-level one. The docs did not say that id is the auth manager user id, that 
3.1.0 to 3.1.5 compare name as well, or that the list is fixed once the review 
row exists.
    
    * Name the default auth manager's user id in the assigned reviewers docs
---
 providers/common/ai/docs/operators/llm.rst         | 18 ++++++++++
 providers/common/ai/docs/operators/llm_branch.rst  |  6 ++--
 .../common/ai/docs/operators/llm_file_analysis.rst |  3 +-
 .../ai/docs/operators/llm_schema_compare.rst       |  5 ++-
 providers/common/ai/docs/operators/llm_sql.rst     |  5 +--
 .../airflow/providers/common/ai/mixins/approval.py | 11 ++++++
 .../airflow/providers/common/ai/operators/llm.py   | 30 ++++++++++++++++
 .../providers/common/ai/operators/llm_branch.py    |  2 +-
 .../common/ai/operators/llm_file_analysis.py       |  2 +-
 .../common/ai/operators/llm_schema_compare.py      |  2 +-
 .../providers/common/ai/operators/llm_sql.py       |  2 +-
 .../tests/unit/common/ai/mixins/test_approval.py   | 19 ++++++++++
 .../ai/tests/unit/common/ai/operators/test_llm.py  | 40 ++++++++++++++++++++++
 13 files changed, 135 insertions(+), 10 deletions(-)

diff --git a/providers/common/ai/docs/operators/llm.rst 
b/providers/common/ai/docs/operators/llm.rst
index 19bb702496c..8011d1f00ae 100644
--- a/providers/common/ai/docs/operators/llm.rst
+++ b/providers/common/ai/docs/operators/llm.rst
@@ -278,6 +278,20 @@ The default ``body`` contains the rendered prompt and the 
output.  Where either
 is sensitive, template only ``{{ task.subject }}`` and a link to the review
 into channels outside Airflow's auth boundary.
 
+By default any user with the permission can answer the review.  Pass
+``approval_assigned_users=[{"id": "<auth-manager-user-id>", "name": 
"<user-name>"}]``
+to restrict it to named reviewers, the way
+:class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with
+``assigned_users``.  ``id`` is the user id reported by the auth manager: with
+the default ``SimpleAuthManager`` it is the username from
+``simple_auth_manager_users``; under the FAB auth manager it is the numeric
+user row id as a string, not the username.  This needs Airflow 3.1+.  On 
Airflow 3.1.0 through 3.1.5 both
+``id`` and ``name`` must match what the auth manager reports, so a wrong
+``name`` blocks the assigned reviewer as well as everyone else; from 3.1.6 only
+``id`` is compared.  The list is stored when the review is first created:
+clearing the task re-runs it against the existing review row, so a changed
+list does not take effect.
+
 Parameters
 ----------
 
@@ -305,6 +319,10 @@ Parameters
   approving.  Default ``False``.
 - ``approval_notifiers``: Notifier, or list of notifiers, called once the 
review
   is open.  Default ``None``.
+- ``approval_assigned_users``: Users allowed to answer the review, as
+  ``{"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.  Needs Airflow 3.1+.
 
 Logging
 -------
diff --git a/providers/common/ai/docs/operators/llm_branch.rst 
b/providers/common/ai/docs/operators/llm_branch.rst
index 621c78efdc2..20ca0d748bf 100644
--- a/providers/common/ai/docs/operators/llm_branch.rst
+++ b/providers/common/ai/docs/operators/llm_branch.rst
@@ -111,8 +111,8 @@ call.
 
 Apart from ``fail_on_reject`` and ``ignore_downstream_trigger_rules``, which
 are specific to this operator, ``approval_timeout``, ``on_approval_timeout``,
-``approval_notifiers``, and the rest of the approval behaviour are inherited 
from
-:ref:`LLMOperator <howto/operator:llm>`.
+``approval_notifiers``, ``approval_assigned_users``, and the rest of the 
approval
+behaviour are inherited from :ref:`LLMOperator <howto/operator:llm>`.
 
 How It Works
 ------------
@@ -152,6 +152,8 @@ Parameters
   branch(es) before approving.  Default ``False``.
 - ``approval_notifiers``: Notifier, or list of notifiers, called once the 
review
   is open.  Default ``None``.
+- ``approval_assigned_users``: Users allowed to answer the review.  ``None``
+  (default) lets any user with the permission respond.  Needs Airflow 3.1+.
 - ``fail_on_reject``: If ``True``, a rejected review fails the task instead of
   skipping the downstream tasks.  Generally discouraged.  Only takes effect
   with ``require_approval=True``.  Default ``False``.
diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst 
b/providers/common/ai/docs/operators/llm_file_analysis.rst
index ad9da9dcc12..c92c8463896 100644
--- a/providers/common/ai/docs/operators/llm_file_analysis.rst
+++ b/providers/common/ai/docs/operators/llm_file_analysis.rst
@@ -163,7 +163,8 @@ Parameters
 
 This operator also inherits ``LLMOperator``'s HITL review parameters --
 ``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
-``allow_modifications``, and ``approval_notifiers`` -- see :doc:`llm` for 
details.
+``allow_modifications``, ``approval_notifiers``, and 
``approval_assigned_users`` --
+see :doc:`llm` for details.
 
 Supported Formats
 -----------------
diff --git a/providers/common/ai/docs/operators/llm_schema_compare.rst 
b/providers/common/ai/docs/operators/llm_schema_compare.rst
index 1132081862e..e7ae94e2c9b 100644
--- a/providers/common/ai/docs/operators/llm_schema_compare.rst
+++ b/providers/common/ai/docs/operators/llm_schema_compare.rst
@@ -133,7 +133,8 @@ returning a ``Sequence[UserContent]`` raises ``TypeError`` 
before the LLM
 call.
 
 ``approval_timeout``, ``on_approval_timeout``, ``allow_modifications``,
-``approval_notifiers``, and the rest of the approval behaviour are inherited 
from
+``approval_notifiers``, ``approval_assigned_users``, and the rest of the 
approval
+behaviour are inherited from
 :ref:`LLMOperator <howto/operator:llm>`.
 
 Conditional ETL Based on Schema Compatibility
@@ -202,6 +203,8 @@ Parameters
   before approving.  Default ``False``.
 - ``approval_notifiers``: Notifier, or list of notifiers, called once the 
review
   is open.  Default ``None``.
+- ``approval_assigned_users``: Users allowed to answer the review.  ``None``
+  (default) lets any user with the permission respond.  Needs Airflow 3.1+.
 
 Logging
 -------
diff --git a/providers/common/ai/docs/operators/llm_sql.rst 
b/providers/common/ai/docs/operators/llm_sql.rst
index 45b64c68b67..ba1d05f36de 100644
--- a/providers/common/ai/docs/operators/llm_sql.rst
+++ b/providers/common/ai/docs/operators/llm_sql.rst
@@ -157,8 +157,9 @@ modified query is re-validated against the same safety 
rules automatically.
     :start-after: [START howto_operator_llm_sql_approval]
     :end-before: [END howto_operator_llm_sql_approval]
 
-``approval_timeout``, ``approval_notifiers``, and the rest of the approval
-behaviour are inherited from :ref:`LLMOperator <howto/operator:llm>`.
+``approval_timeout``, ``approval_notifiers``, ``approval_assigned_users``, and
+the rest of the approval behaviour are inherited from
+:ref:`LLMOperator <howto/operator:llm>`.
 
 SQL Safety Validation
 ---------------------
diff --git 
a/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py 
b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
index 9a011a0d2e8..aa9d85c8707 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
@@ -39,6 +39,7 @@ if TYPE_CHECKING:
 
     from airflow.providers.common.compat.notifier import BaseNotifier
     from airflow.sdk import Context
+    from airflow.sdk.execution_time.hitl import HITLUser
 
 
 class DeferForApprovalProtocol(Protocol):
@@ -48,6 +49,7 @@ class DeferForApprovalProtocol(Protocol):
     allow_modifications: bool
     on_approval_timeout: Literal["fail", "approve", "reject"]
     approval_notifiers: Sequence[BaseNotifier]
+    approval_assigned_users: list[HITLUser]
     prompt: str
     task_id: str
     defer: Any
@@ -84,6 +86,13 @@ class LLMApprovalMixin:
     with the regenerated output, while the open review keeps the original
     subject and body.
 
+    ``approval_assigned_users`` restricts the review to the named users, the 
way
+    :class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with
+    ``assigned_users``.  Leaving it empty lets any user with the permission
+    respond.  The list is stored when the review is first created; clearing
+    the task re-runs it against the existing review row, so a changed list
+    does not take effect.
+
     Operators that use this mixin must set the following attributes:
 
     - ``require_approval`` (``bool``)
@@ -91,6 +100,7 @@ class LLMApprovalMixin:
     - ``approval_timeout`` (``timedelta | None``)
     - ``on_approval_timeout`` (``Literal["fail", "approve", "reject"]``)
     - ``approval_notifiers`` (``Sequence[BaseNotifier]``)
+    - ``approval_assigned_users`` (``list[HITLUser]``)
     - ``prompt`` (``str``)
     """
 
@@ -190,6 +200,7 @@ class LLMApprovalMixin:
             defaults=timeout_defaults,
             multiple=False,
             params=hitl_params,
+            assigned_users=self.approval_assigned_users,
         )
 
         self.subject = subject
diff --git 
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py 
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
index fb813b7744f..ede18548abc 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
@@ -48,6 +48,7 @@ if TYPE_CHECKING:
     from pydantic_ai.usage import UsageLimits
 
     from airflow.sdk import Context
+    from airflow.sdk.execution_time.hitl import HITLUser
 
 
 class LLMOperator(BaseOperator, LLMApprovalMixin):
@@ -118,6 +119,10 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
         ``require_approval=True``.  A retry re-notifies with the regenerated
         output while the open review keeps the original subject and body.
         Default ``None``.
+    :param approval_assigned_users: Users allowed to answer the review, as
+        ``{"id": ..., "name": ...}`` dicts where ``id`` is the auth manager's
+        user id.  ``None`` (default) lets any user with the permission respond.
+        The list is fixed when the review is first created.  Needs Airflow 
3.1+.
     :param serialize_output: If ``True`` and ``output_type`` is a Pydantic
         ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via
         ``model_dump()`` before being pushed to XCom. Default ``False`` --
@@ -152,6 +157,7 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
         on_approval_timeout: Literal["fail", "approve", "reject"] = "fail",
         allow_modifications: bool = False,
         approval_notifiers: BaseNotifier | Iterable[BaseNotifier] | None = 
None,
+        approval_assigned_users: HITLUser | Iterable[HITLUser] | None = None,
         serialize_output: bool = False,
         **kwargs: Any,
     ) -> None:
@@ -203,6 +209,30 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
         for notifier in self.approval_notifiers:
             if not isinstance(notifier, BaseNotifier):
                 raise TypeError(f"approval_notifiers must contain BaseNotifier 
instances, got {notifier!r}")
+        assigned_users: list[Any]
+        if approval_assigned_users is None:
+            assigned_users = []
+        elif isinstance(approval_assigned_users, dict):
+            assigned_users = [approval_assigned_users]
+        elif isinstance(approval_assigned_users, str) or not 
isinstance(approval_assigned_users, Iterable):
+            raise TypeError(
+                "approval_assigned_users must be a {'id': str, 'name': str} 
dict or an iterable of them, "
+                f"got {approval_assigned_users!r}"
+            )
+        else:
+            assigned_users = list(approval_assigned_users)
+        for user in assigned_users:
+            if (
+                not isinstance(user, dict)
+                or not isinstance(user.get("id"), str)
+                or not isinstance(user.get("name"), str)
+            ):
+                raise TypeError(
+                    f"approval_assigned_users entries must be {{'id': str, 
'name': str}} dicts, got {user!r}"
+                )
+        if assigned_users and not AIRFLOW_V_3_1_PLUS:
+            raise 
AirflowOptionalProviderFeatureException("approval_assigned_users needs Airflow 
3.1+.")
+        self.approval_assigned_users: list[HITLUser] = assigned_users
 
     @cached_property
     def llm_hook(self) -> PydanticAIHook:
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 75291ebe9f8..13d2287c3f1 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
@@ -65,7 +65,7 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
     Human-in-the-Loop approval parameters are inherited from
     :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
     (``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
-    ``allow_modifications``, ``approval_notifiers``).
+    ``allow_modifications``, ``approval_notifiers``, 
``approval_assigned_users``).
     The task pauses after the LLM chooses the branch(es) and only skips the
     unselected downstream tasks once a reviewer approves. Rejecting the
     review skips the direct downstream tasks except teardowns, matching
diff --git 
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
 
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
index 8245312018e..31d4635e654 100644
--- 
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
+++ 
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py
@@ -76,7 +76,7 @@ class LLMFileAnalysisOperator(LLMOperator):
     Human-in-the-Loop approval parameters are inherited from
     :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
     (``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
-    ``allow_modifications``, ``approval_notifiers``).
+    ``allow_modifications``, ``approval_notifiers``, 
``approval_assigned_users``).
     The task pauses after the file analysis and only returns the result once a
     reviewer approves.
     """
diff --git 
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
 
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
index f1f494bf263..ccb4a62cb8a 100644
--- 
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
+++ 
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
@@ -115,7 +115,7 @@ class LLMSchemaCompareOperator(LLMOperator):
     Human-in-the-Loop approval parameters are inherited from
     :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
     (``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
-    ``allow_modifications``, ``approval_notifiers``).
+    ``allow_modifications``, ``approval_notifiers``, 
``approval_assigned_users``).
     The task pauses after the comparison and only returns the result once a
     reviewer approves. The review body shows the compatibility verdict, a
     mismatch severity summary, and the full result JSON.
diff --git 
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py 
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
index e4415c2f8a5..96c6453e092 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
@@ -88,7 +88,7 @@ class LLMSQLQueryOperator(LLMOperator):
     Human-in-the-Loop approval parameters are inherited from
     :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
     (``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
-    ``allow_modifications``, ``approval_notifiers``).
+    ``allow_modifications``, ``approval_notifiers``, 
``approval_assigned_users``).
     When ``allow_modifications=True`` and the reviewer edits the SQL, the
     modified query is re-validated against the same safety rules before being
     returned.
diff --git a/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py 
b/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
index c9f18c6b188..960999f0495 100644
--- a/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
+++ b/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
@@ -72,6 +72,7 @@ class FakeOperator(LLMApprovalMixin):
         on_approval_timeout: str = "fail",
         allow_modifications: bool = False,
         approval_notifiers: Sequence[BaseNotifier] = (),
+        approval_assigned_users: list[dict[str, str]] | None = None,
     ):
         self.prompt = prompt
         self.task_id = task_id
@@ -79,6 +80,7 @@ class FakeOperator(LLMApprovalMixin):
         self.on_approval_timeout = on_approval_timeout
         self.allow_modifications = allow_modifications
         self.approval_notifiers = approval_notifiers
+        self.approval_assigned_users = list(approval_assigned_users or [])
 
         self.defer = MagicMock()
         self.log = MagicMock()
@@ -255,6 +257,23 @@ class TestDeferForApproval:
         healthy.assert_called_once_with(context)
         op.defer.assert_called_once()
 
+    @patch(HITL_TRIGGER_PATH, autospec=True)
+    @patch(UPSERT_HITL_PATH)
+    def test_assigned_users_are_forwarded(self, mock_upsert, mock_trigger_cls, 
context):
+        users = [{"id": "u1", "name": "alice"}]
+        op = FakeOperator(approval_assigned_users=users)
+
+        op.defer_for_approval(context, "output")
+
+        assert mock_upsert.call_args[1]["assigned_users"] == users
+
+    @patch(HITL_TRIGGER_PATH, autospec=True)
+    @patch(UPSERT_HITL_PATH)
+    def test_assigned_users_empty_when_unset(self, mock_upsert, 
mock_trigger_cls, approval_op, context):
+        approval_op.defer_for_approval(context, "output")
+
+        assert mock_upsert.call_args[1]["assigned_users"] == []
+
     @patch(HITL_TRIGGER_PATH, autospec=True)
     @patch(UPSERT_HITL_PATH)
     def test_no_modifications_params_empty(self, mock_upsert, 
mock_trigger_cls, approval_op, context):
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py 
b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
index 3bed2f0f16f..aec57fc8a44 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
@@ -362,6 +362,16 @@ class TestLLMOperatorApprovalVersionGate:
         op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c")
         assert op.require_approval is False
 
+    @patch("airflow.providers.common.ai.operators.llm.AIRFLOW_V_3_1_PLUS", 
False)
+    def test_approval_assigned_users_rejected_on_old_core(self):
+        with pytest.raises(AirflowOptionalProviderFeatureException, 
match="needs Airflow 3.1"):
+            LLMOperator(
+                task_id="t",
+                prompt="p",
+                llm_conn_id="c",
+                approval_assigned_users={"id": "u1", "name": "alice"},
+            )
+
 
 @pytest.mark.skipif(
     not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with 
Airflow >= 3.1.0"
@@ -379,6 +389,7 @@ class TestLLMOperatorApproval:
         assert op.approval_timeout is None
         assert op.on_approval_timeout == "fail"
         assert op.approval_notifiers == []
+        assert op.approval_assigned_users == []
 
     def test_unknown_on_approval_timeout_raises(self):
         with pytest.raises(ValueError, match="on_approval_timeout must be"):
@@ -436,6 +447,35 @@ class TestLLMOperatorApproval:
         op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c", 
approval_notifiers=iter(notifiers))
         assert op.approval_notifiers == notifiers
 
+    @pytest.mark.parametrize(
+        "assigned_users",
+        [
+            {"id": "u1", "name": "alice"},
+            [{"id": "u1", "name": "alice"}],
+            iter([{"id": "u1", "name": "alice"}]),
+        ],
+        ids=["single", "list", "generator"],
+    )
+    def test_approval_assigned_users_normalized_to_list(self, assigned_users):
+        op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c", 
approval_assigned_users=assigned_users)
+        assert op.approval_assigned_users == [{"id": "u1", "name": "alice"}]
+
+    @pytest.mark.parametrize(
+        ("assigned_users", "match"),
+        [
+            ("alice", r"dict or an iterable of them, got 'alice'"),
+            (5, r"dict or an iterable of them, got 5"),
+            ({}, r"entries must be \{'id': str, 'name': str\} dicts, got 
\{\}"),
+            ([{"id": "u1"}], r"entries must be .* got \{'id': 'u1'\}"),
+            ([{"id": 1, "name": "alice"}], r"entries must be .* got \{'id': 1, 
'name': 'alice'\}"),
+            (["alice"], r"entries must be .* got 'alice'"),
+        ],
+        ids=["str", "int", "empty_dict", "missing_name", "non_str_id", 
"list_of_str"],
+    )
+    def test_rejects_malformed_approval_assigned_users(self, assigned_users, 
match):
+        with pytest.raises(TypeError, match=match):
+            LLMOperator(task_id="t", prompt="p", llm_conn_id="c", 
approval_assigned_users=assigned_users)
+
     @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", 
autospec=True)
     @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail")
     @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", 
autospec=True)

Reply via email to