kaxil opened a new pull request, #73984: URL: https://github.com/apache/airflow/pull/73984
`AgentOperator` and `@task.agent` now take pydantic-ai [capabilities](https://ai.pydantic.dev/capabilities/) directly: ```python AgentOperator( task_id="reasoner", prompt="...", llm_conn_id="pydanticai_default", capabilities=[Thinking(effort="high"), WebSearch()], ) ``` Capabilities are pydantic-ai's unit for adding behavior to an agent (tools, instructions, model settings, and hooks around each model request or tool call), and the way guardrail packages such as `pydantic-ai-shields` plug in. Until now the only route was `agent_params={"capabilities": [...]}`. `agent_params` is a template field, so every capability ended up in the serialized Dag as its repr, and for a capability holding a function that repr includes a memory address. The same Dag, run once each way: | `agent_params={"capabilities": [...]}` | `capabilities=[...]` | |---|---| |  |  | `capabilities=` is not a template field, like `toolsets=`, so nothing about the capability reaches the serialized Dag; the worker builds it from the Dag file. A `Toolset` capability holding one of this provider's toolsets still has its connection IDs templated per task instance, the same way `toolsets=` does. ## Design rationale **`capabilities` inside `agent_params` keeps working, but not together with the new argument.** pydantic-ai wraps capability hooks in list order, first outermost, so merging two lists would pick an order the Dag author never wrote. Passing both fails the task with a `ValueError`. The check runs when the agent is built rather than in `__init__`, because `agent_params` can still be an unrendered template or an `XComArg` at construction time. **Code mode is now recognized however it is passed.** The `durable=True` check and the per-tool approval check only looked at `code_mode=True`, so a `CodeMode()` capability got past both: `durable=True` accepted it, and approval stayed on alongside code mode. Both now find `CodeMode` at the top level, inside a `CombinedCapability`, or inside a wrapper such as `PrefixTools`. `code_mode=True` plus a `CodeMode()` capability would add two, so that combination is refused too. The lookup never imports `pydantic-ai-harness`, so an agent that does not use code mode is unaffected when the harness is installed without its `code-mode` extra. **Agent-only for now.** `LLMOperator` and its siblings still pass `agent_params` straight to `Agent(...)`, so capabilities there keep working the old way. ## Docs The "Guardrails" page becomes "Capabilities and guardrails" (`guardrails.html` redirects), with a table of what a retry replays under `durable=True` for each kind of capability:  ## Gotchas - A mapped task (`AgentOperator.partial(...).expand(...)`, or a mapped `@task.agent`) stores every `partial` argument, so there `capabilities` is serialized as a repr. - A `CodeMode` built by a capability function when the run starts cannot be seen before the run, so the `durable=True` check does not catch it. -- 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]
