jeff3071 commented on code in PR #72202:
URL: https://github.com/apache/airflow/pull/72202#discussion_r3957220223
##########
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:
Thanks for reviewing! I confirmed that `Agent._get_toolset()` combines the
non-output toolsets before calling `get_wrapper_toolset()`.
When tool logging is enabled, we can pass the configured logging capability
to the agent and remove the eager `wrap_toolsets_for_logging()` call. I’ll
also verify this with an actual Dag.
--
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]
