FrankYang0529 commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4290672883
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1081,6 +1070,20 @@ 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 value to push to XCom, as a dict if a downstream task
could not load the model."""
+ 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.
Review Comment:
Added a paragraph to `structured_output.rst`. The registration sentence now
says `list[...]` types instead of `list` shapes.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1081,6 +1070,20 @@ 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 value to push to XCom, as a dict if a downstream task
could not load the model."""
+ if isinstance(output, BaseModel) and (
Review Comment:
The new `_dump_unregistered_models` function walks lists, tuples and dicts.
If `iter_pydantic_models(output_type)` does not find the class of a model, the
function dumps that model with `mode="json"`.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1081,6 +1070,20 @@ 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 value to push to XCom, as a dict if a downstream task
could not load the model."""
+ 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.
+ return output.model_dump()
Review Comment:
Thanks for the suggestion. The not-registered branch now uses
`model_dump(mode="json")`. Plain `model_dump()` stays when
`_serialize_model_output` is set and the output is a model at the top level.
--
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]
