Support per-run cost limits in common.ai LLM and Agent operators

providers/common/ai/README.rst

@@ -56,7 +56,7 @@ PIP package                                 Version required
 ``apache-airflow``                          ``>=3.0.0``
 ``apache-airflow-providers-common-compat``  ``>=1.15.0``
 ``apache-airflow-providers-standard``       ``>=1.12.1``
-``pydantic-ai-slim``                        ``>=2.0.0``
+``pydantic-ai-slim``                        ``>=2.23.0``
 ==========================================  ==================
 
 Optional cross provider package dependencies

providers/common/ai/docs/index.rst

@@ -204,7 +204,7 @@ PIP package                                 Version required
 ``apache-airflow``                          ``>=3.0.0``
 ``apache-airflow-providers-common-compat``  ``>=1.15.0``
 ``apache-airflow-providers-standard``       ``>=1.12.1``
-``pydantic-ai-slim``                        ``>=2.0.0``
+``pydantic-ai-slim``                        ``>=2.23.0``
 ==========================================  ==================
 
 Optional cross provider package dependencies
@@ -241,11 +241,11 @@ Install them when installing from PyPI. For example:
 ==============  =======================================================================================================================================
 Extra           Dependencies
 ==============  =======================================================================================================================================
-``anthropic``   ``pydantic-ai-slim[anthropic]>=2.0.0``
-``bedrock``     ``pydantic-ai-slim[bedrock]>=2.0.0``
-``google``      ``pydantic-ai-slim[google]>=2.0.0``
-``openai``      ``pydantic-ai-slim[openai]>=2.0.0``
-``mcp``         ``pydantic-ai-slim[mcp]>=2.0.0``
+``anthropic``   ``pydantic-ai-slim[anthropic]>=2.23.0``
+``bedrock``     ``pydantic-ai-slim[bedrock]>=2.23.0``
+``google``      ``pydantic-ai-slim[google]>=2.23.0``
+``openai``      ``pydantic-ai-slim[openai]>=2.23.0``
+``mcp``         ``pydantic-ai-slim[mcp]>=2.23.0``
 ``code-mode``   ``pydantic-ai-harness[codemode]>=0.3.0``
 ``shields``     ``pydantic-ai-shields>=0.3.4``
 ``skills``      ``apache-airflow-providers-git>=0.4.0``, ``pydantic-ai-skills>=1.2.0``

providers/common/ai/docs/observability.rst

@@ -51,6 +51,9 @@ How it works
 * **Content is off by default.** Only token counts, model id, latency, tool
   names, and finish reason are recorded. Prompt and completion text is never
   emitted unless you opt in (see below).
+* **Cost is already on the span.** pydantic-ai's own instrumentation sets a
+  best-effort ``operation.cost`` attribute on the model-call span whenever it
+  can price the response -- no provider configuration is needed for this.
 
 .. note::
 

providers/common/ai/docs/operators/agent.rst

@@ -487,7 +487,25 @@ Parameters
   agent run (initial run, durable replay, and HITL regeneration). Use it to
   cap requests, tokens, or tool calls per task -- agents are particularly
   prone to runaway tool loops, so ``tool_calls_limit`` is a useful guardrail.
-  See :ref:`howto/operator:llm` for an example. Default ``None``.
+  It also supports a per-run USD ``cost_limit``; see :ref:`howto/operator:llm`
+  for the caveats (not a hard guarantee, silently inert for unpriced models)
+  and an example. Default ``None``.
+
+  .. warning::
+     With ``durable=True``, a task retry replays cached model steps instead of
+     re-calling the model -- but pydantic-ai still adds each replayed step's
+     cost to the retry's own usage total, since it cannot distinguish a replay
+     from a live call. A ``cost_limit`` therefore counts already-paid-for
+     replayed cost against every retry's fresh budget, leaving less headroom
+     for the new calls the retry actually makes. And if the limit is lowered
+     between attempts -- easy to do by accident, since ``max_cost`` is
+     templated -- a retry can exceed it with zero new model calls. The
+     ``LLM run cost`` line in the task log reports the run's cumulative cost
+     for the same reason, not what this attempt actually spent.
+- ``max_cost``: Convenience per-run USD cost cap, as a templated alternative to
+  ``usage_limits.cost_limit`` (``usage_limits`` itself cannot be templated). Overrides
+  ``cost_limit`` on ``usage_limits`` if both are set; every other field on
+  ``usage_limits`` is preserved.  Default ``None`` (``usage_limits`` unchanged).
 - ``durable``: When ``True``, enables step-level caching of model responses and
   tool results. On retry, cached steps are replayed instead of re-executing
   expensive LLM calls. On Airflow >= 3.3 the cache uses the task state store (no

providers/common/ai/docs/operators/llm.rst

@@ -119,6 +119,17 @@ calls within a single task.
     :start-after: [START howto_operator_llm_usage_limits]
     :end-before: [END howto_operator_llm_usage_limits]
 
+``usage_limits`` cannot be templated -- it's an object, not a scalar. For a per-run
+cost cap, use ``max_cost`` -- a plain number (``max_cost=0.5``), and templatable:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
+    :language: python
+    :start-after: [START howto_operator_llm_max_cost]
+    :end-before: [END howto_operator_llm_max_cost]
+
+``max_cost`` overrides ``cost_limit`` on ``usage_limits`` (building one if
+``usage_limits`` is ``None``); every other field on ``usage_limits`` is left as-is.
+
 Common knobs on ``UsageLimits``:
 
 - ``request_limit`` — max model requests per run (caps retry/tool-loop blow-ups).
@@ -129,6 +140,19 @@ Common knobs on ``UsageLimits``:
 - ``input_tokens_limit`` / ``output_tokens_limit`` — per-run token caps.
 - ``total_tokens_limit`` — combined input + output cap.
 - ``tool_calls_limit`` — max tool invocations (``AgentOperator`` only).
+- ``cost_limit`` — a ``Decimal`` cap on the run's estimated USD cost. This is **not** a
+  hard guarantee against overspend: the response that crosses the limit has already been
+  produced and billed — pydantic-ai checks the accumulated cost *after* each response and
+  then fails the run with ``UsageLimitExceeded``. It protects you from further spend, not
+  from the request that broke the budget; even a single-request run fails as soon as that
+  request's cost pushes the total over the limit. For self-hosted or unknown
+  models (e.g. Ollama, custom endpoints) pydantic-ai cannot price the response, so cost
+  is ``None`` and ``cost_limit`` silently has no effect (a ``CostNotFoundWarning`` is
+  emitted instead of a failure). And like the other knobs above, setting ``cost_limit``
+  alone still inherits the ``request_limit=50`` default — see the ``request_limit`` note
+  above. Note that ``cost_limit`` (and ``max_cost``) only cap the operator's own LLM calls --
+  the meta-agent that ``LLMRetryPolicy`` runs to classify a failed task is a separate,
+  uncapped LLM call; see :doc:`../retry_policies`.
 
 When the limit is hit pydantic-ai raises ``UsageLimitExceeded``, which
 propagates to Airflow as a task failure — Airflow's standard retry policy
@@ -219,6 +243,10 @@ Parameters
   constructor (e.g. ``retries``, ``model_settings``, ``tools``). Supports Jinja templating.
 - ``usage_limits``: Optional pydantic-ai ``UsageLimits`` enforced on the run.  Fails
   the task when token / request / tool-call budgets are exceeded.  Default ``None``.
+- ``max_cost``: Convenience per-run USD cost cap, as a templated alternative to
+  ``usage_limits.cost_limit`` (``usage_limits`` itself cannot be templated). Overrides
+  ``cost_limit`` on ``usage_limits`` if both are set; every other field on
+  ``usage_limits`` is preserved.  Default ``None`` (``usage_limits`` unchanged).
 - ``require_approval``: If ``True``, the task defers after generating output and waits
   for human review.  Default ``False``.
 - ``approval_timeout``: Maximum time to wait for a review (``timedelta``).  ``None``

providers/common/ai/docs/retry_policies.rst

@@ -79,6 +79,12 @@ When a task fails, ``LLMRetryPolicy``:
 3. Based on the classification, returns RETRY (with a suggested delay) or FAIL
 4. The classification reason is logged in the task logs
 
+This classification call is a separate LLM request, made by ``LLMRetryPolicy``
+itself rather than by an operator -- it is not subject to an operator's
+``usage_limits`` or ``max_cost``, and it runs on every task failure regardless
+of any cost cap configured on the failing task. It is bounded by ``timeout``
+and ``max_exception_length``, but not by a cost limit.
+
 If the LLM call fails (provider down, timeout, bad credentials), the policy
 falls back to ``fallback_rules`` if configured, or to the task's standard
 retry behaviour.

providers/common/ai/pyproject.toml

@@ -69,18 +69,19 @@ dependencies = [
     "apache-airflow>=3.0.0",
     "apache-airflow-providers-common-compat>=1.15.0",
     "apache-airflow-providers-standard>=1.12.1",
-    # Requires the pydantic-ai 2.x agent/instrumentation API (see #69122).
-    "pydantic-ai-slim>=2.0.0",
+    # Requires the pydantic-ai cost API (RunUsage.cost, UsageLimits.cost_limit),
+    # landed in 2.23.0 via https://github.com/pydantic/pydantic-ai/pull/2684.
+    "pydantic-ai-slim>=2.23.0",
 ]
 
 # The optional dependencies should be modified in place in the generated file
 # Any change in the dependencies is preserved when the file is regenerated
 [project.optional-dependencies]
-"anthropic" = ["pydantic-ai-slim[anthropic]>=2.0.0"]
-"bedrock" = ["pydantic-ai-slim[bedrock]>=2.0.0"]
-"google" = ["pydantic-ai-slim[google]>=2.0.0"]
-"openai" = ["pydantic-ai-slim[openai]>=2.0.0"]
-"mcp" = ["pydantic-ai-slim[mcp]>=2.0.0"]
+"anthropic" = ["pydantic-ai-slim[anthropic]>=2.23.0"]
+"bedrock" = ["pydantic-ai-slim[bedrock]>=2.23.0"]
+"google" = ["pydantic-ai-slim[google]>=2.23.0"]
+"openai" = ["pydantic-ai-slim[openai]>=2.23.0"]
+"mcp" = ["pydantic-ai-slim[mcp]>=2.23.0"]
 # Code mode: collapse tool calls into a single `run_code` tool that the model
 # drives by writing Python, executed in the Monty sandbox (pydantic-monty).
 # Enables AgentOperator(code_mode=True). Monty is pre-1.0; pinned here as an
@@ -137,7 +138,7 @@ dev = [
     "apache-airflow-providers-standard",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
     "sqlglot>=30.0.0",
-    "pydantic-ai-slim[mcp]>=2.0.0",
+    "pydantic-ai-slim[mcp]>=2.23.0",
     "pydantic-ai-skills>=1.2.0",
     "apache-airflow-providers-common-sql[datafusion]",
     "langchain>=1.0.0",

providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py

@@ -19,6 +19,7 @@
 from __future__ import annotations
 
 from datetime import timedelta
+from decimal import Decimal
 
 from pydantic import BaseModel
 from pydantic_ai.usage import UsageLimits
@@ -139,6 +140,10 @@ def example_llm_operator_usage_limits():
             request_limit=5,
             input_tokens_limit=4_000,
             output_tokens_limit=1_000,
+            # Fail the task if the run's estimated USD cost exceeds $0.50.
+            # See docs/operators/llm.rst for caveats (not a hard guarantee,
+            # silently inert for models pydantic-ai can't price).
+            cost_limit=Decimal("0.50"),
         ),
     )
 
@@ -148,6 +153,27 @@ def example_llm_operator_usage_limits():
 example_llm_operator_usage_limits()
 
 
+# [START howto_operator_llm_max_cost]
+@dag(tags=["example"])
+def example_llm_operator_max_cost():
+    LLMOperator(
+        task_id="capped_summary",
+        prompt="Summarize the trade-offs of a message queue vs. direct HTTP calls in three bullet points.",
+        llm_conn_id="pydanticai_default",
+        system_prompt="You are a concise technical reviewer.",
+        # Unlike usage_limits (a UsageLimits object), max_cost is a scalar and
+        # can be templated -- e.g. driven by an Airflow Variable so the budget
+        # can change per environment without editing the DAG. This caps a single
+        # task run, not a day's total spend -- each run gets the full budget again.
+        max_cost="{{ var.value.llm_max_cost_per_task }}",
+    )
+
+
+# [END howto_operator_llm_max_cost]
+
+example_llm_operator_max_cost()
+
+
 # [START howto_operator_llm_approval]
 @dag(tags=["example"])
 def example_llm_operator_approval():

providers/common/ai/src/airflow/providers/common/ai/operators/agent.py

@@ -31,6 +31,7 @@
 from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin
 from airflow.providers.common.ai.utils.logging import log_run_summary, wrap_toolsets_for_logging
 from airflow.providers.common.ai.utils.output_type import rehydrate_pydantic_output
+from airflow.providers.common.ai.utils.usage import resolve_usage_limits
 from airflow.providers.common.compat.sdk import (
     AirflowOptionalProviderFeatureException,
     BaseOperator,
@@ -48,6 +49,8 @@
     _CORE_WALKER = False
 
 if TYPE_CHECKING:
+    from decimal import Decimal
+
     from pydantic_ai import Agent
     from pydantic_ai.messages import ModelMessage
     from pydantic_ai.toolsets.abstract import AbstractToolset
@@ -148,6 +151,14 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
         ``UsageLimits(request_limit=..., total_tokens_limit=..., tool_calls_limit=..., ...)``
         to fail the task when the agent exceeds the configured token, request,
         or tool budget. ``None`` (default) means no enforcement.
+    :param max_cost: Convenience per-run USD cost cap, as a templated alternative
+        to ``usage_limits.cost_limit`` (``usage_limits`` itself cannot be
+        templated). When set, overrides ``cost_limit`` on ``usage_limits``
+        (building one if ``usage_limits`` is ``None``); every other field on
+        ``usage_limits`` is left untouched. ``None`` (default) leaves
+        ``usage_limits`` unchanged. See :ref:`howto/operator:llm` for general
+        ``cost_limit`` caveats, and :ref:`howto/operator:agent` for the
+        ``durable=True`` replay double-counting warning.
     :param durable: When ``True``, enables step-level caching of model
         responses and tool results for durable execution.  On retry, cached
         steps are replayed instead of re-executing.  Each cached step is
@@ -232,6 +243,7 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
         "system_prompt",
         "agent_params",
         "message_history",
+        "max_cost",
     )
 
     operator_extra_links = (HITLReviewLink(),)
@@ -248,6 +260,7 @@ def __init__(
         enable_tool_logging: bool = True,
         agent_params: dict[str, Any] | None = None,
         usage_limits: UsageLimits | None = None,
+        max_cost: Decimal | float | str | None = None,
         durable: bool = False,
         code_mode: bool = False,
         message_history: list[ModelMessage] | str | bytes | None = None,
@@ -274,6 +287,7 @@ def __init__(
         self.enable_tool_logging = enable_tool_logging
         self.agent_params = agent_params or {}
         self.usage_limits = usage_limits
+        self.max_cost = max_cost
         self.message_history = message_history
 
         self.durable = durable
@@ -446,7 +460,7 @@ def execute(self, context: Context) -> Any:
 
         agent = self._build_agent()
 
-        run_kwargs: dict[str, Any] = {"usage_limits": self.usage_limits}
+        run_kwargs: dict[str, Any] = {"usage_limits": resolve_usage_limits(self.usage_limits, self.max_cost)}
         history = self._resolve_message_history()
         if history is not None:
             run_kwargs["message_history"] = history
@@ -553,7 +567,11 @@ def regenerate_with_feedback(self, *, feedback: str, message_history: Any) -> tu
         """Re-run the agent with *feedback* appended to the conversation history."""
         agent = self._build_agent()
         messages = message_history or []
-        result = agent.run_sync(feedback, message_history=messages, usage_limits=self.usage_limits)
+        result = agent.run_sync(
+            feedback,
+            message_history=messages,
+            usage_limits=resolve_usage_limits(self.usage_limits, self.max_cost),
+        )
         log_run_summary(self.log, result)
 
         output = result.output

providers/common/ai/src/airflow/providers/common/ai/operators/llm.py

@@ -29,6 +29,7 @@
 from airflow.providers.common.ai.mixins.approval import LLMApprovalMixin
 from airflow.providers.common.ai.utils.logging import log_run_summary
 from airflow.providers.common.ai.utils.output_type import rehydrate_pydantic_output
+from airflow.providers.common.ai.utils.usage import resolve_usage_limits
 from airflow.providers.common.compat.sdk import BaseOperator
 
 try:
@@ -41,6 +42,8 @@
     _CORE_WALKER = False
 
 if TYPE_CHECKING:
+    from decimal import Decimal
+
     from pydantic_ai import Agent
     from pydantic_ai.usage import UsageLimits
 
@@ -80,6 +83,12 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
         ``UsageLimits(request_limit=..., total_tokens_limit=..., ...)`` to fail
         the task when the agent exceeds the configured token, request, or tool
         budget. ``None`` (default) means no enforcement.
+    :param max_cost: Convenience per-run USD cost cap, as a templated alternative
+        to ``usage_limits.cost_limit`` (``usage_limits`` itself cannot be
+        templated). When set, overrides ``cost_limit`` on ``usage_limits``
+        (building one if ``usage_limits`` is ``None``); every other field on
+        ``usage_limits`` is left untouched. ``None`` (default) leaves
+        ``usage_limits`` unchanged. See :ref:`howto/operator:llm` for caveats.
     :param require_approval: If ``True``, the task defers after generating
         output and waits for a human reviewer to approve or reject via the
         HITL interface.  Default ``False``.
@@ -104,6 +113,7 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
         "model_id",
         "system_prompt",
         "agent_params",
+        "max_cost",
     )
 
     def __init__(
@@ -116,6 +126,7 @@ def __init__(
         output_type: type = str,
         agent_params: dict[str, Any] | None = None,
         usage_limits: UsageLimits | None = None,
+        max_cost: Decimal | float | str | None = None,
         require_approval: bool = False,
         approval_timeout: timedelta | None = None,
         allow_modifications: bool = False,
@@ -135,6 +146,7 @@ def __init__(
         self._serialize_model_output = serialize_output or not _CORE_WALKER
         self.agent_params = agent_params or {}
         self.usage_limits = usage_limits
+        self.max_cost = max_cost
         self.require_approval = require_approval
         self.approval_timeout = approval_timeout
         self.allow_modifications = allow_modifications
@@ -161,7 +173,9 @@ def execute(self, context: Context) -> Any:
         agent: Agent[object, Any] = self.llm_hook.create_agent(
             output_type=self.output_type, instructions=self.system_prompt, **self.agent_params
         )
-        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        result = agent.run_sync(
+            self.prompt, usage_limits=resolve_usage_limits(self.usage_limits, self.max_cost)
+        )
         log_run_summary(self.log, result)
         output = result.output
 

providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py

@@ -25,6 +25,7 @@
 
 from airflow.providers.common.ai.operators.llm import LLMOperator
 from airflow.providers.common.ai.utils.logging import log_run_summary
+from airflow.providers.common.ai.utils.usage import resolve_usage_limits
 from airflow.providers.standard.operators.branch import BranchMixIn
 
 if TYPE_CHECKING:
@@ -96,7 +97,9 @@ def execute(self, context: Context) -> str | Iterable[str] | None:
             instructions=self.system_prompt,
             **self.agent_params,
         )
-        result = agent.run_sync(self.prompt, usage_limits=self.usage_limits)
+        result = agent.run_sync(
+            self.prompt, usage_limits=resolve_usage_limits(self.usage_limits, self.max_cost)
+        )
         log_run_summary(self.log, result)
         output = result.output
 

providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py

@@ -26,6 +26,7 @@
 from airflow.providers.common.ai.operators.llm import LLMOperator
 from airflow.providers.common.ai.utils.file_analysis import build_file_analysis_request
 from airflow.providers.common.ai.utils.logging import log_run_summary
+from airflow.providers.common.ai.utils.usage import resolve_usage_limits
 
 if TYPE_CHECKING:
     from pydantic_ai import Agent
@@ -134,7 +135,9 @@ def execute(self, context: Context) -> Any:
             instructions=self._build_system_prompt(),
             **self.agent_params,
         )
-        result = agent.run_sync(request.user_content, usage_limits=self.usage_limits)
+        result = agent.run_sync(
+            request.user_content, usage_limits=resolve_usage_limits(self.usage_limits, self.max_cost)
+        )
         log_run_summary(self.log, result)
         output = result.output