# 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
from typing import TYPE_CHECKING, Annotated, Any
from uuid import UUID
from airflow.plugins_manager import AirflowPlugin
from airflow.providers.common.ai.plugins._plugin_urls import get_bundle_url
from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_4_PLUS
if TYPE_CHECKING:
from airflow.plugins_manager import FastAPIAppDict, ReactAppDict
_PLUGIN_PREFIX = "/hitl-review"
if AIRFLOW_V_3_1_PLUS:
import mimetypes
from pathlib import Path
from fastapi import Depends, FastAPI, HTTPException, Query
from fastapi.staticfiles import StaticFiles
from sqlalchemy import delete, select
from sqlalchemy.orm import Session
from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
from airflow.api_fastapi.core_api.security import requires_access_dag
from airflow.models.dagrun import DagRun
from airflow.models.taskinstance import TaskInstance as TI
from airflow.models.xcom import XComModel
from airflow.providers.common.ai.utils.hitl_review import (
XCOM_AGENT_OUTPUT_PREFIX,
XCOM_AGENT_SESSION,
XCOM_HUMAN_ACTION,
XCOM_HUMAN_FEEDBACK_PREFIX,
AgentSessionData,
HITLReviewResponse,
HumanActionData,
HumanFeedbackRequest,
SessionStatus,
)
from airflow.sdk import TaskInstanceState
from airflow.utils.session import create_session
if AIRFLOW_V_3_4_PLUS:
from airflow.api_fastapi.core_api.services.public.task_coordinates import resolve_task_scope
from airflow.models.dagbag import DBDagBag
from airflow.models.task_coordinates import TaskCoordinateResolver
def _get_session():
with create_session(scoped=False) as session:
yield session
SessionDep = Annotated[Session, Depends(_get_session)]
def _read_xcom(
session: Session,
*,
dag_id: str,
run_id: str,
task_id: str,
map_index: int = -1,
region_id: UUID | None,
key: str,
):
"""Read a single XCom value from the database."""
scope: dict[str, Any] = (
{"region_id": region_id, "include_node_regions": False} if AIRFLOW_V_3_4_PLUS else {}
)
read = XComModel.get_many(
run_id=run_id,
key=key,
dag_ids=dag_id,
task_ids=task_id,
map_indexes=map_index,
limit=1,
**scope,
)
row = session.scalars(read).first()
if row is None:
return None
return row.value
def _read_xcom_by_prefix(
session: Session,
*,
dag_id: str,
run_id: str,
task_id: str,
map_index: int = -1,
region_id: UUID | None,
prefix: str,
) -> dict[int, Any]:
"""Read all iteration-keyed XCom entries matching *prefix* (e.g. ``airflow_hitl_review_agent_output_``)."""
if AIRFLOW_V_3_4_PLUS:
read = XComModel.get_many(
run_id=run_id,
dag_ids=dag_id,
task_ids=task_id,
map_indexes=map_index,
region_id=region_id,
include_node_regions=False,
)
entity = read.column_descriptions[0]["entity"]
query = read.with_only_columns(entity.key, entity.value).where(entity.key.like(f"{prefix}%"))
else:
query = select(XComModel.key, XComModel.value).where(
XComModel.dag_id == dag_id,
XComModel.run_id == run_id,
XComModel.task_id == task_id,
XComModel.map_index == map_index,
XComModel.key.like(f"{prefix}%"),
)
result: dict[int, Any] = {}
for key, value in session.execute(query).all():
suffix = key[len(prefix) :]
if suffix.isdigit():
result[int(suffix)] = value
return result
def _write_xcom(
session: Session,
*,
dag_id: str,
run_id: str,
task_id: str,
map_index: int = -1,
region_id: UUID | None,
key: str,
value,
):
"""Write data to db."""
if AIRFLOW_V_3_4_PLUS:
owner = session.scalar(
select(TI.id).where(
TI.working_set.is_(True),
TI.dag_id == dag_id,
TI.run_id == run_id,
TI.task_id == task_id,
TI.region_index == map_index,
TI.region_id == region_id,
)
)
if owner is None:
raise HTTPException(404, f"Task instance not found on DAG {dag_id!r} with ID {run_id!r}")
XComModel.set_for_attempt(
task_instance_id=owner,
key=key,
value=value,
serialize=False,
session=session,
)
return
# Stores value natively to match worker-written XComs; use XComModel.set(serialize=False) once min Airflow >= 3.2.
dag_run_id = session.scalar(select(DagRun.id).where(DagRun.dag_id == dag_id, DagRun.run_id == run_id))
if dag_run_id is None:
raise HTTPException(404, f"DAG run not found on DAG {dag_id!r} with ID {run_id!r}")
session.execute(
delete(XComModel).where(
XComModel.key == key,
XComModel.run_id == run_id,
XComModel.task_id == task_id,
XComModel.dag_id == dag_id,
XComModel.map_index == map_index,
)
)
session.add(
XComModel(
dag_run_id=dag_run_id,
key=key,
value=value,
run_id=run_id,
task_id=task_id,
dag_id=dag_id,
map_index=map_index,
)
)
session.flush()
_RUNNING_TI_STATES = frozenset(
{
TaskInstanceState.RUNNING,
TaskInstanceState.DEFERRED,
TaskInstanceState.UP_FOR_RETRY,
TaskInstanceState.QUEUED,
TaskInstanceState.SCHEDULED,
}
)
def _is_task_completed(
session: Session,
*,
dag_id: str,
run_id: str,
task_id: str,
map_index: int = -1,
region_id: UUID | None,
) -> bool:
"""Return True if the task instance is no longer running."""
query = select(TI.state).where(
TI.dag_id == dag_id,
TI.run_id == run_id,
TI.task_id == task_id,
)
if AIRFLOW_V_3_4_PLUS:
query = query.where(
TI.working_set.is_(True),
TI.region_index == map_index,
TI.region_id == region_id,
)
else:
query = query.where(TI.map_index == map_index)
state = session.scalar(query)
if state is None:
return True
return state not in _RUNNING_TI_STATES
def _build_session_response(
session: Session,
*,
dag_id: str,
run_id: str,
task_id: str,
map_index: int = -1,
region_id: UUID | None,
) -> HITLReviewResponse | None:
"""Build `HITLReviewResponse` from XCom entries."""
raw = _read_xcom(
session,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=XCOM_AGENT_SESSION,
)
if raw is None:
return None
sess_data = AgentSessionData.model_validate(raw)
outputs = _read_xcom_by_prefix(
session,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
prefix=XCOM_AGENT_OUTPUT_PREFIX,
)
human_responses = _read_xcom_by_prefix(
session,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
prefix=XCOM_HUMAN_FEEDBACK_PREFIX,
)
completed = _is_task_completed(
session,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
)
return HITLReviewResponse.from_xcom(
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
session=sess_data,
outputs=outputs,
human_entries=human_responses,
task_completed=completed,
)
def _get_map_index(q: str = Query("-1", alias="map_index")) -> int:
"""Parse map_index query; use -1 when placeholder unreplaced (e.g. ``{MAP_INDEX}``) or invalid."""
try:
return int(q)
except (ValueError, TypeError):
return -1
MapIndexDep = Annotated[int, Depends(_get_map_index)]
def _get_task_scope(
db: SessionDep,
dag_id: str,
run_id: str,
task_id: str,
map_index: MapIndexDep,
region_id: UUID | None = None,
region_index: Annotated[int | None, Query(ge=-1)] = None,
) -> tuple[UUID | None, int]:
if AIRFLOW_V_3_4_PLUS:
scope = resolve_task_scope(
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
resolver=TaskCoordinateResolver(DBDagBag(), db),
map_index=map_index,
region_id=region_id,
region_index=region_index,
)
return scope.region_id, scope.region_index
if region_id is not None or region_index is not None:
raise HTTPException(400, "Region selectors require Airflow 3.4 or later")
return None, map_index
TaskScopeDep = Annotated[tuple[UUID | None, int], Depends(_get_task_scope)]
hitl_review_app = FastAPI(
title="HITL Review",
description=(
"REST API and chat UI for human-in-the-loop LLM feedback sessions. "
"Sessions are stored in XCom entries on the running task instance."
),
)
@hitl_review_app.get("/health")
async def health() -> dict[str, str]:
"""Liveness check."""
return {"status": "ok"}
@hitl_review_app.get(
"/sessions/find",
response_model=HITLReviewResponse,
dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.HITL_DETAIL))],
)
async def find_session(
db: SessionDep,
dag_id: str,
task_id: str,
run_id: str,
scope: TaskScopeDep,
) -> HITLReviewResponse:
"""Find the feedback session for a specific task instance."""
region_id, map_index = scope
resp = _build_session_response(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
)
if resp is None:
task_active = not _is_task_completed(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
)
raise HTTPException(
status_code=404,
detail={"message": "No matching session found.", "task_active": task_active},
)
return resp
@hitl_review_app.post(
"/sessions/feedback",
response_model=HITLReviewResponse,
dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.HITL_DETAIL))],
)
async def submit_feedback(
body: HumanFeedbackRequest,
db: SessionDep,
dag_id: str,
task_id: str,
run_id: str,
scope: TaskScopeDep,
) -> HITLReviewResponse:
"""Request changes — provide human feedback for the LLM."""
region_id, map_index = scope
if not (body.feedback and body.feedback.strip()):
raise HTTPException(
status_code=400,
detail="Feedback is required when requesting changes.",
)
raw = _read_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=XCOM_AGENT_SESSION,
)
if raw is None:
raise HTTPException(status_code=404, detail="No matching session found.")
sess_data = AgentSessionData.model_validate(raw)
if sess_data.status != SessionStatus.PENDING_REVIEW:
raise HTTPException(
status_code=409,
detail=f"Session is '{sess_data.status.value}', expected 'pending_review'.",
)
iteration = sess_data.iteration
_write_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=f"{XCOM_HUMAN_FEEDBACK_PREFIX}{iteration}",
value=body.feedback,
)
action = HumanActionData(action="changes_requested", feedback=body.feedback, iteration=iteration)
_write_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=XCOM_HUMAN_ACTION,
value=action.model_dump(mode="json"),
)
sess_data.status = SessionStatus.CHANGES_REQUESTED
_write_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=XCOM_AGENT_SESSION,
value=sess_data.model_dump(mode="json"),
)
resp = _build_session_response(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
)
if resp is None:
raise HTTPException(status_code=500, detail="Failed to read session after update.")
return resp
@hitl_review_app.post(
"/sessions/approve",
response_model=HITLReviewResponse,
dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.HITL_DETAIL))],
)
async def approve_session(
db: SessionDep,
dag_id: str,
task_id: str,
run_id: str,
scope: TaskScopeDep,
) -> HITLReviewResponse:
"""Approve the current output."""
region_id, map_index = scope
raw = _read_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=XCOM_AGENT_SESSION,
)
if raw is None:
raise HTTPException(status_code=404, detail="No matching session found.")
sess_data = AgentSessionData.model_validate(raw)
if sess_data.status != SessionStatus.PENDING_REVIEW:
raise HTTPException(
status_code=409,
detail=f"Session is '{sess_data.status.value}', expected 'pending_review'.",
)
action = HumanActionData(action="approve", iteration=sess_data.iteration)
_write_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=XCOM_HUMAN_ACTION,
value=action.model_dump(mode="json"),
)
sess_data.status = SessionStatus.APPROVED
_write_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=XCOM_AGENT_SESSION,
value=sess_data.model_dump(mode="json"),
)
resp = _build_session_response(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
)
if resp is None:
raise HTTPException(status_code=500, detail="Failed to read session after update.")
return resp
@hitl_review_app.post(
"/sessions/reject",
response_model=HITLReviewResponse,
dependencies=[Depends(requires_access_dag(method="PUT", access_entity=DagAccessEntity.HITL_DETAIL))],
)
async def reject_session(
db: SessionDep,
dag_id: str,
task_id: str,
run_id: str,
scope: TaskScopeDep,
) -> HITLReviewResponse:
"""Reject the output."""
region_id, map_index = scope
raw = _read_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=XCOM_AGENT_SESSION,
)
if raw is None:
raise HTTPException(status_code=404, detail="No matching session found.")
sess_data = AgentSessionData.model_validate(raw)
if sess_data.status != SessionStatus.PENDING_REVIEW:
raise HTTPException(
status_code=409,
detail=f"Session is '{sess_data.status.value}', expected 'pending_review'.",
)
action = HumanActionData(action="reject", iteration=sess_data.iteration)
_write_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=XCOM_HUMAN_ACTION,
value=action.model_dump(mode="json"),
)
sess_data.status = SessionStatus.REJECTED
_write_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
key=XCOM_AGENT_SESSION,
value=sess_data.model_dump(mode="json"),
)
resp = _build_session_response(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
region_id=region_id,
)
if resp is None:
raise HTTPException(status_code=500, detail="Failed to read session after update.")
return resp
# Ensure proper MIME types for plugin bundle (FastAPI serves .cjs as text/plain by default)
mimetypes.add_type("application/javascript", ".cjs")
_WWW_DIR = Path(__file__).parent / "www"
_dist_dir = _WWW_DIR / "dist"
if _dist_dir.is_dir():
hitl_review_app.mount(
"/static",
StaticFiles(directory=str(_dist_dir.absolute()), html=True),
name="hitl_review_static",
)
[docs]
class HITLReviewPlugin(AirflowPlugin):
"""Register the HITL Review REST API + chat UI on the Airflow API server."""
[docs]
fastapi_apps: list[FastAPIAppDict] = []
[docs]
react_apps: list[ReactAppDict] = []
if AIRFLOW_V_3_1_PLUS:
fastapi_apps = [
{
"name": "hitl-review",
"app": hitl_review_app,
"url_prefix": _PLUGIN_PREFIX,
}
]
react_apps = [
{
"name": "HITL Review",
"bundle_url": get_bundle_url(_PLUGIN_PREFIX, "main.umd.cjs"),
"destination": "task_instance",
"url_route": "hitl-review",
}
]