This is an automated email from the ASF dual-hosted git repository.
kaxil 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 029cdf8f6cc Return an approved agent answer as `output_type` after
human review (#73904)
029cdf8f6cc is described below
commit 029cdf8f6cce2a1a2729706aaf0b87924c92c66f
Author: Kaxil Naik <[email protected]>
AuthorDate: Tue Sep 29 21:15:15 2026 +0100
Return an approved agent answer as `output_type` after human review (#73904)
With enable_hitl_review, AgentOperator ran the approved text through
json.loads unless output_type was a Pydantic model, so an answer of "42"
came back as an int and one that looked like a JSON object as a dict, even
with the default output_type=str. It now uses the same conversion as
LLMOperator's approval path: str passes through unchanged, and any other
output_type is validated from the JSON, falling back to the text a reviewer
edited into something the type rejects.
---
.../airflow/providers/common/ai/operators/agent.py | 15 +++------
.../tests/unit/common/ai/operators/test_agent.py | 36 ++++++++++++++++++++++
2 files changed, 41 insertions(+), 10 deletions(-)
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 1b6f5d5f0b7..d6ef9ad8100 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
@@ -849,16 +849,11 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin):
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()
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 73881634b2d..1f26bb9751b 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
@@ -968,6 +968,42 @@ class TestAgentOperatorExecute:
assert result == "Approved output"
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible
with Airflow >= 3.1.0"
+ )
+ @pytest.mark.parametrize(
+ ("output_type", "approved", "expected"),
+ [
+ pytest.param(str, "42", "42", id="str-that-parses-as-a-number"),
+ pytest.param(str, '{"total": 1}', '{"total": 1}',
id="str-that-parses-as-an-object"),
+ pytest.param(list[str], '["a", "b"]', ["a", "b"], id="list"),
+ pytest.param(int, "not a number", "not a number",
id="edit-the-type-rejects"),
+ ],
+ )
+
@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_returns_the_approved_output_as_output_type(
+ self, mock_hook_cls, mock_run_hitl, make_mock_run_result, output_type,
approved, expected
+ ):
+ """The approved text comes back as ``output_type``, as it does from
``@task.llm``."""
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ mock_agent.run_sync.return_value = make_mock_run_result("Initial
output")
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+ mock_run_hitl.return_value = approved
+ op = AgentOperator(
+ task_id="test",
+ prompt="Summarize",
+ llm_conn_id="my_llm",
+ output_type=output_type,
+ enable_hitl_review=True,
+ hitl_timeout=timedelta(minutes=5),
+ )
+
+ result = op.execute(context=MagicMock())
+
+ assert result == expected
+ assert type(result) is type(expected)
+
@pytest.mark.skipif(
not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible
with Airflow >= 3.1.0"
)
