kaxil commented on code in PR #70651:
URL: https://github.com/apache/airflow/pull/70651#discussion_r3707019161
##########
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:
The guard closes the reviewer path, thanks. One gap left: the LLM's own
output can be `[]`. `output_type = list[downstream_tasks_enum]` carries no
minimum length, so an empty list is a valid model response -- `branches`
becomes `[]`, `chosen` renders as `[]`, and the form asks a human to approve
selecting nothing. From there, with `allow_modifications=False`, no action
succeeds: Approve hits the new `ValueError` in `_parse_reviewed_branches`,
Reject raises `HITLRejectException`. Either way the task fails after someone
has been waiting on it. Checking for empty in `execute()` before
`defer_for_approval` would fail fast at the LLM step instead of spending a
review that cannot succeed.
--
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]
