# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
import json
import logging
from collections.abc import Iterable
from datetime import timedelta
from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol
from jinja2 import TemplateError
from airflow.providers.common.ai.utils.output_type import dump_output_to_json
from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_3_PLUS
if AIRFLOW_V_3_3_PLUS:
# On Airflow 3.3+ the review parks the task in the first-class AWAITING_INPUT state instead
# of deferring to a trigger. On older Airflow versions this name is absent and defer() is used.
from airflow.sdk.exceptions import TaskAwaitingInput
[docs]
log = logging.getLogger(__name__)
if TYPE_CHECKING:
from collections.abc import Sequence
from airflow.providers.common.compat.notifier import BaseNotifier
from airflow.sdk import Context
from airflow.sdk.execution_time.hitl import HITLUser
def _object_from_the_form_as_text(modified: dict[str, Any], generated_output: str) -> str:
"""
Turn the object the HITL form sent back into review text, keeping ``generated_output`` if untouched.
The form round-trips the value through JavaScript, which rounds integers above 2**53, writes ``1.0``
as ``1`` and may reorder keys (so may the metadata database), so compare values the same way rather
than text; a re-dumped untouched object would otherwise read as an edit and could change a value.
"""
try:
generated = json.loads(generated_output)
except ValueError:
generated = None
if _with_javascript_numbers(modified) == _with_javascript_numbers(generated):
return generated_output
return json.dumps(modified, separators=(",", ":"), ensure_ascii=False)
def _with_javascript_numbers(value: Any) -> Any:
"""Map every number to a float, as ``JSON.parse`` does, so values compare as the browser sees them."""
if isinstance(value, bool):
# Tagged, because Python treats ``False == 0.0`` and ``True == 1.0`` as equal and JSON does not.
return ("bool", value)
if isinstance(value, int | float):
try:
return float(value)
except OverflowError:
return value
if isinstance(value, dict):
return {key: _with_javascript_numbers(item) for key, item in value.items()}
if isinstance(value, list):
return [_with_javascript_numbers(item) for item in value]
return value
[docs]
def normalize_assigned_users(value: Any, *, param: str) -> list[HITLUser]:
"""
Return *value* as a list of ``{'id': str, 'name': str}`` users.
Accepts ``None``, a single user dict, or an iterable of them, and raises ``TypeError``
naming *param* for anything else, so a malformed list fails when the Dag is parsed
rather than when the task first asks for a review.
"""
users: list[Any]
if value is None:
users = []
elif isinstance(value, dict):
users = [value]
elif isinstance(value, str) or not isinstance(value, Iterable):
raise TypeError(
f"{param} must be a {{'id': str, 'name': str}} dict or an iterable of them, got {value!r}"
)
else:
users = list(value)
for user in users:
if (
not isinstance(user, dict)
or not isinstance(user.get("id"), str)
or not isinstance(user.get("name"), str)
):
raise TypeError(f"{param} entries must be {{'id': str, 'name': str}} dicts, got {user!r}")
return users
[docs]
class DeferForApprovalProtocol(Protocol):
"""Protocol for defer for approval mixin."""
[docs]
approval_timeout: timedelta | None
[docs]
allow_modifications: bool
[docs]
on_approval_timeout: Literal["fail", "approve", "reject"]
[docs]
approval_notifiers: Sequence[BaseNotifier]
[docs]
approval_assigned_users: list[HITLUser]
[docs]
def validate_approval_prompt(self) -> None: ...
[docs]
class LLMApprovalMixin:
"""
Mixin that pauses an operator for human review before returning output.
When ``require_approval=True`` on the operator, the generated output is
presented to a human reviewer via the Airflow Human-in-the-Loop (HITL)
interface. The task waits (``awaiting_input`` on Airflow 3.3+, deferred on
older versions) until the reviewer approves or rejects.
If ``allow_modifications=True``, the reviewer can also edit the output
before approving. The (possibly modified) output is then returned as the
task result.
``on_approval_timeout`` decides what happens when ``approval_timeout``
expires without a response: ``"fail"`` raises ``HITLTimeoutError``, while
``"approve"`` and ``"reject"`` answer the review with that option so the
task resumes as if a reviewer had chosen it. The chosen option is also
pre-highlighted for the reviewer in the HITL form.
``approval_notifiers`` are called once the review is open, so a reviewer
learns about it without watching the Required Actions page, the way
:class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with
``notifiers``. The review ``subject`` and ``body`` are exposed as
``{{ task.subject }}`` and ``{{ task.body }}`` in notifier templates. A
notifier whose delivery raises is logged without failing the task, while a
template error fails the task. A retry re-runs the LLM and re-notifies
with the regenerated output. On Airflow 3.4+ the reopened review shows the
regenerated output too; on Airflow 3.1 to 3.3 it keeps the original subject,
body and editable output, so approving it unedited returns the first
attempt's output.
``approval_assigned_users`` restricts the review to the named users, the way
:class:`~airflow.providers.standard.operators.hitl.HITLOperator` does with
``assigned_users``. Leaving it empty lets any user with the permission
respond. On Airflow 3.1 to 3.3 the list is stored when the review is first
created, so clearing the task after changing it does not take effect; on
3.4+ a cleared or retried task replaces it.
Operators that use this mixin must set the following attributes:
- ``require_approval`` (``bool``)
- ``allow_modifications`` (``bool``)
- ``approval_timeout`` (``timedelta | None``)
- ``on_approval_timeout`` (``Literal["fail", "approve", "reject"]``)
- ``approval_notifiers`` (``Sequence[BaseNotifier]``)
- ``approval_assigned_users`` (``list[HITLUser]``)
- ``prompt`` (``str``)
"""
[docs]
TIMEOUT_DEFAULTS: ClassVar[dict[str, list[str]]] = {"approve": [APPROVE], "reject": [REJECT]}
[docs]
def validate_approval_prompt(self: DeferForApprovalProtocol) -> None:
"""Fail fast when the prompt cannot be rendered as text in the approval review body."""
if not isinstance(self.prompt, str):
raise TypeError(
f"{type(self).__name__}: require_approval=True is not supported "
f"with a non-string prompt (got {type(self.prompt).__name__}). "
"The approval review body renders the prompt as text; passing a "
"Sequence[UserContent] would expose object reprs (and any embedded "
"bytes) in the human review UI. Return a str prompt, or disable "
"require_approval."
)
[docs]
def defer_for_approval(
self: DeferForApprovalProtocol,
context: Context,
output: Any,
*,
subject: str | None = None,
body: str | None = None,
modification_schema: dict[str, Any] | None = None,
decision: dict[str, Any] | None = None,
) -> None:
"""
Write HITL detail, then pause the task for human review.
On Airflow 3.3+ the task parks in the ``awaiting_input`` state (no trigger or triggerer
involved); on older versions it defers to :class:`HITLTrigger`. Either way it resumes in
``execute_complete`` once a response (or timeout default) arrives. ``on_approval_timeout``
supplies that timeout default; ``"fail"`` supplies none, so the review times out as an error.
:param context: Airflow task context.
:param output: The generated output to present for review.
:param subject: Headline shown on the Required Actions page.
Defaults to ``"Review output for task `<task_id>`"``.
:param body: Markdown body shown below the headline.
Defaults to the prompt and output wrapped in a code block.
:param modification_schema: JSON schema for the editable ``output`` param
when ``allow_modifications=True``. Defaults to ``{"type": "string"}``.
Pass e.g. ``{"type": "string", "enum": [...]}`` to render a dropdown
of valid values in the review form, or ``{"type": "array", "items":
{"type": "string", "enum": [...]}, "examples": [...]}`` to render a
multi-select (JSON Schema forbids ``enum`` at the array level, so the
options come from ``examples``); a list submitted by the reviewer is
returned from ``execute_complete`` re-serialized as a compact JSON string.
:param decision: The pending decision record, when the operator wrote one. It is carried in
the continuation the task resumes from, next to ``generated_output``, so
``execute_complete`` finalizes the record from what was checkpointed with the pause and
not from a copy a reader could have edited or deleted in the meantime.
"""
from airflow.providers.standard.triggers.hitl import HITLTrigger
from airflow.sdk.execution_time.hitl import upsert_hitl_detail
from airflow.sdk.timezone import utcnow
self.validate_approval_prompt()
raw_output = output
# JSON round-trip: execute_complete validates the string back into output_type.
output = dump_output_to_json(output)
ti_id = context["task_instance"].id
timeout_defaults = LLMApprovalMixin.TIMEOUT_DEFAULTS.get(self.on_approval_timeout)
if subject is None:
subject = f"Review output for task `{self.task_id}`"
if body is None:
body = f"```\nPrompt: {self.prompt}\n\n{output}\n```"
hitl_params: dict[str, dict[str, Any]] = {}
if self.allow_modifications:
# The multi-select rendered for an array schema needs the list, not its JSON string
param_value: Any = output
if (
modification_schema is not None
and modification_schema.get("type") == "array"
and isinstance(raw_output, list)
):
param_value = raw_output
hitl_params = {
"output": {
"value": param_value,
"description": "Edit the output before approving (optional).",
"schema": modification_schema or {"type": "string"},
},
}
upsert_hitl_detail(
ti_id=ti_id,
options=[LLMApprovalMixin.APPROVE, LLMApprovalMixin.REJECT],
subject=subject,
body=body,
defaults=timeout_defaults,
multiple=False,
params=hitl_params,
assigned_users=self.approval_assigned_users,
)
self.subject = subject
self.body = body
for notifier in self.approval_notifiers:
try:
notifier({**context})
except TemplateError:
raise
except Exception:
log.exception("Approval notifier %s failed; the review stays open", notifier)
continuation: dict[str, Any] = {"generated_output": output}
if decision is not None:
continuation["decision"] = decision
if AIRFLOW_V_3_3_PLUS:
# Airflow 3.3+: park the task in AWAITING_INPUT -- no trigger, no triggerer. The
# task is resumed by the Core API response handler or the scheduler timeout sweep.
raise TaskAwaitingInput(
method_name="execute_complete",
kwargs=continuation,
timeout=self.approval_timeout,
)
# Fallback for Airflow versions < 3.3: defer the response check to HITLTrigger on the triggerer.
self.defer(
trigger=HITLTrigger(
ti_id=ti_id,
options=[LLMApprovalMixin.APPROVE, LLMApprovalMixin.REJECT],
defaults=timeout_defaults,
params=hitl_params,
multiple=False,
timeout_datetime=(
utcnow() + self.approval_timeout if self.approval_timeout is not None else None
),
),
method_name="execute_complete",
kwargs=continuation,
)
@staticmethod
def _describe_responder(event: dict[str, Any]) -> str:
responded_by_user = event.get("responded_by_user")
if responded_by_user is None:
return "the approval timeout default"
return responded_by_user["name"]
[docs]
def execute_complete(
self,
context: Context,
generated_output: str,
event: dict[str, Any],
decision: dict[str, Any] | None = None,
) -> str:
"""
Resume after human review.
Called automatically by Airflow when the HITL trigger fires.
Returns the original or reviewer-modified output on approval.
:param context: Airflow task context.
:param generated_output: The output that was deferred for review.
:param event: Trigger event payload containing ``chosen_options``,
``params_input``, ``responded_by_user``, and ``timedout``.
:param decision: The pending decision record passed to ``defer_for_approval``, if any.
The mixin does not read it; an operator that writes a record finalizes it.
:raises HITLRejectException: If the reviewer, or the
``on_approval_timeout="reject"`` default, rejected the output.
:raises HITLTriggerEventError: If the trigger reported an error.
:raises HITLTimeoutError: If the approval timed out.
"""
from airflow.providers.standard.exceptions import (
HITLRejectException,
HITLTimeoutError,
HITLTriggerEventError,
)
if "error" in event:
error_type = event.get("error_type", "unknown")
if error_type == "timeout":
raise HITLTimeoutError(f"Approval timed out: {event['error']}")
raise HITLTriggerEventError(event)
responder = self._describe_responder(event)
chosen = event["chosen_options"]
if self.APPROVE not in chosen:
if event.get("timedout"):
raise HITLRejectException(
"Output was rejected automatically: approval_timeout expired with "
"on_approval_timeout='reject'."
)
raise HITLRejectException(f"Output was rejected by the reviewer {responder}.")
log.info("Output approved by %s.", responder)
output = generated_output
params_input: dict[str, Any] = {} if event.get("timedout") else event.get("params_input") or {}
# Only accept modified output when the operator explicitly allows modifications.
# Without this guard a reviewer could craft a request with params_input even
# when allow_modifications=False, bypassing the read-only approval flow.
if getattr(self, "allow_modifications", False) and params_input:
modified = params_input.get("output")
if "output" in params_input and modified is None:
raise HITLTriggerEventError(
{
"error": "Modified output must not be empty; edit it or reject instead.",
"error_type": "validation",
}
)
if isinstance(modified, list):
for item in modified:
if not isinstance(item, str):
raise HITLTriggerEventError(
{
"error": f"Modified output list items must be strings, "
f"got {type(item).__name__}.",
"error_type": "validation",
}
)
# Compact so an unchanged selection compares equal to generated_output
modified = json.dumps(modified, separators=(",", ":"), ensure_ascii=False)
elif isinstance(modified, dict):
# TODO: Drop once the HITL form stops parsing a string param that looks like JSON into an
# object (airflow-core ui/src/utils/hitl.ts); until then a model output comes back as a dict.
modified = _object_from_the_form_as_text(modified, generated_output)
if modified is not None and not isinstance(modified, str):
# On the awaiting_input path nothing upstream schema-validates params_input
# (HITLTrigger did on the legacy path), so enforce the string contract here
# rather than returning a non-string as the task's output.
raise HITLTriggerEventError(
{
"error": f"Modified output must be a string or a list of strings, "
f"got {type(modified).__name__}.",
"error_type": "validation",
}
)
if modified is not None and modified != generated_output:
log.info("output=%s modified by the reviewer=%s ", modified, responder)
return modified
return output