pankajkoti commented on code in PR #74473:
URL: https://github.com/apache/airflow/pull/74473#discussion_r4227751200
##########
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:
Good catch, thanks. Fixed in 942a11469a. `_to_string` and
`defer_for_approval` now pass `by_alias=True` to both `model_dump_json` and
`dump_json`, so an unedited approval of an aliased model validates back. I
added tests for an aliased model and a list of them through both serializers.
##########
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:
Agreed, this was a regression. Fixed in 942a11469a. When the adapter's
schema type is `call`, the text is returned unchanged and the function is never
invoked. The test asserts it is not called. I kept this narrow on purpose and
left proper handling of output functions and markers to #70132.
##########
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:
Agreed. Fixed in 942a11469a. When the JSON attempt failed only because the
text is not JSON, the message now reports the plain text error, so
`Literal["a", "b"]` with "c" says "Input should be 'a' or 'b'". Otherwise it
still reports the JSON error, which is the useful one for a model missing a
field. The received text is also cut at 200 characters. Both have tests.
##########
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:
Thanks, you are right. Added
`test_regenerate_with_feedback_serializes_list_output_as_json` in 942a11469a,
which asserts `'["a","b"]'`. I checked that reverting the line to `str(output)`
makes it fail.
##########
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:
Good point. The test now also asserts `type(result) is type(expected)`, so
it fails when the raw string comes back unconverted. Fixed in 942a11469a.
##########
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:
Agreed, the review loop has no edit path. I renamed it to
`test_execute_with_hitl_fails_when_review_output_does_not_validate` in
942a11469a. It now returns a payload that is missing a required field and
asserts the "Field required" error. I did not use an aliased model or an output
function here, because after the fixes above both of those now pass an unedited
approval. What is left is a serialized output that really does not validate
back.
--
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]
