FrankYang0529 commented on code in PR #74537:
URL: https://github.com/apache/airflow/pull/74537#discussion_r4244396608
##########
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:
+ # TypeAdapter raises for a list of output types or for a marker. For
an output function, TypeAdapter
+ # builds a schema for the function's arguments, so validating the
reviewed text would call the
+ # function again with that text as its arguments.
+ try:
+ return json.loads(raw)
Review Comment:
Agreed. `AgentOperator` keeps the output of the last run or regeneration and
returns it on approval, so with `output_type=[str, A]` an unedited `"42"`
answer stays a string. There is one difference from returning the output
unchanged. When the approved answer is a model whose class the worker would not
register, the operator returns the model as a dict. The operator decides
whether the worker would register the class by calling `iter_pydantic_models`,
the same function that the worker uses to register classes.
--
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]
