pankajkoti commented on code in PR #75516: URL: https://github.com/apache/airflow/pull/75516#discussion_r4364309988
########## 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=5, until=good_enough) Review Comment: I tried the arithmetic on the shipped values and I might be misreading something. The per-task `usage_limits` cap a pass at 30k tokens and $0.70 (20k and $0.50 for the writer, 10k and $0.20 for the judge). With `max_iterations=5`, the loop can spend at most about 150k tokens and $3.50, which is under `MAX_TOKENS` and `MAX_USD`. So with these constants the budget branches in `good_enough` can't fire, while the intro says the loop stops once it has spent $5.00 or 500,000 tokens. The page does explain that the table rows used lowered caps, which helps. Would a higher `max_iterations` make the default example match its description better? Totally your call, and I can see an argument for keeping the values as placeholders. ########## providers/common/ai/docs/use_cases/refine_a_draft_within_a_budget.rst: ########## @@ -0,0 +1,190 @@ + .. 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=5, 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=5): score 0.7; $0.008118 of $5.00, 1,498 of 500,000 tokens + Loop iteration 2: continue (max_iterations=5): score 0.85; $0.016206 of $5.00, 3,194 of 500,000 tokens + Loop iteration 3: stop (max_iterations=5): passed: score 0.9; $0.024117 of $5.00, 4,823 of 500,000 tokens + 3 passes, score 0.9, 4823 tokens, $0.024117 + +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.024117 of $5.00, 4,823 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, 1,245 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=5: score 0.0; cost unknown, 1,245 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: + +.. code-block:: text + + Failure caused by Exceeded the total_tokens_limit of 100 (total_tokens=122). Consider raising the limit, ... Review Comment: The page introduces this as what the task log shows. I couldn't find "Failure caused by" in Airflow core, the Task SDK or the provider. The only match is the example listener in `airflow-core/src/airflow/example_dags/plugins/event_listener.py`, which Airflow loads when `core.load_examples` is on. If your run had that on, that would explain the line, but a setup without it would show the exception in the traceback instead. Would you like to quote the `UsageLimitExceeded` exception line, or mention that the sample comes from a run with the example listener loaded? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
