kaxil commented on code in PR #75516: URL: https://github.com/apache/airflow/pull/75516#discussion_r4364884491
########## 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: Good find. My run had `load_examples` on, so that line came from the example listener. The page now quotes the provider's own lines (`LLM run failed: ... total_tokens=122` and `Task failed with exception`) and gives the exception's message separately. -- 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]
