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 bbb3c066109 Support require_approval in LLMBranchOperator (#70651)
bbb3c066109 is described below
commit bbb3c0661094ceeec1e549f5d6aa8b0de741b9cc
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Tue Aug 4 12:37:07 2026 +0900
Support require_approval in LLMBranchOperator (#70651)
* Support require_approval in LLMBranchOperator
* Improve LLMBranchOperator approval review form
* List valid branches above the fenced review prompt
* Fail LLM branch runs that select no branches
---
providers/common/ai/docs/operators/llm_branch.rst | 34 +++
.../providers/common/ai/decorators/llm_branch.py | 11 +-
.../common/ai/example_dags/example_llm_branch.py | 35 +++
.../airflow/providers/common/ai/mixins/approval.py | 30 +-
.../airflow/providers/common/ai/operators/llm.py | 9 +-
.../providers/common/ai/operators/llm_branch.py | 72 ++++-
.../providers/common/ai/operators/llm_sql.py | 9 +-
.../unit/common/ai/decorators/test_llm_branch.py | 19 ++
.../tests/unit/common/ai/mixins/test_approval.py | 12 +
.../unit/common/ai/operators/test_llm_branch.py | 314 ++++++++++++++++++++-
10 files changed, 510 insertions(+), 35 deletions(-)
diff --git a/providers/common/ai/docs/operators/llm_branch.rst
b/providers/common/ai/docs/operators/llm_branch.rst
index 42db970427a..52affb9af86 100644
--- a/providers/common/ai/docs/operators/llm_branch.rst
+++ b/providers/common/ai/docs/operators/llm_branch.rst
@@ -75,6 +75,34 @@ With multiple branches:
:start-after: [START howto_decorator_llm_branch_multi]
:end-before: [END howto_decorator_llm_branch_multi]
+Human-in-the-Loop Approval
+--------------------------
+
+Set ``require_approval=True`` to pause the task after the LLM chooses the
+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
+against the downstream task IDs before branching:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py
+ :language: python
+ :start-after: [START howto_operator_llm_branch_approval]
+ :end-before: [END howto_operator_llm_branch_approval]
+
+Rejecting the review, or letting ``approval_timeout`` expire, **fails** the
+task (``HITLRejectException`` / ``HITLTimeoutError``), so downstream tasks
+end up ``upstream_failed`` rather than skipped.
+
+``require_approval=True`` requires a string prompt: a decorated callable
+returning a ``Sequence[UserContent]`` raises ``TypeError`` before the LLM
+call.
+
+``approval_timeout`` and the rest of the approval behaviour are inherited
+from :ref:`LLMOperator <howto/operator:llm>`.
+
How It Works
------------
@@ -99,6 +127,12 @@ Parameters
task ID. When ``True`` the LLM may return one or more task IDs.
- ``agent_params``: Additional keyword arguments passed to the pydantic-ai
``Agent``
constructor (e.g. ``retries``, ``model_settings``). Supports Jinja
templating.
+- ``require_approval``: If ``True``, the task pauses after the LLM chooses the
+ branch(es) and waits for human review before branching. Default ``False``.
+- ``approval_timeout``: Maximum time to wait for a review (``timedelta``).
``None``
+ means wait indefinitely. Default ``None``.
+- ``allow_modifications``: If ``True``, the reviewer can change the chosen
+ branch(es) before approving. Default ``False``.
Logging
-------
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_branch.py
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_branch.py
index 664036ae135..f56f1c4ac0d 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_branch.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_branch.py
@@ -28,7 +28,10 @@ from collections.abc import Callable, Collection, Mapping,
Sequence
from typing import TYPE_CHECKING, Any, ClassVar
from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator
-from airflow.providers.common.ai.utils.validation import validate_prompt
+from airflow.providers.common.ai.utils.validation import (
+ reject_sequence_with_unsupported_feature,
+ validate_prompt,
+)
from airflow.providers.common.compat.sdk import (
DecoratedOperator,
TaskDecorator,
@@ -89,6 +92,12 @@ class _LLMBranchDecoratedOperator(DecoratedOperator,
LLMBranchOperator):
self.prompt = self.python_callable(*self.op_args, **kwargs)
validate_prompt(self.prompt, decorator_name="@task.llm_branch")
+ reject_sequence_with_unsupported_feature(
+ self.prompt,
+ decorator_name="@task.llm_branch",
+ feature_name="require_approval",
+ feature_enabled=self.require_approval,
+ )
self.render_template_fields(context)
return LLMBranchOperator.execute(self, context)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py
index 9dceb685071..c94e984a011 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_branch.py
@@ -18,6 +18,8 @@
from __future__ import annotations
+from datetime import timedelta
+
from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator
from airflow.providers.common.compat.sdk import dag, task
@@ -150,3 +152,36 @@ def example_llm_branch_decorator_multi():
# [END howto_decorator_llm_branch_multi]
example_llm_branch_decorator_multi()
+
+
+# [START howto_operator_llm_branch_approval]
+@dag(tags=["example"])
+def example_llm_branch_approval():
+ route = LLMBranchOperator(
+ task_id="route_with_approval",
+ prompt="User says: 'I was charged twice for my subscription.'",
+ llm_conn_id="pydanticai_default",
+ system_prompt="Route support tickets to the right team.",
+ require_approval=True,
+ approval_timeout=timedelta(hours=24),
+ allow_modifications=True,
+ )
+
+ @task
+ def handle_billing():
+ return "Handling billing issue"
+
+ @task
+ def handle_auth():
+ return "Handling auth issue"
+
+ @task
+ def handle_general():
+ return "Handling general issue"
+
+ route >> [handle_billing(), handle_auth(), handle_general()]
+
+
+# [END howto_operator_llm_branch_approval]
+
+example_llm_branch_approval()
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 9f045c02548..21c2e4c9de4 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
@@ -45,6 +45,8 @@ class DeferForApprovalProtocol(Protocol):
task_id: str
defer: Any
+ def validate_approval_prompt(self) -> None: ...
+
class LLMApprovalMixin:
"""
@@ -70,6 +72,18 @@ class LLMApprovalMixin:
APPROVE = "Approve"
REJECT = "Reject"
+ def validate_approval_prompt(self: DeferForApprovalProtocol) -> None:
+ """Fail fast when the prompt cannot be rendered as text in the
approval review body."""
+ if not isinstance(self.prompt, str):
+ raise TypeError(
+ f"{type(self).__name__}: require_approval=True is not
supported "
+ f"with a non-string prompt (got {type(self.prompt).__name__}).
"
+ "The approval review body renders the prompt as text; passing
a "
+ "Sequence[UserContent] would expose object reprs (and any
embedded "
+ "bytes) in the human review UI. Return a str prompt, or
disable "
+ "require_approval."
+ )
+
def defer_for_approval(
self: DeferForApprovalProtocol,
context: Context,
@@ -77,6 +91,7 @@ class LLMApprovalMixin:
*,
subject: str | None = None,
body: str | None = None,
+ modification_schema: dict[str, Any] | None = None,
) -> None:
"""
Write HITL detail, then pause the task for human review.
@@ -91,19 +106,16 @@ class LLMApprovalMixin:
Defaults to ``"Review output for task `<task_id>`"``.
:param body: Markdown body shown below the headline.
Defaults to the prompt and output wrapped in a code block.
+ :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.
"""
from airflow.providers.standard.triggers.hitl import HITLTrigger
from airflow.sdk.execution_time.hitl import upsert_hitl_detail
from airflow.sdk.timezone import utcnow
- if not isinstance(self.prompt, str):
- raise TypeError(
- "require_approval=True is not supported with a non-string
prompt. "
- "The approval review body renders the prompt as text; passing
a "
- "Sequence[UserContent] would expose object reprs (and any
embedded "
- "bytes) in the human review UI. Return a str prompt, or
disable "
- "require_approval."
- )
+ self.validate_approval_prompt()
if isinstance(output, BaseModel):
output = output.model_dump_json()
@@ -125,7 +137,7 @@ class LLMApprovalMixin:
"output": {
"value": output,
"description": "Edit the output before approving
(optional).",
- "schema": {"type": "string"},
+ "schema": modification_schema or {"type": "string"},
},
}
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 95f33bffffa..36cd616c1b3 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
@@ -155,13 +155,8 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
return PydanticAIHook.get_hook(self.llm_conn_id,
hook_params=hook_params)
def execute(self, context: Context) -> Any:
- if self.require_approval and not isinstance(self.prompt, str):
- raise TypeError(
- f"{type(self).__name__}: require_approval=True is not
supported "
- f"with a non-string prompt (got {type(self.prompt).__name__}).
"
- f"The approval review body renders the prompt as text. Return
a "
- f"str prompt, or disable require_approval."
- )
+ if self.require_approval:
+ self.validate_approval_prompt() # type: ignore[misc]
agent: Agent[object, Any] = self.llm_hook.create_agent(
output_type=self.output_type, instructions=self.system_prompt,
**self.agent_params
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 b6a25a7811e..620b36e6809 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
@@ -18,6 +18,7 @@
from __future__ import annotations
+import json
from collections.abc import Iterable, Sequence
from enum import Enum
from typing import TYPE_CHECKING, Any
@@ -47,6 +48,17 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
single task ID. When ``True`` the LLM may return one or more task IDs.
:param agent_params: Additional keyword arguments passed to the pydantic-ai
``Agent`` constructor (e.g. ``retries``, ``model_settings``,
``tools``).
+
+ Human-in-the-Loop approval parameters are inherited from
+ :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
+ (``require_approval``, ``approval_timeout``, ``allow_modifications``).
+ The task pauses after the LLM chooses the branch(es) and only skips the
+ 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
+ the reviewed branch(es) are validated against the downstream task IDs
+ before branching.
"""
inherits_from_skipmixin = True
@@ -60,12 +72,13 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
**kwargs: Any,
) -> None:
kwargs.pop("output_type", None)
- if kwargs.get("require_approval"):
- raise ValueError("require_approval=True is not supported by
LLMBranchOperator.")
super().__init__(**kwargs)
self.allow_multiple_branches = allow_multiple_branches
def execute(self, context: Context) -> str | Iterable[str] | None:
+ if self.require_approval:
+ self.validate_approval_prompt() # type: ignore[misc]
+
if not self.downstream_task_ids:
raise ValueError(
f"{self.task_id!r} has no downstream tasks. "
@@ -95,4 +108,59 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
else:
branches = str(output)
+ if not branches:
+ raise ValueError(
+ f"LLM selected no branches for {self.task_id!r}, which would
skip every downstream task."
+ )
+
+ if self.require_approval:
+ choices = sorted(self.downstream_task_ids)
+ chosen = branches if isinstance(branches, str) else
json.dumps(branches)
+ body = (
+ f"Valid branches: {', '.join(f'`{c}`' for c in choices)}\n\n"
+ f"```\nPrompt: {self.prompt}\n\nChosen branch(es):
{chosen}\n```"
+ )
+ modification_schema = (
+ None 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
+ )
+
return self.do_branch(context, branches)
+
+ def execute_complete(self, context: Context, generated_output: str, event:
dict[str, Any]) -> Any:
+ """Resume after human review, validating the reviewed choice before
branching."""
+ output = super().execute_complete(context, generated_output, event)
+ branches = self._parse_reviewed_branches(output)
+ selected = {branches} if isinstance(branches, str) else set(branches)
+ invalid = selected - self.downstream_task_ids
+ if invalid:
+ raise ValueError(
+ f"Reviewed branch(es) {sorted(invalid)} are not downstream
tasks of "
+ f"{self.task_id!r}. Valid choices:
{sorted(self.downstream_task_ids)}."
+ )
+ return self.do_branch(context, branches)
+
+ def _parse_reviewed_branches(self, output: str) -> str | list[str]:
+ if not self.allow_multiple_branches:
+ return output
+ try:
+ branches = json.loads(output)
+ except json.JSONDecodeError as e:
+ raise ValueError(
+ f"Reviewed output {output!r} is not valid JSON. With "
+ f"allow_multiple_branches=True the reviewed output must be a "
+ f'JSON list of task IDs, e.g. ["task_a", "task_b"].'
+ ) from e
+ if not isinstance(branches, list) or not all(isinstance(b, str) for b
in branches):
+ raise ValueError(
+ f"Reviewed output {output!r} must be a JSON list of task ID
strings, "
+ f'e.g. ["task_a", "task_b"].'
+ )
+ if not branches:
+ raise ValueError(
+ "Reviewed output selects no branches, which would skip every
downstream "
+ "task. Select at least one task ID, or reject the review
instead."
+ )
+ return branches
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 e8a3b327b42..262d52ff8d0 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
@@ -134,13 +134,8 @@ class LLMSQLQueryOperator(LLMOperator):
return hook
def execute(self, context: Context) -> str:
- if self.require_approval and not isinstance(self.prompt, str):
- raise TypeError(
- f"{type(self).__name__}: require_approval=True is not
supported "
- f"with a non-string prompt (got {type(self.prompt).__name__}).
"
- f"The approval review body renders the prompt as text. Return
a "
- f"str prompt, or disable require_approval."
- )
+ if self.require_approval:
+ self.validate_approval_prompt() # type: ignore[misc]
schema_info = self._get_schema_context()
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
index 37238bff00b..2243315dcf3 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
@@ -83,6 +83,25 @@ class TestLLMBranchDecoratedOperator:
with pytest.raises(TypeError, match="must be"):
op.execute(context={})
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def
test_sequence_prompt_with_require_approval_raises_before_run_sync(self,
mock_hook_cls):
+ """Sequence prompt + require_approval=True fails before the agent
runs."""
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = _LLMBranchDecoratedOperator(
+ task_id="test",
+ python_callable=lambda: ["x",
ImageUrl(url="https://example.com/x.png";)],
+ llm_conn_id="my_llm",
+ require_approval=True,
+ )
+ op.downstream_task_ids = {"positive"}
+
+ with pytest.raises(TypeError, match="require_approval=True"):
+ op.execute(context={})
+
+ mock_agent.run_sync.assert_not_called()
+
@patch.object(LLMBranchOperator, "do_branch")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_accepts_sequence_prompt(self, mock_hook_cls,
mock_do_branch):
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 46120b3c3c3..23c97bd2470 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
@@ -146,6 +146,18 @@ class TestDeferForApproval:
assert param["value"] == "draft text"
assert param["schema"] == {"type": "string"}
+ @patch(HITL_TRIGGER_PATH, autospec=True)
+ @patch(UPSERT_HITL_PATH)
+ def test_modification_schema_overrides_output_param_schema(
+ self, mock_upsert, mock_trigger_cls, approval_op_with_modifications,
context
+ ):
+ schema = {"type": "string", "enum": ["task_a", "task_b"]}
+
+ approval_op_with_modifications.defer_for_approval(context, "task_a",
modification_schema=schema)
+
+ param = mock_upsert.call_args[1]["params"]["output"]
+ assert param["schema"] == schema
+
@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_branch.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
index fccf5e2f837..78107daf138 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
@@ -18,11 +18,24 @@ from __future__ import annotations
from enum import Enum
from unittest.mock import MagicMock, patch
+from uuid import uuid4
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 tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_3_PLUS
+
+if AIRFLOW_V_3_3_PLUS:
+ # On 3.3+ cores require_approval pauses the task in AWAITING_INPUT; older
cores defer to
+ # HITLTrigger. Both signals carry method_name/kwargs/timeout, so the
approval tests assert
+ # against whichever pause signal the running core uses.
+ from airflow.sdk.exceptions import TaskAwaitingInput as ApprovalPauseSignal
+else:
+ ApprovalPauseSignal = TaskDeferred # type: ignore[assignment, misc]
def _make_mock_run_result(output):
@@ -54,15 +67,6 @@ class TestLLMBranchOperator:
# the real output_type is built dynamically from downstream_task_ids
assert op.output_type is str
- def test_require_approval_rejected(self):
- with pytest.raises(ValueError, match="require_approval=True is not
supported"):
- LLMBranchOperator(
- task_id="test",
- prompt="pick a branch",
- llm_conn_id="my_llm",
- require_approval=True,
- )
-
@patch.object(LLMBranchOperator, "do_branch")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_single_branch(self, mock_hook_cls, mock_do_branch):
@@ -117,6 +121,26 @@ class TestLLMBranchOperator:
assert result == ["task_a", "task_c"]
mock_do_branch.assert_called_once_with(ctx, ["task_a", "task_c"])
+ @patch.object(LLMBranchOperator, "do_branch")
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_execute_rejects_empty_branch_selection(self, mock_hook_cls,
mock_do_branch):
+ """LLM returning an empty list fails instead of skipping every
downstream task."""
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.return_value = _make_mock_run_result([])
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMBranchOperator(
+ task_id="test",
+ prompt="Pick branches",
+ llm_conn_id="my_llm",
+ allow_multiple_branches=True,
+ )
+ op.downstream_task_ids = {"task_a", "task_b"}
+
+ with pytest.raises(ValueError, match="selected no branches"):
+ op.execute(MagicMock())
+ mock_do_branch.assert_not_called()
+
@patch.object(LLMBranchOperator, "do_branch")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_system_prompt_forwarded(self, mock_hook_cls, mock_do_branch):
@@ -173,3 +197,275 @@ class TestLLMBranchOperator:
)
with pytest.raises(ValueError, match="no downstream tasks"):
op.execute(MagicMock())
+
+
+def _make_context(ti_id=None):
+ ti_id = ti_id or uuid4()
+ ti = MagicMock()
+ ti.id = ti_id
+ return MagicMock(**{"__getitem__": lambda self, key: {"task_instance":
ti}[key]})
+
+
[email protected](
+ not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with
Airflow >= 3.1.0"
+)
+class TestLLMBranchOperatorApproval:
+ """Tests for LLMBranchOperator with require_approval=True
(LLMApprovalMixin integration)."""
+
+ def test_inherits_llm_approval_mixin(self):
+ assert issubclass(LLMBranchOperator, LLMApprovalMixin)
+
+ @patch.object(LLMBranchOperator, "do_branch")
+ @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_execute_with_approval_pauses_before_branching(
+ self, mock_hook_cls, mock_upsert, mock_trigger_cls, mock_do_branch
+ ):
+ """When require_approval=True, execute() pauses after the LLM choice,
before do_branch."""
+ downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a",
"task_b": "task_b"})
+
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.return_value =
_make_mock_run_result(downstream_enum.task_a)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMBranchOperator(
+ task_id="branch_approval",
+ prompt="Pick a branch",
+ llm_conn_id="my_llm",
+ require_approval=True,
+ )
+ op.downstream_task_ids = {"task_a", "task_b"}
+
+ with pytest.raises(ApprovalPauseSignal) as exc_info:
+ op.execute(_make_context())
+
+ assert exc_info.value.method_name == "execute_complete"
+ assert exc_info.value.kwargs["generated_output"] == "task_a"
+ mock_upsert.assert_called_once()
+ mock_do_branch.assert_not_called()
+
+ @patch.object(LLMBranchOperator, "do_branch")
+ @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_execute_with_approval_serializes_multiple_branches(
+ self, mock_hook_cls, mock_upsert, mock_trigger_cls, mock_do_branch
+ ):
+ """With allow_multiple_branches=True the choice is deferred as a JSON
list."""
+ downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a",
"task_c": "task_c"})
+
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.return_value = _make_mock_run_result(
+ [downstream_enum.task_a, downstream_enum.task_c]
+ )
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMBranchOperator(
+ task_id="branch_approval_multi",
+ prompt="Pick branches",
+ llm_conn_id="my_llm",
+ allow_multiple_branches=True,
+ require_approval=True,
+ )
+ op.downstream_task_ids = {"task_a", "task_b", "task_c"}
+
+ with pytest.raises(ApprovalPauseSignal) as exc_info:
+ op.execute(_make_context())
+
+ assert exc_info.value.kwargs["generated_output"] ==
'["task_a","task_c"]'
+ mock_do_branch.assert_not_called()
+
+ @patch.object(LLMBranchOperator, "do_branch")
+ @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_lists_choices_and_renders_enum_dropdown(
+ self, mock_hook_cls, mock_upsert, mock_trigger_cls, mock_do_branch
+ ):
+ """The review body lists the valid branches and the editable param is
an enum dropdown."""
+ downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a",
"task_b": "task_b"})
+
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.return_value =
_make_mock_run_result(downstream_enum.task_a)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMBranchOperator(
+ task_id="branch_approval",
+ prompt="Pick a branch",
+ llm_conn_id="my_llm",
+ require_approval=True,
+ allow_modifications=True,
+ )
+ op.downstream_task_ids = {"task_b", "task_a"}
+
+ with pytest.raises(ApprovalPauseSignal):
+ op.execute(_make_context())
+
+ call_kwargs = mock_upsert.call_args.kwargs
+ assert call_kwargs["body"].startswith("Valid branches: `task_a`,
`task_b`")
+ assert call_kwargs["params"]["output"]["schema"] == {
+ "type": "string",
+ "enum": ["task_a", "task_b"],
+ }
+
+ @patch.object(LLMBranchOperator, "do_branch")
+ @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(
+ self, mock_hook_cls, mock_upsert, mock_trigger_cls, mock_do_branch
+ ):
+ """With allow_multiple_branches the editable param stays free-text
(JSON list)."""
+ downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a",
"task_b": "task_b"})
+
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.return_value =
_make_mock_run_result([downstream_enum.task_a])
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMBranchOperator(
+ task_id="branch_approval_multi",
+ prompt="Pick branches",
+ llm_conn_id="my_llm",
+ allow_multiple_branches=True,
+ require_approval=True,
+ allow_modifications=True,
+ )
+ op.downstream_task_ids = {"task_a", "task_b"}
+
+ with pytest.raises(ApprovalPauseSignal):
+ op.execute(_make_context())
+
+ 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"}
+
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_execute_rejects_sequence_prompt_with_require_approval(self,
mock_hook_cls):
+ """Non-string prompt + require_approval=True fails before the agent
runs."""
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMBranchOperator(
+ task_id="test",
+ prompt=["describe", b"bytes"], # type: ignore[arg-type]
+ llm_conn_id="my_llm",
+ require_approval=True,
+ )
+ op.downstream_task_ids = {"task_a"}
+
+ with pytest.raises(TypeError, match="require_approval=True"):
+ op.execute(_make_context())
+
+ mock_agent.run_sync.assert_not_called()
+
+ @patch.object(LLMBranchOperator, "do_branch")
+ def test_execute_complete_approved_single_branch(self, mock_do_branch):
+ """execute_complete branches into the approved task."""
+ mock_do_branch.return_value = "task_a"
+ op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c")
+ op.downstream_task_ids = {"task_a", "task_b"}
+ event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+ ctx = _make_context()
+
+ result = op.execute_complete(ctx, generated_output="task_a",
event=event)
+
+ assert result == "task_a"
+ mock_do_branch.assert_called_once_with(ctx, "task_a")
+
+ @patch.object(LLMBranchOperator, "do_branch")
+ def test_execute_complete_approved_multiple_branches(self, mock_do_branch):
+ """execute_complete parses the JSON list back before branching."""
+ mock_do_branch.return_value = ["task_a", "task_c"]
+ op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c",
allow_multiple_branches=True)
+ op.downstream_task_ids = {"task_a", "task_b", "task_c"}
+ event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+ ctx = _make_context()
+
+ result = op.execute_complete(ctx,
generated_output='["task_a","task_c"]', event=event)
+
+ assert result == ["task_a", "task_c"]
+ mock_do_branch.assert_called_once_with(ctx, ["task_a", "task_c"])
+
+ @patch.object(LLMBranchOperator, "do_branch")
+ def test_execute_complete_with_modified_branch(self, mock_do_branch):
+ """A reviewer-modified branch is used when it is a valid downstream
task."""
+ mock_do_branch.return_value = "task_b"
+ op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c",
allow_modifications=True)
+ op.downstream_task_ids = {"task_a", "task_b"}
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": "admin",
+ "params_input": {"output": "task_b"},
+ }
+ ctx = _make_context()
+
+ result = op.execute_complete(ctx, generated_output="task_a",
event=event)
+
+ assert result == "task_b"
+ mock_do_branch.assert_called_once_with(ctx, "task_b")
+
+ @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."""
+ op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c",
allow_modifications=True)
+ op.downstream_task_ids = {"task_a", "task_b"}
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": "admin",
+ "params_input": {"output": "task_x"},
+ }
+
+ with pytest.raises(ValueError, match="not downstream tasks"):
+ op.execute_complete(_make_context(), generated_output="task_a",
event=event)
+
+ mock_do_branch.assert_not_called()
+
+ @patch.object(LLMBranchOperator, "do_branch")
+ def test_execute_complete_rejects_empty_branch_list(self, mock_do_branch):
+ """A reviewed empty list would skip every downstream task and must be
rejected."""
+ 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"}
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": "admin",
+ "params_input": {"output": "[]"},
+ }
+
+ with pytest.raises(ValueError, match="selects no branches"):
+ op.execute_complete(_make_context(),
generated_output='["task_a"]', event=event)
+
+ mock_do_branch.assert_not_called()
+
+ @pytest.mark.parametrize(
+ "modified",
+ ["not json", '{"task_a": 1}', '["task_a", 2]'],
+ ids=["malformed-json", "not-a-list", "non-string-item"],
+ )
+ @patch.object(LLMBranchOperator, "do_branch")
+ def test_execute_complete_rejects_invalid_multi_branch_shapes(self,
mock_do_branch, modified):
+ """With allow_multiple_branches=True the reviewed output must be a
JSON list of strings."""
+ 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"}
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": "admin",
+ "params_input": {"output": modified},
+ }
+
+ with pytest.raises(ValueError, match="JSON list"):
+ op.execute_complete(_make_context(),
generated_output='["task_a"]', event=event)
+
+ mock_do_branch.assert_not_called()
![https://example.com/x.png"](https://example.com/x.png"){kind=link}