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 48a0f8e22be Push LLMOperator's token usage and cost to XCom (#74355)
48a0f8e22be is described below
commit 48a0f8e22be2deb6ee2ee41a4e31763d8d0cca6b
Author: Pierre Jeambrun <[email protected]>
AuthorDate: Mon Oct 19 12:19:16 2026 +0200
Push LLMOperator's token usage and cost to XCom (#74355)
* Push LLMOperator's token usage and cost to XCom
AgentOperator already exposes request/token/cost counts and a run id on
XCom and stamps the task-instance identity on its GenAI spans; LLMOperator
computed the same usage for its log summary but never published it, so
there was no way to see what an LLMOperator run cost, or join it to its
trace, without reading task logs. A run stopped by a cost limit is the one
where this matters most, so the usage is pushed on a failed run too.
* Log LLMOperator's partial usage when a run fails
With do_xcom_push=False, the usage XCom was the only record of what a
failed run spent, so a run stopped by a cost limit left no trace of its
spend anywhere. Logging it matches what AgentOperator already does.
* Cover LLMOperator's untracked failure path and update the AI Model tab
docs
The feature map still said only @task.agent tasks show token counts and
cost in the AI Model tab, which no longer holds now that LLMOperator
pushes the same usage XCom.
* Keep LLMOperator's failed-run usage logging best-effort
Everything that runs after a run raises is reporting only and must never
replace the run's own exception, the same rule AgentOperator follows for
its failed-run report.
---
providers/common/ai/docs/features.rst | 2 +-
providers/common/ai/docs/observability.rst | 43 ++--
providers/common/ai/docs/operators/llm.rst | 8 +
.../airflow/providers/common/ai/operators/llm.py | 119 +++++++--
.../ai/tests/unit/common/ai/decorators/test_llm.py | 10 +-
.../ai/tests/unit/common/ai/operators/test_llm.py | 279 ++++++++++++++++++---
6 files changed, 388 insertions(+), 73 deletions(-)
diff --git a/providers/common/ai/docs/features.rst
b/providers/common/ai/docs/features.rst
index 94fc60c2654..e7f7149c191 100644
--- a/providers/common/ai/docs/features.rst
+++ b/providers/common/ai/docs/features.rst
@@ -299,7 +299,7 @@ See what ran
- 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 for ``@task.agent`` the token counts and an estimated cost.
+ tasks, 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 b7de5af0681..ad5bbb769c8 100644
--- a/providers/common/ai/docs/observability.rst
+++ b/providers/common/ai/docs/observability.rst
@@ -49,10 +49,11 @@ 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
- ``@task.agent``) additionally stamps the task-instance identity on every
GenAI
- span it emits: the five keys core tracing already puts on the task span
- (``airflow.dag_id``, ``airflow.task_id``, ``airflow.dag_run.run_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
+ 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
``airflow.task_instance.id`` as the per-attempt run join key. So a span is
filterable by dag, task, run, attempt, or map index directly, without walking
@@ -62,25 +63,27 @@ 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`` run, the task-instance id (unique
- per attempt, since Airflow regenerates it on each retry) is passed to
+* **Run join key.** For an ``AgentOperator`` or ``LLMOperator`` run (and their
+ ``@task.agent`` / ``@task.llm`` 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
token usage, on XCom under the ``run_id`` and ``usage`` keys. ``usage`` is
this attempt's own usage -- not the cross-attempt cumulative total described
under ``usage_limits`` in :ref:`howto/operator:agent` -- and it is pushed on
a failed attempt too, so a downstream ``all_done`` task or failure callback
- can read what the last attempt spent. With ``durable=True``, steps an
- attempt replays from the cache are not part of it. XCom is cleared at the
start of every
+ can read what the last attempt spent. XCom is cleared at the start of every
attempt, so only the most recent attempt's value survives, not each
historical attempt's. A downstream task can then reference the
- run (``ti.xcom_pull(task_ids="my_agent", key="run_id")``) and a trace
+ run (``ti.xcom_pull(task_ids="my_task", key="run_id")``) and a trace
backend can join a task's output to its agent trace without parsing logs.
- With ``enable_hitl_review`` the ``run_id`` and ``usage`` reflect the initial
- model run, not the human-feedback regenerations. A run that resumes after a
- tool approval (see :doc:`tool_approval`) continues as
- ``<task-instance id>-resumed-<n>`` after pause ``n``, which is the ``run_id``
- the operator pushes; ``usage`` covers every part.
+ ``AgentOperator`` additionally: with ``durable=True``, steps an attempt
replays
+ from the cache are not part of ``usage``; with ``enable_hitl_review`` the
+ ``run_id`` and ``usage`` reflect the initial model run, not the
human-feedback
+ regenerations; and a run that resumes after a tool approval (see
+ :doc:`tool_approval`) continues as ``<task-instance id>-resumed-<n>`` after
+ pause ``n``, which is the ``run_id`` the operator pushes, with ``usage``
+ covering every part.
Airflow 2 has no task-instance id, so there the key is
``<dag_id>/<run_id>/<task_id>/<map_index>/<try_number>``, and spans carry
the five
identity keys without ``airflow.task_instance.id``.
@@ -97,12 +100,12 @@ How it works
``super().execute()`` -- ``LLMBranchOperator``, ``LLMSQLQueryOperator``,
``LLMSchemaCompareOperator``, ``LLMFileAnalysisOperator`` -- don't publish
it yet;
this is upcoming work.
-* **Scope.** The ``run_id`` / ``usage`` XComs come only from ``AgentOperator``
and
- ``@task.agent``, and so do the ``airflow.*`` identity attributes, apart from
a Strands,
- ADK or OpenAI Agents SDK agent run inside ``agent_framework_tracing`` (see
below). ``LLMOperator`` /
- ``@task.llm`` additionally pushes the model name above. The other LLM
- operators still emit GenAI spans correlated to the task span by nesting, but
- without the identity attributes, the run join key, or the model name.
+* **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.
* .. _howto/observability:model-tab:
**Model tab.** Separately from tracing, the provider adds an "AI Model" tab
to the
diff --git a/providers/common/ai/docs/operators/llm.rst
b/providers/common/ai/docs/operators/llm.rst
index aaecd73e7b9..4e24b6f2a1f 100644
--- a/providers/common/ai/docs/operators/llm.rst
+++ b/providers/common/ai/docs/operators/llm.rst
@@ -315,3 +315,11 @@ After each LLM call, the operator logs a summary with
model name, token usage,
and request count at INFO level. At DEBUG level, the LLM output is also logged
(truncated to 500 characters). See :ref:`AgentOperator logging
<howto/operator:agent>`
for details on the log format.
+
+The same request/token/cost counts are also pushed to XCom under the ``usage``
+key, the same shape ``AgentOperator`` uses. Alongside it, the operator pushes a
+``run_id`` XCom and stamps the task's identity on the run's GenAI spans, the
+same way ``AgentOperator`` does -- see :doc:`../observability` for how the join
+key and identity attributes work. Both XComs are pushed on a failed run too, so
+a downstream ``all_done`` task or failure callback can read what was spent, and
+which run, before it raised.
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 4913dcfad46..459479cfc6e 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
@@ -24,11 +24,17 @@ from functools import cached_property
from typing import TYPE_CHECKING, Any, ClassVar, Literal
from pydantic import BaseModel
+from pydantic_ai.usage import RunUsage
from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
from airflow.providers.common.ai.mixins.approval import LLMApprovalMixin,
normalize_assigned_users
from airflow.providers.common.ai.mixins.cancellable_run import
CancellableAgentRunMixin
from airflow.providers.common.ai.mixins.usage_budget import UsageBudgetMixin
+from airflow.providers.common.ai.observability import (
+ build_run_identity_attributes,
+ make_task_instance_run_key,
+ stamp_identity_on_agent_spans,
+)
from airflow.providers.common.ai.policies.decision import DecisionPolicy
from airflow.providers.common.ai.utils.decision import (
DECISION_XCOM_KEY,
@@ -44,9 +50,15 @@ from airflow.providers.common.ai.utils.decision import (
timed_out_record,
validate_decision_policy,
)
-from airflow.providers.common.ai.utils.logging import MODEL_NAME_XCOM_KEY,
log_run_summary
+from airflow.providers.common.ai.utils.logging import (
+ MODEL_NAME_XCOM_KEY,
+ format_usage_for_xcom,
+ log_run_summary,
+ log_run_usage,
+)
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.ai.utils.usage_budget import copy_run_usage,
subtract_run_usage
from airflow.providers.common.compat.notifier import BaseNotifier
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, BaseOperator
from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_1_PLUS
@@ -102,6 +114,16 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
with a fresh budget. This applies to subclasses that override ``execute``
or
``execute_complete``.
+ Alongside the returned output, the run's ``run_id`` and token ``usage``
+ are pushed to XCom under the ``run_id`` and ``usage`` keys, the same shape
+ :class:`~airflow.providers.common.ai.operators.agent.AgentOperator` uses,
+ so a downstream task can reference the run and what it cost. ``usage`` is
this
+ attempt's own usage, not the cross-attempt cumulative total described under
+ ``usage_limits`` below. Both are pushed on a failed run too, so a
downstream
+ ``all_done`` task or failure callback can read what was spent, and which
run,
+ before it raised. The ``run_id`` also ties the task to its GenAI trace
(see the
+ provider's observability docs).
+
:param prompt: The prompt to send to the LLM.
:param llm_conn_id: Connection ID for the LLM provider.
:param model_id: Model identifier (e.g. ``"openai:gpt-5"``).
@@ -344,9 +366,43 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
usage_limits = coerce_usage_limits(self.usage_limits)
agent = self._create_agent(output_type=self.output_type,
instructions=self.system_prompt)
- result = self._run_llm(context, agent, self.prompt,
usage_limits=usage_limits)
+
+ # 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
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)
@@ -435,22 +491,42 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
)
def _run_llm(
- self, context: Context, agent: Agent[Any, Any], user_prompt: Any, *,
usage_limits: UsageLimits | None
+ self,
+ context: Context,
+ agent: Agent[Any, Any],
+ user_prompt: Any,
+ *,
+ usage_limits: UsageLimits | None,
+ run_usage: RunUsage | None = None,
+ **run_kwargs: Any,
) -> AgentRunResult[Any]:
- """Run the agent, bounding ``usage_limits`` across task retries when
the usage budget applies."""
+ """
+ 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
+ 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.
+ """
# Contexts built without a ``task_instance`` stay valid for runs that
have no ``usage_limits``.
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:
- result = self.run_agent_sync(agent, user_prompt,
usage_limits=usage_limits)
+ if run_usage is not None:
+ 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
- run_usage = self._usage_budget.load()
- result, attempt_usage = self._run_tracked(
- agent, user_prompt, run_usage=run_usage, usage_limits=usage_limits
- )
+ cumulative_usage = self._usage_budget.load()
+ base = copy_run_usage(cumulative_usage)
+ try:
+ result, attempt_usage = self._run_tracked(
+ 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))
log_run_summary(self.log, result, usage=attempt_usage)
- self._log_cumulative_usage(run_usage)
+ self._log_cumulative_usage(cumulative_usage)
return result
def _log_review(self, review: ReviewReason, confidence: float | None,
threshold: float | None) -> None:
@@ -486,21 +562,30 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
self._push_decision(context, timed_out_record(decision))
raise
+ def _resolve_task_instance(self, context: Context) -> Any:
+ """
+ Return ``context["task_instance"]`` if it can push XComs, else
``None``.
+
+ A hand-built context (old tests, custom runners) may carry no task
instance, ``None``, or a
+ dict stand-in, none of which has the attributes the run identity is
built from either.
+ """
+ try:
+ ti = context["task_instance"]
+ except (KeyError, TypeError):
+ return None
+ return ti if callable(getattr(ti, "xcom_push", None)) else None
+
def _push_xcom(self, context: Context, key: str, value: Any) -> None:
"""Push ``value`` under ``key`` on XCom, honoring ``do_xcom_push`` and
a missing task instance."""
if not self.do_xcom_push:
return
- try:
- ti = context["task_instance"]
- except (KeyError, TypeError):
- ti = None
- push = getattr(ti, "xcom_push", None)
- if not callable(push):
+ ti = self._resolve_task_instance(context)
+ if ti is None:
# A hand-built context (a dict, or no task instance at all) has
nowhere to push to; the
# value is inspection output, so the run goes on without it.
self.log.warning("No task instance in the context; %r was not
pushed to XCom.", key)
return
- push(key=key, value=value)
+ ti.xcom_push(key=key, value=value)
def _push_decision(self, context: Context, record: dict[str, Any]) -> None:
"""Expose what the model proposed, its confidence, and what the gate
decided, on XCom."""
diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py
index 1e1e656d57b..bc1e4e30d4d 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py
@@ -46,7 +46,7 @@ class TestLLMDecoratedOperator:
assert result == "This is a summary."
assert op.prompt == "Summarize this text"
mock_agent.run_sync.assert_called_once_with(
- "Summarize this text", usage_limits=None, cancellation_token=ANY
+ "Summarize this text", usage=ANY, usage_limits=None,
cancellation_token=ANY
)
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
@@ -100,7 +100,9 @@ class TestLLMDecoratedOperator:
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
+ )
@pytest.mark.skipif(not AIRFLOW_V_3_1_PLUS, reason="require_approval needs
Airflow >= 3.1.0")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
@@ -123,7 +125,7 @@ class TestLLMDecoratedOperator:
@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("done")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -140,5 +142,5 @@ class TestLLMDecoratedOperator:
assert op.prompt == "Summarize quantum computing"
mock_agent.run_sync.assert_called_once_with(
- "Summarize quantum computing", usage_limits=None,
cancellation_token=ANY
+ "Summarize quantum computing", usage=ANY, usage_limits=None,
run_id=ANY, cancellation_token=ANY
)
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 d5d92924904..862b2e119a0 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
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import contextlib
import logging
from datetime import timedelta
from decimal import Decimal
@@ -179,16 +180,20 @@ class TestLLMOperator:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_returns_string_output(self, mock_hook_cls,
make_mock_run_result):
"""Default output_type=str returns the LLM string directly."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("Paris is the
capital of France.")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = LLMOperator(task_id="test", prompt="What is the capital of
France?", llm_conn_id="my_llm")
- result = op.execute(context=MagicMock())
+ result = op.execute(context=_make_context(ti_id="ti-1"))
assert result == "Paris is the capital of France."
mock_agent.run_sync.assert_called_once_with(
- "What is the capital of France?", usage_limits=None,
cancellation_token=ANY
+ "What is the capital of France?",
+ usage=ANY,
+ usage_limits=None,
+ run_id="ti-1",
+ cancellation_token=ANY,
)
mock_hook_cls.get_hook.return_value.create_agent.assert_called_once_with(
output_type=str, instructions="", model_settings=None, retries=None
@@ -215,12 +220,56 @@ class TestLLMOperator:
)
op.execute(context=usage_budget_context)
- # ``usage=`` is the cross-attempt usage budget, which only exists on
Airflow >= 3.3.
- expected_usage = {"usage": ANY} if AIRFLOW_V_3_3_PLUS else {}
mock_agent.run_sync.assert_called_once_with(
- "Summarize", usage_limits=limits, cancellation_token=ANY,
**expected_usage
+ "Summarize", usage=ANY, usage_limits=limits, cancellation_token=ANY
)
+
@patch("airflow.providers.common.ai.operators.llm.stamp_identity_on_agent_spans",
autospec=True)
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_execute_stamps_identity_on_agent_spans(self, mock_hook_cls,
mock_stamp, make_mock_run_result):
+ """execute() derives the identity from the task instance and stamps it
on the agent's spans."""
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ mock_agent.run_sync.return_value = make_mock_run_result("ok")
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMOperator(task_id="test", prompt="p", llm_conn_id="my_llm")
+ op.execute(context=_make_context(ti_id="ti-9"))
+
+ 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-9"
+
+
@patch("airflow.providers.common.ai.operators.llm.stamp_identity_on_agent_spans",
autospec=True)
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_execute_without_task_instance_skips_identity_and_run_id(
+ self, mock_hook_cls, mock_stamp, make_mock_run_result
+ ):
+ """A hand-built context with no task instance runs untracked: no
identity stamp, no run_id."""
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.return_value = make_mock_run_result("ok")
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMOperator(task_id="test", prompt="p", llm_conn_id="my_llm")
+ result = op.execute(context={})
+
+ assert result == "ok"
+ mock_stamp.assert_not_called()
+ mock_agent.run_sync.assert_called_once_with("p", usage=ANY,
usage_limits=None, cancellation_token=ANY)
+
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_failed_run_without_task_instance_reports_no_run_id(self,
mock_hook_cls, caplog):
+ """A failed run with no task instance has no run id to report, and
still raises its own exception."""
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.side_effect = UsageLimitExceeded("boom")
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMOperator(task_id="test", prompt="p", llm_conn_id="my_llm")
+ with caplog.at_level(logging.WARNING),
pytest.raises(UsageLimitExceeded):
+ op.execute(context={})
+
+ assert caplog.messages == ["No task instance in the context; 'usage'
was not pushed to XCom."]
+
@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
@@ -297,10 +346,21 @@ class TestLLMOperator:
FunctionModel(_build_priced_response), **kw
)
+ ti = MagicMock(
+ spec=["id", "max_tries", "xcom_push", "dag_id", "task_id",
"run_id", "try_number", "map_index"]
+ )
+ ti.id = uuid4()
+ ti.max_tries = 0
+ usage_budget_context["task_instance"] = ti
+
op = LLMOperator(task_id="test", prompt="Summarize",
llm_conn_id="my_llm", **cap_kwargs)
with pytest.raises(UsageLimitExceeded, match=r"cost_limit.*0\.05"):
op.execute(context=usage_budget_context)
+ pushes = {c.kwargs["key"]: c.kwargs["value"] for c in
ti.xcom_push.call_args_list}
+ assert pushes["usage"]["requests"] == 1
+ assert pushes["usage"]["cost"] == str(PRICED_COST)
+
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_completes_when_cost_stays_under_cap(self, mock_hook_cls,
usage_budget_context):
"""A real run costing less than the configured ``cost_limit``
completes and returns the model output."""
@@ -335,7 +395,7 @@ class TestLLMOperator:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_structured_output_with_all_params(self, mock_hook_cls,
make_mock_run_result):
"""Structured output returns the Pydantic instance unchanged so
downstream tasks keep the type."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value =
make_mock_run_result(Entities(names=["Alice", "Bob"]))
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -366,7 +426,7 @@ class TestLLMOperator:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_forwards_fallback_conn_ids_to_hook(self, mock_hook_cls,
make_mock_run_result):
"""``fallback_conn_ids`` on the operator overrides the connection's
own extra field."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("ok")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -385,7 +445,7 @@ class TestLLMOperator:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_forwards_empty_fallback_conn_ids_to_hook(self,
mock_hook_cls, make_mock_run_result):
"""An explicit ``[]`` disables a chain configured on the connection,
not just an override."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("ok")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -412,7 +472,7 @@ class TestLLMOperator:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_serialize_output_returns_dict(self, mock_hook_cls,
make_mock_run_result):
"""serialize_output=True dumps the BaseModel to a dict on the wire."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value =
make_mock_run_result(Entities(names=["A", "B"]))
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -452,7 +512,7 @@ class TestLLMOperatorUsageBudget:
raise error
return make_mock_run_result("ok")
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.side_effect = run_sync
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
return mock_agent
@@ -465,7 +525,7 @@ class TestLLMOperatorUsageBudget:
self._make_operator().execute(context=usage_budget_context)
- mock_agent.run_sync.assert_called_once_with("p", usage_limits=None,
cancellation_token=ANY)
+ mock_agent.run_sync.assert_called_once_with("p", usage=ANY,
usage_limits=None, cancellation_token=ANY)
task_state_store_accessor.get.assert_not_called()
task_state_store_accessor.set.assert_not_called()
task_state_store_accessor.delete.assert_not_called()
@@ -480,7 +540,8 @@ class TestLLMOperatorUsageBudget:
)
_, kwargs = mock_agent.run_sync.call_args
- assert "usage" not in kwargs
+ # A fresh count for this attempt, not one loaded from the task state
store.
+ assert kwargs["usage"] == RunUsage()
@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="usage budget needs
Airflow >= 3.3")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
@@ -523,6 +584,30 @@ class TestLLMOperatorUsageBudget:
assert
task_state_store_accessor.get(USAGE_BUDGET_KEY)["usage"]["requests"] == 2
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="usage budget needs
Airflow >= 3.3")
+ @pytest.mark.parametrize("error", [None, RuntimeError("boom")],
ids=["success", "failure"])
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_usage_xcom_is_this_attempts_usage(
+ self, mock_hook_cls, make_mock_run_result, task_state_store_accessor,
seed_usage_budget, error
+ ):
+ """The ``usage`` XCom holds what this attempt spent, not the
cross-attempt total."""
+ seed_usage_budget(task_state_store_accessor, requests=2,
input_tokens=20)
+ self._stub_agent(mock_hook_cls, make_mock_run_result, requests_used=1,
error=error)
+ ti = MagicMock(
+ spec=["id", "max_tries", "xcom_push", "dag_id", "task_id",
"run_id", "try_number", "map_index"]
+ )
+ ti.id = uuid4()
+ ti.max_tries = 0
+ context = {"task_instance": ti, "task_state_store":
task_state_store_accessor}
+ op = self._make_operator(usage_limits=UsageLimits(request_limit=10))
+
+ with contextlib.nullcontext() if error is None else
pytest.raises(RuntimeError, match="boom"):
+ op.execute(context=context)
+
+ pushes = {c.kwargs["key"]: c.kwargs["value"] for c in
ti.xcom_push.call_args_list}
+ assert pushes["usage"]["requests"] == 1
+ assert pushes["usage"]["input_tokens"] == 10
+
@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="usage budget needs
Airflow >= 3.3")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_successful_run_clears_the_budget(
@@ -567,8 +652,8 @@ class TestLLMOperatorUsageBudget:
class TestLLMOperatorConfidenceGate:
"""decision_policy on a structured output, and the decision XCom."""
- def _result(self, make_mock_run_result, output, details):
- result = make_mock_run_result(output)
+ def _result(self, make_mock_run_result, output, details, *, cost=None):
+ result = make_mock_run_result(output, cost=cost)
result.response = ModelResponse(parts=[], model_name="jev-1.13.0",
provider_details=details)
return result
@@ -597,13 +682,33 @@ class TestLLMOperatorConfidenceGate:
assert "No task instance in the context; 'decision' was not pushed to
XCom." in caplog.messages
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
- def test_resolved_model_name_pushed_to_xcom(self, mock_hook_cls,
make_mock_run_result):
- """The model that actually answered is exposed on its own namespaced
XCom key."""
- mock_agent = MagicMock(spec=["run_sync"])
- mock_agent.run_sync.return_value = self._result(make_mock_run_result,
Summary(text="t"), None)
+ def test_resolved_model_name_and_usage_pushed_to_xcom(self, mock_hook_cls,
make_mock_run_result):
+ """The model that actually answered, the run id, and the run's
usage/cost are exposed on XCom."""
+ cost = Decimal("0.0042")
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ result = self._result(make_mock_run_result, Summary(text="t"), None,
cost=cost)
+ result.run_id = "ti-9"
+
+ def _run_sync(*args, **kwargs):
+ # A real run updates the ``usage=`` object it is called with,
which is what the
+ # operator pushes, rather than the result's own ``.usage``.
+ kwargs["usage"].incr(
+ RunUsage(
+ requests=1,
+ input_tokens=120,
+ output_tokens=45,
+ tool_calls=2,
+ cache_read_tokens=5,
+ cache_write_tokens=3,
+ cost=cost,
+ )
+ )
+ return result
+
+ mock_agent.run_sync.side_effect = _run_sync
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
output_type=Summary)
- context = MagicMock(spec=dict)
+ context = _make_context(ti_id="ti-9")
op.execute(context)
@@ -611,6 +716,118 @@ class TestLLMOperatorConfidenceGate:
c.kwargs["key"]: c.kwargs["value"] for c in
context["task_instance"].xcom_push.call_args_list
}
assert pushes[MODEL_NAME_XCOM_KEY] == "jev-1.13.0"
+ assert pushes["run_id"] == "ti-9"
+ assert pushes["usage"] == {
+ "requests": 1,
+ "input_tokens": 120,
+ "output_tokens": 45,
+ "total_tokens": 165,
+ "tool_calls": 2,
+ "cache_read_tokens": 5,
+ "cache_write_tokens": 3,
+ "cost": str(cost),
+ }
+
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_failure_pushes_partial_usage_to_xcom(self, mock_hook_cls):
+ """A run that raises before finishing (e.g. a cost_limit) still
reports the run id and usage incurred."""
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+
+ def _run_sync(*args, **kwargs):
+ kwargs["usage"].incr(RunUsage(requests=1, input_tokens=10,
cost=Decimal("0.1")))
+ raise UsageLimitExceeded("boom")
+
+ mock_agent.run_sync.side_effect = _run_sync
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c")
+ context = _make_context(ti_id="ti-1")
+ with pytest.raises(UsageLimitExceeded):
+ op.execute(context)
+
+ pushes = {
+ c.kwargs["key"]: c.kwargs["value"] for c in
context["task_instance"].xcom_push.call_args_list
+ }
+ assert pushes["run_id"] == "ti-1"
+ assert pushes["usage"] == {
+ "requests": 1,
+ "input_tokens": 10,
+ "output_tokens": 0,
+ "total_tokens": 10,
+ "tool_calls": 0,
+ "cache_read_tokens": 0,
+ "cache_write_tokens": 0,
+ "cost": "0.1",
+ }
+
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_failure_logs_partial_usage_without_xcom(self, mock_hook_cls,
caplog):
+ """The failed run's usage reaches the task log even with
``do_xcom_push=False``."""
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+
+ def _run_sync(*args, **kwargs):
+ kwargs["usage"].incr(RunUsage(requests=1, input_tokens=10,
cost=Decimal("0.1")))
+ raise UsageLimitExceeded("boom")
+
+ mock_agent.run_sync.side_effect = _run_sync
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c",
do_xcom_push=False)
+ context = _make_context(ti_id="ti-1")
+ with caplog.at_level(logging.INFO), pytest.raises(UsageLimitExceeded):
+ op.execute(context)
+
+ assert (
+ "LLM run failed: requests=1, tool_calls=0, input_tokens=10,
output_tokens=0, total_tokens=10"
+ in caplog.messages
+ )
+ assert "LLM run cost: $0.1 (USD, best-effort)" in caplog.messages
+ context["task_instance"].xcom_push.assert_not_called()
+
+ @patch("airflow.providers.common.ai.operators.llm.log_run_usage",
autospec=True)
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def
test_failed_usage_log_while_reporting_does_not_mask_the_original_exception(
+ self, mock_hook_cls, mock_log_run_usage
+ ):
+ """A failure logging the failed run's usage must not replace the run's
own exception."""
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+ mock_agent.run_sync.side_effect = UsageLimitExceeded("boom")
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+ mock_log_run_usage.side_effect = RuntimeError("log down")
+
+ context = _make_context(ti_id="ti-1")
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c")
+ with pytest.raises(UsageLimitExceeded):
+ op.execute(context)
+
+ mock_log_run_usage.assert_called_once()
+ assert {c.kwargs["key"] for c in
context["task_instance"].xcom_push.call_args_list} == {
+ "usage",
+ "run_id",
+ }
+
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def
test_failed_xcom_push_while_reporting_does_not_mask_the_original_exception(self,
mock_hook_cls):
+ """A failing XCom push while reporting the failed run must not replace
the run's own exception."""
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
+
+ def _run_sync(*args, **kwargs):
+ kwargs["usage"].incr(RunUsage(requests=1))
+ raise UsageLimitExceeded("boom")
+
+ mock_agent.run_sync.side_effect = _run_sync
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ context = _make_context(ti_id="ti-1")
+ context["task_instance"].xcom_push.side_effect = RuntimeError("xcom
down")
+
+ op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c")
+ with pytest.raises(UsageLimitExceeded):
+ op.execute(context)
+ assert {c.kwargs["key"] for c in
context["task_instance"].xcom_push.call_args_list} == {
+ "usage",
+ "run_id",
+ }
@pytest.mark.skipif(
not AIRFLOW_V_3_1_PLUS, reason="a reviewing decision_policy needs the
HITL flow, Airflow >= 3.1"
@@ -619,7 +836,7 @@ class TestLLMOperatorConfidenceGate:
def test_confident_output_returns_and_records_per_field_confidence(
self, mock_hook_cls, make_mock_run_result
):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = self._result(
make_mock_run_result,
Summary(text="t"),
@@ -658,7 +875,7 @@ class TestLLMOperatorConfidenceGate:
self, mock_hook_cls, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
"""The least confident of the fields that reported a confidence is
what the bar is compared against."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = self._result(
make_mock_run_result,
Assessment(category="billing", score=0.3),
@@ -701,7 +918,7 @@ class TestLLMOperatorConfidenceGate:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_a_field_that_reports_no_confidence_is_not_gated(self,
mock_hook_cls, make_mock_run_result):
"""A bounded float field reports no confidence (the probability is the
answer); only reported fields count."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = self._result(
make_mock_run_result,
Assessment(category="billing", score=0.3),
@@ -727,7 +944,7 @@ class TestLLMOperatorConfidenceGate:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_nan_confidence_counts_as_not_reported(self, mock_hook_cls,
make_mock_run_result):
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = self._result(
make_mock_run_result,
Summary(text="t"),
@@ -861,7 +1078,7 @@ class TestLLMOperatorConfidenceGate:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_no_bar_keeps_todays_behaviour_and_still_records(self,
mock_hook_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("plain text")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = LLMOperator(task_id="t", prompt="p", llm_conn_id="c")
@@ -886,7 +1103,7 @@ class TestLLMOperatorConfidenceGate:
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_text_model_with_a_bar_can_fail(self, mock_hook_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("plain text")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
op = LLMOperator(
@@ -1085,7 +1302,7 @@ class TestLLMOperatorApproval:
self, mock_hook_cls, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
"""When require_approval=True, execute() defers instead of returning
output."""
- 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
@@ -1112,7 +1329,7 @@ class TestLLMOperatorApproval:
self, mock_hook_cls, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
"""On Airflow versions < 3.3 (flag pinned), execute() falls back to
deferring to HITLTrigger."""
- 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
@@ -1138,7 +1355,7 @@ class TestLLMOperatorApproval:
self, mock_hook_cls, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
"""allow_modifications=True passes an editable 'output' param."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = make_mock_run_result("draft output")
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
@@ -1164,7 +1381,7 @@ class TestLLMOperatorApproval:
self, mock_hook_cls, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
"""approval_timeout is passed to the trigger."""
- 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
@@ -1193,7 +1410,7 @@ class TestLLMOperatorApproval:
self, mock_hook_cls, mock_upsert, mock_trigger_cls,
make_mock_run_result
):
"""Structured (BaseModel) output is serialized before deferring."""
- mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value =
make_mock_run_result(Summary(text="hello"))
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
