kaxil opened a new pull request, #75516: URL: https://github.com/apache/airflow/pull/75516
Adds a Common AI example Dag and a use-case page for a common ask: have an LLM keep improving something until a judge says it is good enough, without the loop as a whole spending more than a set amount. `example_agent_loop_budget` has `@task.agent` write a customer status page update and `@task.llm` score it against a style guide. A task loop repeats both until the score reaches 0.9, or until the loop has spent $5.00 or 500,000 tokens. It needs Airflow 3.4 for task loops and Common AI 1.1.0, the first release in which `@task.llm` pushes the `usage` XCom. Nothing new is added to the API. The pattern uses what already exists: the loop's `until` returning a `LoopDecision`, the `usage` XCom every Common AI operator pushes, and per-task `usage_limits`. A terminal `tally` task pulls the current iteration's `usage` XComs and adds them to the totals in `loop.previous`. The gate checks both caps after every pass and returns a reason that names how the loop ended. ## Design rationale **An unknown cost is never counted as zero.** pydantic-ai reports `cost=None` for a model genai-prices can't price: a private fine-tune, a model newer than the price data, or the `test` model. The first such call sets the loop's dollar total to `None` for the rest of the loop and logs a warning, and the gate's reasons say `cost unknown`. From then on only the token cap can stop the loop. One unpriced model is enough to switch the dollar cap off, so the page tells authors to set a token cap even when they budget in dollars. **The judge uses `serialize_output=True`.** Without it `tally` receives a `Score` instance. That works under the scheduler, but under `airflow dags test` the task fails with `Score was not found in allow list for deserialization imports`, because `dag.test()` never runs the worker-side registration of `output_type` classes. The use-case pages tell readers to run `airflow dags test`, so the example passes a dict. Fixing `dags test` itself is a separate core change. **`loops.rst` gains one sentence:** a `ti.xcom_pull` of another task in the same loop returns that task's value from the current iteration. The example relies on it, and the page did not say what an in-loop pull returns. ## Runs Every row ran the final example under a real scheduler with LocalExecutor. Only the caps, threshold and models changed between runs: Claude Sonnet 4.6 where a cost is shown, pydantic-ai's unpriced `test` model where it says `cost unknown`. | Ending | Passes | Reason on the last gate | Run | |---|---|---|---| | Converged | 3 | `passed: score 0.9; $0.023475 of $5.00, 4,757 of 500,000 tokens` | success | | Dollar budget, $0.01 cap | 2 | `dollar budget spent: score 0.9; $0.016362 of $0.01, 3,226 of 500,000 tokens` | success | | Token budget, 600-token cap, both models unpriced | 3 | `token budget spent: score 0.0; cost unknown, 745 of 600 tokens` | success | | Token budget, 2,000-token cap, priced writer and unpriced judge | 3 | `token budget spent: score 0.0; cost unknown, 2,063 of 2,000 tokens` | success | | `max_iterations` reached | 5 | `score 0.5; $0.048030 of $5.00, 8,882 of 500,000 tokens` | last gate fails with `LoopMaxIterationsExceeded` | In every iteration of every run, `tally`'s token total equals the running sum of that iteration's own `usage` XComs, which confirms the in-loop pull. ## Known issues - Either cap can be overshot by one pass: the pass that crosses it has already run when the gate checks. A 600-token cap stopped at 745 after a 250-token pass. Per-task `usage_limits` bound that pass, give or take one model response. - The `usage` XCom holds the last try only, so the spend of an earlier try of a retried body task is missing from the totals. A forced retry spent 127 tokens on each try and the loop counted 127. - The dollar total is the genai-prices estimate from list prices, not the provider's bill. --- * Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#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](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed. * When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x). * For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number. -- 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]
