kaxil opened a new pull request, #73578:
URL: https://github.com/apache/airflow/pull/73578
A mapped `@task.agent` can now give each task instance its own connection.
The case that needs it is customer-facing analytics: each customer's rows sit
behind their own database role, so each customer's agent has to query through
that customer's connection. `@task.llm_sql` already templates `db_conn_id`, but
toolsets took their connection when the Dag was parsed, so every mapped agent
shared one.
```python
@task.agent(
llm_conn_id="pydanticai_default",
toolsets=[SQLToolset(db_conn_id="analytics_{{ task.op_kwargs.customer
}}")],
)
def report(customer: str) -> str:
return f"Summarize this month's orders for {customer}."
report.expand(customer=customers())
```
What is templated: `SQLToolset.db_conn_id`, `MCPToolset.mcp_conn_id`, and a
`HookToolset`'s hook connection ID. `HookToolset` reads the attribute the
hook's `conn_name_attr` names (`postgres_conn_id`, ...), and falls back to
`conn_id` for hooks such as `WasbHook` that keep it there. Only connection IDs
are templated.
## Design rationale
**Why not add `toolsets` to `AgentOperator.template_fields`?** Template
fields are serialized into the Dag. A toolset's repr is what gets serialized,
and pydantic-ai's wrappers (`.prefixed()`, `.filtered()`) are dataclasses whose
repr can embed function addresses, so the Dag hash would change on every parse.
`toolsets` stays out of `template_fields`, and the operator renders the
toolsets itself.
**Each task instance renders a copy.** `MappedOperator.unmap` passes the
partial's toolset objects straight to every unmapped task, and `dag.test()`
runs every task in one process. Rendering in place would hand map index 0's
connection to map index 1. Each opt-in leaf toolset is copied before rendering
(for `HookToolset`, the hook too). Wrappers and `Toolset` capabilities are
walked with pydantic-ai's `visit_and_replace`, and only when something in them
is templated, so an untemplated custom wrapper is never rebuilt.
**The opt-in attribute is `agent_template_fields`, not `template_fields`.**
Airflow's templater renders any object that carries `template_fields` in place,
wherever it is nested inside another template field. `agent_params` is a
template field, and `agent_params["toolsets"]` is a supported way to pass
toolsets, so the familiar name would bring the leak back through that path.
Third-party toolsets opt in by declaring the same attribute.
**Rendering hangs off `_do_render_template_fields`.** A mapped task never
calls `render_template_fields`: `MappedOperator` renders the unmapped task
through `_do_render_template_fields`. `KubernetesPodOperator` hooks the same
method for the same reason.
The rendered connection is not recorded anywhere else, so each task instance
logs it once, e.g. `Rendered toolset sql-analytics_acme`.
## Screenshots
A two-customer demo on a real scheduler and API server, each customer with
its own SQLite connection, and a `test` model that calls every tool. Map index
0 renders the `acme` connection for both the SQL toolset and the hook toolset;
map index 1 renders `globex`:
Each agent's tool results come from its own customer's database:
## Gotchas
- **Build the connection ID from upstream task output, not `params` or
`dag_run.conf`.** A task can read any connection it names, so a template driven
by trigger input lets whoever triggers the Dag pick the database (for an MCP
`stdio` connection, the command that runs on the worker). The docs say so next
to each example.
- `{{ customer }}` does not work: the task's arguments are not template
variables. Use `{{ task.op_kwargs.customer }}`. With
`AgentOperator.partial(...).expand(prompt=...)`, the connection has to come
from the map index; the docs show that form and its ordering caveat.
- `HookToolset.id` now includes the connection ID
(`hook-PostgresHook-analytics_acme`, previously `hook-PostgresHook`). The
toolset id is part of the durable-execution step fingerprint, so a
`durable=True` task that retries across the upgrade misses its cache once.
- Not templated: `allowed_tables` (validated when the toolset is created, so
a template stays a literal), `DataFusionToolset`, a `Toolset` capability built
from a callable, and hooks that keep their connection ID under some other
attribute. A hook that looks its connection up in `__init__` fails at Dag parse
time, because the template is not a connection ID yet.
---
* 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]
