This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 2ed1196e80e Support require_approval in LLMSchemaCompareOperator
(#71051)
2ed1196e80e is described below
commit 2ed1196e80e8203389e04be26e9e8f48830eb8ea
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Wed Aug 12 17:18:43 2026 +0800
Support require_approval in LLMSchemaCompareOperator (#71051)
Co-authored-by: Wei Lee <[email protected]>
---
.../ai/docs/operators/llm_schema_compare.rst | 27 +++++
.../ai/example_dags/example_llm_schema_compare.py | 21 ++++
.../common/ai/operators/llm_schema_compare.py | 45 +++++++-
.../common/ai/operators/test_llm_schema_compare.py | 123 ++++++++++++++++++++-
4 files changed, 205 insertions(+), 11 deletions(-)
diff --git a/providers/common/ai/docs/operators/llm_schema_compare.rst
b/providers/common/ai/docs/operators/llm_schema_compare.rst
index d5f014adf0c..768d96280be 100644
--- a/providers/common/ai/docs/operators/llm_schema_compare.rst
+++ b/providers/common/ai/docs/operators/llm_schema_compare.rst
@@ -114,6 +114,27 @@ The callable may also return a non-empty
``Sequence[UserContent]`` for
multimodal inputs -- see
:ref:`@task.agent multimodal prompts <howto/operator:agent-multimodal>`.
+Human-in-the-Loop Approval
+--------------------------
+
+Set ``require_approval=True`` to pause the task after the comparison and wait
+for a human reviewer to approve the result before it is returned. The review
+body shows the compatibility verdict, a mismatch severity summary, and the
+full result JSON. Rejecting the review, or letting ``approval_timeout``
+expire, fails the task:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
+ :language: python
+ :start-after: [START howto_operator_llm_schema_compare_approval]
+ :end-before: [END howto_operator_llm_schema_compare_approval]
+
+``require_approval=True`` requires a string prompt: a decorated callable
+returning a ``Sequence[UserContent]`` raises ``TypeError`` before the LLM
+call.
+
+``approval_timeout``, ``allow_modifications``, and the rest of the approval
+behaviour are inherited from :ref:`LLMOperator <howto/operator:llm>`.
+
Conditional ETL Based on Schema Compatibility
----------------------------------------------
@@ -166,6 +187,12 @@ Parameters
catalog-managed sources.
- ``context_strategy``: To fetch primary keys, foreign keys, and
indexes.``full`` or ``basic``,
strongly recommended for cross-system comparisons. default is ``full``
+- ``require_approval``: If ``True``, the task pauses after the comparison and
+ waits for human review before returning the result. Default ``False``.
+- ``approval_timeout``: Maximum time to wait for a review (``timedelta``).
``None``
+ means wait indefinitely. Default ``None``.
+- ``allow_modifications``: If ``True``, the reviewer can edit the result JSON
+ before approving. Default ``False``.
Logging
-------
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
index e884b74a5f8..bf8020e684e 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_schema_compare.py
@@ -18,6 +18,8 @@
from __future__ import annotations
+from datetime import timedelta
+
from airflow.providers.common.ai.operators.llm_schema_compare import
LLMSchemaCompareOperator
from airflow.providers.common.compat.sdk import dag, task
from airflow.providers.common.sql.config import DataSourceConfig
@@ -105,6 +107,25 @@ def example_llm_schema_compare_decorator():
example_llm_schema_compare_decorator()
+# [START howto_operator_llm_schema_compare_approval]
+@dag(tags=["example"])
+def example_llm_schema_compare_approval():
+ LLMSchemaCompareOperator(
+ task_id="detect_schema_drift_with_approval",
+ prompt="Identify schema mismatches that would break data loading
between systems",
+ llm_conn_id="pydanticai_default",
+ db_conn_ids=["postgres_source", "snowflake_target"],
+ table_names=["customers"],
+ require_approval=True,
+ approval_timeout=timedelta(hours=1),
+ )
+
+
+# [END howto_operator_llm_schema_compare_approval]
+
+example_llm_schema_compare_approval()
+
+
# [START howto_operator_llm_schema_compare_conditional]
@dag(tags=["example"])
def example_llm_schema_compare_conditional():
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
index d65ca521876..7992419296f 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
@@ -19,11 +19,12 @@
from __future__ import annotations
import json
+from collections import Counter
from collections.abc import Sequence
from functools import cached_property
from typing import TYPE_CHECKING, Any, Literal
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationError
from airflow.providers.common.ai.operators.llm import LLMOperator
from airflow.providers.common.ai.utils.logging import log_run_summary
@@ -106,6 +107,13 @@ class LLMSchemaCompareOperator(LLMOperator):
:param context_strategy: ``"basic"`` for column names and types only;
``"full"`` to include primary keys, foreign keys, and indexes.
Default ``"full"``.
+
+ Human-in-the-Loop approval parameters are inherited from
+ :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
+ (``require_approval``, ``approval_timeout``, ``allow_modifications``).
+ The task pauses after the comparison and only returns the result once a
+ reviewer approves. The review body shows the compatibility verdict, a
+ mismatch severity summary, and the full result JSON.
"""
template_fields: Sequence[str] = (
@@ -126,9 +134,9 @@ class LLMSchemaCompareOperator(LLMOperator):
system_prompt: str = DEFAULT_SYSTEM_PROMPT,
**kwargs: Any,
) -> None:
- kwargs.pop("output_type", None)
- if kwargs.get("require_approval"):
- raise ValueError("require_approval=True is not supported by
LLMSchemaCompareOperator.")
+ kwargs["output_type"] = SchemaCompareResult
+ # execute() always returns a dict, so the approval resume path must too
+ kwargs["serialize_output"] = True
super().__init__(**kwargs)
self.data_sources = data_sources or []
self.db_conn_ids = db_conn_ids or []
@@ -299,6 +307,9 @@ class LLMSchemaCompareOperator(LLMOperator):
return "".join(parts)
def execute(self, context: Context) -> dict[str, Any]:
+ if self.require_approval:
+ self.validate_approval_prompt() # type: ignore[misc]
+
schema_context = self._build_schema_context()
self.log.info("Schema comparison context:\n%s", schema_context)
@@ -313,8 +324,32 @@ class LLMSchemaCompareOperator(LLMOperator):
self.log.info("Running LLM schema comparison...")
result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
log_run_summary(self.log, result)
+ output = result.output
- output_result = result.output.model_dump()
+ output_result = output.model_dump()
self.log.info("Schema comparison result: \n %s",
json.dumps(output_result, indent=2))
+ if self.require_approval:
+ severity_counts = Counter(mismatch.severity for mismatch in
output.mismatches)
+ summary = ", ".join(
+ f"{severity_counts[severity]} {severity}"
+ for severity in ("critical", "warning", "info")
+ if severity_counts[severity]
+ )
+ body = (
+ f"Compatible: {output.compatible}"
+ + (f" (mismatches: {summary})" if summary else "")
+ + f"\n\n```\nPrompt:
{self.prompt}\n\n{output.model_dump_json(indent=2)}\n```"
+ )
+ self.defer_for_approval(context, output, body=body) # type:
ignore[misc]
+
return output_result
+
+ def execute_complete(self, context: Context, generated_output: str, event:
dict[str, Any]) -> Any:
+ output = super().execute_complete(context, generated_output, event)
+ if isinstance(output, dict):
+ return output
+ try:
+ return SchemaCompareResult.model_validate_json(output).model_dump()
+ except ValidationError as e:
+ raise ValueError(f"Reviewed output is not a valid
SchemaCompareResult: {e}") from e
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
index 01fee1950d1..656457de295 100644
---
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
+++
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
@@ -16,20 +16,36 @@
# under the License.
from __future__ import annotations
+import json
from unittest import mock
from unittest.mock import MagicMock
+from uuid import uuid4
import pytest
from airflow.providers.common.ai.operators.llm_schema_compare import (
LLMSchemaCompareOperator,
SchemaCompareResult,
+ SchemaMismatch,
)
-from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, TaskDeferred
from airflow.providers.common.sql.config import DataSourceConfig
from airflow.providers.common.sql.datafusion.engine import DataFusionEngine
from airflow.providers.common.sql.hooks.sql import DbApiHook
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_3_PLUS
+
+if AIRFLOW_V_3_3_PLUS:
+ from airflow.sdk.exceptions import TaskAwaitingInput
+else:
+ TaskAwaitingInput = TaskDeferred # type: ignore[assignment, misc]
+
+
+def _make_context():
+ ti = MagicMock()
+ ti.id = uuid4()
+ return MagicMock(**{"__getitem__": lambda self, key: {"task_instance":
ti}[key]})
+
def _make_mock_run_result(output):
"""Create a mock AgentRunResult compatible with log_run_summary."""
@@ -97,11 +113,6 @@ class TestLLMSchemaCompareOperator:
"at-least two combinations",
id="one_datasource_only",
),
- pytest.param(
- {"require_approval": True},
- "require_approval=True is not supported",
- id="require_approval",
- ),
],
)
def test_init_validation(self, kwargs, expected_error):
@@ -558,3 +569,103 @@ class TestLLMSchemaCompareOperator:
with pytest.raises(AirflowOptionalProviderFeatureException):
op._introspect_schema_from_datafusion(ds)
+
+
[email protected](
+ not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with
Airflow >= 3.1.0"
+)
+class TestLLMSchemaCompareOperatorApproval:
+ """Tests for LLMSchemaCompareOperator with require_approval=True."""
+
+ _APPROVAL_KWARGS = dict(
+ db_conn_ids=["postgres_default", "snowflake_default"],
+ table_names=["orders"],
+ require_approval=True,
+ )
+
+ def test_require_approval_accepted(self):
+ op = LLMSchemaCompareOperator(**_BASE_KWARGS, **self._APPROVAL_KWARGS)
+ assert op.require_approval is True
+ assert op.output_type is SchemaCompareResult
+
+ @mock.patch(
+
"airflow.providers.common.ai.operators.llm_schema_compare.LLMSchemaCompareOperator._build_schema_context"
+ )
+ @mock.patch(
+
"airflow.providers.common.ai.operators.llm_schema_compare.LLMSchemaCompareOperator._build_system_prompt"
+ )
+ @mock.patch("airflow.providers.standard.triggers.hitl.HITLTrigger",
autospec=True)
+ @mock.patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail")
+ def test_execute_with_approval_pauses_with_summary_body(
+ self, mock_upsert, mock_trigger_cls, mock_build_system_prompt,
mock_build_schema_context, caplog
+ ):
+ mock_build_schema_context.return_value = "schema_context"
+ mock_build_system_prompt.return_value = "system_prompt"
+ result = SchemaCompareResult(
+ compatible=False,
+ mismatches=[
+ SchemaMismatch(
+ source="pg.orders",
+ target="sf.orders",
+ column="total",
+ source_type="numeric(10,2)",
+ target_type="NUMBER(5,0)",
+ severity="critical",
+ description="Precision loss",
+ suggested_action="Widen target column",
+ migration_query="ALTER TABLE orders ...",
+ )
+ ],
+ summary="One critical mismatch",
+ )
+
+ op = LLMSchemaCompareOperator(**_BASE_KWARGS, **self._APPROVAL_KWARGS)
+ mock_agent = mock.Mock()
+ mock_agent.run_sync.return_value = _make_mock_run_result(result)
+ op.llm_hook =
mock.Mock(create_agent=mock.Mock(return_value=mock_agent))
+
+ with pytest.raises(TaskAwaitingInput) as exc_info:
+ op.execute(context=_make_context())
+
+ assert exc_info.value.kwargs["generated_output"] ==
result.model_dump_json()
+ body = mock_upsert.call_args.kwargs["body"]
+ assert body.startswith("Compatible: False (mismatches: 1 critical)")
+ assert "Precision loss" in body
+ assert f"Schema comparison result: \n {json.dumps(result.model_dump(),
indent=2)}" in caplog
+
+ def test_execute_complete_approved_returns_dict(self):
+ result = SchemaCompareResult(compatible=True, mismatches=[],
summary="All good")
+ op = LLMSchemaCompareOperator(**_BASE_KWARGS, **self._APPROVAL_KWARGS)
+ event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+
+ resumed = op.execute_complete({},
generated_output=result.model_dump_json(), event=event)
+
+ assert resumed == result.model_dump()
+
+ @pytest.mark.parametrize(
+ "modified",
+ ['{"compatible": true, "mismatches": []}', "looks fine to me"],
+ ids=["missing-summary", "not-json"],
+ )
+ def test_execute_complete_rejects_invalid_modified_output(self, modified):
+ result = SchemaCompareResult(compatible=True, mismatches=[],
summary="All good")
+ op = LLMSchemaCompareOperator(**_BASE_KWARGS, **self._APPROVAL_KWARGS,
allow_modifications=True)
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": "admin",
+ "params_input": {"output": modified},
+ }
+
+ with pytest.raises(ValueError, match="not a valid
SchemaCompareResult"):
+ op.execute_complete({}, generated_output=result.model_dump_json(),
event=event)
+
+ def test_execute_rejects_sequence_prompt_with_require_approval(self):
+ op = LLMSchemaCompareOperator(
+ task_id="test_task",
+ prompt=["describe", b"bytes"], # type: ignore[arg-type]
+ llm_conn_id="llm_conn",
+ **self._APPROVAL_KWARGS,
+ )
+
+ with pytest.raises(TypeError, match="require_approval=True"):
+ op.execute(context=_make_context())
