kaxil commented on code in PR #70651:
URL: https://github.com/apache/airflow/pull/70651#discussion_r3685253206
##########
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:
This is now the third verbatim copy of the check: `LLMOperator.execute` has
the same block, and `LLMApprovalMixin.defer_for_approval` raises its own
version. The operator-level copies exist to fail before the LLM call, so would
it make sense to put the pre-flight on the mixin and call it from both
`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:
Every other section on this page uses `exampleinclude` from
`example_llm_branch.py`, and both `example_llm.py` and `example_llm_sql.py`
already carry a `require_approval=True` example. Adding one there and including
it here would keep the page consistent and give the feature example-DAG
coverage. The Parameters section below also doesn't list `require_approval` /
`approval_timeout` / `allow_modifications`, which `llm.rst` does.
##########
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:
This catches the wrong type and unknown task IDs, but not `[]`. A reviewer
editing the field to `[]` passes both checks, and `do_branch(context, [])`
lands in `skip_all_except` with an empty branch set, which skips every
downstream task. I ran it to confirm the empty list flows straight through to
`do_branch`. Given the premise that irreversible skipping is what needs the
human gate, that seems like the one value to reject here.
##########
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:
The reviewer never sees the branches they are choosing between. The default
body is ```Prompt: <prompt>\n\n<chosen>``` and the modification field is free
text (`{"type": "string"}`), so with `allow_modifications=True` they have to
type a task ID from memory. I tried `refund` against an operator with
`{refunds, billing}` downstream and got `ValueError: Reviewed branch(es)
['refund'] are not downstream tasks` at resume, after the approval was already
recorded, so the retry burns a second LLM call and needs a fresh review.
`defer_for_approval` accepts `body=`, so passing one that lists
`sorted(self.downstream_task_ids)` covers most of it. Constraining the param
instead with `"enum": sorted(self.downstream_task_ids)` would be better
(`Param.resolve` runs jsonschema, and `FlexibleForm` renders enum params as a
dropdown), but that needs a params override on the mixin.
--
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]
