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")