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. It does not cover the LLM
returning `[]` itself, though: `output_type = list[downstream_tasks_enum]`
carries no minimum length, so an empty selection serializes to `"[]"`, writes a
HITL detail reading `Chosen branch(es): []`, and pauses. Every reviewer action
then fails the task, Approve via this new guard and Reject via
`HITLRejectException`, and with `allow_modifications=False` they cannot fix it
at all. Checking emptiness right after `branches` is computed (line 109) would
cover both paths, and would also close the `require_approval=False` case, where
`do_branch(context, [])` reaches `skip_all_except` with an empty set and skips
every downstream task.
--
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]
