This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new dd7e8ab8321 Fail the task when a reviewed output no longer matches
`output_type` (#74473)
dd7e8ab8321 is described below
commit dd7e8ab83212483c087ef3bd039ae6b97388a99c
Author: Pankaj Koti <[email protected]>
AuthorDate: Wed Oct 14 21:51:15 2026 +0530
Fail the task when a reviewed output no longer matches `output_type`
(#74473)
* Fail the task when a reviewed output no longer matches output_type
When a reviewer edits the output in the approval form, the edited text is
validated back into the operator's output_type. On a validation failure
rehydrate_pydantic_output silently returned the raw string, so the task
succeeded and downstream tasks received a str instead of an int, a list or
a model. It now raises ReviewedOutputValidationError, a ValueError subclass
that names the output type, the text received and the pydantic message.
output_type=str still passes any text through. Types whose value is bare
text, such as a str Literal, an Enum or str | None, are still accepted
as plain text. An output_type that pydantic has no schema for keeps passing
the text through, since there is nothing to validate against.
AgentOperator review serialized non-model outputs with str(), which gives a
Python repr for a list and would now fail an unedited approval. It uses
JSON instead, for the first output and for regenerated ones. Values that
pydantic cannot serialize keep str().
LLMSchemaCompareOperator re-validated the string after the parent returned
and raised its own ValueError. That fallback can no longer be reached, so
the override is removed.
Co-Authored-By: Claude <[email protected]>
* Keep review serialization round-tripping for aliased models and output
functions
An unedited approval of a model with field aliases failed because the
review string was dumped with field names while validation expects the
aliases. The review serializers now dump by alias. An output function used
as output_type was validated as if the text were its arguments, so an
unedited approval raised and the function could be called a second time.
It now returns the text unchanged until it is handled properly.
The validation error now reports the plain text failure when the JSON
attempt only failed because the text is not JSON, and it shortens a long
received text. Tests now pin the regenerated output serialization and the
enum conversion, and the agent review test is renamed to match the only
way that operator can reach the error.
Co-Authored-By: Claude <[email protected]>
* Dump reviewed output by field name and tolerate unedited approvals
The review serializers dumped by alias, which breaks models whose alias is
serialization-only, so an unedited approval failed validation. Both
serializers now share dump_output_to_json, which dumps by field name with
round_trip. Validation accepts a field name or an alias, bare text gets one
retry with surrounding whitespace removed, and an output function is
recognised even when pydantic wraps it in a definitions schema.
An unedited approval of an output that cannot be validated back now returns
the text with a warning, and only an edited output fails the task. Types
pydantic cannot build a schema for keep passing the text through.
Tests cover the no-schema and unresolved forward reference types, models
serialized by alias, computed fields, and the unedited tolerance.
Co-Authored-By: Claude <[email protected]>
* Show reviewers a display dump and keep the JSON error for JSON typos
AgentOperator never validates its review string back, so the round-trip
dump only hid computed fields from the reviewer. The review text now comes
from dump_output_for_display, and dump_output_to_json with round_trip stays
for the approval gate where an edit is validated back. Both share one
helper that keeps the str() fallback and its warning.
A JSON edit with a typo such as a trailing comma reported a top-level type
mismatch from the plain text attempt. The message now keeps the JSON error
when the text error is only a container mismatch, and bare-text types such
as Literal, Enum and int still report the text error.
Co-Authored-By: Claude <[email protected]>
* Keep an approval from failing on form-parsed objects and strict
before-validators
The Airflow HITL form parses a JSON-object value into an object before
sending it back, so
approving a model output in the UI handed the mixin a dict and failed the
task even when
nothing was edited. Dump a dict to compact JSON, as a list already is, so
an untouched
object compares equal to the generated output.
The bare-text retry feeds a string to validators that may only expect the
type's input, so
a mode="before" validator calling data.get raised AttributeError past the
except tuple, on
an unedited approval too. Any failure on that retry now just means "not
bare text", and the
JSON error is reported.
A JSON typo in a union, dataclass, frozenset or NamedTuple edit now reports
the JSON error,
pydantic-ai output markers pass the text through instead of raising
AttributeError, and the
docs and changelog say how a failed edit behaves under retries.
* Treat an untouched object from the review form as unedited
The HITL form round-trips an object output through JavaScript, which rounds
integers
above 2**53, writes 1.0 as 1 and may reorder keys (as may the metadata
database). The
re-dumped text then differed from the generated output, so an untouched
approval read as
an edit: a large id could change silently, and a model whose dump does not
validate back
failed the task. Compare values the way the browser parses them, with a
bool kept distinct
from a number, and keep the generated output unchanged when they match.
A before-validator that raises AttributeError, KeyError or IndexError on a
JSON edit that is
not an object now fails with ReviewedOutputValidationError on the JSON
attempt too, while any
other validator error (a timeout, say) still propagates instead of reading
as a failed
conversion. An object or array the reviewer
typed reports its JSON error even when the type also takes bare text, and
the docs no
longer claim a retry opens a new review on every Airflow version.
* Report the JSON error for plain text sent to a NamedTuple on pydantic 2.14
Pydantic 2.14 raises named_tuple_type instead of arguments_type when a whole
string meets a NamedTuple. The error was missing from the container mismatch
set, so the review message showed the plain text error where it should show
the JSON error. The set now lists both, and the test notes the split.
Co-Authored-By: Claude <[email protected]>
* Show JSON in the decision review body and match container errors by name
pattern
A decision-policy review body showed the output's Python repr, so the docs'
advice to copy
"the JSON shown in the review body" into the edit field did not hold there.
It now shows the
same JSON the default body and the edit field use.
Telling a JSON typo from plain text relied on a list of pydantic error
names, which pydantic
2.14 renamed for a NamedTuple and a deque. Match the "_type" suffix those
errors share, plus
"sequence_str", instead.
Rehydrating no longer forces validation by alias, which overrode a model
that sets
validate_by_alias=False: with an alias named like another field, an
untouched approval
swapped the two values. Field names and, by default, aliases are still
accepted.
A multi-select edit is dumped without escaping non-ASCII, like the
generated output, so an
untouched choice no longer reads as an edit. The changelog no longer says a
Literal or str
Enum keeps accepting free text, and the docs say a union with a str member
keeps an edit
that is not the model as text.
---------
Co-authored-by: Claude <[email protected]>
Co-authored-by: Kaxil Naik <[email protected]>
---
providers/common/ai/docs/approval_gates.rst | 23 ++
providers/common/ai/docs/changelog.rst | 10 +
providers/common/ai/docs/operators/llm.rst | 4 +-
.../ai/docs/operators/llm_schema_compare.rst | 3 +-
.../providers/common/ai/batch/output_schema.py | 9 +-
.../src/airflow/providers/common/ai/exceptions.py | 9 +
.../airflow/providers/common/ai/mixins/approval.py | 49 ++-
.../providers/common/ai/mixins/hitl_review.py | 14 +-
.../airflow/providers/common/ai/operators/agent.py | 7 +-
.../airflow/providers/common/ai/operators/llm.py | 27 +-
.../common/ai/operators/llm_schema_compare.py | 18 +-
.../providers/common/ai/utils/output_type.py | 142 ++++++-
.../unit/common/ai/batch/test_output_schema.py | 13 +-
.../tests/unit/common/ai/mixins/test_approval.py | 133 ++++++-
.../unit/common/ai/mixins/test_hitl_review.py | 58 +++
.../tests/unit/common/ai/operators/test_agent.py | 15 +
.../ai/tests/unit/common/ai/operators/test_llm.py | 49 +++
.../common/ai/operators/test_llm_schema_compare.py | 5 +-
.../tests/unit/common/ai/utils/test_output_type.py | 422 ++++++++++++++++++++-
19 files changed, 913 insertions(+), 97 deletions(-)
diff --git a/providers/common/ai/docs/approval_gates.rst
b/providers/common/ai/docs/approval_gates.rst
index a42ccf2786a..e45d2735d51 100644
--- a/providers/common/ai/docs/approval_gates.rst
+++ b/providers/common/ai/docs/approval_gates.rst
@@ -31,6 +31,29 @@ HITL interface. Optionally allow the reviewer to edit the
output before
approving with ``allow_modifications=True``, and set a deadline with
``approval_timeout``.
+On ``LLMOperator``, ``LLMFileAnalysisOperator`` and
``LLMSchemaCompareOperator``, the edited text is
+converted back into ``output_type`` when the task resumes. Write it as JSON
for the
+type, such as ``31`` for an ``int``, ``["a", "b"]`` for a ``list[str]`` or an
object for
+a Pydantic model. If the edit field shows ``[object Object]`` for an object
output,
+replace its content with the JSON shown in the review body. Types that
pydantic builds
+from a string, such as a ``Literal``, an ``Enum`` or ``str | None``, also
accept the bare text. If the text does not validate,
+the task fails with
+:class:`~airflow.providers.common.ai.exceptions.ReviewedOutputValidationError`
+naming the type and the (truncated) text received, rather than passing a value
of the wrong
+type downstream. Airflow retries it like any other failure, so with
``retries`` set the
+next try runs the model again and the rejected edit is not carried over.
+
+With ``output_type=str`` any edit is accepted as is, and so is an edit to a
union with a ``str``
+member, such as ``Model | str``, which keeps the text when it does not
validate as the model.
+An output function, a pydantic-ai
+output marker such as ``ToolOutput``, or a type pydantic has no schema for, is
returned as
+the text unchanged. An approval without an
+edit whose output cannot be converted back (for example a custom serializer)
also returns
+the text with a warning, instead of failing the task. Computed fields are
left out of the
+text. A model with a serializer that cannot be reversed, or a required field
excluded from
+the dump, shows text that does not validate as shown, so write an edit in the
form the model
+accepts.
+
Human-in-the-loop review needs Airflow 3.1+, whether ``require_approval`` or a
``decision_policy`` with ``on_uncertain="review"`` opens it. On an older
Airflow version the
operator raises ``AirflowOptionalProviderFeatureException`` when it is
constructed, so the Dag file
diff --git a/providers/common/ai/docs/changelog.rst
b/providers/common/ai/docs/changelog.rst
index cc99d7d3996..4314b8fe4a7 100644
--- a/providers/common/ai/docs/changelog.rst
+++ b/providers/common/ai/docs/changelog.rst
@@ -49,6 +49,16 @@ Changelog
headroom. On Airflow < 3.3, or when ``usage_limits`` is ``None``, each
attempt is checked
and counted on its own, unchanged. See :ref:`howto/operator:llm`.
+.. note::
+ A reviewer's edit on ``LLMOperator``, ``LLMFileAnalysisOperator`` or
+ ``LLMSchemaCompareOperator`` (``allow_modifications=True``) is now converted
back into
+ ``output_type``, and an edit that does not validate fails the task with
+ ``ReviewedOutputValidationError`` (a ``ValueError``) instead of returning
the edited text as
+ a ``str``. Use ``output_type=str`` or ``str | None`` to keep accepting free
text; a ``Literal`` or
+ a ``str`` ``Enum`` accepts one of its values typed as bare text, and any
other edit now fails the
+ task. An approval without an edit never fails this way. The review text now
shows field names and leaves out computed fields.
+ See :doc:`approval_gates`.
+
.. note::
Durable execution (``durable=True`` on ``AgentOperator``) is rebuilt on the
``AirflowDurability`` capability. The internal ``durable.caching_model``,
diff --git a/providers/common/ai/docs/operators/llm.rst
b/providers/common/ai/docs/operators/llm.rst
index b789fa0472a..38c4bceacaa 100644
--- a/providers/common/ai/docs/operators/llm.rst
+++ b/providers/common/ai/docs/operators/llm.rst
@@ -297,7 +297,9 @@ Parameters
review path (``require_approval=True`` or a ``decision_policy`` that reviews)
and a positive ``approval_timeout``.
- ``allow_modifications``: If ``True``, the reviewer can edit the output before
- approving. Default ``False``.
+ approving. For a type pydantic can validate, an edit that does not validate
+ fails the task with ``ReviewedOutputValidationError``; a retry runs the model
+ again. Default ``False``.
- ``approval_notifiers``: Notifier, or list of notifiers, called once the
review
is open. Default ``None``.
- ``approval_assigned_users``: Users allowed to answer the review, as
diff --git a/providers/common/ai/docs/operators/llm_schema_compare.rst
b/providers/common/ai/docs/operators/llm_schema_compare.rst
index aaaad00d60a..fa7d9ed8eb7 100644
--- a/providers/common/ai/docs/operators/llm_schema_compare.rst
+++ b/providers/common/ai/docs/operators/llm_schema_compare.rst
@@ -212,7 +212,8 @@ Parameters
review: ``"fail"`` (default), ``"approve"``, or ``"reject"``. Requires
``require_approval=True`` and a positive ``approval_timeout``.
- ``allow_modifications``: If ``True``, the reviewer can edit the result JSON
- before approving. Default ``False``.
+ before approving. An edit that is not a valid result fails the task with
+ ``ReviewedOutputValidationError``. Default ``False``.
- ``approval_notifiers``: Notifier, or list of notifiers, called once the
review
is open. Default ``None``.
- ``approval_assigned_users``: Users allowed to answer the review. ``None``
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/batch/output_schema.py
b/providers/common/ai/src/airflow/providers/common/ai/batch/output_schema.py
index 9d328582117..bdcfb45dd1e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/batch/output_schema.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/batch/output_schema.py
@@ -25,11 +25,10 @@ response back into a validated value or an explicit failure
Deliberately does **not** call
:func:`~airflow.providers.common.ai.utils.output_type.rehydrate_pydantic_output`:
-that helper's failure mode is "validation failed -> return the raw string
-unchanged", which is correct for the HITL round-trip it was built for but
-wrong here -- a batch item that fails validation must be recorded as failed
-(``status: "invalid_output"``), never silently downgraded into a
-string that merely looks like a success.
+that helper raises when a reviewer's edit fails validation, which fails the
whole
+task. That is right for the HITL round-trip it was built for but wrong here --
a
+batch item that fails validation must be recorded as failed
+(``status: "invalid_output"``) while the rest of the batch carries on.
"""
from __future__ import annotations
diff --git a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
index 701917722ae..1be907e6998 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
@@ -90,6 +90,15 @@ class
LLMFileAnalysisMultimodalRequiredError(LLMFileAnalysisUnsupportedFormatErr
"""Raised when image/PDF inputs are used without ``multi_modal=True``."""
+class ReviewedOutputValidationError(ValueError):
+ """
+ Raised when a reviewer-approved output cannot be converted back into the
operator's ``output_type``.
+
+ Typically the reviewer edited the output into something the type rejects,
for example free text
+ for an ``int``. The task fails instead of passing a value of the wrong
type to downstream tasks.
+ """
+
+
class LowConfidenceError(ValueError):
"""
Raised by an LLM operator whose ``DecisionPolicy`` says
``on_uncertain="fail"``.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
index dc33aef7b73..a5854a855b9 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/mixins/approval.py
@@ -24,8 +24,8 @@ from datetime import timedelta
from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol
from jinja2 import TemplateError
-from pydantic import BaseModel, TypeAdapter
+from airflow.providers.common.ai.utils.output_type import dump_output_to_json
from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_3_PLUS
if AIRFLOW_V_3_3_PLUS:
@@ -43,6 +43,40 @@ if TYPE_CHECKING:
from airflow.sdk.execution_time.hitl import HITLUser
+def _object_from_the_form_as_text(modified: dict[str, Any], generated_output:
str) -> str:
+ """
+ Turn the object the HITL form sent back into review text, keeping
``generated_output`` if untouched.
+
+ The form round-trips the value through JavaScript, which rounds integers
above 2**53, writes ``1.0``
+ as ``1`` and may reorder keys (so may the metadata database), so compare
values the same way rather
+ than text; a re-dumped untouched object would otherwise read as an edit
and could change a value.
+ """
+ try:
+ generated = json.loads(generated_output)
+ except ValueError:
+ generated = None
+ if _with_javascript_numbers(modified) ==
_with_javascript_numbers(generated):
+ return generated_output
+ return json.dumps(modified, separators=(",", ":"), ensure_ascii=False)
+
+
+def _with_javascript_numbers(value: Any) -> Any:
+ """Map every number to a float, as ``JSON.parse`` does, so values compare
as the browser sees them."""
+ if isinstance(value, bool):
+ # Tagged, because Python treats ``False == 0.0`` and ``True == 1.0``
as equal and JSON does not.
+ return ("bool", value)
+ if isinstance(value, int | float):
+ try:
+ return float(value)
+ except OverflowError:
+ return value
+ if isinstance(value, dict):
+ return {key: _with_javascript_numbers(item) for key, item in
value.items()}
+ if isinstance(value, list):
+ return [_with_javascript_numbers(item) for item in value]
+ return value
+
+
def normalize_assigned_users(value: Any, *, param: str) -> list[HITLUser]:
"""
Return *value* as a list of ``{'id': str, 'name': str}`` users.
@@ -194,11 +228,8 @@ class LLMApprovalMixin:
self.validate_approval_prompt()
raw_output = output
- 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)
ti_id = context["task_instance"].id
timeout_defaults =
LLMApprovalMixin.TIMEOUT_DEFAULTS.get(self.on_approval_timeout)
@@ -357,7 +388,11 @@ class LLMApprovalMixin:
}
)
# Compact so an unchanged selection compares equal to
generated_output
- modified = json.dumps(modified, separators=(",", ":"))
+ modified = json.dumps(modified, separators=(",", ":"),
ensure_ascii=False)
+ elif isinstance(modified, dict):
+ # TODO: Drop once the HITL form stops parsing a string param
that looks like JSON into an
+ # object (airflow-core ui/src/utils/hitl.ts); until then a
model output comes back as a dict.
+ modified = _object_from_the_form_as_text(modified,
generated_output)
if modified is not None and not isinstance(modified, str):
# On the awaiting_input path nothing upstream schema-validates
params_input
# (HITLTrigger did on the legacy path), so enforce the string
contract here
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py
b/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py
index 8a1327a871e..3a924fbc78f 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py
@@ -21,8 +21,6 @@ import time
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Protocol
-from pydantic import BaseModel
-
from airflow.providers.common.ai.exceptions import HITLMaxIterationsError
from airflow.providers.common.ai.utils.hitl_review import (
XCOM_AGENT_OUTPUT_PREFIX,
@@ -32,6 +30,7 @@ from airflow.providers.common.ai.utils.hitl_review import (
HumanActionData,
SessionStatus,
)
+from airflow.providers.common.ai.utils.output_type import
dump_output_for_display
log = logging.getLogger(__name__)
@@ -104,11 +103,11 @@ class HITLReviewMixin:
Execute the full HITL review loop.
:param context: Airflow task context.
- :param output: Initial LLM output (str or BaseModel).
+ :param output: Initial LLM output (any pydantic-serializable value).
:param message_history: Provider-specific conversation state (e.g.
pydantic-ai ``list[ModelMessage]``). Passed to
:meth:`regenerate_with_feedback` on each iteration.
- :returns: The final approved output as a string.
+ :returns: The approved output as a string (JSON for non-str values
pydantic can serialize).
:raises HITLMaxIterationsError: When max iterations reached without
approval.
:raises HITLRejectException: When the reviewer rejects the output.
:raises HITLTimeoutError: When hitl_timeout elapses with no response.
@@ -285,8 +284,5 @@ class HITLReviewMixin:
@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 the reviewer to read as JSON, or ``str()``
when pydantic cannot."""
+ return dump_output_for_display(output)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
index 8bc303e6494..6aaa78e7812 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
@@ -1410,8 +1410,5 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin, Usa
# shared, is the seeded cumulative object, not what this call alone
contributed.
log_run_summary(self.log, result, usage=regen_usage)
- output = result.output
- self._output_under_review = output
- if isinstance(output, BaseModel):
- output = output.model_dump_json()
- return str(output), result.all_messages()
+ self._output_under_review = result.output
+ return self._to_string(result.output), result.all_messages()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
index ac894d4903d..1033bc6c414 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
@@ -45,7 +45,7 @@ from airflow.providers.common.ai.utils.decision import (
validate_decision_policy,
)
from airflow.providers.common.ai.utils.logging import MODEL_NAME_XCOM_KEY,
log_run_summary
-from airflow.providers.common.ai.utils.output_type import
rehydrate_pydantic_output
+from airflow.providers.common.ai.utils.output_type import dump_output_to_json,
rehydrate_pydantic_output
from airflow.providers.common.ai.utils.usage import coerce_usage_limits
from airflow.providers.common.compat.notifier import BaseNotifier
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, BaseOperator
@@ -162,7 +162,9 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
``decision_policy`` that reviews) and a positive
``approval_timeout``.
:param allow_modifications: If ``True``, the reviewer can edit the output
- before approving. The modified value is returned as the task result.
+ before approving. The edit is converted back into ``output_type`` and
returned
+ as the task result. For a type pydantic can validate, an edit that
does not validate
+ fails the task with ``ReviewedOutputValidationError``; a retry runs
the model again.
Default ``False``.
:param approval_notifiers: Notifiers called once the review is open, so a
reviewer is told about it. Only takes effect when a review is
@@ -375,7 +377,8 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
self._log_review(review, confidence, threshold)
body = None
if review != "require_approval":
- body = f"```\nPrompt: {self.prompt}\n\n{output}\n```\n\n" +
"\n\n".join(
+ # The same JSON text the default body and the edit field show,
not the object's repr.
+ body = f"```\nPrompt:
{self.prompt}\n\n{dump_output_to_json(output)}\n```\n\n" + "\n\n".join(
describe_confidence(model_confidence, name, threshold) for
name in fields
)
self.defer_for_approval(context, output, body=body,
decision=record) # type: ignore[misc]
@@ -396,7 +399,11 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
event: dict[str, Any],
decision: dict[str, Any] | None = None,
) -> Any:
- """Resume after human review and restore the Pydantic model for XCom
consumers."""
+ """
+ Resume after human review and convert the output back into
``output_type``.
+
+ :raises ReviewedOutputValidationError: If a reviewer's edit does not
validate against ``output_type``.
+ """
output = self._complete_review(context, generated_output, event,
decision)
self._clear_usage_budget(context)
return output
@@ -408,11 +415,19 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
event: dict[str, Any],
decision: dict[str, Any] | None,
) -> Any:
- """Resolve the review and rehydrate the output, leaving the usage
budget for the caller to clear."""
+ """
+ Resolve the review and rehydrate the output, leaving the usage budget
for the caller to clear.
+
+ :raises ReviewedOutputValidationError: If a reviewer's edit does not
validate against ``output_type``.
+ """
output = self._resume_after_review(context, generated_output, event,
decision)
self._finalize_decision(context, event, decision, action=None)
return rehydrate_pydantic_output(
- self.output_type, output,
serialize_output=self._serialize_model_output
+ self.output_type,
+ output,
+ serialize_output=self._serialize_model_output,
+ # An unedited approval hands back ``generated_output`` as is, and
only an edit may fail the task.
+ is_edited=output != generated_output,
)
def _run_llm(
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
index 2e2f8d11502..e48a51de444 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
@@ -24,7 +24,7 @@ from collections.abc import Sequence
from functools import cached_property
from typing import TYPE_CHECKING, Any, ClassVar, Literal
-from pydantic import BaseModel, Field, ValidationError
+from pydantic import BaseModel, Field
from airflow.providers.common.ai.operators.llm import LLMOperator
from airflow.providers.common.ai.utils.usage import coerce_usage_limits
@@ -367,19 +367,3 @@ class LLMSchemaCompareOperator(LLMOperator):
self._clear_usage_budget(context)
return output_result
-
- def execute_complete(
- self,
- context: Context,
- generated_output: str,
- event: dict[str, Any],
- decision: dict[str, Any] | None = None,
- ) -> Any:
- output = self._complete_review(context, generated_output, event,
decision)
- if not isinstance(output, dict):
- try:
- output =
SchemaCompareResult.model_validate_json(output).model_dump()
- except ValidationError as e:
- raise ValueError(f"Reviewed output is not a valid
SchemaCompareResult: {e}") from e
- self._clear_usage_budget(context)
- return output
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
index 44ca8505b54..70de2cd73b4 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/output_type.py
@@ -18,9 +18,48 @@
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 CoreSchema, PydanticSerializationError
+
+from airflow.providers.common.ai.exceptions import
ReviewedOutputValidationError
+
+log = logging.getLogger(__name__)
+
+_MAX_SHOWN_CHARS = 200
+# What a value of the wrong shape raises on its way through ``output_type``.
Pydantic wraps a validator's
+# ``ValueError``, but not the ``AttributeError``/``KeyError``/``IndexError``
of a ``mode="before"`` validator
+# that indexes its input (``data.get``) when the review text is not the shape
it expects. Anything else,
+# such as a validator's network or timeout error, propagates instead of
reading as a failed conversion.
+_CONVERSION_ERRORS = (ValidationError, ValueError, TypeError, AttributeError,
KeyError, IndexError)
+
+
+def dump_output_to_json(output: Any) -> str:
+ """Serialize an output as JSON that ``rehydrate_pydantic_output``
validates back, or ``str()`` if pydantic cannot."""
+ # ``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.
+ return _dump_json(output, round_trip=True, by_alias=False)
+
+
+def dump_output_for_display(output: Any) -> str:
+ """Serialize an output as JSON for a reviewer to read, or ``str()`` if
pydantic cannot."""
+ return _dump_json(output)
+
+
+def _dump_json(output: Any, *, round_trip: bool = False, by_alias: bool | None
= None) -> str:
+ if isinstance(output, str):
+ return output
+ try:
+ if isinstance(output, BaseModel):
+ return output.model_dump_json(round_trip=round_trip,
by_alias=by_alias)
+ return TypeAdapter(type(output)).dump_json(output,
round_trip=round_trip, by_alias=by_alias).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(
@@ -28,27 +67,100 @@ def rehydrate_pydantic_output(
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, pydantic-ai output markers such as ``ToolOutput`` 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)
+ # A pydantic-ai output marker (``ToolOutput``, ``NativeOutput``, ...) is
not a type, and
+ # ``TypeAdapter`` fails on it with an ``AttributeError``.
+ except (PydanticUserError, AttributeError):
+ return raw
+ if is_output_function:
return raw
+ # Output is dumped by field name, so also accept a field name. An alias
stays accepted when the model
+ # validates by alias, its default; forcing it would override a model that
turned it off.
+ try:
+ rehydrated = adapter.validate_json(raw, by_name=True)
+ except _CONVERSION_ERRORS 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 _CONVERSION_ERRORS 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 for bare-text types."""
+ 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]}..."
+ reported = text_error if _is_bare_text_failure(raw, json_error,
text_error) else json_error
+ return (
+ f"The reviewed output could not be converted to the output_type
{type_name}. "
+ f"Received {shown!r}. {reported}"
+ )
+
+
+def _is_bare_text_failure(raw: str, json_error: Exception, text_error:
Exception) -> bool:
+ """Tell a text that was never JSON from a JSON edit with a typo, which the
JSON error explains best."""
+ # An object or array the reviewer typed is JSON with a typo, whatever the
type also accepts as text.
+ if not isinstance(text_error, ValidationError) or
raw.lstrip().startswith(("{", "[")):
+ return False
+ is_not_json = isinstance(json_error, ValidationError) and all(
+ error["type"] == "json_invalid" for error in json_error.errors()
+ )
+ # A union adds its member to ``loc``, so match on the error type alone.
+ is_container_mismatch = all(_is_container_mismatch(error["type"]) for
error in text_error.errors())
+ return is_not_json and not is_container_mismatch
+
+
+def _is_container_mismatch(error_type: str) -> bool:
+ """
+ Tell whether pydantic refused the whole text for a structured type, where
the JSON error says more.
+
+ Those errors end in ``_type`` (``model_type``, ``list_type``,
``named_tuple_type``, ...), while a type
+ that accepts text fails with ``*_parsing``, ``*_error`` or ``enum``.
Matching the suffix rather than a
+ list keeps up with pydantic renaming them, as 2.14 did for a NamedTuple
and a deque.
+ """
+ return error_type.endswith("_type") or error_type == "sequence_str"
+
+
+def _is_output_function(schema: CoreSchema) -> bool:
+ """
+ Tell an output function from a type, both of which can have a ``call``
core schema.
+
+ A NamedTuple, wrapped or not, calls a class. A reused parameter model
wraps the schema in ``definitions``.
+ """
+ if schema["type"] == "definitions":
+ schema = schema["schema"]
+ return schema["type"] == "call" and not inspect.isclass(schema["function"])
diff --git
a/providers/common/ai/tests/unit/common/ai/batch/test_output_schema.py
b/providers/common/ai/tests/unit/common/ai/batch/test_output_schema.py
index 9c2d7bceb59..691fd29cf11 100644
--- a/providers/common/ai/tests/unit/common/ai/batch/test_output_schema.py
+++ b/providers/common/ai/tests/unit/common/ai/batch/test_output_schema.py
@@ -30,7 +30,7 @@ from airflow.providers.common.ai.batch.output_schema import (
build_output_spec,
validate_extracted_output,
)
-from airflow.providers.common.ai.exceptions import LLMBatchOutputTypeError
+from airflow.providers.common.ai.exceptions import LLMBatchOutputTypeError,
ReviewedOutputValidationError
from airflow.providers.common.ai.utils.output_type import
rehydrate_pydantic_output
@@ -208,17 +208,16 @@ class TestNotRehydratePydanticOutput:
A regression guard proving this module's failure semantics are the
opposite of
``rehydrate_pydantic_output``.
- ``rehydrate_pydantic_output`` silently downgrades a validation failure to
- the raw string -- correct for the HITL round trip, wrong for batch, where
- a failure must be recorded as ``invalid_output``, not mistaken for a
- successful string result.
+ ``rehydrate_pydantic_output`` raises on a validation failure, which suits
the HITL
+ round trip, where the task should fail. Batch must instead record the item
as
+ ``invalid_output`` and keep going.
"""
def test_same_bad_input_diverges_between_the_two_functions(self):
bad_json = '{"name": "Ann", "age": "not-a-number"}'
- rehydrated = rehydrate_pydantic_output(Diagnosis, bad_json,
serialize_output=True)
- assert rehydrated == bad_json # silently returns the raw string
unchanged
+ with pytest.raises(ReviewedOutputValidationError):
+ rehydrate_pydantic_output(Diagnosis, bad_json,
serialize_output=True)
spec = build_output_spec(Diagnosis)
outcome = validate_extracted_output(ExtractedOutput(kind="json_text",
text=bad_json), spec)
diff --git a/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
b/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
index 834c00d2610..5b727c28d2e 100644
--- a/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
+++ b/providers/common/ai/tests/unit/common/ai/mixins/test_approval.py
@@ -29,7 +29,7 @@ from unittest.mock import MagicMock, patch
from uuid import uuid4
from jinja2 import TemplateError
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, Field, computed_field
from airflow.providers.common.ai.mixins.approval import (
LLMApprovalMixin,
@@ -103,6 +103,25 @@ def context():
return {"task_instance": ti, "dag": DAG("test_dag")}
+class _SerializedByAlias(BaseModel):
+ model_config = ConfigDict(serialize_by_alias=True)
+ user_name: str = Field(serialization_alias="userName")
+
+
+class _NoSchema:
+ def __repr__(self):
+ return "no-schema"
+
+
+class _Doubled(BaseModel):
+ a: int
+
+ @computed_field # type: ignore[prop-decorator]
+ @property
+ def double(self) -> int:
+ return self.a * 2
+
+
# The legacy trigger path is taken on Airflow versions < 3.3; pin the flag so
these tests keep
# exercising the defer() fallback when run against newer Airflow versions.
@patch(AWAIT_INPUT_FLAG_PATH, False)
@@ -143,12 +162,32 @@ class TestDeferForApproval:
@pytest.mark.parametrize(
("output", "expected"),
- [(42, "42"), (True, "true"), ([1, "a"], '[1,"a"]'), ({"k": "v"},
'{"k":"v"}')],
- ids=["int", "bool", "list", "dict"],
+ [
+ (42, "42"),
+ (True, "true"),
+ ([1, "a"], '[1,"a"]'),
+ ({"k": "v"}, '{"k":"v"}'),
+ (_Doubled(a=1), '{"a":1}'),
+ ([_Doubled(a=1)], '[{"a":1}]'),
+ (_SerializedByAlias(user_name="kax"), '{"user_name":"kax"}'),
+ ([_SerializedByAlias(user_name="kax")], '[{"user_name":"kax"}]'),
+ (_NoSchema(), "no-schema"),
+ ],
+ ids=[
+ "int",
+ "bool",
+ "list",
+ "dict",
+ "model-without-computed-fields",
+ "list-of-models-without-computed-fields",
+ "model-serialized-by-alias",
+ "list-of-models-serialized-by-alias",
+ "no-schema-falls-back-to-str",
+ ],
)
@patch(HITL_TRIGGER_PATH, autospec=True)
@patch(UPSERT_HITL_PATH)
- def test_non_string_non_pydantic_output_is_json_encoded(
+ def test_non_string_output_is_dumped_for_review(
self, mock_upsert, mock_trigger_cls, approval_op, context, output,
expected
):
approval_op.defer_for_approval(context, output)
@@ -425,11 +464,11 @@ class TestDeferForApproval:
def test_approved_with_non_string_modified_output_raises(self,
approval_op_with_modifications):
# On the awaiting_input path nothing upstream schema-validates
params_input
# (HITLTrigger did on the legacy path), so execute_complete must
enforce the
- # string contract instead of returning a dict as the task's output.
+ # string contract instead of returning another type as the task's
output.
event = {
"chosen_options": ["Approve"],
"responded_by_user": {"id": "u1", "name": "editor"},
- "params_input": {"output": {"sneaky": "dict"}},
+ "params_input": {"output": 42},
}
with pytest.raises(HITLTriggerEventError, match="must be a string"):
@@ -437,6 +476,88 @@ class TestDeferForApproval:
{}, generated_output="original output", event=event
)
+ @pytest.mark.parametrize(
+ ("generated_output", "modified"),
+ [
+ ('{"title":"a","priority":0}', {"title": "a", "priority": 0}),
+ ('{"id":9007199254740993}', {"id": 9007199254740992}),
+ ('{"score":2.0}', {"score": 2}),
+ ('{"b":1,"10":2}', {"10": 2, "b": 1}),
+ ('{\n "a": 1\n}', {"a": 1}),
+ ('{"ids":[9007199254740993]}', {"ids": [9007199254740992]}),
+ ],
+ ids=[
+ "same",
+ "int-rounded-by-javascript",
+ "integral-float",
+ "reordered-keys",
+ "pretty-printed",
+ "int-rounded-in-a-list",
+ ],
+ )
+ def
test_approved_with_untouched_object_from_the_form_returns_the_generated_output(
+ self, approval_op_with_modifications, generated_output, modified
+ ):
+ # The HITL form parses a JSON-object value into an object with
JavaScript numbers.
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": {"id": "u1", "name": "editor"},
+ "params_input": {"output": modified},
+ }
+
+ result = approval_op_with_modifications.execute_complete(
+ {}, generated_output=generated_output, event=event
+ )
+
+ assert result == generated_output
+
+ @pytest.mark.parametrize(
+ ("generated_output", "modified", "expected"),
+ [
+ ('{"amount":0}', {"amount": False}, '{"amount":false}'),
+ ('{"flags":[0]}', {"flags": [False]}, '{"flags":[false]}'),
+ ],
+ ids=["field", "in-a-list"],
+ )
+ def test_approved_with_a_bool_in_place_of_a_number_counts_as_an_edit(
+ self, approval_op_with_modifications, generated_output, modified,
expected
+ ):
+ # Python treats ``False == 0``, but the reviewer changed the value's
JSON type.
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": {"id": "u1", "name": "editor"},
+ "params_input": {"output": modified},
+ }
+
+ result = approval_op_with_modifications.execute_complete(
+ {}, generated_output=generated_output, event=event
+ )
+
+ assert result == expected
+
+ @pytest.mark.parametrize(
+ ("modified", "expected"),
+ [
+ ({"title": "b", "priority": 3}, '{"title":"b","priority":3}'),
+ ({"title": "café", "priority": 0},
'{"title":"café","priority":0}'),
+ ],
+ ids=["edited", "non-ascii"],
+ )
+ def test_approved_with_edited_object_from_the_form_is_dumped_as_json(
+ self, approval_op_with_modifications, modified, expected
+ ):
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": {"id": "u1", "name": "editor"},
+ "params_input": {"output": modified},
+ }
+
+ result = approval_op_with_modifications.execute_complete(
+ {}, generated_output='{"title":"a","priority":0}', event=event
+ )
+
+ assert result == expected
+
def test_approved_with_list_modified_output_is_serialized(self,
approval_op_with_modifications):
event = {
"chosen_options": ["Approve"],
diff --git
a/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py
b/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py
index fba67629ce1..98839df5ec2 100644
--- a/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py
+++ b/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py
@@ -16,7 +16,10 @@
# under the License.
from __future__ import annotations
+from dataclasses import dataclass
+
import pytest
+from pydantic import BaseModel, Field, computed_field
from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
@@ -110,6 +113,61 @@ def context(mock_ti):
return {"task_instance": mock_ti}
+class _SerializationAliased(BaseModel):
+ user_name: str = Field(serialization_alias="userName")
+
+
+class _Computed(BaseModel):
+ a: int
+
+ @computed_field # type: ignore[prop-decorator]
+ @property
+ def double(self) -> int:
+ return self.a * 2
+
+
+class _NoSchema:
+ def __repr__(self):
+ return "no-schema"
+
+
+@dataclass
+class _UnresolvedForwardReference:
+ value: Missing # type: ignore[name-defined] # noqa: F821
+
+ def __repr__(self):
+ return "unresolved"
+
+
+class TestToString:
+ @pytest.mark.parametrize(
+ ("output", "expected"),
+ [
+ (["tag-a", "tag-b"], '["tag-a","tag-b"]'),
+ ([_SerializationAliased(user_name="kax")],
'[{"user_name":"kax"}]'),
+ (_Computed(a=1), '{"a":1,"double":2}'),
+ (_NoSchema(), "no-schema"),
+ ([_NoSchema()], "[no-schema]"),
+ (_UnresolvedForwardReference(value=1), "unresolved"),
+ ],
+ ids=[
+ "list",
+ "serialization-aliased-list",
+ "computed-field-shown",
+ "no-schema-falls-back-to-str",
+ "unserializable-falls-back-to-str",
+ "unresolved-schema-falls-back-to-str",
+ ],
+ )
+ def test_to_string_dumps_json_or_falls_back_to_str(self, output, expected):
+ assert HITLReviewMixin._to_string(output) == expected
+
+ def test_warns_when_it_falls_back_to_str(self, caplog):
+ HITLReviewMixin._to_string(_NoSchema())
+
+ assert any("Cannot serialize a _NoSchema" in message for message in
caplog.messages)
+
+
class TestHITLReviewMixin:
@patch("airflow.providers.common.ai.mixins.hitl_review.time.sleep",
autospec=True)
def test_pushes_session_and_output_on_start(
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index 27f5cda1ba2..90e376c8f50 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -1569,6 +1569,21 @@ class TestAgentOperatorRegenerateWithFeedback:
usage=ANY,
)
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_regenerate_with_feedback_serializes_list_output_as_json(
+ self, mock_hook_cls, make_mock_run_result
+ ):
+ mock_result = make_mock_run_result(["a", "b"])
+ mock_result.all_messages.return_value = []
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ mock_agent.run_sync.return_value = mock_result
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = AgentOperator(task_id="test", prompt="test",
llm_conn_id="my_llm", output_type=list[str])
+ output, _ = op.regenerate_with_feedback(feedback="Expand",
message_history=[])
+
+ assert output == '["a","b"]'
+
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_regenerate_with_feedback_serializes_base_model_output(self,
mock_hook_cls, make_mock_run_result):
"""regenerate_with_feedback returns JSON string for BaseModel
output."""
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
index 6d57211e248..d5d92924904 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
@@ -30,6 +30,7 @@ from pydantic_ai.messages import ModelMessage, ModelResponse,
TextPart
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits
+from airflow.providers.common.ai.exceptions import
ReviewedOutputValidationError
from airflow.providers.common.ai.mixins.approval import (
LLMApprovalMixin,
)
@@ -679,6 +680,9 @@ class TestLLMOperatorConfidenceGate:
body = mock_upsert.call_args.kwargs["body"]
assert "Confidence: 0.40 (minimum 0.70)" in body
+ # The JSON a reviewer can paste into the edit field, not the model's
repr.
+ assert '{"category":"billing","score":0.3}' in body
+ assert "category='billing'" not in body
pushes = {
c.kwargs["key"]: c.kwargs["value"] for c in
context["task_instance"].xcom_push.call_args_list
}
@@ -1374,6 +1378,51 @@ class TestLLMOperatorApproval:
assert result == expected
+ def test_execute_complete_fails_on_edit_invalid_for_output_type(self):
+ op = LLMOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ output_type=int,
+ require_approval=True,
+ allow_modifications=True,
+ )
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": {"id": "u1", "name": "admin"},
+ "params_input": {"output": "thirty-one"},
+ }
+
+ with pytest.raises(ReviewedOutputValidationError,
match="'thirty-one'"):
+ op.execute_complete({}, generated_output="5", event=event)
+
+ @requires_typed_xcom
+ def
test_execute_complete_restores_a_model_the_form_sent_back_as_an_object(self):
+ # The HITL form parses a JSON-object value into an object, so an
untouched model comes back as a dict.
+ op = LLMOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ output_type=Summary,
+ require_approval=True,
+ allow_modifications=True,
+ )
+ event = {
+ "chosen_options": ["Approve"],
+ "responded_by_user": {"id": "u1", "name": "admin"},
+ "params_input": {"output": {"text": "hello"}},
+ }
+
+ result = op.execute_complete({}, generated_output='{"text":"hello"}',
event=event)
+
+ assert result == Summary(text="hello")
+
+ def
test_execute_complete_returns_an_unedited_output_that_does_not_validate_back(self):
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
output_type=int, require_approval=True)
+ event = {"chosen_options": ["Approve"], "responded_by_user": {"id":
"u1", "name": "admin"}}
+
+ assert op.execute_complete({}, generated_output="null", event=event)
== "null"
+
@pytest.mark.skipif(
not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with
Airflow >= 3.1.0"
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
index c8cc658bb93..ffabb18f47e 100644
---
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
+++
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
@@ -29,6 +29,7 @@ from pydantic_ai.messages import ModelResponse, TextPart
from pydantic_ai.models.function import FunctionModel
from pydantic_ai.usage import UsageLimits
+from airflow.providers.common.ai.exceptions import
ReviewedOutputValidationError
from airflow.providers.common.ai.operators.llm_schema_compare import (
LLMSchemaCompareOperator,
SchemaCompareResult,
@@ -706,7 +707,7 @@ class TestLLMSchemaCompareOperatorUsageBudget:
seed_usage_budget(task_state_store_accessor, requests=1)
op = self._make_operator(require_approval=True,
allow_modifications=True)
- with pytest.raises(ValueError, match="not a valid
SchemaCompareResult"):
+ with pytest.raises(ReviewedOutputValidationError, match="output_type
SchemaCompareResult"):
op.execute_complete(
usage_budget_context,
generated_output="{}",
@@ -823,7 +824,7 @@ class TestLLMSchemaCompareOperatorApproval:
"params_input": {"output": modified},
}
- with pytest.raises(ValueError, match="not a valid
SchemaCompareResult"):
+ with pytest.raises(ReviewedOutputValidationError, match="output_type
SchemaCompareResult"):
op.execute_complete({}, generated_output=result.model_dump_json(),
event=event)
def test_execute_rejects_sequence_prompt_with_require_approval(self):
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
b/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
index f7e1a4799e2..cd22f989c3e 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_output_type.py
@@ -16,16 +16,135 @@
# under the License.
from __future__ import annotations
+from collections import deque
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import Enum
+from typing import Annotated, Literal, NamedTuple
+
import pytest
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, Field, computed_field,
model_validator
+from pydantic_ai import NativeOutput, PromptedOutput, ToolOutput
-from airflow.providers.common.ai.utils.output_type import
rehydrate_pydantic_output
+from airflow.providers.common.ai.exceptions import
ReviewedOutputValidationError
+from airflow.providers.common.ai.utils.output_type import (
+ _is_container_mismatch,
+ dump_output_to_json,
+ rehydrate_pydantic_output,
+)
class A(BaseModel):
x: int
+class B(BaseModel):
+ y: str
+
+
+class FromAttributesModel(BaseModel):
+ model_config = ConfigDict(from_attributes=True)
+ x: int
+
+
+@dataclass
+class DataclassOutput:
+ a: int
+
+
+class Color(str, Enum):
+ RED = "red"
+
+
+class ValidationAlias(BaseModel):
+ user_name: str = Field(alias="userName")
+
+
+class SerializationAlias(BaseModel):
+ user_name: str = Field(serialization_alias="userName")
+
+
+class SerializedByAlias(BaseModel):
+ model_config = ConfigDict(serialize_by_alias=True)
+ user_name: str = Field(serialization_alias="userName")
+
+
+class DictOnlyBeforeValidator(BaseModel):
+ """A model whose before-validator assumes dict input, so it breaks on a
bare string."""
+
+ amount: int
+
+ @model_validator(mode="before")
+ @classmethod
+ def default_amount(cls, data):
+ amount = data.get("amount", 0)
+ return {**data, "amount": amount}
+
+
+class ValidatorWithOperationalError(BaseModel):
+ """A model whose validator fails for a reason that has nothing to do with
the review text."""
+
+ amount: int
+
+ @model_validator(mode="before")
+ @classmethod
+ def look_up_amount(cls, data):
+ raise TimeoutError("lookup service did not answer")
+
+
+class AliasNamedLikeAnotherField(BaseModel):
+ """A model that validates by name only, with an alias that is another
field's name."""
+
+ model_config = ConfigDict(validate_by_alias=False, validate_by_name=True)
+ source: str
+ destination: str = Field(alias="source")
+
+
+class IndexingBeforeValidator(BaseModel):
+ """A model whose before-validator indexes into its input, so a missing key
or item raises."""
+
+ first: str
+
+ @model_validator(mode="before")
+ @classmethod
+ def take_first_item(cls, data):
+ return {"first": data["items"][0]}
+
+
+class ComputedForbidExtra(BaseModel):
+ model_config = ConfigDict(extra="forbid")
+ a: int
+
+ @computed_field # type: ignore[prop-decorator]
+ @property
+ def double(self) -> int:
+ return self.a * 2
+
+
+class Address(BaseModel):
+ city: str
+
+
+class Order(BaseModel):
+ # Reuses Address, so pydantic wraps a function taking an Order in a
``definitions`` schema.
+ billing: Address
+ shipping: Address
+
+
+class NoSchema:
+ """A plain class pydantic cannot build a schema for."""
+
+
+@dataclass
+class UnresolvedForwardReference:
+ value: Missing # type: ignore[name-defined] # noqa: F821
+
+
+class Point(NamedTuple):
+ x: int
+ y: str
+
+
class TestRehydratePydanticOutput:
def test_returns_model_instance(self):
result = rehydrate_pydantic_output(A, '{"x": 7}',
serialize_output=False)
@@ -36,10 +155,6 @@ class TestRehydratePydanticOutput:
result = rehydrate_pydantic_output(A, '{"x": 7}',
serialize_output=True)
assert result == {"x": 7}
- def test_returns_raw_for_str_output_type(self):
- result = rehydrate_pydantic_output(str, "anything",
serialize_output=False)
- assert result == "anything"
-
@pytest.mark.parametrize(
("output_type", "raw", "expected"),
[(int, "5", 5), (bool, "true", True), (list[str], '["a", "b"]', ["a",
"b"])],
@@ -48,15 +163,290 @@ class TestRehydratePydanticOutput:
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),
+ (str | None, "hello\n", "hello\n"),
+ (Literal["a", "b"], "b\n", "b"),
+ (Color, " red ", Color.RED),
+ ],
+ ids=[
+ "literal",
+ "optional-str",
+ "str-enum",
+ "optional-str-keeps-its-newline",
+ "literal-with-newline",
+ "str-enum-with-spaces",
+ ],
+ )
+ def test_accepts_bare_text_for_string_valued_types(self, output_type, raw,
expected):
+ result = rehydrate_pydantic_output(output_type, raw,
serialize_output=False)
+ assert result == expected
+ assert type(result) is type(expected)
+
+ @pytest.mark.parametrize(
+ "takes_an_order", [False, True], ids=["call-schema",
"definitions-wrapped-schema"]
+ )
+ def test_returns_raw_for_output_function_without_calling_it(self,
takes_an_order):
+ calls = []
+
+ def clean(text: str) -> str:
+ calls.append(text)
+ return text
+
+ def route(order: Order) -> str:
+ calls.append(order.billing.city)
+ return order.billing.city
+
+ output_function = route if takes_an_order else clean
- def test_returns_raw_on_invalid_json(self):
- result = rehydrate_pydantic_output(A, "not-json",
serialize_output=False)
- assert result == "not-json"
+ assert rehydrate_pydantic_output(output_function, "Paris",
serialize_output=False) == "Paris"
+ assert calls == []
- def test_returns_raw_on_schema_mismatch(self):
- # ``A`` requires ``x: int`` -- this payload should fail validation
- result = rehydrate_pydantic_output(A, '{"y": "no-x-field"}',
serialize_output=False)
- assert result == '{"y": "no-x-field"}'
+ @pytest.mark.parametrize("output_type", [NoSchema,
UnresolvedForwardReference])
+ def test_returns_raw_for_type_pydantic_cannot_build_a_schema_for(self,
output_type):
+ assert rehydrate_pydantic_output(output_type, "anything",
serialize_output=False) == "anything"
+
+ @pytest.mark.parametrize(
+ "output_type",
+ [ToolOutput(A), NativeOutput(A), PromptedOutput(A)],
+ ids=["tool-output", "native-output", "prompted-output"],
+ )
+ @pytest.mark.parametrize("is_edited", [False, True], ids=["unedited",
"edited"])
+ def test_returns_raw_for_a_pydantic_ai_output_marker(self, output_type,
is_edited):
+ result = rehydrate_pydantic_output(
+ output_type, '{"x": 1}', serialize_output=False,
is_edited=is_edited
+ )
+
+ assert result == '{"x": 1}'
+
+ @pytest.mark.parametrize(
+ "raw",
+ ["anything", '"hi"', "null", "42"],
+ ids=["plain-text", "json-string", "json-null", "json-number"],
+ )
+ def test_returns_any_text_for_a_str_output_type_unparsed(self, raw):
+ assert rehydrate_pydantic_output(str, raw, serialize_output=False) ==
raw
+
+ @pytest.mark.parametrize("output_type", [Point, Annotated[Point, "x"]],
ids=["bare", "annotated"])
+ def test_validates_a_named_tuple_even_though_its_schema_is_a_call(self,
output_type):
+ result = rehydrate_pydantic_output(output_type, '[1, "a"]',
serialize_output=False)
+
+ assert result == Point(1, "a")
+ assert type(result) is Point
+
+ @pytest.mark.parametrize(
+ "raw", ['{"user_name": "kax"}', '{"userName": "kax"}'],
ids=["field-name", "alias"]
+ )
+ def test_accepts_a_field_name_or_an_alias(self, raw):
+ result = rehydrate_pydantic_output(ValidationAlias, raw,
serialize_output=False)
+
+ assert type(result) is ValidationAlias
+ assert result.user_name == "kax"
+
+ def
test_returns_raw_with_a_warning_for_an_unedited_output_that_does_not_validate_back(self,
caplog):
+ assert rehydrate_pydantic_output(int, "null", serialize_output=False,
is_edited=False) == "null"
+ assert any("Returning the output unchanged" in message for message in
caplog.messages)
+
+ def test_error_truncates_long_received_text(self):
+ with pytest.raises(ReviewedOutputValidationError) as exc_info:
+ rehydrate_pydantic_output(int, "x" * 500, serialize_output=False)
+
+ assert "x" * 200 + "..." in str(exc_info.value)
+ assert "x" * 201 not in str(exc_info.value)
+
+ @pytest.mark.parametrize(
+ ("raw", "detail"),
+ [
+ ('{"amount": "ten"}', "valid integer"),
+ ('{"amount": 1,}', "trailing comma"),
+ ('["x"]', "has no attribute 'get'"),
+ ('"x"', "has no attribute 'get'"),
+ ("42", "has no attribute 'get'"),
+ ],
+ ids=["invalid-field", "json-typo", "json-array", "json-string",
"json-number"],
+ )
+ def
test_a_before_validator_failing_on_bare_text_reports_the_json_error(self, raw,
detail):
+ with pytest.raises(ReviewedOutputValidationError, match=detail):
+ rehydrate_pydantic_output(DictOnlyBeforeValidator, raw,
serialize_output=False)
+
+ @pytest.mark.parametrize(
+ ("raw", "detail"),
+ [("{}", "'items'"), ('{"items": []}', "list index out of range")],
+ ids=["key-error", "index-error"],
+ )
+ def
test_a_before_validator_indexing_a_wrong_shape_reports_a_failed_conversion(self,
raw, detail):
+ with pytest.raises(ReviewedOutputValidationError, match=detail):
+ rehydrate_pydantic_output(IndexingBeforeValidator, raw,
serialize_output=False)
+
+ @pytest.mark.parametrize("is_edited", [False, True], ids=["unedited",
"edited"])
+ def test_a_validator_error_unrelated_to_the_text_propagates(self,
is_edited):
+ with pytest.raises(TimeoutError, match="lookup service"):
+ rehydrate_pydantic_output(
+ ValidatorWithOperationalError, '{"amount": 1}',
serialize_output=False, is_edited=is_edited
+ )
+
+ def
test_a_before_validator_failing_on_bare_text_leaves_an_unedited_output_unchanged(self):
+ raw = '{"amount": "ten"}'
+
+ result = rehydrate_pydantic_output(
+ DictOnlyBeforeValidator, raw, serialize_output=False,
is_edited=False
+ )
+
+ assert result == raw
+
+ @pytest.mark.parametrize(
+ ("output_type", "raw", "type_name", "detail"),
+ [
+ (int, "thirty-one", "int", "valid integer"),
+ (list[str], "a, b, c", "list[str]", "Invalid JSON"),
+ (A, '{"y": "no-x-field"}', "A", "Field required"),
+ ],
+ ids=["int", "list", "basemodel"],
+ )
+ def test_raises_when_type_adapter_rejects(self, output_type, raw,
type_name, detail):
+ with pytest.raises(ReviewedOutputValidationError) as exc_info:
+ rehydrate_pydantic_output(output_type, raw, serialize_output=False)
+
+ message = str(exc_info.value)
+ assert f"output_type {type_name}." in message
+ assert repr(raw) in message
+ assert detail in message
+
+ @pytest.mark.parametrize(
+ ("output_type", "output"),
+ [
+ (ValidationAlias, ValidationAlias(userName="kax")),
+ (SerializationAlias, SerializationAlias(user_name="kax")),
+ (SerializedByAlias, SerializedByAlias(user_name="kax")),
+ (list[ValidationAlias], [ValidationAlias(userName="kax")]),
+ (ComputedForbidExtra, ComputedForbidExtra(a=1)),
+ (list[ComputedForbidExtra], [ComputedForbidExtra(a=1)]),
+ ],
+ ids=[
+ "validation-alias",
+ "serialization-alias",
+ "serialize-by-alias",
+ "list-of-aliased",
+ "computed-field-forbid-extra",
+ "list-of-computed-field-forbid-extra",
+ ],
+ )
+ def test_a_dumped_output_validates_back(self, output_type, output):
+ raw = dump_output_to_json(output)
+
+ assert rehydrate_pydantic_output(output_type, raw,
serialize_output=False) == output
+
+ @pytest.mark.parametrize(
+ ("output_type", "raw", "detail"),
+ [
+ (A, '{"x": 1,}', "trailing comma"),
+ (list[str], '["a","b",]', "trailing comma"),
+ (A | B, '{"x": 1,}', "trailing comma"),
+ (list[int] | dict[str, int], "[1,]", "trailing comma"),
+ (DataclassOutput, '{"a": 1,}', "trailing comma"),
+ (frozenset[int], "[1,]", "trailing comma"),
+ (Point, '[1, "a",]', "trailing comma"),
+ (Literal["a", "b"], "c", "Input should be 'a' or 'b'"),
+ (A | int, "hello", "valid integer"),
+ (Literal["a", "b"] | A, '{"x": 1,}', "trailing comma"),
+ (Literal["a", "b"] | A, ' {"x": 1,}', "trailing comma"),
+ (int, "3.5", "fractional part"),
+ ],
+ ids=[
+ "model-typo",
+ "list-typo",
+ "model-union-typo",
+ "container-union-typo",
+ "dataclass-typo",
+ "frozenset-typo",
+ "named-tuple-typo",
+ "bare-text",
+ "bare-text-for-a-model-or-int",
+ "json-typo-for-a-literal-or-model",
+ "json-typo-after-whitespace",
+ "valid-json-of-the-wrong-kind",
+ ],
+ )
+ def
test_error_reports_the_json_error_for_a_json_typo_and_the_text_error_otherwise(
+ self, output_type, raw, detail
+ ):
+ with pytest.raises(ReviewedOutputValidationError, match=detail):
+ rehydrate_pydantic_output(output_type, raw, serialize_output=False)
+
+ @pytest.mark.parametrize(
+ "output_type",
+ [
+ A,
+ FromAttributesModel,
+ DataclassOutput,
+ dict[str, int],
+ list[int],
+ tuple[int, str],
+ set[int],
+ frozenset[int],
+ Point,
+ deque[int],
+ Sequence[int],
+ ],
+ ids=[
+ "model_type",
+ "model_attributes_type",
+ "dataclass_type",
+ "dict_type",
+ "list_type",
+ "tuple_type",
+ "set_type",
+ "frozen_set_type",
+ "arguments_or_named_tuple_type",
+ "list_or_deque_type",
+ "sequence_str",
+ ],
+ )
+ def
test_error_reports_the_json_error_for_plain_text_sent_to_a_structured_type(self,
output_type):
+ # Each id is the error pydantic raises when the whole text meets the
type (pydantic 2.14 renamed
+ # the NamedTuple and deque ones). The JSON error says more.
+ with pytest.raises(ReviewedOutputValidationError, match="Invalid
JSON"):
+ rehydrate_pydantic_output(output_type, "hello",
serialize_output=False)
+
+ def
test_a_union_with_str_keeps_an_edit_that_is_not_the_model_as_text(self):
+ raw = '{"x": "not-a-number"}'
+
+ assert rehydrate_pydantic_output(A | str, raw, serialize_output=False)
== raw
+
+
+def
test_a_model_that_turned_off_validation_by_alias_reads_its_dump_back_unchanged():
+ output = AliasNamedLikeAnotherField(source="SRC", destination="DST")
+
+ result = rehydrate_pydantic_output(
+ AliasNamedLikeAnotherField, dump_output_to_json(output),
serialize_output=False, is_edited=False
+ )
+
+ assert result == output
+
+
+class TestIsContainerMismatch:
+ @pytest.mark.parametrize(
+ "error_type",
+ [
+ "model_type",
+ "dict_type",
+ "list_type",
+ "arguments_type",
+ # Names pydantic 2.14 introduced, which the locked pydantic cannot
raise.
+ "named_tuple_type",
+ "deque_type",
+ "sequence_str",
+ ],
+ )
+ def test_a_structured_type_refusing_the_whole_text(self, error_type):
+ assert _is_container_mismatch(error_type)
+
+ @pytest.mark.parametrize(
+ "error_type", ["int_parsing", "bool_parsing", "uuid_parsing",
"literal_error", "enum", "json_invalid"]
+ )
+ def test_a_type_that_takes_text_rejecting_its_value(self, error_type):
+ assert not _is_container_mismatch(error_type)
