kaxil commented on code in PR #74473:
URL: https://github.com/apache/airflow/pull/74473#discussion_r4243917695
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py:
##########
@@ -195,10 +195,10 @@ def defer_for_approval(
raw_output = output
if isinstance(output, BaseModel):
- output = output.model_dump_json()
+ output = output.model_dump_json(by_alias=True)
Review Comment:
`by_alias=True` fixes `Field(alias=...)` but breaks models whose alias is
serialization-only. `Field(serialization_alias="userName")` or
`AliasGenerator(serialization_alias=to_camel)` now dumps `{"userName": ...}`,
which validation can't read back, so an unedited Approve raises
`ReviewedOutputValidationError ... Field required` where main passed (same for
a `list` of them, and for `_to_string` on the AgentOperator side).
Dumping by field name here and in `_to_string`, and validating with
`validate_json(raw, by_alias=True, by_name=True)` (same flags on
`validate_python`) in `rehydrate_pydantic_output`, round-trips both kinds and
accepts an edit written with either spelling. `by_name` needs pydantic 2.11,
under the 2.12 that pydantic-ai-slim 2.33 already requires. A
dump-then-rehydrate test over both alias styles would catch this; the
dump-string tests can't.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -30,25 +35,50 @@ def rehydrate_pydantic_output(
serialize_output: bool,
) -> Any:
"""
- Turn a JSON string back into a value of ``output_type``.
+ Turn the reviewed string (JSON, or bare text for str-valued types) back
into ``output_type``.
Used by the HITL/approval paths in ``LLMOperator`` and ``AgentOperator``
that round-trip the output through a string when deferring to a human
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.
+ pydantic ``TypeAdapter``, first as JSON and then as bare text.
+ An ``output_type`` pydantic has no schema for, or an output function,
returns ``raw`` unchanged.
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.
+
+ :raises ReviewedOutputValidationError: If ``output_type`` is not ``str``
and ``raw``
+ (typically a reviewer's edit) is neither valid JSON for it nor valid
as plain text.
"""
if output_type is str:
return raw
try:
- rehydrated = TypeAdapter(output_type).validate_json(raw)
- except (ValidationError, ValueError, TypeError):
+ adapter: TypeAdapter[Any] = TypeAdapter(output_type)
+ except PydanticSchemaGenerationError:
+ # Nothing to validate against, so the reviewed text is all there is.
return raw
+ if adapter.core_schema["type"] == "call":
Review Comment:
This only sees a bare `call` schema. When the function's parameters reuse a
model, pydantic wraps it in `definitions`: with `def route(order: Order) ->
str` where `Order` has `billing: Address` and `shipping: Address`, the
top-level type is `definitions`, so an unedited approval of `"Paris"` raises
`ReviewedOutputValidationError ... call[route]` where main returned the text.
Unwrapping one level (`schema["schema"]` when the type is `definitions`) covers
it, and `test_returns_raw_for_output_function_without_calling_it` could take
that function as a second case.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -30,25 +35,50 @@ def rehydrate_pydantic_output(
serialize_output: bool,
) -> Any:
"""
- Turn a JSON string back into a value of ``output_type``.
+ Turn the reviewed string (JSON, or bare text for str-valued types) back
into ``output_type``.
Used by the HITL/approval paths in ``LLMOperator`` and ``AgentOperator``
that round-trip the output through a string when deferring to a human
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.
+ pydantic ``TypeAdapter``, first as JSON and then as bare text.
+ An ``output_type`` pydantic has no schema for, or an output function,
returns ``raw`` unchanged.
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.
+
+ :raises ReviewedOutputValidationError: If ``output_type`` is not ``str``
and ``raw``
+ (typically a reviewer's edit) is neither valid JSON for it nor valid
as plain text.
"""
if output_type is str:
return raw
try:
- rehydrated = TypeAdapter(output_type).validate_json(raw)
- except (ValidationError, ValueError, TypeError):
+ adapter: TypeAdapter[Any] = TypeAdapter(output_type)
+ except PydanticSchemaGenerationError:
Review Comment:
No test passes a type pydantic can't build a schema for, so neither this
branch nor the `str(output)` fallback in `_to_string` is covered. A plain class
as `output_type` here, and an instance of it in `TestToString`, would pin both.
--
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]
