Lee-W commented on code in PR #74695: URL: https://github.com/apache/airflow/pull/74695#discussion_r4348587597
########## 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: `dump_output_for_xcom` now dumps with `model_dump()` when `serialize_output=True` (which includes task-sdk versions without the walker), as 1.0.0 did, so `date`, `Decimal` and `UUID` fields keep their types. `to_jsonable_python` is used only for a model whose class the worker cannot register. One exception: an output with a set, frozenset or dict key holding a model is dumped to JSON values, because pydantic cannot hold a dumped model in a set or as a key. -- 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]
