This is an automated email from the ASF dual-hosted git repository.
Lee-W 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 c3e00461d06 Make LoggingToolset private and update Experimental notes
in common.ai (#74694)
c3e00461d06 is described below
commit c3e00461d06973b0eae03fee203c9eb6a6435d99
Author: Wei Lee <[email protected]>
AuthorDate: Mon Oct 26 15:59:51 2026 +0800
Make LoggingToolset private and update Experimental notes in common.ai
(#74694)
---
providers/common/ai/docs/changelog.rst | 12 ++++++++
providers/common/ai/docs/frameworks/index.rst | 14 ++++++---
providers/common/ai/docs/operators/agent.rst | 15 ++++++----
.../ai/docs/operators/llm_schema_compare.rst | 6 ++++
providers/common/ai/docs/operators/llm_sql.rst | 4 ++-
providers/common/ai/docs/sandbox/index.rst | 2 +-
providers/common/ai/docs/stability.rst | 11 ++++---
providers/common/ai/docs/supported_services.rst | 2 +-
providers/common/ai/docs/toolsets/index.rst | 9 +++---
providers/common/ai/docs/toolsets/logging.rst | 35 ++++++++++++----------
.../airflow/providers/common/ai/operators/agent.py | 9 ++++--
.../common/ai/operators/llm_schema_compare.py | 4 ++-
.../providers/common/ai/operators/llm_sql.py | 4 +++
.../providers/common/ai/toolsets/__init__.py | 2 --
.../providers/common/ai/toolsets/logging.py | 19 +++++-------
.../tests/unit/common/ai/operators/test_agent.py | 2 +-
.../tests/unit/common/ai/test_provider_metadata.py | 2 +-
.../tests/unit/common/ai/toolsets/test_logging.py | 6 ++--
.../in_container/run_provider_yaml_files_check.py | 2 +-
19 files changed, 99 insertions(+), 61 deletions(-)
diff --git a/providers/common/ai/docs/changelog.rst
b/providers/common/ai/docs/changelog.rst
index ee9c5ac9552..1d288a62d4e 100644
--- a/providers/common/ai/docs/changelog.rst
+++ b/providers/common/ai/docs/changelog.rst
@@ -80,6 +80,17 @@ Changelog
keyword, e.g. ``LlamaIndexHook(llm_model="gpt-5")``, which also avoids the
meanings swapping between
the two hooks.
+.. note::
+ ``LoggingToolset`` is removed from the public API. Importing it from either
+ ``airflow.providers.common.ai.toolsets.LoggingToolset`` (exported from the
package since
+ 1.0.0, see #74449 under 1.0.0 Features below) or
+ ``airflow.providers.common.ai.toolsets.logging.LoggingToolset`` (public
since 0.1.0) now
+ raises ``ImportError``. ``AgentOperator``
+ still logs tool calls through ``enable_tool_logging``. To log tool calls of
an agent you
+ build yourself, pass ``capabilities=[ToolLoggingCapability()]`` to the
pydantic-ai ``Agent``
+ (import it from ``airflow.providers.common.ai.toolsets.logging``); its
``logger`` argument
+ sets the logger. See :doc:`stability`.
+
1.1.0
.....
@@ -120,6 +131,7 @@ Misc
``capabilities`` or ``spec_file`` move to ``AgentOperator`` or
``@task.agent``. See
:ref:`howto/operator:llm`.
+.. note::
On Airflow >= 3.3, the ``usage_limits`` of the ``LLM*`` operators
(``LLMOperator``,
``LLMBranchOperator``, ``LLMSQLQueryOperator``, ``LLMSchemaCompareOperator``,
``LLMFileAnalysisOperator``) counts usage across every attempt of the task
instance
diff --git a/providers/common/ai/docs/frameworks/index.rst
b/providers/common/ai/docs/frameworks/index.rst
index e148c21dd53..ff9b6169563 100644
--- a/providers/common/ai/docs/frameworks/index.rst
+++ b/providers/common/ai/docs/frameworks/index.rst
@@ -149,10 +149,16 @@ Tested versions
These are the versions the adapter tests have been run against. Strands and
ADK cap
packages below the versions in Airflow's dependency lock, and current OpenAI
Agents SDK
releases need ``openai`` 3, so CI runs their adapter tests in a job of their
own, one
-framework at a time. Pull requests test each framework against the versions in
-Airflow's CI image, which for the OpenAI Agents SDK means the newest release
that accepts
-``openai`` 2.x; scheduled runs test each framework's newest release with its
own
-dependency pins.
+framework at a time.
+
+Outside canary runs, the job holds every package already in the CI image at its
+installed version with ``--override``, so Strands' and ADK's own dependency
caps are
+overridden. It installs the newest framework release that is older than the
repository's
+``exclude-newer`` window, which for the OpenAI Agents SDK means the newest
release that
+accepts ``openai`` 2.x, not the versions in the table above. It therefore does
not show
+that those versions work in the environment a user gets. Canary runs, which
include
+scheduled and manual runs on ``main`` and pull requests labeled ``canary``,
install each
+framework's newest release with its own dependency pins instead.
Choosing a route
----------------
diff --git a/providers/common/ai/docs/operators/agent.rst
b/providers/common/ai/docs/operators/agent.rst
index 9dbf792c9c9..c1e587c5b42 100644
--- a/providers/common/ai/docs/operators/agent.rst
+++ b/providers/common/ai/docs/operators/agent.rst
@@ -375,13 +375,16 @@ Parameters
``AgentSkillsToolset`` for :ref:`agent-skills`, etc.).
- ``capabilities``: List of pydantic-ai capabilities (``Thinking``,
``WebSearch``, guardrails,
etc.). See :ref:`capabilities`.
-- ``enable_tool_logging``: Wrap the assembled function toolset in
- :class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` so that
- tools supplied through ``toolsets=``, ``agent_params={"tools": [...]}``, and
- capabilities are logged in real time, including tools the model found through
+- ``enable_tool_logging``: Append a
+
:class:`~airflow.providers.common.ai.toolsets.logging.ToolLoggingCapability`,
which
+ wraps the assembled function toolset, so that tools supplied through
``toolsets=``,
+ ``agent_params={"tools": [...]}``, and capabilities are logged in real time,
including
+ tools the model found through
:ref:`tool search <howto/toolset:progressive-disclosure>`. Output tools,
provider-native
- tools, and tools another capability adds through its own wrapper (code mode's
- ``run_code``, ToolSearch's ``search_tools``) are not covered. Default
``True``.
+ tools, and tools another capability adds through its own wrapper (code
mode's ``run_code``,
+ ToolSearch's ``search_tools``) are not covered. If ``capabilities`` already
holds a
+ ``ToolLoggingCapability``, the operator adds none and uses that one's
``logger``.
+ Default ``True``.
- ``agent_params``: Additional keyword arguments passed to the pydantic-ai
``Agent`` constructor (e.g. ``retries``, ``model_settings``).
- .. _agent-usage-budget:
diff --git a/providers/common/ai/docs/operators/llm_schema_compare.rst
b/providers/common/ai/docs/operators/llm_schema_compare.rst
index fa7d9ed8eb7..b09754ac6e4 100644
--- a/providers/common/ai/docs/operators/llm_schema_compare.rst
+++ b/providers/common/ai/docs/operators/llm_schema_compare.rst
@@ -63,6 +63,12 @@ columns alone:
With Object Storage or a Database Table
---------------------------------------
+.. note::
+
+ A source whose connection is not a ``DbApiHook`` runs on the DataFusion
+ engine of ``common.sql``, which pins ``datafusion`` below 52 because the
DataFusion
+ 52 crate is not supported with ``iceberg-core``.
+
Use ``data_sources`` with
:class:`~airflow.providers.common.sql.config.DataSourceConfig` to include
object-storage sources (S3, GCS, Azure Blob Storage, Parquet, CSV, Iceberg,
etc.) in the comparison.
diff --git a/providers/common/ai/docs/operators/llm_sql.rst
b/providers/common/ai/docs/operators/llm_sql.rst
index c843d0a62da..5f457fb016b 100644
--- a/providers/common/ai/docs/operators/llm_sql.rst
+++ b/providers/common/ai/docs/operators/llm_sql.rst
@@ -69,7 +69,9 @@ the schema context.
Object-storage schema introspection requires the ``datafusion`` extra of
``apache-airflow-providers-common-sql``. Install it with
- ``pip install "apache-airflow-providers-common-sql[datafusion]"``.
+ ``pip install "apache-airflow-providers-common-sql[datafusion]"``.
``common.sql`` pins
+ ``datafusion`` below 52 because the DataFusion 52 crate is not supported
with
+ ``iceberg-core``.
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_sql.py
:language: python
diff --git a/providers/common/ai/docs/sandbox/index.rst
b/providers/common/ai/docs/sandbox/index.rst
index aed49b52835..28033a711bc 100644
--- a/providers/common/ai/docs/sandbox/index.rst
+++ b/providers/common/ai/docs/sandbox/index.rst
@@ -611,7 +611,7 @@ Isolation is one of several controls an agent needs, and
each has a limit:
* - Audit
- Can you establish what happened afterwards?
- A record does not prevent the action.
- - :doc:`LoggingToolset <../toolsets/logging>` and
+ - :doc:`ToolLoggingCapability <../toolsets/logging>` and
:doc:`tracing <../observability>`.
A sandbox with no egress can still leak data. A secret printed inside comes
back
diff --git a/providers/common/ai/docs/stability.rst
b/providers/common/ai/docs/stability.rst
index 3fad7b4ca4f..df3ee5898b3 100644
--- a/providers/common/ai/docs/stability.rst
+++ b/providers/common/ai/docs/stability.rst
@@ -102,7 +102,8 @@ Pydantic AI toolsets, but the Pydantic AI class they
inherit from can change.
types, a write nested inside the statement fails it too. A reviewer's
edit under
``allow_modifications=True`` is checked the same way. ``db_conn_id``
with
``table_names``, ``schema_context`` or ``datasource_config`` give the
model the
- schema.
+ schema. ``datasource_config`` reads object storage through the
DataFusion engine of
+ ``common.sql``, which pins ``datafusion`` below 52.
* - ``@task.llm_file_analysis`` and
:class:`~airflow.providers.common.ai.operators.llm_file_analysis.LLMFileAnalysisOperator`
- Reads the file at ``file_path``, or up to ``max_files`` files under a
prefix,
@@ -120,6 +121,8 @@ Pydantic AI toolsets, but the Pydantic AI class they
inherit from can change.
``severity`` of ``critical``, ``warning`` or ``info``.
``context_strategy="basic"``
sends columns only; ``"full"``, the default, adds keys and indexes for
database
sources.
+ ``data_sources`` reads object storage through the DataFusion engine of
+ ``common.sql``, which pins ``datafusion`` below 52.
* - ``@task.llm_batch`` and
:class:`~airflow.providers.common.ai.operators.llm_batch.LLMBatchOperator`
- Submits ``requests`` to the batch API that the ``model_id`` prefix
selects, writes
@@ -254,10 +257,10 @@ Everything this provider ships that is not in the table
above is experimental.
``capture_content``; :doc:`observability`)
- Span names and attributes come from Pydantic AI's instrumentation and
the
OpenTelemetry GenAI conventions, which are still in development.
- * - :class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset`
used directly
+ * -
:class:`~airflow.providers.common.ai.toolsets.logging.ToolLoggingCapability`
(:doc:`toolsets/logging`)
- - The wrapper ``AgentOperator`` applies for ``enable_tool_logging``; its
constructor
- may change.
+ - The capability ``AgentOperator`` applies for ``enable_tool_logging``;
its constructor
+ and behaviour may change.
* - The Strands plugin, the ADK toolset and Strands checkpoint resume
(:doc:`frameworks/index`)
- Written against Strands 1.56 and ADK 2.9.1, which change often. Strands
marks its
diff --git a/providers/common/ai/docs/supported_services.rst
b/providers/common/ai/docs/supported_services.rst
index fb808f6fec4..9bf90e74253 100644
--- a/providers/common/ai/docs/supported_services.rst
+++ b/providers/common/ai/docs/supported_services.rst
@@ -47,7 +47,7 @@ Toolsets
* - :doc:`DataFusionToolset <toolsets/datafusion>`
- Object storage (S3, GCS, Azure Blob Storage, local filesystem, Apache
Iceberg) via
Apache DataFusion
- * - :doc:`LoggingToolset <toolsets/logging>`
+ * - :doc:`ToolLoggingCapability <toolsets/logging>`
- Airflow task logs
* - :doc:`MCPToolset <toolsets/mcp>`
- Any MCP server (user-supplied endpoint)
diff --git a/providers/common/ai/docs/toolsets/index.rst
b/providers/common/ai/docs/toolsets/index.rst
index 1e9351cbb4b..2291ead1a45 100644
--- a/providers/common/ai/docs/toolsets/index.rst
+++ b/providers/common/ai/docs/toolsets/index.rst
@@ -105,7 +105,6 @@ Every toolset imports from the
``airflow.providers.common.ai.toolsets`` package
AgentSkillsToolset,
DataFusionToolset,
HookToolset,
- LoggingToolset,
ManagedAgentToolset,
MCPToolset,
ObjectStorageToolset,
@@ -306,10 +305,10 @@ toolset sets ``max_retries``.
Layering
--------
-:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` is not an
-alternative to anything above. It is a wrapper: it takes a toolset and returns
a
-toolset that logs each call. Choose a route first, then decide whether to wrap
-it. The same applies to
+:class:`~airflow.providers.common.ai.toolsets.logging.ToolLoggingCapability`
is not an
+alternative to anything above. It is a capability: it wraps the agent's
assembled
+function toolset so that each call is logged. Choose a route first, then decide
+whether to log it. The same applies to
:func:`~airflow.providers.common.ai.toolsets.langchain_bridge.airflow_toolset_to_langchain_tools`,
which converts a chosen toolset for a different agent framework rather than
offering another way to reach a system; :doc:`logging` and :doc:`langchain`
cover
diff --git a/providers/common/ai/docs/toolsets/logging.rst
b/providers/common/ai/docs/toolsets/logging.rst
index a83d3b0984f..6efc61dd874 100644
--- a/providers/common/ai/docs/toolsets/logging.rst
+++ b/providers/common/ai/docs/toolsets/logging.rst
@@ -15,21 +15,18 @@
specific language governing permissions and limitations
under the License.
-Tool call logging: ``LoggingToolset``
-=====================================
+Tool call logging
+=================
.. note::
- Experimental: this can change or be removed in a minor release of this
provider.
- See :ref:`howto/stability`.
+ Experimental: the constructor and the behaviour of
+
:class:`~airflow.providers.common.ai.toolsets.logging.ToolLoggingCapability`
can change
+ or be removed in a minor release of this provider. See
:ref:`howto/stability`.
-:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` is a
-``WrapperToolset`` that intercepts ``call_tool()`` to log each tool invocation
-in real time. ``AgentOperator`` applies it automatically (see
-``enable_tool_logging``) through
-:class:`~airflow.providers.common.ai.toolsets.logging.ToolLoggingCapability`.
-Applying the wrapper as a capability means logging covers the assembled
-function toolset, including tools supplied through ``toolsets=``,
+:class:`~airflow.providers.common.ai.toolsets.logging.ToolLoggingCapability`
wraps the
+assembled function toolset so that each tool invocation is logged in real time.
+Logging covers tools supplied through ``toolsets=``,
``agent_params={"tools": [...]}``, and capabilities such as factory-backed
toolsets, nested capabilities, and MCP toolsets. Output tools such as
``final_result`` and provider-native tools, including native MCP, are not
@@ -44,19 +41,25 @@ any other (see :doc:`progressive_disclosure`).
including inside a combined or wrapper capability, the operator does not add
another, so the logger on the supplied instance is used.
-You can also use ``LoggingToolset`` directly with any pydantic-ai ``Agent``:
+For an agent you build yourself, pass the capability to any pydantic-ai
``Agent``:
.. code-block:: python
- from airflow.providers.common.ai.toolsets.logging import LoggingToolset
+ from pydantic_ai import Agent
+
+ from airflow.providers.common.ai.toolsets.logging import
ToolLoggingCapability
from airflow.providers.common.ai.toolsets.sql import SQLToolset
- logged_toolset = LoggingToolset(wrapped=SQLToolset(db_conn_id="my_db"))
+ agent = Agent(
+ "openai:gpt-5",
+ toolsets=[SQLToolset(db_conn_id="my_db")],
+ capabilities=[ToolLoggingCapability()],
+ )
Each tool call logs its name and timing at INFO, inside a collapsible
``::group::`` block in the task log, and its arguments at DEBUG. Control-flow
signals that let the model retry, report a failed tool result, defer or skip a
call, or wait for approval are logged at INFO and re-raised. Other exceptions,
including ``SkipToolValidation`` raised by a tool body, are logged at ERROR
-with a traceback and re-raised. Pass ``logger`` to send the lines to a logger
-other than the toolset module's own.
+with a traceback and re-raised. Pass ``logger`` to ``ToolLoggingCapability``
to send the
+lines to a logger other than the module's own.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
index 407dad59148..37068585f70 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
@@ -326,9 +326,12 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin, Usa
``agent_params`` still works, but stores each capability's repr
in the serialized Dag, and cannot be combined with this argument (the
task fails when it runs).
- :param enable_tool_logging: When ``True`` (default), wraps the agent's
- assembled function toolset in a ``LoggingToolset`` that logs tool calls
- with timing at INFO level and arguments at DEBUG level. This includes
+ :param enable_tool_logging: When ``True`` (default), the operator appends a
+ ``ToolLoggingCapability``, whose ``get_wrapper_toolset`` wraps the
agent's
+ assembled function toolset so that tool calls are logged
+ with timing at INFO level and arguments at DEBUG level. If
``capabilities``
+ already holds a ``ToolLoggingCapability``, the operator adds none and
that
+ one's ``logger`` is used. Logging includes
tools supplied through ``toolsets=``, ``agent_params["tools"]``, and
capabilities, but not output tools, provider-native tools that run
server-side, or tools another capability adds through its own wrapper
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
index e48a51de444..9a56aa936f2 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
@@ -110,7 +110,9 @@ class LLMSchemaCompareOperator(LLMOperator):
``temperature`` or ``max_tokens``. See
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
:param output_retries: How many times the model is asked again when its
output
fails validation. See
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
- :param data_sources: List of DataSourceConfig objects, one per system.
+ :param data_sources: List of DataSourceConfig objects, one per system. A
+ source whose connection is not a ``DbApiHook`` is read through
DataFusion, which
+ ``common.sql`` pins below 52.
:param db_conn_ids: Connection IDs for databases to compare (used with
``table_names``).
:param table_names: Tables to introspect from each ``db_conn_id``.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
index 931ad82f0a1..468d38f9227 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
@@ -89,6 +89,10 @@ class LLMSQLQueryOperator(LLMOperator):
Default: ``(Select, Union, Intersect, Except)``.
:param dialect: SQL dialect for parsing (``postgres``, ``mysql``, etc.).
Auto-detected from the database hook if not set.
+ :param datasource_config: A
+ :class:`~airflow.providers.common.sql.config.DataSourceConfig` for an
object-storage
+ table whose schema is introspected through DataFusion. The DataFusion
engine of
+ ``common.sql`` pins ``datafusion`` below 52.
``usage_limits`` is inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py
index 9c997f50c4d..b95f1c9491c 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py
@@ -19,7 +19,6 @@
from __future__ import annotations
from airflow.providers.common.ai.toolsets.hook import HookToolset
-from airflow.providers.common.ai.toolsets.logging import LoggingToolset
from airflow.providers.common.ai.toolsets.managed_agent import (
BaseManagedAgentToolset,
ManagedAgentToolset,
@@ -33,7 +32,6 @@ __all__ = [
"BaseManagedAgentToolset",
"DataFusionToolset",
"HookToolset",
- "LoggingToolset",
"MCPToolset",
"ManagedAgentToolset",
"ObjectStorageToolset",
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/logging.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/logging.py
index 4708d685d8b..be9c2f52f32 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/logging.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/logging.py
@@ -42,15 +42,8 @@ if TYPE_CHECKING:
@dataclass
-class LoggingToolset(WrapperToolset[Any]):
- """
- Wrap a toolset to log each tool call with timing.
-
- .. note::
-
- Experimental: this can change or be removed in a minor release of this
provider.
- See :ref:`howto/stability`.
- """
+class _LoggingToolset(WrapperToolset[Any]):
+ """Wrap a toolset to log each tool call with timing; applied through
``ToolLoggingCapability``."""
logger: Logger | logging.Logger = field(default_factory=lambda:
logging.getLogger(__name__))
@@ -98,7 +91,11 @@ class LoggingToolset(WrapperToolset[Any]):
@dataclass
class ToolLoggingCapability(AbstractCapability[Any]):
- """Apply tool-call logging to the assembled function toolset for an agent
run."""
+ """
+ Apply tool-call logging to the assembled function toolset for an agent run.
+
+ .. note:: Experimental; see :ref:`howto/stability`.
+ """
logger: Logger | logging.Logger = field(default_factory=lambda:
logging.getLogger(__name__))
@@ -111,4 +108,4 @@ class ToolLoggingCapability(AbstractCapability[Any]):
return CapabilityOrdering(position="innermost")
def get_wrapper_toolset(self, toolset: AbstractToolset[Any]) ->
AbstractToolset[Any]:
- return LoggingToolset(wrapped=toolset, logger=self.logger)
+ return _LoggingToolset(wrapped=toolset, logger=self.logger)
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index 8ee85870a7b..86d36a1e0af 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -2784,7 +2784,7 @@ class TestAgentOperatorSandboxHandleTemplating:
def test_execute_runs_against_the_rendered_handle(self):
"""
- The whole seam: render, wrap in LoggingToolset, pydantic-ai's for_run
copy,
+ The whole seam: render, wrap in _LoggingToolset, pydantic-ai's for_run
copy,
attach on enter, the command on the given handle, release on exit.
"""
diff --git a/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
b/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
index 400829ba0ff..cdc17b932c7 100644
--- a/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
+++ b/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
@@ -423,7 +423,7 @@ TOOLSET_TABLE_ENTRIES: dict[str, tuple[str, str]] = {
"hook": ("HookToolset", "toolsets/hook"),
"sql": ("SQLToolset", "toolsets/sql"),
"datafusion": ("DataFusionToolset", "toolsets/datafusion"),
- "logging": ("LoggingToolset", "toolsets/logging"),
+ "logging": ("ToolLoggingCapability", "toolsets/logging"),
"mcp": ("MCPToolset", "toolsets/mcp"),
"object_storage": ("ObjectStorageToolset", "toolsets/object_storage"),
"sandbox": ("SandboxToolset", "sandbox/index"),
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_logging.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_logging.py
index 909f8ade0f5..4a919887f89 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_logging.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_logging.py
@@ -37,7 +37,7 @@ from pydantic_ai.models.function import FunctionModel
from pydantic_ai.toolsets import FunctionToolset, PrefixedToolset
from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
-from airflow.providers.common.ai.toolsets.logging import LoggingToolset,
ToolLoggingCapability
+from airflow.providers.common.ai.toolsets.logging import
ToolLoggingCapability, _LoggingToolset
@pytest.fixture
@@ -55,7 +55,7 @@ def logger():
@pytest.fixture
def logging_toolset(wrapped_toolset, logger):
- return LoggingToolset(wrapped=wrapped_toolset, logger=logger)
+ return _LoggingToolset(wrapped=wrapped_toolset, logger=logger)
class TestLoggingToolset:
@@ -218,7 +218,7 @@ class TestToolLoggingCapability:
wrapped =
ToolLoggingCapability(logger=logger).get_wrapper_toolset(toolset)
- assert isinstance(wrapped, LoggingToolset)
+ assert isinstance(wrapped, _LoggingToolset)
assert wrapped.wrapped is toolset
assert wrapped.logger is logger
diff --git a/scripts/in_container/run_provider_yaml_files_check.py
b/scripts/in_container/run_provider_yaml_files_check.py
index ab88b9ebd83..3ec8096a288 100755
--- a/scripts/in_container/run_provider_yaml_files_check.py
+++ b/scripts/in_container/run_provider_yaml_files_check.py
@@ -82,7 +82,7 @@ KNOWN_DEPRECATED_CLASSES = [
# a provider's own runtime rather than constructed by users -- and are
intentionally not
# part of the registry's public "toolsets" module category, so they are exempt
from
# check_all_provider_classes_are_registered's registration requirement.
Contrast with
-# e.g. LoggingToolset, which is documented in the toolsets how-to guide and is
registered.
+# e.g. SQLToolset, which is documented in the toolsets how-to guide and is
registered.
INTERNAL_UNREGISTERED_TOOLSET_CLASSES = {
# Masks what a Dag author's toolset returns; applied automatically by
AgentOperator.
"airflow.providers.common.ai.utils.toolset_base.MaskingToolset",
