pankajkoti commented on code in PR #74736: URL: https://github.com/apache/airflow/pull/74736#discussion_r4283189859
########## 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: `list_keys` takes `from_datetime` and `to_datetime`, and `HookToolset` does not convert arguments for `datetime` parameters yet. If the model passes an ISO string, `S3Hook` raises a `TypeError` and the task fails, which has nothing to do with `.defer_loading()`. Could we either leave `list_keys` out of this example, or add a short note on the page that these two parameters are not supported yet? Minor, since the existing `HookToolset` example has the same gap. ########## 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: The "25 `S3Hook` methods" row can't be reproduced because the page doesn't say which 25 were counted. Could we list them, or drop this row and keep the nine-method one? Minor. -- 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]
