Lee-W commented on code in PR #72936: URL: https://github.com/apache/airflow/pull/72936#discussion_r4052562223
########## providers/common/ai/docs/choosing_a_toolset.rst: ########## @@ -0,0 +1,574 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _howto/choosing-a-toolset: + +Choosing a Toolset +================== + +:doc:`toolsets` documents how to configure each toolset. This page answers the +question that comes before that one: you have a system you want an agent to +reach, so which route do you take, and what does each route give up? + +Read the table below by what you already have, not by what a toolset is called. +When two routes both work, the deciding factor is rarely what each one can do — +it is what each one cannot do, and every route has a short list. + +More than one row can be true at once, and the rows are not exclusive: one agent +can carry several toolsets. Two questions break the ties. *Whose credential is +it?* — prefer the route whose credential is an Airflow connection somebody on +your side already reviewed. *Whose tool list is it?* — prefer the route whose +exposed surface you chose rather than inherited. The pair that most often +overlaps is an Airflow hook and a vendor MCP server reaching the same target; +both questions point at the hook, because its credential is the connection and +``allowed_methods`` is a list you write. Reach for the server when its tools +cover work the hook does not expose, or when the alternative is re-wrapping that +API by hand. + +Those two questions do not separate ``HookToolset`` from ``SQLToolset`` when the +target is a DBAPI database, because both answer them the same way. A third one +does: *is the work a fixed operation or an open-ended question?* A named method +you can enumerate in advance is a hook. A question the agent has to express as +SQL is a query, and ``SQLToolset`` answers it with schema discovery, bounded +results and an ``allowed_tables`` walk you can switch on, none of which +``HookToolset`` has an equivalent of. + +Start with what you have +------------------------ + +.. list-table:: + :widths: 50 50 + :header-rows: 1 + + * - What you have + - Route + * - A target that already has an Airflow connection, and a hook method that + already does the thing + - ``HookToolset`` + * - A question that is a query, against a DBAPI database + - ``SQLToolset`` + * - Files on an object store — Parquet, CSV, Avro — or a catalog-managed + table format such as Iceberg, rather than rows in a database + - ``DataFusionToolset`` + * - A vendor that already ships a server built for agents, whose tools you + would otherwise re-wrap by hand + - ``MCPToolset`` + * - Procedural knowledge — how to carry out a task — rather than an endpoint + to call + - ``AgentSkillsToolset`` + * - Work that means running code the model wrote, not calling a tool you chose + - ``SandboxToolset`` + * - Reasoning that should happen on the vendor's own infrastructure + - A subclass of ``BaseManagedAgentToolset`` that you write + +The rest of this page takes those seven in turn. Each entry gives the case for +choosing it, what it cannot do, an example that exists in this repository, and +where its credentials and its work come from. + +``HookToolset`` +--------------- + +**Choose it when** the target already has an Airflow connection and a hook, and +what you want the agent to do is already a method on that hook. This is the +cheapest route — no new server, no new credential, no new query dialect — and +the only one that reaches any provider hook with synchronous methods without +anyone writing an adapter first. :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` is a +reflection-based adapter, so the work is choosing the method list. + +**What it cannot do** + +- It allow-lists method *names*, not arguments. Once ``read_key`` is exposed, + the agent picks the key; the :ref:`defense-layer table <toolset-defense-layers>` + states this outright. Choose methods whose worst case you accept, not methods + you intend to constrain later. +- Its calls act as barriers. The tools are registered with ``sequential=True`` + because hook methods perform synchronous I/O, so a slow call holds up every + other tool the model emitted in that step, not only this toolset's. This is + not specific to ``HookToolset`` — see :ref:`toolset-call-barriers`. +- It returns exactly one shape. Every result goes through ``serialize_for_llm`` + and comes back as a JSON-encoded string; there is no structured error type and + no ``ModelRetry`` wrapper, so a hook exception fails the agent run, and the + task with it, instead of giving the model something it can correct. + ``SQLToolset``, by contrast, hands the database's own error back as a retry. +- Its ``call_tool`` calls the method and serializes what comes back. The code + contains no path that awaits a coroutine result, and none that checks for one, + so an ``async def`` hook method is not a case this adapter is written to + handle. Treat synchronous methods as the supported set. + +**A real example.** The read-only S3 pair from the ``HookToolset`` guidance in +:doc:`toolsets`: + +.. code-block:: python + + HookToolset( + s3_hook, + allowed_methods=["list_keys", "read_key"], + tool_name_prefix="s3_", + ) + +**Credentials and where it runs.** The hook instance is yours, so the credential +is whatever connection that hook resolves — the toolset never looks one up +itself. Calls run in the Airflow worker process. + +``SQLToolset`` +-------------- + +**Choose it when** the question is a query and the data is in a DBAPI database. +:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` gives the agent +four tools — list tables, get schema, query, check query. Set ``allowed_tables`` +and that allow-list is enforced by parsing the SQL rather than by matching +strings; see :ref:`allowed-tables-enforcement` for how the walk handles CTEs, +subqueries and joins. + +**What it cannot do** + +- ``allowed_tables`` is an application-level guardrail, not a replacement for + database permissions. Its own docstring says so, and names the residual gap: + an engine or query the parser reads differently. Point ``db_conn_id`` at a + least-privilege role whose grants match the allow-list. +- It cannot bound the fetch for every driver. Hooks that hand their handler + something other than a DBAPI cursor — ``ExasolHook`` and its pyexasol + statement, for instance — fall back to a full fetch. The payload handed to the + model is still bounded; the transfer is not. See :ref:`bounded-query-results`. +- Its parser-level closure is opt-in, not the default. ``allowed_tables`` + defaults to ``None`` and the table walk returns immediately while it is + unset, so out of the box the agent reaches every table the connection can + see. ``DESCRIBE`` and ``SHOW`` both pass, on dialects that parse them, + while ``allowed_tables`` stays unset. Set ``allowed_tables`` and the walk + turns fail-closed for ``SHOW`` — but not for ``DESCRIBE``: it instead + becomes an ordinary table reference, allowed only when the table it names + is on the list. See :ref:`allowed-tables-enforcement` for what else the + walk rejects once ``allowed_tables`` is active. Statements that modify + data, ``COPY`` among them, are rejected either way while ``allow_writes`` + is ``False``. +- It does not classify failures. A connection error or a typo in a column + name reaching ``list_tables``, ``get_schema`` or ``query`` becomes one + ``ModelRetry``, so the two are treated the same way until the retry budget + runs out and the task fails for Airflow to retry. Two paths do not raise: + ``check_query`` catches its own errors and reports them back as a normal + ``{"valid": false, ...}`` result, and ``get_schema`` returns a normal + ``{"error": ...}`` result instead of raising when the requested table is + outside ``allowed_tables`` — other ``get_schema`` failures still raise and + still become a ``ModelRetry``. + +**A real example.** ``example_pydantic_ai_hook.py`` builds an agent around +``SQLToolset`` inside a plain ``@task`` function, with no operator involved: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py + :language: python + :start-after: [START howto_task_with_toolsets] + :end-before: [END howto_task_with_toolsets] + +**Credentials and where it runs.** ``db_conn_id`` is resolved through +``BaseHook.get_connection``, and the connection must supply a ``DbApiHook``. +Queries run from the Airflow worker process against the database. Its tool +calls act as barriers, as they do for the other routes that build their own +tools; see :ref:`toolset-call-barriers`. + +``DataFusionToolset`` +--------------------- + +**Choose it when** the data is files on an object store rather than rows in a +database — Parquet, CSV or Avro — or a table in a catalog such as Iceberg, and +you want the agent to ask SQL questions of them without loading them anywhere +first. Each ``DataSourceConfig`` registers one table, and several can be +registered so the agent can join across them. The two shapes take different +fields: an object-store format needs a ``uri``, while a catalog format like +Iceberg is looked up by ``db_name`` instead, and ``DataSourceConfig`` raises +``ValueError`` at construction if a catalog format is missing one. + +**What it cannot do** + +- It has no table allow-list. ``allow_writes=False`` is the only guard, and it + blocks non-SELECT statements, not reach: the defense-layer table records that + this toolset "does not prevent the agent from reading any registered data + source". The registration list is therefore the whole boundary — register + exactly what the agent may read. +- It bounds the payload, not the scan. DataFusion has already materialized the + full result before ``max_rows`` and ``max_result_bytes`` apply, so those limits + protect the model's context, not the cost of the query. +- It cannot tell failure kinds apart precisely. The DataFusion Python bindings + expose no native exception types, so the retry decision is made by matching the + error message against regular expressions — which a wording change upstream can + quietly defeat. + +**A real example.** The same bucket as the ``HookToolset`` entry above, reached +the other way. Rather than exposing ``list_keys`` and ``read_key`` and leaving +the agent to reassemble files, this registers the prefix as a table and lets it +write SQL: + +.. code-block:: python + + from airflow.providers.common.ai.toolsets.datafusion import DataFusionToolset + from airflow.providers.common.sql.config import DataSourceConfig + + toolset = DataFusionToolset( + datasource_configs=[ + DataSourceConfig( + conn_id="aws_default", + table_name="sales", + uri="s3://my-bucket/data/sales/", + format="parquet", + ), + ], + max_rows=100, + ) + +Which of the two fits depends on the question. "Read me this object" is a hook +method. "What were last quarter's returns by region" is a query, and expressing +it through ``list_keys`` and ``read_key`` means the model does the aggregation in +its context window instead of the engine doing it. + +An Iceberg table is registered differently. There is no ``uri`` to read files Review Comment: Both extras are named now. (for datafusion and iceberg) -- 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]
