# 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 urllib.parse import urlparse
from airflow.plugins_manager import AirflowPlugin
from airflow.providers.common.compat.sdk import conf
from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_1_PLUS, get_base_airflow_version_tuple
[docs]
AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0)
if TYPE_CHECKING:
from airflow.plugins_manager import FastAPIAppDict, ReactAppDict
_PLUGIN_PREFIX = "/hitl-review"
def _get_base_url_path(path: str) -> str:
"""Construct URL path with webserver base_url prefix for non-root deployments."""
base_url = conf.get("api", "base_url", fallback="/")
if base_url.startswith(("http://", "https://")):
base_path = urlparse(base_url).path
else:
base_path = base_url
base_path = base_path.rstrip("/")
return base_path + path
def _get_bundle_url() -> str:
"""
Return bundle URL for the React plugin.
Uses an absolute URL when api.base_url is a full URL so the bundle loads
correctly in Vite dev mode, where import() resolves relative to the script
origin (5173) rather than the document origin (28080).
"""
path = _get_base_url_path(f"{_PLUGIN_PREFIX}/static/main.umd.cjs")
base_url = conf.get("api", "base_url", fallback="/")
if base_url.startswith(("http://", "https://")):
parsed = urlparse(base_url)
return f"{parsed.scheme}://{parsed.netloc}" + path
return path
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
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, key: str
):
"""Read a single XCom value from the database."""
read = XComModel.get_many(
run_id=run_id,
key=key,
dag_ids=dag_id,
task_ids=task_id,
map_indexes=map_index,
limit=1,
)
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, 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,
)
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, 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.map_index == map_index,
)
)
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
) -> 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,
TI.map_index == map_index,
)
if AIRFLOW_V_3_4_PLUS:
query = query.where(TI.working_set.is_(True))
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
) -> 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,
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,
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,
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,
)
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)]
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,
map_index: MapIndexDep,
) -> HITLReviewResponse:
"""Find the feedback session for a specific task instance."""
resp = _build_session_response(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
)
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,
)
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,
map_index: MapIndexDep,
) -> HITLReviewResponse:
"""Request changes — provide human feedback for the LLM."""
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,
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,
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,
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,
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,
)
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,
map_index: MapIndexDep,
) -> HITLReviewResponse:
"""Approve the current output."""
raw = _read_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
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,
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,
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,
)
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,
map_index: MapIndexDep,
) -> HITLReviewResponse:
"""Reject the output."""
raw = _read_xcom(
db,
dag_id=dag_id,
run_id=run_id,
task_id=task_id,
map_index=map_index,
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,
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,
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,
)
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(),
"destination": "task_instance",
"url_route": "hitl-review",
}
]