guan404ming commented on code in PR #70651:
URL: https://github.com/apache/airflow/pull/70651#discussion_r3708984580
##########
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:
Fixed, empty selection now raises right after branches are computed,
covering both paths.
--
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]
