kaxil commented on code in PR #74736: URL: https://github.com/apache/airflow/pull/74736#discussion_r4284134899
########## providers/common/ai/src/airflow/providers/common/ai/example_dags/example_progressive_tool_disclosure.py: ########## @@ -0,0 +1,98 @@ +# 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. +"""Example Dags: an agent whose S3 tools stay hidden until the model searches for them.""" + +from __future__ import annotations + +from airflow.providers.common.ai.operators.agent import AgentOperator +from airflow.providers.common.ai.toolsets import HookToolset +from airflow.providers.common.compat.sdk import dag + +try: + from airflow.providers.amazon.aws.hooks.s3 import S3Hook +except ImportError: + S3Hook = None # type: ignore[assignment,misc] + + +if S3Hook is not None: + # [START howto_toolset_defer_loading] + # Nine read-only S3Hook methods, one tool each. + REPORT_METHODS = [ + "list_keys", Review Comment: Fixed in #74759, which merged this morning. `HookToolset` now parses ISO 8601 strings for `datetime` parameters, and a value with no UTC offset goes back to the model as a retry, so `list_keys` stays in the example. ########## providers/common/ai/docs/toolsets/progressive_disclosure.rst: ########## @@ -0,0 +1,178 @@ + .. 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/toolset:progressive-disclosure: + +Progressive tool disclosure +=========================== + +Every tool an agent can call costs input tokens on every model request: its name, +description and argument schema are sent each time, whether the model uses the tool or +not. ``.defer_loading()`` hides a toolset's tools until the model searches for them, so +an agent can carry a wide toolset and pay mostly for the tools a run looks for. It is a +Pydantic AI method that every toolset has, including the ones this provider ships. It +pays off on a toolset with many tools, and does not on one with a handful, such as +``SQLToolset`` (see :ref:`progressive-disclosure-when-not`). A tool no run needs is better +left out of the toolset altogether, with ``allowed_methods`` or ``.filtered()``; defer the +ones only some runs need. + +Hide a toolset's tools until the model needs them +------------------------------------------------- + +This agent can call nine read-only ``S3Hook`` methods on one bucket, and the model sees +none of them until it searches: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_progressive_tool_disclosure.py + :language: python + :dedent: 4 + :start-after: [START howto_toolset_defer_loading] + :end-before: [END howto_toolset_defer_loading] + +How the model finds a hidden tool depends on the model: + +- On models with native tool search, the deferred tools are sent to the provider with + ``defer_loading`` set, and the provider keeps their definitions out of the model's + context and searches them itself. Pydantic AI uses native search for Claude Sonnet 4.5, + Opus 4.5 and Haiku 4.5 and the later Claude families it lists, and for GPT-5.4 and later + through the OpenAI Responses API; its + `tool search guide <https://pydantic.dev/docs/ai/tools-toolsets/tools-advanced/#tool-search>`__ + has the details. +- On every other model, the hidden tools are left out of the request. The model is + offered one ``search_tools`` tool that matches keywords against tool names and + descriptions, and the tools it finds are sent with the next request. + +Pydantic AI picks the path from a list of model name prefixes in the installed version, so +a model family that version does not list gets the local search. For example, Pydantic AI +2.48 lists ``claude-haiku-4-5`` but no Haiku 5 model, so ``claude-haiku-5-5`` takes the +local path. + +Either way, a tool the model has found stays available while the search that found it +is in the run's message history. + +Keep the common tools visible +----------------------------- + +A search costs a model request. When most runs use the same few tools, pass the names of +the tools to hide, and leave the rest in the request from the start: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_progressive_tool_disclosure.py + :language: python + :dedent: 4 + :start-after: [START howto_toolset_defer_loading_some] + :end-before: [END howto_toolset_defer_loading_some] + +The names are the tool names the model sees, so they include ``tool_prefix``. A name that +matches no tool is ignored without an error, so a misspelled name leaves that tool visible. + +What it saves +------------- + +These counts are the tokens in the tool definitions of the first model request, counted +with ``tiktoken``'s ``o200k_base`` encoding on a model without native tool search. Each +model provider counts tokens its own way, so its numbers differ somewhat. + +.. list-table:: + :header-rows: 1 + :widths: 52 16 16 16 + + * - Toolset + - All tools sent + - ``.defer_loading()`` + - ``s3_list_keys`` and ``s3_read_key`` visible + * - The nine ``S3Hook`` methods above + - 1,059 + - 163 + - 694 + * - 25 ``S3Hook`` methods, writes included Review Comment: I dropped the whole table rather than just that row. None of the counts said more than "the saving grows with the number of tools hidden", and they'd drift with every hook docstring or Pydantic AI change. The page says that sentence now, and the measurements stay in the PR description. -- 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]
