ColtenOuO commented on code in PR #70132:
URL: https://github.com/apache/airflow/pull/70132#discussion_r3802507017
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -37,16 +65,38 @@ def rehydrate_pydantic_output(
reviewer. ``str`` outputs pass through unchanged; any other ``output_type``
(``BaseModel`` subclass, ``int``, ``list[str]``, ...) is validated with a
pydantic ``TypeAdapter``. When validation fails (reviewer edited the string
- into something the type rejects), returns ``raw`` unchanged.
+ into something the type rejects), returns ``raw`` unchanged. A
+ ``ToolOutput``/``NativeOutput``/``PromptedOutput`` marker wrapping a single
+ type is unwrapped first. An ``output_type`` that is a ``[A, B]`` union
list,
+ a marker wrapping one, or a pydantic-ai *output function* falls back to
+ plain JSON -- ``TypeAdapter`` builds an *arguments* schema for an output
+ function rather than raising, so validating against it would call the
+ function a second time with reviewer-controlled input.
When ``serialize_output`` is ``True``, returns the model dumped to a
``dict`` -- matches the operator's ``serialize_output=True`` opt-in for
consumers that want the dict shape.
"""
if output_type is str:
return raw
+
+ unwrapped = output_type
+ if isinstance(output_type, _OUTPUT_MARKERS):
+ unwrapped = output_type.output if isinstance(output_type, ToolOutput)
else output_type.outputs
+
+ if not isinstance(unwrapped, type):
+ # Not a plain class: an output function, a ``[A, B]`` union list, or a
+ # marker wrapping one. These reach us because the operator passes
+ # output_type straight to Agent(...). Fall back to plain JSON rather
+ # than risk building a schema that re-invokes an output function.
+ try:
+ return json.loads(raw)
Review Comment:
I think that is not intentional. Like a new bug cause this PR : (
`isinstance(unwrapped, type)` is `False` for a generic alias like `list[A]`
(it's `types.GenericAlias`), so it was routed into
the plain `json.loads` fallback meant for output functions / `[A, B]` union
lists, silently downgrading `list[A]` to `list[dict]`.
Fixed in 6280539 by also checking `typing.get_origin(unwrapped) is not None`
so generic aliases still reach `TypeAdapter`. Note `callable(list[A])` is
`True`, so a callable-based check can't distinguish it from an output
function — `get_origin()` is the right signal.
Also generalized the `serialize_output=True` dump step
(`adapter.dump_python(rehydrated, mode="python")` instead of an
`isinstance(rehydrated, BaseModel)` check), since a rehydrated `list[A]` is a
list, not a `BaseModel`, and needs the same dict-shape dump its elements used
to get.
--
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]
