Skip to content

Fix output type loss in AgentOperator's human-in-the-loop review #70132

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ColtenOuO wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix output type loss in AgentOperator's human-in-the-loop review #70132

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Protocol

from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter

from airflow.providers.common.ai.exceptions import HITLMaxIterationsError
from airflow.providers.common.ai.utils.hitl_review import (
Expand Down Expand Up @@ -271,5 +271,5 @@ def _to_string(output: Any) -> str:
if isinstance(output, BaseModel):
return output.model_dump_json()
if not isinstance(output, str):
return str(output)
return TypeAdapter(type(output)).dump_json(output).decode()
return output
22 changes: 9 additions & 13 deletions providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,14 +18,13 @@

from __future__ import annotations

import json
from collections.abc import Sequence
from dataclasses import replace
from datetime import timedelta
from functools import cached_property
from typing import TYPE_CHECKING, Any, ClassVar

from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter

from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin
Expand Down Expand Up @@ -493,16 +492,11 @@ def execute(self, context: Context) -> Any:
output,
message_history=result.all_messages(),
)
if isinstance(self.output_type, type) and issubclass(self.output_type, BaseModel):
return rehydrate_pydantic_output(
self.output_type,
result_str,
serialize_output=self._serialize_model_output,
)
try:
return json.loads(result_str)
except (ValueError, TypeError):
return result_str
return rehydrate_pydantic_output(

@kaxil kaxil Jul 21, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dropping the json.loads fallback here regresses the exact case this PR targets. rehydrate_pydantic_output returns raw unchanged when output_type isn't a BaseModel subclass (its docstring notes the caller is expected to apply its own json.loads), so with output_type=list[str] this returns the JSON string '["tag-a","tag-b"]' instead of the list. The new test_execute_with_hitl_rehydrates_non_base_model_output fails on exactly this (assert '["tag-a", "tag-b"]' == ['tag-a', 'tag-b']). Keeping the previous json.loads/except fallback after the rehydrate call fixes it while preserving the BaseModel handling.

@ColtenOuO ColtenOuO Jul 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In PR #70075 (merged), rehydrate_pydantic_output was updated to handle any non-str output_type, not just BaseModel subclasses:

if output_type is str:
    return raw
try:
    rehydrated = TypeAdapter(output_type).validate_json(raw)
except (ValidationError, ValueError, TypeError):
    return raw
uv run --project providers/common/ai python -c "tic_output
from airflow.providers.common.ai.utils.output_type import rehydrate_pydantic_outputtput=False)
result = rehydrate_pydantic_output(list[str], '[\"tag-a\",\"tag-b\"]', serialize_output=False)
print(repr(result), type(result))
"

Output:

['tag-a', 'tag-b'] <class 'list'>

So I don't think the json.loads fallback needs to come back (?

Let me know what you think, or if I might have missed anything!

Thanks for your review!

self.output_type,
result_str,
serialize_output=self._serialize_model_output,
)

if self._serialize_model_output and isinstance(output, BaseModel):
output = output.model_dump()
Expand Down Expand Up @@ -556,4 +550,6 @@ def regenerate_with_feedback(self, *, feedback: str, message_history: Any) -> tu
output = result.output
if isinstance(output, BaseModel):
output = output.model_dump_json()
return str(output), result.all_messages()
elif not isinstance(output, str):
output = TypeAdapter(type(output)).dump_json(output).decode()
return output, result.all_messages()
26 changes: 26 additions & 0 deletions providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,12 @@
if not AIRFLOW_V_3_1_PLUS:
pytest.skip("Human in the loop is only compatible with Airflow >= 3.1.0", allow_module_level=True)

import json
from datetime import timedelta
from unittest.mock import MagicMock, patch

from pydantic import BaseModel

from airflow.providers.common.ai.exceptions import HITLMaxIterationsError
from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin
from airflow.providers.common.ai.utils.hitl_review import (
Expand Down Expand Up @@ -107,6 +110,29 @@ def context(mock_ti):
return {"task_instance": mock_ti}


class TestToString:
def test_str_output_passes_through_unchanged(self):
assert HITLReviewMixin._to_string("plain text") == "plain text"

def test_base_model_output_uses_model_dump_json(self):
class Summary(BaseModel):
text: str

assert HITLReviewMixin._to_string(Summary(text="hi")) == '{"text":"hi"}'

def test_list_output_is_valid_json_not_python_repr(self):
result = HITLReviewMixin._to_string(["tag-a", "tag-b"])

assert result == '["tag-a","tag-b"]'
assert json.loads(result) == ["tag-a", "tag-b"]

def test_bool_output_is_valid_json_not_python_repr(self):
result = HITLReviewMixin._to_string(True)

assert result == "true"
assert json.loads(result) is True


class TestHITLReviewMixin:
@patch("airflow.providers.common.ai.mixins.hitl_review.time.sleep", autospec=True)
def test_pushes_session_and_output_on_start(
Expand Down
46 changes: 46 additions & 0 deletions providers/common/ai/tests/unit/common/ai/operators/test_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -432,6 +432,31 @@ def test_execute_with_hitl_rehydrates_base_model(self, mock_hook_cls, mock_run_h
assert result.text == "Approved summary"
assert result.score == 0.9

@pytest.mark.skipif(
not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0"
)
@patch("airflow.providers.common.ai.operators.agent.AgentOperator.run_hitl_review", autospec=True)
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
def test_execute_with_hitl_rehydrates_non_base_model_output(self, mock_hook_cls, mock_run_hitl):
mock_result = _make_mock_run_result(["tag-a", "tag-b"])
mock_agent = MagicMock(spec=["run_sync"])
mock_agent.run_sync.return_value = mock_result
mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent
mock_run_hitl.return_value = '["tag-a", "tag-b"]'

op = AgentOperator(
task_id="test",
prompt="Summarize",
llm_conn_id="my_llm",
output_type=list[str],
enable_hitl_review=True,
hitl_timeout=timedelta(minutes=5),
)
context = MagicMock()
result = op.execute(context=context)

assert result == ["tag-a", "tag-b"]

@pytest.mark.skipif(
not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0"
)
Expand Down Expand Up @@ -584,6 +609,27 @@ def test_regenerate_with_feedback_serializes_base_model_output(self, mock_hook_c

assert output == '{"text":"Revised","score":0.0}'

@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
def test_regenerate_with_feedback_serializes_non_base_model_output(self, mock_hook_cls):
mock_result = _make_mock_run_result(["tag-a", "tag-b"])
mock_result.all_messages.return_value = []
mock_agent = MagicMock(spec=["run_sync"])
mock_agent.run_sync.return_value = mock_result
mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent

op = AgentOperator(
task_id="test",
prompt="test",
llm_conn_id="my_llm",
output_type=list[str],
)
output, _ = op.regenerate_with_feedback(
feedback="Expand",
message_history=[],
)

assert output == '["tag-a","tag-b"]'


class TestAgentOperatorDurable:
def test_durable_param_stored(self):
Expand Down