diff --git a/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py b/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py index 612a902a22377..9e85cf23d542d 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py +++ b/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py @@ -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 ( @@ -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 diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index f83a58f0e2aed..aad1d21a2b5bc 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -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 @@ -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( + self.output_type, + result_str, + serialize_output=self._serialize_model_output, + ) if self._serialize_model_output and isinstance(output, BaseModel): output = output.model_dump() @@ -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() diff --git a/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py b/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py index 1bf319f57459d..7b535aee0dce1 100644 --- a/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py +++ b/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py @@ -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 ( @@ -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( diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 2631c544138f5..70031d0b13b50 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -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" ) @@ -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):