kaxil commented on code in PR #72202:
URL: https://github.com/apache/airflow/pull/72202#discussion_r4139394650
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/logging.py:
##########
@@ -58,11 +69,34 @@ async def call_tool(
return result
except ApprovalRequired:
# Not a failure: the run pauses here until a person approves or
rejects the call.
- self.logger.info("Tool %s is waiting to be approved", name)
+ elapsed = time.monotonic() - start
+ self.logger.info("Tool %s is waiting to be approved after %.2fs",
name, elapsed)
+ self.logger.info("::endgroup::")
+ raise
+ except (ModelRetry, ToolFailed, CallDeferred, SkipToolExecution) as e:
+ elapsed = time.monotonic() - start
+ self.logger.info("Tool %s requested %s after %.2fs: %s", name,
type(e).__name__, elapsed, e)
self.logger.info("::endgroup::")
raise
except Exception:
elapsed = time.monotonic() - start
- self.logger.exception("Tool %s failed after %.2fs", name, elapsed)
self.logger.info("::endgroup::")
Review Comment:
Moving `::endgroup::` first only keeps a failure outside the fold when tool
calls run one at a time. pydantic-ai runs parallel tool calls as concurrent
tasks and cancels the siblings when one raises, and `CancelledError` is not an
`Exception`, so no branch closes their group. On 2.33.0 a model calling `boom`
(raises) and `slow` in one turn logs `::group::boom`, `::group::slow`,
`::endgroup::`, then the ERROR, so the failure sits inside `slow`'s fold, which
never closes. An `except asyncio.CancelledError:` that logs `::endgroup::` and
re-raises would close it. Fine as a follow-up, since `toolsets=` tools behaved
this way before the PR too.
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -693,6 +695,75 @@ 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]
+
+ @pytest.mark.parametrize("tool_source", ["tools", "capabilities"])
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_tool_logging_wraps_agent_param_tools(self, mock_hook_cls, caplog,
tool_source):
+ """Tool logging wraps the complete toolset assembled by pydantic-ai."""
+
+ def my_tool() -> str:
+ return "tool-result"
+
+ def model_fn(messages, info):
+ saw_return = any(isinstance(p, ToolReturnPart) for m in messages
for p in getattr(m, "parts", []))
+ if saw_return:
+ return ModelResponse(parts=[TextPart(content="done")])
+ return ModelResponse(parts=[ToolCallPart(tool_name="my_tool",
args={}, tool_call_id="c1")])
+
+ mock_hook_cls.get_hook.return_value.create_agent.side_effect = lambda
**kw: Agent(
+ FunctionModel(model_fn), **kw
+ )
+ configured_tools = [my_tool] if tool_source == "tools" else
[Toolset(FunctionToolset([my_tool]))]
+ op = AgentOperator(
+ task_id="test",
+ prompt="Do something",
+ llm_conn_id="my_llm",
+ agent_params={tool_source: configured_tools},
+ )
+
+ with caplog.at_level(logging.INFO):
+ result = op.execute(context=_make_context())
+
+ assert result == "done"
+ assert any(record.message == "::group::Tool call: my_tool" for record
in caplog.records)
+
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_tool_logging_handles_non_json_mapping_keys(self, mock_hook_cls,
caplog):
+ def my_tool(values: dict[UUID, int]) -> int:
+ return sum(values.values())
+
+ def model_fn(messages, info):
+ saw_return = any(isinstance(p, ToolReturnPart) for m in messages
for p in getattr(m, "parts", []))
+ if saw_return:
+ return ModelResponse(parts=[TextPart(content="done")])
+ return ModelResponse(
+ parts=[
+ ToolCallPart(
+ tool_name="my_tool",
+ args={"values":
{"12345678-1234-5678-1234-567812345678": 1}},
+ tool_call_id="c1",
+ )
+ ]
+ )
+
+ mock_hook_cls.get_hook.return_value.create_agent.side_effect = lambda
**kw: Agent(
+ FunctionModel(model_fn), **kw
+ )
+ op = AgentOperator(
+ task_id="test",
+ prompt="Do something",
+ llm_conn_id="my_llm",
+ agent_params={"tools": [my_tool]},
+ )
+
+ with caplog.at_level(logging.DEBUG):
Review Comment:
This test fails in the Compat 3.0.6 job. On 3.0.x `op.log` is a stdlib
`logging.getLogger("airflow.task.operators...")`, and the default logging
config gives `airflow.task` an explicit INFO level, so
`caplog.at_level(logging.DEBUG)` only lowers the root logger.
`isEnabledFor(logging.DEBUG)` stays False and the `Tool args` line is never
emitted; the captured log in that job shows the run succeeding with no args
line. Pinning the level on that logger with `caplog.at_level(logging.DEBUG,
logger="airflow.task")` should fix it on 3.0.x.
##########
providers/common/ai/docs/toolsets/logging.rst:
##########
@@ -21,8 +21,21 @@ Tool call logging: ``LoggingToolset``
: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`.
+Applying the wrapper as a capability means logging covers the assembled
+function toolset, including tools supplied through ``toolsets=``,
+``agent_params={"tools": [...]}``, and capabilities such as factory-backed
+toolsets, nested capabilities, and MCP toolsets. Output tools such as
+``final_result`` and provider-native tools, including native MCP, are not
+covered by Airflow's real-time tool-call logging.
Review Comment:
One more category belongs in this list: tools another capability adds
through its own wrapper toolset. Logging is innermost, so it never sees them.
With `ToolSearch(strategy="keywords")` the `search_tools` call produced no log
records on 2.33.0, though the post-run summary lists it, and CodeMode's
`run_code` is in the same position by design. A sentence here keeps the
"complete limitations" pointer in guardrails.rst accurate.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/logging.py:
##########
@@ -58,11 +69,34 @@ async def call_tool(
return result
except ApprovalRequired:
# Not a failure: the run pauses here until a person approves or
rejects the call.
- self.logger.info("Tool %s is waiting to be approved", name)
+ elapsed = time.monotonic() - start
+ self.logger.info("Tool %s is waiting to be approved after %.2fs",
name, elapsed)
+ self.logger.info("::endgroup::")
+ raise
+ except (ModelRetry, ToolFailed, CallDeferred, SkipToolExecution) as e:
+ elapsed = time.monotonic() - start
+ self.logger.info("Tool %s requested %s after %.2fs: %s", name,
type(e).__name__, elapsed, e)
self.logger.info("::endgroup::")
raise
except Exception:
elapsed = time.monotonic() - start
- self.logger.exception("Tool %s failed after %.2fs", name, elapsed)
self.logger.info("::endgroup::")
+ self.logger.exception("Tool %s failed after %.2fs", name, elapsed)
raise
+
+
+@dataclass
+class ToolLoggingCapability(AbstractCapability[Any]):
+ """Apply tool-call logging to the assembled function toolset for an agent
run."""
+
+ logger: Logger | logging.Logger = field(default_factory=lambda:
logging.getLogger(__name__))
+
+ @classmethod
+ def get_serialization_name(cls) -> str | None:
+ return None
+
+ def get_ordering(self) -> CapabilityOrdering:
Review Comment:
The docs spellcheck job is red again on this push, this time on
`CombinedCapability`. With no docstring of its own, `get_ordering` inherits
`AbstractCapability.get_ordering`'s, which links `CombinedCapability`, and
autoapi renders that into `_api/.../toolsets/logging/index.rst:85`. A one-line
docstring here (a good spot to say why it has to be innermost) or a wordlist
entry clears it.
--
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]
