kaxil commented on code in PR #72202:
URL: https://github.com/apache/airflow/pull/72202#discussion_r4128213990
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -532,6 +534,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:
+ # Keep logging last because CombinedCapability applies wrappers in
reverse order,
+ # placing LoggingToolset inside CodeModeToolset where code mode
expects wrapped tools.
+ capabilities.append(ToolLoggingCapability(logger=self.log))
Review Comment:
With this on by default, `agent_params={"tools": [...]}` and capability
tools now go through `LoggingToolset.call_tool`, and its `json.dumps(tool_args,
default=str)` at `toolsets/logging.py:58` runs eagerly, even with DEBUG off,
outside the `try`. pydantic-ai hands `call_tool` the validated arguments, so a
tool taking `dict[uuid.UUID, int]` receives `UUID` keys, and `default=str` only
covers values, not keys. On pydantic-ai-slim 2.33.0 such a tool runs fine
without the capability and fails the whole run with it: `TypeError: keys must
be str, int, float, bool or None, not UUID`. So an `agent_params` tool that
works today breaks once this ships. Falling back to `repr(tool_args)` when
`json.dumps` raises fixes it without changing the log line for ordinary
arguments, and a real-run test with a `dict[UUID, int]` tool would pin it.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/logging.py:
##########
@@ -55,8 +63,27 @@ async def call_tool(
self.logger.info("Tool %s returned in %.2fs", name, elapsed)
self.logger.info("::endgroup::")
return result
+ except (ModelRetry, ApprovalRequired, CallDeferred, SkipToolExecution,
SkipToolValidation) as e:
Review Comment:
This tuple is missing one signal and has one it shouldn't. `ToolFailed`
exists at the 2.33.0 floor and pydantic-ai handles it alongside `ModelRetry`,
as a failed result the model sees and adapts to, but it isn't listed, so a tool
raising it still leaves an ERROR traceback in a task that succeeds.
`SkipToolValidation` goes the other way: raised from a tool body it fails the
run, while this branch logs it at INFO as if it were routine. On 2.33.0 I got
run OK plus ERROR with traceback for `ToolFailed`, and a failed run plus a lone
INFO line for `SkipToolValidation`.
Separately, line 68 logs only the class name, so the reason in
`ModelRetry("column X not found")` no longer appears anywhere in the task log.
Adding `: %s` with `e` keeps it. The "Exceptions are logged and re-raised"
sentence in `docs/toolsets/logging.rst` could also mention the INFO/ERROR split.
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -684,6 +685,85 @@ def test_enable_tool_logging_false_skips_wrapping(self,
mock_hook_cls, make_mock
create_call =
mock_hook_cls.get_hook.return_value.create_agent.call_args
assert create_call[1]["toolsets"] == [mock_toolset]
+ assert "capabilities" not in create_call[1]
+
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_tool_logging_wraps_assembled_capability_toolsets(self,
mock_hook_cls, make_mock_run_result):
Review Comment:
This test is named for the case the PR fixes, but the four capability shapes
never run: the hook is mocked, and lines 721-725 call `get_wrapper_toolset` on
a fresh `FunctionToolset` built in the test, which
`TestToolLoggingCapability.test_wraps_assembled_toolset` already covers. If
pydantic-ai stopped passing capability toolsets into `get_wrapper_toolset`, the
suite would stay green. Parametrizing the real-run
`test_tool_logging_wraps_agent_param_tools` over `agent_params={"capabilities":
[Toolset(FunctionToolset([my_tool]))]}` would cover it, and that case fails
without this change, since capability tools were not wrapped before.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py:
##########
@@ -532,6 +534,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:
+ # Keep logging last because CombinedCapability applies wrappers in
reverse order,
Review Comment:
My round-3 comment had the mechanism wrong, and this comment inherited it,
sorry. `CodeMode` declares `get_ordering()` with `position='outermost'`
(already at the harness 0.3.0 floor), and `CombinedCapability` sorts on that
before applying wrappers, so logging lands inside `CodeModeToolset` whichever
order the two are appended in. A capability declaring that same ordering
produced an identical wrapper chain in both orders on pydantic-ai 2.33.0.
Appending last only matters against capabilities that declare no ordering.
Giving `ToolLoggingCapability` a `get_ordering()` that returns
`CapabilityOrdering(position='innermost')` removes that dependency as well (it
stayed inside an unordered capability appended after it), and this comment
could then describe that instead.
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -1208,6 +1300,30 @@ def factory(ctx):
assert result[0] is cap
+ def test_toolset_capability_logging_wraps_durable_cache(self):
+ """Logging stays outside durable caching for a concrete ``Toolset``
capability."""
+ inner = FunctionToolset()
+ op = AgentOperator(
+ task_id="t",
+ prompt="p",
+ llm_conn_id="c",
+ durable=True,
+ agent_params={"capabilities": [Toolset(inner)]},
+ )
+ op._durable_storage = MagicMock(spec=DurableStorageProtocol)
+ op._durable_counter = DurableStepCounter()
+ hook = MagicMock(spec=["create_agent"])
+ op.llm_hook = hook
+
+ op._build_agent()
+
+ capabilities = hook.create_agent.call_args.kwargs["capabilities"]
+ assert isinstance(capabilities[0].toolset, CachingToolset)
+ assert capabilities[0].toolset.wrapped is inner
+ wrapped = capabilities[1].get_wrapper_toolset(capabilities[0].toolset)
Review Comment:
These last three lines call `get_wrapper_toolset` by hand on the toolset the
test just read back, then assert the result wraps it, which holds for any
argument. So they can't catch the ordering the docstring promises. Lines
1321-1322 are the part doing real work; I'd drop these three or narrow the
docstring to what the test checks.
--
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]
