kaxil commented on code in PR #72202:
URL: https://github.com/apache/airflow/pull/72202#discussion_r3963573360
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/logging.py:
##########
@@ -60,3 +61,17 @@ async def call_tool(
self.logger.exception("Tool %s failed after %.2fs", name, elapsed)
self.logger.info("::endgroup::")
raise
+
+
+@dataclass
+class ToolLoggingCapability(AbstractCapability[Any]):
Review Comment:
The docs spellcheck job went red with this commit (it passed on 0e55597).
The generated `_api/.../toolsets/logging/index.rst` now carries
`AbstractCapability`'s inherited docstrings, so `AbstractCapability`,
`FilteredToolset`, `PreparedToolset` and `WrapperToolset` all get flagged, and
`docs/spelling_wordlist.txt` only has `AbstractToolset` today. Adding those
four words there should clear it.
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -242,6 +240,53 @@ def test_enable_tool_logging_false_skips_wrapping(self,
mock_hook_cls):
create_call =
mock_hook_cls.get_hook.return_value.create_agent.call_args
assert create_call[1]["toolsets"] == [mock_toolset]
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_tool_logging_wraps_assembled_capability_toolsets(self,
mock_hook_cls):
Review Comment:
All three new assertions call `get_wrapper_toolset()` by hand on a stand-in
`FunctionToolset`, so the part that can actually drift stays unasserted: that
pydantic-ai calls the hook at all, and that it hands over the fully assembled
non-output toolset. The base hook returns `None`, so if that contract moves
anywhere in the unbounded `>=2.0.0` range the logging degrades to a silent
no-op with the suite still green, and this PR gave up the assertion that used
to catch it (LoggingToolset instances visible in the `toolsets=` list). One
real run through the `Agent(FunctionModel(...))` harness at line 815 with
`enable_tool_logging=True` plus a caplog check for `::group::Tool call:` would
pin it down.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -335,6 +334,10 @@ def _build_agent(self) -> Agent[object, Any]:
capabilities = self._build_durable_capabilities(capabilities,
storage, counter)
if self.code_mode:
capabilities.append(_build_code_mode())
+ if self.enable_tool_logging and (self.toolsets or capabilities):
Review Comment:
The `(self.toolsets or capabilities)` half of this gate skips logging for an
agent whose tools arrive only through `agent_params={"tools": [...]}`, which
agent.rst says is forwarded to the `Agent`. Those tools do reach the assembled
toolset (they show up as `_AgentFunctionToolset` inside the combined toolset
the wrapper receives), so they would be covered once the capability is
appended. Wrapping an empty toolset logs nothing, so `if
self.enable_tool_logging:` on its own looks sufficient.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -335,6 +334,10 @@ def _build_agent(self) -> Agent[object, Any]:
capabilities = self._build_durable_capabilities(capabilities,
storage, counter)
if self.code_mode:
capabilities.append(_build_code_mode())
+ if self.enable_tool_logging and (self.toolsets or capabilities):
+ from airflow.providers.common.ai.toolsets.logging import
ToolLoggingCapability
Review Comment:
This import isn't actually lazy: line 32 pulls in `utils.logging`, which
imports `toolsets.logging` at module level for `LoggingToolset`, so it is
already in `sys.modules` long before `_build_agent` runs. Worth hoisting to the
top. Related, `wrap_toolsets_for_logging` in `utils/logging.py` has no
production caller left, only its own test, so dropping both finishes the
migration and removes that eager chain.
##########
providers/common/ai/docs/operators/agent.rst:
##########
@@ -374,6 +374,10 @@ but anything passed through ``agent_params`` is forwarded
to the underlying
Capabilities compose with toolsets -- pydantic-ai merges tools from both.
+When ``enable_tool_logging=True`` (the default), ``AgentOperator`` applies
+real-time tool-call logging to the complete toolset assembled from
capabilities,
+including factory-backed toolsets, nested capabilities, and MCP tools.
Review Comment:
`MCP` is covered because it defaults to the local toolset, but `WebSearch`,
`WebFetch` and `ImageGeneration`, listed as common capabilities a few lines
above, default to the provider-native tool and contribute no toolset at all
(`NativeOrLocalTool.get_toolset()` returns `None` for them), so those calls run
server-side and never reach the wrapper. Worth a clause putting native tools
outside the claim.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
