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 0aeec0c7962 Render multi-branch review choices as a multi-select
(#71046)
0aeec0c7962 is described below
commit 0aeec0c79626ae0deada0022636601419526a2ef
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Wed Aug 5 02:37:29 2026 +0900
Render multi-branch review choices as a multi-select (#71046)
* Render multi-branch review choices as a multi-select
* Fix multi-select review form validation
* Fix test imports for older Airflow compat
---
providers/common/ai/docs/operators/llm_branch.rst | 4 +-
.../airflow/providers/common/ai/mixins/approval.py | 40 ++++++++++++--
.../providers/common/ai/operators/llm_branch.py | 6 ++-
.../tests/unit/common/ai/mixins/test_approval.py | 62 ++++++++++++++++++++++
.../unit/common/ai/operators/test_llm_branch.py | 42 +++++++++++++--
5 files changed, 143 insertions(+), 11 deletions(-)
diff --git a/providers/common/ai/docs/operators/llm_branch.rst
b/providers/common/ai/docs/operators/llm_branch.rst
index 52affb9af86..21558b2277f 100644
--- a/providers/common/ai/docs/operators/llm_branch.rst
+++ b/providers/common/ai/docs/operators/llm_branch.rst
@@ -83,8 +83,8 @@ branch(es) and wait for a human reviewer to approve the
choice before any
downstream task is skipped. The review form shows the LLM's choice and the
valid downstream task IDs. When ``allow_modifications=True``, the reviewer
can also change the choice — rendered as a dropdown of the downstream task
-IDs, or a free-text JSON list of task IDs (e.g. ``["task_a", "task_b"]``)
-with ``allow_multiple_branches=True``. The reviewed branch(es) are validated
+IDs, or a multi-select of them with ``allow_multiple_branches=True``. The
+reviewed branch(es) are validated
against the downstream task IDs before branching:
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py
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 21c2e4c9de4..cccff32bfd1 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
@@ -17,6 +17,7 @@
from __future__ import annotations
+import json
import logging
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Protocol
@@ -109,7 +110,11 @@ class LLMApprovalMixin:
:param modification_schema: JSON schema for the editable ``output``
param
when ``allow_modifications=True``. Defaults to ``{"type":
"string"}``.
Pass e.g. ``{"type": "string", "enum": [...]}`` to render a
dropdown
- of valid values in the review form.
+ of valid values in the review form, or ``{"type": "array", "items":
+ {"type": "string", "enum": [...]}, "examples": [...]}`` to render a
+ multi-select (JSON Schema forbids ``enum`` at the array level, so
the
+ options come from ``examples``); a list submitted by the reviewer
is
+ returned from ``execute_complete`` re-serialized as a compact JSON
string.
"""
from airflow.providers.standard.triggers.hitl import HITLTrigger
from airflow.sdk.execution_time.hitl import upsert_hitl_detail
@@ -117,6 +122,7 @@ class LLMApprovalMixin:
self.validate_approval_prompt()
+ raw_output = output
if isinstance(output, BaseModel):
output = output.model_dump_json()
elif not isinstance(output, str):
@@ -133,9 +139,17 @@ class LLMApprovalMixin:
hitl_params: dict[str, dict[str, Any]] = {}
if self.allow_modifications:
+ # The multi-select rendered for an array schema needs the list,
not its JSON string
+ param_value: Any = output
+ if (
+ modification_schema is not None
+ and modification_schema.get("type") == "array"
+ and isinstance(raw_output, list)
+ ):
+ param_value = raw_output
hitl_params = {
"output": {
- "value": output,
+ "value": param_value,
"description": "Edit the output before approving
(optional).",
"schema": modification_schema or {"type": "string"},
},
@@ -215,13 +229,33 @@ class LLMApprovalMixin:
# when allow_modifications=False, bypassing the read-only approval
flow.
if getattr(self, "allow_modifications", False) and params_input:
modified = params_input.get("output")
+ if "output" in params_input and modified is None:
+ raise HITLTriggerEventError(
+ {
+ "error": "Modified output must not be empty; edit it
or reject instead.",
+ "error_type": "validation",
+ }
+ )
+ if isinstance(modified, list):
+ for item in modified:
+ if not isinstance(item, str):
+ raise HITLTriggerEventError(
+ {
+ "error": f"Modified output list items must be
strings, "
+ f"got {type(item).__name__}.",
+ "error_type": "validation",
+ }
+ )
+ # Compact so an unchanged selection compares equal to
generated_output
+ modified = json.dumps(modified, separators=(",", ":"))
if modified is not None and not isinstance(modified, str):
# On the awaiting_input path nothing upstream schema-validates
params_input
# (HITLTrigger did on the legacy path), so enforce the string
contract here
# rather than returning a non-string as the task's output.
raise HITLTriggerEventError(
{
- "error": f"Modified output must be a string, got
{type(modified).__name__}.",
+ "error": f"Modified output must be a string or a list
of strings, "
+ f"got {type(modified).__name__}.",
"error_type": "validation",
}
)
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 620b36e6809..88c29943259 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,7 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
unselected downstream tasks once a reviewer approves. The review form
lists the valid downstream task IDs; with ``allow_modifications=True``
the editable choice is rendered as a dropdown of those IDs (single-branch
- mode) or a free-text JSON list (``allow_multiple_branches=True``), and
+ mode) or a multi-select of them (``allow_multiple_branches=True``), and
the reviewed branch(es) are validated against the downstream task IDs
before branching.
"""
@@ -121,7 +121,9 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
f"```\nPrompt: {self.prompt}\n\nChosen branch(es):
{chosen}\n```"
)
modification_schema = (
- None if self.allow_multiple_branches else {"type": "string",
"enum": choices}
+ {"type": "array", "items": {"type": "string", "enum":
choices}, "examples": choices}
+ if self.allow_multiple_branches
+ else {"type": "string", "enum": choices}
)
self.defer_for_approval( # type: ignore[misc]
context, branches, body=body,
modification_schema=modification_schema
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 23c97bd2470..ace3ce7d3d9 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
@@ -158,6 +158,22 @@ class TestDeferForApproval:
param = mock_upsert.call_args[1]["params"]["output"]
assert param["schema"] == schema
+ @patch(HITL_TRIGGER_PATH, autospec=True)
+ @patch(UPSERT_HITL_PATH)
+ def test_array_schema_passes_list_param_value(
+ self, mock_upsert, mock_trigger_cls, approval_op_with_modifications,
context
+ ):
+ choices = ["task_a", "task_b"]
+ schema = {"type": "array", "items": {"type": "string", "enum":
choices}, "examples": choices}
+
+ approval_op_with_modifications.defer_for_approval(context, ["task_a"],
modification_schema=schema)
+
+ param = mock_upsert.call_args[1]["params"]["output"]
+ assert param["value"] == ["task_a"]
+ assert param["schema"] == schema
+ 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_no_modifications_params_empty(self, mock_upsert,
mock_trigger_cls, approval_op, context):
@@ -292,6 +308,52 @@ class TestDeferForApproval:
{}, generated_output="original output", event=event
)
+ def test_approved_with_list_modified_output_is_serialized(self,
approval_op_with_modifications):
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": "editor",
+ "params_input": {"output": ["task_b", "task_c"]},
+ }
+
+ result = approval_op_with_modifications.execute_complete(
+ {}, generated_output='["task_a"]', event=event
+ )
+
+ assert result == '["task_b","task_c"]'
+
+ def test_approved_with_unmodified_list_output_returns_original(self,
approval_op_with_modifications):
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": "editor",
+ "params_input": {"output": ["task_a"]},
+ }
+
+ result = approval_op_with_modifications.execute_complete(
+ {}, generated_output='["task_a"]', event=event
+ )
+
+ assert result == '["task_a"]'
+
+ def test_approved_with_non_string_list_items_raises(self,
approval_op_with_modifications):
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": "editor",
+ "params_input": {"output": ["task_a", 2]},
+ }
+
+ with pytest.raises(HITLTriggerEventError, match="items must be
strings, got int"):
+ approval_op_with_modifications.execute_complete({},
generated_output='["task_a"]', event=event)
+
+ def test_approved_with_cleared_output_raises(self,
approval_op_with_modifications):
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": "editor",
+ "params_input": {"output": None},
+ }
+
+ with pytest.raises(HITLTriggerEventError, match="must not be empty"):
+ approval_op_with_modifications.execute_complete({},
generated_output='["task_a"]', event=event)
+
def test_approved_with_unmodified_output(self,
approval_op_with_modifications):
event = {
"chosen_options": ["Approve"],
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
index 78107daf138..7f2d11bf263 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
@@ -25,7 +25,7 @@ import pytest
from airflow.providers.common.ai.mixins.approval import LLMApprovalMixin
from airflow.providers.common.ai.operators.llm import LLMOperator
from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator
-from airflow.providers.common.compat.sdk import TaskDeferred
+from airflow.providers.common.compat.sdk import Param, ParamValidationError,
TaskDeferred
from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_3_PLUS
@@ -313,10 +313,10 @@ class TestLLMBranchOperatorApproval:
@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)
- def test_review_form_multi_branch_keeps_string_schema(
+ def test_review_form_multi_branch_renders_multiselect(
self, mock_hook_cls, mock_upsert, mock_trigger_cls, mock_do_branch
):
- """With allow_multiple_branches the editable param stays free-text
(JSON list)."""
+ """With allow_multiple_branches the editable param is an array enum
(multi-select)."""
downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a",
"task_b": "task_b"})
mock_agent = MagicMock(spec=["run_sync"])
@@ -338,7 +338,17 @@ class TestLLMBranchOperatorApproval:
call_kwargs = mock_upsert.call_args.kwargs
assert "Valid branches: `task_a`, `task_b`" in call_kwargs["body"]
- assert call_kwargs["params"]["output"]["schema"] == {"type": "string"}
+ assert call_kwargs["params"]["output"]["schema"] == {
+ "type": "array",
+ "items": {"type": "string", "enum": ["task_a", "task_b"]},
+ "examples": ["task_a", "task_b"],
+ }
+ assert call_kwargs["params"]["output"]["value"] == ["task_a"]
+
+ schema = call_kwargs["params"]["output"]["schema"]
+ assert Param(schema=schema).resolve(["task_a"]) == ["task_a"]
+ with pytest.raises(ParamValidationError):
+ Param(schema=schema).resolve(["task_x"])
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_rejects_sequence_prompt_with_require_approval(self,
mock_hook_cls):
@@ -405,6 +415,30 @@ class TestLLMBranchOperatorApproval:
assert result == "task_b"
mock_do_branch.assert_called_once_with(ctx, "task_b")
+ @patch.object(LLMBranchOperator, "do_branch")
+ def test_execute_complete_with_multiselect_modified_branches(self,
mock_do_branch):
+ """A list submitted by the multi-select review form branches into
those tasks."""
+ mock_do_branch.return_value = ["task_b", "task_c"]
+ op = LLMBranchOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ allow_multiple_branches=True,
+ allow_modifications=True,
+ )
+ op.downstream_task_ids = {"task_a", "task_b", "task_c"}
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": "admin",
+ "params_input": {"output": ["task_b", "task_c"]},
+ }
+ ctx = _make_context()
+
+ result = op.execute_complete(ctx, generated_output='["task_a"]',
event=event)
+
+ assert result == ["task_b", "task_c"]
+ mock_do_branch.assert_called_once_with(ctx, ["task_b", "task_c"])
+
@patch.object(LLMBranchOperator, "do_branch")
def test_execute_complete_rejects_invalid_modified_branch(self,
mock_do_branch):
"""A reviewer-modified branch outside downstream_task_ids fails
validation."""
