This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new dd607da5ab3 Add Anthropic tool runner adapter for `common.ai` toolsets
(#73962)
dd607da5ab3 is described below
commit dd607da5ab32b0cfda55f8c508788fba1a4a0924
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Oct 14 11:56:56 2026 +0100
Add Anthropic tool runner adapter for `common.ai` toolsets (#73962)
Agents written on the Anthropic SDK's client.beta.messages.tool_runner can
now
use common.ai's connection-backed toolsets (SQLToolset, HookToolset,
ObjectStorageToolset, sandboxes) through AirflowTools and AsyncAirflowTools,
with the same guarantees as the Strands and ADK adapters: masked results,
correctable failures returned as is_error results within the retry limit,
and any other failure failing the task before the next model request.
run_coroutine_sync now copies the caller's context into its worker thread,
and the anthropic extra's floor moves to 1.1.0, the first release whose tool
runner only runs a turn's tool calls when the turn ends in tool_use.
* Tighten the Anthropic tool runner guide
* Narrow the message block type in the Anthropic adapter tests
* Mention toolset max_retries in the Anthropic tool runner guide
The toolsets now take max_retries (#73957), so the guide no longer says they
allow exactly one correction, and the retry budget docs list the Anthropic
tool runner among the bridges that have no agent budget. The frameworks page
also names the Anthropic SDK among the frameworks a team can keep.
* Fail the Anthropic tool runner adapter's run at max_iterations mid-task
A runner that reaches max_iterations while the model is still calling tools
ends on that tool_use turn, which holds no answer. tools.run now raises
ToolRunnerMaxIterationsError instead of returning it, so the task fails
rather than succeeding with an empty result.
---
providers/common/ai/docs/frameworks/anthropic.rst | 212 ++++++++++++
providers/common/ai/docs/frameworks/index.rst | 23 +-
providers/common/ai/docs/index.rst | 2 +-
providers/common/ai/docs/observability.rst | 3 +-
providers/common/ai/docs/sandbox/index.rst | 10 +-
providers/common/ai/docs/stability.rst | 3 +
providers/common/ai/docs/toolsets/index.rst | 6 +-
.../common/ai/docs/toolsets/object_storage.rst | 3 +-
providers/common/ai/pyproject.toml | 8 +-
.../example_dags/example_anthropic_tool_runner.py | 82 +++++
.../src/airflow/providers/common/ai/exceptions.py | 10 +
.../airflow/providers/common/ai/tools/__init__.py | 4 +-
.../airflow/providers/common/ai/tools/anthropic.py | 288 ++++++++++++++++
.../providers/common/ai/utils/coroutines.py | 13 +-
.../providers/common/ai/utils/tool_metrics.py | 2 +-
.../ai/tests/unit/common/ai/test_exceptions.py | 2 +
.../tests/unit/common/ai/tools/test_anthropic.py | 364 +++++++++++++++++++++
.../tests/unit/common/ai/utils/test_coroutines.py | 14 +
.../observability/metrics/metrics_template.yaml | 4 +-
uv.lock | 4 +-
20 files changed, 1025 insertions(+), 32 deletions(-)
diff --git a/providers/common/ai/docs/frameworks/anthropic.rst
b/providers/common/ai/docs/frameworks/anthropic.rst
new file mode 100644
index 00000000000..214ee136a8d
--- /dev/null
+++ b/providers/common/ai/docs/frameworks/anthropic.rst
@@ -0,0 +1,212 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/frameworks:anthropic:
+
+Anthropic SDK tool runner
+=========================
+
+If your agent is a loop on the Anthropic Python SDK's
+``client.beta.messages.tool_runner``, you can run it in an Airflow task as it
is and
+give it the toolsets this provider ships.
+:class:`~airflow.providers.common.ai.tools.anthropic.AirflowTools` turns
+``SQLToolset``, ``HookToolset`` and most other toolsets into tools the runner
accepts.
+If you have no loop yet, ``AgentOperator`` runs one for you, with durable
replay and
+human review; see :doc:`index`.
+
+.. note::
+
+ ``AirflowTools`` is experimental: it may change in a minor release of this
provider.
+ The framework-neutral tool interface it builds on is stable.
+ See :ref:`howto/stability`.
+
+Run a tool runner in a task
+---------------------------
+
+This task answers a question about a database. The client's API key, and a
gateway URL
+if you route through one, come from an Airflow connection:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_anthropic_tool_runner.py
+ :language: python
+ :start-after: [START example_anthropic_tool_runner]
+ :end-before: [END example_anthropic_tool_runner]
+
+Against a SQLite database with one table, ``orders``, the model called two
tools. Here
+are the task log's tool-call lines, timestamps trimmed, and the value the task
returned:
+
+.. code-block:: text
+
+ [info] ::group::Tool call: list_tables
+ [info] Tool list_tables returned in 0.00s
+ [info] ::group::Tool call: query
+ [info] Tool query returned in 0.67s
+ [info] Done. Returned value was: The database contains **one table**:
+
+ | Table | Row count |
+ |--------|-----------|
+ | `orders` | 4 |
+
+Only ``AirflowTools`` and ``tools.run(runner)`` come from Airflow. The system
prompt,
+model parameters and ``max_iterations`` are the runner's own, and so are tools
you
+write with the SDK's ``@beta_tool`` decorator; pass those alongside Airflow's,
as
+``tools=[*tools.tools, my_tool]``. ``tools.run`` fails the task only for
Airflow's tools:
+the runner still hands an exception from one of your own tools to the model. A
run that
+reaches ``max_iterations`` while the model is still calling tools has no
answer, so
+``tools.run`` raises
+:class:`~airflow.providers.common.ai.exceptions.ToolRunnerMaxIterationsError`
and the
+task fails rather than returning an empty result.
+
+To reach Claude on Amazon Bedrock, Claude Platform on AWS, Google Vertex AI or
Microsoft
+Foundry, build the
+client with ``AnthropicHook(conn_id=...).get_conn()`` from the Anthropic
provider
+instead. It reads the platform, region and credentials from an Anthropic
connection,
+and every client it returns has the same ``tool_runner``.
+
+What each tool does
+-------------------
+
+``AirflowTools`` accepts any toolset that implements
+:class:`~airflow.providers.common.ai.tools.ToolProvider`, and individual
+:class:`~airflow.providers.common.ai.tools.AirflowTool` objects.
``MCPToolset`` is the
+exception; give the runner the MCP server through the SDK instead. Each tool:
+
+- Keeps the toolset tool's name, description and argument schema.
+- Runs through the toolset, so connection resolution, SQL validation,
+ ``allowed_tables``, ``allowed_methods`` and result limits behave as they do
in
+ :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`.
+- Passes its result through Airflow's secret masker before the model sees it. A
+ connection password that turns up in a query result reaches the model as
``***``.
+- Writes a ``Tool call: <name>`` group to the task log. Calls to a toolset are
also
+ counted in the ``common_ai.tool_calls`` metric with ``framework=anthropic``
(see
+ :doc:`../observability`).
+
+When a tool fails
+-----------------
+
+A failure the model can correct, such as a query naming a table outside
+``allowed_tables``, comes back to it as a ``tool_result`` with ``is_error``
set. The
+model received this block when it was told to count the rows of a ``sales``
table that
+the toolset did not allow:
+
+.. code-block:: json
+
+ {
+ "type": "tool_result",
+ "tool_use_id": "toolu_01PNeqBojrc5S7FnNAeP1brL",
+ "content": "The query tool failed: Query references tables that are not
in the allowed tables list: sales. Use list_tables to see the allowed
tables.\nUse the list_tables and get_schema tools to inspect the database, then
fix the query and try again.",
+ "is_error": true
+ }
+
+It then listed the tables, found ``orders``, and answered from it.
``SQLToolset``,
+``HookToolset`` and ``ObjectStorageToolset`` allow one correction unless you
pass the
+toolset ``max_retries``: a tool that fails in two model turns in a row ends
the run with
+``ToolCallError: query kept failing after 1 correction(s): ...``. Each turn
counts once,
+however many of its calls failed. ``SQLToolset`` treats every database error as
+correctable, so a query rejected because the connection's credentials are
wrong ends the
+run on its second failure.
+
+Any other failure, such as a hook raising, ends the run. ``tools.run`` raises
+:class:`~airflow.providers.common.ai.tools.ToolCallError` before the next model
+request, the task fails, and Airflow's own retry takes over. Here is the error
from a
+run where a tool's call to a billing API raised; the masker has already
replaced the
+token in the message:
+
+.. code-block:: text
+
+ airflow.providers.common.ai.tools.ToolCallError: fetch_invoice failed:
RuntimeError: billing API returned 503 for token ***
+
+Drive the runner with ``tools.run``
+-----------------------------------
+
+The tool runner catches every exception a tool raises and sends it to the
model as an
+error result. If you finish the loop with the runner's own ``until_done()``,
or iterate
+it yourself, a hook that raises becomes one more thing for the model to read,
and the
+task can succeed with an answer written around the failure:
+
+.. code-block:: python
+
+ message = runner.until_done() # a failing hook becomes text the model
reads
+ message = tools.run(runner) # a failing hook fails the task
+
+Airflow's tools log a warning when a failure goes to the model this way.
+
+``tools.run`` runs each turn's tool calls itself, which is also how the retry
limit
+knows where one model turn ends. When a call fails in a way that ends the run,
the
+turn's later calls are not run; the calls before it already have.
+
+``tools.run`` takes the runner that ``client.beta.messages.tool_runner``
returns. It
+does not support the streaming runner you get with ``stream=True``: a type
checker
+rejects the call, and at run time it fails with an ``AttributeError`` after
the first
+model request, before any tool runs.
+
+Async clients
+-------------
+
+For an ``AsyncAnthropic`` client, use
+:class:`~airflow.providers.common.ai.tools.anthropic.AsyncAirflowTools` and
await
+``run``, inside an ``async def`` task or any other coroutine:
+
+.. code-block:: python
+
+ from anthropic import AsyncAnthropic
+
+ from airflow.providers.common.ai.tools.anthropic import AsyncAirflowTools
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+ client = AsyncAnthropic()
+ tools = AsyncAirflowTools(SQLToolset(db_conn_id="warehouse"))
+ runner = client.beta.messages.tool_runner(
+ model="claude-opus-5-5",
+ max_tokens=16000,
+ tools=tools.tools,
+ messages=[{"role": "user", "content": "Which tables exist?"}],
+ )
+ message = await tools.run(runner)
+
+Differences from ``AgentOperator``
+----------------------------------
+
+The loop is the SDK's, so what ``AgentOperator`` builds on top of Pydantic AI
does not
+apply:
+
+- A task retry runs the loop again from the start. Tools that change data must
be safe
+ to repeat, because there is no step-level replay as with
+ ``AgentOperator(durable=True)``.
+- There is no human review step like
``AgentOperator(enable_hitl_review=True)``.
+- A toolset's connection ID is used as written. It is not rendered as a
template, so
+ pass the value rather than a Jinja expression.
+
+Installation
+------------
+
+Install this provider with its ``anthropic`` extra:
+
+.. code-block:: bash
+
+ pip install "apache-airflow-providers-common-ai[anthropic]"
+
+The ``anthropic`` extra installs the SDK. Add
``apache-airflow-providers-anthropic`` if
+you build the client with ``AnthropicHook``. ``tool_runner`` is a beta helper
of the
+SDK. The adapter needs ``anthropic`` 1.1.0 or later, which the extra requires,
+and is tested with 1.8.0.
+
+.. seealso::
+
+ - :doc:`../toolsets/index`, for what each toolset exposes and its limits.
+ - :doc:`index`, for the framework-neutral tool interface ``AirflowTools``
is built on.
+ - `Tool Runner
<https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-runner>`__,
+ in Anthropic's documentation, for the runner itself.
diff --git a/providers/common/ai/docs/frameworks/index.rst
b/providers/common/ai/docs/frameworks/index.rst
index 12b5b254074..62daee3fe01 100644
--- a/providers/common/ai/docs/frameworks/index.rst
+++ b/providers/common/ai/docs/frameworks/index.rst
@@ -50,8 +50,8 @@ AI because it does not tie a Dag to either:
Pydantic AI is the recommendation, not a requirement. For a new agent, start
with
``AgentOperator``, which gets the most from this provider. If your team has
standardized
-on Strands, ADK or LangChain, keep that framework: the table below shows what
this
-provider gives an agent written in each one.
+on Strands, ADK, LangChain or the Anthropic SDK, keep that framework: the
table below
+shows what this provider gives an agent written in each one.
Supported frameworks
--------------------
@@ -91,6 +91,12 @@ is Airflow's and which part stays yours.
ADK agent the same tools, with the same masking. Guide: :doc:`adk`.
- You
- ``google-adk``
+ * - Anthropic SDK tool runner
+ - :class:`~airflow.providers.common.ai.tools.anthropic.AirflowTools`
hands the same
+ tools to ``client.beta.messages.tool_runner`` and makes a tool failure
the model
+ cannot fix fail the task. Guide: :doc:`anthropic`.
+ - You
+ - ``[anthropic]`` extra
* - LangChain
- :class:`~airflow.providers.common.ai.hooks.langchain.LangChainHook`
returns chat
and embedding models configured from a connection, and
@@ -107,9 +113,9 @@ is Airflow's and which part stays yours.
- No agent
- ``[llamaindex]`` extra
-The Strands and ADK integrations and the tracing helper are experimental: they
can change
-or be removed in a minor release of this provider. The framework-neutral tool
interface
-under them is stable. See :ref:`howto/stability`.
+The Strands, ADK and Anthropic integrations and the tracing helper are
experimental: they
+can change or be removed in a minor release of this provider. The
framework-neutral tool
+interface under them is stable. See :ref:`howto/stability`.
Tested versions
---------------
@@ -145,6 +151,8 @@ Start from the agent you already have:
``AirflowTools`` plugin. See :doc:`strands`.
- **An existing ADK agent**: keep it, and add Airflow's toolsets with the ADK
``AirflowTools`` toolset. See :doc:`adk`.
+- **A loop written on the Anthropic SDK**: keep it, and pass the tools of
+ ``AirflowTools`` to its tool runner. See :doc:`anthropic`.
- **An existing LangChain agent**: build its model with ``LangChainHook`` and
add
Airflow's toolsets with ``airflow_toolset_to_langchain_tools``.
- **No agent yet, or no preference**: use ``AgentOperator`` or
``@task.agent``. That route
@@ -162,7 +170,7 @@ Airflow also depends on, such as OpenTelemetry or a vendor
SDK.
To give such a framework Airflow's toolsets rather than hand-written tools,
build on
the framework-neutral interface in :mod:`airflow.providers.common.ai.tools`,
which is
-what the Strands and ADK integrations are written against:
+what the Strands, ADK and Anthropic integrations are written against:
- :class:`~airflow.providers.common.ai.tools.AirflowTool` is one operation: a
name, a
description, a JSON Schema for its arguments and an async function. Call it
through
@@ -182,7 +190,7 @@ what the Strands and ADK integrations are written against:
An adapter maps these onto the framework's own tool type and error status, and
makes
sure a ``ToolCallError`` ends the run rather than reaching the model, as the
Strands
plugin does. :func:`~airflow.providers.common.ai.tools.collect_tools` turns
the toolsets and
-tools an adapter is given into one list, as the Strands and ADK adapters do.
+tools an adapter is given into one list, as the Strands, ADK and Anthropic
adapters do.
Blocking hook calls, such as a SQL query, run one at a time in the task's
process, so an
agent that calls two database tools at once gets its answers one after the
other.
@@ -199,3 +207,4 @@ on; see :doc:`../observability`.
Pydantic AI <pydantic_ai>
Strands Agents <strands>
Google ADK <adk>
+ Anthropic SDK <anthropic>
diff --git a/providers/common/ai/docs/index.rst
b/providers/common/ai/docs/index.rst
index f8b384d9bbf..47ec2a7e043 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -225,7 +225,7 @@ Install them when installing from PyPI. For example:
===============
=======================================================================================================================================
Extra Dependencies
===============
=======================================================================================================================================
-``anthropic`` ``pydantic-ai-slim[anthropic]>=2.36.0``, ``anthropic>=1.0.0``
+``anthropic`` ``pydantic-ai-slim[anthropic]>=2.36.0``, ``anthropic>=1.1.0``
``bedrock`` ``pydantic-ai-slim[bedrock]>=2.36.0``
``google`` ``pydantic-ai-slim[google]>=2.36.0``
``openai`` ``pydantic-ai-slim[openai]>=2.36.0``, ``openai>=2.47.0``
diff --git a/providers/common/ai/docs/observability.rst
b/providers/common/ai/docs/observability.rst
index 48f4c24a56e..bd964c7dd12 100644
--- a/providers/common/ai/docs/observability.rst
+++ b/providers/common/ai/docs/observability.rst
@@ -246,7 +246,8 @@ carries three tags:
- The toolset class, such as ``SQLToolset`` or ``ObjectStorageToolset``.
* - ``framework``
- ``pydantic_ai`` for ``AgentOperator`` and your own Pydantic AI agents;
- ``strands``, ``adk`` or ``langchain`` for a call through that
framework's adapter;
+ ``strands``, ``adk``, ``langchain`` or ``anthropic`` for a call through
that
+ framework's adapter;
``none`` for a call to a toolset's ``airflow_tools()`` made without an
adapter.
* - ``outcome``
- ``executed`` when the call returned, ``failed`` when it raised
(including a
diff --git a/providers/common/ai/docs/sandbox/index.rst
b/providers/common/ai/docs/sandbox/index.rst
index 9e58c25f0c3..ce093fe5735 100644
--- a/providers/common/ai/docs/sandbox/index.rst
+++ b/providers/common/ai/docs/sandbox/index.rst
@@ -496,11 +496,11 @@ for the other routes that build their own tools; see
:ref:`toolset-call-barriers
With another agent framework
----------------------------
-A Strands or Google ADK agent can use the same sandbox through ``AirflowTools``
-(see :doc:`../frameworks/index`). Outside a Pydantic AI run nothing ends the
run for the
-toolset, so the task owns the sandbox's life: open the toolset with ``with``
(or
-``async with``) around the agent, and the sandbox it provisions is destroyed
when the
-block ends, however the agent finishes:
+A Strands or Google ADK agent, or a loop on the Anthropic SDK's tool runner,
can use the
+same sandbox through ``AirflowTools`` (see :doc:`../frameworks/index`).
Outside a
+Pydantic AI run nothing ends the run for the toolset, so the task owns the
sandbox's
+life: open the toolset with ``with`` (or ``async with``) around the agent, and
the
+sandbox it provisions is destroyed when the block ends, however the agent
finishes:
.. code-block:: python
diff --git a/providers/common/ai/docs/stability.rst
b/providers/common/ai/docs/stability.rst
index 7877e20b709..5b2e2a61f63 100644
--- a/providers/common/ai/docs/stability.rst
+++ b/providers/common/ai/docs/stability.rst
@@ -257,6 +257,9 @@ Everything this provider ships that is not in the table
above is experimental.
(:doc:`frameworks/index`)
- Written against Strands 1.56 and ADK 2.9.1, which change often. Strands
marks its
own checkpointing experimental.
+ * - The Anthropic tool runner adapter (:doc:`frameworks/anthropic`)
+ - Built on ``tool_runner``, a beta helper of the Anthropic SDK that can
change between
+ SDK releases.
* - The ``common_ai.tool_calls`` metric and
:func:`~airflow.providers.common.ai.tools.tracing.agent_framework_tracing`
(:doc:`observability`)
diff --git a/providers/common/ai/docs/toolsets/index.rst
b/providers/common/ai/docs/toolsets/index.rst
index be80e9313c0..c69bfffddfb 100644
--- a/providers/common/ai/docs/toolsets/index.rst
+++ b/providers/common/ai/docs/toolsets/index.rst
@@ -294,9 +294,9 @@ used to allow exactly one correction whatever ``retries``
said, so an agent that
``retries`` now applies it to them too: ``retries=0`` fails the run on the
first bad
query, and a large integer ``retries`` meant for output validation also lets a
failing
database be queried that many times. Pass ``max_retries=1`` to a toolset to
keep the old
-behaviour. Outside a pydantic-ai agent (the LangChain, Strands and Google ADK
bridges)
-there is no agent budget, so each tool gets one correction unless its toolset
sets
-``max_retries``.
+behaviour. Outside a pydantic-ai agent (the LangChain, Strands, Google ADK and
Anthropic
+tool runner bridges) there is no agent budget, so each tool gets one
correction unless its
+toolset sets ``max_retries``.
Layering
--------
diff --git a/providers/common/ai/docs/toolsets/object_storage.rst
b/providers/common/ai/docs/toolsets/object_storage.rst
index 8a2a0b251a6..6f6c8489d19 100644
--- a/providers/common/ai/docs/toolsets/object_storage.rst
+++ b/providers/common/ai/docs/toolsets/object_storage.rst
@@ -157,6 +157,7 @@ Using it with other agent frameworks
``ObjectStorageToolset`` implements
:class:`~airflow.providers.common.ai.tools.ToolProvider`, so the same three
tools work in
-a Strands or Google ADK agent through ``AirflowTools``; see
:doc:`../frameworks/index`.
+a Strands or Google ADK agent, or the Anthropic SDK's tool runner, through
+``AirflowTools``; see :doc:`../frameworks/index`.
Outside ``AgentOperator``, ``path`` and ``conn_id`` are used as given: they
are not
rendered as templates.
diff --git a/providers/common/ai/pyproject.toml
b/providers/common/ai/pyproject.toml
index 471fe221384..0b93658996e 100644
--- a/providers/common/ai/pyproject.toml
+++ b/providers/common/ai/pyproject.toml
@@ -83,12 +83,14 @@ dependencies = [
# Any change in the dependencies is preserved when the file is regenerated
[project.optional-dependencies]
# @task.llm_batch talks to the provider batch APIs through the vendor SDKs
directly, since
-# pydantic-ai has no batch abstraction. anthropic 1.0.0 is the first stable
major. openai
+# pydantic-ai has no batch abstraction. anthropic 1.1.0 is the first release
whose tool
+# runner runs a turn's tool calls only when the turn ends in tool_use, which
the tool
+# runner adapter in airflow.providers.common.ai.tools.anthropic relies on.
openai
# 2.47.0 is the first release that accepts the httpx2 client
pydantic-ai-slim>=2.33 hands it
# (2.45 and 2.46 raise TypeError on it); pydantic-ai-slim[openai]>=2.33.0
itself asks for
# openai>=3, so this floor only matters where that requirement is overridden,
as in this
# repository's lock.
-"anthropic" = ["pydantic-ai-slim[anthropic]>=2.36.0", "anthropic>=1.0.0"]
+"anthropic" = ["pydantic-ai-slim[anthropic]>=2.36.0", "anthropic>=1.1.0"]
"bedrock" = ["pydantic-ai-slim[bedrock]>=2.36.0"]
"google" = ["pydantic-ai-slim[google]>=2.36.0"]
"openai" = ["pydantic-ai-slim[openai]>=2.36.0", "openai>=2.47.0"]
@@ -184,7 +186,7 @@ dev = [
# @task.llm_batch's OpenAI/Anthropic adapters call these SDKs directly.
The dev group
# needs them so the adapter tests, which build real SDK models, run
instead of skipping.
"openai>=2.47.0",
- "anthropic>=1.0.0",
+ "anthropic>=1.1.0",
"boat-sdk>=1.9.3",
"islo>=0.3.19",
"opensandbox>=1.1.0",
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_anthropic_tool_runner.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_anthropic_tool_runner.py
new file mode 100644
index 00000000000..cebaeab96ef
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_anthropic_tool_runner.py
@@ -0,0 +1,82 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Answer a question about a database with the Anthropic SDK's tool runner, in an
Airflow task.
+
+The agent loop is the Anthropic SDK's own
``client.beta.messages.tool_runner``. Airflow
+supplies the API key, from a connection, and the tools:
+:class:`~airflow.providers.common.ai.tools.anthropic.AirflowTools` turns a
``SQLToolset``
+into runner tools, with read-only SQL validation, bounded results and
Airflow's secret
+masker applied to every result.
+
+Before running:
+
+1. Install this provider's ``anthropic`` extra:
+ ``pip install "apache-airflow-providers-common-ai[anthropic]"``.
+2. Create a connection (``LLM_CONN_ID``, default ``anthropic_default``) whose
password is
+ an Anthropic API key. Set its host only to route through a gateway.
+3. Create a database connection (``DB_CONN_ID``, default ``sql_default``)
whose hook is a
+ ``DbApiHook`` (e.g. SQLite, Postgres, MySQL).
+"""
+
+from __future__ import annotations
+
+import os
+
+from airflow.providers.common.compat.sdk import BaseHook, dag, task
+
+LLM_CONN_ID = os.environ.get("LLM_CONN_ID", "anthropic_default")
+LLM_MODEL = os.environ.get("LLM_MODEL", "claude-opus-5-5")
+DB_CONN_ID = os.environ.get("DB_CONN_ID", "sql_default")
+
+DEFAULT_QUESTION = "Which tables exist, and how many rows does each contain?"
+
+
+# [START example_anthropic_tool_runner]
+@dag(tags=["example"])
+def example_anthropic_tool_runner():
+ """Answer a question about a database with the Anthropic SDK's tool
runner."""
+
+ @task
+ def ask_the_warehouse(question: str = DEFAULT_QUESTION) -> str:
+ from anthropic import Anthropic
+
+ from airflow.providers.common.ai.tools.anthropic import AirflowTools
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+ llm = BaseHook.get_connection(LLM_CONN_ID)
+ client = Anthropic(api_key=llm.password, base_url=llm.host or None)
+ tools = AirflowTools(SQLToolset(db_conn_id=DB_CONN_ID))
+ runner = client.beta.messages.tool_runner(
+ model=LLM_MODEL,
+ max_tokens=16000,
+ system="You are a SQL analyst. Explore the schema before you
query, and only read.",
+ tools=tools.tools,
+ messages=[{"role": "user", "content": question}],
+ # The runner has no limit of its own on model requests.
+ max_iterations=10,
+ )
+ message = tools.run(runner)
+ return "".join(block.text for block in message.content if block.type
== "text")
+
+ ask_the_warehouse()
+
+
+# [END example_anthropic_tool_runner]
+
+
+example_anthropic_tool_runner()
diff --git a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
index a6ecd9a1b5b..701917722ae 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
@@ -23,6 +23,16 @@ class HITLMaxIterationsError(AirflowException):
"""Raised when the HITL review loop exhausts max iterations without
approval or rejection."""
+class ToolRunnerMaxIterationsError(RuntimeError):
+ """
+ Raised when an Anthropic tool runner reaches ``max_iterations`` while the
model is still calling tools.
+
+ The runner's last message asks for tools and holds no answer. Airflow's
default retry
+ behaviour retries it like any other exception, and a fresh attempt may
finish within
+ the limit.
+ """
+
+
class ToolApprovalError(RuntimeError):
"""
Raised when a run paused for tool approval cannot resume safely.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
b/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
index 8f82ec26f1e..2ecea9b3fe1 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
@@ -20,8 +20,8 @@ Framework-neutral tools backed by Airflow connections.
An :class:`AirflowTool` is one operation an agent can call: a name, a
description, a JSON Schema for its arguments, and an async function. Nothing in
it belongs to a particular agent framework, so the same tool can be handed to
-Strands Agents, Google ADK or LangChain through a small adapter, while the
agent
-itself stays native to that framework.
+Strands Agents, Google ADK, LangChain or the Anthropic SDK's tool runner
through a
+small adapter, while the agent itself stays native to that framework.
The toolsets this provider ships, such as
:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/tools/anthropic.py
b/providers/common/ai/src/airflow/providers/common/ai/tools/anthropic.py
new file mode 100644
index 00000000000..a436441bb6a
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/tools/anthropic.py
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Give Airflow tools to an agent loop run by the Anthropic SDK's tool runner.
+
+.. note:: Experimental; see :ref:`howto/stability`.
+"""
+
+from __future__ import annotations
+
+import copy
+import logging
+from collections.abc import Callable, Coroutine, Iterator
+from contextlib import AbstractContextManager, contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any, TypeVar
+
+try:
+ from anthropic.lib.tools import (
+ BetaAsyncFunctionTool,
+ BetaAsyncToolRunner,
+ BetaFunctionTool,
+ BetaToolRunner,
+ ToolError,
+ beta_async_tool,
+ beta_tool,
+ )
+except ImportError as e:
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+ raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.exceptions import ToolRunnerMaxIterationsError
+from airflow.providers.common.ai.tools import AirflowTool, ToolCallError,
collect_tools
+from airflow.providers.common.ai.tools._from_toolset import tool_call_scope
+from airflow.providers.common.ai.utils.coroutines import run_coroutine_sync
+from airflow.providers.common.ai.utils.tool_definition import serialize_for_llm
+from airflow.providers.common.ai.utils.tool_metrics import calling_framework
+
+if TYPE_CHECKING:
+ from anthropic.types.beta import BetaMessage
+ from anthropic.types.beta.parsed_beta_message import ParsedBetaMessage
+
+ from airflow.providers.common.ai.tools import ToolProvider
+
+log = logging.getLogger(__name__)
+
+__all__ = ["AirflowTools", "AsyncAirflowTools"]
+
+# The ``output_format`` a runner parses its final message into.
+_OutputT = TypeVar("_OutputT")
+
+# The failures of the run in progress. ``run`` sets it, the tools append to it.
+_failures: ContextVar[list[ToolCallError] | None] =
ContextVar("common_ai_anthropic_failures", default=None)
+
+
+class AirflowTools:
+ """
+ Airflow's tools for the tool runner of a synchronous ``Anthropic`` client.
+
+ Pass toolsets that implement
+ :class:`~airflow.providers.common.ai.tools.ToolProvider`, such as
``SQLToolset``
+ and ``HookToolset``, or individual
+ :class:`~airflow.providers.common.ai.tools.AirflowTool` objects. Give
:attr:`tools`
+ to ``client.beta.messages.tool_runner`` and drive the runner with
:meth:`run`:
+
+ .. code-block:: python
+
+ from anthropic import Anthropic
+
+ from airflow.providers.common.ai.tools.anthropic import AirflowTools
+ from airflow.providers.common.ai.toolsets import SQLToolset
+
+ client = Anthropic()
+ tools = AirflowTools(SQLToolset("warehouse",
allowed_tables=["orders"]))
+ runner = client.beta.messages.tool_runner(
+ model="claude-opus-5-5",
+ max_tokens=16000,
+ tools=tools.tools,
+ messages=[...],
+ )
+ message = tools.run(runner)
+
+ Each tool keeps the source tool's name, description and argument schema,
and every
+ result passes through Airflow's secret masker first. A failure the model
can
+ correct, such as a query naming a missing column, reaches it as a
``tool_result``
+ with ``is_error: true``, counted against the tool's retry limit once per
model
+ turn. A tool that is still failing once the limit is used up, or a failure
the
+ model cannot fix, such as a hook raising, makes :meth:`run` raise
+ :class:`~airflow.providers.common.ai.tools.ToolCallError` before the next
model
+ request, so the task fails and Airflow retries it. The runner on its own
hands
+ every exception back to the model, so drive it with :meth:`run` rather than
+ ``until_done()``. A runner that reaches its ``max_iterations`` while the
model is still
+ calling tools makes :meth:`run` raise
+
:class:`~airflow.providers.common.ai.exceptions.ToolRunnerMaxIterationsError`,
since
+ its last message holds no answer. A runner created with ``stream=True`` is
not
+ supported.
+
+ For an ``AsyncAnthropic`` client, use :class:`AsyncAirflowTools`.
+
+ :param sources: Toolsets and tools to add, in order.
+ """
+
+ def __init__(self, *sources: ToolProvider | AirflowTool) -> None:
+ self.tools: list[BetaFunctionTool[Callable[..., str]]] = [
+ _sync_tool(tool) for tool in collect_tools(sources)
+ ]
+
+ def run(self, runner: BetaToolRunner[_OutputT]) ->
ParsedBetaMessage[_OutputT]:
+ """
+ Drive ``runner`` to its final message.
+
+ :raises ToolCallError: when a tool failed in a way the model cannot
fix.
+ :raises ToolRunnerMaxIterationsError: when the runner reached
``max_iterations``
+ while the model was still calling tools.
+ """
+ with _run() as turn:
+ for message in runner:
+ with turn(message.id):
+ if _runs_tools(message):
+ runner.generate_tool_call_response()
+ return _answer(runner.until_done())
+
+
+class AsyncAirflowTools:
+ """
+ Airflow's tools for the tool runner of an ``AsyncAnthropic`` client.
+
+ The asynchronous counterpart of :class:`AirflowTools`, with the same
behaviour:
+
+ .. code-block:: python
+
+ tools = AsyncAirflowTools(SQLToolset("warehouse",
allowed_tables=["orders"]))
+ runner = AsyncAnthropic().beta.messages.tool_runner(
+ model="claude-opus-5-5",
+ max_tokens=16000,
+ tools=tools.tools,
+ messages=[...],
+ )
+ message = await tools.run(runner)
+
+ :param sources: Toolsets and tools to add, in order.
+ """
+
+ def __init__(self, *sources: ToolProvider | AirflowTool) -> None:
+ self.tools: list[BetaAsyncFunctionTool[Callable[..., Coroutine[Any,
Any, str]]]] = [
+ _async_tool(tool) for tool in collect_tools(sources)
+ ]
+
+ async def run(self, runner: BetaAsyncToolRunner[_OutputT]) ->
ParsedBetaMessage[_OutputT]:
+ """
+ Drive ``runner`` to its final message.
+
+ :raises ToolCallError: when a tool failed in a way the model cannot
fix.
+ :raises ToolRunnerMaxIterationsError: when the runner reached
``max_iterations``
+ while the model was still calling tools.
+ """
+ with _run() as turn:
+ async for message in runner:
+ with turn(message.id):
+ if _runs_tools(message):
+ await runner.generate_tool_call_response()
+ return _answer(await runner.until_done())
+
+
+def _runs_tools(message: BetaMessage) -> bool:
+ """
+ Whether the runner will run this turn's tool calls, which it does only for
``tool_use``.
+
+ A turn cut off at ``max_tokens`` can hold a tool call with incomplete
arguments, and a
+ ``pause_turn`` or ``compaction`` turn is sent back to the server
unchanged; the runner
+ runs no tool calls for either, so neither may this.
+ """
+ return message.stop_reason == "tool_use"
+
+
+def _answer(message: ParsedBetaMessage[_OutputT]) ->
ParsedBetaMessage[_OutputT]:
+ """
+ Return the runner's final message, unless the runner stopped mid-task.
+
+ The runner ends after a ``tool_use`` turn only when it reaches
``max_iterations``, and
+ that turn holds tool calls rather than an answer.
+ """
+ if _runs_tools(message):
+ raise ToolRunnerMaxIterationsError(
+ f"The tool runner reached max_iterations while the model was still
calling tools "
+ f"(last turn {message.id}). Raise max_iterations, or give the
model a narrower task."
+ )
+ return message
+
+
+@contextmanager
+def _run() -> Iterator[Callable[[str], AbstractContextManager[None]]]:
+ """
+ Scope one run of the runner, and yield the context that runs one model
turn's tool calls.
+
+ The runner caches the tool results it builds and sends them with the next
model
+ request, so building them in the turn's context lets the tools count
retries per turn
+ and lets a failure end the run before that request.
+ """
+ run = object()
+ failures: list[ToolCallError] = []
+ token = _failures.set(failures)
+
+ def raise_first_failure() -> None:
+ if failures:
+ raise failures[0]
+
+ @contextmanager
+ def turn(message_id: str) -> Iterator[None]:
+ # Checked on the way in too: should the runner ever run tools outside
a turn, their
+ # failure ends the run at the next turn rather than once the runner
stops.
+ raise_first_failure()
+ with tool_call_scope(run=run, turn=message_id):
+ yield
+ raise_first_failure()
+
+ try:
+ yield turn
+ raise_first_failure()
+ finally:
+ _failures.reset(token)
+
+
+async def _call(tool: AirflowTool, arguments: dict[str, Any]) -> str:
+ failures = _failures.get()
+ if failures:
+ # The run ends once this turn's calls are done, so start no more of
them.
+ raise ToolError("Not run: an earlier tool call in this turn failed.")
+ try:
+ with calling_framework("anthropic"):
+ result = await tool.call(arguments)
+ except ToolCallError as e:
+ if failures is None:
+ log.warning(
+ "The tool runner was not driven by AirflowTools.run, so it
hands this failure to "
+ "the model instead of failing the task: %s",
+ e,
+ )
+ raise
+ failures.append(e)
+ # ``run`` raises the original once the turn is done. The runner logs
any other
+ # exception with its traceback, so this keeps the task log to one
account of it.
+ raise ToolError(str(e)) from None
+ # Masks what a plain JSON dump would render with str(), such as a
dataclass holding a secret.
+ text = serialize_for_llm(result.content)
+ if result.is_error:
+ raise ToolError(text)
+ return text
+
+
+def _sync_tool(tool: AirflowTool) -> BetaFunctionTool[Callable[..., str]]:
+ def call_airflow_tool(**arguments: Any) -> str:
+ return run_coroutine_sync(_call(tool, arguments))
+
+ # The source schema can be a toolset's module-level constant, so hand over
a copy.
+ return beta_tool(
+ call_airflow_tool,
+ name=tool.name,
+ description=tool.description,
+ input_schema=copy.deepcopy(tool.parameters),
+ )
+
+
+def _async_tool(tool: AirflowTool) -> BetaAsyncFunctionTool[Callable[...,
Coroutine[Any, Any, str]]]:
+ async def call_airflow_tool(**arguments: Any) -> str:
+ return await _call(tool, arguments)
+
+ return beta_async_tool(
+ call_airflow_tool,
+ name=tool.name,
+ description=tool.description,
+ input_schema=copy.deepcopy(tool.parameters),
+ )
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/coroutines.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/coroutines.py
index 7ab97199be1..66b1cc27b55 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/coroutines.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/coroutines.py
@@ -20,6 +20,7 @@ from __future__ import annotations
import asyncio
import concurrent.futures
+import contextvars
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
@@ -32,11 +33,15 @@ def run_coroutine_sync(coroutine: Coroutine[Any, Any, Any])
-> Any:
When no event loop is running in this thread the coroutine is driven with
:func:`asyncio.run`. When one is (an async caller), it runs on a worker
thread
- with its own loop, so the caller's loop is never nested.
+ with its own loop, so the caller's loop is never nested. Either way it
sees the
+ caller's context variables.
"""
try:
asyncio.get_running_loop()
except RuntimeError:
- return asyncio.run(coroutine)
- with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
- return pool.submit(asyncio.run, coroutine).result()
+ pass
+ else:
+ with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+ return pool.submit(contextvars.copy_context().run, asyncio.run,
coroutine).result()
+ # Outside the except block, so what the coroutine raises is not chained to
the RuntimeError.
+ return asyncio.run(coroutine)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/tool_metrics.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/tool_metrics.py
index e2caf09afc5..8f1c663cc37 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/tool_metrics.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/tool_metrics.py
@@ -28,7 +28,7 @@ if TYPE_CHECKING:
from collections.abc import Iterator
#: The agent framework a tool call came through, as the ``framework`` tag
reports it.
-Framework = Literal["pydantic_ai", "strands", "adk", "langchain", "none"]
+Framework = Literal["pydantic_ai", "strands", "adk", "langchain", "anthropic",
"none"]
# Set by each framework adapter around the calls it makes.
_framework: ContextVar[Framework | None] =
ContextVar("common_ai_tool_framework", default=None)
diff --git a/providers/common/ai/tests/unit/common/ai/test_exceptions.py
b/providers/common/ai/tests/unit/common/ai/test_exceptions.py
index b7b95d4cbec..c6dfa6d06b2 100644
--- a/providers/common/ai/tests/unit/common/ai/test_exceptions.py
+++ b/providers/common/ai/tests/unit/common/ai/test_exceptions.py
@@ -25,6 +25,7 @@ from airflow.providers.common.ai.exceptions import (
LLMFileAnalysisMultimodalRequiredError,
LLMFileAnalysisUnsupportedFormatError,
ManagedAgentInvocationError,
+ ToolRunnerMaxIterationsError,
)
from airflow.providers.common.compat.sdk import AirflowException
@@ -35,6 +36,7 @@ EXCEPTION_CASES = [
(LLMFileAnalysisLimitExceededError, LLMFileAnalysisError),
(LLMFileAnalysisMultimodalRequiredError,
LLMFileAnalysisUnsupportedFormatError),
(ManagedAgentInvocationError, RuntimeError),
+ (ToolRunnerMaxIterationsError, RuntimeError),
]
diff --git a/providers/common/ai/tests/unit/common/ai/tools/test_anthropic.py
b/providers/common/ai/tests/unit/common/ai/tools/test_anthropic.py
new file mode 100644
index 00000000000..b8a2d736f6d
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/tools/test_anthropic.py
@@ -0,0 +1,364 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+import copy
+import json
+from typing import TYPE_CHECKING, Any
+from unittest.mock import MagicMock, patch
+from uuid import uuid4
+
+import pytest
+
+pytest.importorskip("anthropic")
+
+import httpx2
+from anthropic import Anthropic, AsyncAnthropic, DefaultAsyncHttpxClient,
DefaultHttpxClient
+from anthropic.types.beta import BetaTextBlock
+from pydantic_ai.exceptions import ModelRetry
+
+from airflow.providers.common.ai.exceptions import ToolRunnerMaxIterationsError
+from airflow.providers.common.ai.tools import AirflowTool, ToolCallError,
ToolResult
+from airflow.providers.common.ai.tools.anthropic import AirflowTools,
AsyncAirflowTools
+from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+from unit.common.ai.tools.test__from_toolset import _scripted
+from unit.common.ai.toolsets.test_sql import _make_mock_db_hook
+
+if TYPE_CHECKING:
+ from anthropic.lib.tools import BetaToolRunner
+ from anthropic.types.beta.parsed_beta_message import ParsedBetaMessage
+
+
+class _ScriptedModel:
+ """
+ Stands in for the Messages API: asks for ``tool_name`` until it has
``calls`` results,
+ then answers with the last result it got. Keeps every request body it
receives.
+
+ With a ``stop_reason`` other than ``tool_use``, the first turn asks for
the tool but ends
+ with that reason, and any later turn answers without asking again.
``per_turn`` is how
+ many calls each asking turn makes at once.
+ """
+
+ def __init__(
+ self,
+ tool_name: str,
+ tool_input: dict[str, Any],
+ *,
+ calls: int = 1,
+ stop_reason: str = "tool_use",
+ per_turn: int = 1,
+ ) -> None:
+ self.tool_name = tool_name
+ self.tool_input = tool_input
+ self.calls = calls
+ self.stop_reason = stop_reason
+ self.per_turn = per_turn
+ self.requests: list[dict[str, Any]] = []
+ # Unique across instances, since the retry limit tells turns apart by
message id.
+ self._prefix = uuid4().hex[:8]
+
+ def results(self) -> list[dict[str, Any]]:
+ """Every ``tool_result`` block sent to the model, in order."""
+ messages = self.requests[-1]["messages"] if self.requests else []
+ return [
+ block
+ for message in messages
+ if message["role"] == "user" and isinstance(message["content"],
list)
+ for block in message["content"]
+ if block.get("type") == "tool_result"
+ ]
+
+ def __call__(self, request: httpx2.Request) -> httpx2.Response:
+ self.requests.append(json.loads(request.content))
+ turn = len(self.requests)
+ results = self.results()
+ if len(results) < self.calls and (turn == 1 or self.stop_reason ==
"tool_use"):
+ content = [
+ {
+ "type": "tool_use",
+ "id": f"toolu_{self._prefix}_{turn}_{i}",
+ "name": self.tool_name,
+ "input": self.tool_input,
+ }
+ for i in range(self.per_turn)
+ ]
+ stop_reason = self.stop_reason
+ else:
+ content = [{"type": "text", "text": json.dumps(results[-1] if
results else None)}]
+ stop_reason = "end_turn"
+ return httpx2.Response(
+ 200,
+ json={
+ "id": f"msg_{self._prefix}_{turn}",
+ "type": "message",
+ "role": "assistant",
+ "model": "claude-test",
+ "content": content,
+ "stop_reason": stop_reason,
+ "stop_sequence": None,
+ "usage": {"input_tokens": 1, "output_tokens": 1},
+ },
+ )
+
+
+def _sync_runner(model: _ScriptedModel, tools: AirflowTools, **runner_options:
Any) -> BetaToolRunner:
+ client = Anthropic(api_key="test",
http_client=DefaultHttpxClient(transport=httpx2.MockTransport(model)))
+ return client.beta.messages.tool_runner(
+ model="claude-test",
+ max_tokens=100,
+ tools=tools.tools,
+ messages=[{"role": "user", "content": "go"}],
+ **runner_options,
+ )
+
+
+def _drive_sync(model: _ScriptedModel, tools: AirflowTools, **runner_options:
Any) -> ParsedBetaMessage:
+ return tools.run(_sync_runner(model, tools, **runner_options))
+
+
+def _drive_async(model: _ScriptedModel, tools: AsyncAirflowTools,
**runner_options: Any) -> ParsedBetaMessage:
+ async def run() -> ParsedBetaMessage:
+ client = AsyncAnthropic(
+ api_key="test",
http_client=DefaultAsyncHttpxClient(transport=httpx2.MockTransport(model))
+ )
+ runner = client.beta.messages.tool_runner(
+ model="claude-test",
+ max_tokens=100,
+ tools=tools.tools,
+ messages=[{"role": "user", "content": "go"}],
+ **runner_options,
+ )
+ return await tools.run(runner)
+
+ return asyncio.run(run())
+
+
+def _answer(message: ParsedBetaMessage) -> Any:
+ """The tool result the scripted model echoed back as its answer."""
+ block = message.content[0]
+ assert isinstance(block, BetaTextBlock)
+ return json.loads(block.text)
+
+
[email protected](
+ params=[
+ pytest.param((AirflowTools, _drive_sync), id="sync"),
+ pytest.param((AsyncAirflowTools, _drive_async), id="async"),
+ ]
+)
+def flavour(request):
+ """The adapter class and the function that drives a runner with it."""
+ return request.param
+
+
[email protected]
+def run_agent(flavour):
+ tools_class, drive = flavour
+
+ def run(model: _ScriptedModel, *sources, **runner_options: Any) ->
ParsedBetaMessage:
+ return drive(model, tools_class(*sources), **runner_options)
+
+ return run
+
+
+def _tool(result: ToolResult | Exception, ran: list[dict[str, Any]] | None =
None) -> AirflowTool:
+ """A ``lookup`` tool that returns or raises ``result``, recording its
calls in ``ran``."""
+
+ async def function(arguments: dict[str, Any]) -> ToolResult:
+ if ran is not None:
+ ran.append(arguments)
+ if isinstance(result, Exception):
+ raise result
+ return result
+
+ return AirflowTool(
+ name="lookup",
+ description="Look something up.",
+ parameters={"type": "object", "properties": {"key": {"type":
"string"}}, "required": ["key"]},
+ function=function,
+ )
+
+
[email protected]("tools_class", [AirflowTools, AsyncAirflowTools])
+def test_each_toolset_tool_becomes_a_runner_tool_with_its_schema(tools_class):
+ tools = tools_class(_tool(ToolResult("ok")), SQLToolset("pg_default"))
+
+ definitions = [tool.to_dict() for tool in tools.tools]
+
+ assert [d["name"] for d in definitions] == ["lookup", "list_tables",
"get_schema", "query", "check_query"]
+ assert definitions[3]["input_schema"]["required"] == ["sql"]
+
+
[email protected]("tools_class", [AirflowTools, AsyncAirflowTools])
+def test_a_runner_tool_does_not_share_the_source_schema(tools_class):
+ """A toolset's schema can be a module-level constant that the pydantic-ai
path also uses."""
+ source = _tool(ToolResult("ok"))
+ before = copy.deepcopy(source.parameters)
+
+
tools_class(source).tools[0].to_dict()["input_schema"]["properties"]["key"]["type"]
= "integer"
+
+ assert source.parameters == before
+
+
+class TestRun:
+ def test_the_model_reads_a_toolset_result(self, run_agent):
+ ts = SQLToolset("pg_default")
+ ts._hook = _make_mock_db_hook(records=[(1, "Ada")],
last_description=[("id",), ("name",)])
+ model = _ScriptedModel("query", {"sql": "SELECT id, name FROM users"})
+
+ result = _answer(run_agent(model, ts))
+
+ assert "is_error" not in result
+ assert json.loads(result["content"])["rows"] == [[1, "Ada"]]
+
+ def test_a_correctable_failure_reaches_the_model_as_an_error(self,
run_agent):
+ model = _ScriptedModel("lookup", {"key": "a"})
+
+ result = _answer(run_agent(model, _tool(ToolResult("Unknown column
'nme'", is_error=True))))
+
+ assert result["is_error"] is True
+ assert "Unknown column 'nme'" in json.dumps(result["content"])
+
+ def
test_any_other_failure_fails_the_run_before_the_next_model_request(self,
run_agent, caplog):
+ """The runner alone would hand the exception to the model and ask it
again."""
+ model = _ScriptedModel("lookup", {"key": "a"})
+
+ with pytest.raises(ToolCallError, match="PermissionError: role cannot
read orders"):
+ run_agent(model, _tool(PermissionError("role cannot read orders")))
+
+ assert len(model.requests) == 1
+ # The runner logs the exceptions it catches; this one reaches the task
log once, as the failure.
+ assert "Error occurred while calling tool" not in caplog.text
+
+ def
test_a_failure_goes_to_the_model_with_a_warning_when_run_is_not_used(self,
caplog):
+ model = _ScriptedModel("lookup", {"key": "a"})
+ runner = _sync_runner(model, AirflowTools(_tool(PermissionError("role
cannot read orders"))))
+
+ runner.until_done()
+
+ assert model.results()[0]["is_error"] is True
+ assert "not driven by AirflowTools.run" in caplog.text
+
+ def
test_a_tool_that_keeps_failing_fails_the_run_once_its_retries_are_spent(self,
run_agent):
+ model = _ScriptedModel("step", {}, calls=5)
+ step = _scripted(ModelRetry("no such key"), ModelRetry("no such key"),
max_retries=1)
+
+ with pytest.raises(ToolCallError, match="kept failing after 1
correction"):
+ run_agent(model, step)
+
+ # One failure the model may correct, then the second ends the run.
+ assert len(model.requests) == 2
+
+ @pytest.mark.parametrize(
+ ("stop_reason", "model_requests"),
+ [
+ pytest.param("max_tokens", 1, id="cut-off-turn-ends-the-run"),
+ pytest.param("pause_turn", 2, id="paused-turn-is-resumed"),
+ ],
+ )
+ def test_a_turn_that_did_not_end_in_tool_use_runs_no_tools(self,
run_agent, stop_reason, model_requests):
+ """The runner itself only runs a turn's tool calls when it ended in
``tool_use``."""
+ ran: list[dict[str, Any]] = []
+ model = _ScriptedModel("lookup", {"key": "a"}, stop_reason=stop_reason)
+
+ run_agent(model, _tool(ToolResult("ran"), ran))
+
+ assert ran == []
+ assert len(model.requests) == model_requests
+
+ def test_a_run_that_reaches_max_iterations_while_calling_tools_fails(self,
run_agent):
+ """The runner's last message is then a turn asking for tools, with no
answer in it."""
+ model = _ScriptedModel("lookup", {"key": "a"}, calls=5)
+
+ with pytest.raises(ToolRunnerMaxIterationsError, match="reached
max_iterations"):
+ run_agent(model, _tool(ToolResult("ran")), max_iterations=2)
+
+ assert len(model.requests) == 2
+
+ def test_a_run_that_answers_on_its_last_iteration_returns_the_answer(self,
run_agent):
+ model = _ScriptedModel("lookup", {"key": "a"})
+
+ result = _answer(run_agent(model, _tool(ToolResult("ran")),
max_iterations=2))
+
+ assert result["content"] == "ran"
+
+ def test_calls_after_a_failure_in_the_same_turn_do_not_run(self,
run_agent):
+ """The turn's results are never sent, so a later call's side effects
would be for nothing."""
+ ran: list[dict[str, Any]] = []
+ model = _ScriptedModel("lookup", {"key": "a"}, per_turn=3)
+
+ with pytest.raises(ToolCallError, match="PermissionError"):
+ run_agent(model, _tool(PermissionError("role cannot read orders"),
ran))
+
+ assert len(ran) == 1
+
+ def
test_the_sync_tools_fail_the_run_when_called_under_a_running_event_loop(self):
+ """The sync tools run on a worker thread there, which must still see
the run's state."""
+ model = _ScriptedModel("lookup", {"key": "a"})
+
+ async def caller() -> None:
+ _drive_sync(model, AirflowTools(_tool(PermissionError("role cannot
read orders"))))
+
+ with pytest.raises(ToolCallError, match="PermissionError"):
+ asyncio.run(caller())
+
+ assert len(model.requests) == 1
+
+ def test_a_structured_result_reaches_the_model_as_json(self, run_agent):
+ model = _ScriptedModel("lookup", {"key": "a"})
+
+ result = _answer(run_agent(model, _tool(ToolResult({"rows": [[1,
None]]}))))
+
+ assert json.loads(result["content"]) == {"rows": [[1, None]]}
+
+ def test_failed_calls_in_one_turn_count_once_against_the_retry_limit(self,
run_agent):
+ model = _ScriptedModel("step", {}, calls=2, per_turn=2)
+
+ run_agent(model, _scripted(ModelRetry("no such key"), ModelRetry("no
such key"), max_retries=1))
+
+ assert [r.get("is_error") for r in model.results()] == [True, True]
+
+ def
test_each_run_of_a_reused_instance_starts_with_a_fresh_retry_budget(self,
flavour):
+ tools_class, drive = flavour
+ tools = tools_class(_scripted(ModelRetry("no such key"),
ModelRetry("no such key"), max_retries=1))
+
+ first, second = _ScriptedModel("step", {}), _ScriptedModel("step", {})
+ drive(first, tools)
+ drive(second, tools)
+
+ assert second.results()[0]["is_error"] is True
+
+ @pytest.mark.enable_redact
+ def test_a_secret_in_a_result_reaches_the_model_masked(self, run_agent,
registered_secret):
+ model = _ScriptedModel("lookup", {"key": "a"})
+
+ run_agent(model, _tool(ToolResult(f"key={registered_secret}")))
+
+ sent = json.dumps(model.requests)
+ assert registered_secret not in sent
+ assert "key=***" in sent
+
+ def test_the_toolsets_calls_are_counted_as_anthropic(self, run_agent):
+ ts = SQLToolset("pg_default")
+ ts._hook = _make_mock_db_hook()
+
+ with patch("airflow.providers.common.ai.utils.tool_metrics.Stats",
MagicMock(spec=["incr"])) as stats:
+ run_agent(_ScriptedModel("list_tables", {}), ts)
+
+ assert stats.incr.call_args.kwargs["tags"]["framework"] == "anthropic"
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_coroutines.py
b/providers/common/ai/tests/unit/common/ai/utils/test_coroutines.py
index d4a467f56a0..f96cbecd691 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_coroutines.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_coroutines.py
@@ -18,14 +18,21 @@ from __future__ import annotations
import asyncio
import threading
+from contextvars import ContextVar
from airflow.providers.common.ai.utils.coroutines import run_coroutine_sync
+_value: ContextVar[str] = ContextVar("value", default="unset")
+
async def _thread_name() -> str:
return threading.current_thread().name
+async def _read_value() -> str:
+ return _value.get()
+
+
class TestRunCoroutineSync:
def test_runs_on_the_calling_thread_when_no_loop_is_running(self):
assert run_coroutine_sync(_thread_name()) ==
threading.current_thread().name
@@ -35,3 +42,10 @@ class TestRunCoroutineSync:
return run_coroutine_sync(_thread_name())
assert asyncio.run(caller()) != threading.current_thread().name
+
+ def test_sees_the_callers_context_variables_from_a_running_loop(self):
+ async def caller() -> str:
+ _value.set("caller")
+ return run_coroutine_sync(_read_value())
+
+ assert asyncio.run(caller()) == "caller"
diff --git
a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
index 4782f69b7d7..428ae4e3e35 100644
---
a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
+++
b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
@@ -428,8 +428,8 @@ metrics:
- name: "common_ai.tool_calls"
description: "Number of calls to a toolset of the common.ai provider.
Metric with toolset
- (the toolset class), framework (``pydantic_ai``, ``strands``, ``adk``,
``langchain`` or
- ``none``) and outcome (``executed``, ``failed`` or ``replayed``) tagging."
+ (the toolset class), framework (``pydantic_ai``, ``strands``, ``adk``,
``langchain``,
+ ``anthropic`` or ``none``) and outcome (``executed``, ``failed`` or
``replayed``) tagging."
type: "counter"
legacy_name: "-"
name_variables: ["toolset", "framework", "outcome"]
diff --git a/uv.lock b/uv.lock
index 5d262af491d..61072416bf8 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4673,7 +4673,7 @@ docs = [
[package.metadata]
requires-dist = [
- { name = "anthropic", marker = "extra == 'anthropic'", specifier =
">=1.0.0" },
+ { name = "anthropic", marker = "extra == 'anthropic'", specifier =
">=1.1.0" },
{ name = "apache-airflow", editable = "." },
{ name = "apache-airflow-providers-common-compat", editable =
"providers/common/compat" },
{ name = "apache-airflow-providers-common-sql", marker = "extra ==
'common-sql'", editable = "providers/common/sql" },
@@ -4715,7 +4715,7 @@ provides-extras = ["anthropic", "bedrock", "google",
"openai", "typesafe", "mcp"
[package.metadata.requires-dev]
dev = [
- { name = "anthropic", specifier = ">=1.0.0" },
+ { name = "anthropic", specifier = ">=1.1.0" },
{ name = "apache-airflow", editable = "." },
{ name = "apache-airflow-devel-common", editable = "devel-common" },
{ name = "apache-airflow-providers-common-compat", editable =
"providers/common/compat" },
