kaxil commented on code in PR #72202:
URL: https://github.com/apache/airflow/pull/72202#discussion_r3951299261
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -394,6 +397,22 @@ def _build_durable_capabilities(
rewrapped.append(capability)
return rewrapped
+ def _build_logging_capabilities(self, capabilities: list[Any]) ->
list[Any]:
+ """Wrap concrete toolsets provided via a pydantic-ai ``Toolset``
capability for logging."""
+ # Keep capability imports out of Dag-parse-time code paths; they are
only
+ # needed while constructing the runtime agent.
+ from pydantic_ai.capabilities import Toolset
+ from pydantic_ai.toolsets.abstract import AbstractToolset
+
+ rewrapped: list[Any] = []
+ for capability in capabilities:
+ if isinstance(capability, Toolset) and
isinstance(capability.toolset, AbstractToolset):
Review Comment:
The `isinstance(capability, Toolset)` gate catches only one of the shapes a
capability contributes a toolset in. Running four shapes through this logic,
only the first one logs: concrete `Toolset(ts)` yes, `Toolset(lambda ctx: ts)`
no, `PrefixTools(wrapped=Toolset(ts))` no, `MCP(url=...)` no.
`MCP` stands out because agent.rst lists it among the common capabilities
and `MCP.get_toolset()` does hand back a concrete `MCPToolset`, it just extends
`NativeOrLocalTool` rather than `Toolset`. MCP calls are also the ones where a
timing log earns the most.
pydantic-ai has a hook aimed at this:
`AbstractCapability.get_wrapper_toolset(toolset)` receives the combined
non-output toolset per run (dispatched in `Agent._get_toolset`,
`pydantic_ai/agent/__init__.py`) and is documented for "cross-cutting toolset
wrappers ... custom `WrapperToolset` subclasses". A one-method capability
appended when `enable_tool_logging` is on covered all four shapes in a local
probe, factory-backed included, and it is present in pydantic-ai-slim 2.0.0 so
the floor stays put.
Any reason not to go that route instead of matching on capability type?
##########
providers/common/ai/docs/operators/agent.rst:
##########
@@ -374,6 +374,13 @@ 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 a concrete toolset supplied through a
``Toolset``
+capability, just as it does for entries in ``toolsets=``. Other capability
+types and ``Toolset`` capabilities backed by a callable factory are left
+unchanged because they do not expose a concrete toolset when the operator is
Review Comment:
This reason isn't right for all of them. `MCP.get_toolset()` and
`PrefixTools.get_toolset()` both return a concrete toolset at build time; they
get skipped because they aren't `Toolset` instances, not because there's
nothing concrete to wrap. Since MCP is listed as a common capability a few
lines up, someone reading this will assume their MCP tool calls are covered.
--
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]
