kaxil commented on code in PR #70132:
URL: https://github.com/apache/airflow/pull/70132#discussion_r3755990602
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -18,11 +18,29 @@
from __future__ import annotations
+import json
from typing import Any
from pydantic import BaseModel, TypeAdapter, ValidationError
+def dump_output_to_json(output: Any) -> str:
+ """
+ Serialize an LLM output into the string carried through human review.
+
+ The inverse of :func:`rehydrate_pydantic_output`: both sides of the review
+ round-trip live here so they cannot drift apart.
+ """
+ if isinstance(output, BaseModel):
+ return output.model_dump_json()
+ if isinstance(output, str):
+ return output
+ try:
+ return TypeAdapter(type(output)).dump_json(output).decode()
+ except Exception:
Review Comment:
This turns what used to be a crash into a silent repr. At the merge-base
`approval.py` had no `try`/`except` here, so a non-serializable output raised
before the reviewer saw anything; now `str(output)` runs and the reviewer is
shown a Python object address -- `dump_output_to_json(Widget(3))` returns
`'<__main__.Widget object at 0x1023430b0>'`.
Reachable on the path the new test comment already names:
`Agent(TestModel(), output_type=make_widget)` builds on just a `UserWarning`
("Could not generate return schema for 'make_widget': unsupported return type
... Falling back to unconstrained schema") and returns a `Widget`. That string
goes into the approval body and the `airflow_hitl_review_agent_output_N` XComs,
and comes back out of `rehydrate_pydantic_output` as the same plain string --
the repr-instead-of-JSON failure this PR is fixing. Preferring the fallback
over the crash is right, but it should say so in the log; the module has no
logger, and `utils/file_analysis.py` in the same package already has one.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -46,7 +65,19 @@ def rehydrate_pydantic_output(
if output_type is str:
return raw
try:
- rehydrated = TypeAdapter(output_type).validate_json(raw)
+ adapter = TypeAdapter(output_type)
Review Comment:
`TypeAdapter` doesn't raise for an output function -- it builds a `call`
schema -- so this guard never fires for them, and `adapter.validate_json(raw)`
below validates `raw` as the function's *arguments* and invokes it. Verified on
both the provider's `pydantic-ai>=2.0.0` floor and 2.22.0, with
```python
def open_ticket(severity: int) -> Ticket:
SIDE_EFFECTS.append(...) # e.g. POST to Jira
return Ticket(severity=severity)
```
`rehydrate_pydantic_output(open_ticket, '{"severity":0}')` returns
`Ticket(severity=0)` **and runs `open_ticket` a second time**. If the reviewer
edits the approved string it re-runs with the edited arguments (`'{"severity":
9999}'` gives a second call with `severity=9999`).
That is new here: `Agent(TestModel(), output_type=open_ticket)` is accepted
by pydantic-ai, and at the merge-base `isinstance(open_ticket, type)` was False
so `execute()` took `json.loads` and the function was never re-invoked (same
probe against the pre-PR source: no new side effects). Output functions are
pydantic-ai's tool-shaped final-output hook, so the ones with side effects now
fire twice, after the reviewer has already approved.
The comment just below lists output functions as a case this branch handles,
which is the tell: "did `TypeAdapter` raise" isn't the right test. A positive
check works -- reject a non-class callable before adapting, and unwrap the
markers instead of falling back for them (`ToolOutput(A).output`,
`NativeOutput(A).outputs`, `PromptedOutput(A).outputs` all give `A`). That also
fixes the smaller problem that `ToolOutput(A)` currently returns `{'x': 7}`
where a bare `A` returns `A(x=7)`, with `serialize_output` ignored on that
path, and it leaves `json.loads` serving only the genuinely un-buildable `[A,
B]` union. Worth a test either way:
`test_falls_back_to_json_when_output_type_has_no_schema` parametrizes only `[A,
B]` and `ToolOutput(A)`, so the output-function case is asserted nowhere.
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py:
##########
@@ -105,11 +104,8 @@ def defer_for_approval(
"require_approval."
)
- if isinstance(output, BaseModel):
- output = output.model_dump_json()
- elif not isinstance(output, str):
- # JSON round-trip: execute_complete validates the string back into
output_type.
- output = TypeAdapter(type(output)).dump_json(output).decode()
+ # JSON round-trip: execute_complete validates the string back into
output_type.
+ output = dump_output_to_json(output)
Review Comment:
Heads-up for the rebase, since this is the only file that conflicts with
`main`: #71046 added `raw_output = output` immediately above this block
(`approval.py:125` on main) and reads it further down --
`isinstance(raw_output, list)` then `param_value = raw_output` -- to render the
multi-select. Resolving the conflict by taking this side wholesale drops it and
quietly breaks that path. The `pydantic` import removal is still correct after
the rebase: `BaseModel` and `TypeAdapter` have no other use in the file.
--
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]
