pankajastro commented on code in PR #74736: URL: https://github.com/apache/airflow/pull/74736#discussion_r4282235162
########## 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] Review Comment: Should `example_defer_loading()` move above `# [END howto_toolset_defer_loading]`, so the rendered snippet registers the Dag? Same for `example_defer_loading_some()` at the end of the file. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## 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, Review Comment: `REPORT_METHODS` isn't defined in this snippet. Could the list be inlined here? --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## 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 Review Comment: Could this intro end with "Use it on toolsets with many tools." instead of the "does not on one with a handful... left out" caveats? Could "When it does not help" also become "Where it fits", cut to that one sentence? --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## 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, Review Comment: Could the model list and the Pydantic AI 2.48 / Haiku 5 example be replaced with "Models with native tool search use it; others get a local `search_tools`", keeping the upstream link? The list will go stale when Pydantic AI updates its prefixes. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## 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>` Review Comment: Could this row come out of the toolsets table? `.defer_loading()` isn't a toolset, and `toolsets/index.rst` already links to the page. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting -- 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]
