This is an automated email from the ASF dual-hosted git repository.
kaxilnaik 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 fe9bc5b2fd3 Add toolset design guidance to common AI AGENTS.md (#62870)
fe9bc5b2fd3 is described below
commit fe9bc5b2fd3bfaff1cba365b3b7f9c764795d30c
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Mar 4 13:27:14 2026 +0000
Add toolset design guidance to common AI AGENTS.md (#62870)
The existing AGENTS.md covers hooks and operators but doesn't mention
toolsets. As we add more toolset backends (DataFusion, etc.) alongside
the existing SQLToolset, having clear guidance upfront will help
contributors follow consistent patterns.
Adds:
- "One backend per toolset" design principle
- "Adding a New Toolset" section with the four-tool pattern and
constructor hygiene guidelines
- Toolsets path in Key Paths
---
providers/common/ai/AGENTS.md | 15 +++++++++++++++
1 file changed, 15 insertions(+)
diff --git a/providers/common/ai/AGENTS.md b/providers/common/ai/AGENTS.md
index 2c068d70891..8e8711c8ab7 100644
--- a/providers/common/ai/AGENTS.md
+++ b/providers/common/ai/AGENTS.md
@@ -19,6 +19,10 @@ The hook is a thin bridge between Airflow connections and
pydantic-ai's model/pr
code path. Wait until there are 3+ concrete use cases before introducing an
abstraction.
- **Operators stay focused.** Each operator does one thing: `LLMOperator`
(prompt → output),
`LLMBranchOperator` (prompt → branch decision), `LLMSQLOperator` (prompt →
validated SQL).
+- **One backend per toolset.** A toolset wraps a single execution backend
(e.g. `DbApiHook`,
+ `DataFusionEngine`). If a new backend needs the same tool interface, create
a new toolset class —
+ don't add mutually exclusive constructor params and branch in every method.
Users compose toolsets
+ via `AgentOperator(toolsets=[...])`.
## Adding Support for a New LLM Provider
@@ -34,6 +38,16 @@ If pydantic-ai already supports the provider (check [models
docs](https://ai.pyd
If pydantic-ai does *not* support the provider, contribute upstream to
pydantic-ai rather than
building a wrapper here.
+## Adding a New Toolset
+
+1. If the toolset provides SQL-like access (tables, schemas, queries), follow
the four-tool pattern
+ from `SQLToolset`: `list_tables`, `get_schema`, `query`, `check_query`.
+2. Extract shared logic (result truncation, SQL validation, `allowed_tables`
filtering) to helpers
+ in the toolsets package rather than inheritance.
+3. Keep constructors focused — only accept params that apply to the backend.
Don't add params that
+ are silently ignored in certain modes.
+4. Mark tools `sequential=True` when the backend uses synchronous I/O.
+
## Security
- **No dynamic imports from connection extras.** Never use
`importlib.import_module()` on
@@ -56,5 +70,6 @@ building a wrapper here.
- Hook: `src/airflow/providers/common/ai/hooks/pydantic_ai.py`
- Operators: `src/airflow/providers/common/ai/operators/`
- Decorators: `src/airflow/providers/common/ai/decorators/`
+- Toolsets: `src/airflow/providers/common/ai/toolsets/`
- Tests: `tests/unit/common/ai/`
- Docs: `docs/`
