This is an automated email from the ASF dual-hosted git repository.
guan404ming 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 69509289c96 Preserve output_type through human approval in LLM
operators (#70075)
69509289c96 is described below
commit 69509289c9682c6efec578731b52f810402bd1d9
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Sun Jul 19 21:51:31 2026 +0800
Preserve output_type through human approval in LLM operators (#70075)
* Preserve output_type through human approval in LLM operators
* Cover BaseModel in output_type restore test
---
.../airflow/providers/common/ai/mixins/approval.py | 8 ++++----
.../airflow/providers/common/ai/utils/output_type.py | 20 ++++++++++----------
.../ai/tests/unit/common/ai/mixins/test_approval.py | 13 +++++++++----
.../ai/tests/unit/common/ai/operators/test_llm.py | 19 +++++++++++++++++++
.../tests/unit/common/ai/utils/test_output_type.py | 15 ++++++++++++++-
5 files changed, 56 insertions(+), 19 deletions(-)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
index 5ebd679efcd..9f045c02548 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
@@ -21,7 +21,7 @@ import logging
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Protocol
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_3_PLUS
@@ -107,9 +107,9 @@ class LLMApprovalMixin:
if isinstance(output, BaseModel):
output = output.model_dump_json()
- if not isinstance(output, str):
- # Always make string output so that when comparing in the
execute_complete matches
- output = str(output)
+ elif not isinstance(output, str):
+ # JSON round-trip: execute_complete validates the string back into
output_type.
+ output = TypeAdapter(type(output)).dump_json(output).decode()
ti_id = context["task_instance"].id
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
index 2ae9b6a5bc3..44ca8505b54 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
@@ -20,7 +20,7 @@ from __future__ import annotations
from typing import Any
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, TypeAdapter, ValidationError
def rehydrate_pydantic_output(
@@ -30,25 +30,25 @@ def rehydrate_pydantic_output(
serialize_output: bool,
) -> Any:
"""
- Turn a JSON string back into the ``output_type`` Pydantic model.
+ Turn a JSON string back into a value of ``output_type``.
Used by the HITL/approval paths in ``LLMOperator`` and ``AgentOperator``
- that round-trip the model through a string when deferring to a human
- reviewer. When ``output_type`` is not a ``BaseModel`` subclass, returns
- ``raw`` unchanged so the caller can apply its own fallback (e.g.
- ``json.loads``). When validation fails (reviewer edited the string into
- something the schema rejects), also returns ``raw`` unchanged.
+ that round-trip the output through a string when deferring to a human
+ reviewer. ``str`` outputs pass through unchanged; any other ``output_type``
+ (``BaseModel`` subclass, ``int``, ``list[str]``, ...) is validated with a
+ pydantic ``TypeAdapter``. When validation fails (reviewer edited the string
+ into something the type rejects), returns ``raw`` unchanged.
When ``serialize_output`` is ``True``, returns the model dumped to a
``dict`` -- matches the operator's ``serialize_output=True`` opt-in for
consumers that want the dict shape.
"""
- if not (isinstance(output_type, type) and issubclass(output_type,
BaseModel)):
+ if output_type is str:
return raw
try:
- rehydrated = output_type.model_validate_json(raw)
+ rehydrated = TypeAdapter(output_type).validate_json(raw)
except (ValidationError, ValueError, TypeError):
return raw
- if serialize_output:
+ if serialize_output and isinstance(rehydrated, BaseModel):
return rehydrated.model_dump()
return rehydrated
diff --git a/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
b/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
index 54b675da723..46120b3c3c3 100644
--- a/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
+++ b/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
@@ -118,15 +118,20 @@ class TestDeferForApproval:
defer_kwargs = approval_op.defer.call_args[1]
assert defer_kwargs["kwargs"]["generated_output"] ==
'{"text":"Paris","confidence":0.95}'
+ @pytest.mark.parametrize(
+ ("output", "expected"),
+ [(42, "42"), (True, "true"), ([1, "a"], '[1,"a"]'), ({"k": "v"},
'{"k":"v"}')],
+ ids=["int", "bool", "list", "dict"],
+ )
@patch(HITL_TRIGGER_PATH, autospec=True)
@patch(UPSERT_HITL_PATH)
- def test_non_string_non_pydantic_output_is_stringified(
- self, mock_upsert, mock_trigger_cls, approval_op, context
+ def test_non_string_non_pydantic_output_is_json_encoded(
+ self, mock_upsert, mock_trigger_cls, approval_op, context, output,
expected
):
- approval_op.defer_for_approval(context, 42)
+ approval_op.defer_for_approval(context, output)
defer_kwargs = approval_op.defer.call_args[1]
- assert defer_kwargs["kwargs"]["generated_output"] == "42"
+ assert defer_kwargs["kwargs"]["generated_output"] == expected
@patch(HITL_TRIGGER_PATH, autospec=True)
@patch(UPSERT_HITL_PATH)
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
index f9f3bf09099..e2004b4031c 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
@@ -379,6 +379,25 @@ class TestLLMOperatorApproval:
assert isinstance(result, Summary)
assert result.text == "hello"
+ @pytest.mark.parametrize(
+ ("output_type", "generated_output", "expected"),
+ [
+ (int, "5", 5),
+ (list[str], '["a","b"]', ["a", "b"]),
+ pytest.param(Summary, '{"text":"hello"}', Summary(text="hello"),
marks=requires_typed_xcom),
+ ],
+ ids=["int", "list", "basemodel"],
+ )
+ def test_execute_complete_restores_non_str_output_type(self, output_type,
generated_output, expected):
+ op = LLMOperator(
+ task_id="t", prompt="p", llm_conn_id="c", output_type=output_type,
require_approval=True
+ )
+ event = {"chosen_options": ["Approve"], "responded_by_user": "admin"}
+
+ result = op.execute_complete({}, generated_output=generated_output,
event=event)
+
+ assert result == expected
+
@pytest.mark.skipif(
not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with
Airflow >= 3.1.0"
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
b/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
index 4c3dcae5ee2..f7e1a4799e2 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import pytest
from pydantic import BaseModel
from airflow.providers.common.ai.utils.output_type import
rehydrate_pydantic_output
@@ -35,10 +36,22 @@ class TestRehydratePydanticOutput:
result = rehydrate_pydantic_output(A, '{"x": 7}',
serialize_output=True)
assert result == {"x": 7}
- def test_returns_raw_for_non_basemodel(self):
+ def test_returns_raw_for_str_output_type(self):
result = rehydrate_pydantic_output(str, "anything",
serialize_output=False)
assert result == "anything"
+ @pytest.mark.parametrize(
+ ("output_type", "raw", "expected"),
+ [(int, "5", 5), (bool, "true", True), (list[str], '["a", "b"]', ["a",
"b"])],
+ ids=["int", "bool", "list"],
+ )
+ def test_validates_other_types_with_type_adapter(self, output_type, raw,
expected):
+ assert rehydrate_pydantic_output(output_type, raw,
serialize_output=False) == expected
+
+ def test_returns_raw_when_type_adapter_rejects(self):
+ result = rehydrate_pydantic_output(int, "not-a-number",
serialize_output=False)
+ assert result == "not-a-number"
+
def test_returns_raw_on_invalid_json(self):
result = rehydrate_pydantic_output(A, "not-json",
serialize_output=False)
assert result == "not-json"
