Add BaseAIHook and Update usage

Author: gopidesupavan

providers/common/ai/AGENTS.md

@@ -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_ai.py`
 - Operators: `src/airflow/providers/common/ai/operators/`
 - Decorators: `src/airflow/providers/common/ai/decorators/`
 - Toolsets: `src/airflow/providers/common/ai/toolsets/`

providers/common/ai/docs/changelog.rst

@@ -25,6 +25,15 @@
 Changelog
 ---------
 
+Next release
+............
+
+Features
+^^^^^^^^
+
+* Add ``BaseAIHook`` contract with framework-agnostic ``create_agent`` / ``run_agent`` /
+  ``get_model`` interface so ``AgentOperator`` selects the agent backend via connection type.
+
 0.3.0
 .....
 

providers/common/ai/docs/hooks/index.rst

@@ -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

providers/common/ai/docs/operators/agent.rst

@@ -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
@@ -251,11 +254,17 @@ 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``,
-  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_ai.BaseToolset` subclasses
+  (``SQLToolset``), pydantic-ai ``AbstractToolset`` implementations
+  (``HookToolset``, ``MCPToolset``, ``DataFusionToolset``, third-party toolsets),
+  plain Python callables, or native pydantic-ai ``Tool`` objects. Mixed lists
+  are supported.
+- ``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_ai.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``).
 - ``usage_limits``: Optional pydantic-ai ``UsageLimits`` enforced on every
@@ -272,9 +281,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:

providers/common/ai/docs/toolsets.rst

@@ -34,29 +34,41 @@ 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_ai.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``).
+    - :class:`~airflow.providers.common.ai.hooks.base_ai.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.
 
-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:
+
+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_ai.AgentRunRequest`
+and call the hook yourself:
 
 .. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py
     :language: python
@@ -67,11 +79,13 @@ 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
+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 ``AgentOperator`` provides.
+integration, and the automatic tool call logging and routing that
+``AgentOperator`` provides via
+:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset`.
 
 
 ``HookToolset``

providers/common/ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py

@@ -14,13 +14,13 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""Example DAGs demonstrating PydanticAIHook and direct pydantic-ai Agent usage."""
+"""Example DAGs demonstrating BaseAIHook and AgentRunRequest usage."""
 
 from __future__ import annotations
 
 from pydantic import BaseModel
 
-from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, BaseAIHook
 from airflow.providers.common.compat.sdk import dag, task
 
 
@@ -29,9 +29,10 @@
 def example_pydantic_ai_hook():
     @task
     def generate_summary(text: str) -> str:
-        hook = PydanticAIHook(llm_conn_id="pydanticai_default")
-        agent = hook.create_agent(output_type=str, instructions="Summarize concisely.")
-        result = agent.run_sync(text)
+        hook = BaseAIHook.get_agent_hook("pydanticai_default")
+        request = AgentRunRequest(prompt=text, output_type=str, instructions="Summarize concisely.")
+        agent = hook.create_agent(request)
+        result = hook.run_agent(agent, request)
         return result.output
 
     generate_summary("Apache Airflow is a platform for programmatically authoring...")
@@ -51,12 +52,14 @@ class SQLResult(BaseModel):
             query: str
             explanation: str
 
-        hook = PydanticAIHook(llm_conn_id="pydanticai_default")
-        agent = hook.create_agent(
+        hook = BaseAIHook.get_agent_hook("pydanticai_default")
+        request = AgentRunRequest(
+            prompt=prompt,
             output_type=SQLResult,
             instructions="Generate a SQL query and explain it.",
         )
-        result = agent.run_sync(prompt)
+        agent = hook.create_agent(request)
+        result = hook.run_agent(agent, request)
         return result.output.model_dump()
 
     generate_sql("Find the top 10 customers by revenue")
@@ -76,8 +79,9 @@ def example_task_with_toolsets():
     def analyze_revenue() -> str:
         from airflow.providers.common.ai.toolsets.sql import SQLToolset
 
-        hook = PydanticAIHook(llm_conn_id="pydanticai_default")
-        agent = hook.create_agent(
+        hook = BaseAIHook.get_agent_hook("pydanticai_default")
+        request = AgentRunRequest(
+            prompt="Which customers have spent the most? Show the top 5.",
             output_type=str,
             instructions=(
                 "You are a sales analytics assistant. "
@@ -91,7 +95,8 @@ def analyze_revenue() -> str:
                 ),
             ],
         )
-        result = agent.run_sync("Which customers have spent the most? Show the top 5.")
+        agent = hook.create_agent(request)
+        result = hook.run_agent(agent, request)
         return result.output
 
     analyze_revenue()