Lee-W commented on code in PR #74695: URL: https://github.com/apache/airflow/pull/74695#discussion_r4323872505
########## providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py: ########## @@ -38,6 +45,24 @@ _CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError, KeyError, IndexError) +def is_model_dump_required(output: Any, output_type: Any, *, serialize_output: bool) -> bool: + """ + Tell whether a ``BaseModel`` output must be pushed to XCom as 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 those shapes are dumped. + Without ``iter_pydantic_models`` there is no evidence either way, so the caller's ``serialize_output`` decides. + """ + if not isinstance(output, BaseModel): + return False + if serialize_output: + return True + if iter_pydantic_models is None: + return False + return type(output) not in iter_pydantic_models(output_type) Review Comment: Switched to a shared `dump_output_for_xcom` that dumps with `pydantic_core.to_jsonable_python`, so `HttpUrl`, `Path` and `Enum` fields serialize. A test feeds the dumped model to `serde.serialize`. Models inside a list, tuple, set or dict output are dumped too when the worker registers no class for them, so `[list[Item], str]` is covered. Field names are dumped, not aliases, like `dump_output_to_json`. Dict keys, dataclass and NamedTuple contents are not walked. -- 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]
