You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "FrankYang0529 (via GitHub)" <gi...@apache.org> on 2026/10/14 13:09:56 UTC
[PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
FrankYang0529 opened a new pull request, #74729:
URL: https://github.com/apache/airflow/pull/74729
Follow-up for https://github.com/apache/airflow/pull/74537#pullrequestreview-5490873096.
- Move the condition that #74537 added to the HITL review path into a new method, `_convert_output_for_xcom`. Both returns of `_complete_run` now call this method. Without review, `AgentOperator` now also pushes a model as a dict when the worker would not register the model's class, for example with `output_type=[Summary, int]`.
- Rename `test_execute_with_hitl_rehydrates_base_model` to `test_execute_with_hitl_returns_the_model_instance` and remove its outdated comment about the JSON string.
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: https://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
For user-facing UI changes, please attach before/after screenshots (or a short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [X] Yes - Claude Code
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on PR #74729:
URL: https://github.com/apache/airflow/pull/74729#issuecomment-6361813148
Thanks! I fixed all three nits:
- The changelog note now leads with the five output types and the worker, like `structured_output.rst`.
- Added `timedelta` to the HITL lists in `structured_output.rst` and the changelog note.
- `test_execute_with_hitl_returns_the_model_instance` now has `run_hitl_review` return `'{"text": "Edited"}'`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4289952933
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1081,6 +1070,20 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the value to push to XCom, as a dict if a downstream task could not load the model."""
+ if isinstance(output, BaseModel) and (
Review Comment:
This only catches a model at the top level. `ToolOutput(list[Summary])` or `[list[Summary], int]` returns a list of `Summary` instances, and `iter_pydantic_models` finds nothing in those shapes either, so the downstream pull still hits the allow-list error. Convert nested models too, or narrow the docstring to a top-level model?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1081,6 +1070,20 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the value to push to XCom, as a dict if a downstream task could not load the model."""
+ if isinstance(output, BaseModel) and (
+ self._serialize_model_output or type(output) not in iter_pydantic_models(self.output_type)
+ ):
+ # A downstream task can load a model from XCom only if the worker registered the model's class.
+ # The worker registers the classes that iter_pydantic_models finds in output_type.
+ # iter_pydantic_models finds no class in a list of output types, an output marker, an output
+ # function, a TypeAliasType or a NewType, so the operator pushes the model as a dict.
+ # When iter_pydantic_models could not be imported, _serialize_model_output is always True, so
+ # the condition above never reaches the iter_pydantic_models call.
+ return output.model_dump()
Review Comment:
This now runs on the default path for those five shapes, and python-mode `model_dump()` leaves `HttpUrl` and `Path` objects in the dict, which serde can't serialize (`TypeError: Cannot serialize object of type <class 'pydantic.networks.HttpUrl'>`). Pushing the instance worked because serde's pydantic serializer dumps with `mode="json"`, so a `ToolOutput(Report)` whose model has an `HttpUrl` field used to push fine and now fails at the push, after the model calls are paid for. With `durable=True` the journal is already cleaned up by then, so the retry re-runs every step.
Could the not-registered branch use `model_dump(mode="json")`, keeping plain `model_dump()` when `_serialize_model_output` is set so the documented `serialize_output` behaviour stays as is? A parametrized case with an `HttpUrl` field pushed through `airflow.sdk.serde.serialize` would pin it.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1081,6 +1070,20 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the value to push to XCom, as a dict if a downstream task could not load the model."""
+ if isinstance(output, BaseModel) and (
+ self._serialize_model_output or type(output) not in iter_pydantic_models(self.output_type)
+ ):
+ # A downstream task can load a model from XCom only if the worker registered the model's class.
Review Comment:
Someone who added `Summary` to `[core] allowed_deserialization_classes` for one of these shapes got the instance back before and gets a dict now, so `result.text` raises `AttributeError`. Could `structured_output.rst` say that a model from a list of output types, a marker or a function arrives as a dict? It currently says `list` shapes are registered, which reads as covering `[Summary, int]`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4344107613
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1120,6 +1129,23 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the XCom value, with dicts in place of models that a downstream task could not load."""
+ if self._serialize_model_output and isinstance(output, BaseModel):
+ return output.model_dump()
+ if not _CORE_WALKER:
+ # iter_pydantic_models could not be imported, so the worker registers no class.
+ # _serialize_model_output is always True when that import fails, so a model at the top level was
+ # already dumped above. The class of a model inside a list, tuple or dict is not registered
+ # either, so this call passes an empty set of registered classes to dump that model too.
+ return _dump_unregistered_models(output, set())
+ # A downstream task can load a model from XCom only if the worker registered the model's class.
+ # The worker registers the classes that iter_pydantic_models finds in output_type.
+ # iter_pydantic_models finds no class in a list of output types, an output marker, an output
+ # function, a TypeAliasType or a NewType. For these output types the operator pushes each model
+ # as a dict, including models inside a list, tuple or dict.
+ return _dump_unregistered_models(output, set(iter_pydantic_models(self.output_type)))
Review Comment:
Added the early return at the top of `_convert_output_for_xcom`. With `do_xcom_push=False`, the method now returns the output unchanged and dumps nothing. Thanks for the review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4339928561
##########
providers/common/ai/docs/structured_output.rst:
##########
@@ -50,6 +50,34 @@ loaded Dag's tasks at startup and registers each declared class, so it also work
for mapped producers (``.expand(...)``) and for workers that load Dags from a
cache that bypasses operator construction.
+The worker finds no class to register when ``output_type`` is a list of output
+types such as ``[Summary, int]``, an output marker such as
+``ToolOutput(Summary)``, an output function, a ``TypeAliasType`` or a
+``NewType``. For these output types, ``AgentOperator`` and ``@task.agent`` push
+each model in the output as a JSON-compatible dict, including models inside a
+list, tuple or dict. A model is pushed as a dict even if its class is listed in
+``[core] allowed_deserialization_classes``, or if another task in the same Dag
+declares the class as its ``output_type``. A downstream task therefore reads
+``result["field"]``, not ``result.field``.
+
+In a JSON-compatible dict, values such as ``datetime``, ``date``, ``UUID`` and
+``Decimal`` arrive as strings, tuples and sets arrive as lists, and ``Enum``
+members arrive as their values. In HITL review, the released common.ai 1.0.0
+dumped a top-level model of these output types with plain ``model_dump()``
+after the reviewer approved the model. On Airflow 3.3 and later without
+``serialize_output=True``, the ``datetime`` fields of such a model therefore
+arrive as strings after an upgrade, not as ``datetime`` objects.
Review Comment:
Unless this makes it before 1.1.0 is released.. It is currently being voted on https://lists.apache.org/thread.html/sn28295qls3pk7lv5nx3smpgqtmz1m1g
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil merged PR #74729:
URL: https://github.com/apache/airflow/pull/74729
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4321416289
##########
providers/common/ai/docs/structured_output.rst:
##########
@@ -50,6 +50,15 @@ loaded Dag's tasks at startup and registers each declared class, so it also work
for mapped producers (``.expand(...)``) and for workers that load Dags from a
cache that bypasses operator construction.
+The worker finds no class to register when ``output_type`` is a list of output
+types such as ``[Summary, int]``, an output marker such as
+``ToolOutput(Summary)``, an output function, a ``TypeAliasType`` or a
+``NewType``. For these output types, ``AgentOperator`` and ``@task.agent`` push
+each model in the output as a JSON-compatible dict, including models inside a
Review Comment:
Added two paragraphs. The first paragraph gives examples of values that change type in JSON mode. It also covers HITL review. The second paragraph says that a top-level model is still dumped with plain `model_dump()` in either of two cases: `serialize_output=True` is set, or Airflow is older than 3.3.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4321405440
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1077,6 +1087,22 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the XCom value, with dicts in place of models that a downstream task could not load."""
+ if self._serialize_model_output and isinstance(output, BaseModel):
+ return output.model_dump()
+ if not _CORE_WALKER:
+ # iter_pydantic_models could not be imported, so the worker registers no class.
+ # _serialize_model_output is always True when that import fails, so a model at the top level
+ # was already dumped above. Any other output is pushed unchanged.
+ return output
Review Comment:
Changed to use `_dump_unregistered_models(output, set())`. Thanks for the review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4290672883
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1081,6 +1070,20 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the value to push to XCom, as a dict if a downstream task could not load the model."""
+ if isinstance(output, BaseModel) and (
+ self._serialize_model_output or type(output) not in iter_pydantic_models(self.output_type)
+ ):
+ # A downstream task can load a model from XCom only if the worker registered the model's class.
Review Comment:
Added a paragraph to `structured_output.rst`. The registration sentence now says `list[...]` types instead of `list` shapes.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1081,6 +1070,20 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the value to push to XCom, as a dict if a downstream task could not load the model."""
+ if isinstance(output, BaseModel) and (
Review Comment:
The new `_dump_unregistered_models` function walks lists, tuples and dicts. If `iter_pydantic_models(output_type)` does not find the class of a model, the function dumps that model with `mode="json"`.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1081,6 +1070,20 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the value to push to XCom, as a dict if a downstream task could not load the model."""
+ if isinstance(output, BaseModel) and (
+ self._serialize_model_output or type(output) not in iter_pydantic_models(self.output_type)
+ ):
+ # A downstream task can load a model from XCom only if the worker registered the model's class.
+ # The worker registers the classes that iter_pydantic_models finds in output_type.
+ # iter_pydantic_models finds no class in a list of output types, an output marker, an output
+ # function, a TypeAliasType or a NewType, so the operator pushes the model as a dict.
+ # When iter_pydantic_models could not be imported, _serialize_model_output is always True, so
+ # the condition above never reaches the iter_pydantic_models call.
+ return output.model_dump()
Review Comment:
Thanks for the suggestion. The not-registered branch now uses `model_dump(mode="json")`. Plain `model_dump()` stays when `_serialize_model_output` is set and the output is a model at the top level.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4320045120
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1077,6 +1087,22 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the XCom value, with dicts in place of models that a downstream task could not load."""
+ if self._serialize_model_output and isinstance(output, BaseModel):
+ return output.model_dump()
+ if not _CORE_WALKER:
+ # iter_pydantic_models could not be imported, so the worker registers no class.
+ # _serialize_model_output is always True when that import fails, so a model at the top level
+ # was already dumped above. Any other output is pushed unchanged.
+ return output
Review Comment:
Below 3.3 this still pushes `ToolOutput(list[Summary])` or `[list[Summary], int]` as a list of model instances, and since nothing is registered there, the downstream pull hits the same allow-list error. A top-level model is already dumped on this path, so `return _dump_unregistered_models(output, set())` would make nested ones match. The provider floor is still 2.11.
##########
providers/common/ai/docs/structured_output.rst:
##########
@@ -50,6 +50,15 @@ loaded Dag's tasks at startup and registers each declared class, so it also work
for mapped producers (``.expand(...)``) and for workers that load Dags from a
cache that bypasses operator construction.
+The worker finds no class to register when ``output_type`` is a list of output
+types such as ``[Summary, int]``, an output marker such as
+``ToolOutput(Summary)``, an output function, a ``TypeAliasType`` or a
+``NewType``. For these output types, ``AgentOperator`` and ``@task.agent`` push
+each model in the output as a JSON-compatible dict, including models inside a
Review Comment:
In JSON mode datetimes, UUIDs and Decimals arrive as strings and tuples as lists. 1.0.0 dumped the HITL output in python mode, so a HITL user with a datetime field sees the type change on upgrade. Could the paragraph say that, and that `serialize_output=True` still dumps a top-level model with plain `model_dump()`, so Enum or HttpUrl fields stay as objects there?
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -1032,6 +1046,99 @@ def test_execute_structured_output(self, mock_hook_cls, make_mock_run_result):
assert result.text == "Great"
assert result.score == 0.95
+ @pytest.mark.parametrize(
+ "output_type",
+ [
+ pytest.param([Summary, int], id="list-of-output-types"),
+ pytest.param(ToolOutput(Summary), id="output-marker"),
+ pytest.param(_summarize, id="output-function"),
+ pytest.param(SummaryAlias, id="type-alias"),
+ pytest.param(SummaryNewType, id="new-type"),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
+ def test_execute_dumps_a_model_the_worker_does_not_register(
+ self, mock_hook_cls, make_mock_run_result, output_type
+ ):
+ mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent(
+ Summary(text="ok"), make_mock_run_result
+ )
+ op = AgentOperator(
+ task_id="test", prompt="Analyze this", llm_conn_id="my_llm", output_type=output_type
+ )
+
+ result = op.execute(context=MagicMock())
+
+ assert result == {"text": "ok", "score": 0.0}
+
+ @requires_typed_xcom
+ @pytest.mark.parametrize(
+ ("output_type", "output", "expected", "serialize_output"),
+ [
+ pytest.param(
+ ToolOutput(list[Summary]),
+ [Summary(text="ok")],
+ [{"text": "ok", "score": 0.0}],
+ False,
+ id="list-in-output-marker",
+ ),
+ pytest.param(
+ [list[Summary], int],
+ [Summary(text="ok")],
+ [{"text": "ok", "score": 0.0}],
+ False,
+ id="list-in-list-of-output-types",
+ ),
+ pytest.param(
+ ToolOutput(tuple[Summary, int]),
+ (Summary(text="ok"), 1),
+ ({"text": "ok", "score": 0.0}, 1),
+ False,
+ id="tuple-in-output-marker",
+ ),
+ pytest.param(
+ ToolOutput(dict[str, Summary]),
+ {"a": Summary(text="ok")},
+ {"a": {"text": "ok", "score": 0.0}},
+ False,
+ id="dict-in-output-marker",
+ ),
+ pytest.param(
+ ToolOutput(list[Summary]),
+ [Summary(text="ok")],
+ [{"text": "ok", "score": 0.0}],
+ True,
+ id="list-in-output-marker-with-serialize-output",
+ ),
+ pytest.param(
+ ToolOutput(LinkedSummary),
+ LinkedSummary(kind=SummaryKind.SHORT, url=HttpUrl("https://example.com/a")),
+ {"kind": "short", "url": "https://example.com/a"},
+ False,
+ id="enum-and-url-fields",
+ ),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
+ def test_execute_pushes_each_model_the_worker_does_not_register_as_json(
Review Comment:
Every row here has an empty registered set, so nothing checks that `list[Summary]` or `dict[str, Summary]` keeps the instances. Passing `set()` instead of `registered` in the three container branches leaves the whole file green. Could you add a `list[Summary]` row that expects `[Summary(text="ok")]` and assert `type(result[0]) is Summary`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4321419701
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -1032,6 +1046,99 @@ def test_execute_structured_output(self, mock_hook_cls, make_mock_run_result):
assert result.text == "Great"
assert result.score == 0.95
+ @pytest.mark.parametrize(
+ "output_type",
+ [
+ pytest.param([Summary, int], id="list-of-output-types"),
+ pytest.param(ToolOutput(Summary), id="output-marker"),
+ pytest.param(_summarize, id="output-function"),
+ pytest.param(SummaryAlias, id="type-alias"),
+ pytest.param(SummaryNewType, id="new-type"),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
+ def test_execute_dumps_a_model_the_worker_does_not_register(
+ self, mock_hook_cls, make_mock_run_result, output_type
+ ):
+ mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent(
+ Summary(text="ok"), make_mock_run_result
+ )
+ op = AgentOperator(
+ task_id="test", prompt="Analyze this", llm_conn_id="my_llm", output_type=output_type
+ )
+
+ result = op.execute(context=MagicMock())
+
+ assert result == {"text": "ok", "score": 0.0}
+
+ @requires_typed_xcom
+ @pytest.mark.parametrize(
+ ("output_type", "output", "expected", "serialize_output"),
+ [
+ pytest.param(
+ ToolOutput(list[Summary]),
+ [Summary(text="ok")],
+ [{"text": "ok", "score": 0.0}],
+ False,
+ id="list-in-output-marker",
+ ),
+ pytest.param(
+ [list[Summary], int],
+ [Summary(text="ok")],
+ [{"text": "ok", "score": 0.0}],
+ False,
+ id="list-in-list-of-output-types",
+ ),
+ pytest.param(
+ ToolOutput(tuple[Summary, int]),
+ (Summary(text="ok"), 1),
+ ({"text": "ok", "score": 0.0}, 1),
+ False,
+ id="tuple-in-output-marker",
+ ),
+ pytest.param(
+ ToolOutput(dict[str, Summary]),
+ {"a": Summary(text="ok")},
+ {"a": {"text": "ok", "score": 0.0}},
+ False,
+ id="dict-in-output-marker",
+ ),
+ pytest.param(
+ ToolOutput(list[Summary]),
+ [Summary(text="ok")],
+ [{"text": "ok", "score": 0.0}],
+ True,
+ id="list-in-output-marker-with-serialize-output",
+ ),
+ pytest.param(
+ ToolOutput(LinkedSummary),
+ LinkedSummary(kind=SummaryKind.SHORT, url=HttpUrl("https://example.com/a")),
+ {"kind": "short", "url": "https://example.com/a"},
+ False,
+ id="enum-and-url-fields",
+ ),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
+ def test_execute_pushes_each_model_the_worker_does_not_register_as_json(
Review Comment:
Added `test_execute_keeps_each_model_the_worker_registers` with `list[Summary]`, `tuple[Summary, int]` and `dict[str, Summary]` rows.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "FrankYang0529 (via GitHub)" <gi...@apache.org>.
FrankYang0529 commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4344106225
##########
providers/common/ai/docs/structured_output.rst:
##########
@@ -50,6 +50,34 @@ loaded Dag's tasks at startup and registers each declared class, so it also work
for mapped producers (``.expand(...)``) and for workers that load Dags from a
cache that bypasses operator construction.
+The worker finds no class to register when ``output_type`` is a list of output
+types such as ``[Summary, int]``, an output marker such as
+``ToolOutput(Summary)``, an output function, a ``TypeAliasType`` or a
+``NewType``. For these output types, ``AgentOperator`` and ``@task.agent`` push
+each model in the output as a JSON-compatible dict, including models inside a
+list, tuple or dict. A model is pushed as a dict even if its class is listed in
+``[core] allowed_deserialization_classes``, or if another task in the same Dag
+declares the class as its ``output_type``. A downstream task therefore reads
+``result["field"]``, not ``result.field``.
+
+In a JSON-compatible dict, values such as ``datetime``, ``date``, ``UUID`` and
+``Decimal`` arrive as strings, tuples and sets arrive as lists, and ``Enum``
+members arrive as their values. In HITL review, the released common.ai 1.0.0
+dumped a top-level model of these output types with plain ``model_dump()``
+after the reviewer approved the model. On Airflow 3.3 and later without
+``serialize_output=True``, the ``datetime`` fields of such a model therefore
+arrive as strings after an upgrade, not as ``datetime`` objects.
Review Comment:
Changed the text to say 1.0.0 and 1.1.0. If a later 1.1.0 candidate includes this PR, I'll change both places (the doc and the changelog note) to say only 1.0.0.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4339910692
##########
providers/common/ai/docs/structured_output.rst:
##########
@@ -50,6 +50,34 @@ loaded Dag's tasks at startup and registers each declared class, so it also work
for mapped producers (``.expand(...)``) and for workers that load Dags from a
cache that bypasses operator construction.
+The worker finds no class to register when ``output_type`` is a list of output
+types such as ``[Summary, int]``, an output marker such as
+``ToolOutput(Summary)``, an output function, a ``TypeAliasType`` or a
+``NewType``. For these output types, ``AgentOperator`` and ``@task.agent`` push
+each model in the output as a JSON-compatible dict, including models inside a
+list, tuple or dict. A model is pushed as a dict even if its class is listed in
+``[core] allowed_deserialization_classes``, or if another task in the same Dag
+declares the class as its ``output_type``. A downstream task therefore reads
+``result["field"]``, not ``result.field``.
+
+In a JSON-compatible dict, values such as ``datetime``, ``date``, ``UUID`` and
+``Decimal`` arrive as strings, tuples and sets arrive as lists, and ``Enum``
+members arrive as their values. In HITL review, the released common.ai 1.0.0
+dumped a top-level model of these output types with plain ``model_dump()``
+after the reviewer approved the model. On Airflow 3.3 and later without
+``serialize_output=True``, the ``datetime`` fields of such a model therefore
+arrive as strings after an upgrade, not as ``datetime`` objects.
Review Comment:
`providers-common-ai/1.1.0rc1` still has the plain `model_dump()` on this path, so this should say 1.0.0 and 1.1.0. The non-HITL change isn't covered either: on 3.3+ these output types pushed the instance before, so someone who allow-listed the class (or declared it as another task's `output_type`) got `result.field` working and now gets a dict and an `AttributeError`. Since both change the type a downstream task receives, could this also get a `.. note::` at the top of `changelog.rst`, like the 1.0.0 ones?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1120,6 +1129,23 @@ def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsag
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the XCom value, with dicts in place of models that a downstream task could not load."""
+ if self._serialize_model_output and isinstance(output, BaseModel):
+ return output.model_dump()
+ if not _CORE_WALKER:
+ # iter_pydantic_models could not be imported, so the worker registers no class.
+ # _serialize_model_output is always True when that import fails, so a model at the top level was
+ # already dumped above. The class of a model inside a list, tuple or dict is not registered
+ # either, so this call passes an empty set of registered classes to dump that model too.
+ return _dump_unregistered_models(output, set())
+ # A downstream task can load a model from XCom only if the worker registered the model's class.
+ # The worker registers the classes that iter_pydantic_models finds in output_type.
+ # iter_pydantic_models finds no class in a list of output types, an output marker, an output
+ # function, a TypeAliasType or a NewType. For these output types the operator pushes each model
+ # as a dict, including models inside a list, tuple or dict.
+ return _dump_unregistered_models(output, set(iter_pydantic_models(self.output_type)))
Review Comment:
With `do_xcom_push=False` nothing gets pushed, but this still runs `model_dump(mode="json")`, which raises for a model with non-UTF-8 `bytes` (`UnicodeDecodeError`) or an `arbitrary_types_allowed` field (`PydanticSerializationError`). Such a task with `ToolOutput(...)` on 3.3+ succeeded before and now fails after the model calls are paid for. Returning `output` early when `not self.do_xcom_push` would avoid it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix pulling AgentOperator output for list, marker and function types [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #74729:
URL: https://github.com/apache/airflow/pull/74729#issuecomment-6350161896
@FrankYang0529 Could you fix those so we can merge the PR, plz
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org