Lee-W commented on code in PR #76916:
URL: https://github.com/apache/airflow/pull/76916#discussion_r4549392275
##########
providers/common/ai/docs/hitl_review.rst:
##########
@@ -152,6 +171,60 @@ navigate to
:start-after: [START howto_operator_agent_hitl_review]
:end-before: [END howto_operator_agent_hitl_review]
+.. _hitl-review-durable:
+
+With durable execution
+----------------------
+
+With ``durable=True`` (see :doc:`durable_execution`), the first output and
every
+regeneration are agent runs in the task's journal, so a retry replays them
instead of
+calling the model and the tools again (on Airflow 3.1 and 3.2, only the first
output; see the
+end of this section). Each regeneration logs its own ``Durable:`` line in the
task log, after
+the first run's.
+
+On Airflow 3.3 and later, each request for changes is also kept in the task
state store
+beside the journal, with the round it was given in and a digest of the output
the reviewer
+read. When
+the task fails after the review has started (a regeneration hits a failing
tool or an
+unavailable model, or the review times out) the retry replays the first
output, applies the
Review Comment:
```suggestion
read. When the task fails after the review has started (a regeneration hits
a failing tool
or an unavailable model, or the review times out) the retry replays the
first output, applies the
```
the breakline feels weird
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py:
##########
@@ -57,23 +133,56 @@ class HITLReviewProtocol(Protocol):
@property
def prompt(self) -> str | Sequence[Any]: ...
+ def regenerate_with_feedback(self, *, feedback: str, message_history: Any)
-> tuple[str, Any]: ...
+
+ def dump_review_state(self, message_history: Any) -> str: ...
+
+ def load_review_state(self, context: Context, state: str) -> Any: ...
+
+ def complete_hitl_review(self, context: Context, approved_output: str, *,
state: str | None) -> Any: ...
+
+ def _parks_between_rounds(self) -> bool: ...
+
+ def _recorded_review_feedback(self, context: Context, session:
AgentSessionData) -> str | None: ...
Review Comment:
```suggestion
def _get_recorded_review_feedback(self, context: Context, session:
AgentSessionData) -> str | None: ...
```
##########
providers/common/ai/src/airflow/providers/common/ai/durable/journal.py:
##########
@@ -581,30 +592,116 @@ def pause(self) -> None:
replay, this attempt has run live from there: what it recorded becomes
trusted, and
the earlier entries at positions it left empty (a call still waiting
for approval, a
tool that raised) are deleted. Positions from here on still run live
after the pause.
+ Otherwise it only notes how far this attempt got, when that is further
than any before.
"""
self._load_meta()
start = min((p for p in (self._diverge_from, self._diverged_at) if p
is not None), default=None)
if start is None or self._position <= start:
- # This attempt has not run anything live past the untrusted
entries yet.
+ # This attempt has not run anything live past the untrusted
entries yet. Past a step
+ # it did not record (a tool set to run on every attempt), a later
completion or
+ # cleanup cannot find how far it got, so note that and leave the
rest as it was.
+ if self._position > self._high_water:
+ self._save_meta(
+ self._diverge_from,
+ model_at_fault=self._model_at_fault,
+ paused=self._paused,
+ tool_failure=self._tool_failure or None,
+ last_model=self._noted_last_model,
+ )
+ return
+ survived = self._delete_unrecorded(range(start, self._position))
+ trusted_until = survived[0] if survived else self._position
+ self._diverge_from, self._model_at_fault = trusted_until, False
+ # A pause is not a failure, and the task resumes in the same attempt,
so the count of
+ # attempts in a row a tool call failed carries over to a failure after
the resume.
+ self._save_meta(
+ trusted_until,
+ model_at_fault=False,
+ paused=True,
+ tool_failure=self._tool_failure or None,
+ last_model=self._last_model_position,
+ )
+
+ def complete(self) -> None:
Review Comment:
A killed attempt leaves no `high_water` note, so `_position < _high_water`
is false and a non-diverging shorter completion writes nothing, keeping the
killed attempt’s entries past its end. A journal-layer reproduction (attempt A
killed after a refused write at position 2, B answers at position 2, then
reject) has the retry return B’s rejected answer. Could `complete()` always
write `last_model` and the end position when the journal is kept, or maybe we
can document it as a limitation with a tracking issue and an
`xfail(strict=True)` test?
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py:
##########
@@ -57,23 +133,56 @@ class HITLReviewProtocol(Protocol):
@property
def prompt(self) -> str | Sequence[Any]: ...
+ def regenerate_with_feedback(self, *, feedback: str, message_history: Any)
-> tuple[str, Any]: ...
+
+ def dump_review_state(self, message_history: Any) -> str: ...
+
+ def load_review_state(self, context: Context, state: str) -> Any: ...
+
+ def complete_hitl_review(self, context: Context, approved_output: str, *,
state: str | None) -> Any: ...
+
+ def _parks_between_rounds(self) -> bool: ...
+
+ def _recorded_review_feedback(self, context: Context, session:
AgentSessionData) -> str | None: ...
+
+ def _record_review_feedback(self, context: Context, session:
AgentSessionData, feedback: str) -> None: ...
+
+ def _start_review_round(
+ self, context: Context, session: AgentSessionData, *, message_history:
Any, deadline: datetime | None
+ ) -> NoReturn: ...
+
+ def _park_for_review(
+ self, context: Context, session: AgentSessionData, *, deadline:
datetime | None, state_sha256: str
+ ) -> NoReturn: ...
+
+ def _next_round(
Review Comment:
```suggestion
def _regenerate_next_round(
```
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py:
##########
@@ -287,6 +603,53 @@ def regenerate_with_feedback(self, *, feedback: str,
message_history: Any) -> tu
"Operators using HITLReviewMixin must implement
regenerate_with_feedback()."
)
+ def dump_review_state(self, message_history: Any) -> str:
+ """
+ Serialize what :meth:`load_review_state` needs to regenerate after the
task resumes.
+
+ Called on Airflow 3.3+ before every park; the string is kept in the
task state store.
+
+ :param message_history: Provider-specific conversation state.
+ """
+ raise NotImplementedError("Operators using HITLReviewMixin must
implement dump_review_state().")
+
+ def load_review_state(self, context: Context, state: str) -> Any:
+ """
+ Restore what :meth:`dump_review_state` saved, in the resumed task.
+
+ Called before each regeneration.
+
+ :returns: The ``message_history`` to pass to
:meth:`regenerate_with_feedback`.
+ """
+ raise NotImplementedError("Operators using HITLReviewMixin must
implement load_review_state().")
+
+ def _recorded_review_feedback(self, context: Context, session:
AgentSessionData) -> str | None:
Review Comment:
```suggestion
def _get_recorded_review_feedback(self, context: Context, session:
AgentSessionData) -> str | None:
```
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1824,6 +1852,141 @@ def _emit_run_metadata(self, context: Context, result:
Any, *, usage: RunUsage)
if (model_name := getattr(result.response, "model_name", None)) is not
None:
ti.xcom_push(key=MODEL_NAME_XCOM_KEY, value=model_name)
+ def dump_review_state(self, message_history: Any) -> str:
+ """Save the conversation, and what the task returns if the output
under review is approved."""
+ output = self._output_under_review
+ state: dict[str, Any] = {
+ "messages": ModelMessagesTypeAdapter.dump_python(message_history,
mode="json")
+ }
+ output_json = dump_output_to_json(output)
+ if (
+ rehydrate_pydantic_output(self.output_type, output_json,
serialize_output=False, is_edited=False)
+ == output
+ ):
+ state["output_json"] = output_json
+ else:
+ # Reading the JSON back would give a different value: an output
function's or an output
+ # marker's result stays text, and "42" under ``str | int`` becomes
42. Save the result instead.
+ state["result"] =
to_jsonable_python(self._convert_output_for_xcom(output))
+ return json.dumps(state)
+
+ def load_review_state(self, context: Context, state: str) -> Any:
+ """Restore the conversation and the usage budget before regenerating
in the resumed task."""
+ ti = context["task_instance"]
+ # The regeneration shares the budget the paused rounds spent, as it
does without a pause.
+ self._usage_budget = self._build_usage_budget(context,
coerce_usage_limits(self.usage_limits), ti=ti)
+ self._run_usage = self._usage_budget.load() if self._usage_budget else
None
+ self._run_identity_attrs = build_run_identity_attributes(ti)
+ return
ModelMessagesTypeAdapter.validate_python(json.loads(state)["messages"])
+
+ def complete_hitl_review(self, context: Context, approved_output: str, *,
state: str | None) -> Any:
+ """Return the approved output as the agent returned it;
``approved_output`` is only its display text."""
+ # The reviewer can approve, reject, or ask for a new answer, but
cannot edit the answer text.
+ # So the approved answer is the last answer the agent returned.
+ if state is None:
+ result = self._convert_output_for_xcom(self._output_under_review)
+ elif "result" in (saved := json.loads(state)):
+ result = saved["result"]
+ else:
+ result = self._convert_output_for_xcom(
+ rehydrate_pydantic_output(
+ self.output_type, saved["output_json"],
serialize_output=False, is_edited=False
+ )
+ )
+ if self._durable_journal is not None:
+ # The approved output is the task's result, so nothing replays the
review again.
+ self._durable_journal.cleanup()
+ self._delete_review_answers(context)
+ self._clear_usage_budget(context)
+ return result
+
+ def resume_hitl_review(
+ self,
+ context: Context,
+ iteration: int,
+ deadline: str | None,
+ state_sha256: str,
+ event: dict[str, Any],
+ ) -> Any:
+ if self.durable:
+ # The regeneration after round ``iteration`` is the journal's run
``iteration``, in this
+ # process as on a retry that replays the rounds before it; run 0
made the first output.
+ self._durable_journal =
DurableJournal(self._build_durable_storage(context), first_run=iteration)
+ with self._durable_review(context):
+ return super().resume_hitl_review(context, iteration, deadline,
state_sha256, event)
+
+ @contextmanager
+ def _durable_review(self, context: Context) -> Iterator[None]:
+ """Under durable, make a review that ends in a reject or
``max_hitl_iterations`` start over on retry."""
+ try:
+ yield
+ except (HITLRejectException, HITLMaxIterationsError):
+ if (journal := self._durable_journal) is not None:
+ # Replaying would show the reviewer the outputs of this review
again. The retry asks
+ # the model again for the first output, replaying the tool
calls before it, and
+ # regenerates from there.
+ self.log.info(
+ "The review ended without an approved output; a retry asks
the model again for "
+ "the first output and starts the review over."
+ )
+ journal.ask_model_again(0)
+ self._delete_review_answers(context)
+ raise
+
+ @staticmethod
+ def _stored_review_answers(context: Context) -> list[_ReviewAnswer]:
Review Comment:
```suggestion
def _load_review_answers(context: Context) -> list[_ReviewAnswer]:
```
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py:
##########
@@ -144,15 +275,218 @@ def run_hitl_review(
deadline = time.monotonic() + self.hitl_timeout.total_seconds() if
self.hitl_timeout else None
return self._poll_loop( # type: ignore[attr-defined]
+ context=context,
ti=ti,
session=session,
message_history=message_history,
deadline=deadline,
)
+ def _start_review_round(
+ self: HITLReviewProtocol,
+ context: Context,
+ session: AgentSessionData,
+ *,
+ message_history: Any,
+ deadline: datetime | None,
+ ) -> NoReturn:
+ """
+ Save the conversation, show ``session`` to the reviewer, and park
until they act.
+
+ A round the reviewer already asked to change in an earlier attempt of
the task (see
+ :meth:`_recorded_review_feedback`) is not shown again: its feedback is
applied, and the
+ next round starts.
+ """
+ ti = context["task_instance"]
+ while session.iteration < self.max_hitl_iterations and (
+ feedback := self._recorded_review_feedback(context, session)
+ ):
+ self.log.info(
+ "Applying the feedback given on iteration %d in an earlier
attempt: %s",
+ session.iteration,
+ feedback,
+ )
+ # The HITL Review tab shows each round from these.
+ ti.xcom_push(key=f"{XCOM_AGENT_OUTPUT_PREFIX}{session.iteration}",
value=session.current_output)
+
ti.xcom_push(key=f"{XCOM_HUMAN_FEEDBACK_PREFIX}{session.iteration}",
value=feedback)
+ session, message_history = self._next_round(session, feedback,
message_history)
+ state = self.dump_review_state(message_history)
+ retention = (
+ deadline - utcnow() + _REVIEW_STATE_RETENTION_MARGIN if deadline
is not None else NEVER_EXPIRE
+ )
+ context["task_state_store"].set(_REVIEW_STATE_KEY, state,
retention=retention)
+ self._park_for_review(
+ context, session, deadline=deadline,
state_sha256=hashlib.sha256(state.encode()).hexdigest()
+ )
+
+ def _next_round(
Review Comment:
```suggestion
def _regenerate_next_round(
```
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1824,6 +1852,141 @@ def _emit_run_metadata(self, context: Context, result:
Any, *, usage: RunUsage)
if (model_name := getattr(result.response, "model_name", None)) is not
None:
ti.xcom_push(key=MODEL_NAME_XCOM_KEY, value=model_name)
+ def dump_review_state(self, message_history: Any) -> str:
+ """Save the conversation, and what the task returns if the output
under review is approved."""
+ output = self._output_under_review
+ state: dict[str, Any] = {
+ "messages": ModelMessagesTypeAdapter.dump_python(message_history,
mode="json")
+ }
+ output_json = dump_output_to_json(output)
+ if (
+ rehydrate_pydantic_output(self.output_type, output_json,
serialize_output=False, is_edited=False)
+ == output
+ ):
+ state["output_json"] = output_json
+ else:
+ # Reading the JSON back would give a different value: an output
function's or an output
+ # marker's result stays text, and "42" under ``str | int`` becomes
42. Save the result instead.
+ state["result"] =
to_jsonable_python(self._convert_output_for_xcom(output))
+ return json.dumps(state)
+
+ def load_review_state(self, context: Context, state: str) -> Any:
+ """Restore the conversation and the usage budget before regenerating
in the resumed task."""
+ ti = context["task_instance"]
+ # The regeneration shares the budget the paused rounds spent, as it
does without a pause.
+ self._usage_budget = self._build_usage_budget(context,
coerce_usage_limits(self.usage_limits), ti=ti)
+ self._run_usage = self._usage_budget.load() if self._usage_budget else
None
+ self._run_identity_attrs = build_run_identity_attributes(ti)
+ return
ModelMessagesTypeAdapter.validate_python(json.loads(state)["messages"])
+
+ def complete_hitl_review(self, context: Context, approved_output: str, *,
state: str | None) -> Any:
+ """Return the approved output as the agent returned it;
``approved_output`` is only its display text."""
+ # The reviewer can approve, reject, or ask for a new answer, but
cannot edit the answer text.
+ # So the approved answer is the last answer the agent returned.
+ if state is None:
+ result = self._convert_output_for_xcom(self._output_under_review)
+ elif "result" in (saved := json.loads(state)):
+ result = saved["result"]
+ else:
+ result = self._convert_output_for_xcom(
+ rehydrate_pydantic_output(
+ self.output_type, saved["output_json"],
serialize_output=False, is_edited=False
+ )
+ )
+ if self._durable_journal is not None:
+ # The approved output is the task's result, so nothing replays the
review again.
+ self._durable_journal.cleanup()
+ self._delete_review_answers(context)
+ self._clear_usage_budget(context)
+ return result
+
+ def resume_hitl_review(
+ self,
+ context: Context,
+ iteration: int,
+ deadline: str | None,
+ state_sha256: str,
+ event: dict[str, Any],
+ ) -> Any:
+ if self.durable:
+ # The regeneration after round ``iteration`` is the journal's run
``iteration``, in this
+ # process as on a retry that replays the rounds before it; run 0
made the first output.
+ self._durable_journal =
DurableJournal(self._build_durable_storage(context), first_run=iteration)
+ with self._durable_review(context):
+ return super().resume_hitl_review(context, iteration, deadline,
state_sha256, event)
+
+ @contextmanager
+ def _durable_review(self, context: Context) -> Iterator[None]:
+ """Under durable, make a review that ends in a reject or
``max_hitl_iterations`` start over on retry."""
+ try:
+ yield
+ except (HITLRejectException, HITLMaxIterationsError):
+ if (journal := self._durable_journal) is not None:
+ # Replaying would show the reviewer the outputs of this review
again. The retry asks
+ # the model again for the first output, replaying the tool
calls before it, and
+ # regenerates from there.
+ self.log.info(
+ "The review ended without an approved output; a retry asks
the model again for "
+ "the first output and starts the review over."
+ )
+ journal.ask_model_again(0)
+ self._delete_review_answers(context)
+ raise
+
+ @staticmethod
+ def _stored_review_answers(context: Context) -> list[_ReviewAnswer]:
+ return _REVIEW_ANSWERS_ADAPTER.validate_python(
+ context["task_state_store"].get(_HITL_REVIEW_ANSWERS_KEY) or []
+ )
+
+ def _recorded_review_feedback(self, context: Context, session:
AgentSessionData) -> str | None:
Review Comment:
```suggestion
def _get_recorded_review_feedback(self, context: Context, session:
AgentSessionData) -> str | None:
```
--
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]
