This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun 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 226de8a1911 Publish usage, run id and model name from LLMOperator's
subclasses (#75208)
226de8a1911 is described below
commit 226de8a19117856d2ea8587b03692031aa6b7239
Author: Pierre Jeambrun <[email protected]>
AuthorDate: Thu Oct 29 10:58:32 2026 +0100
Publish usage, run id and model name from LLMOperator's subclasses (#75208)
* Publish usage, run id and model name from LLMOperator's subclasses
LLMBranchOperator, LLMSQLQueryOperator, LLMSchemaCompareOperator and
LLMFileAnalysisOperator override execute() without calling
super().execute(), so they never pushed the usage, run_id and model name
XComs, never stamped the task-instance identity on their GenAI spans,
and had no AI Model tab. Calling super() is not an option: it would also
apply LLMOperator's decision record, confidence gate and review body,
which these operators replace or don't have. The reporting now lives in
_run_llm, the one place every execute() runs the agent, so the
subclasses, and any future one, report the same metadata as LLMOperator.
* Clarify which operators publish the run telemetry
The Scope bullet read as if agents traced through agent_framework_tracing
were excluded from everything, when they still get the airflow.* identity
attributes and only miss the XComs.
* Rename the agent call under _run_llm to _run_agent_within_budget
Both methods were named after running, which hid that _run_llm is the
entry point every execute() calls and the other one only runs the agent
under the usage budget.
* Make the subclass telemetry tests catch a lost run id or identity stamp
The subclass tests passed even when _run_llm forwarded no run id or skipped
the
identity stamp, and the AI Model tab's operator list could miss a new
LLMOperator
subclass unnoticed. The tests now fail in each of those cases.
* Point the subclass docs at the run XComs and spec the test task instance
The SQL, branch, file-analysis and schema-compare pages stopped at the log
summary, so a reader of those pages could not tell the operators also push
usage, run_id and the model name. The XCom tests used an unspecced task
instance, so any attribute the identity stamping read would silently pass.
* Fix the docs build on the run XComs pointers
The anchor sits inside a bullet rather than before a section, so Sphinx has
no title to use for a bare reference to it.
* Link the subclass pages to the observability page instead of a bullet
anchor
Every other operator page, llm.rst included, points at the observability
page as a whole; an anchor inside a bullet was a one-off pattern.
---
providers/common/ai/docs/features.rst | 4 +-
providers/common/ai/docs/observability.rst | 43 ++++-----
providers/common/ai/docs/operators/llm_branch.rst | 6 +-
.../common/ai/docs/operators/llm_file_analysis.rst | 10 ++
.../ai/docs/operators/llm_schema_compare.rst | 3 +
providers/common/ai/docs/operators/llm_sql.rst | 3 +
.../airflow/providers/common/ai/operators/llm.py | 101 ++++++++++++---------
.../providers/common/ai/plugins/model_panel.py | 14 ++-
.../common/ai/tests/unit/common/ai/conftest.py | 57 ++++++++++++
.../unit/common/ai/decorators/test_llm_branch.py | 12 ++-
.../common/ai/decorators/test_llm_file_analysis.py | 8 +-
.../ai/decorators/test_llm_schema_compare.py | 2 +-
.../unit/common/ai/decorators/test_llm_sql.py | 12 ++-
.../unit/common/ai/operators/test_llm_branch.py | 84 +++++++++++------
.../common/ai/operators/test_llm_file_analysis.py | 42 +++++++--
.../common/ai/operators/test_llm_schema_compare.py | 26 +++++-
.../tests/unit/common/ai/operators/test_llm_sql.py | 18 +++-
.../unit/common/ai/plugins/test_model_panel.py | 34 ++++---
18 files changed, 339 insertions(+), 140 deletions(-)
diff --git a/providers/common/ai/docs/features.rst
b/providers/common/ai/docs/features.rst
index 7a461dec8ae..9ecee96b711 100644
--- a/providers/common/ai/docs/features.rst
+++ b/providers/common/ai/docs/features.rst
@@ -302,8 +302,8 @@ See what ran
- Exports model and tool calls as OpenTelemetry spans.
- None
* - :ref:`AI Model tab <howto/observability:model-tab>`
- - Shows the model on the task instance page of ``@task.llm`` and
``@task.agent``
- tasks, and their token counts and an estimated cost.
+ - Shows the model on the task instance page of ``@task.agent`` and every
``@task.llm*``
+ task except ``@task.llm_batch``, and their token counts and an
estimated cost.
- None
.. toctree::
diff --git a/providers/common/ai/docs/observability.rst
b/providers/common/ai/docs/observability.rst
index 8068bdf7a4b..b1a6f4752db 100644
--- a/providers/common/ai/docs/observability.rst
+++ b/providers/common/ai/docs/observability.rst
@@ -49,9 +49,9 @@ How it works
no GenAI spans are emitted.
* **Correlation.** The worker opens a task span before the operator runs, so
the
agent's spans nest under it and share its ``trace_id``.
- :class:`~airflow.providers.common.ai.operators.agent.AgentOperator` and
- :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` (and their
- ``@task.agent`` / ``@task.llm`` decorator forms) additionally stamp the
+ :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`,
+ :class:`~airflow.providers.common.ai.operators.llm.LLMOperator` and its SQL
/ branch /
+ file-analysis / schema-compare subclasses (and their decorator forms)
additionally stamp the
task-instance identity on every GenAI span they emit: the five keys core
tracing already puts on the task span (``airflow.dag_id``,
``airflow.task_id``, ``airflow.dag_run.run_id``,
``airflow.task_instance.try_number``, ``airflow.task_instance.map_index``)
plus
@@ -63,8 +63,8 @@ How it works
attempts share one trace and appear as repeated task-run spans on it,
distinguished by ``try number``. Only a manual clear or rerun regenerates the
context and starts a new trace.
-* **Run join key.** For an ``AgentOperator`` or ``LLMOperator`` run (and their
- ``@task.agent`` / ``@task.llm`` decorator equivalents), the task-instance id
+* **Run join key.** For an ``AgentOperator`` run or a run of ``LLMOperator``
or one of
+ its subclasses (and their decorator equivalents), the task-instance id
(unique per attempt, since Airflow regenerates it on each retry) is passed to
pydantic-ai as the run's ``run_id``. It surfaces on the run's GenAI spans as
``gen_ai.agent.call.id``, and the operator also exposes it, alongside the
run's
@@ -92,31 +92,26 @@ How it works
``usage`` above, so a downstream task, a Jinja template, or a REST client
can read which
model responded without parsing the decision record or a span
(``ti.xcom_pull(task_ids="my_agent",
key="__AIRFLOW__COMMON_AI_MODEL_NAME__")``).
- Only ``AgentOperator`` /
- ``@task.agent`` and ``LLMOperator`` / ``@task.llm`` push it today; like
``run_id`` and
- ``usage``, with ``enable_hitl_review`` it is the initial run's model, not a
- human-feedback regeneration's, since ``regenerate_with_feedback`` doesn't
re-push it.
- ``LLMOperator``'s subclasses that override ``execute()`` instead of calling
- ``super().execute()`` -- ``LLMBranchOperator``, ``LLMSQLQueryOperator``,
- ``LLMSchemaCompareOperator``, ``LLMFileAnalysisOperator`` -- don't publish
it yet;
- this is upcoming work.
-* **Scope.** The ``run_id`` / ``usage`` XComs, the model name, and the
``airflow.*``
- identity attributes come from ``AgentOperator`` / ``@task.agent`` and
``LLMOperator`` /
- ``@task.llm``, apart from a Strands, ADK or OpenAI Agents SDK agent run
inside
- ``agent_framework_tracing`` (see below). The other LLM operators still emit
GenAI spans
- correlated to the task span by nesting, but without the identity attributes,
the run
- join key, the model name, or a ``usage`` XCom.
+ ``AgentOperator`` and ``LLMOperator`` and its subclasses (and their
decorator forms)
+ push it; like ``run_id`` and ``usage``, with ``enable_hitl_review`` it is
the initial
+ run's model, not a human-feedback regeneration's, since
``regenerate_with_feedback``
+ doesn't re-push it.
+* **Scope.** The ``run_id`` / ``usage`` XComs and the model name come from
+ ``AgentOperator`` and ``LLMOperator`` and its subclasses (and their
decorator forms).
+ Those operators also stamp the ``airflow.*`` identity attributes on their
spans, and so
+ does a Strands, ADK or OpenAI Agents SDK agent run inside
``agent_framework_tracing``
+ (see below).
* .. _howto/observability:model-tab:
**Model tab.** Separately from tracing, the provider adds an "AI Model" tab
to the
task-instance page showing the resolved model name and the ``usage`` XCom
above as a
readable table, including an estimated cost. It needs Airflow >= 3.4, since
it relies on
the ``applies_to`` operator scoping core added in that version to show up
only for
- ``LLMOperator`` / ``AgentOperator`` runs (and their ``@task.llm`` /
``@task.agent``
- decorator forms) -- the match is on the exact operator name, so a subclass
of either does
- not get the tab automatically. Like the span-level cost above, the tab's
number is a
- best-effort ``genai-prices`` estimate, not provider-reported billing. With a
custom XCom
- backend that offloads large values to external storage (e.g. ``common.io``
with
+ ``AgentOperator`` runs and runs of ``LLMOperator`` and its SQL / branch /
file-analysis /
+ schema-compare subclasses (and their decorator forms) -- the match is on the
exact
+ operator name, so a subclass of your own does not get the tab automatically.
Like the
+ span-level cost above, the tab's number is a best-effort ``genai-prices``
estimate, not
+ provider-reported billing. With a custom XCom backend that offloads large
values to external storage (e.g. ``common.io`` with
``xcom_objectstorage_threshold`` set low), the tab reads the stored
reference rather than
the real value, since the public XCom endpoint it calls never runs the
backend's
``deserialize_value``: the model name badge would show the reference string
(e.g. an
diff --git a/providers/common/ai/docs/operators/llm_branch.rst
b/providers/common/ai/docs/operators/llm_branch.rst
index b0c260db5be..b7405200037 100644
--- a/providers/common/ai/docs/operators/llm_branch.rst
+++ b/providers/common/ai/docs/operators/llm_branch.rst
@@ -231,7 +231,8 @@ The decision record
^^^^^^^^^^^^^^^^^^^
Whether or not a bar is set, the operator pushes a ``decision`` XCom next to
-its return value (suppressed by ``do_xcom_push=False`` like any other):
+its return value and the run XComs listed under `Logging`_ (suppressed by
+``do_xcom_push=False`` like any other):
.. code-block:: json
@@ -350,3 +351,6 @@ Logging
After each LLM call, the operator logs a summary with model name, token usage,
and request count at INFO level. See :ref:`AgentOperator logging
<howto/operator:agent>`
for details on the log format.
+
+Like :doc:`llm`, the operator also pushes ``usage``, ``run_id`` and the model
name to
+XCom, on a failed run too. See :doc:`../observability`.
diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst
b/providers/common/ai/docs/operators/llm_file_analysis.rst
index 1a666b16893..e4efec17dae 100644
--- a/providers/common/ai/docs/operators/llm_file_analysis.rst
+++ b/providers/common/ai/docs/operators/llm_file_analysis.rst
@@ -196,3 +196,13 @@ Parquet and Avro readers require their corresponding
optional extras:
pip install "apache-airflow-providers-common-ai[parquet]"
pip install "apache-airflow-providers-common-ai[avro]"
+
+Logging
+-------
+
+After each LLM call, the operator logs a summary with model name, token usage,
+and request count at INFO level. See :ref:`AgentOperator logging
<howto/operator:agent>`
+for details on the log format.
+
+Like :doc:`llm`, the operator also pushes ``usage``, ``run_id`` and the model
name to
+XCom, on a failed run too. See :doc:`../observability`.
diff --git a/providers/common/ai/docs/operators/llm_schema_compare.rst
b/providers/common/ai/docs/operators/llm_schema_compare.rst
index c7a2e729d95..d7e821acb3a 100644
--- a/providers/common/ai/docs/operators/llm_schema_compare.rst
+++ b/providers/common/ai/docs/operators/llm_schema_compare.rst
@@ -227,3 +227,6 @@ Logging
After each LLM call, the operator logs a summary with model name, token usage,
and request count at INFO level. See :ref:`AgentOperator logging
<howto/operator:agent>`
for details on the log format.
+
+Like :doc:`llm`, the operator also pushes ``usage``, ``run_id`` and the model
name to
+XCom, on a failed run too. See :doc:`../observability`.
diff --git a/providers/common/ai/docs/operators/llm_sql.rst
b/providers/common/ai/docs/operators/llm_sql.rst
index 5f457fb016b..89a9d9bc987 100644
--- a/providers/common/ai/docs/operators/llm_sql.rst
+++ b/providers/common/ai/docs/operators/llm_sql.rst
@@ -185,3 +185,6 @@ After each LLM call, the operator logs a summary with model
name, token usage,
and request count at INFO level. At DEBUG level, the generated SQL is also
logged (truncated to 500 characters). See :ref:`AgentOperator logging
<howto/operator:agent>`
for details on the log format.
+
+Like :doc:`llm`, the operator also pushes ``usage``, ``run_id`` and the model
name to
+XCom, on a failed run too. See :doc:`../observability`.
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 3ad36fffdb9..2cc5c822e06 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
@@ -368,45 +368,10 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
agent = self._create_agent(output_type=self.output_type,
instructions=self.system_prompt)
- # A per-attempt key (the task-instance id on Airflow 3, which is
regenerated on each
- # retry; dag/run/task/map/try on Airflow 2) ties this run to its GenAI
trace, the same
- # join key AgentOperator uses. A hand-built context with no task
instance runs untracked.
- ti = self._resolve_task_instance(context)
- run_kwargs: dict[str, Any] = {}
- if ti is not None:
- stamp_identity_on_agent_spans(agent,
build_run_identity_attributes(ti))
- run_kwargs["run_id"] = make_task_instance_run_key(ti)
-
- attempt_usage = RunUsage()
- try:
- result = self._run_llm(
- context, agent, self.prompt, usage_limits=usage_limits,
run_usage=attempt_usage, **run_kwargs
- )
- except BaseException:
- # A run that raises partway through (e.g. a cost_limit-triggered
UsageLimitExceeded)
- # still reports what it spent: the capped run is the one cost
tracking most needs.
- try:
- log_run_usage(self.log, attempt_usage, outcome="failed")
- except Exception:
- self.log.warning("Failed to log partial usage for the failed
run", exc_info=True)
- try:
- self._push_xcom(context, "usage",
format_usage_for_xcom(attempt_usage))
- except Exception:
- self.log.warning("Failed to push usage XCom for the failed
run", exc_info=True)
- if ti is not None:
- try:
- self._push_xcom(context, "run_id", run_kwargs["run_id"])
- except Exception:
- self.log.warning("Failed to push run_id XCom for the
failed run", exc_info=True)
- raise
+ result = self._run_llm(context, agent, self.prompt,
usage_limits=usage_limits)
output = result.output
- self._push_xcom(context, "usage", format_usage_for_xcom(attempt_usage))
- if ti is not None:
- self._push_xcom(context, "run_id", result.run_id)
model_confidence = ModelConfidence.from_result(result)
- if model_confidence.model is not None:
- self._push_xcom(context, MODEL_NAME_XCOM_KEY,
model_confidence.model)
# Gated by the least confident of the fields that reported a
confidence; a field whose type
# reports none is not gated. A bare output type is the one field
``response``.
fields = list(model_confidence.confidence) or ["response"]
@@ -499,13 +464,67 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
user_prompt: Any,
*,
usage_limits: UsageLimits | None,
- run_usage: RunUsage | None = None,
+ ) -> AgentRunResult[Any]:
+ """
+ Run the agent and publish the run's ``usage``, ``run_id`` and model
name on XCom.
+
+ Every ``execute`` of this operator and its subclasses runs the agent
through here, so they
+ all report the same run metadata.
+ """
+ # A per-attempt key (the task-instance id on Airflow 3, which is
regenerated on each
+ # retry; dag/run/task/map/try on Airflow 2) ties this run to its GenAI
trace, the same
+ # join key AgentOperator uses. A hand-built context with no task
instance runs untracked.
+ ti = self._resolve_task_instance(context)
+ run_kwargs: dict[str, Any] = {}
+ if ti is not None:
+ stamp_identity_on_agent_spans(agent,
build_run_identity_attributes(ti))
+ run_kwargs["run_id"] = make_task_instance_run_key(ti)
+
+ attempt_usage = RunUsage()
+ try:
+ result = self._run_agent_within_budget(
+ context, agent, user_prompt, usage_limits=usage_limits,
run_usage=attempt_usage, **run_kwargs
+ )
+ except BaseException:
+ # A run that raises partway through (e.g. a cost_limit-triggered
UsageLimitExceeded)
+ # still reports what it spent: the capped run is the one cost
tracking most needs.
+ try:
+ log_run_usage(self.log, attempt_usage, outcome="failed")
+ except Exception:
+ self.log.warning("Failed to log partial usage for the failed
run", exc_info=True)
+ try:
+ self._push_xcom(context, "usage",
format_usage_for_xcom(attempt_usage))
+ except Exception:
+ self.log.warning("Failed to push usage XCom for the failed
run", exc_info=True)
+ if ti is not None:
+ try:
+ self._push_xcom(context, "run_id", run_kwargs["run_id"])
+ except Exception:
+ self.log.warning("Failed to push run_id XCom for the
failed run", exc_info=True)
+ raise
+
+ self._push_xcom(context, "usage", format_usage_for_xcom(attempt_usage))
+ if ti is not None:
+ self._push_xcom(context, "run_id", result.run_id)
+ model_name = ModelConfidence.from_result(result).model
+ if model_name is not None:
+ self._push_xcom(context, MODEL_NAME_XCOM_KEY, model_name)
+ return result
+
+ def _run_agent_within_budget(
+ self,
+ context: Context,
+ agent: Agent[Any, Any],
+ user_prompt: Any,
+ *,
+ usage_limits: UsageLimits | None,
+ run_usage: RunUsage,
**run_kwargs: Any,
) -> AgentRunResult[Any]:
"""
Run the agent, bounding ``usage_limits`` across task retries when the
usage budget applies.
- ``run_usage``, when given, receives this attempt's own usage, also
when the run raises
+ ``run_usage`` receives this attempt's own usage, also when the run
raises
partway through. With the usage budget, pydantic-ai is handed the
cross-attempt total
instead, so the attempt's share is credited to ``run_usage`` once the
run ends.
"""
@@ -513,8 +532,7 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
ti = context["task_instance"] if usage_limits is not None else None
self._usage_budget = self._build_usage_budget(context, usage_limits,
ti=ti)
if self._usage_budget is None:
- if run_usage is not None:
- run_kwargs["usage"] = run_usage
+ run_kwargs["usage"] = run_usage
result = self.run_agent_sync(agent, user_prompt,
usage_limits=usage_limits, **run_kwargs)
log_run_summary(self.log, result)
return result
@@ -525,8 +543,7 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
agent, user_prompt, run_usage=cumulative_usage,
usage_limits=usage_limits, **run_kwargs
)
finally:
- if run_usage is not None:
- run_usage.incr(subtract_run_usage(cumulative_usage, base))
+ run_usage.incr(subtract_run_usage(cumulative_usage, base))
log_run_summary(self.log, result, usage=attempt_usage)
self._log_cumulative_usage(cumulative_usage)
return result
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py
b/providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py
index f3a4295cbf7..91b594db37e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/plugins/model_panel.py
@@ -80,16 +80,24 @@ class ModelPanelPlugin(AirflowPlugin):
"bundle_url": get_bundle_url(_PLUGIN_PREFIX, "model.umd.cjs"),
"destination": "task_instance",
"url_route": "ai-model",
- # Only LLMOperator and AgentOperator (and their
@task.llm/@task.agent
- # TaskFlow forms) push model_name/usage; other operators would
show
- # a permanently-empty tab. `operator_name` alone covers both
the raw
+ # Only AgentOperator, LLMOperator and its subclasses (and
their TaskFlow
+ # forms) push model_name/usage; other operators would show a
+ # permanently-empty tab. `operator_name` alone covers both the
raw
# class name and the decorator's display name, since
TaskInstance falls
# back to the class name when there's no custom_operator_name.
"applies_to": {
"operator_name": [
"LLMOperator",
+ "LLMBranchOperator",
+ "LLMSQLQueryOperator",
+ "LLMSchemaCompareOperator",
+ "LLMFileAnalysisOperator",
"AgentOperator",
"@task.llm",
+ "@task.llm_branch",
+ "@task.llm_sql",
+ "@task.llm_schema_compare",
+ "@task.llm_file_analysis",
"@task.agent",
],
},
diff --git a/providers/common/ai/tests/unit/common/ai/conftest.py
b/providers/common/ai/tests/unit/common/ai/conftest.py
index c002170da34..e99710eea19 100644
--- a/providers/common/ai/tests/unit/common/ai/conftest.py
+++ b/providers/common/ai/tests/unit/common/ai/conftest.py
@@ -20,8 +20,10 @@ from unittest.mock import MagicMock
from uuid import uuid4
import pytest
+from pydantic_ai.messages import ModelResponse
from pydantic_ai.usage import RunUsage
+from airflow.providers.common.ai.utils.logging import MODEL_NAME_XCOM_KEY
from airflow.providers.common.ai.utils.usage_budget import
TaskStateStoreUsageBudget
from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
@@ -104,6 +106,61 @@ def make_mock_run_result():
return _make
[email protected]
+def reporting_task_instance():
+ """A task instance with id ``ti-1``, specced to what a run reads off it or
pushes to."""
+ return MagicMock(
+ spec=["id", "dag_id", "task_id", "run_id", "try_number", "map_index",
"xcom_push"],
+ id="ti-1",
+ dag_id="dag",
+ task_id="test",
+ run_id="run",
+ try_number=1,
+ map_index=-1,
+ )
+
+
[email protected]
+def make_reporting_agent(make_mock_run_result):
+ """Factory fixture creating a mock agent whose run spends a known usage,
as ``(agent, expected_pushes)``.
+
+ ``expected_pushes`` maps each XCom key the run must report (``usage``,
``run_id``, and the
+ model name unless ``model_name`` is ``None``) to its value, for a task
instance with id ``ti-1``.
+ """
+
+ def _make(output, *, model_name):
+ result = make_mock_run_result(output)
+ result.response = ModelResponse(parts=[], model_name=model_name)
+
+ def _run_sync(*args, **kwargs):
+ # A real run updates the ``usage=`` object it is called with,
which is what gets pushed,
+ # and reports back the ``run_id=`` it is called with.
+ kwargs["usage"].incr(RunUsage(requests=1, input_tokens=10,
output_tokens=5))
+ result.run_id = kwargs["run_id"]
+ return result
+
+ agent = MagicMock(spec=["run_sync", "instrument"])
+ agent.run_sync.side_effect = _run_sync
+ expected_pushes = {
+ "usage": {
+ "requests": 1,
+ "input_tokens": 10,
+ "output_tokens": 5,
+ "total_tokens": 15,
+ "tool_calls": 0,
+ "cache_read_tokens": 0,
+ "cache_write_tokens": 0,
+ "cost": None,
+ },
+ "run_id": "ti-1",
+ }
+ if model_name is not None:
+ expected_pushes[MODEL_NAME_XCOM_KEY] = model_name
+ return agent, expected_pushes
+
+ return _make
+
+
@pytest.fixture
def task_state_store_accessor():
"""A ``MagicMock(spec=TaskStateStoreAccessor)`` backed by a plain dict.
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
index 09f2dda4ccb..747bbd46717 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py
@@ -36,7 +36,7 @@ class TestLLMBranchDecoratedOperator:
"""The callable's return value becomes the LLM prompt, LLM output goes
through do_branch."""
downstream_enum = Enum("DownstreamTasks", {"positive": "positive",
"negative": "negative"})
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value =
make_mock_run_result(downstream_enum.positive)
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
mock_do_branch.return_value = "positive"
@@ -56,7 +56,7 @@ class TestLLMBranchDecoratedOperator:
assert result == "positive"
assert op.prompt == "Route this review"
mock_agent.run_sync.assert_called_once_with(
- "Route this review", usage_limits=None, cancellation_token=ANY
+ "Route this review", usage=ANY, usage_limits=None,
cancellation_token=ANY
)
mock_do_branch.assert_called_once()
@@ -81,7 +81,7 @@ class TestLLMBranchDecoratedOperator:
"""A non-empty Sequence[UserContent] return value is forwarded to
run_sync as-is."""
downstream_enum = Enum("DownstreamTasks", {"positive": "positive"})
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value =
make_mock_run_result(downstream_enum.positive)
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
mock_do_branch.return_value = "positive"
@@ -101,7 +101,9 @@ class TestLLMBranchDecoratedOperator:
op.execute(context={})
assert op.prompt == prompt
- mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None,
cancellation_token=ANY)
+ mock_agent.run_sync.assert_called_once_with(
+ prompt, usage=ANY, usage_limits=None, cancellation_token=ANY
+ )
@patch.object(LLMBranchOperator, "do_branch")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
@@ -111,7 +113,7 @@ class TestLLMBranchDecoratedOperator:
"""op_kwargs are resolved by the callable to build the prompt."""
downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a"})
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value =
make_mock_run_result(downstream_enum.task_a)
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py
index aa45df1458b..5903138f817 100644
---
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py
+++
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py
@@ -58,7 +58,7 @@ class TestLLMFileAnalysisDecoratedOperator:
resolved_paths=["/tmp/app.log"],
total_size_bytes=10,
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _make_mock_run_result("This is a
summary.")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -76,7 +76,7 @@ class TestLLMFileAnalysisDecoratedOperator:
assert result == "This is a summary."
assert op.prompt == "Summarize this text"
mock_agent.run_sync.assert_called_once_with(
- "prepared prompt", usage_limits=None, cancellation_token=ANY
+ "prepared prompt", usage=ANY, usage_limits=None,
cancellation_token=ANY
)
@pytest.mark.parametrize(
@@ -104,7 +104,7 @@ class TestLLMFileAnalysisDecoratedOperator:
resolved_paths=["/tmp/app.log"],
total_size_bytes=10,
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _make_mock_run_result("done")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -122,5 +122,5 @@ class TestLLMFileAnalysisDecoratedOperator:
assert op.prompt == "Summarize system logs"
mock_agent.run_sync.assert_called_once_with(
- "prepared prompt", usage_limits=None, cancellation_token=ANY
+ "prepared prompt", usage=ANY, usage_limits=None,
cancellation_token=ANY
)
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
index e6004561167..73adba3f2d3 100644
---
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
+++
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
@@ -38,7 +38,7 @@ def _make_compare_result():
def _make_mock_agent(output: SchemaCompareResult, make_mock_run_result):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result(output)
return mock_agent
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
index ecee8aac631..aed2f559586 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
@@ -43,7 +43,7 @@ class TestLLMSQLDecoratedOperator:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_calls_callable_and_uses_result_as_prompt(self,
mock_hook_cls, make_mock_run_result):
"""The user's callable return value becomes the LLM prompt."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("SELECT 1")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -56,7 +56,7 @@ class TestLLMSQLDecoratedOperator:
assert result == "SELECT 1"
assert op.prompt == "Get all users"
mock_agent.run_sync.assert_called_once_with(
- "Get all users", usage_limits=None, cancellation_token=ANY
+ "Get all users", usage=ANY, usage_limits=None,
cancellation_token=ANY
)
@pytest.mark.parametrize(
@@ -77,7 +77,7 @@ class TestLLMSQLDecoratedOperator:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_accepts_sequence_prompt(self, mock_hook_cls,
make_mock_run_result):
"""A non-empty Sequence[UserContent] return value is forwarded to
run_sync as-is."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("SELECT 1")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -91,12 +91,14 @@ class TestLLMSQLDecoratedOperator:
op.execute(context={})
assert op.prompt == prompt
- mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None,
cancellation_token=ANY)
+ mock_agent.run_sync.assert_called_once_with(
+ prompt, usage=ANY, usage_limits=None, cancellation_token=ANY
+ )
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_merges_op_kwargs_into_callable(self, mock_hook_cls,
make_mock_run_result):
"""op_kwargs are resolved by the callable to build the prompt."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("SELECT 1")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
index e6f799b581c..baa07b6e18e 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
@@ -72,7 +72,7 @@ class TestLLMBranchOperator:
def test_execute_builds_agent_with_model_settings_and_output_retries_only(
self, mock_hook_cls, mock_do_branch, make_mock_run_result
):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("task_a")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
mock_do_branch.return_value = "task_a"
@@ -97,7 +97,7 @@ class TestLLMBranchOperator:
def test_execute_single_branch(self, mock_hook_cls, mock_do_branch,
make_mock_run_result):
"""LLM returns a single enum member → do_branch receives a string."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("task_a")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
mock_do_branch.return_value = "task_a"
@@ -115,9 +115,35 @@ class TestLLMBranchOperator:
assert result == "task_a"
mock_do_branch.assert_called_once_with(ctx, "task_a")
mock_agent.run_sync.assert_called_once_with(
- "Pick a branch", usage_limits=None, cancellation_token=ANY
+ "Pick a branch", usage=ANY, usage_limits=None, run_id=ANY,
cancellation_token=ANY
)
+ @pytest.mark.parametrize("model_name", ["gpt-5", None])
+
@patch("airflow.providers.common.ai.operators.llm.stamp_identity_on_agent_spans",
autospec=True)
+ @patch.object(LLMBranchOperator, "do_branch")
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_run_usage_run_id_and_model_name_pushed_to_xcom(
+ self, mock_hook_cls, mock_do_branch, mock_stamp, make_reporting_agent,
model_name
+ ):
+ mock_agent, expected_pushes = make_reporting_agent("task_a",
model_name=model_name)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+ op = LLMBranchOperator(task_id="test", prompt="Pick a branch",
llm_conn_id="my_llm")
+ op.downstream_task_ids = {"task_a", "task_b"}
+ context = _make_context(ti_id="ti-1")
+
+ op.execute(context)
+
+ pushes = {
+ c.kwargs["key"]: c.kwargs["value"]
+ for c in context["task_instance"].xcom_push.call_args_list
+ if c.kwargs["key"] != "decision"
+ }
+ assert pushes == expected_pushes
+ mock_stamp.assert_called_once()
+ agent_arg, attrs = mock_stamp.call_args.args
+ assert agent_arg is mock_agent
+ assert attrs["airflow.task_instance.id"] == "ti-1"
+
@patch.object(LLMBranchOperator, "do_branch")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_coerces_usage_limits_dict_before_run_sync(
@@ -125,7 +151,7 @@ class TestLLMBranchOperator:
):
"""A dict ``usage_limits`` is coerced into a real ``UsageLimits``
before ``run_sync``."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("task_a")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
mock_do_branch.return_value = "task_a"
@@ -148,7 +174,7 @@ class TestLLMBranchOperator:
def test_execute_multi_branch(self, mock_hook_cls, mock_do_branch,
make_mock_run_result):
"""allow_multiple_branches=True → LLM returns list of enums →
do_branch receives list."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result(["task_a",
"task_c"])
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
mock_do_branch.return_value = ["task_a", "task_c"]
@@ -173,7 +199,7 @@ class TestLLMBranchOperator:
self, mock_hook_cls, mock_do_branch, make_mock_run_result
):
"""LLM returning an empty list fails instead of skipping every
downstream task."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result([])
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -194,7 +220,7 @@ class TestLLMBranchOperator:
def test_system_prompt_forwarded(self, mock_hook_cls, mock_do_branch,
make_mock_run_result):
"""system_prompt is passed to create_agent(instructions=...)."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("task_a")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -216,7 +242,7 @@ class TestLLMBranchOperator:
def test_downstream_task_ids_used_for_enum(self, mock_hook_cls,
mock_do_branch, make_mock_run_result):
"""The dynamic enum is built from self.downstream_task_ids."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("billing")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -245,7 +271,7 @@ class TestLLMBranchOperator:
unlucky set order; without ``sorted()`` the enum comes out reversed
and this fails.
"""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("task_a")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -268,7 +294,7 @@ class TestLLMBranchOperator:
a per-option description from: ``anyOf`` of ``{const, description}``,
not a bare ``enum``.
pydantic-ai's ``Choices`` builds it; a bare string and a
``BranchOption`` describe alike.
"""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("handle_auth")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -399,7 +425,7 @@ class TestLLMBranchOperator:
self, mock_hook_cls, mock_do_branch, make_mock_run_result
):
"""Without descriptions the schema is the bare enum it always was."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
mock_agent.run_sync.return_value = make_mock_run_result("task_a")
op = LLMBranchOperator(task_id="route", prompt="Pick",
llm_conn_id="my_llm")
@@ -441,7 +467,7 @@ class TestLLMBranchOperator:
self, mock_hook_cls, mock_do_branch, make_mock_run_result
):
"""On pydantic-ai without ``Choices`` the Enum fallback carries the
same anyOf and the pick unwraps."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = LLMBranchOperator(
task_id="route",
@@ -578,7 +604,7 @@ class TestLLMBranchOperatorUsageBudget:
usage_budget_context,
):
seed_usage_budget(task_state_store_accessor, requests=1)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("task_a")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
key_present_in_do_branch = []
@@ -677,7 +703,7 @@ class TestLLMBranchOperatorConfidenceGate:
def test_confident_pick_branches_and_records_the_decision(
self, mock_hook_cls, mock_do_branch, make_mock_run_result
):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(
make_mock_run_result,
"rerun",
@@ -718,7 +744,7 @@ class TestLLMBranchOperatorConfidenceGate:
def test_pick_below_the_bar_goes_to_review_before_branching(
self, mock_hook_cls, mock_do_branch, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(
make_mock_run_result,
"page_oncall",
@@ -763,7 +789,7 @@ class TestLLMBranchOperatorConfidenceGate:
self, mock_hook_cls, mock_do_branch, make_mock_run_result
):
"""0.52 is under page_oncall's 0.9 bar but rerun's bar is 0.6, so a
rerun at 0.65 branches."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(make_mock_run_result,
"rerun", confidence=0.65)
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = self._op(
@@ -791,7 +817,7 @@ class TestLLMBranchOperatorConfidenceGate:
self, mock_hook_cls, mock_do_branch, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
"""A branch the overrides do not name is still gated, by
min_confidence."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(make_mock_run_result,
"ignore", confidence=0.3)
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = self._op(
@@ -820,7 +846,7 @@ class TestLLMBranchOperatorConfidenceGate:
self, mock_hook_cls, mock_do_branch, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
"""A model that reports no confidence must not silently bypass a
configured bar."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(make_mock_run_result,
"rerun", model="claude-sonnet-5")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = self._op(decision_policy=DecisionPolicy(min_confidence=0.7))
@@ -840,7 +866,7 @@ class TestLLMBranchOperatorConfidenceGate:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_low_confidence_can_fail_the_task(self, mock_hook_cls,
mock_do_branch, make_mock_run_result):
"""on_uncertain="fail": the record is written first so the failure is
explained."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(make_mock_run_result,
"rerun", confidence=0.4)
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = self._op(decision_policy=DecisionPolicy(min_confidence=0.7,
on_uncertain="fail"))
@@ -859,7 +885,7 @@ class TestLLMBranchOperatorConfidenceGate:
@patch.object(LLMBranchOperator, "do_branch")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_text_model_with_a_bar_can_fail(self, mock_hook_cls,
mock_do_branch, make_mock_run_result):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(make_mock_run_result,
"rerun", model="claude-sonnet-5")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = self._op(decision_policy=DecisionPolicy(min_confidence=0.7,
on_uncertain="fail"))
@@ -874,7 +900,7 @@ class TestLLMBranchOperatorConfidenceGate:
def test_no_bar_and_no_confidence_is_todays_behaviour(
self, mock_hook_cls, mock_do_branch, make_mock_run_result
):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(make_mock_run_result,
"rerun", model="claude-sonnet-5")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = self._op()
@@ -905,7 +931,7 @@ class TestLLMBranchOperatorConfidenceGate:
def test_do_xcom_push_false_suppresses_the_record(
self, mock_hook_cls, mock_do_branch, make_mock_run_result
):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(make_mock_run_result,
"rerun", confidence=0.9)
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = self._op(decision_policy=DecisionPolicy(min_confidence=0.7),
do_xcom_push=False)
@@ -923,7 +949,7 @@ class TestLLMBranchOperatorConfidenceGate:
def test_multiple_branches_use_the_strictest_picked_bar(
self, mock_hook_cls, mock_do_branch, make_mock_run_result
):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(
make_mock_run_result, ["rerun", "page_oncall"], confidence=0.95
)
@@ -953,7 +979,7 @@ class TestLLMBranchOperatorConfidenceGate:
def test_require_approval_still_asks_at_high_confidence(
self, mock_hook_cls, mock_do_branch, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _jev_result(make_mock_run_result,
"rerun", confidence=0.99)
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = self._op(require_approval=True,
decision_policy=DecisionPolicy(min_confidence=0.5))
@@ -1089,7 +1115,7 @@ class TestLLMBranchOperatorApproval:
):
"""When require_approval=True, execute() pauses after the LLM choice,
before do_branch."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("task_a")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -1118,7 +1144,7 @@ class TestLLMBranchOperatorApproval:
):
"""With allow_multiple_branches=True the choice is deferred as a JSON
list."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result(["task_a",
"task_c"])
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -1146,7 +1172,7 @@ class TestLLMBranchOperatorApproval:
):
"""The review body lists the valid branches and the editable param is
an enum dropdown."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("task_a")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -1178,7 +1204,7 @@ class TestLLMBranchOperatorApproval:
):
"""With allow_multiple_branches the editable param is an array enum
(multi-select)."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result(["task_a"])
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -1215,7 +1241,7 @@ class TestLLMBranchOperatorApproval:
def test_execute_with_approval_reviews_a_sequence_prompt(
self, mock_hook_cls, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("task_a")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
index 74e2f8f2b59..61feeafa241 100644
---
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
+++
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py
@@ -112,7 +112,7 @@ class TestLLMFileAnalysisOperator:
mock_build_request.return_value = FileAnalysisRequest(
user_content="prepared prompt", resolved_paths=["/tmp/app.log"],
total_size_bytes=10
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _make_mock_run_result("Analysis
complete")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -141,7 +141,7 @@ class TestLLMFileAnalysisOperator:
resolved_paths=["/tmp/app.log"],
total_size_bytes=10,
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _make_mock_run_result("Analysis
complete")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -166,9 +166,31 @@ class TestLLMFileAnalysisOperator:
sample_rows=10,
)
mock_agent.run_sync.assert_called_once_with(
- "prepared prompt", usage_limits=None, cancellation_token=ANY
+ "prepared prompt", usage=ANY, usage_limits=None,
cancellation_token=ANY
)
+ @pytest.mark.parametrize("model_name", ["gpt-5", None])
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ @patch(
+
"airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request",
autospec=True
+ )
+ def test_run_usage_run_id_and_model_name_pushed_to_xcom(
+ self, mock_build_request, mock_hook_cls, make_reporting_agent,
reporting_task_instance, model_name
+ ):
+ mock_build_request.return_value = FileAnalysisRequest(
+ user_content="prepared prompt", resolved_paths=["/tmp/app.log"],
total_size_bytes=10
+ )
+ mock_agent, expected_pushes = make_reporting_agent("Analysis
complete", model_name=model_name)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+ op = LLMFileAnalysisOperator(
+ task_id="test", prompt="Summarize the file", llm_conn_id="my_llm",
file_path="/tmp/app.log"
+ )
+ ti = reporting_task_instance
+
+ op.execute(context={"task_instance": ti})
+
+ assert {c.kwargs["key"]: c.kwargs["value"] for c in
ti.xcom_push.call_args_list} == expected_pushes
+
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
@patch(
"airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request",
autospec=True
@@ -182,7 +204,7 @@ class TestLLMFileAnalysisOperator:
resolved_paths=["/tmp/app.log"],
total_size_bytes=10,
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _make_mock_run_result("Analysis
complete")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -209,7 +231,7 @@ class TestLLMFileAnalysisOperator:
resolved_paths=["/tmp/app.log"],
total_size_bytes=10,
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value =
_make_mock_run_result(Summary(findings=["error spike"]))
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -236,7 +258,7 @@ class TestLLMFileAnalysisOperator:
resolved_paths=["/tmp/app.log"],
total_size_bytes=10,
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value =
_make_mock_run_result(Summary(findings=["error spike"]))
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -335,7 +357,7 @@ class TestLLMFileAnalysisOperatorUsageBudget:
mock_build_request.return_value = FileAnalysisRequest(
user_content="prepared prompt", resolved_paths=["/tmp/app.log"],
total_size_bytes=10
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _make_mock_run_result("analysis")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = LLMFileAnalysisOperator(
@@ -369,7 +391,7 @@ class TestLLMFileAnalysisOperatorApproval:
resolved_paths=["/tmp/app.log"],
total_size_bytes=10,
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _make_mock_run_result("LLM
response")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -403,7 +425,7 @@ class TestLLMFileAnalysisOperatorApproval:
resolved_paths=["/tmp/app.log"],
total_size_bytes=10,
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value =
_make_mock_run_result(Summary(findings=["error spike"]))
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -479,7 +501,7 @@ class TestLLMFileAnalysisOperatorApproval:
resolved_paths=["/tmp/app.log"],
total_size_bytes=10,
)
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = _make_mock_run_result("output")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
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 fa9c7a08384..51be164cc14 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
@@ -296,10 +296,33 @@ class TestLLMSchemaCompareOperator:
retries=None,
)
mock_agent.run_sync.assert_called_once_with(
- "user_prompt", usage_limits=None, cancellation_token=mock.ANY
+ "user_prompt", usage=mock.ANY, usage_limits=None,
cancellation_token=mock.ANY
)
assert result == {"compatible": True, "mismatches": [], "summary":
"All good"}
+ @pytest.mark.parametrize("model_name", ["gpt-5", None])
+ @mock.patch.object(LLMSchemaCompareOperator, "_build_schema_context",
return_value="schema_context")
+ def test_run_usage_run_id_and_model_name_pushed_to_xcom(
+ self, mock_build_schema_context, make_reporting_agent,
reporting_task_instance, model_name
+ ):
+ mock_agent, expected_pushes = make_reporting_agent(
+ SchemaCompareResult(compatible=True, mismatches=[], summary="All
good"), model_name=model_name
+ )
+ op = LLMSchemaCompareOperator(
+ task_id="test",
+ prompt="Compare",
+ llm_conn_id="llm_conn",
+ db_conn_ids=["a", "b"],
+ table_names=["t"],
+ )
+ op.llm_hook = MagicMock(spec=PydanticAIHook)
+ op.llm_hook.create_agent.return_value = mock_agent
+ ti = reporting_task_instance
+
+ op.execute(context={"task_instance": ti})
+
+ assert {c.kwargs["key"]: c.kwargs["value"] for c in
ti.xcom_push.call_args_list} == expected_pushes
+
@mock.patch(
"airflow.providers.common.ai.operators.llm_schema_compare.LLMSchemaCompareOperator._build_schema_context"
)
@@ -415,6 +438,7 @@ class TestLLMSchemaCompareOperator:
mock_agent.run_sync.assert_called_once_with(
"Compare S3 Parquet schema against the Postgres table and flag
breaking changes",
+ usage=mock.ANY,
usage_limits=None,
cancellation_token=mock.ANY,
)
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
index cc3fd903ee8..8e56ceb6d5a 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
@@ -53,7 +53,7 @@ else:
def _make_mock_agent(output: str, make_mock_run_result):
"""Create a mock agent that returns the given output string."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result(output)
return mock_agent
@@ -252,9 +252,23 @@ class TestLLMSQLQueryOperator:
assert result == "SELECT id, name FROM users WHERE active = true"
mock_agent.run_sync.assert_called_once_with(
- "Get active users", usage_limits=None, cancellation_token=ANY
+ "Get active users", usage=ANY, usage_limits=None, run_id=ANY,
cancellation_token=ANY
)
+ @pytest.mark.parametrize("model_name", ["gpt-5", None])
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_run_usage_run_id_and_model_name_pushed_to_xcom(
+ self, mock_hook_cls, make_reporting_agent, reporting_task_instance,
model_name
+ ):
+ mock_agent, expected_pushes = make_reporting_agent("SELECT 1",
model_name=model_name)
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+ op = LLMSQLQueryOperator(task_id="test", prompt="Get one",
llm_conn_id="my_llm")
+ ti = reporting_task_instance
+
+ op.execute(context={"task_instance": ti})
+
+ assert {c.kwargs["key"]: c.kwargs["value"] for c in
ti.xcom_push.call_args_list} == expected_pushes
+
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_coerces_usage_limits_dict_before_run_sync(
self, mock_hook_cls, make_mock_run_result, usage_budget_context
diff --git
a/providers/common/ai/tests/unit/common/ai/plugins/test_model_panel.py
b/providers/common/ai/tests/unit/common/ai/plugins/test_model_panel.py
index 94c0363da14..9b339cec15d 100644
--- a/providers/common/ai/tests/unit/common/ai/plugins/test_model_panel.py
+++ b/providers/common/ai/tests/unit/common/ai/plugins/test_model_panel.py
@@ -17,6 +17,8 @@
from __future__ import annotations
import copy
+import importlib
+import pkgutil
import pytest
@@ -27,8 +29,7 @@ if not AIRFLOW_V_3_1_PLUS:
pytest.skip("The AI Model panel is only compatible with Airflow >= 3.1.0",
allow_module_level=True)
from airflow import plugins_manager
-from airflow.providers.common.ai.decorators.agent import
_AgentDecoratedOperator
-from airflow.providers.common.ai.decorators.llm import _LLMDecoratedOperator
+from airflow.providers.common.ai import decorators, operators
from airflow.providers.common.ai.operators.agent import AgentOperator
from airflow.providers.common.ai.operators.llm import LLMOperator
from airflow.providers.common.ai.plugins.model_panel import ModelPanelPlugin
@@ -38,6 +39,23 @@ requires_3_4 = pytest.mark.skipif(
)
+def get_in_tree_model_operator_names() -> set[str]:
+ """Return the operator name a task instance records for each in-tree class
that publishes the model name."""
+ for package in (operators, decorators):
+ for module in pkgutil.iter_modules(package.__path__):
+ importlib.import_module(f"{package.__name__}.{module.name}")
+
+ names = set()
+ pending: list[type] = [LLMOperator, AgentOperator]
+ while pending:
+ cls = pending.pop()
+ pending.extend(cls.__subclasses__())
+ # Subclasses defined by other test modules are not shipped, so only
the provider's own count.
+ if cls.__module__.startswith("airflow.providers.common.ai."):
+ names.add(getattr(cls, "custom_operator_name", None) or
cls.__name__)
+ return names
+
+
class TestModelPanelPlugin:
def test_plugin_name(self):
assert ModelPanelPlugin.name == "ai_model_panel"
@@ -64,16 +82,10 @@ class TestModelPanelPlugin:
@requires_3_4
def test_applies_to_scopes_to_operators_that_publish_model_name(self):
- """Tied to the actual operator/decorator names a task instance would
record,
- so a rename of either doesn't silently desync the tab's scoping from
reality."""
+ """Built from the in-tree subclasses, so a new LLMOperator subclass or
a rename can't silently
+ desync the tab's scoping from the operators that publish the model
name."""
app = ModelPanelPlugin.react_apps[0]
- operator_names = app["applies_to"]["operator_name"]
- assert set(operator_names) == {
- LLMOperator.__name__,
- AgentOperator.__name__,
- _LLMDecoratedOperator.custom_operator_name,
- _AgentDecoratedOperator.custom_operator_name,
- }
+ assert set(app["applies_to"]["operator_name"]) ==
get_in_tree_model_operator_names()
@requires_3_4
def test_react_app_survives_core_applies_to_validation(self):
