From 3030164b60e5eb4486f336a12e1a34fb01c9b512 Mon Sep 17 00:00:00 2001 From: "Guan-Ming (Wesley) Chiu" <105915352+guan404ming@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:18:11 +0800 Subject: [PATCH] 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. --- providers/common/ai/docs/operators/llm.rst | 9 +++++++ .../common/ai/docs/operators/llm_branch.rst | 6 +++-- .../ai/docs/operators/llm_file_analysis.rst | 4 +-- .../ai/docs/operators/llm_schema_compare.rst | 7 +++-- .../providers/common/ai/mixins/approval.py | 15 +++++++++++ .../providers/common/ai/operators/llm.py | 15 ++++++++++- .../common/ai/operators/llm_branch.py | 3 ++- .../common/ai/operators/llm_file_analysis.py | 3 ++- .../common/ai/operators/llm_schema_compare.py | 3 ++- .../providers/common/ai/operators/llm_sql.py | 3 ++- .../unit/common/ai/mixins/test_approval.py | 19 ++++++++++++++ .../unit/common/ai/operators/test_llm.py | 26 ++++++++++++++++++- 12 files changed, 101 insertions(+), 12 deletions(-) diff --git a/providers/common/ai/docs/operators/llm.rst b/providers/common/ai/docs/operators/llm.rst index 426a2573206fc..a51a385b524b3 100644 --- a/providers/common/ai/docs/operators/llm.rst +++ b/providers/common/ai/docs/operators/llm.rst @@ -205,6 +205,12 @@ approving with ``allow_modifications=True``, and set a deadline with :start-after: [START howto_operator_llm_approval] :end-before: [END howto_operator_llm_approval] +By default any user with the permission can answer the review. Pass +``approval_assigned_users=[{"id": "", "name": ""}]`` to +restrict it to named reviewers, the way +:class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with +``assigned_users``. This needs Airflow 3.1+. + Parameters ---------- @@ -225,6 +231,9 @@ Parameters means wait indefinitely. Default ``None``. - ``allow_modifications``: If ``True``, the reviewer can edit the output before approving. Default ``False``. +- ``approval_assigned_users``: Users allowed to answer the review, as + ``{"id": ..., "name": ...}`` dicts. ``None`` (default) lets any user with the + permission respond. 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 1ef8c557abcf2..63e516b7a62fc 100644 --- a/providers/common/ai/docs/operators/llm_branch.rst +++ b/providers/common/ai/docs/operators/llm_branch.rst @@ -106,8 +106,8 @@ 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 `. +``approval_timeout``, ``approval_assigned_users``, and the rest of the approval +behaviour are inherited from :ref:`LLMOperator `. How It Works ------------ @@ -139,6 +139,8 @@ Parameters means wait indefinitely. Default ``None``. - ``allow_modifications``: If ``True``, the reviewer can change the chosen branch(es) before approving. Default ``False``. +- ``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. 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 ac688ebfaf415..61014e77fe16f 100644 --- a/providers/common/ai/docs/operators/llm_file_analysis.rst +++ b/providers/common/ai/docs/operators/llm_file_analysis.rst @@ -159,8 +159,8 @@ Parameters downstream consumer needs the dict shape. This operator also inherits ``LLMOperator``'s HITL review parameters -- -``require_approval``, ``approval_timeout``, and ``allow_modifications`` -- see -:doc:`llm` for details. +``require_approval``, ``approval_timeout``, ``allow_modifications``, 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 768d96280be55..afc79b3a51e58 100644 --- a/providers/common/ai/docs/operators/llm_schema_compare.rst +++ b/providers/common/ai/docs/operators/llm_schema_compare.rst @@ -132,8 +132,9 @@ expire, fails the task: returning a ``Sequence[UserContent]`` raises ``TypeError`` before the LLM call. -``approval_timeout``, ``allow_modifications``, and the rest of the approval -behaviour are inherited from :ref:`LLMOperator `. +``approval_timeout``, ``allow_modifications``, ``approval_assigned_users``, and +the rest of the approval behaviour are inherited from +:ref:`LLMOperator `. Conditional ETL Based on Schema Compatibility ---------------------------------------------- @@ -193,6 +194,8 @@ Parameters means wait indefinitely. Default ``None``. - ``allow_modifications``: If ``True``, the reviewer can edit the result JSON before approving. Default ``False``. +- ``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/src/airflow/providers/common/ai/mixins/approval.py b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py index cccff32bfd150..c506cf934542f 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 @@ -35,6 +35,7 @@ if TYPE_CHECKING: from airflow.sdk import Context + from airflow.sdk.execution_time.hitl import HITLUser class DeferForApprovalProtocol(Protocol): @@ -42,6 +43,7 @@ class DeferForApprovalProtocol(Protocol): approval_timeout: timedelta | None allow_modifications: bool + approval_assigned_users: list[HITLUser] | None prompt: str task_id: str defer: Any @@ -62,11 +64,17 @@ class LLMApprovalMixin: before approving. The (possibly modified) output is then returned as the task result. + ``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 unset lets any user with the permission + respond. + Operators that use this mixin must set the following attributes: - ``require_approval`` (``bool``) - ``allow_modifications`` (``bool``) - ``approval_timeout`` (``timedelta | None``) + - ``approval_assigned_users`` (``list[HITLUser] | None``) - ``prompt`` (``str``) """ @@ -155,6 +163,12 @@ def defer_for_approval( }, } + # Only pass assigned_users when set: cores before 3.2 have no such argument, and the + # operator has already rejected the parameter on those versions. + assignee_kwargs: dict[str, Any] = ( + {"assigned_users": self.approval_assigned_users} if self.approval_assigned_users else {} + ) + upsert_hitl_detail( ti_id=ti_id, options=[LLMApprovalMixin.APPROVE, LLMApprovalMixin.REJECT], @@ -163,6 +177,7 @@ def defer_for_approval( defaults=None, multiple=False, params=hitl_params, + **assignee_kwargs, ) if AIRFLOW_V_3_3_PLUS: 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 36cd616c1b386..617d19911b4ed 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 @@ -29,7 +29,8 @@ from airflow.providers.common.ai.mixins.approval import LLMApprovalMixin from airflow.providers.common.ai.utils.logging import log_run_summary from airflow.providers.common.ai.utils.output_type import rehydrate_pydantic_output -from airflow.providers.common.compat.sdk import BaseOperator +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, BaseOperator +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_1_PLUS try: # New enough cores register an operator's declared ``output_type`` classes for @@ -45,6 +46,7 @@ from pydantic_ai.usage import UsageLimits from airflow.sdk import Context + from airflow.sdk.execution_time.hitl import HITLUser class LLMOperator(BaseOperator, LLMApprovalMixin): @@ -88,6 +90,9 @@ class LLMOperator(BaseOperator, LLMApprovalMixin): :param allow_modifications: If ``True``, the reviewer can edit the output before approving. The modified value is returned as the task result. Default ``False``. + :param approval_assigned_users: Users allowed to answer the review, as + ``{"id": ..., "name": ...}`` dicts. ``None`` (default) lets any user + with the permission respond. 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`` -- @@ -119,6 +124,7 @@ def __init__( require_approval: bool = False, approval_timeout: timedelta | None = None, allow_modifications: bool = False, + approval_assigned_users: HITLUser | list[HITLUser] | None = None, serialize_output: bool = False, **kwargs: Any, ) -> None: @@ -135,9 +141,16 @@ def __init__( self._serialize_model_output = serialize_output or not _CORE_WALKER self.agent_params = agent_params or {} self.usage_limits = usage_limits + if approval_assigned_users and not AIRFLOW_V_3_1_PLUS: + raise AirflowOptionalProviderFeatureException("approval_assigned_users needs Airflow 3.1+.") self.require_approval = require_approval self.approval_timeout = approval_timeout self.allow_modifications = allow_modifications + self.approval_assigned_users = ( + [approval_assigned_users] + if isinstance(approval_assigned_users, dict) + else approval_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 5de42fb2db666..4672f44a444c1 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 @@ -56,7 +56,8 @@ 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``, ``allow_modifications``). + (``require_approval``, ``approval_timeout``, ``allow_modifications``, + ``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 7b757ba22caf7..93a021cb95784 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 @@ -71,7 +71,8 @@ class LLMFileAnalysisOperator(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_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 7992419296fa2..4043c8f5beef4 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 @@ -110,7 +110,8 @@ class LLMSchemaCompareOperator(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_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 262d52ff8d0f4..127eec309ecc4 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 @@ -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_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 ace3ce7d3d98e..f83f2a923991a 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 @@ -53,11 +53,13 @@ def __init__( task_id: str = "test_task", approval_timeout: timedelta | None = None, allow_modifications: bool = False, + approval_assigned_users: list[dict[str, str]] | None = None, ): self.prompt = prompt self.task_id = task_id self.approval_timeout = approval_timeout self.allow_modifications = allow_modifications + self.approval_assigned_users = approval_assigned_users self.defer = MagicMock() self.log = MagicMock() @@ -174,6 +176,23 @@ 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_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_omitted_when_unset(self, mock_upsert, mock_trigger_cls, approval_op, context): + approval_op.defer_for_approval(context, "output") + + assert "assigned_users" not in mock_upsert.call_args[1] + @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 e2004b4031c32..bb6945b6bfae1 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 @@ -36,7 +36,10 @@ except ImportError: _CORE_WALKER = False -from airflow.providers.common.compat.sdk import TaskDeferred +from airflow.providers.common.compat.sdk import ( + AirflowOptionalProviderFeatureException, + TaskDeferred, +) if AIRFLOW_V_3_3_PLUS: # On 3.3+ cores require_approval pauses the task in AWAITING_INPUT; older cores defer @@ -193,6 +196,27 @@ def test_default_approval_flags(self): assert op.require_approval is False assert op.allow_modifications is False assert op.approval_timeout is None + assert op.approval_assigned_users is None + + @pytest.mark.skipif(not AIRFLOW_V_3_1_PLUS, reason="assigned_users needs Airflow 3.1+") + @pytest.mark.parametrize( + "assigned_users", + [{"id": "u1", "name": "alice"}, [{"id": "u1", "name": "alice"}]], + ids=["single", "list"], + ) + 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.skipif(AIRFLOW_V_3_1_PLUS, reason="guard only fires on cores before 3.1") + 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"}, + ) @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail")