kaxil opened a new pull request, #73532:
URL: https://github.com/apache/airflow/pull/73532
Airflow has five ways to invoke a vendor-hosted agent (Bedrock AgentCore,
Vertex AI Agent Engine, Azure AI Foundry, Snowflake Cortex, Anthropic Managed
Agents) and none of them share an argument name, a result shape or an error
class. `BaseManagedAgentToolset` from #71946 was supposed to fix that. A month
later it has zero implementations. The reason is structural: adopting it means
subclassing a pydantic-ai toolset from inside a provider that still supports
Airflow 2, which nobody can do.
So this moves the common surface to where `DbApiHook` put it: the hook.
`airflow.providers.common.ai.managed_agents` defines `BaseManagedAgentHook`, a
mixin a vendor hook adopts next to its own base. The agent is an argument, like
a statement is to `DbApiHook.run`. Three methods (`resolve_agent`,
`agent_capabilities`, `invoke_agent`), typed request/response dataclasses, no
pydantic-ai import. Amazon adopts it for AgentCore runtimes, Google for Agent
Engine deployments. On top, written once: `ManagedAgentToolset` (any client of
the contract becomes one tool for a calling model) and
`FailoverManagedAgentClient` (clients composed across clouds).
```python
claims = BedrockAgentCoreManagedAgentHook(aws_conn_id="aws_prod",
region_name="us-east-1").agent(RUNTIME_ARN)
AgentOperator(
task_id="triage",
llm_conn_id="anthropic_default",
prompt="Review claim 4411 and decide whether to pay it.",
toolsets=[ManagedAgentToolset(claims, tool_name="ask_claims_agent",
description="Reviews an insurance claim.")],
)
```
## Design rationale
**Vendors take common.ai as an optional extra, exactly as they take
common.messaging today.** common.ai needs Airflow 3.0; amazon and google still
floor 2.11. There is already a pattern in the repo for a 3.x-only common base
adopted by 2.x providers: `BaseMessageQueueProvider`. Guarded import raising
`AirflowOptionalProviderFeatureException`, an extra in `pyproject.toml`,
`pytest.importorskip` in the tests. The two adoptions here follow
`aws/queues/sqs.py`. Plain `amazon` installs are untouched; pydantic-ai only
arrives with `amazon[common.ai]`.
**Type the intersection, pass the rest through.** Identity, prompt or
messages, an optional session, a timeout, a text answer: those mean the same
thing everywhere and are typed. The rest rides in `vendor_options` in and `raw`
out. Each hook rejects options that would point the call somewhere else
(another ARN, account, engine). The model sees `text`. Python callers get `raw`.
**Why `ManagedAgentRejected` instead of `ModelRetry`.** The contract has to
stay free of pydantic-ai or the guarded import stops being cheap. The toolset
translates at the boundary. As it turns out neither adopted platform has a
rephrase-class error anyway: AgentCore returns a container's complaints inside
a 200, and Agent Engine's `INVALID_ARGUMENT` is an author mistake (wrong input
key), not something a model can fix by rewording. Both hooks raise terminal
errors or let transient ones through.
**Failover over clients, not toolsets.** The removed
`FailoverManagedAgentToolset` took toolsets as members but only ever called
`invoke` on them, so a member's `tool_name`, `description` and `max_retries`
did nothing. It also resolved every member's identity before the call, so a
standby with a bad connection failed calls the primary would have answered. The
new group takes clients, resolves identity only for logs and metric tags, logs
which member answered, and refuses any request carrying a `session_id` (a
failover starts a fresh conversation on the standby). A bound agent applies the
same rule for a single hook whose capabilities do not include sessions.
**`replayable` now does something.** #71946 declared it; nothing read it;
`durable=True` replayed managed-agent calls from cache regardless.
`CachingToolset` now runs the tool instead of replaying when the wrapped
toolset says `replayable=False`, and it looks through `prefixed`, `filtered`
and `CombinedToolset`, which do not carry the attribute.
## Migration path
common.ai is 0.x and the old base had no subclasses in tree. Still, for
anyone who tried it:
- `FailoverManagedAgentToolset` is gone. Use
`ManagedAgentToolset(FailoverManagedAgentClient([...]), ...)`.
- `BaseManagedAgentToolset.agent_ref` returns a `ManagedAgentRef` rather
than a `dict[str, str]`.
- `managed_agent.served` loses `role` and `position`, gains `tool`.
`managed_agent.failover` loses `tool`. `served` counts answers, `failover`
counts transitions; read them next to each other. Which member answered is in
the task log.
The new amazon and google extras carry the `# use next version` marker for
the release manager.
## Gotchas
- Bedrock: one boto3 client per distinct request timeout, reused, so the
toolset's `timeout` becomes connect/read timeout without a new session per tool
call. Retries off unless the connection or caller set them. The connection's
`config_kwargs` apply as for any AWS hook. Session ids are checked against
AgentCore's 33 to 256 character bound before the call. Non-JSON, malformed or
oversize bodies are terminal errors that name the agent and connection.
- Agent Engine: the query path keeps no conversation state, so a
`session_id` is refused rather than silently sent as a fresh call. Query jobs
stay on `RunQueryJobOperator`, which can defer.
- Not touched: Azure Foundry, Cortex, Anthropic. Cortex has PRs in flight
and is the obvious next adoption. Anthropic's session-shaped agent wants an
operator, not a toolset.
Supersedes the toolset-only design of #71946.
---
* 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]
