kaxil commented on code in PR #74537:
URL: https://github.com/apache/airflow/pull/74537#discussion_r4243924342
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -1133,6 +1133,31 @@ def
test_execute_with_hitl_returns_the_approved_output_as_output_type(
assert result == expected
assert type(result) is type(expected)
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible
with Airflow >= 3.1.0"
+ )
+
@patch("airflow.providers.common.ai.operators.agent.AgentOperator.run_hitl_review",
autospec=True)
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def
test_execute_with_hitl_parses_the_approved_output_for_a_list_of_output_types(
Review Comment:
This has the same setup as
`test_execute_with_hitl_returns_the_approved_output_as_output_type` just above,
so the `[Summary, int]` case could be one more `pytest.param` in that list.
That also brings in the `type(result) is type(expected)` check.
##########
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:
For `AgentOperator` the reviewer can't edit the text. The approve endpoint
only writes the action, and `_poll_loop` returns `session.current_output`,
which is the operator's own string. So with `output_type=[str, A]` an unedited
`"42"` answer comes back here as the int `42`. (For plain `str`,
`test_execute_with_hitl_returns_the_approved_output_as_output_type` pins `"42"`
staying a string.)
Would it be simpler for `AgentOperator` to keep the typed output from the
last run or regeneration and return it on approve? Then this branch is only
needed by `LLMOperator`, whose reviewers can edit.
##########
providers/common/ai/tests/unit/common/ai/utils/test_output_type.py:
##########
@@ -60,3 +65,26 @@ def test_returns_raw_on_schema_mismatch(self):
# ``A`` requires ``x: int`` -- this payload should fail validation
result = rehydrate_pydantic_output(A, '{"y": "no-x-field"}',
serialize_output=False)
assert result == '{"y": "no-x-field"}'
+
+ @pytest.mark.parametrize(
+ ("output_type", "raw", "expected"),
+ [
+ pytest.param([A, B], '{"x": 7}', {"x": 7},
id="list-of-output-types"),
Review Comment:
Nothing pins the `get_origin(output_type) is None` half of the gate.
Deleting it keeps every test green, because the only generic alias tested is
`list[str]`, where `TypeAdapter` and `json.loads` return the same value. A
`pytest.param(list[A], '[{"x": 7}]', [A(x=7)])` case would catch it.
##########
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:
This check also catches a `TypeAliasType` (`type LA = list[A]`) and a
`NewType`, which `TypeAdapter` already handled: neither is a class, and
`get_origin` returns `None` for both. Against this head, `TypeAliasType("LA",
list[A])` with `'[{"x":1}]'` now returns `[{'x': 1}]` instead of `[A(x=1)]`,
and `NewType("NA", A)` returns a dict instead of `A(x=1)`. pydantic-ai 2.33
accepts both as `output_type`, and `LLMOperator` shares this helper, so for it
that's a regression from 0.10.0. Could the gate match the shapes it is meant
for instead: a `list`/`tuple`, the pydantic-ai marker classes, and
`inspect.isfunction`/`ismethod`? A plain `callable()` check won't work because
a `NewType` is callable.
#74473 also rewrites this function, so the two PRs need reconciling before
either lands.
--
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]
