Add BaseAIHook and Update usages

[gopidesupavan]

Add BaseAIHook and Update usage

Summary

Introduce BaseAIHook, a backend-neutral contract for multi-turn LLM agents in the common-ai provider. AgentOperator and @task.agent now resolve the agent runtime from the connection conn_type (for example pydanticai, pydanticai-bedrock, pydanticai-azure) and delegate all framework-specific work to the hook.

PydanticAIHook is the first implementation. All LLM operators and LLMRetryPolicy are migrated to a shared AgentRunRequest / run_agent API. SQLToolset is migrated to the new framework-agnostic BaseToolset interface.

This lays the foundation for additional agent backends without adding parallel operator classes per framework. A follow-up PR will add AWS Strands as the next hook implementation; this contract also opens the door for Google ADK and other agent runtimes behind the same AgentOperator / @task.agent surface.

---

Motivation

Before this change:

• AgentOperator contained pydantic-ai-specific logic (tool wrapping, durable caching, agent construction).
• PydanticAIHook.create_agent() / run_agent() used ad-hoc keyword arguments.
• Tool logging and durable execution were handled in the operator layer via LoggingToolset / CachingToolset wrappers.

The operator should stay framework-agnostic. Hooks should own agent lifecycle, tool resolution, durable execution, and normalized results.

---

Design

BaseAIHook contract

New abstract hook with:

Method / property	Purpose
get_model()	Return backend model/client
get_conn()	Compatibility shim → get_model()
create_agent(request)	Build (but do not run) the agent
run_agent(agent, request)	Execute and return AgentRunResult
_tool_spec_to_native(spec)	Convert ToolSpec → native tool representation
get_agent_hook(conn_id)	Resolve hook from connection conn_type

Capability flags: supports_toolsets, supports_durable, supports_usage_limits.

Parameter objects

• AgentRunRequest — prompt, output type, instructions, toolsets, usage limits, message history, durable context, agent params
• AgentRunResult — output, message history, model name, usage, tool names, durable stats
• ToolSpec — framework-neutral tool descriptor (name, description, JSON schema, callable)
• BaseToolset — abstract as_tools() → list[ToolSpec]
• DurableContext / DurableStats — durable execution identity and cache statistics

Shared hook helpers

Moved into BaseAIHook:

• _resolve_tools() — converts BaseToolset, plain callables, and native tool objects
• _logged_callable() — per-tool real-time logging
• _cached_callable() — per-tool durable step caching
• _init_durable() — DurableStorage / DurableStepCounter setup

PydanticAIHook implementation

• Implements full BaseAIHook contract
• Splits toolsets into two paths:
  • AbstractToolset (HookToolset, MCPToolset, DataFusionToolset, third-party) → Agent(toolsets=[...]) with LoggingToolset / CachingToolset wrapping when enabled
  • BaseToolset / callables / native Tool → resolved via _resolve_tools → Agent(tools=[...])
• Durable model caching via CachingModel in run_agent
• get_model() replaces direct get_conn() usage; get_conn() delegates for backward compatibility

AgentOperator thinning

Operator execution is now:

request=self._build_request(prompt=self.prompt)
agent=self.llm_hook.create_agent(request)
run_result=self.llm_hook.run_agent(agent, request)

No pydantic-ai imports at runtime (except UsageLimits under TYPE_CHECKING).

Early validation via _validate_hook_capabilities() checks hook support for toolsets, durable, and usage limits.

SQLToolset → BaseToolset

SQLToolset no longer implements pydantic-ai's AbstractToolset. It implements BaseToolset.as_tools() returning four ToolSpec objects with JSON schemas (list_tables, get_schema, query, check_query).

HookToolset, MCPToolset, and DataFusionToolset remain AbstractToolset and continue to work unchanged through the pydantic-ai routing path.

---

Other changes

All LLM operators migrated

These now use BaseAIHook.get_agent_hook() and AgentRunRequest:

• LLMOperator
• LLMBranchOperator
• LLMSQLOperator
• LLMSchemaCompareOperator
• LLMFileAnalysisOperator
• LLMRetryPolicy

Logging utilities

• log_run_summary() now accepts AgentRunResult directly
• Removed wrap_toolsets_for_logging() from the operator path; logging is handled in the hook layer

Examples and docs

• Updated example_pydantic_ai_hook.py to use BaseAIHook.get_agent_hook() + AgentRunRequest
• Updated docs/operators/agent.rst, docs/toolsets.rst, AGENTS.md
• Changelog entry for the new contract

Tests

• New test_base_ai.py — dataclasses, _resolve_tools, logging/caching wrappers
• Expanded test_pydantic_ai.py — contract, durable init, AbstractToolset routing/wrapping
• Updated operator, decorator, and policy tests to mock BaseAIHook and assert AgentRunRequest forwarding
• Rewritten test_sql.py for BaseToolset.as_tools() API

---

Breaking changes

PydanticAIHook API

Before:

agent=hook.create_agent(output_type=str, instructions="...", toolsets=[...])
result=hook.run_agent(agent, prompt="hello", usage_limits=limits)

After:

request=AgentRunRequest(prompt="hello", output_type=str, instructions="...", toolsets=[...], usage_limits=limits)
agent=hook.create_agent(request)
result=hook.run_agent(agent, request)

get_conn() still works (delegates to get_model()).

SQLToolset direct pydantic-ai usage

Before: pass SQLToolset(...) directly to pydantic-ai Agent(toolsets=[...]).

After: use via AgentOperator / @task.agent, or build through the hook:

request=AgentRunRequest(prompt="...", toolsets=[SQLToolset(db_conn_id="my_db")])
agent=hook.create_agent(request)
result=hook.run_agent(agent, request)

SQLToolset is now a BaseToolset, not an AbstractToolset.

---

Migration guide

Custom code calling PydanticAIHook directly

Replace kwargs-style create_agent / run_agent with AgentRunRequest:

fromairflow.providers.common.ai.hooks.base_aiimportAgentRunRequest, BaseAIHookhook=BaseAIHook.get_agent_hook("pydanticai_default", hook_params={"model_id": "openai:gpt-5"})
request=AgentRunRequest(
prompt="Analyze this dataset",
output_type=str,
instructions="You are a data analyst.",
toolsets=[SQLToolset(db_conn_id="postgres_default")],
)
agent=hook.create_agent(request)
result=hook.run_agent(agent, request)
print(result.output)

DAG authors using operators / decorators

No DAG changes required for:

• AgentOperator / @task.agent
• LLMOperator / @task.llm
• Other LLM decorators

Connection conn_type continues to select the backend.

Adding a new agent backend

Subclass BaseAIHook and implement:

1. get_model()
2. create_agent(request)
3. run_agent(agent, request)
4. _tool_spec_to_native(spec)

Register the hook in provider.yaml. Reuse shared helpers (_resolve_tools, _logged_callable, _cached_callable, _init_durable) where applicable.

---

Known limitations / follow-ups

• HookToolset / DataFusionToolset could be migrated to BaseToolset in a follow-up; they work today via the AbstractToolset pass-through path.

---

Test plan

• tests/unit/common/ai/hooks/test_base_ai.py

• tests/unit/common/ai/hooks/test_pydantic_ai.py

• tests/unit/common/ai/operators/test_agent.py

• tests/unit/common/ai/operators/test_llm.py

• tests/unit/common/ai/operators/test_llm_branch.py

• tests/unit/common/ai/operators/test_llm_sql.py

• tests/unit/common/ai/operators/test_llm_schema_compare.py

• tests/unit/common/ai/operators/test_llm_file_analysis.py

• tests/unit/common/ai/decorators/test_agent.py

• tests/unit/common/ai/decorators/test_llm*.py

• tests/unit/common/ai/policies/test_retry.py

• tests/unit/common/ai/toolsets/test_sql.py

• tests/unit/common/ai/utils/test_logging.py

• Follow-up: AWS Strands agent hook (StrandsAIHook implementing BaseAIHook)

• Follow-up: Google ADK agent hook (same contract, new conn_type registration)

• Follow-up: migrate HookToolset, MCPToolset, and DataFusionToolset from pydantic-ai AbstractToolset to BaseToolset

Was generative AI tooling used to co-author this PR?

• Yes
• No

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[Copilot]

Pull request overview

This PR introduces a new BaseAIHook contract in the common-ai provider to make multi-turn agent execution backend-neutral. Operators/decorators now construct an AgentRunRequest, resolve the runtime hook from the connection conn_type, and delegate agent lifecycle/tool resolution/durable execution to the hook implementation (starting with PydanticAIHook).

Changes:

• Add BaseAIHook + shared request/response/tool abstractions (AgentRunRequest, AgentRunResult, ToolSpec, BaseToolset) and shared helper logic (tool resolution, logging, caching).
• Refactor AgentOperator, LLM operators/decorators, and LLMRetryPolicy to use get_agent_hook() and the shared create_agent(request) / run_agent(agent, request) flow.
• Migrate SQLToolset to the framework-agnostic BaseToolset interface and update logging utilities, docs, examples, and tests accordingly.

Reviewed changes

Copilot reviewed 34 out of 35 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
uv.lock	Lockfile update (adds additional jpype1 wheel entries).
providers/common/ai/tests/unit/common/ai/utils/test_logging.py	Update logging tests to validate AgentRunResult-based summaries.
providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py	Rewrite SQLToolset tests for BaseToolset.as_tools() + direct tool callables.
providers/common/ai/tests/unit/common/ai/policies/test_retry.py	Update retry policy tests to mock BaseAIHook.get_agent_hook() + request forwarding.
providers/common/ai/tests/unit/common/ai/operators/test_llm.py	Update LLMOperator tests to assert AgentRunRequest construction and run_agent usage.
providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py	Update LLMSQL operator tests for request-based hook invocation.
providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py	Update schema compare operator tests to validate request contents + run_agent call.
providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py	Update file analysis operator tests for BaseAIHook + request flow.
providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py	Update branch operator tests to use BaseAIHook and request-based execution.
providers/common/ai/tests/unit/common/ai/operators/test_agent.py	Update AgentOperator tests for capability validation + request/durable context forwarding.
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py	Expand PydanticAIHook tests for BaseAIHook contract, tool routing, and durable behavior.
providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py	Add new unit tests covering BaseAIHook dataclasses and tool/log/cache helpers.
providers/common/ai/tests/unit/common/ai/decorators/test_llm.py	Update @task.llm tests to mock BaseAIHook.get_agent_hook() and validate request prompt.
providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py	Update @task.llm_sql tests for request-based hook execution.
providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py	Update schema compare decorator tests for BaseAIHook.get_agent_hook() flow.
providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py	Update file analysis decorator tests for request-based execution.
providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py	Update branch decorator tests to mock BaseAIHook.get_agent_hook() and validate behavior.
providers/common/ai/tests/unit/common/ai/decorators/test_agent.py	Update @task.agent tests for request forwarding and toolset passthrough.
providers/common/ai/src/airflow/providers/common/ai/utils/logging.py	Make logging backend-neutral by consuming AgentRunResult directly.
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py	Convert SQLToolset from pydantic-ai AbstractToolset to framework-neutral BaseToolset.
providers/common/ai/src/airflow/providers/common/ai/policies/retry.py	Migrate LLMRetryPolicy to use BaseAIHook + AgentRunRequest and run_agent.
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py	Refactor LLMOperator to build AgentRunRequest and call hook create_agent/run_agent.
providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py	Refactor LLMSQL operator to use AgentRunRequest and hook execution.
providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py	Refactor schema compare operator to request-based hook execution.
providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py	Refactor file analysis operator to request-based hook execution.
providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py	Refactor branch operator to request-based hook execution.
providers/common/ai/src/airflow/providers/common/ai/operators/agent.py	Thin AgentOperator: capability validation + request building + hook-driven execution/durable/tool logging.
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py	Implement BaseAIHook contract for pydantic-ai, including tool routing and durable execution support.
providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py	Add new BaseAIHook contract, request/result dataclasses, tool abstraction, and shared helpers.
providers/common/ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py	Update example DAG to use BaseAIHook.get_agent_hook() + AgentRunRequest.
providers/common/ai/docs/toolsets.rst	Update toolset docs to describe mixed toolset routing and BaseToolset usage.
providers/common/ai/docs/operators/agent.rst	Update AgentOperator docs to explain backend selection via connection conn_type and new toolset shapes.
providers/common/ai/docs/hooks/index.rst	Update hook selection docs to reflect conn_type-driven backend selection for agents.
providers/common/ai/docs/changelog.rst	Add changelog entry for BaseAIHook contract introduction.
providers/common/ai/AGENTS.md	Update contributor guidance to describe BaseAIHook and backend-neutral agent design.

[gopidesupavan]

Example dag which has all the variants to validate

"""
Example DAGs demonstrating bound methods, functools.partial, and callable objects as agent tools.
These patterns are supported natively — no BaseToolset subclass needed.
"""
from __future__ import annotations
import functools
from pydantic_ai.tools import Tool
from airflow.providers.common.ai.operators.agent import AgentOperator
from airflow.providers.common.compat.sdk import dag, task
# ---------------------------------------------------------------------------
# 1. Bound method: methods on a service class passed directly as tools
# ---------------------------------------------------------------------------
# [START howto_agent_bound_method_tools]
@dag(schedule=None, tags=["example"])
def example_agent_bound_method_tools():
"""Pass bound methods of a service class directly as agent tools."""
class InventoryService:
"""Thin wrapper around an inventory data source."""
def __init__(self, warehouse_id: str) -> None:
self._warehouse_id = warehouse_id
def get_stock_level(self, product_id: str) -> int:
"""Return the current stock count for a product in this warehouse."""
# Replace with a real DB/API call in production.
mock_stock = {"SKU-001": 42, "SKU-002": 0, "SKU-003": 17}
return mock_stock.get(product_id, -1)
def list_low_stock(self, threshold: int = 10) -> list[str]:
"""Return product IDs whose stock is at or below *threshold*."""
mock_stock = {"SKU-001": 42, "SKU-002": 0, "SKU-003": 17}
return [pid for pid, qty in mock_stock.items() if qty <= threshold]
service = InventoryService(warehouse_id="WH-EU-01")
AgentOperator(
task_id="inventory_analyst",
prompt="Which products are running low and what are their exact stock levels?",
llm_conn_id="pydanticai_default",
system_prompt=(
"You are a warehouse inventory assistant. "
"Use the tools to identify low-stock products and report their quantities."
),
# Bound methods are passed directly — __name__ and __doc__ are picked up automatically.
toolsets=[service.get_stock_level, service.list_low_stock],
)
# [END howto_agent_bound_method_tools]
example_agent_bound_method_tools()
# ---------------------------------------------------------------------------
# 2. functools.partial: pre-configure a generic function for a specific context
# ---------------------------------------------------------------------------
# [START howto_agent_partial_tools]
@dag(schedule=None, tags=["example"])
def example_agent_partial_tools():
"""Pre-configure generic functions with functools.partial before passing as tools."""
def fetch_metric(environment: str, metric_name: str) -> float:
"""Fetch a named metric value from the given environment."""
# Replace with a real metrics API call in production.
mock = {
("prod", "error_rate"): 0.012,
("prod", "p99_latency_ms"): 145.0,
("prod", "requests_per_second"): 3200.0,
}
return mock.get((environment, metric_name), 0.0)
def list_available_metrics(environment: str) -> list[str]:
"""List the metric names available in the given environment."""
return ["error_rate", "p99_latency_ms", "requests_per_second"]
# Pre-bind the environment so the agent only needs to supply metric_name.
prod_fetch_metric = functools.partial(fetch_metric, "prod")
prod_list_metrics = functools.partial(list_available_metrics, "prod")
AgentOperator(
task_id="sre_analyst",
prompt="Is the production service healthy? Check error rate and latency.",
llm_conn_id="pydanticai_default",
system_prompt=(
"You are an SRE assistant. "
"Use the tools to inspect production metrics and summarise service health."
),
# functools.partial — tool name is taken from the underlying function (__func__.__name__).
toolsets=[prod_fetch_metric, prod_list_metrics],
)
# [END howto_agent_partial_tools]
example_agent_partial_tools()
# ---------------------------------------------------------------------------
# 3. Callable object: a class with __call__ encapsulating shared state
# ---------------------------------------------------------------------------
# [START howto_agent_callable_object_tools]
@dag(schedule=None, tags=["example"])
def example_agent_callable_object_tools():
"""Pass a callable object (class with __call__) directly as an agent tool."""
class CustomerLookup:
"""Look up customer details from a shared in-memory store."""
def __init__(self, customer_data: dict) -> None:
self._data = customer_data
def __call__(self, customer_id: str) -> dict:
"""Return name, tier, and lifetime value for the given customer ID."""
return self._data.get(customer_id, {"error": f"Customer {customer_id!r} not found"})
lookup = CustomerLookup(
customer_data={
"C-001": {"name": "Acme Corp", "tier": "enterprise", "ltv_usd": 85000},
"C-002": {"name": "Globex Ltd", "tier": "pro", "ltv_usd": 12000},
"C-003": {"name": "Initech", "tier": "starter", "ltv_usd": 900},
}
)
@task.agent(
llm_conn_id="pydanticai_default",
system_prompt=(
"You are a customer success assistant. "
"Use the CustomerLookup tool to retrieve customer details and answer questions. "
"Always call CustomerLookup with the customer_id from the question before answering. "
"Do not guess customer attributes without a tool lookup."
),
# Callable object — tool name defaults to the class name (CustomerLookup).
toolsets=[lookup],
)
def analyse(question: str) -> str:
return question
analyse("Call CustomerLookup for customer C-001 and report that customer's tier and lifetime value.")
# [END howto_agent_callable_object_tools]
example_agent_callable_object_tools()
# ---------------------------------------------------------------------------
# 4. Mixed: combine all three callable patterns in one agent
# ---------------------------------------------------------------------------
# [START howto_agent_mixed_callable_tools]
@dag(schedule=None, tags=["example"])
def example_agent_mixed_callable_tools():
"""Mix bound methods, functools.partial, and callable objects in a single agent."""
# --- bound method ---
class OrderService:
def get_order(self, order_id: str) -> dict:
"""Fetch order details by order ID."""
mock = {
"ORD-1": {"status": "shipped", "items": 3, "total_usd": 299.0},
"ORD-2": {"status": "pending", "items": 1, "total_usd": 49.0},
}
return mock.get(order_id, {"error": "not found"})
order_service = OrderService()
# --- functools.partial ---
def send_notification(channel: str, message: str) -> str:
"""Send *message* to a notification *channel* and return a confirmation."""
# Replace with a real Slack/email call in production.
return f"Sent to {channel!r}: {message}"
notify_ops = functools.partial(send_notification, "ops-alerts")
# --- callable object ---
class ExchangeRate:
def __call__(self, currency: str) -> float:
"""Return the current USD exchange rate for the given currency code."""
rates = {"EUR": 1.08, "GBP": 1.27, "JPY": 0.0067}
return rates.get(currency.upper(), 1.0)
exchange_rate = ExchangeRate()
AgentOperator(
task_id="order_ops_agent",
prompt=(
"Check orders ORD-1 and ORD-2. Convert ORD-1's total to EUR and send a summary to ops-alerts."
),
llm_conn_id="pydanticai_default",
system_prompt=(
"You are an order operations assistant. "
"Use the available tools to look up orders, convert currencies, and send notifications."
),
toolsets=[
order_service.get_order, # bound method
notify_ops, # functools.partial
exchange_rate, # callable object
],
)
# [END howto_agent_mixed_callable_tools]
example_agent_mixed_callable_tools()
# ---------------------------------------------------------------------------
# 5. Mixed native Tool with Airflow-resolved callables
# ---------------------------------------------------------------------------
# [START howto_agent_mixed_native_tool]
@dag(schedule=None, tags=["example"])
def example_agent_mixed_native_tool():
"""Mix plain Python callables with a native pydantic-ai Tool."""
def get_customer(customer_id: str) -> dict:
"""Fetch a customer profile by customer ID."""
customers = {
"C-001": {"name": "Acme Corp", "tier": "enterprise", "renewal_risk": "medium"},
"C-002": {"name": "Globex Ltd", "tier": "pro", "renewal_risk": "low"},
}
return customers.get(customer_id, {"error": f"Customer {customer_id!r} not found"})
def calculate_discount(customer_id: str, order_total_usd: float) -> float:
"""Calculate an approved discount percentage for a renewal order."""
if customer_id == "C-001" and order_total_usd >= 10000:
return 12.5
return 5.0
def escalate_account(customer_id: str, reason: str) -> str:
"""Create an escalation note for a customer account."""
return f"Escalated {customer_id}: {reason}"
escalation_tool = Tool(
escalate_account,
name="escalate_account",
description="Create an escalation note when a customer needs human follow-up.",
)
AgentOperator(
task_id="customer_retention_agent",
prompt=(
"Look up customer C-001, calculate the approved discount for a 12000 USD renewal, "
"and escalate the account with the reason 'high-value renewal'."
),
llm_conn_id="pydanticai_default",
system_prompt=(
"You are a customer retention assistant. Use get_customer for account data, "
"calculate_discount for pricing guidance, and escalate_account when follow-up is needed."
),
toolsets=[
get_customer, # Airflow-resolved callable
calculate_discount, # Airflow-resolved callable
escalation_tool, # native pydantic-ai Tool, passed through unchanged
],
)
# [END howto_agent_mixed_native_tool]
example_agent_mixed_native_tool()