kaxil commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4340279710
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,56 @@
_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError,
KeyError, IndexError)
+def is_model_dump_required(output: Any, output_type: Any, *, serialize_output:
bool) -> bool:
+ """
+ Tell whether an output holding ``BaseModel`` instances must be pushed to
XCom dumped.
+
+ The instances are looked up in the output itself, including the elements
of a list, tuple, set or
+ frozenset and the values of a dict. A downstream task can load a model
from XCom only if the worker
+ registered the model's class, and the worker registers the classes
``iter_pydantic_models`` finds in
+ ``output_type``. It finds none in a list of output types, an output
marker, an output function, a
+ ``TypeAliasType`` or a ``NewType``, so an instance of a class it did not
find is dumped.
+ Without ``iter_pydantic_models`` there is no evidence either way, so the
caller's ``serialize_output`` decides.
+ """
+ if not (found := list(_iter_models_in_output(output))):
+ return False
+ if serialize_output:
+ return True
+ if iter_pydantic_models is None:
+ return False
+ registered = set(iter_pydantic_models(output_type))
+ return any(type(model) not in registered for model in found)
+
+
+def dump_output_for_xcom(output: Any, output_type: Any, *, serialize_output:
bool) -> Any:
+ """
+ Return the output as it is pushed to XCom: dumped when
``is_model_dump_required``, else unchanged.
+
+ The dump is JSON-compatible, as the serde ``BaseModel`` serializer's is,
so a ``HttpUrl``, ``Path`` or
+ ``Enum`` field is a string that XCom accepts. Field names, not aliases,
match ``dump_output_to_json``.
+ """
+ if is_model_dump_required(output, output_type,
serialize_output=serialize_output):
+ return to_jsonable_python(output, by_alias=False)
Review Comment:
I asked for `to_jsonable_python` with the newly dumped models in mind, but
this also runs on the paths 1.0.0 dumped with python-mode `model_dump()`:
`serialize_output=True`, and every model output before Airflow 3.3
(`serialize_output or not _CORE_WALKER`). There a `date`, `datetime`, `Decimal`
or `UUID` field now reaches the downstream task as a string instead of the
object serde used to round-trip, so code like `inv["due"] < today` starts
raising after a minor upgrade of a stable operator. Could the JSON-mode dump
apply only to models dumped because the worker can't register their class, keep
`model_dump()` for the other two, and add a test that a `date` field survives
`serialize_output=True`?
##########
providers/common/ai/docs/structured_output.rst:
##########
@@ -64,4 +67,7 @@ that pattern work.
If a downstream consumer needs the dict shape (e.g. forwarding to an external
system that expects JSON-style payloads), pass ``serialize_output=True`` and
the
-operator calls ``model_dump()`` before pushing to XCom.
+operator dumps the model with ``pydantic_core.to_jsonable_python`` before
pushing to XCom,
+so ``HttpUrl``, ``Path`` and ``Enum`` fields become plain JSON values. The
operator also dumps an
+output that is a list, tuple, set or dict holding models, and any model of a
shape the worker
+registers no class for (such as a list of types), even without
``serialize_output``.
Review Comment:
This reads as if a list or dict of models is always dumped, but a registered
`output_type=list[MyModel]` still pushes the instances unless
`serialize_output=True`. Only a model whose class the worker can't register is
dumped without it. Could the sentence split those two cases?
##########
providers/common/ai/tests/unit/common/ai/utils/test_output_type.py:
##########
@@ -16,28 +16,53 @@
# under the License.
from __future__ import annotations
+import json
from collections import deque
from collections.abc import Sequence
from dataclasses import dataclass
+from datetime import UTC, datetime
from enum import Enum
+from pathlib import Path
from typing import Annotated, Literal, NamedTuple
+from unittest import mock
import pytest
-from pydantic import BaseModel, ConfigDict, Field, computed_field,
model_validator
+from pydantic import BaseModel, ConfigDict, Field, HttpUrl, computed_field,
model_validator
from pydantic_ai import NativeOutput, PromptedOutput, ToolOutput
from airflow.providers.common.ai.exceptions import
ReviewedOutputValidationError
from airflow.providers.common.ai.utils.output_type import (
_is_container_mismatch,
+ dump_output_for_xcom,
dump_output_to_json,
+ is_model_dump_required,
rehydrate_pydantic_output,
)
+from airflow.sdk.serde import deserialize, serialize
Review Comment:
This import fails collection in the 2.11.1, 3.0.6 and 3.1.8 compat jobs (`No
module named 'airflow.sdk.serde'`, "Interrupted: 2 errors during collection"),
so none of the provider's tests run there. `test_llm.py:61` has the same one.
Moving both into the `try` block just below works, since the tests that use
them are already gated by `requires_walker` / `requires_typed_xcom`.
--
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]
