aaron-y-chen commented on code in PR #70132:
URL: https://github.com/apache/airflow/pull/70132#discussion_r3800796996
##########
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:
If `output_type=list[A]`, where `A` is a user-defined Pydantic model,
`rehydrate_pydantic_output()` returns `list[dict]` instead of `list[A]`.
I would expect the declared output type to be preserved, with the result
being `[A(x=7)]`. Is this behavior intentional?
```console
uv run --project providers/common/ai python -c '
from pydantic import BaseModel
from airflow.providers.common.ai.utils.output_type import
rehydrate_pydantic_output
class A(BaseModel):
x: int
result = rehydrate_pydantic_output(
list[A],
"[{\"x\": 7}]",
serialize_output=False,
)
print(result)
print(type(result[0]))
assert isinstance(result[0], A)
'
>>> [{'x': 7}]
>>> <class 'dict'>
>>> Traceback (most recent call last):
>>> File "<string>", line 15, in <module>
>>> AssertionError
```
--
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]
