kaxil commented on code in PR #74473:
URL: https://github.com/apache/airflow/pull/74473#discussion_r4252208477
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py:
##########
@@ -285,8 +284,5 @@ def regenerate_with_feedback(self, *, feedback: str,
message_history: Any) -> tu
@staticmethod
def _to_string(output: Any) -> str:
- if isinstance(output, BaseModel):
- return output.model_dump_json()
- if not isinstance(output, str):
- return str(output)
- return output
+ """Serialize an output for review as JSON, or ``str()`` when pydantic
cannot."""
+ return dump_output_to_json(output)
Review Comment:
`AgentOperator` never parses this string back (it returns
`_output_under_review` since #74537), but the review text now goes through the
round-trip dump, which drops computed fields. A model with `amount: int` and a
computed `risk` shows the reviewer `{"amount":150}` where main showed
`{"amount":150,"risk":"DENY"}`, and the downstream task still gets `risk` from
`model_dump()`. Could `_to_string` keep a display dump (`model_dump_json()`
plus the new `str()` fallback) and leave `round_trip=True` to
`defer_for_approval`? The `_Computed` cases in
`test_an_unedited_review_string_validates_back` would go too, since nothing on
this path validates back.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -18,37 +18,110 @@
from __future__ import annotations
-from typing import Any
+import inspect
+import logging
+from typing import Any, get_origin
from pydantic import BaseModel, TypeAdapter, ValidationError
+from pydantic.errors import PydanticUserError
+from pydantic_core import PydanticSerializationError
+
+from airflow.providers.common.ai.exceptions import
ReviewedOutputValidationError
+
+log = logging.getLogger(__name__)
+
+_MAX_SHOWN_CHARS = 200
+
+
+def dump_output_to_json(output: Any) -> str:
+ """Serialize an output for review as JSON that validates back into its
type, or ``str()`` if pydantic cannot."""
+ if isinstance(output, str):
+ return output
+ # ``round_trip`` drops computed fields, and ``by_alias=False`` keeps field
names even when a model
+ # sets ``serialize_by_alias``, so a field name or validation alias reads
back.
+ try:
+ if isinstance(output, BaseModel):
+ return output.model_dump_json(round_trip=True, by_alias=False)
+ return TypeAdapter(type(output)).dump_json(output, round_trip=True,
by_alias=False).decode()
+ except (PydanticUserError, PydanticSerializationError):
+ log.warning("Cannot serialize a %s to JSON, so the reviewer sees its
str().", type(output).__name__)
+ return str(output)
def rehydrate_pydantic_output(
output_type: Any,
raw: str,
*,
serialize_output: bool,
+ is_edited: bool = True,
) -> Any:
"""
- Turn a JSON string back into a value of ``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.
-
- 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.
+ Turn the reviewed string back into ``output_type``.
+
+ ``raw`` is JSON, or bare text for types pydantic builds from a string.
``str`` outputs, output
+ functions and types pydantic has no schema for return ``raw`` unchanged,
and so does an unedited
+ output that does not validate back (with a warning).
+
+ :param is_edited: ``False`` when ``raw`` is the unmodified generated
output.
+
+ :raises ReviewedOutputValidationError: If a reviewer's edit does not
validate against ``output_type``.
"""
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)
+ is_output_function = _is_output_function(adapter.core_schema)
+ except PydanticUserError:
return raw
+ if is_output_function:
+ return raw
+ # Output is dumped by field name, so accept a field name or a validation
alias.
+ try:
+ rehydrated = adapter.validate_json(raw, by_alias=True, by_name=True)
+ except (ValidationError, ValueError, TypeError) as json_error:
+ # Bare text is how a str-valued Literal, Enum or ``str | None`` output
is carried through review.
+ try:
+ rehydrated = _validate_bare_text(adapter, raw)
+ except (ValidationError, ValueError, TypeError) as text_error:
+ message = _build_failure_message(output_type, raw, json_error,
text_error)
+ if not is_edited:
+ log.warning("%s Returning the output unchanged.", message)
+ return raw
+ raise ReviewedOutputValidationError(message) from json_error
if serialize_output and isinstance(rehydrated, BaseModel):
return rehydrated.model_dump()
return rehydrated
+
+
+def _validate_bare_text(adapter: TypeAdapter[Any], raw: str) -> Any:
+ """Validate ``raw`` as is, then without the surrounding whitespace a
reviewer's textarea tends to add."""
+ try:
+ return adapter.validate_python(raw)
+ except ValidationError:
+ if raw.strip() == raw:
+ raise
+ return adapter.validate_python(raw.strip())
+
+
+def _build_failure_message(output_type: Any, raw: str, json_error: Exception,
text_error: Exception) -> str:
+ """Describe why ``raw`` does not match ``output_type``, reporting the text
error when it is not JSON."""
+ is_not_json = isinstance(json_error, ValidationError) and all(
+ error["type"] == "json_invalid" for error in json_error.errors()
+ )
+ type_name = output_type if get_origin(output_type) else
getattr(output_type, "__name__", output_type)
+ shown = raw if len(raw) <= _MAX_SHOWN_CHARS else
f"{raw[:_MAX_SHOWN_CHARS]}..."
+ return (
+ f"The reviewed output could not be converted to the output_type
{type_name}. "
+ f"Received {shown!r}. {text_error if is_not_json else json_error}"
Review Comment:
`is_not_json` is also true for a typo in a JSON edit, and then the text
error is a top-level type mismatch. A model edit with a trailing comma reports
`Input should be a valid dictionary or instance of R [type=model_type]`, and
`["a","b",]` for `list[str]` reports `list_type`, where the JSON error says
`Invalid JSON: trailing comma at line 1 column 10`. The removed
`LLMSchemaCompareOperator` override showed the JSON error for its edit-the-JSON
flow. Keeping `json_error` when the text error is a container mismatch at `loc
== ()` (`model_type`, `list_type`, `dict_type`, `tuple_type`, `set_type`)
covers both, and `Literal`, `Enum` and `int` still get the text message.
--
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]
