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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
