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 task loop 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¶
Task loops –
refine.loop(max_iterations=20, until=good_enough)repeats the task group, andgood_enoughreturns aLoopDecisionwith a reason.The
usageXCom that@task.agentand@task.llmpush, read back inside the loop to keep a running total of tokens and estimated cost.Usage limits on each task, which caps what one pass can spend.
Structured output and XCom – 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.
Install the provider with the extra for your vendor:
pip install "apache-airflow-providers-common-ai[anthropic]>=1.1.0"
Trigger the Dag:
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:
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:
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()
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:
Ending |
Reason on the last gate |
Result |
|---|---|---|
Converged |
|
The gate succeeds and |
Dollar budget spent |
|
The gate succeeds and |
Token budget spent |
|
The gate succeeds and |
|
|
The last gate fails and |
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:
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:
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
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:
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:
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_iterationsfor harder drafts. Reaching it fails the loop, so keep it high enough that the threshold or a budget normally ends the loop first.