kaxil commented on code in PR #74695:
URL: https://github.com/apache/airflow/pull/74695#discussion_r4268554807
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py:
##########
@@ -84,12 +84,12 @@ def __init__(
self,
mcp_conn_id: str,
*,
- tool_prefix: str | None = None,
+ tool_prefix: str = "",
token_provider: Callable[[], str] | None = None,
env_provider: Callable[[], dict[str, str]] | None = None,
) -> None:
self._mcp_conn_id = mcp_conn_id
- self._tool_prefix = tool_prefix
+ self._tool_prefix = tool_prefix or ""
Review Comment:
`tool_prefix` is annotated `str` now, but the docstring, `mcp.rst` and the
changelog still promise `None` is accepted, and the new tests need `# type:
ignore[arg-type]` to exercise that. Since the point here is to settle this
signature, should the `None` promise go along with the `or ""` here and in
`MCPHook`? `hooks/mcp.py:254` only checks truthiness, so a stray `None` still
works without being part of the contract.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -429,10 +444,11 @@ async def execute_tool(
ctx: RunContext[Any],
tool: ToolsetTool[Any],
) -> Any:
- if name not in ("list_tables", "get_schema", "query", "check_query"):
+ base = name.removeprefix(f"{self._tool_prefix}_") if self._tool_prefix
else name
+ if base not in ("list_tables", "get_schema", "query", "check_query"):
raise ValueError(f"Unknown tool: {name!r}")
try:
- return await self.run_blocking(self._run_tool, name, tool_args)
+ return await self.run_blocking(self._run_tool, base, tool_args)
Review Comment:
The retry text just below still says "Use the list_tables and get_schema
tools", and the same bare names appear in the `allowed_tables` refusals (lines
612 and 622), the DataFusion retry at `datafusion.py:296`, and the shared
`get_schema` description and hints in `utils/query_results.py` (96, 261, 312).
With `tool_prefix="pg"` no tool has those names, and in the setup the docstring
recommends a prefix for (a second `SQLToolset`), a bare `list_tables` can
belong to the other database. Building these strings from
`self._tool_name(...)`, and passing the name into the `query_results` helpers,
would keep the hints pointing at this toolset's own tools.
##########
providers/common/ai/docs/toolsets/sql.rst:
##########
@@ -329,6 +329,9 @@ Parameters
See :ref:`bounded-query-results` and :ref:`bounded-schema-results`.
- ``max_columns``: Maximum columns ``get_schema`` returns in full. Default
``100``.
Above it the result becomes a bounded summary. See
:ref:`bounded-schema-results`.
+- ``tool_prefix``: Optional prefix joined to each tool name with an underscore
(e.g. ``"pg"`` produces ``"pg_query"``).
Review Comment:
With a prefix set, the `allow_writes` approval recipe a few lines up stops
gating writes. `.approval_required()` passes the user's function the toolset's
own `ToolDefinition`, so with `tool_prefix="pg"` it sees `pg_query`,
`tool_def.name == "query"` is False, and the write runs without asking anyone.
Could this entry (or the recipe itself) say the function sees the prefixed
name, like `mcp.rst` already does for its tool filter?
--
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]
