kaxil commented on code in PR #74736: URL: https://github.com/apache/airflow/pull/74736#discussion_r4283234886
########## 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: I'd keep it after `[END]`. That's where 18 of the sibling example files put the call (4 put it inside), and the snippet shows the Dag definition a reader copies. ########## 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: Done, the list is inlined so the snippet stands on its own. ########## 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: I took the intro change: it now ends with "Use it on toolsets with many tools", plus the line about leaving out tools no run needs, and the cross-reference is gone. I kept "When it does not help", though. The skills paragraph is the one readers wouldn't guess, since `.defer_loading()` leaves the skill catalog in the instructions untouched. The SQLToolset numbers stop people deferring a four-tool toolset. ########## 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: Done, the model list and the 2.48 / Haiku 5 example are gone. One version-free sentence stays: Pydantic AI matches model-name prefixes, so a model its list doesn't match gets the local `search_tools`. -- 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]
