This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 72d612c357d Add restricted-agent examples to common.ai toolset guides
(#74379)
72d612c357d is described below
commit 72d612c357dfeb73108713cba6f2afc4d07a18d6
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Oct 7 06:18:17 2026 +0100
Add restricted-agent examples to common.ai toolset guides (#74379)
---
providers/common/ai/docs/toolsets/datafusion.rst | 50 +++++++++++
providers/common/ai/docs/toolsets/hook.rst | 57 ++++++++++++-
.../common/ai/docs/toolsets/object_storage.rst | 38 +++++++++
providers/common/ai/docs/toolsets/skills.rst | 38 +++++++++
providers/common/ai/docs/toolsets/sql.rst | 97 ++++++++++++++++++++++
.../common/ai/example_dags/example_agent_skills.py | 29 +++++++
.../ai/example_dags/example_datafusion_toolset.py | 72 ++++++++++++++++
.../common/ai/example_dags/example_hook_toolset.py | 57 +++++++++++++
.../example_dags/example_object_storage_toolset.py | 35 ++++++++
.../common/ai/example_dags/example_sql_toolset.py | 65 +++++++++++++++
10 files changed, 536 insertions(+), 2 deletions(-)
diff --git a/providers/common/ai/docs/toolsets/datafusion.rst
b/providers/common/ai/docs/toolsets/datafusion.rst
index 5c494caae74..a1aa337b38c 100644
--- a/providers/common/ai/docs/toolsets/datafusion.rst
+++ b/providers/common/ai/docs/toolsets/datafusion.rst
@@ -75,6 +75,56 @@ The ``DataFusionEngine`` is created lazily on the first tool
call. This
toolset requires the ``datafusion`` extra of
``apache-airflow-providers-common-sql``.
+.. _datafusion-toolset-restricted:
+
+Restricting the agent
+---------------------
+
+With ``allow_writes=False`` (the default), the tables you register are the
only ones
+the agent can query; there is no separate allow-list. This agent can query two
tables,
+cannot write, and gets bounded results:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_datafusion_toolset.py
+ :language: python
+ :start-after: [START howto_toolset_datafusion_restricted]
+ :end-before: [END howto_toolset_datafusion_restricted]
+
+Run against an S3 endpoint where an ``acme-payroll`` bucket sits beside the
reports,
+these queries were refused, and the model got the error back to correct:
+
+``SELECT * FROM payroll``
+ ``error: Error while executing query: DataFusion error:
Diagnostic(Diagnostic {
+ kind: Error, message: "table 'payroll' not found", ...``
+
+``SELECT * FROM 's3://acme-payroll/salaries.csv'``
+ ``error: Error while executing query: DataFusion error:
Diagnostic(Diagnostic {
+ kind: Error, message: "table 's3://acme-payroll/salaries.csv' not found",
...``
+
+``CREATE TABLE copy AS SELECT * FROM sales``
+ ``error: Statement type 'Create' is not allowed. Allowed types: Select,
Union,
+ Intersect, Except. Only read-only SELECT-family queries are allowed unless
+ allow_writes is enabled; check the SQL syntax and statement type, then try
+ again.``
+
+The second query shows that a URL in the SQL is looked up as a table name, not
read
+as a file. A query over the registered tables runs:
+
+.. code-block:: sql
+
+ SELECT s.region, CAST(sum(r.amount) AS DOUBLE) / sum(s.amount) AS
return_rate
+ FROM sales s JOIN returns r ON r.region = s.region
+ GROUP BY s.region ORDER BY return_rate DESC
+
+It returns:
+
+.. code-block:: json
+
+
{"columns":["region","return_rate"],"rows":[["AMER",0.5],["EMEA",0.2]],"row_count":2}
+
+The object store is created for the whole bucket, not the prefix each
+``DataSourceConfig`` registers, so give ``aws_reports_reader`` credentials
that can
+read only those prefixes.
+
Parameters
----------
diff --git a/providers/common/ai/docs/toolsets/hook.rst
b/providers/common/ai/docs/toolsets/hook.rst
index 15a6ad0d2b8..84d93871000 100644
--- a/providers/common/ai/docs/toolsets/hook.rst
+++ b/providers/common/ai/docs/toolsets/hook.rst
@@ -107,8 +107,11 @@ the Dag author's decision, such as which bucket a storage
hook reads, pin them:
)
A pinned argument is left out of the schema the model sees and passed to every
allowed
-method. If the model supplies it anyway, the call is refused and the model is
told the
-argument is fixed.
+method. If the model supplies it anyway, the call is refused. For a method
with named
+parameters, such as ``S3Hook.read_key``, argument validation refuses it as an
extra
+input (see :ref:`hook-toolset-restricted`). A method that names the parameter
and also
+takes ``**kwargs`` would accept it, so the toolset refuses it itself and tells
the
+model the argument is fixed.
A pin binds one parameter name, so every allowed method has to take it by that
name.
When one does not, the toolset raises ``ValueError`` when it is created: a
method that
@@ -123,6 +126,56 @@ Pinned values are passed as written: they are not rendered
as templates, and the
part of what ``AgentOperator(durable=True)`` fingerprints, so change one only
between Dag
runs, not between the tries of one.
+.. _hook-toolset-restricted:
+
+Restricting the agent
+---------------------
+
+``allowed_methods`` decides which hook methods become tools, and
``pinned_arguments``
+decides which of their arguments the model cannot set. This agent can list and
read
+one bucket through ``S3Hook``, and nothing else:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_hook_toolset.py
+ :language: python
+ :start-after: [START howto_toolset_hook_restricted]
+ :end-before: [END howto_toolset_hook_restricted]
+
+The model is offered two tools: ``s3_read_key``, which takes only ``key``, and
+``s3_list_keys``, whose parameters include ``prefix`` but not ``bucket_name``.
Run
+against an S3 endpoint that also holds an ``acme-payroll`` bucket, a call that
names
+that bucket was refused before it reached S3, with this message to the model:
+
+.. code-block:: text
+
+ 1 validation error:
+ ```json
+ [
+ {
+ "type": "extra_forbidden",
+ "loc": [
+ "bucket_name"
+ ],
+ "msg": "Extra inputs are not permitted",
+ "input": "acme-payroll"
+ }
+ ]
+ ```
+
+ Fix the errors and try again.
+
+A call to a method that is not listed, such as ``s3_delete_objects``, got
+``Unknown tool name: 's3_delete_objects'. Available tools: 's3_list_keys',
+'s3_read_key'``. The model can correct both kinds of call and carry on.
+
+An exception from the hook is different: it fails the run, and the task with
it.
+Reading a key that does not exist ended the run with ``ClientError: An error
occurred
+(404) when calling the HeadObject operation: Not Found``.
+
+The pin fixes the bucket and leaves every key in it to the model. Give
+``aws_reports_reader`` credentials that can read only that bucket, so the
connection
+holds the same limit if a method you expose later reaches another bucket some
other
+way.
+
Parameters
----------
diff --git a/providers/common/ai/docs/toolsets/object_storage.rst
b/providers/common/ai/docs/toolsets/object_storage.rst
index 67984d575e7..ed3d6577724 100644
--- a/providers/common/ai/docs/toolsets/object_storage.rst
+++ b/providers/common/ai/docs/toolsets/object_storage.rst
@@ -72,6 +72,44 @@ the task: an image or PDF, a binary file, a file larger than
``max_read_bytes``
by default), a corrupt file, a path that does not exist, or one the connection
may not
read.
+.. _object-storage-toolset-restricted:
+
+Restricting the agent
+---------------------
+
+This agent can read only under ``s3://acme-reports/finance/``, and every list
and read
+is bounded:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_object_storage_toolset.py
+ :language: python
+ :start-after: [START howto_toolset_object_storage_restricted]
+ :end-before: [END howto_toolset_object_storage_restricted]
+
+Run against an S3 endpoint where an ``acme-payroll`` bucket sits beside the
reports,
+these reads were refused. Each refusal comes back to the model as the tool's
result,
+so it does not use ``max_retries``, and the run carries on:
+
+``../../acme-payroll/salaries.csv``
+ ``'../../acme-payroll/salaries.csv' leaves the storage root; '..' is not
allowed.``
+
+``s3://acme-payroll/salaries.csv``
+ ``'s3://acme-payroll/salaries.csv' is not a relative path. Name files
relative to
+ the storage root.``
+
+``/etc/passwd``
+ ``'/etc/passwd' is not a relative path. Name files relative to the storage
root.``
+
+``2026-09/chart.png``
+ ``'2026-09/chart.png' is a png file, which this tool cannot read as text.``
+
+``2026-09/ledger.csv``, a 1.5 MB file
+ ``'2026-09/ledger.csv' is larger than the 1.0MB this tool reads.``
+
+Because refused reads cost nothing from ``max_retries``, bound a run that keeps
+asking with the operator's ``usage_limits``. The path check is the toolset's
only
+boundary on location; credentials that can read only
``s3://acme-reports/finance/`` keep that limit
+if the check has a gap.
+
Parameters
----------
diff --git a/providers/common/ai/docs/toolsets/skills.rst
b/providers/common/ai/docs/toolsets/skills.rst
index ed32a2fb2c7..bb1cdcc6bf4 100644
--- a/providers/common/ai/docs/toolsets/skills.rst
+++ b/providers/common/ai/docs/toolsets/skills.rst
@@ -82,6 +82,44 @@ need strict isolation.
``branch`` to a trusted ref, and treat skill contents as code that runs in
your environment.
+.. _agent-skills-restricted:
+
+Restricting the agent
+---------------------
+
+This agent can read skills but cannot run any script a skill ships, and files
that
+match the ``exclude_resources`` patterns stay out of its reach. The example
skills ship
+no scripts, so the exclusion matters once a skill adds one:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
+ :language: python
+ :start-after: [START howto_operator_agent_skills_restricted]
+ :end-before: [END howto_operator_agent_skills_restricted]
+
+The model is offered ``list_skills``, ``load_skill`` and
``read_skill_resource``. A
+call to the excluded tool got ``Unknown tool name: 'run_skill_script'.
Available
+tools: 'list_skills', 'load_skill', 'read_skill_resource'``.
+
+The example skills contain no file the patterns match, so the next run used a
copy of
+them with ``warehouse.env``, ``secrets/token.txt`` and ``reference.md`` added
to
+``sql-reporting``. ``load_skill`` listed only ``reference.md`` as a resource,
and
+reading an excluded file got:
+
+.. code-block:: text
+
+ Resource 'warehouse.env' not found in skill 'sql-reporting'. Available
resources: ['reference.md']. Use the exact name from load_skill output.
+
+The skills tools allow the model one correction each, and the agent's
``retries``
+does not change that. A second refused read in a row failed the run with
+``UnexpectedModelBehavior``, even with ``retries`` set to ``{"tools": 3}``:
+
+.. code-block:: text
+
+ Tool 'read_skill_resource' exceeded max retries count of 1. Consider
raising the retry limit, or see the docs on tool retries:
https://pydantic.dev/docs/ai/tools-toolsets/tools-advanced/#tool-retries
+
+``exclude_resources`` hides files from the resource tools only. A skill's
scripts can
+still read them, which is why the example excludes ``run_skill_script`` as
well.
+
Parameters
----------
diff --git a/providers/common/ai/docs/toolsets/sql.rst
b/providers/common/ai/docs/toolsets/sql.rst
index e59a0e1a88a..6d0a85b6d71 100644
--- a/providers/common/ai/docs/toolsets/sql.rst
+++ b/providers/common/ai/docs/toolsets/sql.rst
@@ -65,6 +65,103 @@ rejects by scanning the parsed statement for write
operations. When
table, so its target must be on the list, while ``SHOW`` enumerates objects
beyond
any single table and is rejected outright (see
:ref:`allowed-tables-enforcement`).
+.. _sql-toolset-restricted:
+
+Restricting the agent
+---------------------
+
+Every limit can be set on one task: the toolset holds most of them, and the
+operator holds the tool-call limit. This agent can query two tables, cannot
write, gets
+bounded results, and has a budget for both refused and successful calls:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_sql_toolset.py
+ :language: python
+ :start-after: [START howto_toolset_sql_restricted]
+ :end-before: [END howto_toolset_sql_restricted]
+
+A refused query never reaches the database. The model gets the reason back as
an
+error it can correct, and the run carries on. Run against a Postgres warehouse
that
+also holds a ``secrets`` table, these queries were refused with these messages:
+
+``SELECT token FROM secrets``
+ ``The query tool failed: Query references tables that are not in the
allowed
+ tables list: secrets. Use list_tables to see the allowed tables.``
+
+``SELECT pg_read_file('/etc/passwd')``
+ ``The query tool failed: Query uses a data source that cannot be checked
against
+ allowed_tables: function(s) the parser cannot verify against allowed_tables
+ (pg_read_file); if these functions are trusted, permit them via
+ allowed_functions. Query the allowed tables directly: use list_tables to
see
+ them.``
+
+``DELETE FROM orders``
+ ``The query tool failed: Statement type 'Delete' is not allowed. Allowed
types:
+ Select, Union, Intersect, Except, Describe, Show``
+
+Each message ends with the same two lines:
+
+.. code-block:: text
+
+ Use the list_tables and get_schema tools to inspect the database, then fix
the query and try again.
+
+ Fix the errors and try again.
+
+This query runs, because the example lists ``json_build_object`` in
+``allowed_functions``:
+
+.. code-block:: sql
+
+ SELECT c.region, json_build_object('revenue', sum(o.amount)) AS revenue
+ FROM orders o JOIN customers c ON c.id = o.customer_id
+ GROUP BY c.region ORDER BY c.region
+
+It returns:
+
+.. code-block:: json
+
+
{"columns":["region","revenue"],"rows":[["AMER",{"revenue":50}],["EMEA",{"revenue":200}]],"row_count":2}
+
+Without ``allowed_functions``, the same query is refused with the same message
as
+``pg_read_file``, naming ``json_build_object``.
+
+The two budgets count different calls, and running out of either fails the
task.
+``max_retries`` counts refused and failed calls, and a fourth refused
``query`` call
+in a row ends the run with ``UnexpectedModelBehavior``:
+
+.. code-block:: text
+
+ Tool 'query' exceeded max retries count of 3. Consider raising the retry
limit, or see the docs on tool retries:
https://pydantic.dev/docs/ai/tools-toolsets/tools-advanced/#tool-retries
+
+``tool_calls_limit`` counts successful calls only, and fails the run with
+``UsageLimitExceeded`` before any call that would take the count past 20:
+
+.. code-block:: text
+
+ The next tool call(s) would exceed the tool_calls_limit of 20
(tool_calls=21). Consider raising the limit, or see the docs on usage limits
for budget-aware patterns:
https://pydantic.dev/docs/ai/core-concepts/agent/#usage-limits
+
+The task's own ``retries`` then decide whether it runs again, and the budgets
differ
+there too. ``max_retries`` starts again on each attempt. On Airflow 3.3 and
later,
+``usage_limits`` counts across every attempt, so a retry after
``tool_calls_limit``
+ran out fails at its first tool call; see
+:ref:`the usage budget <agent-usage-budget>`. ``UsageLimits`` also keeps
+pydantic-ai's default ``request_limit`` of 50 model requests, which refused
calls
+use up too.
+
+``allowed_tables`` works by parsing the SQL, so a query the parser reads
differently
+from the database, or a function listed in ``allowed_functions``, can get past
it.
+The connection's role is the limit that holds regardless. The example's
+``warehouse_agent_reader`` connection logs in as a role created with:
+
+.. code-block:: sql
+
+ CREATE ROLE warehouse_agent_reader LOGIN PASSWORD '...';
+ GRANT SELECT ON orders, customers TO warehouse_agent_reader;
+
+With that role and no ``allowed_tables``, ``SELECT token FROM secrets`` reaches
+Postgres, which refuses it, and the model gets ``The query tool failed:
permission
+denied for table secrets``. :ref:`allowed-tables-enforcement` lists what the
parser
+checks and where it stops.
+
Multi-schema warehouses
-------------------------
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
index d7523e608a0..f922b40993d 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py
@@ -109,3 +109,32 @@ def example_agent_skills_git():
# [END howto_operator_agent_skills_git]
example_agent_skills_git()
+
+
+# ---------------------------------------------------------------------------
+# 3. Skills the model can read but not run
+# ---------------------------------------------------------------------------
+
+
+# [START howto_operator_agent_skills_restricted]
+@dag(tags=["example"])
+def example_agent_skills_restricted():
+ AgentOperator(
+ task_id="report_writer",
+ prompt="Review this query against our SQL reporting conventions:
SELECT * FROM orders",
+ llm_conn_id="pydanticai_default",
+ toolsets=[
+ AgentSkillsToolset(
+ sources=[str(SKILLS_DIR)],
+ # Hide the tool that runs a skill's scripts on the worker.
+ exclude_tools={"run_skill_script"},
+ # Keep matching files out of the resources the model can list
and read.
+ exclude_resources=["*.env", "secrets/*"],
+ )
+ ],
+ )
+
+
+# [END howto_operator_agent_skills_restricted]
+
+example_agent_skills_restricted()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_datafusion_toolset.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_datafusion_toolset.py
new file mode 100644
index 00000000000..9e8670f108a
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_datafusion_toolset.py
@@ -0,0 +1,72 @@
+# 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 Dag: an agent that runs read-only SQL over two registered tables of
files on S3."""
+
+from __future__ import annotations
+
+from airflow.providers.common.ai.operators.agent import AgentOperator
+from airflow.providers.common.compat.sdk import dag
+
+try:
+ from airflow.providers.common.ai.toolsets.datafusion import
DataFusionToolset
+ from airflow.providers.common.sql.config import DataSourceConfig
+except Exception:
+ DataFusionToolset = None # type: ignore[assignment,misc]
+
+
+if DataFusionToolset is not None:
+ # [START howto_toolset_datafusion_restricted]
+ @dag(tags=["example"])
+ def example_datafusion_toolset_restricted():
+ AgentOperator(
+ task_id="returns_by_region",
+ prompt="Which region had the highest return rate in September?",
+ llm_conn_id="pydanticai_default",
+ toolsets=[
+ DataFusionToolset(
+ # Register only what the agent may query: there is no
other allow-list.
+ datasource_configs=[
+ DataSourceConfig(
+ conn_id="aws_reports_reader",
+ table_name="sales",
+ uri="s3://acme-reports/finance/sales/",
+ format="parquet",
+ ),
+ DataSourceConfig(
+ conn_id="aws_reports_reader",
+ table_name="returns",
+ uri="s3://acme-reports/finance/returns/",
+ format="csv",
+ ),
+ ],
+ # The default: allow only SELECT-family statements.
+ allow_writes=False,
+ # Return at most 100 rows and 32 KiB from one query.
+ max_rows=100,
+ max_result_bytes=32 * 1024,
+ # Summarize a table wider than 50 columns instead of
listing every column.
+ max_columns=50,
+ # Let the model correct a refused query, or one naming an
unknown table or
+ # column, up to 3 times.
+ max_retries=3,
+ )
+ ],
+ )
+
+ # [END howto_toolset_datafusion_restricted]
+
+ example_datafusion_toolset_restricted()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_hook_toolset.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_hook_toolset.py
new file mode 100644
index 00000000000..88dc99aa108
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_hook_toolset.py
@@ -0,0 +1,57 @@
+# 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 Dag: an agent that can list and read one S3 bucket through
``S3Hook``, and nothing else."""
+
+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_hook_restricted]
+ @dag(tags=["example"])
+ def example_hook_toolset_restricted():
+ AgentOperator(
+ task_id="read_reports",
+ prompt="Summarize the September finance report.",
+ llm_conn_id="pydanticai_default",
+ toolsets=[
+ HookToolset(
+ # Log in with credentials scoped to the reports bucket.
+ S3Hook(aws_conn_id="aws_reports_reader"),
+ # Offer these two methods as tools; no other S3Hook method
is reachable.
+ allowed_methods=["list_keys", "read_key"],
+ # Fix the bucket: the model never sees the argument and
cannot set it.
+ pinned_arguments={"bucket_name": "acme-reports"},
+ # Name the tools s3_list_keys and s3_read_key.
+ tool_name_prefix="s3_",
+ # Let the model correct an invalid call up to 2 times.
+ max_retries=2,
+ )
+ ],
+ )
+
+ # [END howto_toolset_hook_restricted]
+
+ example_hook_toolset_restricted()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_object_storage_toolset.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_object_storage_toolset.py
index 3cd86eb9fb5..74a5000e784 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_object_storage_toolset.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_object_storage_toolset.py
@@ -18,6 +18,7 @@
from __future__ import annotations
+from airflow.providers.common.ai.operators.agent import AgentOperator
from airflow.providers.common.ai.toolsets import ObjectStorageToolset
from airflow.providers.common.compat.sdk import dag, task
@@ -46,3 +47,37 @@ def example_object_storage_toolset():
example_object_storage_toolset()
+
+
+# [START howto_toolset_object_storage_restricted]
+@dag(tags=["example"])
+def example_object_storage_toolset_restricted():
+ AgentOperator(
+ task_id="read_finance_reports",
+ prompt="Summarize the September finance report.",
+ llm_conn_id="pydanticai_default",
+ toolsets=[
+ ObjectStorageToolset(
+ # Resolve every path the model names under this root. The
toolset never writes.
+ "s3://acme-reports/finance/",
+ # Credentials scoped to that prefix keep the limit if the path
check has a gap.
+ conn_id="aws_reports_reader",
+ # List at most 50 entries per call; the model pages through
the rest.
+ max_files=50,
+ # Refuse a file larger than 1 MiB, measured after
decompression.
+ max_read_bytes=1024 * 1024,
+ # Return at most 16 KiB from one read; the model pages through
a longer text file.
+ max_output_bytes=16 * 1024,
+ # Name the tools reports_list_files, reports_get_file_info and
reports_read_file.
+ tool_prefix="reports",
+ # Let the model correct an invalid call up to 2 times.
+ max_retries=2,
+ )
+ ],
+ )
+
+
+# [END howto_toolset_object_storage_restricted]
+
+
+example_object_storage_toolset_restricted()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_sql_toolset.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_sql_toolset.py
new file mode 100644
index 00000000000..33d4cc33b71
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_sql_toolset.py
@@ -0,0 +1,65 @@
+# 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 Dag: a SQL agent held to two tables, read-only, with bounded
results and a call budget."""
+
+from __future__ import annotations
+
+from pydantic_ai.usage import UsageLimits
+
+from airflow.providers.common.ai.operators.agent import AgentOperator
+from airflow.providers.common.compat.sdk import dag
+
+try:
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+except Exception:
+ SQLToolset = None # type: ignore[assignment,misc]
+
+
+if SQLToolset is not None:
+ # [START howto_toolset_sql_restricted]
+ @dag(tags=["example"])
+ def example_sql_toolset_restricted():
+ AgentOperator(
+ task_id="revenue_by_region",
+ prompt="What was last week's revenue by region?",
+ llm_conn_id="pydanticai_default",
+ toolsets=[
+ SQLToolset(
+ # Log in as a role granted SELECT on orders and customers
only.
+ db_conn_id="warehouse_agent_reader",
+ # Refuse any query that reaches another table or a source
the parser cannot check.
+ allowed_tables=["orders", "customers"],
+ # Accept this function, which the SQL parser does not
recognize.
+ allowed_functions=["json_build_object"],
+ # The default: allow only SELECT-family, DESCRIBE and SHOW
statements.
+ allow_writes=False,
+ # Return at most 100 rows and 32 KiB from one query.
+ max_rows=100,
+ max_result_bytes=32 * 1024,
+ # Summarize a table wider than 50 columns instead of
listing every column.
+ max_columns=50,
+ # Let the model correct a refused or failed call up to 3
times.
+ max_retries=3,
+ )
+ ],
+ # Fail the run before a 21st successful tool call, counted across
retries on Airflow 3.3+.
+ usage_limits=UsageLimits(tool_calls_limit=20),
+ )
+
+ # [END howto_toolset_sql_restricted]
+
+ example_sql_toolset_restricted()
