This is an automated email from the ASF dual-hosted git repository.
kaxil 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 57a4751c967 Add an LLM task loop example that stops on a dollar or
token budget (#75516)
57a4751c967 is described below
commit 57a4751c9678432858662aa3223ff7688c07650f
Author: Kaxil Naik <[email protected]>
AuthorDate: Sat Oct 24 00:36:03 2026 +0100
Add an LLM task loop example that stops on a dollar or token budget (#75516)
The Common AI example Dag example_agent_loop_budget has an agent write a
status page update and a second model score it, repeating both in a task
loop until the score passes or the loop has spent a dollar or token budget.
A terminal tally task sums each pass's usage XCom into loop.previous, and
the gate returns a LoopDecision whose reason names the ending.
A model with no price data turns the dollar cap off for the rest of the
loop instead of counting as zero, so the token cap always applies.
The new use-case page documents the four endings, the one-pass overshoot,
estimated versus billed cost, unpriced models and retries. The core loops
page links to it and states that an in-loop xcom_pull returns the current
iteration's value.
* Keep the loop example's judge rules on separate lines
Move the judge's system prompt to a module-level JUDGE_INSTRUCTIONS
constant, next to INCIDENT_NOTES, so the numbered rules reach the model
as a list instead of one run-on line.
* Let the loop example's budget stop the loop before max_iterations
With max_iterations=5, the per-task usage limits capped the loop at 150,000
tokens and $3.50, below its $5.00 and 500,000-token budget, so neither
budget
stop could fire. Raise max_iterations to 20 and recapture the converged and
max_iterations runs the page shows.
Quote the provider's own log lines for a tripped per-task limit instead of a
line printed by the example event listener, and drop the use-case count from
the docs index so it does not go stale.
---
.../docs/authoring-and-scheduling/loops.rst | 10 +-
providers/common/ai/docs/examples.rst | 4 +
providers/common/ai/docs/index.rst | 2 +-
providers/common/ai/docs/use_cases/index.rst | 4 +
.../use_cases/refine_a_draft_within_a_budget.rst | 202 +++++++++++++++++++++
.../ai/example_dags/example_agent_loop_budget.py | 141 ++++++++++++++
6 files changed, 359 insertions(+), 4 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst
b/airflow-core/docs/authoring-and-scheduling/loops.rst
index 9e414563e59..42f45dc57ea 100644
--- a/airflow-core/docs/authoring-and-scheduling/loops.rst
+++ b/airflow-core/docs/authoring-and-scheduling/loops.rst
@@ -130,6 +130,9 @@ stopping.
The ``accurate_enough`` condition in the ``refine_estimate`` example above
returns a
``LoopDecision``, so each iteration of that loop shows its error next to its
outcome.
+:doc:`Refine a draft within a budget
<apache-airflow-providers-common-ai:use_cases/refine_a_draft_within_a_budget>`
+uses the same pattern for LLM calls, stopping with a reason when a judge
passes the draft or when
+the loop has spent its dollar or token budget.
Read the loop context
=====================
@@ -270,9 +273,10 @@ gate's state. A trigger rule such as ``always`` does not
wait for the loop and
may see only the iterations that exist so far.
Inside the loop nothing changes: use ``loop.result`` and ``loop.previous``. A
-task outside the loop cannot pull a single value from a loop task, because the
-task ran once for each iteration, and cannot use its output to expand a mapped
-task.
+``ti.xcom_pull`` of another task in the same loop returns that task's value
from the
+current iteration. A task outside the loop cannot pull a single value from a
loop
+task, because the task ran once for each iteration, and cannot use its output
to
+expand a mapped task.
.. _loops-mapped-tasks:
diff --git a/providers/common/ai/docs/examples.rst
b/providers/common/ai/docs/examples.rst
index 0672b0d4b67..67015bec505 100644
--- a/providers/common/ai/docs/examples.rst
+++ b/providers/common/ai/docs/examples.rst
@@ -100,6 +100,10 @@ Agents & tools
- A LangChain ReAct agent that decides its own tool calls, composed with
``LLMOperator`` for
report formatting and AIP-90 HITL review
(`example_langchain_tool_agent.py
<https://github.com/apache/airflow/blob/providers-common-ai/|version|/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_langchain_tool_agent.py>`__).
+ * - :doc:`use_cases/refine_a_draft_within_a_budget`
+ - An agent and a judge in a task loop that stops when the draft passes or
a dollar or token
+ budget is spent
+ (`example_agent_loop_budget.py
<https://github.com/apache/airflow/blob/providers-common-ai/|version|/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_loop_budget.py>`__).
Retrieval & document processing
--------------------------------
diff --git a/providers/common/ai/docs/index.rst
b/providers/common/ai/docs/index.rst
index 63ced48e7b0..8b2322bca01 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -35,7 +35,7 @@ Start here
- :doc:`features`: what the provider can do, with the page, extra, minimum
Airflow
version and stability of each feature.
- :doc:`quickstart`: install, connect a vendor, run a two-task Dag and check
its output.
-- :doc:`use_cases/index`: ten jobs a data team already has, each with the Dag
that does it.
+- :doc:`use_cases/index`: jobs a data team already has, each with the Dag that
does it.
- :doc:`model_providers`: which vendors work, and the extra, connection and
prefix for each.
This is the Dag the quick start runs. The ``summarize`` task sends the release
notes to the
diff --git a/providers/common/ai/docs/use_cases/index.rst
b/providers/common/ai/docs/use_cases/index.rst
index 7793e4e4c5d..6da286663ea 100644
--- a/providers/common/ai/docs/use_cases/index.rst
+++ b/providers/common/ai/docs/use_cases/index.rst
@@ -70,6 +70,9 @@ exception: it reads a decision-model connection,
``decision_default``, instead.
* - :doc:`research_agent_with_review`
- Decides which tools to call to answer a research question
- Human edits the question first, separate formatting step, approval
before delivery
+ * - :doc:`refine_a_draft_within_a_budget`
+ - Writes a status page update, then scores it and gives feedback
+ - Repeats both in a task loop until the score passes or a dollar or token
budget is spent
More ideas
----------
@@ -122,3 +125,4 @@ can be built from the patterns shown on the pages above.
Weekly status report <weekly_status_report>
Classify reviews in bulk <classify_reviews_in_bulk>
Research agent with review <research_agent_with_review>
+ Refine a draft within a budget <refine_a_draft_within_a_budget>
diff --git
a/providers/common/ai/docs/use_cases/refine_a_draft_within_a_budget.rst
b/providers/common/ai/docs/use_cases/refine_a_draft_within_a_budget.rst
new file mode 100644
index 00000000000..c65f8a96ad4
--- /dev/null
+++ b/providers/common/ai/docs/use_cases/refine_a_draft_within_a_budget.rst
@@ -0,0 +1,202 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+Refine a draft within a budget
+==============================
+
+A model's first draft of a customer status page update is rarely the one you
publish. This Dag
+has an agent write the update and a second model score it against a style
guide. A
+:ref:`task loop <apache-airflow:loops>` repeats both until the score reaches
0.9, and stops
+sooner once the loop as a whole has spent $5.00 or 500,000 tokens. Each pass
shows in the UI with
+the reason the loop continued or stopped.
+
+The budget is checked between passes, so it limits how many passes run, not
what one pass
+spends. To cap a pass, set ``usage_limits`` on each task.
+
+What this demonstrates
+----------------------
+
+* :ref:`Task loops <apache-airflow:loops>` -- ``refine.loop(max_iterations=20,
until=good_enough)``
+ repeats the task group, and ``good_enough`` returns a ``LoopDecision`` with
a reason.
+* The ``usage`` XCom that ``@task.agent`` and ``@task.llm`` push, read back
inside the loop to
+ keep a running total of tokens and estimated cost.
+* :ref:`howto/operator:llm_usage_limits` on each task, which caps what one
pass can spend.
+* :doc:`../structured_output` -- the judge's output is validated against
``Score``.
+
+Run it
+------
+
+The Dag needs Airflow 3.4.0 or later for task loops, and Common AI 1.1.0 or
later, the first
+release in which ``@task.llm`` pushes the ``usage`` XCom. Both tasks use the
model set on the
+``pydanticai_default`` connection.
+
+1. Install the provider with the extra for your vendor:
+
+ .. code-block:: bash
+
+ pip install "apache-airflow-providers-common-ai[anthropic]>=1.1.0"
+
+2. Trigger the Dag:
+
+ .. code-block:: bash
+
+ airflow dags test example_agent_loop_budget
+
+The gate logs its decision for each iteration, and ``publish`` logs the
result. From a run with
+``anthropic:claude-sonnet-4-6``:
+
+.. code-block:: text
+
+ Loop iteration 1: continue (max_iterations=20): score 0.5; $0.007317 of
$5.00, 1,427 of 500,000 tokens
+ Loop iteration 2: continue (max_iterations=20): score 0.7; $0.016605 of
$5.00, 3,131 of 500,000 tokens
+ Loop iteration 3: continue (max_iterations=20): score 0.5; $0.025290 of
$5.00, 4,870 of 500,000 tokens
+ Loop iteration 4: continue (max_iterations=20): score 0.5; $0.033585 of
$5.00, 6,543 of 500,000 tokens
+ Loop iteration 5: continue (max_iterations=20): score 0.5; $0.042198 of
$5.00, 8,226 of 500,000 tokens
+ Loop iteration 6: continue (max_iterations=20): score 0.7; $0.051069 of
$5.00, 9,919 of 500,000 tokens
+ Loop iteration 7: continue (max_iterations=20): score 0.5; $0.060381 of
$5.00, 11,643 of 500,000 tokens
+ Loop iteration 8: continue (max_iterations=20): score 0.7; $0.069204 of
$5.00, 13,376 of 500,000 tokens
+ Loop iteration 9: continue (max_iterations=20): score 0.5; $0.079605 of
$5.00, 15,175 of 500,000 tokens
+ Loop iteration 10: continue (max_iterations=20): score 0.5; $0.087213 of
$5.00, 16,883 of 500,000 tokens
+ Loop iteration 11: stop (max_iterations=20): passed: score 0.9; $0.096552
of $5.00, 18,624 of 500,000 tokens
+ 11 passes, score 0.9, 18624 tokens, $0.096552
+
+The Dag
+-------
+
+The whole loop, with the budget in ``tally`` and ``good_enough``:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_agent_loop_budget.py
+ :language: python
+ :start-after: [START howto_loop_budget]
+ :end-before: [END howto_loop_budget]
+
+Each pass runs ``write_draft``, then ``judge``, then ``tally``. ``tally`` is
the loop body's only
+terminal task, so what it returns is ``loop.result`` for the gate and
``loop.previous`` for the
+next pass: the draft, the judge's feedback, and the running totals.
+
+Inside a loop, ``ti.xcom_pull(..., key="usage")`` returns the ``usage`` XCom
of the current
+iteration only, so ``tally`` adds this pass's spend to the totals it finds in
``loop.previous``.
+Money is added as ``Decimal`` and stored as a string in XCom, so the total
does not pick up
+``float`` rounding.
+
+How the loop ends
+-----------------
+
+Every reason carries the score and both totals, and a reason that stops the
loop starts with the
+ending:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 22 40 38
+
+ * - Ending
+ - Reason on the last gate
+ - Result
+ * - Converged
+ - ``passed: score 0.9; $0.096552 of $5.00, 18,624 of 500,000 tokens``
+ - The gate succeeds and ``publish`` runs.
+ * - Dollar budget spent
+ - ``dollar budget spent: score 0.9; $0.016362 of $0.01, 3,226 of 500,000
tokens``
+ - The gate succeeds and ``publish`` runs.
+ * - Token budget spent
+ - ``token budget spent: score 0.0; cost unknown, 745 of 600 tokens``
+ - The gate succeeds and ``publish`` runs.
+ * - ``max_iterations`` reached
+ - ``score 0.0; cost unknown, 4,995 of 500,000 tokens``
+ - The last gate fails and ``publish`` is ``upstream_failed``.
+
+The converged row is the run above. The dollar row ran the same models with a
$0.01 cap and a
+threshold of 0.99, so the budget ran out first. The token and
``max_iterations`` rows ran
+pydantic-ai's ``test`` model, which has no price, the token row with a
600-token cap.
+
+Reaching ``max_iterations`` fails the last gate with
``LoopMaxIterationsExceeded``, carrying the
+last reason:
+
+.. code-block:: text
+
+ Loop condition remained false at max_iterations=20: score 0.0; cost
unknown, 4,995 of 500,000 tokens
+
+Stopping on a budget is a successful end of the loop, so ``publish`` runs and
returns the last
+draft even though it did not pass. Check ``passes[-1]["score"]`` in
``publish`` if a draft below
+the threshold should not go out.
+
+A per-task ``usage_limits`` that trips is not a budget stop. The task fails
with
+``UsageLimitExceeded``, the gate does not run, and the loop fails. The task
log shows what the
+call spent before it stopped:
+
+.. code-block:: text
+
+ LLM run failed: requests=0, tool_calls=0, input_tokens=118,
output_tokens=4, total_tokens=122
+ Task failed with exception
+
+The exception reads ``Exceeded the total_tokens_limit of 100
(total_tokens=122)``.
+
+Set each task's limits above what a normal pass spends, so they only catch a
runaway call.
+
+What the totals count
+---------------------
+
+**Either cap can be overshot by one pass.** The pass that crosses a cap has
already run when the
+gate checks it: a 600-token cap stopped at 745 after a pass of 250 tokens. The
two tasks'
+``usage_limits`` bound one pass, give or take one model response per task,
since each limit is
+checked after every response; the 122 tokens against a limit of 100 above are
one such response.
+
+**The dollar total is an estimate from list prices.** ``cost`` in the
``usage`` XCom is
+pydantic-ai's ``RunUsage.cost``, which `genai-prices
<https://github.com/pydantic/genai-prices>`__
+computes from the token counts and the model name at the provider's list
prices. The provider is
+matched by its API URL first and by its name after that, so a call through a
gateway that
+``genai-prices`` does not know is still priced at the vendor's list rates.
When a provider reports
+what it billed, as OpenRouter does, pydantic-ai keeps that amount in the
response's
+``provider_details`` and the ``usage`` XCom still carries the estimate. Token
counts are the ones
+the provider reported.
+
+**An unpriced model switches the dollar cap off.** A model ``genai-prices``
does not know, such
+as a private fine-tune, a model newer than your ``genai-prices`` data, or
pydantic-ai's ``test``
+model, reports ``cost=None``. The first pass with such a call sets the loop's
``usd`` to ``None``
+for the rest of the loop and logs a warning in ``tally``:
+
+.. code-block:: text
+
+ A model call had no price, so the dollar cap no longer applies to this
loop.
+
+From then on the gate's reasons say ``cost unknown``, and only the token cap
can stop the loop.
+One unpriced model is enough: with a priced writer and the unpriced ``test``
judge, ``usd`` was
+``None`` from the first pass and the loop ended on
+``token budget spent: score 0.0; cost unknown, 2,063 of 2,000 tokens``.
+
+The same call also leaves that task's own ``cost_limit`` unenforced. The task
log warns and the
+task carries on:
+
+.. code-block:: text
+
+ A `cost_limit` is set but cannot be enforced because no cost was
calculated for this run. ...
+
+Set a token cap even when the budget you care about is dollars. Counting the
unknown cost as zero
+would report less than was spent, and with no priced model at all the dollar
cap would never trip.
+
+**Spend from an earlier try of a task is not in the totals.** The ``usage``
XCom holds the last
+try's usage only. In a run where ``write_draft`` failed after its model call
and passed on retry,
+each try used 127 tokens and the loop counted 127. A retry of ``write_draft``
or ``judge`` after a
+try that reached the model leaves that try's spend out of the totals.
+
+Adapting it
+-----------
+
+* Give the writer and the judge different models with ``model_id``, for
example a cheaper judge.
+ Mixing a priced and an unpriced model switches the dollar cap off, as
described above.
+* Raise ``max_iterations`` for harder drafts. Reaching it fails the loop, so
keep it high enough
+ that the threshold or a budget normally ends the loop first.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_loop_budget.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_loop_budget.py
new file mode 100644
index 00000000000..1c8faef123f
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_loop_budget.py
@@ -0,0 +1,141 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Revise a draft until a judge passes it, within a dollar and a token budget,
using a task loop."""
+
+from __future__ import annotations
+
+# [START howto_loop_budget]
+import logging
+from decimal import Decimal
+
+from pydantic import BaseModel, Field
+
+from airflow.sdk import LoopDecision, dag, task, task_group
+
+log = logging.getLogger(__name__)
+
+THRESHOLD = 0.9
+MAX_USD = Decimal("5.00")
+MAX_TOKENS = 500_000
+
+INCIDENT_NOTES = """\
+09:12 UTC p95 latency on the orders API jumps from 180 ms to 4 s.
+09:20 Checkout errors at 12% in eu-west. Other regions normal.
+09:41 Cause: a connection pool limit lowered in this morning's config deploy.
+09:55 Config rolled back. Error rate falling, 2% at 10:05.
+Next update at 10:30 UTC.
+"""
+
+JUDGE_INSTRUCTIONS = """\
+Score a customer status page update from 0 to 1 against this style guide, and
give one
+sentence of feedback on what to fix. Score 1 only if every rule holds.
+
+1. Every fact matches the incident notes.
+2. The first sentence says which customers are affected and how.
+3. No internal terms: no metric names, config, deploys, connection pools or
rollbacks.
+4. At most 80 words.
+5. It ends with the time of the next update.
+"""
+
+
+class Score(BaseModel):
+ """The judge's verdict on one draft."""
+
+ score: float = Field(ge=0, le=1)
+ feedback: str
+
+
+@dag(tags=["example"])
+def example_agent_loop_budget():
+ @task_group
+ def refine():
+ @task.agent(
+ llm_conn_id="pydanticai_default",
+ system_prompt="You write status page updates for customers. Plain
words, no internal jargon.",
+ # Bounds one pass: the loop's own caps are only checked between
passes.
+ usage_limits={"total_tokens_limit": 20_000, "cost_limit": "0.50"},
+ )
+ def write_draft(*, loop):
+ if loop.index == 1:
+ return f"Write a status page update from these incident
notes:\n{INCIDENT_NOTES}"
+ return (
+ f"Revise this status page update.\nFeedback:
{loop.previous['feedback']}\n"
+ f"Draft:\n{loop.previous['draft']}\nIncident
notes:\n{INCIDENT_NOTES}"
+ )
+
+ @task.llm(
+ llm_conn_id="pydanticai_default",
+ output_type=Score,
+ # A dict in XCom, so ``tally`` can read it under ``airflow dags
test`` too.
+ serialize_output=True,
+ system_prompt=JUDGE_INSTRUCTIONS,
+ usage_limits={"total_tokens_limit": 10_000, "cost_limit": "0.20"},
+ )
+ def judge(draft):
+ return f"Incident notes:\n{INCIDENT_NOTES}\nUpdate:\n{draft}"
+
+ @task
+ def tally(draft, verdict, *, loop, ti):
+ # Inside a loop, this pulls the usage of the current iteration
only.
+ usage = ti.xcom_pull(task_ids=["refine.write_draft",
"refine.judge"], key="usage")
+ spent = {"tokens": 0, "usd": "0"} if loop.index == 1 else
loop.previous
+ tokens = spent["tokens"] + sum(u["total_tokens"] for u in usage)
+ costs = [u["cost"] for u in usage]
+ if spent["usd"] is None:
+ usd = None
+ elif None in costs:
+ # An unpriced call makes the dollar total unknown, not lower.
+ log.warning("A model call had no price, so the dollar cap no
longer applies to this loop.")
+ usd = None
+ else:
+ usd = str(Decimal(spent["usd"]) + sum(map(Decimal, costs)))
+ return {
+ "draft": draft,
+ "score": verdict["score"],
+ "feedback": verdict["feedback"],
+ "tokens": tokens,
+ "usd": usd,
+ }
+
+ draft = write_draft()
+ tally(draft, judge(draft))
+
+ def good_enough(*, loop):
+ result = loop.result
+ usd = "cost unknown" if result["usd"] is None else f"${result['usd']}
of ${MAX_USD}"
+ progress = f"score {result['score']}; {usd}, {result['tokens']:,} of
{MAX_TOKENS:,} tokens"
+ if result["score"] >= THRESHOLD:
+ return LoopDecision(stop=True, reason=f"passed: {progress}")
+ if result["usd"] is not None and Decimal(result["usd"]) >= MAX_USD:
+ return LoopDecision(stop=True, reason=f"dollar budget spent:
{progress}")
+ if result["tokens"] >= MAX_TOKENS:
+ return LoopDecision(stop=True, reason=f"token budget spent:
{progress}")
+ return LoopDecision(stop=False, reason=progress)
+
+ @task
+ def publish(passes):
+ last = passes[-1]
+ usd = "cost unknown" if last["usd"] is None else f"${last['usd']}"
+ log.info("%d passes, score %s, %s tokens, %s", len(passes),
last["score"], last["tokens"], usd)
+ return last["draft"]
+
+ refinement = refine.loop(max_iterations=20, until=good_enough)
+ publish(refinement["tally"].output)
+
+
+example_agent_loop_budget()
+# [END howto_loop_budget]
