apache · gopidesupavan · May 23, 2026 · May 24, 2026 · May 25, 2026 · May 25, 2026
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
@@ -25,6 +25,7 @@ adsinsights
 afterall
 agentcore
 agentic
+AgentT
 ai
 aio
 aiobotocore
@@ -677,6 +678,7 @@ fsGroup
 fsspec
 fullname
 func
+functools
 ga
 Gantt
 gantt
@@ -831,6 +833,7 @@ init
 initdb
 initialisation
 initialiser
+initialises
 initialising
 initializer
 inout

diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json
@@ -555,7 +555,7 @@
       "apache-airflow-providers-common-compat>=1.14.1",
       "apache-airflow-providers-standard>=1.12.1",
       "apache-airflow>=3.0.0",
-      "pydantic-ai-slim>=1.71.0"
+      "pydantic-ai-slim>=1.96.0"
     ],
     "devel-deps": [
       "langchain>=1.0.0",

diff --git a/generated/provider_dependencies.json.sha256sum b/generated/provider_dependencies.json.sha256sum
@@ -1 +1 @@
-86e39c620f3926c99e1c702a496d6161032e1a3ac69eba7da10214a2c4ba24f1
+507d80ebe713a42b294363cf05e50b01c0727a7e384056e4b8b23d5102bd11b5
@@ -12,11 +12,13 @@ The hook is a thin bridge between Airflow connections and pydantic-ai's model/pr
   Bedrock, Ollama, etc.) via `infer_model()` and provider classes like `AzureProvider`, `BedrockProvider`.
   Do not re-implement provider-specific logic that pydantic-ai handles.
   Before writing new code, check: https://ai.pydantic.dev/models/
-- **Keep the hook thin.** `PydanticAIHook.get_conn()` maps Airflow connection fields to pydantic-ai
-  constructors. That is the hook's entire job. Do not add abstraction layers (builders, factories,
-  registries, Protocols) on top of pydantic-ai's own abstractions.
-- **No premature abstraction.** Do not add Protocols, builder patterns, or plugin systems for a single
-  code path. Wait until there are 3+ concrete use cases before introducing an abstraction.
+- **Keep LLM hooks thin.** `PydanticAIHook.get_conn()` maps Airflow connection fields to pydantic-ai
+  constructors. That is the hook's entire job for one-shot LLM operators.
+- **Agent backends use `BaseAIHook`.** `AgentOperator` / `@task.agent` resolve
+  `BaseAIHook.get_agent_hook(conn_id)` so the connection ``conn_type`` selects the runtime
+  (``pydanticai``, ``pydanticai-bedrock``, ``pydanticai-azure``, …). New agent frameworks subclass
+  `BaseAIHook` and implement `get_model`, `create_agent`, `run_agent`, and `_tool_spec_to_native`;
+  do not add parallel operator classes per framework.
 - **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`,
@@ -67,7 +69,8 @@ building a wrapper here.
 
 ## Key Paths
 
-- Hook: `src/airflow/providers/common/ai/hooks/pydantic_ai.py`
+- Hooks: `src/airflow/providers/common/ai/hooks/pydantic_ai.py` (pydantic-ai)
+- Base hook contract: `src/airflow/providers/common/ai/hooks/base.py`
 - Operators: `src/airflow/providers/common/ai/operators/`
 - Decorators: `src/airflow/providers/common/ai/decorators/`
 - Toolsets: `src/airflow/providers/common/ai/toolsets/`

@@ -56,7 +56,7 @@ PIP package                                 Version required
 ``apache-airflow``                          ``>=3.0.0``
 ``apache-airflow-providers-common-compat``  ``>=1.14.1``
 ``apache-airflow-providers-standard``       ``>=1.12.1``
-``pydantic-ai-slim``                        ``>=1.71.0``
+``pydantic-ai-slim``                        ``>=1.96.0``
 ==========================================  ==================
 
 Cross provider package dependencies

@@ -47,6 +47,20 @@ name added to ``[core] allowed_deserialization_classes`` -- the consumer DAG's
 worker only loads its own DAG. On Airflow versions whose worker does not register
 declared classes, the operators dump to ``dict`` instead.
 
+Direct callers of :meth:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook.create_agent`
+and :meth:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook.run_agent` must use
+:class:`~airflow.providers.common.ai.hooks.base.AgentRunRequest` instead of keyword arguments.
+DAG authors using :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`,
+``@task.agent``, and the other LLM operators are unaffected.
+
+``SQLToolset`` now implements the framework-neutral
+:class:`~airflow.providers.common.ai.hooks.base.BaseToolset` interface instead of
+pydantic-ai's ``AbstractToolset`` interface. DAG authors using ``SQLToolset``
+with ``AgentOperator`` or ``@task.agent`` are unaffected. Direct pydantic-ai
+``Agent(toolsets=[SQLToolset(...)])`` callers should use
+``AgentOperator(toolsets=[SQLToolset(...)])`` or pass the SQL tool callables
+through an Airflow agent hook request.
+
 0.4.0
 .....
 

@@ -34,8 +34,9 @@ Choosing a hook
    * - Hook
      - When to use
    * - :class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook`
-     - Default for ``common.ai`` operators (``LLMOperator``, ``AgentOperator``,
-       ``LLMBranchOperator``, ...). Returns a pydantic-ai ``Agent`` / ``Model``.
+     - Default for one-shot LLM operators (``LLMOperator``, ``LLMBranchOperator``, ...).
+       Also used by ``AgentOperator`` when ``conn_type`` is ``pydanticai`` (or
+       ``pydanticai-bedrock``, ``pydanticai-azure``, ``pydanticai-vertex``).
    * - :class:`~airflow.providers.common.ai.hooks.langchain.LangChainHook`
      - Direct LangChain access for tasks that compose ``Runnable``\\s, use the
        LangChain agent surface, or need LangChain-native chat / embedding model

@@ -109,7 +109,7 @@ PIP package                                 Version required
 ``apache-airflow``                          ``>=3.0.0``
 ``apache-airflow-providers-common-compat``  ``>=1.14.1``
 ``apache-airflow-providers-standard``       ``>=1.12.1``
-``pydantic-ai-slim``                        ``>=1.71.0``
+``pydantic-ai-slim``                        ``>=1.96.0``
 ==========================================  ==================
 
 Cross provider package dependencies

@@ -31,7 +31,10 @@ a single prompt and returns the output. ``AgentOperator`` manages a stateful
 tool-call loop where the LLM decides which tools to call and when to stop.
 
 .. seealso::
-    :ref:`Connection configuration <howto/connection:pydanticai>`
+    :ref:`Pydantic AI connection <howto/connection:pydanticai>`
+
+The agent backend is selected by the Airflow connection ``conn_type`` (for example
+``pydanticai``, ``pydanticai-bedrock``, or ``pydanticai-azure``). You do not choose a different operator class.
 
 
 SQL Agent
@@ -304,11 +307,21 @@ Parameters
   templating.
 - ``output_type``: Expected output type (default: ``str``). Set to a Pydantic
   ``BaseModel`` for structured output.
-- ``toolsets``: List of pydantic-ai toolsets (``SQLToolset``, ``HookToolset``,
-  ``AgentSkillsToolset`` for :ref:`agent-skills`, etc.).
-- ``enable_tool_logging``: Wrap each toolset in
-  :class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` so that
-  every tool call is logged in real time. Default ``True``.
+- ``toolsets``: List of toolsets the agent can use. Accepts
+  :class:`~airflow.providers.common.ai.hooks.base.BaseToolset` subclasses
+  (``SQLToolset``), pydantic-ai ``AbstractToolset`` implementations
+  (``HookToolset``, ``MCPToolset``, ``DataFusionToolset``,
+  ``AgentSkillsToolset`` for :ref:`agent-skills`, third-party toolsets),
+  pydantic-ai ``DynamicToolset`` instances, plain Python callables, or native
+  pydantic-ai ``Tool`` objects. Mixed lists are supported. Bare Python
+  callables are treated as callable tools; wrap pydantic-ai ``ToolsetFunc``
+  factories with ``DynamicToolset`` to pass them through as native dynamic
+  toolsets.
+- ``enable_tool_logging``: When ``True`` (default), wraps each tool call with
+  real-time logging. For pydantic-ai ``AbstractToolset`` items this is done via
+  :class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset`; for
+  plain callables and :class:`~airflow.providers.common.ai.hooks.base.BaseToolset`
+  items it is applied at the callable level.
 - ``agent_params``: Additional keyword arguments passed to the pydantic-ai
   ``Agent`` constructor (e.g. ``retries``, ``model_settings``, ``capabilities``).
   See :ref:`capabilities-passthrough` for how to enable pydantic-ai capabilities
@@ -327,9 +340,9 @@ Parameters
 Logging
 -------
 
-All AI operators automatically log a post-run summary after ``run_sync()``
-completes. ``AgentOperator`` additionally wraps toolsets for real-time
-per-tool-call logging (controlled by ``enable_tool_logging``).
+All AI operators automatically log a post-run summary after the agent run
+completes. ``AgentOperator`` additionally provides real-time per-tool-call
+logging (controlled by ``enable_tool_logging``).
 
 **Real-time tool call logging** (AgentOperator only) — each tool call is
 logged as it happens:

@@ -34,29 +34,64 @@ Three toolsets are included:
   `MCP servers <https://modelcontextprotocol.io/>`__ configured via Airflow
   connections.
 
-All three implement pydantic-ai's
-`AbstractToolset <https://ai.pydantic.dev/toolsets/>`__ interface and can be
-passed to any pydantic-ai ``Agent``, including via
-:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`.
+:class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` and
+:class:`~airflow.providers.common.ai.toolsets.mcp.MCPToolset` implement pydantic-ai's
+`AbstractToolset <https://ai.pydantic.dev/toolsets/>`__ interface.
+:class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` implements the
+framework-agnostic :class:`~airflow.providers.common.ai.hooks.base.BaseToolset` interface.
+All three can be passed to
+:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`, which routes each
+toolset to the correct agent parameter automatically.
 
 .. note::
 
-    ``AgentOperator`` accepts **any** ``AbstractToolset`` implementation — not
-    just the Airflow-native toolsets above. PydanticAI's own MCP server
-    classes (``MCPServerStreamableHTTP``, ``MCPServerSSE``, ``MCPServerStdio``)
-    and third-party toolsets work too. The Airflow-native toolsets add
-    connection management, secret backend integration, and the connection UI,
-    but you are not locked in.
+    ``AgentOperator`` accepts a mixed ``toolsets`` list containing any
+    combination of:
 
+    - pydantic-ai ``AbstractToolset`` implementations (``HookToolset``,
+      ``MCPToolset``, ``DataFusionToolset``).
+    - Any third-party ``AbstractToolset``, including PydanticAI's own MCP
+      server classes (``MCPServerStreamableHTTP``, ``MCPServerSSE``,
+      ``MCPServerStdio``).
+    - pydantic-ai dynamic toolsets, by wrapping a ``ToolsetFunc`` factory with
+      ``DynamicToolset``.
+    - :class:`~airflow.providers.common.ai.hooks.base.BaseToolset`
+      subclasses (``SQLToolset``).
+    - Plain Python callables (``def my_tool(...): ...``).
+    - Native pydantic-ai ``Tool`` objects.
 
-Using Toolsets Directly with PydanticAI
----------------------------------------
+    The hook routes each item to the correct agent parameter automatically.
+    Bare Python callables are treated as callable tools. To pass a pydantic-ai
+    ``ToolsetFunc`` factory through as a native dynamic toolset, wrap it with
+    ``DynamicToolset``:
 
-Toolsets are standard pydantic-ai ``AbstractToolset`` implementations with no
-dependency on ``AgentOperator`` or ``@task.agent``. You can use them anywhere
-you can run Python within Airflow -- ``@task`` functions, ``PythonOperator``
-callables, or any custom operator's ``execute()`` method -- by creating a
-``pydantic_ai.Agent`` yourself:
+    .. code-block:: python
+
+        from pydantic_ai import RunContext
+        from pydantic_ai.agent import DynamicToolset
+        from pydantic_ai.toolsets import AbstractToolset
+
+
+        def select_toolset(ctx: RunContext) -> AbstractToolset | None:
+            return None
+
+
+        AgentOperator(
+            task_id="agent",
+            prompt="Answer with the tools available for this run.",
+            llm_conn_id="pydanticai_default",
+            toolsets=[DynamicToolset(select_toolset)],
+        )
+
+
+Using Toolsets Directly
+-----------------------
+
+Toolsets can be used anywhere you can run Python within Airflow — ``@task``
+functions, ``PythonOperator`` callables, or any custom operator's
+``execute()`` method — without needing ``AgentOperator`` or ``@task.agent``.
+Pass toolsets via :class:`~airflow.providers.common.ai.hooks.base.AgentRunRequest`
+and call the hook yourself:
 
 .. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py
     :language: python
@@ -67,11 +102,18 @@ This works because toolsets resolve Airflow connections lazily via
 ``BaseHook.get_connection()``, which is available in any task execution
 context.
 
-This approach gives you full control over the agent lifecycle -- you can call
-``agent.run_sync()`` multiple times, swap models at runtime, or combine
-results from several agents in a single task. The tradeoff is that you lose
-the durable execution (step-level caching with retry replay), HITL review
-integration, and automatic tool call logging that ``AgentOperator`` provides.
+This approach gives you direct control over the agent lifecycle — you can
+build and run multiple agents in a single task, or combine results from
+several runs. The tradeoff is that you lose the durable execution
+(step-level caching with retry replay), HITL review integration, and
+automatic tool call logging that
+:class:`~airflow.providers.common.ai.operators.agent.AgentOperator` provides
+via the agent hook (:class:`~airflow.providers.common.ai.hooks.base.BaseAIHook`):
+callable-level logging and caching for
+:class:`~airflow.providers.common.ai.hooks.base.BaseToolset` tools and plain
+callables, and :class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` /
+:class:`~airflow.providers.common.ai.durable.caching_toolset.CachingToolset`
+wrapping for pydantic-ai ``AbstractToolset`` items.
 
 
 ``HookToolset``

@@ -62,6 +62,9 @@ integrations:
     tags: [ai]
 
 hooks:
+  - integration-name: Common AI
+    python-modules:
+      - airflow.providers.common.ai.hooks.base
   - integration-name: Pydantic AI
     python-modules:
       - airflow.providers.common.ai.hooks.pydantic_ai

@@ -69,7 +69,7 @@ dependencies = [
     "apache-airflow>=3.0.0",
     "apache-airflow-providers-common-compat>=1.14.1",
     "apache-airflow-providers-standard>=1.12.1",
-    "pydantic-ai-slim>=1.71.0",
+    "pydantic-ai-slim>=1.96.0",
 ]
 
 # The optional dependencies should be modified in place in the generated file
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		86e39c620f3926c99e1c702a496d6161032e1a3ac69eba7da10214a2c4ba24f1
		507d80ebe713a42b294363cf05e50b01c0727a7e384056e4b8b23d5102bd11b5