SIGN IN SIGN UP
apache / airflow UNCLAIMED

Apache Airflow - A platform to programmatically author, schedule, and monitor workflows

0 0 267 Python

Return Pydantic model instances through XCom for structured output (#67644)

`LLMOperator`, `LLMAgentOperator`, `LLMFileAnalysisOperator`, and their
`@task.llm` / `@task.agent` / `@task.llm_file_analysis` decorators stop
calling `model_dump()` on Pydantic outputs before pushing to XCom. Downstream
tasks now receive the model instance directly, so they can type-hint the
class (`def downstream(result: MyModel)`) and use attribute access
(`result.field`) instead of subscript access on a dict.

To avoid forcing every DAG author to edit `[core] allowed_deserialization_classes`,
the operators auto-register their `output_type` (and any `BaseModel`
reachable from `Union`/`Optional`/`list` shapes) via a new
`airflow.sdk.serde.allow_class(cls)` helper. The registration is process-local
and runs in each worker's `__init__` -- same-DAG downstream tasks parse the
DAG file when they start up, which re-runs the constructor and re-populates
the per-process allow-list.

The helper rejects classes that cannot be re-imported by qualname (defined
in a function body, nested in another class, dynamically built with a
mismatched `__name__`, or parametrised generics) so the failure surfaces at
DAG parse time rather than at XCom-consume time.

UI XCom viewer and cross-DAG `xcom_pull` are still gated by
`[core] allowed_deserialization_classes` because the API server and other
DAGs' workers don't import the producing DAG. Documented explicitly in the
operator guides.

Older Airflow versions that lack `allow_class` continue to get the dict
form via a try/except fallback in each operator, so the provider keeps
working on `apache-airflow>=3.0.0`.

* common-ai: Correct docs on what UI XCom viewer shows for Pydantic outputs

The UI's XCom viewer renders structured-output Pydantic instances via the
``stringify`` path (``airflow.serialization.stringify``) rather than the
``deserialize`` path, so user classes outside the ``airflow.*`` glob do not
hit the allow-list gate -- they show up as ``module.MyModel@version=1(...)``
without any config change. Only cross-DAG ``xcom_pull`` is still gated.

Also hoist a ``pydantic.create_model`` import to module scope in the serde
test that was using it inline.

* Make XCom stringify readable for user Pydantic/dataclass classes

Strip DagBag's ``unusual_prefix_<sha>_`` module prefix from the displayed
classname and repr-quote string field values inside the ``classname@version=N(...)``
form. Before this change, an XCom value carrying a user-defined Pydantic class
rendered in the UI as:

   unusual_prefix_9ce9eb..._typed_xcom_demo.TicketAnalysis@version=1(
     priority=high,category=bug,summary=Nightly ETL...)

After:

   typed_xcom_demo.TicketAnalysis@version=1(
     priority='high', category='bug', summary='Nightly ETL...')

The prefix is a DagBag artifact (added to avoid ``sys.modules`` clashes
between same-named DAG files in different bundles) and has no value in the
human-readable XCom display. Quoting strings disambiguates ``field=value``
from a bare token and matches Pydantic/dataclass repr conventions.

* common-ai: Fix CI failures from PR review

Three CI failures fixed:

1. Compat tests against Airflow 3.0.6 / 3.1.8: new tests assumed
   allow_class is importable and asserted on Pydantic instance shape.
   Gate the new tests behind a requires_allow_class marker so they skip
   cleanly on older Airflow (operators already fall back to model_dump
   there via the import-safe import).

2. Docs build failed with 12 RST errors in autoapi-generated index.rst
   for example_dags modules. Pydantic BaseModel's inherited docstring
   leaks through autoapi rendering and breaks the Definition list. An
   explicit docstring on each module-level Pydantic class overrides the
   inherited one and keeps the RST valid.

3. Spell-check: qualname is a Python attribute name; backtick it in
   prose so the spell-checker treats it as code. Switched 'parametrised'
   (British) to 'parameterized' (American) to match wordlist.

* common-ai: Add serialize_output flag for opt-in dict shape

Per Jed's review: some downstream consumers want the dict shape (e.g.
forwarding the value to an external system that expects JSON-style payloads).

Add serialize_output: bool = False to LLMOperator and AgentOperator (and via
inheritance, LLMFileAnalysisOperator). When True the operator calls
model_dump() before pushing to XCom, restoring the pre-PR behavior on demand
without giving up the typed default. The class is not registered in
_extra_allowed in that mode since the wire carries a plain dict and never
hits the allow-list gate.
K
Kaxil Naik committed
9318bd62502bc005491c3a72710b0ffe5e5968c6
Parent: 1974666
Committed by GitHub <noreply@github.com> on 5/29/2026, 1:06:07 AM