Support per-run cost limits in common.ai LLM and Agent operators by Lee-W · Pull Request #71403 · apache/airflow

Lee-W · 2026-08-11T08:19:17Z

Why

usage_limits=UsageLimits(...) already caps requests and tokens. But we still have the following 3 gaps

Cost could not be expressed at all. UsageLimits.cost_limit and RunUsage.cost landed in pydantic-ai 2.23.0, but this provider's floor was >=2.0.0. On any resolved 2.0–2.22 install, cost_limit does not exist — a Dag author can cap tokens but has no way to cap spend, and no way to see what a run cost afterwards.
No UsageLimits field could be templated. A UsageLimits instance has neither resolve nor template_fields, so Airflow's template walk is a no-op on it. That blocks every field, not just cost: a request_limit or input_tokens_limit that follows the environment had no way to be written.
Where the cap stops being a guarantee was undocumented. Each of those cases reads as a broken promise to anyone who assumed otherwise.

What

Raise the pydantic-ai floor to 2.23.0 so cost_limit / RunUsage.cost are guaranteed present. That alone unlocks cost enforcement: the operators already forward usage_limits untouched to run_sync, so there is no interception logic of our own.
usage_limits now accepts a dict and is a template_fields entry. Passing {"cost_limit": "{{ params.budget }}", "request_limit": 5} makes every field templatable in one shot — a future tool_calls_limit needs no new operator parameter. A UsageLimits instance still passes through by identity: an author who built the object owns its values.
coerce_usage_limits() builds the real UsageLimits at execute time, parsing each value per UsageLimits's own dataclass field type rather than a hand-maintained list. A rendered value the author did not intend — a Variable that exists but is empty renders to "", and a typo renders to an arbitrary non-numeric string — raises a ValueError naming the offending field before any model call, instead of silently disabling that limit or blowing up deep inside pydantic-ai.
Log the run cost when the model reports one.
Document where the cap is not a guarantee: the request that crosses the limit still completes and still bills; self-hosted or unpriced models report no cost at all, so the limit silently does nothing; and a durable replay re-counts cached spend against a fresh attempt's budget.

Was generative AI tooling used to co-author this PR?

Yes (please specify the tool below)

Generated-by: [Claude] following the guidelines

Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
When adding dependency, check compliance with the ASF 3rd Party License Policy.
For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

uranusjr · 2026-08-17T04:36:33Z

We should add some cases on operator tests that actually use usage_limit. The resolve_usage_limit function itself looks fine, but there’s no evidence how it’s used is also fine.

Lee-W · 2026-08-18T13:14:54Z

Sounds good. Just added operator-level tests so the cap is enforced by pydantic-ai instead of a mock

github-actions · 2026-08-20T05:56:36Z

uv.lock on main just moved via #72314 ("[main] Upgrade important CI environment"), commit f044973 and this PR currently conflicts.

Quickest fix:

git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-lease

Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future uv.lock bumps.

vikramkoka · 2026-08-25T17:27:28Z

@Lee-W
I am struggling to understand the why here at least based on the PR description.

The token limits are enforced and I ran into this with the AIP-tracker example, which is why this piece of code was explicitly added by me there
usage_limits=UsageLimits(
request_limit=5,
input_tokens_limit=20_000,
output_tokens_limit=4_000,
)

There are a lot of files touched in this PR, so I didn't go through and check all the code changes, but at least the PR description seems off?

Lee-W · 2026-08-27T11:41:00Z

Yes, the PR description is outdated after several iteration. Just update it to the latest change

ColtenOuO

Overall, I think the approach looks great and the test coverage is very thorough.

Because max_cost is templated, its value is often not in the Dag author's control -- an unset Variable or a typo renders to a string the operator then has to make sense of. An "inf" that passed validation left a cap that could never be exceeded, which is worse than no cap at all: it reads as protection that isn't there. "nan" escaped as a bare decimal.InvalidOperation from the negativity comparison, contradicting the documented ValueError contract. The classification call LLMRetryPolicy makes on every task failure is not an operator run, so no cost cap has ever applied to it. That is pre-existing and deliberate, but someone who sets max_cost reasonably expects it to bound everything the task spends, so the exemption is now stated in the docs. The provider requires pydantic-ai-slim 2.23.0 for the cost API, so the extras now say so too rather than leaving a floor the resolver silently overrides.

The operator tests only asserted that a resolved UsageLimits reached a mocked run_sync. A mock accepts any value, so those assertions could not show that pydantic-ai honours the limit, that a bad templated max_cost surfaces its error, or that a run under the cap still completes.

`usage_limits` accepted only a `UsageLimits` instance, which is neither a scalar nor a member of any operator's `template_fields`, so no limit could come from a Variable, an XCom, or params. The `max_cost` parameter added earlier in this branch worked around that for the single most common knob, but it covered cost alone: a templated `request_limit` would have needed a second `max_requests`, and so on for every remaining field. Widening `usage_limits` itself closes the gap once, for every knob, and follows the pattern operators such as KubernetesPodOperator already use -- accept a templatable dict, build the real object at execute time. `max_cost` is dropped rather than kept as a shorthand: it never appeared in a released version, so there is nothing to deprecate, and keeping both entry points would only add surface to document and test. Because rendered values are not in the Dag author's control -- an unset Variable renders to an empty string, a typo to an arbitrary one -- coercion is deliberately loud. Each field is parsed per its declared type and rejected with a message naming both the field and the value received, and non-finite values are refused outright so that a cap can never be silently disabled. Validation happens after rendering rather than in `__init__`, which keeps a whole-field template such as `usage_limits="{{ params.limits }}"` usable under native rendering.

usage_limits is a templated field, so a Jinja render (or a hand-built dict) can produce a non-numeric value for a Decimal/int field, e.g. usage_limits={"cost_limit": []}. That value fell through the coercers straight into the finiteness check, which only handles Decimal/int/ float and raised an undocumented TypeError from math.isfinite instead of the ValueError this module otherwise promises for bad input.

The templated-usage-limits example rendered cost_limit from {{ var.value.llm_cost_cap_per_task }}, but nothing in the repo defines that Airflow Variable. Rendering a Jinja var.value reference to a Variable that doesn't exist raises instead of returning an empty string, so copying this example as-is fails immediately. Use var.value.get() with a default so it runs out of the box while still showing the templated pattern.

`bool` is a subclass of `int`, so a `usage_limits` dict carrying `False` where `null` was meant -- easy to produce from `Variable.get(..., deserialize_json=True)` -- built `UsageLimits(request_limit=False)` and killed the run before its first model call with an error pointing at pydantic-ai internals. A plain float slipped through the same way and left `request_limit` non-integral. Failing loudly here, naming the field and the offending value, is the one promise this module exists to keep. An integral float (`2.0`) is still accepted and narrowed to `int`, since native-mode Jinja renders one for what the author wrote as a whole number.

Resolving every `UsageLimits` field type at import time meant a single unsupported annotation upstream -- a `Literal[...]` or `list[int] | None` on a field nobody uses -- would raise while the module was still loading. All six operator modules import it at module scope, so that failure reached every Dag using the provider, including Dags that never set `usage_limits`. The pydantic-ai floor here has no upper bound and it ships a minor most weekdays, so this is exposure rather than theory. The failure stays loud; it just stops taking bystanders with it. Only an author who actually passes the offending field sees it.

Three statements did not match what the code does, and each would mislead a Dag author reading them while debugging a real failure. An unset Variable does not render to an empty string: `VariableAccessor` raises, and Jinja's sandbox only swallows `AttributeError`, so rendering fails long before the value reaches coercion. The empty string comes from a Variable that exists and is empty. That sentence had been copied into four places. Jinja does produce a real `None` under `render_template_as_native_obj`, and at that point it is indistinguishable from one the author wrote, so the limit is silently disabled. Worth stating as a limitation rather than implying the code detects it. Field-type resolution is scoped per field only for annotations it cannot support; `get_type_hints` parses the whole class in one call, so a failure there still reaches every caller.

The test asserted only that a retry raises `UsageLimitExceeded`, never the other half of what it claims to show: that the retry made no new model calls. The counter went out with the return value when the exception propagated, so the test passed just as well against a cold cache that had made a live call -- the one scenario it exists to rule out. `_run_one_attempt` now takes an optional counter so the caller keeps a reference across the raise. The three existing call sites are unchanged.

`%s` on a `Decimal` switches to scientific notation below a certain magnitude, so a cheap single-call run reported `LLM run cost: $7.5E-7`. Anyone reading the run summary to find out what a task actually cost then has to convert it themselves -- which is most of the point of logging it.

`cost_limit` was documented as silently having no effect for models pydantic-ai cannot price, in the same breath as saying a `CostNotFoundWarning` is emitted. A reader chasing why their cap did nothing would go looking for a signal the docs had just told them not to expect. A reviewer's suggestion fixing the wording was agreed to but never made it into a push. The same claim had meanwhile spread to the agent operator docs and the example Dag, so all three now say the limit has no effect without halting execution.

Found by the same sweep as the cost-limit wording and unrelated to it: the self-hosted model docs and a hook comment both describe the `except TypeError` fallback in `get_conn()` as silent. It logs a warning naming the provider and the rejected kwargs before falling back, so a reader is told to expect nothing while the log is in fact telling them exactly what happened. The hook comment also read as though the fallback avoided authenticating as the wrong identity, when that is its consequence.

The helper in conftest.py had no callers: all twelve test modules carried their own copy, which is why the same `cost=None` fix had to be applied by hand twelve times on this branch. Left that way, the next module copied from the local pattern would log `LLM run cost: $<MagicMock ...>` with nothing to catch it. Ten modules now take it as a fixture, matching how this repo already hands out factories elsewhere. The two file-analysis modules keep their own: they build every mock with an explicit `spec=[...]`, and the shared helper would loosen that deliberately strict pattern for no gain here.

`cost_limit` caps one run, so every retry and every scheduled run starts with the full allowance again. A reader who wanted a daily or monthly budget had nothing in these docs pointing anywhere useful. Kept to a single sentence on purpose: pydantic-ai-harness is not a default dependency here and nothing in this repo tests against it, so spelling out its options would go stale without anyone noticing.

count_tokens_before_request had no shape check at all: a list, dict, or plain float landed straight on UsageLimits and pydantic-ai read it by truthiness, so a container silently disabled the pre-flight token check the field exists to enforce. Integer fields had a matching gap for a native Decimal -- request_limit=Decimal("3.5") skipped the non-integral check already in place for a plain float.

"has no effect without halting execution" parses as "has no effect unless execution halts" -- the opposite of what's true. cost_limit is simply not enforced when pydantic-ai can't price a response; a CostNotFoundWarning is logged instead of failing the run. Same fix in all three places the phrase was copied.

The sentence pointed authors at pydantic-ai-harness's SpendLimits capability for a budget tracked across runs, but its default InMemorySpendStore does not enforce anything across processes -- every Airflow worker keeps its own count, so the cross-run claim only holds once a caller passes a shared store, which the sentence didn't say. Cross-task/cross-run budgets are being deferred past 1.0 rather than documented here now.

…ings Three gaps found reviewing the templated usage_limits coercion. _resolve_field_type reduced an annotation on argument count alone, so a single-parameter generic resolved to its element type: list[int] gave int and set[str] gave str. Only dict[str, int] raised, and only because it has two arguments, which is why the test covering this passed for the wrong reason. The origin is now checked, both on the annotation and on the member a union reduces to: list[int] | None is a union at the top level and reduces to list[int], which is still a type instance on Python 3.10 and only stopped being one in 3.11 (gh-101162), so checking only the top level left the oldest supported Python resolving a container field to its element type. _coerce_int accepted only an integer literal, so the templated spelling of a value the native path accepts was rejected: 5.0 became 5 but "5.0" raised, and the same for 1e3 and Decimal("5.0"). Jinja's / is true division, so {{ a / b }} renders "5.0" for a whole-number result, making the outcome depend on the Dag's render_template_as_native_obj flag. It now parses through Decimal and applies the same integral test the native path uses. int(Decimal) carries none of the digit ceiling CPython applies to int(str), so that fallback needed its own bound: "1E+100000" otherwise spent unbounded CPU building the integer and then failed inside its own error message with CPython's "Exceeds the limit (4300 digits)" rather than the field-named error this module promises. The native Decimal branch had the same hole.

The coercion sat on the run_sync call, so a usage_limits value that cannot be parsed only failed the task after the work needed to reach the model had already run: the file read and encode in LLMFileAnalysisOperator, the schema introspection in LLMSQLQueryOperator and LLMSchemaCompareOperator, and the durable-storage and agent construction in AgentOperator. A typo'd or empty Airflow Variable paid for all of that before reporting itself. It now runs ahead of the first expensive call in each execute(), and in regenerate_with_feedback(), with the result passed through to run_sync. The cheap guard raises already at the top of execute() keep their precedence, and the coercion still happens after __init__ because usage_limits is a template field whose value is not rendered yet when the constructor runs.

The cost_limit bullet said pydantic-ai cannot price self-hosted or custom endpoints, so the cap is not enforced there. Pricing keys on the model name, not the endpoint: a vLLM deployment served as openai:<model>, which self_hosted_models.rst gives as the only working prefix for vLLM, is priced at that model's public list rates rather than at what the deployment costs. The bullet now says which models are priced and which are not, matching how agent.rst already words it. Says that each limit bounds one agent run rather than a task, naming the Airflow task retry and the AgentOperator HITL regeneration that each start a fresh count. It deliberately stops short of quoting a total, because the same bullet already says the cap is not a hard guarantee against overspend, and it says "Airflow task retries" because the paragraph above uses retries in pydantic-ai's sense. Also refreshes the self_hosted_models.rst verification baseline, which still cited pydantic-ai 2.5.0 and 2.10.0, both below the 2.23.0 floor this change sets, and bumps the README extras table, where five rows still read >=2.0.0 while the Requirements table above them and docs/index.rst say >=2.23.0.

kaxil · 2026-09-14T22:42:38Z

Pushed three commits covering the open threads:

c8c326659 usage.py coercion: reject subscripted generics, including list[int] | None, which reduced to list[int] and resolved as a container type on Python 3.10; accept integral templated values so "5.0" behaves like a bare 5.0; and bound the Decimal path so a huge exponent cannot outrun CPython's digit ceiling.
588bc6214 hoist the coercion above the expensive setup in each execute(), so an unparsable value fails before the file read, the schema introspection, or the agent build.
d3f898d2b docs: pricing keys on the model name rather than the endpoint, each limit bounds one agent run, plus the self_hosted_models.rst baseline and the README extras table.

Cost cap against a real Anthropic model, 16 input and 4 output tokens, provider-reported cost $0.00018 against a cost_limit of 1E-7:

Fail-fast with no connection configured at all. The task fails on the named ValueError from llm.py:180, the hoisted coercion, rather than on the missing connection:

The docs spellcheck rejects "recognises". prek does not catch it, so it only surfaced in the Sphinx spellcheck job.

boring-cyborg Bot added area:providers kind:documentation provider:common-ai labels Aug 11, 2026

Lee-W force-pushed the common-ai-cost-limit branch 5 times, most recently from eb460d9 to acb43e9 Compare August 12, 2026 13:47

Lee-W marked this pull request as ready for review August 13, 2026 02:17

Lee-W requested review from gopidesupavan and kaxil as code owners August 13, 2026 02:17

Lee-W commented Aug 13, 2026

View reviewed changes

Comment thread providers/common/ai/docs/index.rst

Lee-W force-pushed the common-ai-cost-limit branch from acb43e9 to 04872a2 Compare August 13, 2026 08:45

Lee-W self-assigned this Aug 17, 2026

Lee-W force-pushed the common-ai-cost-limit branch from 04872a2 to d64570f Compare August 18, 2026 06:55

Lee-W requested a review from ashb as a code owner August 18, 2026 06:55

Lee-W force-pushed the common-ai-cost-limit branch 2 times, most recently from 08f6936 to 0249af2 Compare August 18, 2026 13:09

Lee-W requested a review from amoghrajesh August 18, 2026 13:20

Lee-W force-pushed the common-ai-cost-limit branch from 0249af2 to 8970257 Compare August 20, 2026 06:01

amoghrajesh reviewed Aug 20, 2026

View reviewed changes

Comment thread providers/common/ai/src/airflow/providers/common/ai/utils/usage.py Outdated

Lee-W force-pushed the common-ai-cost-limit branch 2 times, most recently from 2aa5193 to 097e847 Compare August 25, 2026 13:59

Lee-W force-pushed the common-ai-cost-limit branch from 097e847 to d07e532 Compare August 27, 2026 11:31

ColtenOuO reviewed Aug 30, 2026

View reviewed changes

Comment thread providers/common/ai/src/airflow/providers/common/ai/utils/usage.py

Comment thread providers/common/ai/docs/operators/llm.rst Outdated

Lee-W added 14 commits September 12, 2026 15:14

Lee-W force-pushed the common-ai-cost-limit branch from 93cbc48 to 607786c Compare September 12, 2026 07:14

Lee-W added 3 commits September 12, 2026 15:15

Lee-W force-pushed the common-ai-cost-limit branch from 607786c to 989b0b5 Compare September 12, 2026 07:15

kaxil added 3 commits September 14, 2026 21:52

Use US spelling in the common.ai cost_limit pricing caveat

ffbf1b2

The docs spellcheck rejects "recognises". prek does not catch it, so it only surfaced in the Sphinx spellcheck job.

kaxil approved these changes Sep 15, 2026

View reviewed changes

kaxil merged commit 4c3ec9c into apache:main Sep 15, 2026
95 checks passed

kaxil deleted the common-ai-cost-limit branch September 15, 2026 00:28

potiuk mentioned this pull request Sep 20, 2026

Reject non-string prompts in LLMFileAnalysisOperator before reading files #71734

Merged

1 task done

shahar1 mentioned this pull request Sep 24, 2026

Status of testing Providers that were prepared on September 22, 2026 #73658

Closed

ColtenOuO mentioned this pull request Oct 7, 2026

Add ColtenOuO's April to September contribution report for apache-airflow opensource4you/readme#26

Open

Navigation Menu

Repository navigation

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

Support per-run cost limits in common.ai LLM and Agent operators#71403
kaxil merged 22 commits into
apache:mainfrom
astronomer:common-ai-cost-limit

Lee-W commented Aug 11, 2026 •

edited

Loading

Uh oh!

Uh oh!

uranusjr commented Aug 17, 2026

Uh oh!

Lee-W commented Aug 18, 2026

Uh oh!

github-actions Bot commented Aug 20, 2026 •

edited

Loading

Uh oh!

Uh oh!

vikramkoka commented Aug 25, 2026

Uh oh!

Lee-W commented Aug 27, 2026

Uh oh!

ColtenOuO left a comment

Uh oh!

Uh oh!

Uh oh!

kaxil commented Sep 14, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

Lee-W commented Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why

What

Was generative AI tooling used to co-author this PR?

Uh oh!

Uh oh!

uranusjr commented Aug 17, 2026

Uh oh!

Lee-W commented Aug 18, 2026

Uh oh!

github-actions Bot commented Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

vikramkoka commented Aug 25, 2026

Uh oh!

Lee-W commented Aug 27, 2026

Uh oh!

ColtenOuO left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

kaxil commented Sep 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Lee-W commented Aug 11, 2026 •

edited

Loading

github-actions Bot commented Aug 20, 2026 •

edited

Loading

kaxil commented Sep 14, 2026 •

edited

Loading