kaxil commented on code in PR #72202:
URL: https://github.com/apache/airflow/pull/72202#discussion_r3968170058
##########
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)
Review Comment:
This `except Exception` catches pydantic-ai's control-flow signals, so a run
that succeeds now carries an ERROR-level traceback. `ModelRetry`,
`ApprovalRequired`, `CallDeferred`, `SkipToolExecution` and
`SkipToolValidation` are all plain `Exception` subclasses, and `ModelRetry` is
the documented way for a tool to ask the model to fix its arguments. Loading
this module against pydantic-ai 2.0.0 with a tool that raises `ModelRetry` on
its first call gives `ERROR Tool flaky failed after 0.00s` with a full
traceback, then a clean retry, then `output='done'`.
The handler itself is untouched here, but the diff changes what reaches it:
previously only entries of `toolsets=` were wrapped, so a `ModelRetry` from
`agent_params={"tools": [...]}` or from `capabilities=[Toolset(...)]` never got
here. `agent.py:338` now routes every assembled tool through it with
`enable_tool_logging` defaulting to `True`, which is what turns a narrow wart
into an every-agent-run one. Catching the control-flow signals first and
re-raising after an INFO line would keep the group markers balanced without the
false error.
While you are in this handler: `logger.exception` fires before the
`::endgroup::` on the next line, so a genuine failure's traceback lands
*inside* the fold, and the log viewer renders folds collapsed by default
(`Logs.tsx:86` defaults `expanded` to `false`, and `useLogGroups.tsx:52-53`
then seeds `expandedGroups` empty). Emitting `::endgroup::` before the
`logger.exception` call would leave real tracebacks visible without a click.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -335,6 +335,8 @@ 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:
+ capabilities.append(ToolLoggingCapability(logger=self.log))
Review Comment:
Worth a line of comment that this has to stay appended last.
`CombinedCapability.get_wrapper_toolset` iterates
`reversed(self.capabilities)`, so last-appended ends up the innermost wrapper,
and that is the position that makes `code_mode` work:
`CodeModeToolset.call_tool` delegates native tools to `self.wrapped` and builds
its sandbox `ToolManager` over the wrapped toolset, so inner sandboxed calls
flow through `LoggingToolset`. Move this append above `_build_code_mode()` and
you would log one `run_code` call carrying the whole code blob and lose every
tool call inside it.
##########
providers/common/ai/docs/operators/agent.rst:
##########
@@ -374,6 +374,11 @@ 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,
Review Comment:
Output tools are the other gap in "the complete toolset":
`get_wrapper_toolset` is handed the combined *non-output* toolset, so with
`output_type` set to a `BaseModel` the `final_result` call never reaches the
wrapper. On pydantic-ai 2.0.0 a structured-output agent with no function tools
logs zero `::group::` lines, while `_extract_tool_sequence` walks every
`ToolCallPart` and so the post-run summary still reports `final_result`. Worth
saying function tools here, and in the matching sentence in `toolsets.rst`, so
the two log layers don't look like they disagree by accident.
##########
providers/common/ai/docs/operators/agent.rst:
##########
@@ -476,7 +481,8 @@ Parameters
``BaseModel`` for structured output.
- ``toolsets``: List of pydantic-ai toolsets (``SQLToolset``, ``HookToolset``,
``AgentSkillsToolset`` for :ref:`agent-skills`, etc.).
-- ``enable_tool_logging``: Wrap each toolset in
+- ``enable_tool_logging``: Wrap the toolsets supplied through ``toolsets=`` and
Review Comment:
This bullet still describes per-toolset wrapping of two named sources, so it
now reads narrower than the docstring at `agent.py:141`, which says "the
agent's assembled toolset". One `LoggingToolset` goes around the whole
assembled non-output toolset, so it also covers tools arriving via
`agent_params={"tools": [...]}`, the path
`test_tool_logging_wraps_agent_param_tools` exercises, not just `toolsets=` and
capabilities.
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -242,6 +242,35 @@ 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]
Review Comment:
Now that logging no longer touches the `toolsets=` list, this is the same
assertion `test_execute_passes_toolsets_in_agent_kwargs` makes 20 lines up with
logging left on, so this test would still pass if the `enable_tool_logging`
gate were dropped entirely. The disabled path is only covered incidentally by
`test_code_mode_default_off_no_capabilities`; an `assert "capabilities" not in
create_call[1]` here would pin it where the name says it is.
##########
providers/common/ai/docs/toolsets.rst:
##########
@@ -351,8 +351,14 @@ Parameters
: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`.
Review Comment:
Since `enable_tool_logging` defaults to `True`, a reader who follows this
`:class:` link and adds `ToolLoggingCapability()` to their own `capabilities=`
list gets two of them. Two instances in one list log every call twice with
nested markers, and the second copy's default logger is this module's
`logging.getLogger(__name__)` rather than the task logger, so those lines may
not land in the task log at all. A clause saying `AgentOperator` adds it for
you, and to set `enable_tool_logging=False` before supplying your own, would
head that off.
--
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]
