FrankYang0529 commented on code in PR #74537:
URL: https://github.com/apache/airflow/pull/74537#discussion_r4244376082
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -39,12 +40,26 @@ def rehydrate_pydantic_output(
pydantic ``TypeAdapter``. When validation fails (reviewer edited the string
into something the type rejects), returns ``raw`` unchanged.
+ pydantic-ai also accepts an ``output_type`` that is neither a class nor a
+ generic alias: a list of output types, an output marker such as
+ ``ToolOutput``, or an output function. For these output types, ``raw`` is
+ parsed with ``json.loads`` instead, and ``raw`` is returned unchanged when
+ it is not JSON.
+
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
+ if not isinstance(output_type, type) and get_origin(output_type) is None:
Review Comment:
Thanks for the review. I reverted the change for
`rehydrate_pydantic_output`. `AgentOperator` doesn't convert the answer text,
so it does not need to match the shapes you listed. The operator now only
decides whether to return an approved model as a dict. To make that decision,
it calls `iter_pydantic_models`, the same function that the worker uses to
register classes for XCom. The worker finds no class to register in a
`TypeAliasType` or a `NewType`, so for those types the operator also returns an
approved model as a dict.
--
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]
