kaxil opened a new pull request, #73367:
URL: https://github.com/apache/airflow/pull/73367
`LLMBranchOperator` shows the model each branch as its task ID and nothing
else. That works when the IDs speak for themselves and the prompt clearly fits
one. It stops working when two branches could plausibly own the same input: in
the provider's own example, "my password reset email never arrived" is a
sign-in problem to one team and an email problem to another, and nothing tells
the model which team owns it. The only place to explain the branches today is
`system_prompt`, as prose the model has to match back to a task ID by name.
This adds `branch_descriptions`, a mapping of downstream task ID to a short
description of what choosing that branch means. The descriptions travel in the
output schema next to the option they describe, so the model reads each option
together with its meaning.
```python
LLMBranchOperator(
task_id="route_ticket",
prompt="User says: 'My password reset email never arrived.'",
llm_conn_id="pydanticai_default",
system_prompt="Route the ticket to the team responsible for resolving
it.",
branch_descriptions={
"handle_auth": "Sign-in, passwords, 2FA and lockouts. This team owns
missing password-reset emails.",
"handle_billing": "Invoices, charges, refunds and plan changes.",
"handle_general": "General support triage and tickets that need
clarification.",
},
)
```
Stacked on #73366, which sorts the option order the enum is built in; this
change builds on the same lines.
## Design rationale
**The encoding is the one JSON Schema shape that carries a description per
value.** A JSON Schema `enum` has no slot for it. `anyOf` of `{const,
description}` does, and it is what both sides read: a text model receives it in
the output tool's parameters schema, and pydantic-ai's TypeSafe adapter reads
an option's meaning from exactly that shape (`_options()` in
`pydantic_ai/models/typesafe.py`: `enum` gives options with no descriptions,
`anyOf` of `const` gives one per option). So the enum the operator already
builds gets a `__get_pydantic_json_schema__` that renders `anyOf`, and nothing
else changes: validation is still the enum, the output is still an enum member,
and the `response` wrapper pydantic-ai puts around a bare output type stays as
it was. Without the parameter the schema is the bare `enum` it always was.
**Not a union of described `Literal`s, and not a wrapper model.** The union
was the obvious route and it fails: pydantic-ai treats a top-level `X | Y`
output type as two output tools and keeps only the first, so a two-branch Dag
would validate only the first branch. Wrapping the union in a one-field model
works but renames the tool parameter for every existing user and changes what a
classifier model sees as the question. Keeping the enum and only changing its
schema avoids both.
**Explicit mapping only, no `doc_md` fallback.** Reading `doc_md` off the
downstream tasks was tempting, and under TaskFlow most tasks already have one
because the decorator copies the docstring in. But that text was written for
the UI, carries Markdown and run notes, and a tidied docstring would move a
branch decision with nothing in the diff near the branch operator. With an
explicit mapping, what the model reads is what the Dag author wrote for it, and
a rename shows up in the diff.
**A bad key fails before the model call.** A key that is not a downstream
task ID raises `ValueError` naming the key and the valid IDs, at the top of
`execute`, because downstream tasks are attached after construction and cannot
be checked in `__init__`. It is an error rather than a warning because a
misspelled key that silently becomes an option with no description is the case
this parameter is meant to remove. Unlisted tasks are presented by ID alone, as
today, so a mapping can cover only the branches that need explaining.
**A description changes what the model is asked, not how sure it is.** The
docs say so, because a text model's structured output carries no confidence
value and there is nothing to gate on.
## Verified against running models
Four Dags through a local Airflow, the guide's ticket with and without
descriptions, on `anthropic:claude-sonnet-5` and on `typesafe:jev-1.13.0`. All
four branched to `handle_auth` and skipped the other two. The descriptions
reaching the model shows in the input token counts: 802 with them against 671
without on Sonnet, 403 against 328 on Jev. Jev's probabilities for this ticket
were 1.00 / 0.00 / 0.00 both ways (confidence 0.99 plain, 1.00 described), so
on this prompt the descriptions confirmed a pick that was already certain
rather than changed it. That is one prompt and one run per variant, and the
docs make no accuracy claim from it.
## Gotchas
`branch_descriptions` is a template field, so Jinja in the descriptions
renders at execute time like `system_prompt`. Keys are full task IDs, so a
branch inside a TaskGroup is `group.task`.
The unit tests capture the schema through a real pydantic-ai `Agent` with
`FunctionModel`, which sits where every provider adapter sits and receives the
same `output_tools` schema. The TypeSafe adapter itself is not exercised in CI
because the image carries pydantic-ai 2.31.1; its reading of
`anyOf`/`const`/`description` is from its source at 2.45.0 and the live run
above.
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* 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]
