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 06e7e284e32 Log capability tool calls in AgentOperator (#72202)
06e7e284e32 is described below
commit 06e7e284e323fd8e7787b54e62aa38db3e236db3
Author: Jeff(Wei-Hao) Lu <[email protected]>
AuthorDate: Mon Oct 12 09:10:31 2026 +0900
Log capability tool calls in AgentOperator (#72202)
---
docs/spelling_wordlist.txt | 4 +
providers/common/ai/docs/capabilities.rst | 10 +-
providers/common/ai/docs/operators/agent.rst | 11 +-
providers/common/ai/docs/toolsets/logging.rst | 30 ++-
.../airflow/providers/common/ai/operators/agent.py | 44 +++--
.../providers/common/ai/toolsets/logging.py | 51 ++++-
.../airflow/providers/common/ai/utils/logging.py | 11 --
.../tests/unit/common/ai/decorators/test_agent.py | 12 +-
.../tests/unit/common/ai/operators/test_agent.py | 219 +++++++++++++++++++--
.../tests/unit/common/ai/toolsets/test_logging.py | 135 ++++++++++++-
.../ai/tests/unit/common/ai/utils/test_logging.py | 17 --
11 files changed, 459 insertions(+), 85 deletions(-)
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index c9ec654ce71..d7e3d0cdfe2 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -288,6 +288,7 @@ Colour
colour
colours
ColumnLineageDatasetFacet
+CombinedCapability
CommandType
comparator
compat
@@ -678,6 +679,7 @@ Filesystem
filesystem
filesystems
filetype
+FilteredToolset
finalStatus
fips
firebase
@@ -1332,6 +1334,7 @@ preflight
prek
Preload
preload
+PreparedToolset
prepend
prepended
prepends
@@ -1971,6 +1974,7 @@ winrm
WIT
workgroup
workspaces
+WrapperToolset
writeable
wsman
WTF
diff --git a/providers/common/ai/docs/capabilities.rst
b/providers/common/ai/docs/capabilities.rst
index 0284396dff6..1a2b76e9959 100644
--- a/providers/common/ai/docs/capabilities.rst
+++ b/providers/common/ai/docs/capabilities.rst
@@ -85,9 +85,13 @@ templated inside a ``Toolset`` capability the same way as in
``toolsets=``:
A ``Toolset`` capability built from a function is resolved when the run
starts, so its
connection IDs are not templated.
-Tool results from a ``Toolset`` capability are masked like those from
``toolsets=``, but
-``enable_tool_logging`` only logs calls to ``toolsets=``. Pass a toolset in
``toolsets=`` unless
-you need it inside the capability list, for example to order it against a
guardrail.
+Tool results from a ``Toolset`` capability are masked like those from
``toolsets=``.
+When ``enable_tool_logging=True`` (the default), ``AgentOperator`` logs calls
to
+these toolsets and other tools contributed by capabilities, including MCP
+toolsets. Output tools and provider-native tools, including native MCP, are not
+covered. See :doc:`toolsets/logging` for configuration details and the complete
+limitations. Put a toolset in the capability list when you need to order it
+against another capability, such as a guardrail.
With durable execution
----------------------
diff --git a/providers/common/ai/docs/operators/agent.rst
b/providers/common/ai/docs/operators/agent.rst
index 1f03bf00d0b..08f01e4f475 100644
--- a/providers/common/ai/docs/operators/agent.rst
+++ b/providers/common/ai/docs/operators/agent.rst
@@ -373,9 +373,12 @@ Parameters
``AgentSkillsToolset`` for :ref:`agent-skills`, etc.).
- ``capabilities``: List of pydantic-ai capabilities (``Thinking``,
``WebSearch``, guardrails,
etc.). See :ref:`capabilities`.
-- ``enable_tool_logging``: Wrap each toolset in
+- ``enable_tool_logging``: Wrap the assembled function toolset in
:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` so that
- every tool call is logged in real time. Default ``True``.
+ tools supplied through ``toolsets=``, ``agent_params={"tools": [...]}``, and
+ capabilities are logged in real time. 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``.
- ``agent_params``: Additional keyword arguments passed to the pydantic-ai
``Agent`` constructor (e.g. ``retries``, ``model_settings``).
- .. _agent-usage-budget:
@@ -458,8 +461,8 @@ Logging
-------
All AI operators automatically log a post-run summary after ``run_sync()``
-completes. ``AgentOperator`` additionally wraps toolsets for real-time
-per-tool-call logging (controlled by ``enable_tool_logging``).
+completes. ``AgentOperator`` additionally wraps the assembled function toolset
+for real-time per-tool-call logging (controlled by ``enable_tool_logging``).
**Real-time tool call logging** (AgentOperator only): each tool call is
logged as it happens:
diff --git a/providers/common/ai/docs/toolsets/logging.rst
b/providers/common/ai/docs/toolsets/logging.rst
index 76e300bdbe4..6cba50b98b5 100644
--- a/providers/common/ai/docs/toolsets/logging.rst
+++ b/providers/common/ai/docs/toolsets/logging.rst
@@ -26,8 +26,23 @@ Tool call logging: ``LoggingToolset``
: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``), but you can also use it directly with any pydantic-ai
-``Agent``:
+``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=``,
+``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
+covered by Airflow's real-time tool-call logging. Tools added by another
+capability through its own wrapper toolset, such as ToolSearch's
+``search_tools`` and CodeMode's ``run_code``, are also not covered.
+
+``AgentOperator`` adds ``ToolLoggingCapability`` automatically when
+``enable_tool_logging=True``. If ``capabilities=`` already contains one,
+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``:
.. code-block:: python
@@ -36,7 +51,10 @@ in real time. ``AgentOperator`` applies it automatically (see
logged_toolset = LoggingToolset(wrapped=SQLToolset(db_conn_id="my_db"))
-Each call logs the tool's name and how long it took at INFO, inside a
collapsible
-``::group::`` block in the task log, and its arguments at DEBUG. A call that
raises
-is logged with its traceback and the exception is re-raised. Pass ``logger`` to
-send the lines to a logger other than the toolset module's own.
+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.
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 f856756c05e..cc073097bb2 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
@@ -51,13 +51,13 @@ from airflow.providers.common.ai.observability import (
make_task_instance_run_key,
stamp_identity_on_agent_spans,
)
+from airflow.providers.common.ai.toolsets.logging import ToolLoggingCapability
from airflow.providers.common.ai.toolsets.sandbox import SandboxToolset
from airflow.providers.common.ai.utils.logging import (
MODEL_NAME_XCOM_KEY,
format_usage_for_xcom,
log_run_summary,
log_run_usage,
- wrap_toolsets_for_logging,
)
from airflow.providers.common.ai.utils.output_type import
rehydrate_pydantic_output
from airflow.providers.common.ai.utils.prompt_cache import PromptCaching
@@ -168,24 +168,19 @@ def _replace_capability_toolset(
return capability
-def _contains_code_mode(capabilities: Iterable[AgentCapability[Any]]) -> bool:
+def _contains_capability(
+ capabilities: Iterable[AgentCapability[Any]], capability_type:
type[AbstractCapability[Any]]
+) -> bool:
"""
- Whether any capability, or one nested inside a combined or wrapper
capability, is ``CodeMode``.
+ Whether any capability, or one nested inside a combined or wrapper
capability, has the given type.
A capability function, or a ``DynamicCapability``, builds its capability
when the run
starts, so there is nothing to inspect here.
"""
- # CodeMode's own module is in sys.modules once CodeMode has been imported,
and only then.
- # The pydantic_ai_harness package root is not a safe place to look: it
exports CodeMode
- # through a module __getattr__ that imports that module, which fails
without the
- # ``code-mode`` extra.
- code_mode_cls = getattr(sys.modules.get("pydantic_ai_harness.code_mode"),
"CodeMode", None)
- if code_mode_cls is None:
- return False
pending = [capability for capability in capabilities if
isinstance(capability, AbstractCapability)]
while pending:
capability = pending.pop()
- if isinstance(capability, code_mode_cls):
+ if isinstance(capability, capability_type):
return True
if isinstance(capability, WrapperCapability):
# apply() does not visit a wrapper's single wrapped capability,
only a combined one's children.
@@ -197,6 +192,16 @@ def _contains_code_mode(capabilities:
Iterable[AgentCapability[Any]]) -> bool:
return False
+def _contains_code_mode(capabilities: Iterable[AgentCapability[Any]]) -> bool:
+ """Whether the declared capabilities contain ``CodeMode``."""
+ # CodeMode's own module is in sys.modules once CodeMode has been imported,
and only then.
+ # The pydantic_ai_harness package root is not a safe place to look: it
exports CodeMode
+ # through a module __getattr__ that imports that module, which fails
without the
+ # ``code-mode`` extra.
+ code_mode_cls = getattr(sys.modules.get("pydantic_ai_harness.code_mode"),
"CodeMode", None)
+ return code_mode_cls is not None and _contains_capability(capabilities,
code_mode_cls)
+
+
def _declares_agent_template_fields(toolset: Any) -> bool:
"""Whether *toolset*, or a toolset it wraps or combines, has connection
IDs to render."""
return isinstance(toolset, AbstractToolset) and any(
@@ -269,9 +274,14 @@ 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 each toolset in
a
- ``LoggingToolset`` that logs tool calls with timing at INFO level and
- arguments at DEBUG level. Set to ``False`` to disable.
+ :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
+ 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
+ (code mode's ``run_code``, ToolSearch's ``search_tools``). Set to
+ ``False`` to disable.
:param agent_params: Additional keyword arguments passed to the pydantic-ai
``Agent`` constructor (e.g. ``retries``, ``model_settings``).
:param usage_limits: Optional pydantic-ai
@@ -712,8 +722,6 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator,
HITLReviewMixin, Usa
toolsets: list[AbstractToolset] = [ensure_masked(ts) for ts in
self.toolsets]
if self.durable and storage is not None and counter is not None:
toolsets = self._build_durable_toolsets(toolsets, storage,
counter)
- if self.enable_tool_logging:
- toolsets = wrap_toolsets_for_logging(toolsets, self.log)
extra_kwargs["toolsets"] = toolsets
elif extra_kwargs.get("toolsets"):
extra_kwargs["toolsets"] = [ensure_masked(ts) for ts in
extra_kwargs["toolsets"]]
@@ -728,6 +736,10 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin, Usa
capabilities = self._build_durable_capabilities(capabilities,
storage, counter)
if self.cache_prompt:
capabilities.append(PromptCaching())
+ if self.enable_tool_logging and not _contains_capability(capabilities,
ToolLoggingCapability):
+ # ToolLoggingCapability's innermost ordering keeps logging inside
capability wrappers,
+ # including CodeModeToolset where code mode expects the wrapped
tools.
+ capabilities.append(ToolLoggingCapability(logger=self.log))
if capabilities:
extra_kwargs["capabilities"] = capabilities
return self.llm_hook.create_agent(
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 805f7cf32e2..4708d685d8b 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
@@ -18,17 +18,25 @@
from __future__ import annotations
+import asyncio
import json
import logging
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
-from pydantic_ai.exceptions import ApprovalRequired
+from pydantic_ai.capabilities import AbstractCapability, CapabilityOrdering
+from pydantic_ai.exceptions import (
+ ApprovalRequired,
+ CallDeferred,
+ ModelRetry,
+ SkipToolExecution,
+ ToolFailed,
+)
from pydantic_ai.toolsets.wrapper import WrapperToolset
if TYPE_CHECKING:
- from pydantic_ai.toolsets.abstract import ToolsetTool
+ from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
from airflow.sdk.types import Logger
@@ -54,8 +62,12 @@ class LoggingToolset(WrapperToolset[Any]):
tool: ToolsetTool[Any],
) -> Any:
self.logger.info("::group::Tool call: %s", name)
- if tool_args:
- self.logger.debug("Tool args: %s", json.dumps(tool_args,
default=str))
+ if tool_args and self.logger.isEnabledFor(logging.DEBUG):
+ try:
+ serialized_args = json.dumps(tool_args, default=str)
+ except (TypeError, ValueError):
+ serialized_args = repr(tool_args)
+ self.logger.debug("Tool args: %s", serialized_args)
start = time.monotonic()
try:
result = await self.wrapped.call_tool(name, tool_args, ctx, tool)
@@ -63,13 +75,40 @@ class LoggingToolset(WrapperToolset[Any]):
self.logger.info("Tool %s returned in %.2fs", name, elapsed)
self.logger.info("::endgroup::")
return result
+ except asyncio.CancelledError:
+ self.logger.info("::endgroup::")
+ raise
except ApprovalRequired:
# Not a failure: the run pauses here until a person approves or
rejects the call.
- self.logger.info("Tool %s is waiting to be approved", name)
+ elapsed = time.monotonic() - start
+ self.logger.info("Tool %s is waiting to be approved after %.2fs",
name, elapsed)
+ self.logger.info("::endgroup::")
+ raise
+ except (ModelRetry, ToolFailed, CallDeferred, SkipToolExecution) as e:
+ elapsed = time.monotonic() - start
+ self.logger.info("Tool %s requested %s after %.2fs: %s", name,
type(e).__name__, elapsed, e)
self.logger.info("::endgroup::")
raise
except Exception:
elapsed = time.monotonic() - start
- self.logger.exception("Tool %s failed after %.2fs", name, elapsed)
self.logger.info("::endgroup::")
+ self.logger.exception("Tool %s failed after %.2fs", name, elapsed)
raise
+
+
+@dataclass
+class ToolLoggingCapability(AbstractCapability[Any]):
+ """Apply tool-call logging to the assembled function toolset for an agent
run."""
+
+ logger: Logger | logging.Logger = field(default_factory=lambda:
logging.getLogger(__name__))
+
+ @classmethod
+ def get_serialization_name(cls) -> str | None:
+ return None
+
+ def get_ordering(self) -> CapabilityOrdering:
+ """Place logging innermost so outer capability wrappers delegate
base-tool calls through it."""
+ return CapabilityOrdering(position="innermost")
+
+ def get_wrapper_toolset(self, toolset: AbstractToolset[Any]) ->
AbstractToolset[Any]:
+ return LoggingToolset(wrapped=toolset, logger=self.logger)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
index 57b1380705e..7cd0ecaafff 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py
@@ -24,11 +24,8 @@ from typing import TYPE_CHECKING, Any
from pydantic import BaseModel
from pydantic_ai.messages import ToolCallPart
-from airflow.providers.common.ai.toolsets.logging import LoggingToolset
-
if TYPE_CHECKING:
from pydantic_ai.result import AgentRunResult
- from pydantic_ai.toolsets.abstract import AbstractToolset
from pydantic_ai.usage import RunUsage
from airflow.sdk.types import Logger
@@ -137,11 +134,3 @@ def _extract_tool_sequence(result: AgentRunResult[Any]) ->
list[str]:
for part in getattr(message, "parts", [])
if isinstance(part, ToolCallPart)
]
-
-
-def wrap_toolsets_for_logging(
- toolsets: list[AbstractToolset[Any]],
- logger: Logger | logging.Logger,
-) -> list[AbstractToolset[Any]]:
- """Wrap each toolset in a LoggingToolset."""
- return [LoggingToolset(wrapped=ts, logger=logger) for ts in toolsets]
diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
index fffdd51a15f..a1c4a5330fa 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py
@@ -25,7 +25,7 @@ from pydantic_ai.messages import ImageUrl
from pydantic_ai.toolsets.function import FunctionToolset
from airflow.providers.common.ai.decorators.agent import
_AgentDecoratedOperator
-from airflow.providers.common.ai.toolsets.logging import LoggingToolset
+from airflow.providers.common.ai.toolsets.logging import ToolLoggingCapability
from airflow.providers.common.ai.utils.prompt_cache import PromptCaching
from airflow.providers.common.ai.utils.toolset_base import MaskingToolset
@@ -176,10 +176,9 @@ class TestAgentDecoratedOperator:
op.execute(context=_make_context())
create_call =
mock_hook_cls.get_hook.return_value.create_agent.call_args
- passed_toolsets = create_call[1]["toolsets"]
- assert len(passed_toolsets) == 1
- assert isinstance(passed_toolsets[0], LoggingToolset)
- assert passed_toolsets[0].wrapped == MaskingToolset(wrapped=toolset)
+ assert create_call[1]["toolsets"] == [MaskingToolset(wrapped=toolset)]
+ assert create_call[1]["capabilities"][0] == PromptCaching()
+ assert isinstance(create_call[1]["capabilities"][1],
ToolLoggingCapability)
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_passes_capabilities_through(self, mock_hook_cls,
make_mock_run_result):
@@ -197,7 +196,8 @@ class TestAgentDecoratedOperator:
op.execute(context=_make_context())
create_call =
mock_hook_cls.get_hook.return_value.create_agent.call_args
- assert create_call.kwargs["capabilities"] == [thinking,
PromptCaching()]
+ assert create_call.kwargs["capabilities"][:2] == [thinking,
PromptCaching()]
+ assert isinstance(create_call.kwargs["capabilities"][2],
ToolLoggingCapability)
@requires_typed_xcom
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
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 2b4b2bd3bbe..4f09990822c 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
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import asyncio
import dataclasses
import json
import logging
@@ -26,6 +27,7 @@ from decimal import Decimal
from types import ModuleType, SimpleNamespace
from typing import Any
from unittest.mock import ANY, MagicMock, PropertyMock, call, patch
+from uuid import UUID
import pytest
from pydantic import BaseModel
@@ -73,7 +75,7 @@ from airflow.providers.common.ai.sandbox.base import (
SandboxExecResult,
)
from airflow.providers.common.ai.toolsets.hook import HookToolset
-from airflow.providers.common.ai.toolsets.logging import LoggingToolset
+from airflow.providers.common.ai.toolsets.logging import ToolLoggingCapability
from airflow.providers.common.ai.toolsets.mcp import MCPToolset
from airflow.providers.common.ai.toolsets.sandbox import SandboxToolset
from airflow.providers.common.ai.toolsets.sql import SQLToolset
@@ -757,9 +759,12 @@ class TestAgentOperatorExecute:
)
# On 3.3+ the agent may also end on a tool call awaiting approval.
expected_output_type = [str, DeferredToolRequests] if
AIRFLOW_V_3_3_PLUS else str
-
mock_hook_cls.get_hook.return_value.create_agent.assert_called_once_with(
- output_type=expected_output_type, instructions="You are helpful.",
capabilities=[PromptCaching()]
- )
+ create_agent_kwargs =
mock_hook_cls.get_hook.return_value.create_agent.call_args.kwargs
+ assert create_agent_kwargs["output_type"] == expected_output_type
+ assert create_agent_kwargs["instructions"] == "You are helpful."
+ assert len(create_agent_kwargs["capabilities"]) == 2
+ assert create_agent_kwargs["capabilities"][0] == PromptCaching()
+ assert isinstance(create_agent_kwargs["capabilities"][1],
ToolLoggingCapability)
mock_agent.run_sync.assert_called_once_with(
"What is the answer?", usage_limits=None, run_id="ti-1",
cancellation_token=ANY, usage=ANY
)
@@ -795,11 +800,9 @@ class TestAgentOperatorExecute:
op.execute(context=MagicMock())
create_call =
mock_hook_cls.get_hook.return_value.create_agent.call_args
- passed_toolsets = create_call[1]["toolsets"]
- assert len(passed_toolsets) == 1
- assert isinstance(passed_toolsets[0], LoggingToolset)
- assert isinstance(passed_toolsets[0].wrapped, MaskingToolset)
- assert passed_toolsets[0].wrapped.wrapped is mock_toolset
+ assert create_call[1]["toolsets"] ==
[MaskingToolset(wrapped=mock_toolset)]
+ assert create_call[1]["capabilities"][0] == PromptCaching()
+ assert isinstance(create_call[1]["capabilities"][1],
ToolLoggingCapability)
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_enable_tool_logging_false_skips_wrapping(self, mock_hook_cls,
make_mock_run_result):
@@ -820,6 +823,126 @@ class TestAgentOperatorExecute:
create_call =
mock_hook_cls.get_hook.return_value.create_agent.call_args
assert create_call[1]["toolsets"] ==
[MaskingToolset(wrapped=mock_toolset)]
+ assert create_call[1]["capabilities"] == [PromptCaching()]
+
+ @pytest.mark.parametrize("tool_source", ["tools", "agent_params",
"capabilities"])
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_tool_logging_wraps_tools_from_all_sources(self, mock_hook_cls,
caplog, tool_source):
+ """Tool logging wraps the complete toolset assembled by pydantic-ai."""
+
+ def my_tool() -> str:
+ return "tool-result"
+
+ def model_fn(messages, info):
+ saw_return = any(isinstance(p, ToolReturnPart) for m in messages
for p in getattr(m, "parts", []))
+ if saw_return:
+ return ModelResponse(parts=[TextPart(content="done")])
+ return ModelResponse(parts=[ToolCallPart(tool_name="my_tool",
args={}, tool_call_id="c1")])
+
+ mock_hook_cls.get_hook.return_value.create_agent.side_effect = lambda
**kw: Agent(
+ FunctionModel(model_fn), **kw
+ )
+ if tool_source == "tools":
+ kwargs = {"agent_params": {"tools": [my_tool]}}
+ else:
+ capability = Toolset(FunctionToolset([my_tool]))
+ kwargs = (
+ {"agent_params": {"capabilities": [capability]}}
+ if tool_source == "agent_params"
+ else {"capabilities": [capability]}
+ )
+ op = AgentOperator(
+ task_id="test",
+ prompt="Do something",
+ llm_conn_id="my_llm",
+ **kwargs,
+ )
+
+ with caplog.at_level(logging.INFO):
+ result = op.execute(context=_make_context())
+
+ assert result == "done"
+ assert any(record.message == "::group::Tool call: my_tool" for record
in caplog.records)
+
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_tool_logging_closes_groups_when_parallel_tool_fails(self,
mock_hook_cls, caplog):
+ slow_cancelled = False
+
+ async def boom() -> str:
+ await asyncio.sleep(0)
+ raise RuntimeError("boom")
+
+ async def slow() -> None:
+ nonlocal slow_cancelled
+ try:
+ await asyncio.sleep(60)
+ except asyncio.CancelledError:
+ slow_cancelled = True
+ raise
+
+ def model_fn(messages, info):
+ return ModelResponse(
+ parts=[
+ ToolCallPart(tool_name="boom", args={},
tool_call_id="c-boom"),
+ ToolCallPart(tool_name="slow", args={},
tool_call_id="c-slow"),
+ ]
+ )
+
+ mock_hook_cls.get_hook.return_value.create_agent.side_effect = lambda
**kw: Agent(
+ FunctionModel(model_fn), **kw
+ )
+ op = AgentOperator(
+ task_id="test",
+ prompt="Do something",
+ llm_conn_id="my_llm",
+ agent_params={"tools": [boom, slow]},
+ )
+
+ with caplog.at_level(logging.INFO):
+ with pytest.raises(RuntimeError, match="boom"):
+ op.execute(context=_make_context())
+
+ assert slow_cancelled
+ messages = [record.message for record in caplog.records]
+ assert sum(message.startswith("::group::Tool call:") for message in
messages) == 2
+ assert messages.count("::endgroup::") == 2
+
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_tool_logging_handles_non_json_mapping_keys(self, mock_hook_cls,
caplog):
+ def my_tool(values: dict[UUID, int]) -> int:
+ return sum(values.values())
+
+ def model_fn(messages, info):
+ saw_return = any(isinstance(p, ToolReturnPart) for m in messages
for p in getattr(m, "parts", []))
+ if saw_return:
+ return ModelResponse(parts=[TextPart(content="done")])
+ return ModelResponse(
+ parts=[
+ ToolCallPart(
+ tool_name="my_tool",
+ args={"values":
{"12345678-1234-5678-1234-567812345678": 1}},
+ tool_call_id="c1",
+ )
+ ]
+ )
+
+ mock_hook_cls.get_hook.return_value.create_agent.side_effect = lambda
**kw: Agent(
+ FunctionModel(model_fn), **kw
+ )
+ op = AgentOperator(
+ task_id="test",
+ prompt="Do something",
+ llm_conn_id="my_llm",
+ agent_params={"tools": [my_tool]},
+ )
+
+ with caplog.at_level(logging.DEBUG, logger="airflow.task"):
+ result = op.execute(context=_make_context())
+
+ assert result == "done"
+ assert any(
+ "UUID('12345678-1234-5678-1234-567812345678')" in record.message
for record in caplog.records
+ )
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_passes_agent_params(self, mock_hook_cls,
make_mock_run_result):
@@ -853,6 +976,7 @@ class TestAgentOperatorExecute:
llm_conn_id="my_llm",
toolsets=[MagicMock(spec=AbstractToolset)],
cache_prompt=False,
+ enable_tool_logging=False,
)
op.execute(context=_make_context())
@@ -867,7 +991,11 @@ class TestAgentOperatorExecute:
)
op = AgentOperator(
- task_id="t", prompt="hi", llm_conn_id="my_llm",
agent_params={"capabilities": ["existing"]}
+ task_id="t",
+ prompt="hi",
+ llm_conn_id="my_llm",
+ agent_params={"capabilities": ["existing"]},
+ enable_tool_logging=False,
)
op.execute(context=_make_context())
@@ -1191,12 +1319,51 @@ class TestAgentOperatorCapabilities:
)
thinking, search = Thinking(effort="high"), WebSearch()
- op = AgentOperator(task_id="t", prompt="p", llm_conn_id="llm",
capabilities=[thinking, search])
+ op = AgentOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="llm",
+ capabilities=[thinking, search],
+ enable_tool_logging=False,
+ )
op.execute(context=_make_context())
create_call =
mock_hook_cls.get_hook.return_value.create_agent.call_args
assert create_call.kwargs["capabilities"] == [thinking, search,
PromptCaching()]
+ @pytest.mark.parametrize(
+ "wrap",
+ [
+ pytest.param(lambda capability: capability, id="direct"),
+ pytest.param(lambda capability: CombinedCapability([Thinking(),
capability]), id="combined"),
+ pytest.param(lambda capability: PrefixTools(wrapped=capability,
prefix="custom"), id="wrapped"),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_declared_tool_logging_capability_prevents_default(
+ self, mock_hook_cls, make_mock_run_result, wrap
+ ):
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
+ "ok", make_mock_run_result
+ )
+ custom_logging =
ToolLoggingCapability(logger=MagicMock(spec=logging.Logger))
+ declared = wrap(custom_logging)
+ op = AgentOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="llm",
+ capabilities=[declared],
+ )
+
+ op.execute(context=_make_context())
+
+ capabilities =
mock_hook_cls.get_hook.return_value.create_agent.call_args.kwargs["capabilities"]
+ assert capabilities[0] is declared
+ assert not any(
+ isinstance(capability, ToolLoggingCapability) and capability is
not custom_logging
+ for capability in capabilities
+ )
+
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_capabilities_in_both_places_are_refused(self, mock_hook_cls):
op = AgentOperator(
@@ -1504,6 +1671,36 @@ class TestAgentOperatorDurable:
assert result[0] is cap
+ @pytest.mark.parametrize("passed_as", ["capabilities", "agent_params"])
+ def test_toolset_capability_logging_wraps_durable_cache(self, passed_as):
+ """Concrete ``Toolset`` capabilities receive durable caching before
agent creation."""
+ inner = FunctionToolset()
+ capability = Toolset(inner)
+ kwargs = (
+ {"capabilities": [capability]}
+ if passed_as == "capabilities"
+ else {"agent_params": {"capabilities": [capability]}}
+ )
+ op = AgentOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ durable=True,
+ **kwargs,
+ )
+ op._durable_storage = MagicMock(spec=DurableStorageProtocol)
+ op._durable_counter = DurableStepCounter()
+ hook = MagicMock(spec=["create_agent"])
+ op.llm_hook = hook
+
+ op._build_agent()
+
+ capabilities = hook.create_agent.call_args.kwargs["capabilities"]
+ assert isinstance(capabilities[0].toolset, CachingToolset)
+ assert isinstance(capabilities[0].toolset.wrapped, MaskingToolset)
+ assert capabilities[0].toolset.wrapped.wrapped is inner
+ assert isinstance(capabilities[-1], ToolLoggingCapability)
+
def test_toolset_capability_tool_replayed_on_retry(self):
"""A tool supplied via a ``Toolset`` capability is cached and replayed
on a
retry instead of re-executing. Regression: such tools bypassed the
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 6ffecbf4926..909f8ade0f5 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
@@ -16,15 +16,28 @@
# under the License.
from __future__ import annotations
+import asyncio
import logging
+from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
-from pydantic_ai import RunContext
-from pydantic_ai.exceptions import ApprovalRequired
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.exceptions import (
+ ApprovalRequired,
+ CallDeferred,
+ ModelRetry,
+ SkipToolExecution,
+ SkipToolValidation,
+ ToolFailed,
+)
+from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart,
ToolReturnPart
+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
+from airflow.providers.common.ai.toolsets.logging import LoggingToolset,
ToolLoggingCapability
@pytest.fixture
@@ -91,8 +104,78 @@ class TestLoggingToolset:
with pytest.raises(RuntimeError, match="boom"):
await logging_toolset.call_tool("bad_tool", {}, ctx, tool)
- assert any("Tool bad_tool failed after" in r.message for r in
caplog.records)
- assert any("::endgroup::" in r.message for r in caplog.records)
+ end_group_index = next(
+ i for i, record in enumerate(caplog.records) if record.message ==
"::endgroup::"
+ )
+ failure_index = next(
+ i for i, record in enumerate(caplog.records) if "Tool bad_tool
failed after" in record.message
+ )
+ assert end_group_index < failure_index
+ assert caplog.records[failure_index].levelno == logging.ERROR
+
+ @pytest.mark.asyncio
+ async def test_closes_group_when_call_is_cancelled(
+ self, logging_toolset, wrapped_toolset, logger, caplog
+ ):
+ wrapped_toolset.call_tool =
AsyncMock(side_effect=asyncio.CancelledError())
+
+ with caplog.at_level(logging.INFO, logger="test.logging_toolset"):
+ with pytest.raises(asyncio.CancelledError):
+ await logging_toolset.call_tool("slow_tool", {}, MagicMock(),
MagicMock())
+
+ assert caplog.records[-1].message == "::endgroup::"
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ "signal",
+ [
+ ModelRetry("retry"),
+ ToolFailed("failed"),
+ CallDeferred(),
+ SkipToolExecution("result"),
+ ],
+ ids=lambda signal: type(signal).__name__,
+ )
+ async def test_logs_control_flow_at_info(self, logging_toolset,
wrapped_toolset, logger, caplog, signal):
+ wrapped_toolset.call_tool = AsyncMock(side_effect=signal)
+
+ with caplog.at_level(logging.INFO, logger="test.logging_toolset"):
+ with pytest.raises(type(signal)):
+ await logging_toolset.call_tool("retrying_tool", {},
MagicMock(), MagicMock())
+
+ assert any(
+ record.message.startswith(f"Tool retrying_tool requested
{type(signal).__name__} after")
+ and record.levelno == logging.INFO
+ for record in caplog.records
+ )
+ assert not any(record.levelno >= logging.ERROR for record in
caplog.records)
+ assert caplog.records[-1].message == "::endgroup::"
+
+ @pytest.mark.asyncio
+ async def test_logs_control_flow_reason(self, logging_toolset,
wrapped_toolset, logger, caplog):
+ wrapped_toolset.call_tool = AsyncMock(side_effect=ModelRetry("column X
not found"))
+
+ with caplog.at_level(logging.INFO, logger="test.logging_toolset"):
+ with pytest.raises(ModelRetry, match="column X not found"):
+ await logging_toolset.call_tool("retrying_tool", {},
MagicMock(), MagicMock())
+
+ assert any("column X not found" in record.message for record in
caplog.records)
+
+ @pytest.mark.asyncio
+ async def test_skip_tool_validation_is_logged_as_error(
+ self, logging_toolset, wrapped_toolset, logger, caplog
+ ):
+ signal = SkipToolValidation({"value": 1})
+ wrapped_toolset.call_tool = AsyncMock(side_effect=signal)
+
+ with caplog.at_level(logging.INFO, logger="test.logging_toolset"):
+ with pytest.raises(SkipToolValidation):
+ await logging_toolset.call_tool("invalid_tool", {"value": 1},
MagicMock(), MagicMock())
+
+ assert any(
+ record.message.startswith("Tool invalid_tool failed after") and
record.levelno == logging.ERROR
+ for record in caplog.records
+ )
@pytest.mark.asyncio
async def test_a_call_waiting_for_approval_is_not_logged_as_a_failure(
@@ -127,3 +210,45 @@ class TestLoggingToolset:
await logging_toolset.call_tool("list_tables", {}, ctx, tool)
assert not any("Tool args:" in r.message for r in caplog.records)
+
+
+class TestToolLoggingCapability:
+ def test_wraps_assembled_toolset(self, logger):
+ toolset = FunctionToolset()
+
+ wrapped =
ToolLoggingCapability(logger=logger).get_wrapper_toolset(toolset)
+
+ assert isinstance(wrapped, LoggingToolset)
+ assert wrapped.wrapped is toolset
+ assert wrapped.logger is logger
+
+ def test_ordering_keeps_logging_inside_other_capability_wrappers(self,
logger, caplog):
+ class PrefixingCapability(AbstractCapability[Any]):
+ def get_wrapper_toolset(self, toolset):
+ return PrefixedToolset(toolset, prefix="outer")
+
+ def my_tool() -> str:
+ return "tool-result"
+
+ def model_fn(messages, info):
+ if any(isinstance(part, ToolReturnPart) for message in messages
for part in message.parts):
+ return ModelResponse(parts=[TextPart(content="done")])
+ return ModelResponse(
+ parts=[ToolCallPart(tool_name="outer_my_tool", args={},
tool_call_id="call-1")]
+ )
+
+ agent = Agent(
+ FunctionModel(model_fn),
+ tools=[my_tool],
+ capabilities=[PrefixingCapability(),
ToolLoggingCapability(logger=logger)],
+ )
+
+ with caplog.at_level(logging.INFO, logger="test.logging_toolset"):
+ result = agent.run_sync("run the tool")
+
+ assert result.output == "done"
+ assert any(record.message == "::group::Tool call: my_tool" for record
in caplog.records)
+ assert not any(record.message == "::group::Tool call: outer_my_tool"
for record in caplog.records)
+
+ def test_is_not_serializable(self):
+ assert ToolLoggingCapability.get_serialization_name() is None
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
b/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
index 6733315f66b..49273a49341 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_logging.py
@@ -34,13 +34,11 @@ from pydantic_ai.models.function import FunctionModel
from pydantic_ai.models.test import TestModel
from pydantic_ai.usage import RunUsage
-from airflow.providers.common.ai.toolsets.logging import LoggingToolset
from airflow.providers.common.ai.utils.logging import (
_log_output_debug,
format_usage_for_xcom,
log_run_summary,
log_run_usage,
- wrap_toolsets_for_logging,
)
@@ -342,18 +340,3 @@ class TestLogOutputDebug:
debug_records = [r for r in caplog.records if r.levelno ==
logging.DEBUG]
assert len(debug_records) == 0
-
-
-class TestWrapToolsetsForLogging:
- def test_wraps_each_toolset(self):
- ts_a = MagicMock()
- ts_b = MagicMock()
- logger = logging.getLogger("test.wrap")
-
- wrapped = wrap_toolsets_for_logging([ts_a, ts_b], logger)
-
- assert len(wrapped) == 2
- assert all(isinstance(w, LoggingToolset) for w in wrapped)
- assert wrapped[0].wrapped is ts_a
- assert wrapped[1].wrapped is ts_b
- assert wrapped[0].logger is logger
