kaxil commented on code in PR #74473:
URL: https://github.com/apache/airflow/pull/74473#discussion_r4225097899
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -30,25 +33,43 @@ 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 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:
This passthrough doesn't catch output functions: `TypeAdapter(fn)` builds a
`call` schema instead of raising, so both attempts validate the text as the
function's arguments. With `def clean(text: str) -> str` as `output_type`, an
unedited approval of `"A"` now raises (`call[clean] Invalid JSON`), where main
returned the text. The description leaves output functions to #70132, but this
change takes them from working to failing on every approval. Could a `call`
schema (`adapter.core_schema["type"] == "call"`) return `raw` here until #70132
lands?
##########
providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py:
##########
@@ -268,8 +270,13 @@ def regenerate_with_feedback(self, *, feedback: str,
message_history: Any) -> tu
@staticmethod
def _to_string(output: Any) -> str:
+ """Serialize an output for review as JSON, or ``str()`` when pydantic
cannot."""
if isinstance(output, BaseModel):
return output.model_dump_json()
Review Comment:
An unedited approval now fails for any model with field aliases.
`model_dump_json()` writes field names (`{"user_name": "kax"}`) but
`validate_json` expects the alias (`userName`), so `rehydrate_pydantic_output`
raises `ReviewedOutputValidationError` on a plain Approve. Before this PR that
case passed a string downstream; now the task fails, a retry re-runs the model
and fails the same way, and `on_approval_timeout="approve"` fails with nobody
involved.
The serializer `LLMOperator` uses is the one in `approval.py` (L198-201),
which has the same gap. Passing `by_alias=True` to both `model_dump_json` and
`dump_json` in both places fixes it (I checked a single aliased model and a
list of them), and a test with an aliased model and an unedited approval would
pin it.
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -1037,6 +1038,31 @@ def
test_execute_with_enable_hitl_review_delegates_to_run_hitl_review(
assert result == "Approved output"
mock_run_hitl.assert_called_once_with(op, context, "Initial output",
message_history=msg_history)
+ @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_fails_on_edit_invalid_for_output_type(
Review Comment:
`AgentOperator`'s review loop has no edit path (`HumanActionData` only
carries action, feedback and iteration), so `run_hitl_review` can't return
reviewer prose like this. On this operator the error only fires when the
serialized output doesn't validate back, such as an aliased model or an output
function. Could the test be named for that and use one of those cases?
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -1428,7 +1428,4 @@ def regenerate_with_feedback(self, *, feedback: str,
message_history: Any) -> tu
# shared, is the seeded cumulative object, not what this call alone
contributed.
log_run_summary(self.log, result, usage=regen_usage)
- output = result.output
- if isinstance(output, BaseModel):
- output = output.model_dump_json()
- return str(output), result.all_messages()
+ return self._to_string(result.output), result.all_messages()
Review Comment:
Nothing tests this line. The existing `regenerate_with_feedback` tests only
use `str` and `BaseModel` outputs, which the old body handled too, so reverting
it to `str(output)` leaves the suite green. A `list[str]` case asserting
`'["a","b"]'` would cover the regenerated-output half of the fix.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py:
##########
@@ -30,25 +33,43 @@ 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 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
+ try:
+ rehydrated = adapter.validate_json(raw)
+ except (ValidationError, ValueError, TypeError) as exc:
+ # Bare text is how a str-valued Literal, Enum or ``str | None`` output
is carried through review.
+ try:
+ rehydrated = adapter.validate_python(raw)
+ except (ValidationError, ValueError, TypeError):
+ type_name = (
+ output_type if get_origin(output_type) else
getattr(output_type, "__name__", output_type)
+ )
+ raise ReviewedOutputValidationError(
+ f"The reviewed output could not be converted to the
output_type {type_name}. "
+ f"Received {raw!r}. {exc}"
Review Comment:
`exc` is the `validate_json` error, so for the bare-text types the docs now
describe, the message points the reviewer at JSON. `Literal["a", "b"]` with the
edit `"c"` reports `Invalid JSON: expected value at line 1 column 1` and drops
the python-mode error, which says `Input should be 'a' or 'b'`. Reporting the
second error when the first one is only `json_invalid` would be more useful.
`{raw!r}` also puts the whole LLM output into the task error, so truncating it
would keep a long output from flooding the log.
##########
providers/common/ai/tests/unit/common/ai/utils/test_output_type.py:
##########
@@ -48,15 +56,31 @@ def test_returns_raw_for_str_output_type(self):
def test_validates_other_types_with_type_adapter(self, output_type, raw,
expected):
assert rehydrate_pydantic_output(output_type, raw,
serialize_output=False) == expected
- def test_returns_raw_when_type_adapter_rejects(self):
- result = rehydrate_pydantic_output(int, "not-a-number",
serialize_output=False)
- assert result == "not-a-number"
+ @pytest.mark.parametrize(
+ ("output_type", "raw", "expected"),
+ [
+ (Literal["a", "b"], "a", "a"),
+ (str | None, "hello", "hello"),
+ (Color, "red", Color.RED),
Review Comment:
`Color` subclasses `str`, so `Color.RED == "red"` and this case passes even
when the raw string comes back unconverted (it passes on main). Asserting
`type(result) is Color` would make it check the conversion.
--
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]
