kaxil commented on code in PR #72202:
URL: https://github.com/apache/airflow/pull/72202#discussion_r4190186683
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -292,9 +292,12 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin):
``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 or provider-native tools that run
+ server-side. Set to ``False`` to disable.
Review Comment:
`toolsets/logging.rst` now lists a third exclusion that this docstring
leaves out: tools another capability adds through its own wrapper. That
includes `run_code` from this operator's own `code_mode=True`, so a reader of
this parameter would expect a `Tool call: run_code` group and only see the
inner calls.
```suggestion
: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.
```
##########
providers/common/ai/tests/unit/common/ai/toolsets/test_logging.py:
##########
@@ -127,3 +206,20 @@ async def test_empty_args_not_logged(self,
logging_toolset, wrapped_toolset, cap
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_is_innermost(self, logger):
Review Comment:
This is the only test that covers the `innermost` ordering, and it only
reads back the value `get_ordering()` returns. With `position` changed to
`"outermost"`, the rest of the suite still passes.
A behavioural test would pin it: run a real `Agent(FunctionModel)` with
`tools=[my_tool]` and a test-local capability whose `get_wrapper_toolset`
returns `PrefixedToolset(toolset, prefix="outer")`, listed before
`ToolLoggingCapability`, and have the model call `outer_my_tool`. With
`innermost` the log shows `::group::Tool call: my_tool`; with `outermost` it
shows `::group::Tool call: outer_my_tool`.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -771,6 +772,10 @@ def _build_agent(self) -> Agent[object, Any]:
capabilities.append(_build_code_mode())
if self.cache_prompt:
capabilities.append(PromptCaching())
+ if self.enable_tool_logging:
+ # ToolLoggingCapability's innermost ordering keeps logging inside
capability wrappers,
+ # including CodeModeToolset where code mode expects the wrapped
tools.
+ capabilities.append(ToolLoggingCapability(logger=self.log))
Review Comment:
A user who passes their own `ToolLoggingCapability` in `capabilities=` (for
a custom logger) and leaves `enable_tool_logging` at its default gets every
tool call logged twice, with one `::group::` fold nested inside the other. The
docs warn about it, but the operator could rule it out the way it already
handles code mode. `_contains_code_mode` walks the declared capabilities,
including wrapper and combined ones; taking the capability class as a parameter
would let this append be skipped when one is already declared, so the user's
logger wins. Raising, as the code-mode check does, would also work.
pydantic-ai's id-based capability dedup isn't a substitute: it isn't in the
2.33.0 floor, and on newer versions the later instance wins, which would
replace the user's logger with `self.log`.
##########
providers/common/ai/docs/operators/agent.rst:
##########
@@ -372,9 +372,11 @@ 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 and provider-native tools
+ are not covered. Default ``True``.
Review Comment:
Same gap as the `enable_tool_logging` docstring: wrapper-added tools such as
`run_code` are missing here.
```suggestion
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``.
```
--
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]
