This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch commonai-progressive-tool-disclosure in repository https://gitbox.apache.org/repos/asf/airflow.git
commit b38c98ecb1e7daef8abdf29227e04bdf87335c1e Author: Kaxil Naik <[email protected]> AuthorDate: Wed Oct 14 13:52:33 2026 +0100 Document progressive tool disclosure for Common AI toolsets --- providers/common/ai/docs/features.rst | 4 + providers/common/ai/docs/operators/agent.rst | 7 +- providers/common/ai/docs/stability.rst | 5 +- providers/common/ai/docs/toolsets/index.rst | 6 +- providers/common/ai/docs/toolsets/logging.rst | 4 +- .../ai/docs/toolsets/progressive_disclosure.rst | 178 +++++++++++++++++++++ providers/common/ai/docs/toolsets/skills.rst | 10 +- providers/common/ai/pyproject.toml | 5 +- .../example_progressive_tool_disclosure.py | 98 ++++++++++++ 9 files changed, 302 insertions(+), 15 deletions(-) diff --git a/providers/common/ai/docs/features.rst b/providers/common/ai/docs/features.rst index 2526b79ae3c..737d9c8c2ef 100644 --- a/providers/common/ai/docs/features.rst +++ b/providers/common/ai/docs/features.rst @@ -135,6 +135,10 @@ checklist. * - :doc:`ManagedAgentToolset <toolsets/managed_agent>` - Sends a prompt to an agent on Bedrock AgentCore or Vertex AI Agent Engine. - The Amazon or Google provider's ``common.ai`` extra + * - :doc:`Progressive tool disclosure <toolsets/progressive_disclosure>` + - Hides a toolset's tools from the model until it searches for them, with + ``.defer_loading()`` on any toolset. + - None Control output and spend ------------------------ diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index 7e9b41d463f..fbb7b542989 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -378,9 +378,10 @@ Parameters - ``enable_tool_logging``: Wrap the assembled function toolset in :class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` so that tools supplied through ``toolsets=``, ``agent_params={"tools": [...]}``, and - capabilities are logged in real time. Output tools, provider-native tools, and - tools another capability adds through its own wrapper (code mode's ``run_code``, - ToolSearch's ``search_tools``) are not covered. Default ``True``. + capabilities are logged in real time, including tools the model found through + :ref:`tool search <howto/toolset:progressive-disclosure>`. Output tools, provider-native + tools, and tools another capability adds through its own wrapper (code mode's + ``run_code``, ToolSearch's ``search_tools``) are not covered. Default ``True``. - ``agent_params``: Additional keyword arguments passed to the pydantic-ai ``Agent`` constructor (e.g. ``retries``, ``model_settings``). - .. _agent-usage-budget: diff --git a/providers/common/ai/docs/stability.rst b/providers/common/ai/docs/stability.rst index 5b2e2a61f63..c9733255cd5 100644 --- a/providers/common/ai/docs/stability.rst +++ b/providers/common/ai/docs/stability.rst @@ -223,8 +223,9 @@ Everything this provider ships that is not in the table above is experimental. - Thin integrations of ``pydantic-ai-harness`` and ``pydantic-ai-shields``, which are before their 1.0 releases. * - The Agent Skills toolset (:doc:`toolsets/skills`) - - It may move from ``pydantic-ai-skills`` to the on-demand loading Pydantic AI now has - itself (`pydantic/pydantic-ai#5230 <https://github.com/pydantic/pydantic-ai/pull/5230>`__). + - It is built on ``pydantic-ai-skills``, a community package. Whether it moves to + ``pydantic-ai-harness``'s ``Skills`` once that loads skill resources and scripts is + still open. * - :class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset` and its backends (:doc:`sandbox/index`) - ``durable=True`` cannot be used with it yet, the ``sbx`` backend leaves its sandbox diff --git a/providers/common/ai/docs/toolsets/index.rst b/providers/common/ai/docs/toolsets/index.rst index c69bfffddfb..faa2085fa6d 100644 --- a/providers/common/ai/docs/toolsets/index.rst +++ b/providers/common/ai/docs/toolsets/index.rst @@ -67,13 +67,14 @@ with is resolved in the worker and is not part of any of that. - ``timeout`` for each call; what the remote agent may touch is set at the vendor -Two controls work on every toolset. ``.approval_required()`` pauses the task before a +Three controls work on every toolset. ``.approval_required()`` pauses the task before a matching call runs, until a person approves or rejects it on the **Required Actions** page. It needs Airflow 3.3 or later, pauses at every step that calls a matching tool on Airflow 3.4 and later (once per Dag run on 3.3), and does not combine with ``durable=True`` or a ``SandboxToolset`` that provisions its own sandbox (:doc:`../tool_approval` lists every limit). ``.filtered()`` drops tools from what the model is offered, as the :ref:`MCP guide <howto/toolset:mcp-filtered>` -shows. +shows. ``.defer_loading()`` hides a toolset's tools from the model until it searches for +them, which pays off on a toolset with many tools (see :doc:`progressive_disclosure`). :doc:`logging` and :doc:`langchain` do not reach a system of their own: one wraps a toolset to log its calls, the other converts a toolset for a LangChain agent. @@ -92,6 +93,7 @@ toolset to log its calls, the other converts a toolset for a LangChain agent. Vendor-managed agents <managed_agent> LangChain tools <langchain> Tool call logging <logging> + Progressive tool disclosure <progressive_disclosure> Importing toolsets ------------------ diff --git a/providers/common/ai/docs/toolsets/logging.rst b/providers/common/ai/docs/toolsets/logging.rst index 6cba50b98b5..a83d3b0984f 100644 --- a/providers/common/ai/docs/toolsets/logging.rst +++ b/providers/common/ai/docs/toolsets/logging.rst @@ -35,7 +35,9 @@ 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. Tools added by another capability through its own wrapper toolset, such as ToolSearch's -``search_tools`` and CodeMode's ``run_code``, are also not covered. +``search_tools`` and CodeMode's ``run_code``, are also not covered. A tool hidden with +``.defer_loading()`` is logged once the model finds it, and its result is masked like +any other (see :doc:`progressive_disclosure`). ``AgentOperator`` adds ``ToolLoggingCapability`` automatically when ``enable_tool_logging=True``. If ``capabilities=`` already contains one, diff --git a/providers/common/ai/docs/toolsets/progressive_disclosure.rst b/providers/common/ai/docs/toolsets/progressive_disclosure.rst new file mode 100644 index 00000000000..24b1a62fcff --- /dev/null +++ b/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 + - 4,609 + - 163 + - 761 + +The 163 tokens are ``search_tools`` itself, so the saving grows with the number of tools +hidden. On a model with native tool search there is no ``search_tools``, and the +provider keeps the deferred tools out of the model's context until they are found. + +What a run looks like +--------------------- + +The first example, run on ``claude-haiku-5-5`` (which took the local search), searched +once, then called ``s3_list_keys`` once and ``s3_read_key`` twice. This is the +``AgentOperator`` part of its task log, with timestamps and logger names trimmed: + +.. code-block:: text + + ::group::Tool call: s3_list_keys + Tool s3_list_keys returned in 0.20s + ::endgroup:: + ::group::Tool call: s3_read_key + Tool s3_read_key returned in 0.08s + ::endgroup:: + ::group::Tool call: s3_read_key + Tool s3_read_key returned in 0.01s + ::endgroup:: + ::group::LLM run complete: model=claude-haiku-5-5, requests=4, tool_calls=4, input_tokens=6683, output_tokens=474, total_tokens=7157 + LLM prompt cache: cache_read_tokens=3894, cache_write_tokens=2779 + Tool call sequence: search_tools -> s3_list_keys -> s3_read_key -> s3_read_key + ::endgroup:: + +``search_tools`` has no ``Tool call`` group of its own, because tool call logging does not +cover tools another capability adds. It shows only in the sequence at the end of the run, +and it is one of the four tool calls counted. + +The second example, with ``s3_list_keys`` and ``s3_read_key`` visible from the start, made +the same three S3 calls without a search, in one model request fewer: + +.. code-block:: text + + ::group::LLM run complete: model=claude-haiku-5-5, requests=3, tool_calls=3, input_tokens=4633, output_tokens=394, total_tokens=5027 + LLM prompt cache: cache_read_tokens=2880, cache_write_tokens=1745 + Tool call sequence: s3_list_keys -> s3_read_key -> s3_read_key + ::endgroup:: + +The local search matches each keyword on its own, so one search can reveal more tools than +the model goes on to use. In a run where the model searched for ``list keys read key``, +the search revealed six of the nine S3 tools, and every request after it carried all six. + +What stays the same +------------------- + +A tool the model found is the same tool, called the same way: + +- Its result passes through Airflow's secret masker before the model sees it, as every + tool result does (see :doc:`../agent_security`). +- Its calls are logged like any other tool's (see :doc:`logging`). +- With ``durable=True``, a retry replays the model requests and tool calls that came + after a search; only the step that failed runs again (see :doc:`../durable_execution`). +- A templated connection ID on the toolset is rendered for each task instance, as it is + without ``.defer_loading()``. + +.. _progressive-disclosure-when-not: + +When it does not help +--------------------- + +``SQLToolset`` has four tools whatever the database holds: 488 tokens of tool +definitions, against 163 for ``search_tools`` in their place. Hiding them saves 325 +tokens a request and costs a search request in nearly every run, since the model needs +the tools to answer anything. On a warehouse with hundreds of tables, ``list_tables`` +output grows with the table count; limit it with ``allowed_tables`` (see +:ref:`sql-toolset-restricted`). + +``AgentSkillsToolset`` has four tools, and its list of skills goes into the agent's +instructions, which ``.defer_loading()`` does not touch. For 30 skills that list was +2,129 tokens against 907 for the tools. The toolset already loads each skill's body only +when the model asks for it. diff --git a/providers/common/ai/docs/toolsets/skills.rst b/providers/common/ai/docs/toolsets/skills.rst index 52398c8a064..ea7b33c3245 100644 --- a/providers/common/ai/docs/toolsets/skills.rst +++ b/providers/common/ai/docs/toolsets/skills.rst @@ -33,11 +33,11 @@ prompt until the model decides it needs one, so a large skill library costs few tokens until used (progressive disclosure). It is backed by the community `pydantic-ai-skills -<https://github.com/DougTrajano/pydantic-ai-skills>`__ package (MIT). Pydantic AI now -has on-demand loading of its own (`pydantic/pydantic-ai#5230 -<https://github.com/pydantic/pydantic-ai/pull/5230>`__), and this toolset may move to it. -Install the optional -extra to use it: +<https://github.com/DougTrajano/pydantic-ai-skills>`__ package (MIT). ``pydantic-ai-harness``'s +`Skills <https://github.com/pydantic/pydantic-ai-harness/tree/main/pydantic_ai_harness/skills>`__ +capability also loads Agent Skills, as on-demand capabilities, but only their ``SKILL.md`` +instructions: no resources and no ``run_skill_script``. +Install the optional extra to use this toolset: .. code-block:: bash diff --git a/providers/common/ai/pyproject.toml b/providers/common/ai/pyproject.toml index 9aebf776494..71a3e1ac8e3 100644 --- a/providers/common/ai/pyproject.toml +++ b/providers/common/ai/pyproject.toml @@ -132,8 +132,9 @@ dependencies = [ "shields" = ["pydantic-ai-shields>=0.3.4"] # Agent Skills (agentskills.io) support. pydantic-ai-skills provides the toolset; # the git provider supplies GitHook + GitPython for cloning GitSkills with -# credentials from a `git` connection. Native progressive disclosure is tracked -# upstream in pydantic/pydantic-ai#5230; revisit this extra once that lands. +# credentials from a `git` connection. pydantic-ai-harness's Skills capability is not +# a replacement: it loads only SKILL.md instructions, with no skill resources and no +# run_skill_script. "skills" = [ "apache-airflow-providers-git>=0.4.0", "pydantic-ai-skills>=1.2.0", diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_progressive_tool_disclosure.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_progressive_tool_disclosure.py new file mode 100644 index 00000000000..12085b470a5 --- /dev/null +++ b/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", + "list_prefixes", + "check_for_key", + "check_for_prefix", + "check_for_wildcard_key", + "check_for_bucket", + "head_object", + "read_key", + "get_bucket_tagging", + ] + + @dag(tags=["example"]) + def example_defer_loading(): + AgentOperator( + task_id="answer_from_reports", + prompt="Summarize the reports under 2026/09/.", + llm_conn_id="pydanticai_default", + toolsets=[ + # Hide every tool until the model searches for it. + HookToolset( + S3Hook(aws_conn_id="aws_reports_reader"), + allowed_methods=REPORT_METHODS, + pinned_arguments={"bucket_name": "acme-reports"}, + tool_prefix="s3", + ).defer_loading() + ], + ) + + # [END howto_toolset_defer_loading] + + example_defer_loading() + + # [START howto_toolset_defer_loading_some] + @dag(tags=["example"]) + def example_defer_loading_some(): + AgentOperator( + task_id="answer_from_reports", + prompt="Summarize the reports under 2026/09/.", + llm_conn_id="pydanticai_default", + toolsets=[ + HookToolset( + S3Hook(aws_conn_id="aws_reports_reader"), + # The nine methods of the previous example. + allowed_methods=REPORT_METHODS, + pinned_arguments={"bucket_name": "acme-reports"}, + tool_prefix="s3", + ).defer_loading( + # Hide these seven; s3_list_keys and s3_read_key, which most runs use, stay visible. + [ + "s3_list_prefixes", + "s3_check_for_key", + "s3_check_for_prefix", + "s3_check_for_wildcard_key", + "s3_check_for_bucket", + "s3_head_object", + "s3_get_bucket_tagging", + ] + ) + ], + ) + + # [END howto_toolset_defer_loading_some] + + example_defer_loading_some()
