guan404ming commented on code in PR #70651:
URL: https://github.com/apache/airflow/pull/70651#discussion_r3700955370
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py:
##########
@@ -95,4 +111,38 @@ def execute(self, context: Context) -> str | Iterable[str]
| None:
else:
branches = str(output)
+ if self.require_approval:
+ self.defer_for_approval(context, branches) # type: ignore[misc]
Review Comment:
Done — body now lists valid branches; single-branch mode renders an enum
dropdown via `modification_schema`.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py:
##########
@@ -95,4 +111,38 @@ def execute(self, context: Context) -> str | Iterable[str]
| None:
else:
branches = str(output)
+ if self.require_approval:
+ self.defer_for_approval(context, branches) # type: ignore[misc]
+
+ 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):
Review Comment:
Done — empty reviewed lists now raise `ValueError` before `do_branch`, so
nothing gets skipped silently.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py:
##########
@@ -60,12 +70,18 @@ def __init__(
**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 and not isinstance(self.prompt, str):
Review Comment:
Done — moved the check to `LLMApprovalMixin.validate_approval_prompt`,
called from all three `execute()` bodies.
##########
providers/common/ai/docs/operators/llm_branch.rst:
##########
@@ -75,6 +75,30 @@ 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. When ``allow_modifications=True``, the reviewer
+can also change the choice — the modified branch(es) are validated against
+the downstream task IDs before branching. With
+``allow_multiple_branches=True`` the reviewed value is a JSON list of task
+IDs (e.g. ``["task_a", "task_b"]``):
+
+.. code-block:: python
Review Comment:
Done — added `require_approval` example DAG, switched to `exampleinclude`,
and documented the approval parameters.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
