kaxil commented on code in PR #74729:
URL: https://github.com/apache/airflow/pull/74729#discussion_r4289952933
##########
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:
This only catches a model at the top level. `ToolOutput(list[Summary])` or
`[list[Summary], int]` returns a list of `Summary` instances, and
`iter_pydantic_models` finds nothing in those shapes either, so the downstream
pull still hits the allow-list error. Convert nested models too, or narrow the
docstring to a top-level model?
##########
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:
This now runs on the default path for those five shapes, and python-mode
`model_dump()` leaves `HttpUrl` and `Path` objects in the dict, which serde
can't serialize (`TypeError: Cannot serialize object of type <class
'pydantic.networks.HttpUrl'>`). Pushing the instance worked because serde's
pydantic serializer dumps with `mode="json"`, so a `ToolOutput(Report)` whose
model has an `HttpUrl` field used to push fine and now fails at the push, after
the model calls are paid for. With `durable=True` the journal is already
cleaned up by then, so the retry re-runs every step.
Could the not-registered branch use `model_dump(mode="json")`, keeping plain
`model_dump()` when `_serialize_model_output` is set so the documented
`serialize_output` behaviour stays as is? A parametrized case with an `HttpUrl`
field pushed through `airflow.sdk.serde.serialize` would pin it.
##########
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:
Someone who added `Summary` to `[core] allowed_deserialization_classes` for
one of these shapes got the instance back before and gets a dict now, so
`result.text` raises `AttributeError`. Could `structured_output.rst` say that a
model from a list of output types, a marker or a function arrives as a dict? It
currently says `list` shapes are registered, which reads as covering `[Summary,
int]`.
--
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]
