kaxil commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4353422934
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,70 @@
_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 := [model for model, _ in _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.
+
+ With ``serialize_output`` the dump holds Python values, as the
``model_dump()`` of 1.0.0 does, so a
+ ``date`` or ``Decimal`` field keeps its type; an output with a set,
frozenset or dict key holding a model
+ is dumped to JSON values instead. A model the worker registers no class
for, without ``serialize_output``,
+ is dumped to JSON values as well, field names rather than aliases, so a
``HttpUrl``, ``Path`` or ``Enum``
+ field is a string that XCom accepts.
+ """
+ if not is_model_dump_required(output, output_type,
serialize_output=serialize_output):
+ return output
+ # Python mode is what 1.0.0 pushed, a minor release must not change it. A
model dumped to a dict cannot
+ # be a set member, so python mode raises on a set holding models.
+ if serialize_output and not any(in_set for _, in_set in
_iter_models_in_output(output)):
+ if isinstance(output, BaseModel):
+ return output.model_dump()
+ return TypeAdapter(Any).dump_python(output, mode="python")
Review Comment:
In 1.0.0, `model_dump()` ran only when the output itself was a `BaseModel`
(`llm.py:390` at the tag), so a `list[Link]` or `dict[str, Link]` went to XCom
as instances and serde's pydantic serializer handled each one. Here the whole
container is dumped in python mode, which leaves the `HttpUrl` inside a plain
dict, and the push fails with `Cannot serialize object of type HttpUrl`. I
checked this in breeze with `serialize_output=True`; it also hits Airflow
before 3.3 without the flag, since `_serialize_model_output` is always on
there. Could this branch keep 1.0.0's top-level-only `model_dump()`, with a
`serde.serialize` test for `list[Link]`?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -214,12 +221,14 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
connection does not silently switch off a control the author set.
Independent of
``require_approval``, which always asks. ``on_uncertain="review"``
needs Airflow 3.1+,
like ``require_approval``. Default ``None``: no gate, behaviour
unchanged.
- :param serialize_output: If ``True`` and ``output_type`` is a Pydantic
- ``BaseModel`` subclass, the model instance is dumped to a ``dict`` via
- ``model_dump()`` before being pushed to XCom. Default ``False`` --
- the Pydantic instance flows through XCom unchanged. Set to ``True``
- when a downstream consumer needs the dict shape (e.g. sending to an
- external system that expects JSON-style payloads).
+ :param serialize_output: If ``True``, a Pydantic ``BaseModel`` output, or a
+ list, tuple, set or dict holding models, is dumped to JSON-compatible
+ values with ``pydantic_core.to_jsonable_python`` before being pushed to
Review Comment:
This docstring and `AgentOperator`'s still say the dump uses
`to_jsonable_python`, while `agent.rst` and `structured_output.rst` now say
`model_dump()`, so the API reference and the guide disagree. The `LLMOperator`
row in `stability.rst` also doesn't mention that a `[MyModel, int]` output is
pushed as a dict without `serialize_output`.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1106,21 +1114,13 @@ def _complete_run(self, context: Context, result: Any,
*, attempt_usage: RunUsag
# So the approved answer is the last answer the agent returned,
and there is no need to parse
# it back from that text.
output = self._output_under_review
- 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.
- output = output.model_dump()
+ output = dump_output_for_xcom(
Review Comment:
1.0.0 dumped an unregistered model on this review path with `model_dump()`
(`agent.py:1097-1106` at the tag), so with `output_type=[Invoice, str]` a
`date` or `Decimal` field now reaches the downstream task as a string after an
approval. JSON mode also fixes `HttpUrl` here, so I'd be fine either way, but
the comment at `output_type.py:82` says python mode is what 1.0.0 pushed, which
isn't the case for this path.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -38,6 +46,70 @@
_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 := [model for model, _ in _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.
+
+ With ``serialize_output`` the dump holds Python values, as the
``model_dump()`` of 1.0.0 does, so a
+ ``date`` or ``Decimal`` field keeps its type; an output with a set,
frozenset or dict key holding a model
+ is dumped to JSON values instead. A model the worker registers no class
for, without ``serialize_output``,
+ is dumped to JSON values as well, field names rather than aliases, so a
``HttpUrl``, ``Path`` or ``Enum``
+ field is a string that XCom accepts.
+ """
+ if not is_model_dump_required(output, output_type,
serialize_output=serialize_output):
+ return output
+ # Python mode is what 1.0.0 pushed, a minor release must not change it. A
model dumped to a dict cannot
+ # be a set member, so python mode raises on a set holding models.
+ if serialize_output and not any(in_set for _, in_set in
_iter_models_in_output(output)):
+ if isinstance(output, BaseModel):
+ return output.model_dump()
+ return TypeAdapter(Any).dump_python(output, mode="python")
+ return to_jsonable_python(output, by_alias=False)
+
+
+def _iter_models_in_output(output: Any) -> Iterator[tuple[BaseModel, bool]]:
+ """
+ Yield the ``BaseModel`` instances in an output, each with whether it must
be hashable.
+
+ Looks inside list, tuple, set, frozenset and dict. A set member or a dict
key is hashable, so it is
+ flagged, and so is whatever is nested in one. Dataclass and NamedTuple
contents and model fields
Review Comment:
Small one: a NamedTuple is a `tuple`, so line 106 does walk its contents.
Only dataclass contents and model fields are skipped.
--
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]
