diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 00c557d1eb2a0..8e21d1e37c3cf 100644 --- 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 index d9c5dac2752bc..501dcf5a4fd42 100644 --- 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 index 8a645f848dd33..9b311a0a6a4ac 100644 --- a/generated/provider_dependencies.json.sha256sum +++ b/generated/provider_dependencies.json.sha256sum @@ -1 +1 @@ -86e39c620f3926c99e1c702a496d6161032e1a3ac69eba7da10214a2c4ba24f1 +507d80ebe713a42b294363cf05e50b01c0727a7e384056e4b8b23d5102bd11b5 diff --git a/providers/common/ai/AGENTS.md b/providers/common/ai/AGENTS.md index 8e8711c8ab78f..15dbd055a95e2 100644 --- a/providers/common/ai/AGENTS.md +++ b/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.py` - Operators: `src/airflow/providers/common/ai/operators/` - Decorators: `src/airflow/providers/common/ai/decorators/` - Toolsets: `src/airflow/providers/common/ai/toolsets/` diff --git a/providers/common/ai/README.rst b/providers/common/ai/README.rst index 8be522e21cf19..1b13fdda6fd53 100644 --- a/providers/common/ai/README.rst +++ b/providers/common/ai/README.rst @@ -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 diff --git a/providers/common/ai/docs/changelog.rst b/providers/common/ai/docs/changelog.rst index 9a9e4cfc81276..5e8190a1995f1 100644 --- a/providers/common/ai/docs/changelog.rst +++ b/providers/common/ai/docs/changelog.rst @@ -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 ..... diff --git a/providers/common/ai/docs/hooks/index.rst b/providers/common/ai/docs/hooks/index.rst index 2786cd1b0ced5..caa47d7d19b36 100644 --- a/providers/common/ai/docs/hooks/index.rst +++ b/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 diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index 4b3014ca1ffbb..029ce4bd97299 100644 --- a/providers/common/ai/docs/index.rst +++ b/providers/common/ai/docs/index.rst @@ -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 diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index 8e7c8ad5f3983..52f2e2a3bf3a4 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/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 ` + :ref:`Pydantic AI connection ` + +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: diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index b5e868abea209..5355bcf498d59 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -34,29 +34,64 @@ Three toolsets are included: `MCP servers `__ configured via Airflow connections. -All three implement pydantic-ai's -`AbstractToolset `__ 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 `__ 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`` diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index 695a88ec2c103..78da663821091 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -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 diff --git a/providers/common/ai/pyproject.toml b/providers/common/ai/pyproject.toml index 910ef1401c978..cc9a79775bae2 100644 --- a/providers/common/ai/pyproject.toml +++ b/providers/common/ai/pyproject.toml @@ -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 diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_callable_toolsets.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_callable_toolsets.py new file mode 100644 index 0000000000000..b843f1dcb3714 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_callable_toolsets.py @@ -0,0 +1,231 @@ +# 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 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 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() diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py index d1790dcaba6b1..1825f44463eae 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py +++ b/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 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() diff --git a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py index 8bc03c266cb14..9407941cba00e 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py +++ b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py @@ -67,6 +67,7 @@ def get_provider_info(): }, ], "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"], diff --git a/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py b/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py new file mode 100644 index 0000000000000..f85ce7d4ae509 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py @@ -0,0 +1,414 @@ +# 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. +"""Shared contract for agent-framework hooks used by :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`.""" + +from __future__ import annotations + +import functools +import json +import time +from abc import ABCMeta, abstractmethod +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, ClassVar, Generic, TypeVar + +from airflow.providers.common.ai.utils.callables import is_async_callable +from airflow.providers.common.ai.utils.function_schema import callable_to_tool_spec +from airflow.providers.common.compat.sdk import BaseHook + +AgentT = TypeVar("AgentT") + + +class Capability(str, Enum): + """ + Capability tokens declared by concrete hook classes. + + A hook advertises its support by including the relevant tokens in its + :attr:`BaseAIHook.capabilities` frozenset. + :meth:`BaseAIHook.validate_run_request` rejects requests that use a + feature whose token is absent. + """ + + TOOLSETS = "toolsets" + USAGE_LIMITS = "usage_limits" + DURABLE = "durable" + + +@dataclass +class AgentUsage: + """Token and request usage from an agent run, when the backend exposes it.""" + + requests: int | None = None + tool_calls: int | None = None + input_tokens: int | None = None + output_tokens: int | None = None + total_tokens: int | None = None + + +@dataclass +class DurableStats: + """Step-level cache statistics from a durable agent run.""" + + replayed_model: int = 0 + replayed_tool: int = 0 + cached_model: int = 0 + cached_tool: int = 0 + + +@dataclass +class AgentRunResult: + """ + Backend-neutral result from :meth:`BaseAIHook.run_agent`. + + :param output: Final agent output (``str``, Pydantic model instance, etc.). + :param message_history: Opaque conversation state for HITL regeneration; only pass back to the + same hook implementation that produced it. + :param model_name: Resolved model identifier, when available. + :param usage: Usage counters when the backend exposes them. + :param tool_names: Ordered tool names invoked during the run, when known. + :param durable_stats: Durable step-cache statistics, populated when durable execution is enabled. + """ + + output: Any + message_history: Any = None + model_name: str | None = None + usage: AgentUsage | None = None + tool_names: list[str] | None = None + durable_stats: DurableStats | None = None + + +@dataclass +class ToolSpec: + """ + Framework-neutral tool descriptor. + + Toolsets produce :class:`ToolSpec` objects; each hook converts them to its + native tool representation via :meth:`BaseAIHook._tool_spec_to_native`. + + :param name: Tool name exposed to the LLM. + :param description: Human-readable description used by the LLM to decide when to call this tool. + :param parameters: JSON Schema ``object`` describing the tool's parameters. + :param fn: Callable that implements the tool. Must accept keyword arguments matching *parameters*. + :param sequential: When ``True``, the backend must not invoke this tool concurrently with others + in the same turn (for example when tools share a non-thread-safe connection). + """ + + name: str + description: str + parameters: dict[str, Any] + fn: Callable[..., Any] + sequential: bool = False + + +@dataclass +class DurableContext: + """Framework-neutral identity of the running task, used to locate the durable cache file.""" + + dag_id: str + task_id: str + run_id: str + map_index: int = -1 + + +@dataclass +class AgentRunRequest: + """ + Parameter object passed to :meth:`BaseAIHook.create_agent` and :meth:`BaseAIHook.run_agent`. + + Encapsulates everything the hook needs to build and run an agent in a single + framework-neutral structure, so that :class:`~airflow.providers.common.ai.operators.agent.AgentOperator` + has zero framework-specific imports. This contract is currently validated by + the pydantic-ai hook family and may evolve as more framework backends are added. + + :param prompt: User prompt for this invocation (plain ``str`` or a multimodal + ``Sequence`` accepted by the backend agent's run API). + :param output_type: Expected structured output type (default: ``str``). + :param instructions: System-level instructions for the agent. + :param toolsets: List of tools/toolsets the agent may call (BaseToolset instances, plain callables, or backend-native tool objects). + :param usage_limits: Backend-specific usage limits; ignored if the hook does not support them. + :param message_history: Prior conversation state from a previous :class:`AgentRunResult`. + :param enable_tool_logging: When ``True`` (default), wraps Airflow-resolved tool callables with + a logging shim. Backend-native tool objects may be passed through unchanged by the concrete + hook and might not receive this wrapper. + :param durable_context: When set, enables step-level durable caching for the run. + :param agent_params: Extra keyword arguments forwarded to the underlying agent constructor. + Use this escape hatch for framework-specific options. + """ + + prompt: str | Sequence[Any] + output_type: type[Any] | None = str + instructions: str = "" + toolsets: list[Any] | None = None + usage_limits: Any = None + message_history: Any = None + enable_tool_logging: bool = True + durable_context: DurableContext | None = None + agent_params: dict[str, Any] = field(default_factory=dict) + + +class BaseToolset(metaclass=ABCMeta): + """ + Abstract base for framework-agnostic toolsets. + + Subclasses implement :meth:`as_tools` to return a list of :class:`ToolSpec` + objects. Each hook converts those specs to its native tool representation + via :meth:`BaseAIHook._tool_spec_to_native`. + """ + + @abstractmethod + def as_tools(self) -> list[ToolSpec]: + """Return the list of tools this toolset exposes.""" + + +class BaseAIHook(BaseHook, Generic[AgentT], metaclass=ABCMeta): + """ + Abstract hook for multi-turn LLM agents. + + :class:`~airflow.providers.common.ai.operators.agent.AgentOperator` resolves the concrete hook + from the Airflow connection ``conn_type`` (for example ``pydanticai`` or ``pydanticai-bedrock``). + + :param llm_conn_id: Optional connection ID override (subclasses may apply a default). + :param model_id: Optional model override; not all backends use this parameter. + + Subclasses implement :meth:`get_model`, :meth:`_build_agent`, :meth:`run_agent`, and + :meth:`_tool_spec_to_native`. + + Shared helpers :meth:`_resolve_tools` and :meth:`_logged_callable` are provided for all hooks. + Durable cache helpers live in ``DurableAgentMixin`` so non-durable hooks do not inherit + backend-specific durable mechanics. + """ + + conn_name_attr = "llm_conn_id" + + capabilities: ClassVar[frozenset[Capability]] = frozenset() + + def __init__( + self, + llm_conn_id: str | None = None, + model_id: str | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.llm_conn_id = llm_conn_id + self.model_id = model_id + + @classmethod + def get_agent_hook(cls, conn_id: str, *, hook_params: dict[str, Any] | None = None) -> BaseAIHook[Any]: + """ + Return an agent hook for *conn_id*, verifying it implements this contract. + + Uses the connection's ``conn_type`` to select the hook class registered in + ``provider.yaml``. + """ + hook = cls.get_hook(conn_id, hook_params=hook_params) + if not isinstance(hook, BaseAIHook): + raise TypeError( + f"Connection {conn_id!r} resolved to {type(hook).__name__}, which is not a BaseAIHook. " + "Use a connection type registered for agent frameworks (e.g. pydanticai, pydanticai-bedrock)." + ) + return hook + + @abstractmethod + def get_model(self) -> Any: + """Return the backend model/client used to construct agents.""" + + def get_conn(self) -> Any: + """ + Return the backend model/client for :class:`~airflow.hooks.base.BaseHook` compatibility. + + Agent hooks use :meth:`get_model` internally; this shim keeps the traditional ``get_conn()`` + hook API available for callers that expect it. + """ + return self.get_model() + + def create_agent(self, request: AgentRunRequest) -> AgentT: + """ + Build (but do not run) the agent described by *request*. + + Responsible for resolving :attr:`AgentRunRequest.toolsets` via + :meth:`_resolve_tools` and constructing the framework-native agent object + with the model, tools, instructions, and output type from *request*. + + :param request: All parameters needed to configure the agent. + :returns: Framework-native agent handle, ready to be passed to :meth:`run_agent`. + """ + self.validate_run_request(request) + return self._build_agent(request) + + @abstractmethod + def _build_agent(self, request: AgentRunRequest) -> AgentT: + """Build the framework-native agent handle after :meth:`validate_run_request` succeeds.""" + + @abstractmethod + def run_agent(self, agent: AgentT, request: AgentRunRequest) -> AgentRunResult: + """ + Execute *agent* for *request* and return a normalized :class:`AgentRunResult`. + + Implementations with durable execution should keep durable state on their + concrete agent handle, apply it during the run, and call ``storage.cleanup()`` + only after a successful run (keep the cache file when the run raises so + Airflow retries can replay cached steps). + + :param agent: Framework-native agent handle produced by :meth:`create_agent`. + :param request: The same request used to create the agent (prompt, usage + limits, message history, etc.). + """ + + @abstractmethod + def _tool_spec_to_native(self, spec: ToolSpec) -> Any: + """ + Convert a :class:`ToolSpec` to the agent framework's native tool representation. + + Called once per tool inside :meth:`_resolve_tools`. The returned object + is collected into a list and passed to the underlying agent constructor. + + :param spec: Universal tool descriptor, with the callable already wrapped + by any enabled logging / caching shims. + """ + + def validate_run_request(self, request: AgentRunRequest) -> None: + """ + Raise if *request* uses features this hook implementation does not support. + + :meth:`create_agent` calls this before delegating to the hook implementation. + """ + hook_name = type(self).__name__ + conn_id = self.llm_conn_id or "unknown" + if request.toolsets and Capability.TOOLSETS not in self.capabilities: + raise ValueError( + f"toolsets not supported for connection {conn_id!r} (conn_type resolves to {hook_name})." + ) + if request.usage_limits is not None and Capability.USAGE_LIMITS not in self.capabilities: + raise ValueError( + f"usage_limits not supported for connection {conn_id!r} (conn_type resolves to {hook_name})." + ) + if request.durable_context is not None and Capability.DURABLE not in self.capabilities: + raise ValueError( + f"durable execution not supported for connection {conn_id!r} (conn_type resolves to {hook_name})." + ) + + def _resolve_tools( + self, + toolsets: list[Any], + enable_logging: bool, + cache_wrapper: Callable[[Callable[..., Any]], Callable[..., Any]] | None = None, + *, + force_sequential: bool = False, + ) -> list[Any]: + """ + Convert a mixed list of toolsets / callables / native tools into framework-native tools. + + Three cases per item: + + * :class:`BaseToolset` — calls ``as_tools()`` and processes each :class:`ToolSpec`. + * Any callable (plain function, bound method, :func:`functools.partial`, or callable + object) — auto-wraps into a :class:`ToolSpec` using ``__name__`` and ``__doc__`` + (with sensible fallbacks for partials and callable objects), then processes it the + same way. + * Anything else — passed through unchanged (assumed to be a native tool object already + constructed for the target framework). + + The processing pipeline for ``BaseToolset`` and callable items is built inside-out: + *fn* → optional log wrapper → optional cache wrapper → :meth:`_tool_spec_to_native`. + At execution time, the outer cache wrapper runs first, so durable cache hits skip + the logging wrapper. + + :param toolsets: Mix of :class:`BaseToolset` instances, callables (functions, bound + methods, :func:`functools.partial`, or callable objects), and native tool objects. + :param enable_logging: When ``True``, wrap each callable with :meth:`_logged_callable`. + :param cache_wrapper: Optional wrapper used by durable hooks to cache Airflow-resolved + callable tools. Native backend tool objects are passed through unchanged. + :param force_sequential: When ``True``, mark all Airflow-resolved callable tools as + sequential. Native backend tool objects are passed through unchanged. + """ + native: list[Any] = [] + for ts in toolsets: + if isinstance(ts, BaseToolset): + specs = ts.as_tools() + elif callable(ts): + specs = [callable_to_tool_spec(ts)] + else: + native.append(ts) + continue + for spec in specs: + fn = spec.fn + if enable_logging: + fn = self._logged_callable(fn, self.log, name=spec.name) + if cache_wrapper is not None: + fn = cache_wrapper(fn) + adapted = ToolSpec( + name=spec.name, + description=spec.description, + parameters=spec.parameters, + fn=fn, + sequential=spec.sequential or force_sequential, + ) + native.append(self._tool_spec_to_native(adapted)) + return native + + @staticmethod + def _logged_callable( + fn: Callable[..., Any], + logger: Any, + *, + name: str | None = None, + ) -> Callable[..., Any]: + """Wrap *fn* to log tool name, args, timing, and exceptions.""" + _tool_name = name or getattr(fn, "__name__", type(fn).__name__) + + if is_async_callable(fn): + + @functools.wraps(fn) + async def async_wrapper(*args, **kwargs): + logger.info("::group::Tool call: %s", _tool_name) + if kwargs: + logger.debug("Tool args: %s", json.dumps(kwargs, default=str)) + start = time.monotonic() + try: + result = await fn(*args, **kwargs) + elapsed = time.monotonic() - start + logger.info("Tool %s returned in %.2fs", _tool_name, elapsed) + return result + except Exception: + elapsed = time.monotonic() - start + logger.exception("Tool %s failed after %.2fs", _tool_name, elapsed) + raise + finally: + logger.info("::endgroup::") + + return async_wrapper + + @functools.wraps(fn) + def sync_wrapper(*args, **kwargs): + logger.info("::group::Tool call: %s", _tool_name) + if kwargs: + logger.debug("Tool args: %s", json.dumps(kwargs, default=str)) + start = time.monotonic() + try: + result = fn(*args, **kwargs) + elapsed = time.monotonic() - start + logger.info("Tool %s returned in %.2fs", _tool_name, elapsed) + return result + except Exception: + elapsed = time.monotonic() - start + logger.exception("Tool %s failed after %.2fs", _tool_name, elapsed) + raise + finally: + logger.info("::endgroup::") + + return sync_wrapper diff --git a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py index 44e2436576f21..32a158d04de74 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py @@ -16,22 +16,45 @@ # under the License. from __future__ import annotations -from typing import TYPE_CHECKING, Any, TypeVar, overload +import functools +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any from pydantic_ai import Agent +from pydantic_ai.messages import ToolCallPart from pydantic_ai.models import infer_model from pydantic_ai.providers import infer_provider, infer_provider_class - +from pydantic_ai.tools import Tool +from pydantic_ai.toolsets.abstract import AbstractToolset + +from airflow.providers.common.ai.durable.caching_model import CachingModel +from airflow.providers.common.ai.durable.caching_toolset import CachingToolset +from airflow.providers.common.ai.hooks.base import ( + AgentRunRequest, + AgentRunResult, + AgentUsage, + BaseAIHook, + Capability, + DurableStats, + ToolSpec, +) +from airflow.providers.common.ai.mixins.durable import DurableAgentMixin, DurableState from airflow.providers.common.ai.observability import genai_instrumentation_settings -from airflow.providers.common.compat.sdk import BaseHook - -OutputT = TypeVar("OutputT") +from airflow.providers.common.ai.toolsets.logging import LoggingToolset if TYPE_CHECKING: from pydantic_ai.models import KnownModelName, Model -class PydanticAIHook(BaseHook): +@dataclass +class PydanticAgentHandle: + """Pydantic-ai agent plus optional durable cache state for one run.""" + + agent: Agent[None, Any] + durable_state: DurableState | None = None + + +class PydanticAIHook(DurableAgentMixin, BaseAIHook[PydanticAgentHandle]): """ Hook for LLM access via pydantic-ai. @@ -57,19 +80,20 @@ class PydanticAIHook(BaseHook): conn_type = "pydanticai" hook_name = "Pydantic AI" + capabilities = frozenset({Capability.TOOLSETS, Capability.USAGE_LIMITS, Capability.DURABLE}) + def __init__( self, llm_conn_id: str | None = None, model_id: str | None = None, **kwargs, ) -> None: - super().__init__(**kwargs) # Resolve at runtime so each subclass uses its own default_conn_name. # A bare `llm_conn_id: str = default_conn_name` would bind the *base* # class value for all subclasses because Python evaluates default # argument values at class-definition time. - self.llm_conn_id = llm_conn_id if llm_conn_id is not None else self.default_conn_name - self.model_id = model_id + resolved_conn_id = llm_conn_id if llm_conn_id is not None else self.default_conn_name + super().__init__(llm_conn_id=resolved_conn_id, model_id=model_id, **kwargs) self._model: Model | None = None @staticmethod @@ -117,7 +141,7 @@ def _get_provider_kwargs( kwargs["base_url"] = base_url return kwargs - def get_conn(self) -> Model: + def get_model(self) -> Model: """ Return a configured pydantic-ai ``Model``. @@ -134,7 +158,7 @@ def get_conn(self) -> Model: if self._model is not None: return self._model - conn = self.get_connection(self.llm_conn_id) + conn = self.get_connection(self.llm_conn_id or self.default_conn_name) extra: dict[str, Any] = conn.extra_dejson model_name: str | KnownModelName = self.model_id or extra.get("model", "") @@ -172,41 +196,201 @@ def _provider_factory(pname: str) -> Any: self._model = infer_model(model_name) return self._model - @overload - def create_agent( - self, output_type: type[OutputT], *, instructions: str, **agent_kwargs - ) -> Agent[None, OutputT]: ... + # ------------------------------------------------------------------ + # BaseAIHook abstract interface + # ------------------------------------------------------------------ - @overload - def create_agent(self, *, instructions: str, **agent_kwargs) -> Agent[None, str]: ... + def _tool_spec_to_native(self, spec: ToolSpec) -> Any: + """Convert a :class:`~airflow.providers.common.ai.hooks.base.ToolSpec` to a pydantic-ai ``Tool``.""" + return Tool.from_schema( + spec.fn, + name=spec.name, + description=spec.description, + sequential=spec.sequential, + json_schema=spec.parameters, + ) - def create_agent( - self, output_type: type[Any] = str, *, instructions: str, **agent_kwargs - ) -> Agent[None, Any]: + def _build_agent(self, request: AgentRunRequest) -> PydanticAgentHandle: """ - Create a pydantic-ai Agent configured with this hook's model. + Build a pydantic-ai ``Agent`` handle from *request*. - When ``[common.ai] otel_export_enabled`` is set and the worker has an - OpenTelemetry exporter configured, the agent is instrumented to emit - GenAI spans through Airflow's tracing pipeline. See + When :attr:`~AgentRunRequest.durable_context` is set, initialises durable + storage and step counter and returns them alongside the native agent for use + by :meth:`run_agent`. When ``[common.ai] otel_export_enabled`` is set and the + worker has an OpenTelemetry exporter configured, the agent is instrumented to + emit GenAI spans through Airflow's tracing pipeline. See :mod:`airflow.providers.common.ai.observability`. - :param output_type: The expected output type from the agent (default: ``str``). - :param instructions: System-level instructions for the agent. - :param agent_kwargs: Additional keyword arguments passed to the Agent constructor. + :param request: Agent configuration — output type, instructions, toolsets, extra params. + + Native pydantic-ai ``Tool`` instances supplied in ``request.toolsets`` are passed through + unchanged. Airflow tool logging and durable tool-result caching are applied to + framework-neutral callables / ``BaseToolset`` specs and pydantic-ai ``AbstractToolset`` + instances, but not to native ``Tool`` instances. """ - agent = Agent(self.get_conn(), output_type=output_type, instructions=instructions, **agent_kwargs) - if "instrument" not in agent_kwargs: - # Set the public ``agent.instrument`` surface rather than the + durable_state = None + if request.durable_context is not None: + durable_state = self._init_durable(request.durable_context) + + extra_kwargs = dict(request.agent_params or {}) + if request.toolsets: + if "tools" in extra_kwargs: + raise ValueError( + "agent_params must not include 'tools' when toolsets= is set on AgentRunRequest." + ) + if "toolsets" in extra_kwargs: + raise ValueError( + "agent_params must not include 'toolsets' when toolsets= is set on AgentRunRequest." + ) + + abstract_items = [ts for ts in request.toolsets if isinstance(ts, AbstractToolset)] + pipeline_items = [ts for ts in request.toolsets if not isinstance(ts, AbstractToolset)] + + if pipeline_items: + resolved: list[Any] = [] + for item in pipeline_items: + if isinstance(item, Tool): + # Native pydantic-ai Tool objects are callable, so this check must happen + # before _resolve_tools(); otherwise they would be rebuilt from their + # __call__ signature as Airflow-resolved callables instead of preserving + # their original schema/configuration. + resolved.append(item) + else: + cache_wrapper = None + if durable_state is not None: + cache_wrapper = functools.partial( + self._cached_callable, + storage=durable_state.storage, + counter=durable_state.counter, + ) + resolved.extend( + self._resolve_tools( + [item], + request.enable_tool_logging, + cache_wrapper, + force_sequential=durable_state is not None, + ) + ) + extra_kwargs["tools"] = resolved + self.log.info( + "Agent tools configured: count=%d names=%s", + len(resolved), + [getattr(tool, "name", type(tool).__name__) for tool in resolved], + ) + + if abstract_items: + processed: list[Any] = list(abstract_items) + if request.enable_tool_logging: + processed = [LoggingToolset(wrapped=ts, logger=self.log) for ts in processed] + if durable_state is not None: + processed = [ + CachingToolset( + wrapped=ts, + storage=durable_state.storage, + counter=durable_state.counter, + ) + for ts in processed + ] + extra_kwargs["toolsets"] = processed + self.log.info( + "Agent abstract toolsets configured: count=%d types=%s", + len(processed), + [type(toolset).__name__ for toolset in processed], + ) + + if isinstance(request.output_type, dict): + raise ValueError( + "PydanticAIHook does not support raw JSON schema mappings for output_type. " + "Pass a Python type, such as a Pydantic BaseModel subclass." + ) + + agent_kwargs: dict[str, Any] = {"instructions": request.instructions, **extra_kwargs} + if request.output_type is not None: + agent_kwargs["output_type"] = request.output_type + + agent = Agent( + self.get_model(), + **agent_kwargs, + ) + if "instrument" not in extra_kwargs: + # Set the public ``agent.instrument`` property rather than the # ``Agent(instrument=...)`` constructor kwarg, which is deprecated in - # current pydantic-ai. Assigning ``agent.instrument`` works across the - # provider's ``pydantic-ai-slim>=1.71`` floor (a plain instance - # attribute on older versions, a property on newer ones). A caller - # that passed its own ``instrument`` wins. + # current pydantic-ai. A caller that passed its own ``instrument`` + # via agent_params wins. settings = genai_instrumentation_settings() if settings is not None: agent.instrument = settings - return agent + return PydanticAgentHandle(agent=agent, durable_state=durable_state) + + def run_agent(self, agent: PydanticAgentHandle, request: AgentRunRequest) -> AgentRunResult: + """Run *agent* synchronously for *request* and return a normalized :class:`~airflow.providers.common.ai.hooks.base.AgentRunResult`.""" + if not isinstance(agent, PydanticAgentHandle): + raise TypeError("PydanticAIHook.run_agent() requires a PydanticAgentHandle from create_agent().") + + native_agent = agent.agent + durable_state = agent.durable_state + if request.durable_context is None and durable_state is not None: + raise ValueError( + "PydanticAIHook.run_agent() received durable state, but request.durable_context is not set." + ) + if request.durable_context is not None and durable_state is None: + raise ValueError("Durable execution requires a PydanticAgentHandle with durable state.") + + run_kwargs: dict[str, Any] = {} + if request.message_history is not None: + run_kwargs["message_history"] = request.message_history + if request.usage_limits is not None: + run_kwargs["usage_limits"] = request.usage_limits + + storage = durable_state.storage if durable_state is not None else None + counter = durable_state.counter if durable_state is not None else None + + if storage is not None and counter is not None: + if native_agent.model is None: + raise ValueError("Agent model must be set when durable=True") + model = native_agent.model + resolved_model = infer_model(model) if isinstance(model, str) else model + caching_model = CachingModel( + resolved_model, + storage=storage, + counter=counter, + ) + with native_agent.override(model=caching_model): + result = native_agent.run_sync(request.prompt, **run_kwargs) + else: + result = native_agent.run_sync(request.prompt, **run_kwargs) + usage = result.usage + tool_names: list[str] = [] + for message in result.all_messages(): + for part in getattr(message, "parts", []): + if isinstance(part, ToolCallPart): + tool_names.append(part.tool_name) + + run_result = AgentRunResult( + output=result.output, + message_history=result.all_messages(), + model_name=getattr(result.response, "model_name", None), + usage=AgentUsage( + requests=usage.requests, + tool_calls=usage.tool_calls, + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, + total_tokens=usage.total_tokens, + ), + tool_names=tool_names or None, + ) + + if counter is not None: + run_result.durable_stats = DurableStats( + replayed_model=counter.replayed_model, + replayed_tool=counter.replayed_tool, + cached_model=counter.cached_model, + cached_tool=counter.cached_tool, + ) + + if storage is not None: + storage.cleanup() + return run_result def test_connection(self) -> tuple[bool, str]: """ @@ -218,7 +402,7 @@ def test_connection(self) -> tuple[bool, str]: connectivity (quotas, billing, rate limits). """ try: - self.get_conn() + self.get_model() return True, "Model resolved successfully." except Exception as e: return False, str(e) diff --git a/providers/common/ai/src/airflow/providers/common/ai/mixins/durable.py b/providers/common/ai/src/airflow/providers/common/ai/mixins/durable.py new file mode 100644 index 0000000000000..34027a9de79a7 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/mixins/durable.py @@ -0,0 +1,97 @@ +# 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. +"""Durable cache helpers for agent hooks.""" + +from __future__ import annotations + +import functools +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.ai.durable.step_counter import DurableStepCounter +from airflow.providers.common.ai.durable.storage import DurableStorage +from airflow.providers.common.ai.utils.callables import is_async_callable + +if TYPE_CHECKING: + from airflow.providers.common.ai.hooks.base import DurableContext + + +@dataclass +class DurableState: + """Durable cache storage and counters for one agent run.""" + + storage: DurableStorage + counter: DurableStepCounter + + +class DurableAgentMixin: + """Reusable durable step-cache helpers for hooks that support durable execution.""" + + def _init_durable(self, ctx: DurableContext) -> DurableState: + """ + Create and return durable state for *ctx*. + + Hooks call this inside ``_build_agent`` when + :attr:`~airflow.providers.common.ai.hooks.base.AgentRunRequest.durable_context` is set. + """ + storage = DurableStorage( + dag_id=ctx.dag_id, + task_id=ctx.task_id, + run_id=ctx.run_id, + map_index=ctx.map_index, + ) + return DurableState(storage=storage, counter=DurableStepCounter()) + + @staticmethod + def _cached_callable( + fn: Callable[..., Any], + storage: DurableStorage, + counter: DurableStepCounter, + ) -> Callable[..., Any]: + """Wrap *fn* to cache its result in *storage* using a monotonic step counter.""" + if is_async_callable(fn): + + @functools.wraps(fn) + async def async_wrapper(*args, **kwargs): + step = counter.next_step() + key = f"tool_step_{step}" + found, cached = storage.load_tool_result(key) + if found: + counter.replayed_tool += 1 + return cached + result = await fn(*args, **kwargs) + storage.save_tool_result(key, result) + counter.cached_tool += 1 + return result + + return async_wrapper + + @functools.wraps(fn) + def sync_wrapper(*args, **kwargs): + step = counter.next_step() + key = f"tool_step_{step}" + found, cached = storage.load_tool_result(key) + if found: + counter.replayed_tool += 1 + return cached + result = fn(*args, **kwargs) + storage.save_tool_result(key, result) + counter.cached_tool += 1 + return result + + return sync_wrapper diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index b41a0c54d8b2f..0eb50735b6e1d 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Operator for running pydantic-ai agents with tools and multi-turn reasoning.""" +"""Operator for running LLM agents with tools and multi-turn reasoning.""" from __future__ import annotations @@ -26,10 +26,11 @@ from pydantic import BaseModel -from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook +from airflow.providers.common.ai.hooks.base import AgentRunRequest, BaseAIHook, DurableContext from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin -from airflow.providers.common.ai.utils.logging import log_run_summary, wrap_toolsets_for_logging +from airflow.providers.common.ai.utils.logging import log_run_summary from airflow.providers.common.ai.utils.output_type import rehydrate_pydantic_output +from airflow.providers.common.ai.utils.validation import reject_sequence_with_unsupported_feature from airflow.providers.common.compat.sdk import ( AirflowOptionalProviderFeatureException, BaseOperator, @@ -47,12 +48,8 @@ _CORE_WALKER = False if TYPE_CHECKING: - from pydantic_ai import Agent - from pydantic_ai.toolsets.abstract import AbstractToolset from pydantic_ai.usage import UsageLimits - from airflow.providers.common.ai.durable.step_counter import DurableStepCounter - from airflow.providers.common.ai.durable.storage import DurableStorage from airflow.providers.common.compat.sdk import TaskInstanceKey from airflow.sdk import Context @@ -92,14 +89,17 @@ def get_link( class AgentOperator(BaseOperator, HITLReviewMixin): """ - Run a pydantic-ai Agent with tools and multi-turn reasoning. + Run an LLM agent with tools and multi-turn reasoning. Provide ``llm_conn_id`` and optional ``toolsets`` to let the operator build and run the agent. The agent reasons about the prompt, calls tools in a multi-turn loop, and returns a final answer. + The agent backend is selected by the connection ``conn_type`` (for example + ``pydanticai``, ``pydanticai-bedrock``, or ``pydanticai-azure``). + :param prompt: The prompt to send to the agent. - :param llm_conn_id: Connection ID for the LLM provider. + :param llm_conn_id: Connection ID for the agent provider. :param model_id: Model identifier (e.g. ``"openai:gpt-5"``). Overrides the model stored in the connection's extra field. :param system_prompt: System-level instructions for the agent. @@ -108,22 +108,25 @@ class AgentOperator(BaseOperator, HITLReviewMixin): returned to XCom unchanged so downstream tasks can type-hint it directly. The class must be defined at module scope -- nested classes cannot be deserialized from XCom. - :param toolsets: List of pydantic-ai toolsets the agent can use - (e.g. ``SQLToolset``, ``HookToolset``). - :param enable_tool_logging: When ``True`` (default), wraps each toolset in a - ``LoggingToolset`` that logs tool calls with timing at INFO level and - arguments at DEBUG level. Set to ``False`` to disable. - :param agent_params: Additional keyword arguments passed to the pydantic-ai - ``Agent`` constructor (e.g. ``retries``, ``model_settings``). - :param usage_limits: Optional pydantic-ai + :param toolsets: List of :class:`~airflow.providers.common.ai.hooks.base.BaseToolset` + instances the agent can use. + :param enable_tool_logging: When ``True`` (default), wraps Airflow-resolved tool callables + with a logging shim that logs calls with timing at INFO level and arguments at DEBUG level. + Backend-native tool objects may be passed through unchanged by the selected hook and might + not receive this wrapper. Set to ``False`` to disable. + :param agent_params: Additional keyword arguments passed to the underlying agent + constructor (e.g. ``retries``, ``model_settings``). + :param usage_limits: Optional :class:`~pydantic_ai.usage.UsageLimits` enforced on every agent run (initial run, durable replay, and HITL regeneration). Pass ``UsageLimits(request_limit=..., total_tokens_limit=..., tool_calls_limit=..., ...)`` to fail the task when the agent exceeds the configured token, request, or tool budget. ``None`` (default) means no enforcement. :param durable: When ``True``, enables step-level caching of model - responses and tool results for durable execution. On retry, cached - steps are replayed instead of re-executing. Default ``False``. + responses and Airflow-resolved tool results for durable execution. On retry, cached + steps are replayed instead of re-executing. Backend-native tool objects may be passed + through unchanged by the selected hook and might not receive tool-result caching. + Default ``False``. Requires ``[common.ai] durable_cache_path`` to be set. **HITL Review parameters** (requires the ``hitl_review`` plugin): @@ -170,7 +173,7 @@ def __init__( model_id: str | None = None, system_prompt: str = "", output_type: type = str, - toolsets: list[AbstractToolset] | None = None, + toolsets: list[Any] | None = None, enable_tool_logging: bool = True, agent_params: dict[str, Any] | None = None, usage_limits: UsageLimits | None = None, @@ -215,86 +218,53 @@ def __init__( ) @cached_property - def llm_hook(self) -> PydanticAIHook: - """Return PydanticAIHook for the configured LLM connection.""" + def llm_hook(self) -> BaseAIHook: + """Return the agent hook for the configured connection (resolved from ``conn_type``).""" hook_params = { "model_id": self.model_id, } - return PydanticAIHook.get_hook(self.llm_conn_id, hook_params=hook_params) - - def _build_agent(self) -> Agent[None, Any]: - """Build and return a pydantic-ai Agent from the operator's config.""" - extra_kwargs = dict(self.agent_params) - if self.toolsets: - toolsets = self.toolsets - if self.durable and self._durable_storage is not None and self._durable_counter is not None: - toolsets = self._build_durable_toolsets( - toolsets, self._durable_storage, self._durable_counter - ) - if self.enable_tool_logging: - toolsets = wrap_toolsets_for_logging(toolsets, self.log) - extra_kwargs["toolsets"] = toolsets - return self.llm_hook.create_agent( - output_type=self.output_type, - instructions=self.system_prompt, - **extra_kwargs, - ) - - def _build_durable_toolsets( - self, toolsets: list[AbstractToolset], storage: DurableStorage, counter: DurableStepCounter - ) -> list[AbstractToolset]: - """Wrap each toolset with CachingToolset for durable execution.""" - from airflow.providers.common.ai.durable.caching_toolset import CachingToolset - - return [CachingToolset(wrapped=ts, storage=storage, counter=counter) for ts in toolsets] - - def execute(self, context: Context) -> Any: - if self.enable_hitl_review and not isinstance(self.prompt, str): - raise TypeError( - f"{type(self).__name__}: enable_hitl_review=True is not supported " - f"with a non-string prompt (got {type(self.prompt).__name__}). " - f"The HITL session model requires a string prompt. Return a str " - f"prompt, or disable enable_hitl_review." - ) - - self._durable_storage = None - self._durable_counter = None - - if self.durable: - from airflow.providers.common.ai.durable.step_counter import DurableStepCounter - from airflow.providers.common.ai.durable.storage import DurableStorage - - ti = context["task_instance"] - self._durable_storage = DurableStorage( + return BaseAIHook.get_agent_hook(self.llm_conn_id, hook_params=hook_params) + + def _build_request(self, *, prompt: str, message_history: Any = None) -> AgentRunRequest: + """Build an :class:`~airflow.providers.common.ai.hooks.base.AgentRunRequest` from operator config.""" + durable_context: DurableContext | None = None + if self.durable and hasattr(self, "_durable_ti") and self._durable_ti is not None: + ti = self._durable_ti + durable_context = DurableContext( dag_id=ti.dag_id, task_id=ti.task_id, run_id=ti.run_id, map_index=ti.map_index if ti.map_index is not None else -1, ) - self._durable_counter = DurableStepCounter() - - agent = self._build_agent() - - storage = self._durable_storage - counter = self._durable_counter - if self.durable and storage is not None and counter is not None: - from pydantic_ai.models import infer_model + return AgentRunRequest( + prompt=prompt, + output_type=self.output_type, + instructions=self.system_prompt, + toolsets=self.toolsets, + usage_limits=self.usage_limits, + message_history=message_history, + enable_tool_logging=self.enable_tool_logging, + durable_context=durable_context, + agent_params=dict(self.agent_params), + ) - from airflow.providers.common.ai.durable.caching_model import CachingModel + def execute(self, context: Context) -> Any: + reject_sequence_with_unsupported_feature( + self.prompt, + decorator_name=type(self).__name__, + feature_name="enable_hitl_review", + feature_enabled=self.enable_hitl_review, + ) + self._durable_ti = context["task_instance"] if self.durable else None - if agent.model is None: - raise ValueError("Agent model must be set when durable=True") - resolved_model = infer_model(agent.model) - caching_model = CachingModel(resolved_model, storage=storage, counter=counter) - with agent.override(model=caching_model): - result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) - else: - result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) + request = self._build_request(prompt=self.prompt) + agent = self.llm_hook.create_agent(request) + run_result = self.llm_hook.run_agent(agent, request) - log_run_summary(self.log, result) + log_run_summary(self.log, run_result) - if self._durable_counter is not None: - c = self._durable_counter + if run_result.durable_stats is not None: + c = run_result.durable_stats replayed = c.replayed_model + c.replayed_tool cached = c.cached_model + c.cached_tool if replayed: @@ -309,16 +279,13 @@ def execute(self, context: Context) -> Any: c.cached_tool, ) - if self._durable_storage is not None: - self._durable_storage.cleanup() - - output = result.output + output = run_result.output if self.enable_hitl_review: result_str = self.run_hitl_review( # type: ignore[misc] context, output, - message_history=result.all_messages(), + message_history=run_result.message_history, ) if isinstance(self.output_type, type) and issubclass(self.output_type, BaseModel): return rehydrate_pydantic_output( @@ -337,12 +304,12 @@ def execute(self, context: Context) -> Any: def regenerate_with_feedback(self, *, feedback: str, message_history: Any) -> tuple[str, Any]: """Re-run the agent with *feedback* appended to the conversation history.""" - agent = self._build_agent() - messages = message_history or [] - result = agent.run_sync(feedback, message_history=messages, usage_limits=self.usage_limits) - log_run_summary(self.log, result) + request = self._build_request(prompt=feedback, message_history=message_history) + agent = self.llm_hook.create_agent(request) + run_result = self.llm_hook.run_agent(agent, request) + log_run_summary(self.log, run_result) - output = result.output + output = run_result.output if isinstance(output, BaseModel): output = output.model_dump_json() - return str(output), result.all_messages() + return str(output), run_result.message_history diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py index c9a22632f2e3d..076ab1b425122 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py @@ -25,7 +25,7 @@ from pydantic import BaseModel -from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook +from airflow.providers.common.ai.hooks.base import AgentRunRequest, BaseAIHook from airflow.providers.common.ai.mixins.approval import LLMApprovalMixin from airflow.providers.common.ai.utils.logging import log_run_summary from airflow.providers.common.ai.utils.output_type import rehydrate_pydantic_output @@ -41,7 +41,6 @@ _CORE_WALKER = False if TYPE_CHECKING: - from pydantic_ai import Agent from pydantic_ai.usage import UsageLimits from airflow.sdk import Context @@ -51,7 +50,7 @@ class LLMOperator(BaseOperator, LLMApprovalMixin): """ Call an LLM with a prompt and return the output. - Uses a :class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook` + Uses a :class:`~airflow.providers.common.ai.hooks.base.BaseAIHook` for LLM access. Supports plain string output (default) and structured output via a Pydantic ``BaseModel``. When ``output_type`` is a ``BaseModel`` subclass, the model instance is returned to XCom unchanged so downstream tasks can @@ -140,19 +139,9 @@ def __init__( self.allow_modifications = allow_modifications @cached_property - def llm_hook(self) -> PydanticAIHook: - """ - Return the correct PydanticAIHook subclass for the configured connection. - - Delegates to :meth:`~PydanticAIHook.get_hook` which looks up - the connection's ``conn_type`` and instantiates the matching subclass - (e.g. :class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook` - for ``pydanticai-azure`` connections). - """ - hook_params = { - "model_id": self.model_id, - } - return PydanticAIHook.get_hook(self.llm_conn_id, hook_params=hook_params) + def llm_hook(self) -> BaseAIHook: + """Return the agent hook for the configured connection.""" + return BaseAIHook.get_agent_hook(self.llm_conn_id, hook_params={"model_id": self.model_id}) def execute(self, context: Context) -> Any: if self.require_approval and not isinstance(self.prompt, str): @@ -163,10 +152,15 @@ def execute(self, context: Context) -> Any: f"str prompt, or disable require_approval." ) - agent: Agent[None, Any] = self.llm_hook.create_agent( - output_type=self.output_type, instructions=self.system_prompt, **self.agent_params + request = AgentRunRequest( + prompt=self.prompt, + output_type=self.output_type, + instructions=self.system_prompt, + usage_limits=self.usage_limits, + agent_params=dict(self.agent_params), ) - result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) + agent = self.llm_hook.create_agent(request) + result = self.llm_hook.run_agent(agent, request) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py index 0395040852f53..93b740c0ab6fb 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py @@ -22,6 +22,7 @@ from enum import Enum from typing import TYPE_CHECKING, Any +from airflow.providers.common.ai.hooks.base import AgentRunRequest from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.utils.logging import log_run_summary from airflow.providers.standard.operators.branch import BranchMixIn @@ -76,12 +77,15 @@ def execute(self, context: Context) -> str | Iterable[str] | None: ) output_type = list[downstream_tasks_enum] if self.allow_multiple_branches else downstream_tasks_enum - agent = self.llm_hook.create_agent( + request = AgentRunRequest( + prompt=self.prompt, output_type=output_type, instructions=self.system_prompt, - **self.agent_params, + usage_limits=self.usage_limits, + agent_params=dict(self.agent_params), ) - result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) + agent = self.llm_hook.create_agent(request) + result = self.llm_hook.run_agent(agent, request) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py index 0c9d1df1bdd20..2fc82b043cdd6 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_file_analysis.py @@ -23,13 +23,12 @@ from pydantic import BaseModel +from airflow.providers.common.ai.hooks.base import AgentRunRequest from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.utils.file_analysis import build_file_analysis_request from airflow.providers.common.ai.utils.logging import log_run_summary if TYPE_CHECKING: - from pydantic_ai import Agent - from airflow.sdk import Context @@ -129,12 +128,15 @@ def execute(self, context: Context) -> Any: self.sample_rows, ) self.log.debug("Resolved file analysis paths: %s", request.resolved_paths) - agent: Agent[None, Any] = self.llm_hook.create_agent( + run_request = AgentRunRequest( + prompt=request.user_content, output_type=self.output_type, instructions=self._build_system_prompt(), - **self.agent_params, + usage_limits=self.usage_limits, + agent_params=dict(self.agent_params), ) - result = agent.run_sync(request.user_content, usage_limits=self.usage_limits) + agent = self.llm_hook.create_agent(run_request) + result = self.llm_hook.run_agent(agent, run_request) log_run_summary(self.log, result) output = result.output diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py index b51ea7c1f7c3b..a2dc69202ca68 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py @@ -25,6 +25,7 @@ from pydantic import BaseModel, Field +from airflow.providers.common.ai.hooks.base import AgentRunRequest from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.utils.logging import log_run_summary from airflow.providers.common.compat.sdk import AirflowException, BaseHook @@ -303,13 +304,16 @@ def execute(self, context: Context) -> dict[str, Any]: full_system_prompt = self._build_system_prompt(schema_context) - agent = self.llm_hook.create_agent( + request = AgentRunRequest( + prompt=self.prompt, output_type=SchemaCompareResult, instructions=full_system_prompt, - **self.agent_params, + usage_limits=self.usage_limits, + agent_params=dict(self.agent_params), ) self.log.info("Running LLM schema comparison...") - result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) + agent = self.llm_hook.create_agent(request) + result = self.llm_hook.run_agent(agent, request) log_run_summary(self.log, result) output_result = result.output.model_dump() diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py index 7342be2b7e34f..9910f3ae209c5 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py @@ -34,6 +34,7 @@ raise AirflowOptionalProviderFeatureException(e) +from airflow.providers.common.ai.hooks.base import AgentRunRequest from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.utils.logging import log_run_summary from airflow.providers.common.compat.sdk import BaseHook @@ -147,10 +148,15 @@ def execute(self, context: Context) -> str: full_system_prompt = self._build_system_prompt(schema_info) - agent = self.llm_hook.create_agent( - output_type=str, instructions=full_system_prompt, **self.agent_params + request = AgentRunRequest( + prompt=self.prompt, + output_type=str, + instructions=full_system_prompt, + usage_limits=self.usage_limits, + agent_params=dict(self.agent_params), ) - result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) + agent = self.llm_hook.create_agent(request) + result = self.llm_hook.run_agent(agent, request) log_run_summary(self.log, result) sql = self._strip_llm_output(result.output) diff --git a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py index f92e4e0d64f90..14e13db2f440d 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py +++ b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py @@ -81,7 +81,7 @@ class LLMRetryPolicy(RetryPolicy): """ Retry policy that uses an LLM to classify errors and decide retry behaviour. - Uses :class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook` + Uses :class:`~airflow.providers.common.ai.hooks.base.BaseAIHook` to call any configured LLM provider (OpenAI, Anthropic, Bedrock, Vertex, Ollama, etc.) for error classification with structured output. @@ -139,13 +139,11 @@ def _classify( try_number: int, max_tries: int, ) -> RetryDecision: - from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook + from pydantic_ai.settings import ModelSettings - hook = PydanticAIHook(llm_conn_id=self.llm_conn_id, model_id=self.model_id) - agent = hook.create_agent( - output_type=ErrorClassification, - instructions=self.instructions, - ) + from airflow.providers.common.ai.hooks.base import AgentRunRequest, BaseAIHook + + hook = BaseAIHook.get_agent_hook(self.llm_conn_id, hook_params={"model_id": self.model_id}) prompt = ( f"Classify this error from a data pipeline task " @@ -153,12 +151,14 @@ def _classify( f"{type(exception).__name__}: {exception}" ) - from pydantic_ai.settings import ModelSettings - - result = agent.run_sync( - prompt, - model_settings=ModelSettings(timeout=self.timeout), + request = AgentRunRequest( + prompt=prompt, + output_type=ErrorClassification, + instructions=self.instructions, + agent_params={"model_settings": ModelSettings(timeout=self.timeout)}, ) + agent = hook.create_agent(request) + result = hook.run_agent(agent, request) classification = result.output log.info( diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/langchain_bridge.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/langchain_bridge.py index 3f5762679ae04..e16f246dded93 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/langchain_bridge.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/langchain_bridge.py @@ -40,12 +40,17 @@ from pydantic_ai.models.test import TestModel from pydantic_ai.usage import RunUsage +from airflow.providers.common.ai.hooks.base import BaseToolset +from airflow.providers.common.ai.utils.callables import is_async_callable + if TYPE_CHECKING: from collections.abc import Coroutine from langchain_core.tools import StructuredTool from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool + from airflow.providers.common.ai.hooks.base import ToolSpec + def _run_coro_sync(coro: Coroutine[Any, Any, Any]) -> Any: """ @@ -66,7 +71,7 @@ def _run_coro_sync(coro: Coroutine[Any, Any, Any]) -> Any: def airflow_toolset_to_langchain_tools( - toolset: AbstractToolset[Any], + toolset: AbstractToolset[Any] | BaseToolset, *, deps: Any = None, ) -> list[StructuredTool]: @@ -119,6 +124,9 @@ def airflow_toolset_to_langchain_tools( raise AirflowOptionalProviderFeatureException(e) + if isinstance(toolset, BaseToolset): + return [_build_structured_tool_from_spec(spec, StructuredTool) for spec in toolset.as_tools()] + # An inert placeholder context. The curated common.ai toolsets ignore it; # TestModel satisfies RunContext's required `model` field without reaching a # real LLM (the bridge never runs the model, only the tools). @@ -170,3 +178,34 @@ async def _async_call(**kwargs: Any) -> Any: description=tool_def.description or name, args_schema=tool_def.parameters_json_schema, ) + + +def _build_structured_tool_from_spec( + spec: ToolSpec, + structured_tool_cls: type[StructuredTool], +) -> StructuredTool: + """Build a single LangChain ``StructuredTool`` from an Airflow ``ToolSpec``.""" + + def _sync_call(**kwargs: Any) -> Any: + try: + if is_async_callable(spec.fn): + return _run_coro_sync(spec.fn(**kwargs)) + return spec.fn(**kwargs) + except ModelRetry as e: + return str(e) + + async def _async_call(**kwargs: Any) -> Any: + try: + if is_async_callable(spec.fn): + return await spec.fn(**kwargs) + return spec.fn(**kwargs) + except ModelRetry as e: + return str(e) + + return structured_tool_cls.from_function( + func=_sync_call, + coroutine=_async_call, + name=spec.name, + description=spec.description or spec.name, + args_schema=spec.parameters, + ) diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py index ee3128705a1f9..75b807ad81a00 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py @@ -21,7 +21,7 @@ import json import sqlite3 from contextlib import suppress -from typing import TYPE_CHECKING, Any +from typing import Any try: from airflow.providers.common.ai.utils.sql_validation import ( @@ -35,18 +35,10 @@ raise AirflowOptionalProviderFeatureException(e) from pydantic_ai.exceptions import ModelRetry -from pydantic_ai.tools import ToolDefinition -from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool -from pydantic_core import SchemaValidator, core_schema +from airflow.providers.common.ai.hooks.base import BaseToolset, ToolSpec from airflow.providers.common.compat.sdk import BaseHook -if TYPE_CHECKING: - from pydantic_ai._run_context import RunContext - -_PASSTHROUGH_VALIDATOR = SchemaValidator(core_schema.any_schema()) - -# JSON Schemas for the four SQL tools. _LIST_TABLES_SCHEMA: dict[str, Any] = { "type": "object", "properties": {}, @@ -102,7 +94,7 @@ _SQLALCHEMY_RETRYABLE_EXCEPTIONS = (_SQLAlchemyProgrammingError,) -class SQLToolset(AbstractToolset[Any]): +class SQLToolset(BaseToolset): """ Curated toolset that gives an LLM agent safe access to a SQL database. @@ -208,50 +200,40 @@ def _get_db_hook(self) -> DbApiHook: return self._hook # ------------------------------------------------------------------ - # AbstractToolset interface + # BaseToolset interface # ------------------------------------------------------------------ - async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]: - tools: dict[str, ToolsetTool[Any]] = {} - - for name, description, schema in ( - ("list_tables", "List available table names in the database.", _LIST_TABLES_SCHEMA), - ("get_schema", "Get column names and types for a table.", _GET_SCHEMA_SCHEMA), - ("query", "Execute a SQL query and return rows as JSON.", _QUERY_SCHEMA), - ("check_query", "Validate SQL syntax without executing it.", _CHECK_QUERY_SCHEMA), - ): - # sequential=True because all tools use a shared DbApiHook with - # synchronous I/O — they must not run concurrently. - tool_def = ToolDefinition( - name=name, - description=description, - parameters_json_schema=schema, + def as_tools(self) -> list[ToolSpec]: + return [ + ToolSpec( + name="list_tables", + description="List available table names in the database.", + parameters=_LIST_TABLES_SCHEMA, + fn=self._list_tables, sequential=True, - ) - tools[name] = ToolsetTool( - toolset=self, - tool_def=tool_def, - max_retries=1, - args_validator=_PASSTHROUGH_VALIDATOR, - ) - return tools - - async def call_tool( - self, - name: str, - tool_args: dict[str, Any], - ctx: RunContext[Any], - tool: ToolsetTool[Any], - ) -> Any: - if name == "list_tables": - return self._list_tables() - if name == "get_schema": - return self._get_schema(tool_args["table_name"]) - if name == "query": - return self._query(tool_args["sql"]) - if name == "check_query": - return self._check_query(tool_args["sql"]) - raise ValueError(f"Unknown tool: {name!r}") + ), + ToolSpec( + name="get_schema", + description="Get column names and types for a table.", + parameters=_GET_SCHEMA_SCHEMA, + fn=self._get_schema, + sequential=True, + ), + ToolSpec( + name="query", + description="Execute a SQL query and return rows as JSON.", + parameters=_QUERY_SCHEMA, + fn=self._query, + sequential=True, + ), + ToolSpec( + name="check_query", + description="Validate SQL syntax without executing it.", + parameters=_CHECK_QUERY_SCHEMA, + fn=self._check_query, + sequential=True, + ), + ] # ------------------------------------------------------------------ # Tool implementations @@ -325,7 +307,6 @@ def _query(self, sql: str) -> str: f"error: {e!s}, Use get_schema and list_tables tools for more details." ) from e raise - # Fetch column names from cursor description. col_names: list[str] | None = None if hook.last_description: col_names = [desc[0] for desc in hook.last_description] diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/callables.py b/providers/common/ai/src/airflow/providers/common/ai/utils/callables.py new file mode 100644 index 0000000000000..edba22bf16e48 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/callables.py @@ -0,0 +1,30 @@ +# 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. +"""Helpers for callable introspection.""" + +from __future__ import annotations + +import functools +import inspect +from collections.abc import Callable +from typing import Any + + +def is_async_callable(fn: Callable[..., Any]) -> bool: + """Return whether *fn* should be wrapped as an async callable.""" + target = fn.func if isinstance(fn, functools.partial) else fn + return inspect.iscoroutinefunction(target) or inspect.iscoroutinefunction(type(target).__call__) diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/function_schema.py b/providers/common/ai/src/airflow/providers/common/ai/utils/function_schema.py new file mode 100644 index 0000000000000..d7836dc786a66 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/function_schema.py @@ -0,0 +1,186 @@ +# 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. +"""Helpers for extracting JSON Schema and tool metadata from plain Python callables.""" + +from __future__ import annotations + +import functools +import inspect +from collections.abc import Callable +from typing import TYPE_CHECKING, Annotated, Any, get_args, get_origin + +from pydantic import Field, create_model +from typing_extensions import get_type_hints + +if TYPE_CHECKING: + from airflow.providers.common.ai.hooks.base import ToolSpec + +_EMPTY_OBJECT_SCHEMA: dict[str, Any] = {"type": "object", "properties": {}} +_SKIP_PARAMS = frozenset({"self", "cls"}) +_DOCSTRING_SECTION_PREFIXES = ( + "args:", + "arguments:", + "parameters:", + "params:", + "returns:", + "return:", + "yields:", + "yield:", + "raises:", + "raise:", + "except:", + "exceptions:", + "example:", + "examples:", + "note:", + "notes:", + "see also:", + "references:", +) + + +def _extract_docstring_summary(obj: Any) -> str: + """Return leading descriptive docstring text before Args/Returns-style sections.""" + doc = inspect.getdoc(obj) + if not doc: + return "" + result: list[str] = [] + for line in doc.split("\n"): + if line.strip().lower().startswith(_DOCSTRING_SECTION_PREFIXES): + break + result.append(line) + return "\n".join(result).strip() + + +def extract_function_description(fn: Callable[..., Any]) -> str: + """Return a short description for *fn* from its leading docstring text.""" + # Unwrap partials to get the underlying function's docstring. + if isinstance(fn, functools.partial): + return extract_function_description(fn.func) + + # Callable objects (class instances) have no __name__. + # Prefer __call__ docstring (what calling does), then class docstring, then class name. + if not hasattr(fn, "__name__"): + return ( + _extract_docstring_summary(type(fn).__call__) + or _extract_docstring_summary(fn) + or type(fn).__name__ + ) + + return _extract_docstring_summary(fn) or fn.__name__ + + +def build_function_json_schema(fn: Callable[..., Any]) -> dict[str, Any]: + """ + Build a JSON Schema ``object`` for the parameters of *fn*. + + Reads type hints (including ``Annotated[T, "description"]``) and default + values to produce a schema suitable for LLM tool binding. + Falls back to an empty object schema on any introspection failure. + + ``self``, ``cls``, ``*args``, and ``**kwargs`` are excluded. + Positional-only params are rejected because tool callables must accept + keyword arguments matching the generated schema. + For ``functools.partial``, only the remaining free parameters appear. + """ + # Partials: sig from partial (bound args already removed), hints from inner fn. + hint_source: Callable[..., Any] = fn + if isinstance(fn, functools.partial): + hint_source = fn.func + while isinstance(hint_source, functools.partial): + hint_source = hint_source.func + + try: + sig = inspect.signature(fn) + except (ValueError, TypeError): + return _EMPTY_OBJECT_SCHEMA + + try: + hints = get_type_hints(hint_source, include_extras=True) + except Exception: + hints = {} + + field_defs: dict[str, Any] = {} + for param_name, param in sig.parameters.items(): + if param_name in _SKIP_PARAMS: + continue + if param.kind is inspect.Parameter.POSITIONAL_ONLY: + # Auto-generated tool schemas describe named JSON object fields, + # and tool frameworks invoke the callable with keyword arguments + # derived from those fields. A positional-only parameter cannot be + # satisfied by that contract, so fail fast with a clear error. + name = getattr(fn, "__name__", type(fn).__name__) + raise ValueError( + f"Cannot build a tool schema for {name}: " + f"parameter {param_name!r} is positional-only. " + "Tool parameters must be callable by keyword." + ) + if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + continue + + annotation = hints.get(param_name, param.annotation) + if annotation is inspect.Parameter.empty: + annotation = Any + default = ... if param.default is inspect.Parameter.empty else param.default + + if get_origin(annotation) is Annotated: + type_args = get_args(annotation) + actual_type = type_args[0] + desc: str | None = next((a for a in type_args[1:] if isinstance(a, str)), None) + else: + actual_type = annotation + desc = None + + field_defs[param_name] = (actual_type, Field(default=default, description=desc)) + + if not field_defs: + return _EMPTY_OBJECT_SCHEMA + + try: + schema = create_model(f"_{getattr(fn, '__name__', 'tool')}", **field_defs).model_json_schema() + except Exception: + return _EMPTY_OBJECT_SCHEMA + + schema.pop("title", None) + schema.pop("additionalProperties", None) + for prop in schema.get("properties", {}).values(): + prop.pop("title", None) + return schema + + +def callable_to_tool_spec(fn: Callable[..., Any]) -> ToolSpec: + """ + Build a :class:`~airflow.providers.common.ai.hooks.base.ToolSpec` from a plain callable. + + Combines :func:`extract_function_description` and :func:`build_function_json_schema` + so callers get name, description, and a full parameter schema in one call. + """ + # Lazy import avoids a circular dependency: base imports this module, + # this module imports ToolSpec from base. + from airflow.providers.common.ai.hooks.base import ToolSpec + + inner: Callable[..., Any] = fn + while isinstance(inner, functools.partial): + inner = inner.func + name = getattr(inner, "__name__", type(inner).__name__) + + return ToolSpec( + name=name, + description=extract_function_description(fn), + parameters=build_function_json_schema(fn), + fn=fn, + ) diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py index 47cabf7ce6ccc..ffe45eb06e4f3 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py @@ -14,44 +14,40 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Logging utilities for pydantic-ai agent runs.""" +"""Logging utilities for agent runs.""" from __future__ import annotations import logging from typing import TYPE_CHECKING, Any -from pydantic_ai.messages import ToolCallPart - -from airflow.providers.common.ai.toolsets.logging import LoggingToolset - if TYPE_CHECKING: - from pydantic_ai.result import AgentRunResult - from pydantic_ai.toolsets.abstract import AbstractToolset - + from airflow.providers.common.ai.hooks.base import AgentRunResult from airflow.sdk.types import Logger _MAX_OUTPUT_LEN = 500 -def log_run_summary(logger: Logger | logging.Logger, result: AgentRunResult[Any]) -> None: +def log_run_summary(logger: Logger | logging.Logger, result: AgentRunResult) -> None: """Log model name, token usage, and tool call sequence from an agent run.""" - usage = result.usage() - model_name = getattr(result.response, "model_name", "unknown") - logger.info( - "::group::LLM run complete: model=%s, requests=%s, tool_calls=%s, " - "input_tokens=%s, output_tokens=%s, total_tokens=%s", - model_name, - usage.requests, - usage.tool_calls, - usage.input_tokens, - usage.output_tokens, - usage.total_tokens, - ) + model_name = result.model_name or "unknown" + usage = result.usage + if usage is not None: + logger.info( + "::group::LLM run complete: model=%s, requests=%s, tool_calls=%s, " + "input_tokens=%s, output_tokens=%s, total_tokens=%s", + model_name, + usage.requests, + usage.tool_calls, + usage.input_tokens, + usage.output_tokens, + usage.total_tokens, + ) + else: + logger.info("::group::LLM run complete: model=%s", model_name) - tool_names = _extract_tool_sequence(result) - if tool_names: - logger.info("Tool call sequence: %s", " -> ".join(tool_names)) + if result.tool_names: + logger.info("Tool call sequence: %s", " -> ".join(result.tool_names)) _log_output_debug(logger, result.output) logger.info("::endgroup::") @@ -70,21 +66,3 @@ def _log_output_debug(logger: Logger | logging.Logger, output: Any) -> None: if len(text) > _MAX_OUTPUT_LEN: text = text[:_MAX_OUTPUT_LEN] + "..." logger.debug("Output: %s", text) - - -def _extract_tool_sequence(result: AgentRunResult[Any]) -> list[str]: - """Extract ordered tool names from the message history.""" - tool_names: list[str] = [] - for message in result.all_messages(): - for part in getattr(message, "parts", []): - if isinstance(part, ToolCallPart): - tool_names.append(part.tool_name) - return tool_names - - -def wrap_toolsets_for_logging( - toolsets: list[AbstractToolset[Any]], - logger: Logger | logging.Logger, -) -> list[AbstractToolset[Any]]: - """Wrap each toolset in a LoggingToolset.""" - return [LoggingToolset(wrapped=ts, logger=logger) for ts in toolsets] diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py index 25a176e3297a6..84ee150839a6c 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py @@ -23,7 +23,7 @@ from pydantic_ai.messages import ImageUrl from airflow.providers.common.ai.decorators.agent import _AgentDecoratedOperator -from airflow.providers.common.ai.toolsets.logging import LoggingToolset +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook, Capability try: from airflow.sdk.serde import SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as _CORE_WALKER @@ -40,38 +40,37 @@ class Summary(BaseModel): text: str -def _make_mock_run_result(output): - """Create a mock AgentRunResult compatible with log_run_summary.""" - mock_result = MagicMock() - mock_result.output = output - mock_result.usage.return_value = MagicMock( - requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0 - ) - mock_result.response = MagicMock(model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result +def _make_run_result(output): + return AgentRunResult(output=output, model_name="test-model", usage=AgentUsage(requests=1)) + + +def _make_mock_hook(run_result): + mock_hook = MagicMock(spec=BaseAIHook) + mock_hook.capabilities = frozenset({Capability.TOOLSETS, Capability.USAGE_LIMITS}) + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = run_result + return mock_hook class TestAgentDecoratedOperator: def test_custom_operator_name(self): assert _AgentDecoratedOperator.custom_operator_name == "@task.agent" - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_execute_calls_callable_and_returns_output(self, mock_hook_cls): - """The callable's return value becomes the agent prompt.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("The top customer is Acme Corp.") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + def test_execute_calls_callable_and_returns_output(self): + mock_hook = _make_mock_hook(_make_run_result("The top customer is Acme Corp.")) def my_prompt(): return "Who is our top customer?" op = _AgentDecoratedOperator(task_id="test", python_callable=my_prompt, llm_conn_id="my_llm") - result = op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + result = op.execute(context={}) assert result == "The top customer is Acme Corp." assert op.prompt == "Who is our top customer?" - mock_agent.run_sync.assert_called_once_with("Who is our top customer?", usage_limits=None) + + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == "Who is our top customer?" @pytest.mark.parametrize( "return_value", @@ -88,36 +87,31 @@ def test_execute_raises_on_invalid_prompt(self, return_value): with pytest.raises(TypeError, match="must be"): op.execute(context={}) - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_execute_accepts_sequence_prompt(self, mock_hook_cls): - """A non-empty Sequence[UserContent] return value is forwarded to run_sync as-is.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("ok") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + def test_execute_accepts_sequence_prompt(self): + """A non-empty Sequence[UserContent] return value is forwarded as-is.""" image = ImageUrl(url="https://example.com/x.png") prompt = ["Describe this:", image] + mock_hook = _make_mock_hook(_make_run_result("ok")) def my_prompt(): return prompt op = _AgentDecoratedOperator(task_id="test", python_callable=my_prompt, llm_conn_id="my_llm") - op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={}) assert op.prompt == prompt - mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None) + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == prompt - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_sequence_prompt_with_hitl_review_raises_before_run_sync(self, mock_hook_cls): + def test_sequence_prompt_with_hitl_review_raises(self): """Sequence prompt + enable_hitl_review=True fails before the agent runs.""" from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS if not AIRFLOW_V_3_1_PLUS: pytest.skip("enable_hitl_review requires Airflow >= 3.1.0") - mock_agent = MagicMock(spec=["run_sync"]) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - op = _AgentDecoratedOperator( task_id="test", python_callable=lambda: ["x", ImageUrl(url="https://example.com/x.png")], @@ -127,14 +121,8 @@ def test_sequence_prompt_with_hitl_review_raises_before_run_sync(self, mock_hook with pytest.raises(TypeError, match="enable_hitl_review=True"): op.execute(context={}) - mock_agent.run_sync.assert_not_called() - - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_execute_merges_op_kwargs_into_callable(self, mock_hook_cls): - """op_kwargs are resolved by the callable to build the prompt.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("done") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + def test_execute_merges_op_kwargs_into_callable(self): + mock_hook = _make_mock_hook(_make_run_result("done")) def my_prompt(topic): return f"Analyze {topic}" @@ -145,18 +133,16 @@ def my_prompt(topic): llm_conn_id="my_llm", op_kwargs={"topic": "revenue trends"}, ) - op.execute(context={"task_instance": MagicMock()}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={"task_instance": MagicMock()}) assert op.prompt == "Analyze revenue trends" - mock_agent.run_sync.assert_called_once_with("Analyze revenue trends", usage_limits=None) - - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_execute_passes_toolsets_through(self, mock_hook_cls): - """Toolsets passed to the decorator are forwarded to the agent.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("result") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == "Analyze revenue trends" + def test_execute_passes_toolsets_through(self): + """Toolsets passed to the decorator are forwarded verbatim in AgentRunRequest.""" + mock_hook = _make_mock_hook(_make_run_result("result")) mock_toolset = MagicMock() op = _AgentDecoratedOperator( @@ -165,21 +151,16 @@ def test_execute_passes_toolsets_through(self, mock_hook_cls): llm_conn_id="my_llm", toolsets=[mock_toolset], ) - op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={}) - create_call = mock_hook_cls.get_hook.return_value.create_agent.call_args - passed_toolsets = create_call[1]["toolsets"] - assert len(passed_toolsets) == 1 - assert isinstance(passed_toolsets[0], LoggingToolset) - assert passed_toolsets[0].wrapped is mock_toolset + request = mock_hook.create_agent.call_args[0][0] + assert request.toolsets == [mock_toolset] @requires_typed_xcom - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_execute_structured_output(self, mock_hook_cls): + def test_execute_structured_output(self): """BaseModel output flows through XCom as the Pydantic instance.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(Summary(text="Great results")) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = _make_mock_hook(_make_run_result(Summary(text="Great results"))) op = _AgentDecoratedOperator( task_id="test", @@ -187,13 +168,13 @@ def test_execute_structured_output(self, mock_hook_cls): llm_conn_id="my_llm", output_type=Summary, ) - result = op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + result = op.execute(context={}) assert isinstance(result, Summary) assert result.text == "Great results" def test_durable_kwarg_passes_through_to_operator(self): - """durable=True is forwarded to AgentOperator via **kwargs.""" op = _AgentDecoratedOperator( task_id="test", python_callable=lambda: "prompt", @@ -203,7 +184,6 @@ def test_durable_kwarg_passes_through_to_operator(self): assert op.durable is True def test_durable_default_false_through_decorator(self): - """durable defaults to False when not specified.""" op = _AgentDecoratedOperator( task_id="test", python_callable=lambda: "prompt", diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py index 67ea067d160c0..262028fcc60a0 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm.py @@ -22,40 +22,43 @@ from pydantic_ai.messages import ImageUrl from airflow.providers.common.ai.decorators.llm import _LLMDecoratedOperator +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook -def _make_mock_run_result(output): - """Create a mock AgentRunResult compatible with log_run_summary.""" - mock_result = MagicMock() - mock_result.output = output - mock_result.usage.return_value = MagicMock( - requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0 +def _make_run_result(output): + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1), ) - mock_result.response = MagicMock(model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result + + +def _make_mock_hook(run_result): + mock_hook = MagicMock() + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = run_result + return mock_hook class TestLLMDecoratedOperator: def test_custom_operator_name(self): assert _LLMDecoratedOperator.custom_operator_name == "@task.llm" - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_calls_callable_and_returns_output(self, mock_hook_cls): + def test_execute_calls_callable_and_returns_output(self): """The callable's return value becomes the LLM prompt.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("This is a summary.") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = _make_mock_hook(_make_run_result("This is a summary.")) def my_prompt(): return "Summarize this text" op = _LLMDecoratedOperator(task_id="test", python_callable=my_prompt, llm_conn_id="my_llm") - result = op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + result = op.execute(context={}) assert result == "This is a summary." assert op.prompt == "Summarize this text" - mock_agent.run_sync.assert_called_once_with("Summarize this text", usage_limits=None) + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == "Summarize this text" @pytest.mark.parametrize( "return_value", @@ -72,31 +75,25 @@ def test_execute_raises_on_invalid_prompt(self, return_value): with pytest.raises(TypeError, match="must be"): op.execute(context={}) - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_accepts_sequence_prompt(self, mock_hook_cls): - """A non-empty Sequence[UserContent] return value is forwarded to run_sync as-is.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("ok") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - + def test_execute_accepts_sequence_prompt(self): + """A non-empty Sequence[UserContent] return value is forwarded as-is.""" image = ImageUrl(url="https://example.com/x.png") prompt = ["Describe this:", image] + mock_hook = _make_mock_hook(_make_run_result("ok")) def my_prompt(): return prompt op = _LLMDecoratedOperator(task_id="test", python_callable=my_prompt, llm_conn_id="my_llm") - op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={}) assert op.prompt == prompt - mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None) + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == prompt - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_sequence_prompt_with_require_approval_raises_before_run_sync(self, mock_hook_cls): + def test_sequence_prompt_with_require_approval_raises(self): """Sequence prompt + require_approval=True fails before the agent runs.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - op = _LLMDecoratedOperator( task_id="test", python_callable=lambda: ["x", ImageUrl(url="https://example.com/x.png")], @@ -106,14 +103,9 @@ def test_sequence_prompt_with_require_approval_raises_before_run_sync(self, mock with pytest.raises(TypeError, match="require_approval=True"): op.execute(context={}) - mock_agent.run_sync.assert_not_called() - - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_merges_op_kwargs_into_callable(self, mock_hook_cls): + def test_execute_merges_op_kwargs_into_callable(self): """op_kwargs are resolved by the callable to build the prompt.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("done") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = _make_mock_hook(_make_run_result("done")) def my_prompt(topic): return f"Summarize {topic}" @@ -124,7 +116,7 @@ def my_prompt(topic): llm_conn_id="my_llm", op_kwargs={"topic": "quantum computing"}, ) - op.execute(context={"task_instance": MagicMock()}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={"task_instance": MagicMock()}) assert op.prompt == "Summarize quantum computing" - mock_agent.run_sync.assert_called_once_with("Summarize quantum computing", usage_limits=None) diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py index 023af790d36da..a39780fbb094d 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_branch.py @@ -23,19 +23,23 @@ from pydantic_ai.messages import ImageUrl from airflow.providers.common.ai.decorators.llm_branch import _LLMBranchDecoratedOperator +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator -def _make_mock_run_result(output): - """Create a mock AgentRunResult compatible with log_run_summary.""" - mock_result = MagicMock() - mock_result.output = output - mock_result.usage.return_value = MagicMock( - requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0 +def _make_run_result(output): + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1), ) - mock_result.response = MagicMock(model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result + + +def _make_mock_hook(run_result): + mock_hook = MagicMock() + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = run_result + return mock_hook class TestLLMBranchDecoratedOperator: @@ -43,14 +47,10 @@ def test_custom_operator_name(self): assert _LLMBranchDecoratedOperator.custom_operator_name == "@task.llm_branch" @patch.object(LLMBranchOperator, "do_branch") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_calls_callable_and_branches(self, mock_hook_cls, mock_do_branch): + def test_execute_calls_callable_and_branches(self, mock_do_branch): """The callable's return value becomes the LLM prompt, LLM output goes through do_branch.""" downstream_enum = Enum("DownstreamTasks", {"positive": "positive", "negative": "negative"}) - - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(downstream_enum.positive) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = _make_mock_hook(_make_run_result(downstream_enum.positive)) mock_do_branch.return_value = "positive" def my_prompt(): @@ -63,11 +63,11 @@ def my_prompt(): ) op.downstream_task_ids = {"positive", "negative"} - result = op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + result = op.execute(context={}) assert result == "positive" assert op.prompt == "Route this review" - mock_agent.run_sync.assert_called_once_with("Route this review", usage_limits=None) mock_do_branch.assert_called_once() @pytest.mark.parametrize( @@ -86,18 +86,13 @@ def test_execute_raises_on_invalid_prompt(self, return_value): op.execute(context={}) @patch.object(LLMBranchOperator, "do_branch") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_accepts_sequence_prompt(self, mock_hook_cls, mock_do_branch): - """A non-empty Sequence[UserContent] return value is forwarded to run_sync as-is.""" + def test_execute_accepts_sequence_prompt(self, mock_do_branch): + """A non-empty Sequence[UserContent] return value is forwarded as-is.""" downstream_enum = Enum("DownstreamTasks", {"positive": "positive"}) - - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(downstream_enum.positive) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - mock_do_branch.return_value = "positive" - image = ImageUrl(url="https://example.com/x.png") prompt = ["Route based on this image:", image] + mock_hook = _make_mock_hook(_make_run_result(downstream_enum.positive)) + mock_do_branch.return_value = "positive" def my_prompt(): return prompt @@ -108,20 +103,19 @@ def my_prompt(): llm_conn_id="my_llm", ) op.downstream_task_ids = {"positive"} - op.execute(context={}) + + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={}) assert op.prompt == prompt - mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None) + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == prompt @patch.object(LLMBranchOperator, "do_branch") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_merges_op_kwargs_into_callable(self, mock_hook_cls, mock_do_branch): + def test_execute_merges_op_kwargs_into_callable(self, mock_do_branch): """op_kwargs are resolved by the callable to build the prompt.""" downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a"}) - - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(downstream_enum.task_a) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = _make_mock_hook(_make_run_result(downstream_enum.task_a)) def my_prompt(ticket_type): return f"Route this {ticket_type} ticket" @@ -134,6 +128,7 @@ def my_prompt(ticket_type): ) op.downstream_task_ids = {"task_a"} - op.execute(context={"task_instance": MagicMock()}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={"task_instance": MagicMock()}) assert op.prompt == "Route this billing ticket" diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py index 41fc750d5051c..0a1e5e79b0c62 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_file_analysis.py @@ -21,42 +21,39 @@ import pytest from airflow.providers.common.ai.decorators.llm_file_analysis import _LLMFileAnalysisDecoratedOperator +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.utils.file_analysis import FileAnalysisRequest -def _make_mock_run_result(output): - mock_result = MagicMock(spec=["output", "usage", "response", "all_messages"]) - mock_result.output = output - mock_result.usage.return_value = MagicMock( - spec=["requests", "tool_calls", "input_tokens", "output_tokens", "total_tokens"], - requests=1, - tool_calls=0, - input_tokens=0, - output_tokens=0, - total_tokens=0, +def _make_run_result(output): + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1), ) - mock_result.response = MagicMock(spec=["model_name"], model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result + + +def _make_mock_hook(run_result): + mock_hook = MagicMock() + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = run_result + return mock_hook class TestLLMFileAnalysisDecoratedOperator: def test_custom_operator_name(self): assert _LLMFileAnalysisDecoratedOperator.custom_operator_name == "@task.llm_file_analysis" - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) @patch( "airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request", autospec=True ) - def test_execute_calls_callable_and_returns_output(self, mock_build_request, mock_hook_cls): + def test_execute_calls_callable_and_returns_output(self, mock_build_request): mock_build_request.return_value = FileAnalysisRequest( user_content="prepared prompt", resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("This is a summary.") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = _make_mock_hook(_make_run_result("This is a summary.")) def my_prompt(): return "Summarize this text" @@ -67,11 +64,13 @@ def my_prompt(): llm_conn_id="my_llm", file_path="/tmp/app.log", ) - result = op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + result = op.execute(context={}) assert result == "This is a summary." assert op.prompt == "Summarize this text" - mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None) + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == "prepared prompt" @pytest.mark.parametrize( "return_value", @@ -88,19 +87,16 @@ def test_execute_raises_on_invalid_prompt(self, return_value): with pytest.raises(TypeError, match="non-empty string"): op.execute(context={}) - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) @patch( "airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request", autospec=True ) - def test_execute_merges_op_kwargs_into_callable(self, mock_build_request, mock_hook_cls): + def test_execute_merges_op_kwargs_into_callable(self, mock_build_request): mock_build_request.return_value = FileAnalysisRequest( user_content="prepared prompt", resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("done") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = _make_mock_hook(_make_run_result("done")) def my_prompt(topic): return f"Summarize {topic}" @@ -112,7 +108,7 @@ def my_prompt(topic): file_path="/tmp/app.log", op_kwargs={"topic": "system logs"}, ) - op.execute(context={"task_instance": MagicMock(spec=["task_id"])}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={"task_instance": MagicMock(spec=["task_id"])}) assert op.prompt == "Summarize system logs" - mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None) diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py index df3c3f571c2e7..856ecee55c617 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py @@ -22,22 +22,26 @@ from pydantic_ai.messages import ImageUrl from airflow.providers.common.ai.decorators.llm_schema_compare import _LLMSchemaCompareDecoratedOperator +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.operators.llm_schema_compare import ( LLMSchemaCompareOperator, SchemaCompareResult, ) -def _make_mock_run_result(output): - """Create a mock AgentRunResult compatible with log_run_summary.""" - mock_result = MagicMock() - mock_result.output = output - mock_result.usage.return_value = MagicMock( - requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0 +def _make_run_result(output): + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1), ) - mock_result.response = MagicMock(model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result + + +def _make_mock_hook(run_result): + mock_hook = MagicMock() + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = run_result + return mock_hook def _make_compare_result(): @@ -48,23 +52,14 @@ def _make_compare_result(): ) -def _make_mock_agent(output: SchemaCompareResult): - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(output) - return mock_agent - - class TestLLMSchemaCompareDecoratedOperator: def test_custom_operator_name(self): assert _LLMSchemaCompareDecoratedOperator.custom_operator_name == "@task.llm_schema_compare" - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) @patch.object(LLMSchemaCompareOperator, "_build_schema_context", return_value="mocked schema") - def test_execute_calls_callable_and_uses_result_as_prompt(self, mock_build_ctx, mock_hook_cls): + def test_execute_calls_callable_and_uses_result_as_prompt(self, mock_build_ctx): """The user's callable return value becomes the LLM prompt.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent( - _make_compare_result() - ) + mock_hook = _make_mock_hook(_make_run_result(_make_compare_result())) def my_prompt_fn(): return "Compare schemas and flag breaking changes" @@ -76,7 +71,8 @@ def my_prompt_fn(): db_conn_ids=["postgres_default", "snowflake_default"], table_names=["test_table"], ) - result = op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + result = op.execute(context={}) assert result["compatible"] is True assert op.prompt == "Compare schemas and flag breaking changes" @@ -98,15 +94,12 @@ def test_execute_raises_on_invalid_prompt(self, return_value): with pytest.raises(TypeError, match="must be"): op.execute(context={}) - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) @patch.object(LLMSchemaCompareOperator, "_build_schema_context", return_value="mocked schema") - def test_execute_accepts_sequence_prompt(self, mock_build_ctx, mock_hook_cls): - """A non-empty Sequence[UserContent] return value is forwarded to run_sync as-is.""" - mock_agent = _make_mock_agent(_make_compare_result()) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - + def test_execute_accepts_sequence_prompt(self, mock_build_ctx): + """A non-empty Sequence[UserContent] return value is forwarded as-is.""" image = ImageUrl(url="https://example.com/x.png") prompt = ["Compare these schemas:", image] + mock_hook = _make_mock_hook(_make_run_result(_make_compare_result())) def my_prompt_fn(): return prompt @@ -118,19 +111,17 @@ def my_prompt_fn(): db_conn_ids=["postgres_default", "snowflake_default"], table_names=["test_table"], ) - op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={}) assert op.prompt == prompt - forwarded_prompt = mock_agent.run_sync.call_args[0][0] - assert forwarded_prompt == prompt + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == prompt - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) @patch.object(LLMSchemaCompareOperator, "_build_schema_context", return_value="mocked schema") - def test_execute_merges_op_kwargs_into_callable(self, mock_build_ctx, mock_hook_cls): + def test_execute_merges_op_kwargs_into_callable(self, mock_build_ctx): """op_kwargs are resolved by the callable to build the prompt.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent( - _make_compare_result() - ) + mock_hook = _make_mock_hook(_make_run_result(_make_compare_result())) def my_prompt_fn(target_env): return f"Compare schemas for {target_env} environment" @@ -143,6 +134,7 @@ def my_prompt_fn(target_env): db_conn_ids=["postgres_default", "snowflake_default"], table_names=["test_table"], ) - op.execute(context={"task_instance": MagicMock()}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={"task_instance": MagicMock()}) assert op.prompt == "Compare schemas for production environment" diff --git a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py index 5b2e4b6e6e3bf..bdabf5abfa742 100644 --- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py +++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py @@ -22,40 +22,43 @@ from pydantic_ai.messages import ImageUrl from airflow.providers.common.ai.decorators.llm_sql import _LLMSQLDecoratedOperator +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook -def _make_mock_run_result(output): - """Create a mock AgentRunResult compatible with log_run_summary.""" - mock_result = MagicMock() - mock_result.output = output - mock_result.usage.return_value = MagicMock( - requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0 +def _make_run_result(output): + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1), ) - mock_result.response = MagicMock(model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result + + +def _make_mock_hook(run_result): + mock_hook = MagicMock() + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = run_result + return mock_hook class TestLLMSQLDecoratedOperator: def test_custom_operator_name(self): assert _LLMSQLDecoratedOperator.custom_operator_name == "@task.llm_sql" - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_calls_callable_and_uses_result_as_prompt(self, mock_hook_cls): + def test_execute_calls_callable_and_uses_result_as_prompt(self): """The user's callable return value becomes the LLM prompt.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("SELECT 1") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = _make_mock_hook(_make_run_result("SELECT 1")) def my_prompt_fn(): return "Get all users" op = _LLMSQLDecoratedOperator(task_id="test", python_callable=my_prompt_fn, llm_conn_id="my_llm") - result = op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + result = op.execute(context={}) assert result == "SELECT 1" assert op.prompt == "Get all users" - mock_agent.run_sync.assert_called_once_with("Get all users", usage_limits=None) + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == "Get all users" @pytest.mark.parametrize( "return_value", @@ -72,31 +75,25 @@ def test_execute_raises_on_invalid_prompt(self, return_value): with pytest.raises(TypeError, match="must be"): op.execute(context={}) - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_accepts_sequence_prompt(self, mock_hook_cls): - """A non-empty Sequence[UserContent] return value is forwarded to run_sync as-is.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("SELECT 1") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - + def test_execute_accepts_sequence_prompt(self): + """A non-empty Sequence[UserContent] return value is forwarded as-is.""" image = ImageUrl(url="https://example.com/x.png") prompt = ["Write SQL for this diagram:", image] + mock_hook = _make_mock_hook(_make_run_result("SELECT 1")) def my_prompt_fn(): return prompt op = _LLMSQLDecoratedOperator(task_id="test", python_callable=my_prompt_fn, llm_conn_id="my_llm") - op.execute(context={}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={}) assert op.prompt == prompt - mock_agent.run_sync.assert_called_once_with(prompt, usage_limits=None) + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == prompt - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_sequence_prompt_with_require_approval_raises_before_run_sync(self, mock_hook_cls): + def test_sequence_prompt_with_require_approval_raises(self): """Sequence prompt + require_approval=True fails before the agent runs.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - op = _LLMSQLDecoratedOperator( task_id="test", python_callable=lambda: ["x", ImageUrl(url="https://example.com/x.png")], @@ -106,14 +103,9 @@ def test_sequence_prompt_with_require_approval_raises_before_run_sync(self, mock with pytest.raises(TypeError, match="require_approval=True"): op.execute(context={}) - mock_agent.run_sync.assert_not_called() - - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_merges_op_kwargs_into_callable(self, mock_hook_cls): + def test_execute_merges_op_kwargs_into_callable(self): """op_kwargs are resolved by the callable to build the prompt.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("SELECT 1") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = _make_mock_hook(_make_run_result("SELECT 1")) def my_prompt_fn(table_name): return f"Get all rows from {table_name}" @@ -124,6 +116,7 @@ def my_prompt_fn(table_name): llm_conn_id="my_llm", op_kwargs={"table_name": "users"}, ) - op.execute(context={"task_instance": MagicMock()}) + with patch.object(BaseAIHook, "get_agent_hook", return_value=mock_hook): + op.execute(context={"task_instance": MagicMock()}) assert op.prompt == "Get all rows from users" diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_base.py b/providers/common/ai/tests/unit/common/ai/hooks/test_base.py new file mode 100644 index 0000000000000..94831ff89809e --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_base.py @@ -0,0 +1,733 @@ +# 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. +from __future__ import annotations + +import asyncio +import functools +import inspect +from unittest.mock import MagicMock, patch + +import pytest + +from airflow.providers.common.ai.hooks.base import ( + AgentRunRequest, + AgentRunResult, + AgentUsage, + BaseAIHook, + BaseToolset, + Capability, + DurableContext, + DurableStats, + ToolSpec, +) +from airflow.providers.common.compat.sdk import BaseHook + + +class TestBaseAIHookGetAgentHook: + @patch("airflow.providers.common.ai.hooks.base.BaseHook.get_hook", autospec=True) + def test_returns_hook_when_instance_is_base_hook(self, mock_get_hook): + mock_hook = MagicMock(spec=BaseAIHook) + mock_get_hook.return_value = mock_hook + + result = BaseAIHook.get_agent_hook("my_conn") + + assert result is mock_hook + mock_get_hook.assert_called_once_with("my_conn", hook_params=None) + + @patch("airflow.providers.common.ai.hooks.base.BaseHook.get_hook", autospec=True) + def test_raises_when_hook_is_not_base_hook(self, mock_get_hook): + mock_get_hook.return_value = MagicMock(spec=BaseHook) + + with pytest.raises(TypeError, match="not a BaseAIHook"): + BaseAIHook.get_agent_hook("my_conn") + + +class TestBaseAIHookInit: + def test_stores_model_id_and_conn_id(self): + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return spec.fn + + hook = ConcreteHook(llm_conn_id="my_conn", model_id="openai:gpt-5") + assert hook.llm_conn_id == "my_conn" + assert hook.model_id == "openai:gpt-5" + + +class TestValidateRunRequest: + def test_rejects_toolsets_when_unsupported(self): + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + capabilities = frozenset({Capability.USAGE_LIMITS, Capability.DURABLE}) + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return spec.fn + + hook = ConcreteHook(llm_conn_id="test_conn") + request = AgentRunRequest(prompt="hi", toolsets=[MagicMock()]) + with pytest.raises(ValueError, match="toolsets not supported"): + hook.validate_run_request(request) + + def test_rejects_usage_limits_when_unsupported(self): + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + capabilities = frozenset({Capability.TOOLSETS, Capability.DURABLE}) + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return spec.fn + + hook = ConcreteHook(llm_conn_id="test_conn") + request = AgentRunRequest(prompt="hi", usage_limits=MagicMock()) + with pytest.raises(ValueError, match="usage_limits not supported"): + hook.validate_run_request(request) + + def test_rejects_durable_when_unsupported(self): + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + capabilities = frozenset({Capability.TOOLSETS, Capability.USAGE_LIMITS}) + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return spec.fn + + hook = ConcreteHook(llm_conn_id="test_conn") + request = AgentRunRequest( + prompt="hi", + durable_context=DurableContext(dag_id="d", task_id="t", run_id="r"), + ) + + with pytest.raises(ValueError, match="durable execution not supported"): + hook.validate_run_request(request) + + def test_create_agent_validates_before_building(self): + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + capabilities = frozenset() # no capabilities — toolsets will be rejected + + def __init__(self): + super().__init__(llm_conn_id="test_conn") + self.built = False + + def get_model(self): + return None + + def _build_agent(self, request): + self.built = True + return "agent" + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return spec.fn + + hook = ConcreteHook() + + with pytest.raises(ValueError, match="toolsets not supported"): + hook.create_agent(AgentRunRequest(prompt="hi", toolsets=[MagicMock()])) + + assert hook.built is False + + +class TestAgentRunResult: + def test_agent_usage_defaults_to_none(self): + assert AgentUsage() == AgentUsage( + requests=None, + tool_calls=None, + input_tokens=None, + output_tokens=None, + total_tokens=None, + ) + + def test_dataclass_fields(self): + usage = AgentUsage(requests=1, tool_calls=2, total_tokens=10) + result = AgentRunResult( + output="answer", + message_history=["msg"], + model_name="test-model", + usage=usage, + tool_names=["query"], + ) + assert result.output == "answer" + assert result.message_history == ["msg"] + assert result.model_name == "test-model" + assert result.usage == usage + assert result.tool_names == ["query"] + assert result.durable_stats is None + + def test_durable_stats_field(self): + stats = DurableStats(replayed_model=2, cached_model=3) + result = AgentRunResult(output="x", durable_stats=stats) + assert result.durable_stats is stats + + +class TestAgentRunRequest: + def test_defaults(self): + req = AgentRunRequest(prompt="hello") + assert req.prompt == "hello" + assert req.output_type is str + assert req.instructions == "" + assert req.toolsets is None + assert req.usage_limits is None + assert req.message_history is None + assert req.enable_tool_logging is True + assert req.durable_context is None + assert req.agent_params == {} + + def test_with_all_fields(self): + ctx = DurableContext(dag_id="d", task_id="t", run_id="r", map_index=2) + req = AgentRunRequest( + prompt="test", + output_type=int, + instructions="sys", + toolsets=["ts"], + usage_limits="limits", + message_history=["h"], + enable_tool_logging=False, + durable_context=ctx, + agent_params={"retries": 3}, + ) + assert req.output_type is int + assert req.instructions == "sys" + assert req.durable_context is ctx + assert req.agent_params == {"retries": 3} + + +class TestBaseAIHookResolveTools: + def test_resolve_tools_calls_spec_to_native(self): + """_resolve_tools converts each ToolSpec via _tool_spec_to_native.""" + + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return {"name": spec.name, "fn": spec.fn} + + hook = ConcreteHook.__new__(ConcreteHook) + + def my_tool(x: int) -> str: + return str(x) + + class MyToolset(BaseToolset): + def as_tools(self): + return [ToolSpec(name="my_tool", description="desc", parameters={}, fn=my_tool)] + + result = hook._resolve_tools([MyToolset()], enable_logging=False) + + assert len(result) == 1 + assert result[0]["name"] == "my_tool" + + def test_resolve_tools_wraps_with_logging(self): + """When enable_logging=True, callable is wrapped.""" + mock_log = MagicMock() + + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + @property + def log(self): + return mock_log + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return spec.fn + + hook = ConcreteHook.__new__(ConcreteHook) + + calls = [] + + def original(): + calls.append("original") + return "result" + + class SimpleToolset(BaseToolset): + def as_tools(self): + return [ToolSpec(name="original", description="", parameters={}, fn=original)] + + [wrapped_fn] = hook._resolve_tools([SimpleToolset()], enable_logging=True) + wrapped_fn() + + assert calls == ["original"] + mock_log.info.assert_called() + + def test_resolve_tools_wraps_plain_callable(self): + """A plain function is auto-wrapped using __name__ and __doc__.""" + + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return {"name": spec.name, "description": spec.description, "fn": spec.fn} + + hook = ConcreteHook.__new__(ConcreteHook) + + def roll_dice() -> str: + """Roll a six-sided die and return the result.""" + return "4" + + result = hook._resolve_tools([roll_dice], enable_logging=False) + + assert len(result) == 1 + assert result[0]["name"] == "roll_dice" + assert result[0]["description"] == "Roll a six-sided die and return the result." + assert result[0]["fn"] is roll_dice + + def test_resolve_tools_wraps_bound_method(self): + """A bound method is auto-wrapped using __name__ and __doc__.""" + + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return {"name": spec.name, "description": spec.description, "fn": spec.fn} + + hook = ConcreteHook.__new__(ConcreteHook) + + class MyHelper: + def search(self, query: str) -> str: + """Search for data.""" + return query + + helper = MyHelper() + bound_method = helper.search + result = hook._resolve_tools([bound_method], enable_logging=False) + + assert len(result) == 1 + assert result[0]["name"] == "search" + assert result[0]["description"] == "Search for data." + assert result[0]["fn"] is bound_method + + def test_resolve_tools_wraps_partial(self): + """A functools.partial is auto-wrapped using the underlying function's name and doc.""" + + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return {"name": spec.name, "description": spec.description, "fn": spec.fn} + + hook = ConcreteHook.__new__(ConcreteHook) + + def query_db(db: str, query: str) -> str: + """Query the database.""" + return f"{db}: {query}" + + partial_tool = functools.partial(query_db, db="prod") + result = hook._resolve_tools([partial_tool], enable_logging=False) + + assert len(result) == 1 + assert result[0]["name"] == "query_db" + assert result[0]["description"] == "Query the database." + assert result[0]["fn"] is partial_tool + + def test_resolve_tools_wraps_callable_object(self): + """A callable object is auto-wrapped using the class name.""" + + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return {"name": spec.name, "fn": spec.fn} + + hook = ConcreteHook.__new__(ConcreteHook) + + class Searcher: + def __call__(self, query: str) -> str: + return query + + searcher = Searcher() + result = hook._resolve_tools([searcher], enable_logging=False) + + assert len(result) == 1 + assert result[0]["name"] == "Searcher" + assert result[0]["fn"] is searcher + + def test_resolve_tools_passes_non_function_non_toolset_through(self): + """Items that are not BaseToolset and not plain functions are passed through unchanged.""" + + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return spec.fn + + hook = ConcreteHook.__new__(ConcreteHook) + + native_tool_obj = object() # not a function, not a BaseToolset + result = hook._resolve_tools([native_tool_obj], enable_logging=True) + + assert result == [native_tool_obj] + + def test_resolve_tools_mixes_base_toolset_and_native(self): + """BaseToolset items are converted; non-function native items are passed through in order.""" + + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return f"converted:{spec.name}" + + hook = ConcreteHook.__new__(ConcreteHook) + + native_tool = object() # not a function, passes through unchanged + + class MyToolset(BaseToolset): + def as_tools(self): + return [ToolSpec(name="greet", description="", parameters={}, fn=lambda: "hi")] + + result = hook._resolve_tools([MyToolset(), native_tool], enable_logging=False) + + assert result == ["converted:greet", native_tool] + + def test_resolve_tools_applies_cache_wrapper_and_forces_sequential_for_resolved_specs(self): + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return spec + + hook = ConcreteHook.__new__(ConcreteHook) + + def original(): + return "result" + + def cached(fn): + def wrapper(): + return fn() + + return wrapper + + class SimpleToolset(BaseToolset): + def as_tools(self): + return [ToolSpec(name="original", description="", parameters={}, fn=original)] + + [spec] = hook._resolve_tools( + [SimpleToolset()], + enable_logging=False, + cache_wrapper=cached, + force_sequential=True, + ) + + assert spec.sequential is True + assert spec.fn is not original + + def test_resolve_tools_cache_hit_skips_logging_wrapper(self): + mock_log = MagicMock() + + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + @property + def log(self): + return mock_log + + def get_model(self): + return None + + def _build_agent(self, request): + return None + + def run_agent(self, agent, request): + return AgentRunResult(output="") + + def _tool_spec_to_native(self, spec): + return spec.fn + + hook = ConcreteHook.__new__(ConcreteHook) + calls = [] + + def original(): + calls.append("original") + return "computed" + + def cache_hit_wrapper(fn): + def wrapper(): + return "cached" + + return wrapper + + class SimpleToolset(BaseToolset): + def as_tools(self): + return [ToolSpec(name="original", description="", parameters={}, fn=original)] + + [wrapped] = hook._resolve_tools( + [SimpleToolset()], + enable_logging=True, + cache_wrapper=cache_hit_wrapper, + ) + + assert wrapped() == "cached" + assert calls == [] + mock_log.info.assert_not_called() + + +class TestBaseAIHookLoggedCallable: + def test_logged_callable_logs_and_returns(self): + logger = MagicMock() + calls = [] + + def fn(x): + calls.append(x) + return x * 2 + + wrapped = BaseAIHook._logged_callable(fn, logger) + result = wrapped(x=5) + + assert result == 10 + assert calls == [5] + logger.info.assert_called() + + def test_logged_callable_logs_exception(self): + logger = MagicMock() + + def failing(): + raise RuntimeError("boom") + + wrapped = BaseAIHook._logged_callable(failing, logger) + with pytest.raises(RuntimeError, match="boom"): + wrapped() + + logger.exception.assert_called_once() + + def test_logged_callable_uses_explicit_name_over_introspection(self): + logger = MagicMock() + + def fn(): + return "ok" + + wrapped = BaseAIHook._logged_callable(fn, logger, name="my_tool") + wrapped() + + logger.info.assert_any_call("::group::Tool call: %s", "my_tool") + logger.info.assert_any_call("Tool %s returned in %.2fs", "my_tool", pytest.approx(0.0, abs=1.0)) + + def test_logged_callable_partial_logs_correct_name_without_explicit_name(self): + """Without an explicit name, a partial falls back to type(fn).__name__ = 'partial'.""" + logger = MagicMock() + + def fetch_metric(environment: str, metric_name: str) -> float: + return 1.0 + + partial_fn = functools.partial(fetch_metric, "prod") + wrapped = BaseAIHook._logged_callable(partial_fn, logger) + wrapped(metric_name="cpu") + + # Without name= the fallback is type(partial).__name__ = "partial", not "fetch_metric". + logger.info.assert_any_call("::group::Tool call: %s", "partial") + + def test_logged_callable_partial_logs_correct_name_with_explicit_name(self): + """Passing name= fixes the 'partial' log name for functools.partial tools.""" + logger = MagicMock() + + def fetch_metric(environment: str, metric_name: str) -> float: + return 1.0 + + partial_fn = functools.partial(fetch_metric, "prod") + wrapped = BaseAIHook._logged_callable(partial_fn, logger, name="fetch_metric") + wrapped(metric_name="cpu") + + logger.info.assert_any_call("::group::Tool call: %s", "fetch_metric") + + def test_logged_callable_preserves_partial_introspection(self): + logger = MagicMock() + + def fetch_metric(environment: str, metric_name: str) -> float: + return 1.0 + + wrapped = BaseAIHook._logged_callable(functools.partial(fetch_metric, "prod"), logger) + + assert inspect.signature(wrapped) == inspect.signature(functools.partial(fetch_metric, "prod")) + + def test_logged_callable_preserves_callable_object_introspection(self): + logger = MagicMock() + + class CustomerLookup: + def __call__(self, customer_id: str) -> dict[str, str]: + return {"customer_id": customer_id} + + wrapped = BaseAIHook._logged_callable(CustomerLookup(), logger) + + signature = inspect.signature(wrapped) + assert tuple(signature.parameters) == ("customer_id",) + + def test_logged_callable_preserves_async_function_behavior(self): + logger = MagicMock() + + async def fn(value): + return value * 2 + + wrapped = BaseAIHook._logged_callable(fn, logger) + + assert inspect.iscoroutinefunction(wrapped) + assert asyncio.run(wrapped(3)) == 6 + logger.info.assert_called() + + def test_logged_callable_preserves_async_partial_behavior(self): + logger = MagicMock() + + async def fn(prefix, value): + return f"{prefix}:{value}" + + wrapped = BaseAIHook._logged_callable(functools.partial(fn, "prod"), logger) + + assert inspect.iscoroutinefunction(wrapped) + assert asyncio.run(wrapped("cpu")) == "prod:cpu" + logger.info.assert_called() + + def test_logged_callable_preserves_async_callable_object_behavior(self): + logger = MagicMock() + + class Lookup: + async def __call__(self, value): + return value.upper() + + wrapped = BaseAIHook._logged_callable(Lookup(), logger) + + assert inspect.iscoroutinefunction(wrapped) + assert asyncio.run(wrapped("abc")) == "ABC" + logger.info.assert_called() diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py index fb63b1f3e4813..9e467aaea5584 100644 --- a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py @@ -16,21 +16,102 @@ # under the License. from __future__ import annotations +import functools import json import sys +from typing import TYPE_CHECKING from unittest.mock import MagicMock, patch import pytest +from pydantic_ai import Agent, RunContext +from pydantic_ai.agent import DynamicToolset +from pydantic_ai.messages import ModelResponse, TextPart from pydantic_ai.models import Model from pydantic_ai.models.test import TestModel +from pydantic_ai.run import AgentRunResult as PydanticAgentRunResult +from pydantic_ai.usage import RunUsage, UsageLimits from airflow.models.connection import Connection +from airflow.providers.common.ai.hooks.base import ( + AgentRunRequest, + AgentRunResult, + BaseAIHook, + Capability, + ToolSpec, +) from airflow.providers.common.ai.hooks.pydantic_ai import ( + PydanticAgentHandle, PydanticAIAzureHook, PydanticAIBedrockHook, PydanticAIHook, PydanticAIVertexHook, ) +from airflow.providers.common.ai.mixins.durable import DurableState + +if TYPE_CHECKING: + from pydantic_ai.toolsets import AbstractToolset + + +def _test_agent() -> Agent[None, str]: + return Agent(TestModel()) + + +def _test_handle() -> PydanticAgentHandle: + return PydanticAgentHandle(agent=_test_agent()) + + +def _pydantic_run_result( + output: str, + *, + model_name: str = "test-model", + message_history: list | None = None, + requests: int = 1, + tool_calls: int = 0, + input_tokens: int = 0, + output_tokens: int = 0, +) -> MagicMock: + mock_result = MagicMock(spec=PydanticAgentRunResult) + mock_result.output = output + mock_result.usage = RunUsage( + requests=requests, + tool_calls=tool_calls, + input_tokens=input_tokens, + output_tokens=output_tokens, + ) + mock_result.response = ModelResponse( + parts=[TextPart(content=str(output))], + model_name=model_name, + ) + mock_result.all_messages.return_value = message_history or [] + return mock_result + + +def _noop_override_context() -> MagicMock: + ctx = MagicMock() + ctx.__enter__ = MagicMock(return_value=None) + ctx.__exit__ = MagicMock(return_value=False) + return ctx + + +class _PydanticAIHookWithTestModel(PydanticAIHook): + """Concrete hook that uses a real TestModel without patching Agent construction.""" + + def __init__(self, model: TestModel): + super().__init__(llm_conn_id="test_conn", model_id="test-model") + self._test_model = model + + def get_model(self) -> TestModel: + return self._test_model + + +class TestPydanticAIHookBaseContract: + def test_is_base_hook(self): + assert issubclass(PydanticAIHook, BaseAIHook) + + def test_capability_flags(self): + assert Capability.TOOLSETS in PydanticAIHook.capabilities + assert Capability.DURABLE in PydanticAIHook.capabilities + assert Capability.USAGE_LIMITS in PydanticAIHook.capabilities class TestPydanticAIHookInit: @@ -57,11 +138,16 @@ def test_vertex_hook_uses_own_default_conn_name(self): hook = PydanticAIVertexHook() assert hook.llm_conn_id == "pydanticai_vertex_default" + def test_durable_state_not_stored_on_hook_instance(self): + hook = PydanticAIHook() + assert not hasattr(hook, "_durable_storage") + assert not hasattr(hook, "_durable_counter") + -class TestPydanticAIHookGetConn: +class TestPydanticAIHookGetModel: @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class", autospec=True) - def test_get_conn_with_api_key_and_base_url(self, mock_infer_provider_class, mock_infer_model): + def test_get_model_with_api_key_and_base_url(self, mock_infer_provider_class, mock_infer_model): """Credentials are injected via provider_factory, not as direct kwargs.""" mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model @@ -76,16 +162,14 @@ def test_get_conn_with_api_key_and_base_url(self, mock_infer_provider_class, moc host="https://api.openai.com/v1", ) with patch.object(hook, "get_connection", return_value=conn): - result = hook.get_conn() + result = hook.get_model() assert result is mock_model mock_infer_model.assert_called_once() call_args = mock_infer_model.call_args assert call_args[0][0] == "openai:gpt-5.3" - # provider_factory should be passed as keyword arg assert "provider_factory" in call_args[1] - # Call the factory to verify it creates the provider with credentials factory = call_args[1]["provider_factory"] factory("openai") mock_infer_provider_class.assert_called_with("openai") @@ -95,7 +179,7 @@ def test_get_conn_with_api_key_and_base_url(self, mock_infer_provider_class, moc @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class", autospec=True) - def test_get_conn_with_model_from_extra(self, mock_infer_provider_class, mock_infer_model): + def test_get_model_with_model_from_extra(self, mock_infer_provider_class, mock_infer_model): mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model mock_infer_provider_class.return_value = MagicMock(return_value=MagicMock()) @@ -108,7 +192,7 @@ def test_get_conn_with_model_from_extra(self, mock_infer_provider_class, mock_in extra='{"model": "anthropic:claude-opus-4-6"}', ) with patch.object(hook, "get_connection", return_value=conn): - result = hook.get_conn() + result = hook.get_model() assert result is mock_model assert mock_infer_model.call_args[0][0] == "anthropic:claude-opus-4-6" @@ -127,12 +211,11 @@ def test_model_id_param_overrides_extra(self, mock_infer_provider_class, mock_in extra='{"model": "anthropic:claude-opus-4-6"}', ) with patch.object(hook, "get_connection", return_value=conn): - hook.get_conn() + hook.get_model() - # model_id param takes priority over extra assert mock_infer_model.call_args[0][0] == "openai:gpt-5.3" - def test_get_conn_raises_when_no_model(self): + def test_get_model_raises_when_no_model(self): hook = PydanticAIHook(llm_conn_id="test_conn") conn = Connection( conn_id="test_conn", @@ -141,10 +224,10 @@ def test_get_conn_raises_when_no_model(self): ) with patch.object(hook, "get_connection", return_value=conn): with pytest.raises(ValueError, match="No model specified"): - hook.get_conn() + hook.get_model() @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) - def test_get_conn_without_credentials_uses_default_provider(self, mock_infer_model): + def test_get_model_without_credentials_uses_default_provider(self, mock_infer_model): """No api_key or base_url means env-based auth (Bedrock, Vertex, etc.).""" mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model @@ -155,14 +238,13 @@ def test_get_conn_without_credentials_uses_default_provider(self, mock_infer_mod conn_type="pydanticai", ) with patch.object(hook, "get_connection", return_value=conn): - hook.get_conn() + hook.get_model() - # No provider_factory — uses default infer_provider which reads env vars mock_infer_model.assert_called_once_with("bedrock:us.anthropic.claude-v2") @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class", autospec=True) - def test_get_conn_with_base_url_only(self, mock_infer_provider_class, mock_infer_model): + def test_get_model_with_base_url_only(self, mock_infer_provider_class, mock_infer_model): """Ollama / vLLM: base_url but no API key.""" mock_infer_model.return_value = MagicMock(spec=Model) mock_infer_provider_class.return_value = MagicMock(return_value=MagicMock()) @@ -174,30 +256,166 @@ def test_get_conn_with_base_url_only(self, mock_infer_provider_class, mock_infer host="http://localhost:11434/v1", ) with patch.object(hook, "get_connection", return_value=conn): - hook.get_conn() + hook.get_model() - # provider_factory should be used since base_url is set factory = mock_infer_model.call_args[1]["provider_factory"] factory("openai") mock_infer_provider_class.return_value.assert_called_with(base_url="http://localhost:11434/v1") @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) - def test_get_conn_caches_model(self, mock_infer_model): - """get_conn() should resolve the model once and cache it.""" + def test_get_model_caches_result(self, mock_infer_model): + """get_model() should resolve the model once and cache it.""" mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") conn = Connection(conn_id="test_conn", conn_type="pydanticai") with patch.object(hook, "get_connection", return_value=conn): - first = hook.get_conn() - second = hook.get_conn() + first = hook.get_model() + second = hook.get_model() assert first is second mock_infer_model.assert_called_once() + def test_get_conn_delegates_to_get_model(self): + """get_conn() is a compatibility shim that calls get_model().""" + hook = PydanticAIHook() + mock_model = MagicMock() + with patch.object(hook, "get_model", return_value=mock_model): + result = hook.get_conn() + assert result is mock_model + class TestPydanticAIHookCreateAgent: + def test_create_agent_runs_callable_object_tool_with_real_schema(self): + """Callable objects should produce a real pydantic-ai function tool and execute successfully.""" + model = TestModel(call_tools="all") + hook = _PydanticAIHookWithTestModel(model) + calls: list[str] = [] + + class CustomerLookup: + def __call__(self, customer_id: str) -> dict[str, str]: + calls.append(customer_id) + return {"customer_id": customer_id} + + request = AgentRunRequest( + prompt="Look up a customer", + toolsets=[CustomerLookup()], + enable_tool_logging=True, + ) + + agent = hook.create_agent(request) + run_result = hook.run_agent(agent, request) + + assert run_result.usage is not None + assert run_result.usage.tool_calls == 1 + assert len(calls) == 1 + assert isinstance(calls[0], str) + + [tool_def] = model.last_model_request_parameters.function_tools + assert tool_def.name == "CustomerLookup" + assert set(tool_def.parameters_json_schema["properties"]) == {"customer_id"} + assert "environment" not in tool_def.parameters_json_schema["properties"] + + def test_create_agent_runs_partial_tool_with_bound_argument_removed_from_schema(self): + """functools.partial should expose only remaining parameters and preserve bound args at runtime.""" + model = TestModel(call_tools="all") + hook = _PydanticAIHookWithTestModel(model) + calls: list[tuple[str, str]] = [] + + def fetch_metric(environment: str, metric_name: str) -> float: + calls.append((environment, metric_name)) + return 1.0 + + request = AgentRunRequest( + prompt="Fetch a metric", + toolsets=[functools.partial(fetch_metric, "prod")], + enable_tool_logging=True, + ) + + agent = hook.create_agent(request) + run_result = hook.run_agent(agent, request) + + assert run_result.usage is not None + assert run_result.usage.tool_calls == 1 + assert len(calls) == 1 + assert calls[0][0] == "prod" + assert isinstance(calls[0][1], str) + + [tool_def] = model.last_model_request_parameters.function_tools + assert tool_def.name == "fetch_metric" + assert set(tool_def.parameters_json_schema["properties"]) == {"metric_name"} + assert "environment" not in tool_def.parameters_json_schema["properties"] + + def test_create_agent_runs_bound_method_tool_with_real_schema(self): + """Bound methods should expose method parameters without leaking ``self``.""" + model = TestModel(call_tools="all") + hook = _PydanticAIHookWithTestModel(model) + calls: list[str] = [] + + class InventoryClient: + def get_stock_level(self, sku: str) -> int: + """Return stock level for a SKU.""" + calls.append(sku) + return 42 + + request = AgentRunRequest( + prompt="Check stock", + toolsets=[InventoryClient().get_stock_level], + enable_tool_logging=True, + ) + + handle = hook.create_agent(request) + run_result = hook.run_agent(handle, request) + + assert run_result.usage is not None + assert run_result.usage.tool_calls == 1 + assert len(calls) == 1 + assert isinstance(calls[0], str) + + [tool_def] = model.last_model_request_parameters.function_tools + assert tool_def.name == "get_stock_level" + assert set(tool_def.parameters_json_schema["properties"]) == {"sku"} + assert "self" not in tool_def.parameters_json_schema["properties"] + + def test_create_agent_runs_mixed_callable_tool_patterns(self): + """Bound methods, partials, and callable objects can be mixed in one request.""" + model = TestModel(call_tools="all") + hook = _PydanticAIHookWithTestModel(model) + calls: list[str] = [] + + class MetricsClient: + def ping(self) -> str: + calls.append("bound") + return "bound" + + def fetch_metric(environment: str) -> str: + calls.append(environment) + return environment + + class CustomerLookup: + def __call__(self) -> str: + calls.append("callable") + return "callable" + + request = AgentRunRequest( + prompt="Run all tools", + toolsets=[ + MetricsClient().ping, + functools.partial(fetch_metric, "prod"), + CustomerLookup(), + ], + enable_tool_logging=True, + ) + + handle = hook.create_agent(request) + run_result = hook.run_agent(handle, request) + + assert run_result.usage is not None + assert run_result.usage.tool_calls == 3 + # TestModel is not a real model and does not guarantee tool call order. + assert sorted(calls) == ["bound", "callable", "prod"] + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) def test_create_agent_defaults(self, mock_agent_cls, mock_infer_model): @@ -205,12 +423,10 @@ def test_create_agent_defaults(self, mock_agent_cls, mock_infer_model): mock_infer_model.return_value = mock_model hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") - conn = Connection( - conn_id="test_conn", - conn_type="pydanticai", - ) + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest(prompt="hi", instructions="You are a helpful assistant.") with patch.object(hook, "get_connection", return_value=conn): - hook.create_agent(instructions="You are a helpful assistant.") + hook.create_agent(request) mock_agent_cls.assert_called_once_with( mock_model, @@ -220,21 +436,20 @@ def test_create_agent_defaults(self, mock_agent_cls, mock_infer_model): @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) - def test_create_agent_with_params(self, mock_agent_cls, mock_infer_model): + def test_create_agent_with_agent_params(self, mock_agent_cls, mock_infer_model): mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") - conn = Connection( - conn_id="test_conn", - conn_type="pydanticai", + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest( + prompt="hi", + output_type=dict, + instructions="Be helpful.", + agent_params={"retries": 3}, ) with patch.object(hook, "get_connection", return_value=conn): - hook.create_agent( - output_type=dict, - instructions="Be helpful.", - retries=3, - ) + hook.create_agent(request) mock_agent_cls.assert_called_once_with( mock_model, @@ -243,6 +458,508 @@ def test_create_agent_with_params(self, mock_agent_cls, mock_infer_model): retries=3, ) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) + def test_create_agent_rejects_raw_json_schema_output_type(self, mock_agent_cls, mock_infer_model): + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest( + prompt="hi", + output_type={"type": "object", "properties": {}}, # type: ignore[arg-type] + ) + + with patch.object(hook, "get_connection", return_value=conn): + with pytest.raises(ValueError, match="raw JSON schema mappings"): + hook.create_agent(request) + + mock_agent_cls.assert_not_called() + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) + def test_create_agent_rejects_tools_in_agent_params_with_toolsets(self, mock_agent_cls, mock_infer_model): + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest( + prompt="hi", + toolsets=[lambda: "ok"], + agent_params={"tools": [MagicMock()]}, + ) + with patch.object(hook, "get_connection", return_value=conn): + with pytest.raises(ValueError, match="agent_params must not include 'tools'"): + hook.create_agent(request) + + mock_agent_cls.assert_not_called() + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) + def test_create_agent_rejects_toolsets_in_agent_params_with_toolsets( + self, mock_agent_cls, mock_infer_model + ): + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest( + prompt="hi", + toolsets=[lambda: "ok"], + agent_params={"toolsets": [MagicMock()]}, + ) + with patch.object(hook, "get_connection", return_value=conn): + with pytest.raises(ValueError, match="agent_params must not include 'toolsets'"): + hook.create_agent(request) + + mock_agent_cls.assert_not_called() + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + def test_create_agent_inits_durable_when_context_set(self, mock_infer_model): + from airflow.providers.common.ai.hooks.base import DurableContext + + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + ctx = DurableContext(dag_id="d", task_id="t", run_id="r") + request = AgentRunRequest(prompt="hi", durable_context=ctx) + + mock_storage = MagicMock() + mock_counter = MagicMock() + durable_state = DurableState(storage=mock_storage, counter=mock_counter) + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + with ( + patch.object(hook, "get_connection", return_value=conn), + patch.object(hook, "_init_durable", return_value=durable_state), + ): + handle = hook.create_agent(request) + + assert handle.durable_state is durable_state + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + def test_create_agent_returns_handle_without_durable_when_no_context(self, mock_infer_model): + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + + request = AgentRunRequest(prompt="hi") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + with patch.object(hook, "get_connection", return_value=conn): + handle = hook.create_agent(request) + + assert isinstance(handle, PydanticAgentHandle) + assert handle.durable_state is None + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) + def test_create_agent_passes_native_tools_through_directly(self, mock_agent_cls, mock_infer_model): + """Native pydantic-ai Tool objects bypass Airflow callable wrappers.""" + from pydantic_ai.tools import Tool + + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + native_tool = MagicMock(spec=Tool) + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest(prompt="hi", toolsets=[native_tool]) + with patch.object(hook, "get_connection", return_value=conn): + hook.create_agent(request) + + call_kwargs = mock_agent_cls.call_args[1] + assert any(t is native_tool for t in call_kwargs["tools"]) + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) + def test_create_agent_mixes_base_toolset_and_native_tool(self, mock_agent_cls, mock_infer_model): + """BaseToolset items are expanded; native Tool objects are passed through unchanged.""" + from pydantic_ai.tools import Tool + + from airflow.providers.common.ai.hooks.base import BaseToolset + + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + def first_fn() -> str: + return "first" + + def second_fn() -> str: + return "second" + + class MyToolset(BaseToolset): + def as_tools(self): + return [ + ToolSpec(name="first_fn", description="desc", parameters={}, fn=first_fn), + ToolSpec(name="second_fn", description="desc", parameters={}, fn=second_fn), + ] + + native_tool = MagicMock(spec=Tool) + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest(prompt="hi", toolsets=[MyToolset(), native_tool], enable_tool_logging=False) + with patch.object(hook, "get_connection", return_value=conn): + hook.create_agent(request) + + call_kwargs = mock_agent_cls.call_args[1] + tools = call_kwargs["tools"] + assert len(tools) == 3 + assert tools[2] is native_tool + assert [tool.name for tool in tools[:2]] == ["first_fn", "second_fn"] + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) + def test_create_agent_durable_forces_airflow_tools_sequential_but_preserves_native_tool( + self, mock_agent_cls, mock_infer_model + ): + """Durable cache serialization only applies to Airflow-resolved callables, not native Tool objects.""" + from pydantic_ai.tools import Tool + + from airflow.providers.common.ai.hooks.base import DurableContext + + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + native_tool = Tool(lambda: "native", name="native_tool") + storage = MagicMock() + counter = MagicMock() + durable_state = DurableState(storage=storage, counter=counter) + + def airflow_tool() -> str: + return "airflow" + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest( + prompt="hi", + toolsets=[airflow_tool, native_tool], + durable_context=DurableContext(dag_id="d", task_id="t", run_id="r"), + enable_tool_logging=False, + ) + with ( + patch.object(hook, "get_connection", return_value=conn), + patch.object(hook, "_init_durable", return_value=durable_state), + ): + hook.create_agent(request) + + tools = mock_agent_cls.call_args[1]["tools"] + assert len(tools) == 2 + assert tools[0].sequential is True + assert tools[1] is native_tool + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) + def test_create_agent_routes_abstract_toolset_to_toolsets_kwarg(self, mock_agent_cls, mock_infer_model): + """AbstractToolset items must go in Agent(toolsets=[...]), not Agent(tools=[...]).""" + from pydantic_ai.toolsets.abstract import AbstractToolset + + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + abstract_ts = MagicMock(spec=AbstractToolset) + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest(prompt="hi", toolsets=[abstract_ts], enable_tool_logging=False) + with patch.object(hook, "get_connection", return_value=conn): + hook.create_agent(request) + + call_kwargs = mock_agent_cls.call_args[1] + assert "tools" not in call_kwargs + assert "toolsets" in call_kwargs + assert any(ts is abstract_ts for ts in call_kwargs["toolsets"]) + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) + def test_create_agent_routes_dynamic_toolset_to_toolsets_kwarg(self, mock_agent_cls, mock_infer_model): + """DynamicToolset-wrapped factories must pass through as native pydantic-ai toolsets.""" + + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + def select_toolset(ctx: RunContext) -> AbstractToolset | None: + return None + + dynamic_toolset = DynamicToolset(select_toolset) + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest(prompt="hi", toolsets=[dynamic_toolset], enable_tool_logging=False) + with patch.object(hook, "get_connection", return_value=conn): + hook.create_agent(request) + + call_kwargs = mock_agent_cls.call_args[1] + assert "tools" not in call_kwargs + assert "toolsets" in call_kwargs + assert any(ts is dynamic_toolset for ts in call_kwargs["toolsets"]) + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) + def test_create_agent_wraps_abstract_toolset_with_logging(self, mock_agent_cls, mock_infer_model): + """AbstractToolset items are wrapped with LoggingToolset when enable_tool_logging=True.""" + from pydantic_ai.toolsets.abstract import AbstractToolset + + from airflow.providers.common.ai.toolsets.logging import LoggingToolset + + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + abstract_ts = MagicMock(spec=AbstractToolset) + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest(prompt="hi", toolsets=[abstract_ts], enable_tool_logging=True) + with patch.object(hook, "get_connection", return_value=conn): + hook.create_agent(request) + + call_kwargs = mock_agent_cls.call_args[1] + toolsets = call_kwargs["toolsets"] + assert len(toolsets) == 1 + assert isinstance(toolsets[0], LoggingToolset) + assert toolsets[0].wrapped is abstract_ts + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) + def test_create_agent_wraps_abstract_toolset_with_caching_when_durable( + self, mock_agent_cls, mock_infer_model + ): + """AbstractToolset items use CachingToolset outside LoggingToolset for durable runs.""" + from pydantic_ai.toolsets.abstract import AbstractToolset + + from airflow.providers.common.ai.durable.caching_toolset import CachingToolset + from airflow.providers.common.ai.hooks.base import DurableContext + from airflow.providers.common.ai.toolsets.logging import LoggingToolset + + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + abstract_ts = MagicMock(spec=AbstractToolset) + mock_storage = MagicMock() + mock_counter = MagicMock() + durable_state = DurableState(storage=mock_storage, counter=mock_counter) + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + ctx = DurableContext(dag_id="d", task_id="t", run_id="r") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + request = AgentRunRequest( + prompt="hi", toolsets=[abstract_ts], durable_context=ctx, enable_tool_logging=True + ) + with ( + patch.object(hook, "get_connection", return_value=conn), + patch.object(hook, "_init_durable", return_value=durable_state), + ): + hook.create_agent(request) + + call_kwargs = mock_agent_cls.call_args[1] + toolsets = call_kwargs["toolsets"] + assert len(toolsets) == 1 + outer = toolsets[0] + assert isinstance(outer, CachingToolset) + assert isinstance(outer.wrapped, LoggingToolset) + assert outer.wrapped.wrapped is abstract_ts + + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) + def test_create_agent_returns_durable_state_per_handle_not_on_hook(self, mock_infer_model): + """Second create_agent must not overwrite durable state for the first agent.""" + from airflow.providers.common.ai.hooks.base import DurableContext + + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + hook = PydanticAIHook(llm_conn_id="test_conn", model_id="openai:gpt-5.3") + ctx_a = DurableContext(dag_id="d", task_id="t", run_id="r1") + ctx_b = DurableContext(dag_id="d", task_id="t", run_id="r2") + storage_a, counter_a = MagicMock(), MagicMock() + storage_b, counter_b = MagicMock(), MagicMock() + durable_a = DurableState(storage=storage_a, counter=counter_a) + durable_b = DurableState(storage=storage_b, counter=counter_b) + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + + with patch.object(hook, "get_connection", return_value=conn): + with patch.object(hook, "_init_durable", side_effect=[durable_a, durable_b]): + handle_a = hook.create_agent(AgentRunRequest(prompt="a", durable_context=ctx_a)) + handle_b = hook.create_agent(AgentRunRequest(prompt="b", durable_context=ctx_b)) + + assert handle_a is not handle_b + assert handle_a.durable_state is durable_a + assert handle_b.durable_state is durable_b + + +class TestPydanticAIHookRunAgent: + def test_run_agent_returns_agent_run_result(self): + hook = PydanticAIHook() + agent = _test_agent() + handle = PydanticAgentHandle(agent=agent) + mock_result = _pydantic_run_result( + "done", + model_name="openai:gpt-5", + input_tokens=5, + output_tokens=10, + ) + + request = AgentRunRequest(prompt="hello") + with patch.object(agent, "run_sync", return_value=mock_result) as mock_run_sync: + run_result = hook.run_agent(handle, request) + + assert isinstance(run_result, AgentRunResult) + assert run_result.output == "done" + assert run_result.model_name == "openai:gpt-5" + assert run_result.usage.total_tokens == 15 + mock_run_sync.assert_called_once_with("hello") + + def test_create_agent_rejects_unsupported_usage_limits(self): + hook = PydanticAIHook() + hook.capabilities = frozenset() # strip all capabilities + with pytest.raises(ValueError, match="usage_limits not supported"): + hook.create_agent(AgentRunRequest(prompt="hi", usage_limits=UsageLimits())) + + def test_run_agent_forwards_message_history_and_usage_limits(self): + hook = PydanticAIHook() + agent = _test_agent() + handle = PydanticAgentHandle(agent=agent) + mock_result = _pydantic_run_result("ok", model_name="m", message_history=["history"]) + limits = UsageLimits() + history = ["prior"] + + request = AgentRunRequest(prompt="more", message_history=history, usage_limits=limits) + with patch.object(agent, "run_sync", return_value=mock_result) as mock_run_sync: + hook.run_agent(handle, request) + + mock_run_sync.assert_called_once_with("more", message_history=history, usage_limits=limits) + + @patch.object(Agent, "override") + @patch.object(Agent, "run_sync") + @patch("airflow.providers.common.ai.hooks.pydantic_ai.CachingModel") + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", side_effect=lambda m: m) + def test_run_agent_durable_applies_caching_model( + self, + mock_infer_model, + mock_caching_model_cls, + mock_run_sync, + mock_override, + ): + """When durable state is set, run_agent wraps model with CachingModel.""" + from airflow.providers.common.ai.hooks.base import DurableContext + + hook = PydanticAIHook() + agent = _test_agent() + mock_run_sync.return_value = _pydantic_run_result("ok", model_name="m") + mock_override.return_value = _noop_override_context() + mock_caching_model_cls.return_value = MagicMock() + + mock_storage = MagicMock() + mock_counter = MagicMock() + mock_counter.replayed_model = 1 + mock_counter.replayed_tool = 0 + mock_counter.cached_model = 0 + mock_counter.cached_tool = 0 + handle = PydanticAgentHandle( + agent=agent, + durable_state=DurableState(storage=mock_storage, counter=mock_counter), + ) + + request = AgentRunRequest( + prompt="hi", + durable_context=DurableContext(dag_id="d", task_id="t", run_id="r"), + ) + run_result = hook.run_agent(handle, request) + + mock_caching_model_cls.assert_called_once() + mock_override.assert_called_once() + mock_run_sync.assert_called_once_with("hi") + assert run_result.durable_stats is not None + mock_storage.cleanup.assert_called_once() + + @patch.object(Agent, "override") + @patch.object(Agent, "run_sync", side_effect=RuntimeError("boom")) + @patch("airflow.providers.common.ai.hooks.pydantic_ai.CachingModel") + @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", side_effect=lambda m: m) + def test_run_agent_preserves_durable_cache_on_exception( + self, + mock_infer_model, + mock_caching_model_cls, + mock_run_sync, + mock_override, + ): + from airflow.providers.common.ai.hooks.base import DurableContext + + hook = PydanticAIHook() + agent = _test_agent() + mock_override.return_value = _noop_override_context() + + mock_storage = MagicMock() + mock_counter = MagicMock() + handle = PydanticAgentHandle( + agent=agent, + durable_state=DurableState(storage=mock_storage, counter=mock_counter), + ) + + with pytest.raises(RuntimeError, match="boom"): + hook.run_agent( + handle, + AgentRunRequest( + prompt="hi", + durable_context=DurableContext(dag_id="d", task_id="t", run_id="r"), + ), + ) + + mock_storage.cleanup.assert_not_called() + + def test_run_agent_rejects_native_agent_without_handle(self): + hook = PydanticAIHook() + + with pytest.raises(TypeError, match="requires a PydanticAgentHandle"): + hook.run_agent(_test_agent(), AgentRunRequest(prompt="hi")) + + def test_run_agent_rejects_durable_request_without_durable_state(self): + from airflow.providers.common.ai.hooks.base import DurableContext + + hook = PydanticAIHook() + handle = _test_handle() + request = AgentRunRequest( + prompt="hi", + durable_context=DurableContext(dag_id="d", task_id="t", run_id="r"), + ) + + with pytest.raises(ValueError, match="requires a PydanticAgentHandle with durable state"): + hook.run_agent(handle, request) + + def test_run_agent_rejects_durable_handle_without_durable_request(self): + hook = PydanticAIHook() + handle = PydanticAgentHandle( + agent=_test_agent(), + durable_state=DurableState(storage=MagicMock(), counter=MagicMock()), + ) + + with pytest.raises(ValueError, match="durable state, but request.durable_context is not set"): + hook.run_agent(handle, AgentRunRequest(prompt="hi")) + + def test_tool_spec_to_native_tools_called(self): + hook = PydanticAIHook() + + def fn(customer_id: int) -> str: + """test function""" + return "ok" + + with patch("airflow.providers.common.ai.hooks.pydantic_ai.Tool") as mock_tool_cls: + hook._resolve_tools(toolsets=[fn], enable_logging=False) + mock_tool_cls.from_schema.assert_called_once_with( + fn, + name="fn", + description="test function", + sequential=False, + json_schema={ + "properties": {"customer_id": {"type": "integer"}}, + "required": ["customer_id"], + "type": "object", + }, + ) + class TestPydanticAIHookCreateAgentInstrumentation: """create_agent() wires OpenTelemetry instrumentation from observability.""" @@ -256,31 +973,30 @@ def test_instrument_set_when_settings_returned(self, mock_settings): sentinel = MagicMock(name="InstrumentationSettings") mock_settings.return_value = sentinel hook = self._hook() - with patch.object(hook, "get_conn", return_value=TestModel()): - agent = hook.create_agent(instructions="hi") + with patch.object(hook, "get_model", return_value=TestModel()): + handle = hook.create_agent(AgentRunRequest(prompt="test", instructions="hi")) - assert agent.instrument is sentinel + assert handle.agent.instrument is sentinel @patch("airflow.providers.common.ai.hooks.pydantic_ai.genai_instrumentation_settings") def test_no_instrument_when_settings_none(self, mock_settings): mock_settings.return_value = None hook = self._hook() - with patch.object(hook, "get_conn", return_value=TestModel()): - agent = hook.create_agent(instructions="hi") + with patch.object(hook, "get_model", return_value=TestModel()): + handle = hook.create_agent(AgentRunRequest(prompt="test", instructions="hi")) mock_settings.assert_called_once() - assert agent.instrument is None + assert handle.agent.instrument is None @patch("airflow.providers.common.ai.hooks.pydantic_ai.Agent", autospec=True) @patch("airflow.providers.common.ai.hooks.pydantic_ai.genai_instrumentation_settings") - @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) - def test_caller_instrument_short_circuits(self, mock_infer_model, mock_settings, mock_agent_cls): - """A caller that passes its own ``instrument`` wins; we don't override it.""" - mock_infer_model.return_value = MagicMock(spec=Model) + def test_caller_instrument_short_circuits(self, mock_settings, mock_agent_cls): + """A caller that passes its own ``instrument`` via agent_params wins; we don't override it.""" hook = self._hook() - conn = Connection(conn_id="test_conn", conn_type="pydanticai") - with patch.object(hook, "get_connection", return_value=conn): - hook.create_agent(instructions="hi", instrument=False) + with patch.object(hook, "get_model", return_value=TestModel()): + hook.create_agent( + AgentRunRequest(prompt="test", instructions="hi", agent_params={"instrument": False}) + ) mock_settings.assert_not_called() @@ -387,12 +1103,11 @@ def test_get_provider_kwargs_empty_without_api_version(self): "https://myresource.openai.azure.com", {"model": "azure:gpt-4o"}, ) - # api_version should not appear if not in extra assert "api_version" not in result @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class", autospec=True) - def test_get_conn_uses_azure_endpoint(self, mock_infer_provider_class, mock_infer_model): + def test_get_model_uses_azure_endpoint(self, mock_infer_provider_class, mock_infer_model): mock_infer_model.return_value = MagicMock(spec=Model) mock_provider_cls = MagicMock(return_value=MagicMock()) mock_infer_provider_class.return_value = mock_provider_cls @@ -406,7 +1121,7 @@ def test_get_conn_uses_azure_endpoint(self, mock_infer_provider_class, mock_infe extra=json.dumps({"model": "azure:gpt-4o", "api_version": "2024-07-01-preview"}), ) with patch.object(hook, "get_connection", return_value=conn): - hook.get_conn() + hook.get_model() factory = mock_infer_model.call_args[1]["provider_factory"] factory("azure") @@ -417,7 +1132,7 @@ def test_get_conn_uses_azure_endpoint(self, mock_infer_provider_class, mock_infe ) @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) - def test_get_conn_falls_back_to_env_auth_when_no_kwargs(self, mock_infer_model): + def test_get_model_falls_back_to_env_auth_when_no_kwargs(self, mock_infer_model): """No host + no password → env-var auth path (empty _get_provider_kwargs).""" mock_infer_model.return_value = MagicMock(spec=Model) hook = PydanticAIAzureHook(llm_conn_id="azure_test") @@ -427,7 +1142,7 @@ def test_get_conn_falls_back_to_env_auth_when_no_kwargs(self, mock_infer_model): extra=json.dumps({"model": "azure:gpt-4o"}), ) with patch.object(hook, "get_connection", return_value=conn): - hook.get_conn() + hook.get_model() mock_infer_model.assert_called_once_with("azure:gpt-4o") @@ -471,7 +1186,7 @@ def test_get_provider_kwargs_returns_empty_for_env_auth(self): assert result == {} @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) - def test_get_conn_falls_back_to_env_auth(self, mock_infer_model): + def test_get_model_falls_back_to_env_auth(self, mock_infer_model): mock_infer_model.return_value = MagicMock(spec=Model) hook = PydanticAIBedrockHook(llm_conn_id="bedrock_test") conn = Connection( @@ -480,13 +1195,13 @@ def test_get_conn_falls_back_to_env_auth(self, mock_infer_model): extra=json.dumps({"model": "bedrock:us.anthropic.claude-opus-4-5"}), ) with patch.object(hook, "get_connection", return_value=conn): - hook.get_conn() + hook.get_model() mock_infer_model.assert_called_once_with("bedrock:us.anthropic.claude-opus-4-5") @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class", autospec=True) - def test_get_conn_uses_explicit_keys(self, mock_infer_provider_class, mock_infer_model): + def test_get_model_uses_explicit_keys(self, mock_infer_provider_class, mock_infer_model): mock_infer_model.return_value = MagicMock(spec=Model) mock_provider_cls = MagicMock(return_value=MagicMock()) mock_infer_provider_class.return_value = mock_provider_cls @@ -505,7 +1220,7 @@ def test_get_conn_uses_explicit_keys(self, mock_infer_provider_class, mock_infer ), ) with patch.object(hook, "get_connection", return_value=conn): - hook.get_conn() + hook.get_model() factory = mock_infer_model.call_args[1]["provider_factory"] factory("bedrock") @@ -552,8 +1267,8 @@ def test_get_provider_kwargs_float_timeouts(self): None, { "model": "bedrock:us.anthropic.claude-opus-4-5", - "aws_read_timeout": 60, # int from JSON - "aws_connect_timeout": 10.5, # float already + "aws_read_timeout": 60, + "aws_connect_timeout": 10.5, }, ) assert result["aws_read_timeout"] == 60.0 @@ -656,7 +1371,7 @@ def test_get_provider_kwargs_returns_empty_for_adc(self): assert result == {} @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) - def test_get_conn_falls_back_to_adc(self, mock_infer_model): + def test_get_model_falls_back_to_adc(self, mock_infer_model): mock_infer_model.return_value = MagicMock(spec=Model) hook = PydanticAIVertexHook(llm_conn_id="vertex_test") conn = Connection( @@ -665,13 +1380,13 @@ def test_get_conn_falls_back_to_adc(self, mock_infer_model): extra=json.dumps({"model": "google-vertex:gemini-2.0-flash"}), ) with patch.object(hook, "get_connection", return_value=conn): - hook.get_conn() + hook.get_model() mock_infer_model.assert_called_once_with("google-vertex:gemini-2.0-flash") @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_provider_class", autospec=True) - def test_get_conn_uses_explicit_project(self, mock_infer_provider_class, mock_infer_model): + def test_get_model_uses_explicit_project(self, mock_infer_provider_class, mock_infer_model): mock_infer_model.return_value = MagicMock(spec=Model) mock_provider_cls = MagicMock(return_value=MagicMock()) mock_infer_provider_class.return_value = mock_provider_cls @@ -689,7 +1404,7 @@ def test_get_conn_uses_explicit_project(self, mock_infer_provider_class, mock_in ), ) with patch.object(hook, "get_connection", return_value=conn): - hook.get_conn() + hook.get_model() factory = mock_infer_model.call_args[1]["provider_factory"] factory("google-vertex") diff --git a/providers/common/ai/tests/unit/common/ai/mixins/test_durable.py b/providers/common/ai/tests/unit/common/ai/mixins/test_durable.py new file mode 100644 index 0000000000000..7213933d7aae2 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/mixins/test_durable.py @@ -0,0 +1,121 @@ +# 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. +from __future__ import annotations + +import asyncio +import functools +import inspect +from unittest.mock import MagicMock + +from airflow.providers.common.ai.durable.step_counter import DurableStepCounter +from airflow.providers.common.ai.durable.storage import DurableStorage +from airflow.providers.common.ai.mixins.durable import DurableAgentMixin + + +class TestDurableAgentMixinCachedCallable: + def test_cached_callable_saves_and_returns(self): + storage = MagicMock(spec=DurableStorage) + counter = MagicMock(spec=DurableStepCounter) + counter.next_step.return_value = 1 + counter.cached_tool = 0 + storage.load_tool_result.return_value = (False, None) + + calls = [] + + def fn(): + calls.append(1) + return "computed" + + wrapped = DurableAgentMixin._cached_callable(fn, storage, counter) + result = wrapped() + + assert result == "computed" + assert calls == [1] + storage.save_tool_result.assert_called_once_with("tool_step_1", "computed") + + def test_cached_callable_replays_on_hit(self): + storage = MagicMock(spec=DurableStorage) + counter = MagicMock(spec=DurableStepCounter) + counter.replayed_tool = 0 + counter.next_step.return_value = 1 + storage.load_tool_result.return_value = (True, "cached_value") + + calls = [] + + def fn(): + calls.append(1) + return "computed" + + wrapped = DurableAgentMixin._cached_callable(fn, storage, counter) + result = wrapped() + + assert result == "cached_value" + assert calls == [] + assert counter.replayed_tool == 1 + storage.save_tool_result.assert_not_called() + + def test_cached_callable_preserves_async_function_behavior(self): + storage = MagicMock(spec=DurableStorage) + counter = MagicMock(spec=DurableStepCounter) + counter.cached_tool = 0 + counter.next_step.return_value = 1 + storage.load_tool_result.return_value = (False, None) + + async def fn(value): + return value * 2 + + wrapped = DurableAgentMixin._cached_callable(fn, storage, counter) + + assert inspect.iscoroutinefunction(wrapped) + assert asyncio.run(wrapped(3)) == 6 + storage.save_tool_result.assert_called_once_with("tool_step_1", 6) + assert counter.cached_tool == 1 + + def test_cached_callable_preserves_async_partial_behavior(self): + storage = MagicMock(spec=DurableStorage) + counter = MagicMock(spec=DurableStepCounter) + counter.cached_tool = 0 + counter.next_step.return_value = 1 + storage.load_tool_result.return_value = (False, None) + + async def fn(prefix, value): + return f"{prefix}:{value}" + + wrapped = DurableAgentMixin._cached_callable(functools.partial(fn, "prod"), storage, counter) + + assert inspect.iscoroutinefunction(wrapped) + assert asyncio.run(wrapped("cpu")) == "prod:cpu" + storage.save_tool_result.assert_called_once_with("tool_step_1", "prod:cpu") + assert counter.cached_tool == 1 + + def test_cached_callable_preserves_async_callable_object_behavior(self): + storage = MagicMock(spec=DurableStorage) + counter = MagicMock(spec=DurableStepCounter) + counter.cached_tool = 0 + counter.next_step.return_value = 1 + storage.load_tool_result.return_value = (False, None) + + class Lookup: + async def __call__(self, value): + return value.upper() + + wrapped = DurableAgentMixin._cached_callable(Lookup(), storage, counter) + + assert inspect.iscoroutinefunction(wrapped) + assert asyncio.run(wrapped("abc")) == "ABC" + storage.save_tool_result.assert_called_once_with("tool_step_1", "ABC") + assert counter.cached_tool == 1 diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 4a0b08cf1483e..1fdd677e86e8b 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -23,8 +23,14 @@ from pydantic import BaseModel from pydantic_ai.usage import UsageLimits +from airflow.providers.common.ai.hooks.base import ( + AgentRunRequest, + AgentRunResult, + AgentUsage, + BaseAIHook, + Capability, +) from airflow.providers.common.ai.operators.agent import AgentOperator, HITLReviewLink -from airflow.providers.common.ai.toolsets.logging import LoggingToolset from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS @@ -44,23 +50,48 @@ class Summary(BaseModel): score: float = 0.0 -def _make_mock_run_result(output): - """Create a mock AgentRunResult compatible with log_run_summary.""" - mock_result = MagicMock() - mock_result.output = output - mock_result.usage.return_value = MagicMock( - requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0 +def _make_agent_run_result(output, *, message_history=None) -> AgentRunResult: + return AgentRunResult( + output=output, + message_history=[] if message_history is None else message_history, + model_name="test-model", + usage=AgentUsage(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0), ) - mock_result.response = MagicMock(model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result -def _make_mock_agent(output): - """Create a mock agent that returns the given output.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(output) - return mock_agent +def _make_mock_hook(output, *, message_history=None): + """Return (mock_hook, mock_agent) wired for AgentOperator.execute.""" + mock_hook = MagicMock(spec=BaseAIHook) + mock_hook.llm_conn_id = "my_llm" + mock_hook.capabilities = frozenset({Capability.TOOLSETS, Capability.DURABLE, Capability.USAGE_LIMITS}) + mock_agent = MagicMock() + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_agent_run_result(output, message_history=message_history) + return mock_hook, mock_agent + + +class TestAgentOperatorHookCapabilities: + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) + def test_execute_rejects_toolsets_when_hook_does_not_support_them(self, mock_hook_cls): + mock_hook = MagicMock(spec=BaseAIHook) + mock_hook.llm_conn_id = "strands_conn" + mock_hook.capabilities = frozenset() # no capabilities — all features rejected + + def create_agent(request): + BaseAIHook.validate_run_request(mock_hook, request) + return MagicMock() + + mock_hook.create_agent.side_effect = create_agent + mock_hook_cls.get_agent_hook.return_value = mock_hook + + op = AgentOperator( + task_id="test", + prompt="test", + llm_conn_id="strands_conn", + toolsets=[MagicMock()], + ) + with pytest.raises(ValueError, match="toolsets not supported"): + op.execute(context=MagicMock()) class TestAgentOperatorValidation: @@ -95,11 +126,11 @@ def test_template_fields(self): class TestAgentOperatorExecute: - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_forwards_usage_limits_to_run_sync(self, mock_hook_cls): - """``usage_limits`` is forwarded to ``agent.run_sync`` on the non-durable path.""" - mock_agent = _make_mock_agent("ok") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + """``usage_limits`` is forwarded in the AgentRunRequest on the non-durable path.""" + mock_hook, mock_agent = _make_mock_hook("ok") + mock_hook_cls.get_agent_hook.return_value = mock_hook limits = UsageLimits(request_limit=3, tool_calls_limit=5) op = AgentOperator( @@ -110,13 +141,15 @@ def test_execute_forwards_usage_limits_to_run_sync(self, mock_hook_cls): ) op.execute(context=MagicMock()) - mock_agent.run_sync.assert_called_once_with("run", usage_limits=limits) + request = mock_hook.run_agent.call_args[0][1] + assert isinstance(request, AgentRunRequest) + assert request.usage_limits is limits - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_regenerate_with_feedback_forwards_usage_limits(self, mock_hook_cls): """``usage_limits`` is also forwarded by ``regenerate_with_feedback``.""" - mock_agent = _make_mock_agent("revised") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook, mock_agent = _make_mock_hook("revised") + mock_hook_cls.get_agent_hook.return_value = mock_hook limits = UsageLimits(request_limit=1) op = AgentOperator( @@ -127,16 +160,16 @@ def test_regenerate_with_feedback_forwards_usage_limits(self, mock_hook_cls): ) op.regenerate_with_feedback(feedback="Add detail", message_history=[]) - mock_agent.run_sync.assert_called_once_with( - "Add detail", - message_history=[], - usage_limits=limits, - ) + request = mock_hook.run_agent.call_args[0][1] + assert isinstance(request, AgentRunRequest) + assert request.prompt == "Add detail" + assert request.message_history == [] + assert request.usage_limits is limits - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_creates_agent_from_hook(self, mock_hook_cls): - mock_agent = _make_mock_agent("The answer is 42.") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook, mock_agent = _make_mock_hook("The answer is 42.") + mock_hook_cls.get_agent_hook.return_value = mock_hook op = AgentOperator( task_id="test", @@ -147,16 +180,24 @@ def test_execute_creates_agent_from_hook(self, mock_hook_cls): result = op.execute(context=MagicMock()) assert result == "The answer is 42." - mock_hook_cls.get_hook.assert_called_once_with("my_llm", hook_params={"model_id": None}) - mock_hook_cls.get_hook.return_value.create_agent.assert_called_once_with( - output_type=str, instructions="You are helpful." - ) - mock_agent.run_sync.assert_called_once_with("What is the answer?", usage_limits=None) + mock_hook_cls.get_agent_hook.assert_called_once_with("my_llm", hook_params={"model_id": None}) + + create_request = mock_hook.create_agent.call_args[0][0] + assert isinstance(create_request, AgentRunRequest) + assert create_request.output_type is str + assert create_request.instructions == "You are helpful." + assert create_request.prompt == "What is the answer?" + assert create_request.usage_limits is None - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + run_args = mock_hook.run_agent.call_args[0] + assert run_args[0] is mock_agent + assert run_args[1] is create_request + + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_passes_toolsets_in_agent_kwargs(self, mock_hook_cls): - """Toolsets are passed through to the agent constructor.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("done") + """Toolsets are passed through to create_agent in the request.""" + mock_hook, _ = _make_mock_hook("done") + mock_hook_cls.get_agent_hook.return_value = mock_hook mock_toolset = MagicMock() op = AgentOperator( @@ -167,16 +208,15 @@ def test_execute_passes_toolsets_in_agent_kwargs(self, mock_hook_cls): ) op.execute(context=MagicMock()) - create_call = mock_hook_cls.get_hook.return_value.create_agent.call_args - passed_toolsets = create_call[1]["toolsets"] - assert len(passed_toolsets) == 1 - assert isinstance(passed_toolsets[0], LoggingToolset) - assert passed_toolsets[0].wrapped is mock_toolset + request = mock_hook.create_agent.call_args[0][0] + assert request.toolsets == [mock_toolset] + assert request.enable_tool_logging is True - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_enable_tool_logging_false_skips_wrapping(self, mock_hook_cls): - """enable_tool_logging=False passes toolsets through unwrapped.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("done") + """enable_tool_logging=False is set on the request.""" + mock_hook, _ = _make_mock_hook("done") + mock_hook_cls.get_agent_hook.return_value = mock_hook mock_toolset = MagicMock() op = AgentOperator( @@ -188,13 +228,15 @@ def test_enable_tool_logging_false_skips_wrapping(self, mock_hook_cls): ) op.execute(context=MagicMock()) - create_call = mock_hook_cls.get_hook.return_value.create_agent.call_args - assert create_call[1]["toolsets"] == [mock_toolset] + request = mock_hook.create_agent.call_args[0][0] + assert request.toolsets == [mock_toolset] + assert request.enable_tool_logging is False - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_passes_agent_params(self, mock_hook_cls): - """agent_params are unpacked into create_agent.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("ok") + """agent_params are included in the request.""" + mock_hook, _ = _make_mock_hook("ok") + mock_hook_cls.get_agent_hook.return_value = mock_hook op = AgentOperator( task_id="test", @@ -204,17 +246,15 @@ def test_execute_passes_agent_params(self, mock_hook_cls): ) op.execute(context=MagicMock()) - create_call = mock_hook_cls.get_hook.return_value.create_agent.call_args - assert create_call[1]["retries"] == 3 - assert create_call[1]["model_settings"] == {"temperature": 0} + request = mock_hook.create_agent.call_args[0][0] + assert request.agent_params == {"retries": 3, "model_settings": {"temperature": 0}} @requires_typed_xcom - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_structured_output(self, mock_hook_cls): """Structured output keeps the Pydantic instance so downstream tasks can type-hint it.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent( - Summary(text="Great", score=0.95) - ) + mock_hook, _ = _make_mock_hook(Summary(text="Great", score=0.95)) + mock_hook_cls.get_agent_hook.return_value = mock_hook op = AgentOperator( task_id="test", @@ -232,10 +272,11 @@ def test_declares_output_type_for_deserialization(self): """Declares ``output_type`` so the worker-side DAG walk registers it for deserialization.""" assert "output_type" in AgentOperator.deserialization_allowed_class_fields - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_with_model_id(self, mock_hook_cls): - """model_id is passed to PydanticAIHook.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("ok") + """model_id is passed to the agent hook.""" + mock_hook, _ = _make_mock_hook("ok") + mock_hook_cls.get_agent_hook.return_value = mock_hook op = AgentOperator( task_id="test", @@ -245,21 +286,20 @@ def test_execute_with_model_id(self, mock_hook_cls): ) op.execute(context=MagicMock()) - mock_hook_cls.get_hook.assert_called_once_with("my_llm", hook_params={"model_id": "openai:gpt-5"}) + mock_hook_cls.get_agent_hook.assert_called_once_with( + "my_llm", hook_params={"model_id": "openai:gpt-5"} + ) @pytest.mark.skipif( not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" ) @patch("airflow.providers.common.ai.operators.agent.AgentOperator.run_hitl_review", autospec=True) - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_with_enable_hitl_review_delegates_to_run_hitl_review(self, mock_hook_cls, mock_run_hitl): """When enable_hitl_review=True, execute delegates to run_hitl_review with output and message_history.""" msg_history = [MagicMock()] - mock_result = _make_mock_run_result("Initial output") - mock_result.all_messages.return_value = msg_history - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = mock_result - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook, _ = _make_mock_hook("Initial output", message_history=msg_history) + mock_hook_cls.get_agent_hook.return_value = mock_hook mock_run_hitl.return_value = "Approved output" op = AgentOperator( @@ -280,14 +320,11 @@ def test_execute_with_enable_hitl_review_delegates_to_run_hitl_review(self, mock not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" ) @patch("airflow.providers.common.ai.operators.agent.AgentOperator.run_hitl_review", autospec=True) - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_with_hitl_rehydrates_base_model(self, mock_hook_cls, mock_run_hitl): """When enable_hitl_review=True and output_type is BaseModel, execute returns the model instance.""" - mock_result = _make_mock_run_result(Summary(text="Approved summary", score=0.9)) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = mock_result - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - # run_hitl_review returns JSON string (as stored in session.current_output) + mock_hook, _ = _make_mock_hook(Summary(text="Approved summary", score=0.9)) + mock_hook_cls.get_agent_hook.return_value = mock_hook mock_run_hitl.return_value = '{"text": "Approved summary", "score": 0.9}' op = AgentOperator( @@ -309,13 +346,11 @@ def test_execute_with_hitl_rehydrates_base_model(self, mock_hook_cls, mock_run_h not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" ) @patch("airflow.providers.common.ai.operators.agent.AgentOperator.run_hitl_review", autospec=True) - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_with_hitl_returns_string_unchanged(self, mock_hook_cls, mock_run_hitl): """When enable_hitl_review=True and output_type is str, execute returns string as-is.""" - mock_result = _make_mock_run_result("Initial output") - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = mock_result - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook, _ = _make_mock_hook("Initial output") + mock_hook_cls.get_agent_hook.return_value = mock_hook mock_run_hitl.return_value = "Approved output" op = AgentOperator( @@ -335,15 +370,13 @@ def test_execute_with_hitl_returns_string_unchanged(self, mock_hook_cls, mock_ru not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" ) @patch("airflow.providers.common.ai.operators.agent.AgentOperator.run_hitl_review", autospec=True) - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_propagates_hitl_max_iterations_error(self, mock_hook_cls, mock_run_hitl): """When run_hitl_review raises HITLMaxIterationsError, execute propagates it.""" from airflow.providers.common.ai.exceptions import HITLMaxIterationsError - mock_result = _make_mock_run_result("Initial output") - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = mock_result - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook, _ = _make_mock_hook("Initial output") + mock_hook_cls.get_agent_hook.return_value = mock_hook mock_run_hitl.side_effect = HITLMaxIterationsError("Task exceeded max iterations.") op = AgentOperator( @@ -407,42 +440,42 @@ def test_get_link_returns_url_with_params_when_hitl_enabled(self): not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" ) class TestAgentOperatorRegenerateWithFeedback: - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_regenerate_with_feedback_calls_agent_with_feedback_and_history(self, mock_hook_cls): - """regenerate_with_feedback builds agent and calls run_sync with feedback and message_history.""" + """regenerate_with_feedback builds request with feedback and message_history.""" msg_history = [MagicMock()] - mock_result = _make_mock_run_result("Revised output") - mock_result.all_messages.return_value = msg_history + [MagicMock()] - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = mock_result - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + new_history = msg_history + [MagicMock()] + mock_hook, mock_agent = _make_mock_hook("Revised output", message_history=new_history) + mock_hook_cls.get_agent_hook.return_value = mock_hook op = AgentOperator( task_id="test", prompt="Summarize", llm_conn_id="my_llm", ) - output, new_history = op.regenerate_with_feedback( + output, returned_history = op.regenerate_with_feedback( feedback="Add more detail", message_history=msg_history, ) assert output == "Revised output" - assert new_history == mock_result.all_messages.return_value - mock_agent.run_sync.assert_called_once_with( - "Add more detail", - message_history=msg_history, - usage_limits=None, - ) + assert returned_history == new_history - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + request = mock_hook.create_agent.call_args[0][0] + assert isinstance(request, AgentRunRequest) + assert request.prompt == "Add more detail" + assert request.message_history is msg_history + assert request.usage_limits is None + + run_args = mock_hook.run_agent.call_args[0] + assert run_args[0] is mock_agent + assert run_args[1] is request + + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_regenerate_with_feedback_serializes_base_model_output(self, mock_hook_cls): """regenerate_with_feedback returns JSON string for BaseModel output.""" - mock_result = _make_mock_run_result(Summary(text="Revised")) - mock_result.all_messages.return_value = [] - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = mock_result - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook, _ = _make_mock_hook(Summary(text="Revised")) + mock_hook_cls.get_agent_hook.return_value = mock_hook op = AgentOperator( task_id="test", @@ -467,68 +500,53 @@ def test_durable_default_false(self): op = AgentOperator(task_id="test", prompt="test", llm_conn_id="my_llm") assert op.durable is False - @patch("pydantic_ai.models.wrapper.infer_model", side_effect=lambda m: m) - @patch("pydantic_ai.models.infer_model", autospec=True) - @patch("airflow.providers.common.ai.durable.storage._get_base_path") - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_execute_durable_wraps_model_and_cleans_up( - self, mock_hook_cls, mock_base_path, mock_infer, _, tmp_path - ): - """durable=True wraps model with CachingModel and cleans up on success.""" - from airflow.sdk import ObjectStoragePath - - mock_base_path.return_value = ObjectStoragePath(f"file://{tmp_path.as_posix()}") - - mock_agent = MagicMock() - mock_agent.run_sync.return_value = _make_mock_run_result("ok") - mock_agent.model = "test-model" - mock_agent.override = MagicMock() - mock_agent.override.return_value.__enter__ = MagicMock(return_value=None) - mock_agent.override.return_value.__exit__ = MagicMock(return_value=False) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - - mock_resolved = MagicMock() - mock_infer.return_value = mock_resolved - + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) + def test_execute_durable_passes_durable_context_in_request(self, mock_hook_cls): + """durable=True builds DurableContext from task instance and passes it in request.""" + mock_hook, _ = _make_mock_hook("ok") + mock_hook_cls.get_agent_hook.return_value = mock_hook + + ti = MagicMock() + ti.dag_id = "my_dag" + ti.task_id = "my_task" + ti.run_id = "my_run" + ti.map_index = -1 context = MagicMock() - context.__getitem__ = MagicMock( - return_value=MagicMock(dag_id="d", task_id="t", run_id="r", map_index=-1) - ) + context.__getitem__ = MagicMock(return_value=ti) op = AgentOperator(task_id="test", prompt="test", llm_conn_id="my_llm", durable=True) result = op.execute(context=context) assert result == "ok" - mock_agent.override.assert_called_once() - override_kwargs = mock_agent.override.call_args[1] - assert "model" in override_kwargs - - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) + request = mock_hook.create_agent.call_args[0][0] + assert request.durable_context is not None + assert request.durable_context.dag_id == "my_dag" + assert request.durable_context.task_id == "my_task" + assert request.durable_context.run_id == "my_run" + assert request.durable_context.map_index == -1 + + @patch("airflow.providers.common.ai.operators.agent.BaseAIHook", autospec=True) def test_execute_non_durable_does_not_wrap(self, mock_hook_cls): - """Default (durable=False) does not use override.""" - mock_agent = _make_mock_agent("ok") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + """Default (durable=False) sets durable_context to None in request.""" + mock_hook, _ = _make_mock_hook("ok") + mock_hook_cls.get_agent_hook.return_value = mock_hook op = AgentOperator(task_id="test", prompt="test", llm_conn_id="my_llm") op.execute(context=MagicMock()) - # run_sync called directly, no override - mock_agent.run_sync.assert_called_once_with("test", usage_limits=None) + request = mock_hook.run_agent.call_args[0][1] + assert request.prompt == "test" + assert request.durable_context is None @pytest.mark.skipif( not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0" ) class TestAgentOperatorMultimodalPromptGuard: - """AgentOperator.execute raises before agent.run_sync when enable_hitl_review=True - and self.prompt is not a string -- covering direct construction and the native - template rendering escape (where a string template renders to a Sequence).""" - - @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) - def test_execute_rejects_sequence_prompt_with_hitl_review(self, mock_hook_cls): - mock_agent = MagicMock(spec=["run_sync"]) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + """AgentOperator.execute raises before run_agent when enable_hitl_review=True + and self.prompt is not a string.""" + def test_execute_rejects_sequence_prompt_with_hitl_review(self): op = AgentOperator( task_id="t", prompt="placeholder", @@ -539,5 +557,3 @@ def test_execute_rejects_sequence_prompt_with_hitl_review(self, mock_hook_cls): with pytest.raises(TypeError, match="enable_hitl_review=True"): op.execute(context=MagicMock()) - - mock_agent.run_sync.assert_not_called() diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py index 2a707752fdfe3..80da669cc1c2c 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py @@ -24,6 +24,12 @@ from pydantic import BaseModel from pydantic_ai.usage import UsageLimits +from airflow.providers.common.ai.hooks.base import ( + AgentRunRequest, + AgentRunResult, + AgentUsage, + BaseAIHook, +) from airflow.providers.common.ai.mixins.approval import ( LLMApprovalMixin, ) @@ -54,15 +60,12 @@ class Summary(BaseModel): def _make_mock_run_result(output): - """Create a mock AgentRunResult compatible with log_run_summary.""" - mock_result = MagicMock() - mock_result.output = output - mock_result.usage.return_value = MagicMock( - requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0 + """Create an AgentRunResult compatible with log_run_summary.""" + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0), ) - mock_result.response = MagicMock(model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result class TestLLMOperator: @@ -70,29 +73,50 @@ def test_template_fields(self): expected = {"prompt", "llm_conn_id", "model_id", "system_prompt", "agent_params"} assert set(LLMOperator.template_fields) == expected - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_returns_string_output(self, mock_hook_cls): """Default output_type=str returns the LLM string directly.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("Paris is the capital of France.") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock() + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_mock_run_result("Paris is the capital of France.") op = LLMOperator(task_id="test", prompt="What is the capital of France?", llm_conn_id="my_llm") result = op.execute(context=MagicMock()) assert result == "Paris is the capital of France." - mock_agent.run_sync.assert_called_once_with("What is the capital of France?", usage_limits=None) - mock_hook_cls.get_hook.return_value.create_agent.assert_called_once_with( - output_type=str, instructions="" + request = mock_hook.create_agent.call_args[0][0] + assert isinstance(request, AgentRunRequest) + assert request.prompt == "What is the capital of France?" + assert request.output_type is str + assert request.instructions == "" + assert request.usage_limits is None + mock_hook.run_agent.assert_called_once_with(mock_agent, request) + mock_hook_cls.get_agent_hook.assert_called_once_with("my_llm", hook_params={"model_id": None}) + + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) + def test_execute_rejects_usage_limits_when_hook_unsupported(self, mock_hook_cls): + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.llm_conn_id = "my_llm" + mock_hook.capabilities = frozenset() # no capabilities — usage_limits will be rejected + mock_hook.create_agent.side_effect = lambda req: BaseAIHook.validate_run_request(mock_hook, req) + + op = LLMOperator( + task_id="test", + prompt="Summarize", + llm_conn_id="my_llm", + usage_limits=UsageLimits(request_limit=1), ) - mock_hook_cls.get_hook.assert_called_once_with("my_llm", hook_params={"model_id": None}) + with pytest.raises(ValueError, match="usage_limits not supported"): + op.execute(context=MagicMock()) - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_forwards_usage_limits_to_run_sync(self, mock_hook_cls): - """``usage_limits`` is forwarded verbatim to ``agent.run_sync``.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("ok") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) + def test_execute_forwards_usage_limits_to_run_agent(self, mock_hook_cls): + """``usage_limits`` is forwarded verbatim to ``hook.run_agent`` via AgentRunRequest.""" + mock_agent = MagicMock() + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_mock_run_result("ok") limits = UsageLimits(request_limit=2, output_tokens_limit=100) op = LLMOperator( @@ -103,15 +127,17 @@ def test_execute_forwards_usage_limits_to_run_sync(self, mock_hook_cls): ) op.execute(context=MagicMock()) - mock_agent.run_sync.assert_called_once_with("Summarize", usage_limits=limits) + request = mock_hook.create_agent.call_args[0][0] + assert request.usage_limits is limits @requires_typed_xcom - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_structured_output_with_all_params(self, mock_hook_cls): """Structured output returns the Pydantic instance unchanged so downstream tasks keep the type.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(Entities(names=["Alice", "Bob"])) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock() + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_mock_run_result(Entities(names=["Alice", "Bob"])) op = LLMOperator( task_id="test", @@ -126,13 +152,13 @@ def test_execute_structured_output_with_all_params(self, mock_hook_cls): assert isinstance(result, Entities) assert result.names == ["Alice", "Bob"] - mock_hook_cls.get_hook.assert_called_once_with("my_llm", hook_params={"model_id": "openai:gpt-5"}) - mock_hook_cls.get_hook.return_value.create_agent.assert_called_once_with( - output_type=Entities, - instructions="You are an extractor.", - retries=3, - model_settings={"temperature": 0.9}, + mock_hook_cls.get_agent_hook.assert_called_once_with( + "my_llm", hook_params={"model_id": "openai:gpt-5"} ) + request = mock_hook.create_agent.call_args[0][0] + assert request.output_type is Entities + assert request.instructions == "You are an extractor." + assert request.agent_params == {"retries": 3, "model_settings": {"temperature": 0.9}} def test_declares_output_type_for_deserialization(self): """Declares ``output_type`` so the worker-side DAG walk registers it for deserialization. @@ -142,12 +168,13 @@ def test_declares_output_type_for_deserialization(self): """ assert "output_type" in LLMOperator.deserialization_allowed_class_fields - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_serialize_output_returns_dict(self, mock_hook_cls): """serialize_output=True dumps the BaseModel to a dict on the wire.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(Entities(names=["A", "B"])) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock() + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_mock_run_result(Entities(names=["A", "B"])) op = LLMOperator( task_id="t", @@ -186,14 +213,15 @@ def test_default_approval_flags(self): @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_with_approval_defers(self, mock_hook_cls, mock_upsert, mock_trigger_cls): """When require_approval=True, execute() defers instead of returning output.""" from airflow.providers.common.compat.sdk import TaskDeferred - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("LLM response") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock() + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_mock_run_result("LLM response") op = LLMOperator( task_id="approval_test", @@ -212,14 +240,15 @@ def test_execute_with_approval_defers(self, mock_hook_cls, mock_upsert, mock_tri @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_with_approval_and_modifications(self, mock_hook_cls, mock_upsert, mock_trigger_cls): """allow_modifications=True passes an editable 'output' param.""" from airflow.providers.common.compat.sdk import TaskDeferred - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("draft output") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock() + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_mock_run_result("draft output") op = LLMOperator( task_id="mod_test", @@ -238,14 +267,15 @@ def test_execute_with_approval_and_modifications(self, mock_hook_cls, mock_upser @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_with_approval_and_timeout(self, mock_hook_cls, mock_upsert, mock_trigger_cls): """approval_timeout is passed to the trigger.""" from airflow.providers.common.compat.sdk import TaskDeferred - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("output") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock() + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_mock_run_result("output") timeout = timedelta(hours=1) op = LLMOperator( @@ -264,14 +294,15 @@ def test_execute_with_approval_and_timeout(self, mock_hook_cls, mock_upsert, moc @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_with_approval_structured_output(self, mock_hook_cls, mock_upsert, mock_trigger_cls): """Structured (BaseModel) output is serialized before deferring.""" from airflow.providers.common.compat.sdk import TaskDeferred - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(Summary(text="hello")) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock() + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_mock_run_result(Summary(text="hello")) op = LLMOperator( task_id="struct_test", @@ -287,12 +318,13 @@ def test_execute_with_approval_structured_output(self, mock_hook_cls, mock_upser assert exc_info.value.kwargs["generated_output"] == '{"text":"hello"}' - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_without_approval_returns_normally(self, mock_hook_cls): """When require_approval=False, execute() returns output directly.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("plain output") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock() + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_mock_run_result("plain output") op = LLMOperator(task_id="no_approval", prompt="p", llm_conn_id="my_llm", require_approval=False) result = op.execute(context={}) @@ -361,11 +393,8 @@ class TestLLMOperatorMultimodalPromptGuard: and self.prompt is not a string -- covering direct-operator construction and the native template rendering escape (where a string template renders to a Sequence).""" - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_rejects_sequence_prompt_with_require_approval(self, mock_hook_cls): - mock_agent = MagicMock(spec=["run_sync"]) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook.get_agent_hook", autospec=True) + def test_execute_rejects_sequence_prompt_with_require_approval(self, mock_get_agent_hook): op = LLMOperator( task_id="t", prompt="placeholder", @@ -377,4 +406,4 @@ def test_execute_rejects_sequence_prompt_with_require_approval(self, mock_hook_c with pytest.raises(TypeError, match="require_approval=True"): op.execute(context=_make_context()) - mock_agent.run_sync.assert_not_called() + mock_get_agent_hook.assert_not_called() diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py index 82f710cc10a0b..2092d4b7cd63c 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py @@ -21,20 +21,17 @@ import pytest +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator -def _make_mock_run_result(output): - """Create a mock AgentRunResult compatible with log_run_summary.""" - mock_result = MagicMock() - mock_result.output = output - mock_result.usage.return_value = MagicMock( - requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0 +def _make_run_result(output): + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0), ) - mock_result.response = MagicMock(model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result class TestLLMBranchOperator: @@ -57,14 +54,14 @@ def test_output_type_ignored(self): assert op.output_type is str @patch.object(LLMBranchOperator, "do_branch") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_single_branch(self, mock_hook_cls, mock_do_branch): """LLM returns a single enum member → do_branch receives a string.""" downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a", "task_b": "task_b"}) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(downstream_enum.task_a) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = _make_run_result(downstream_enum.task_a) mock_do_branch.return_value = "task_a" op = LLMBranchOperator( @@ -79,21 +76,18 @@ def test_execute_single_branch(self, mock_hook_cls, mock_do_branch): assert result == "task_a" mock_do_branch.assert_called_once_with(ctx, "task_a") - mock_agent.run_sync.assert_called_once_with("Pick a branch", usage_limits=None) @patch.object(LLMBranchOperator, "do_branch") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_multi_branch(self, mock_hook_cls, mock_do_branch): """allow_multiple_branches=True → LLM returns list of enums → do_branch receives list.""" downstream_enum = Enum( "DownstreamTasks", {"task_a": "task_a", "task_b": "task_b", "task_c": "task_c"} ) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result( - [downstream_enum.task_a, downstream_enum.task_c] - ) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = _make_run_result([downstream_enum.task_a, downstream_enum.task_c]) mock_do_branch.return_value = ["task_a", "task_c"] op = LLMBranchOperator( @@ -111,14 +105,14 @@ def test_execute_multi_branch(self, mock_hook_cls, mock_do_branch): mock_do_branch.assert_called_once_with(ctx, ["task_a", "task_c"]) @patch.object(LLMBranchOperator, "do_branch") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_system_prompt_forwarded(self, mock_hook_cls, mock_do_branch): - """system_prompt is passed to create_agent(instructions=...).""" + """system_prompt is passed as AgentRunRequest.instructions.""" downstream_enum = Enum("DownstreamTasks", {"task_a": "task_a"}) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(downstream_enum.task_a) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = _make_run_result(downstream_enum.task_a) op = LLMBranchOperator( task_id="test", @@ -130,20 +124,20 @@ def test_system_prompt_forwarded(self, mock_hook_cls, mock_do_branch): op.execute(MagicMock()) - call_kwargs = mock_hook_cls.get_hook.return_value.create_agent.call_args - assert call_kwargs.kwargs["instructions"] == "Route tickets to the right team." + request = mock_hook.create_agent.call_args[0][0] + assert request.instructions == "Route tickets to the right team." @patch.object(LLMBranchOperator, "do_branch") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_downstream_task_ids_used_for_enum(self, mock_hook_cls, mock_do_branch): """The dynamic enum is built from self.downstream_task_ids.""" downstream_enum = Enum( "DownstreamTasks", {"billing": "billing", "auth": "auth", "general": "general"} ) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(downstream_enum.billing) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = _make_run_result(downstream_enum.billing) op = LLMBranchOperator( task_id="test", @@ -154,8 +148,8 @@ def test_downstream_task_ids_used_for_enum(self, mock_hook_cls, mock_do_branch): op.execute(MagicMock()) - output_type = mock_hook_cls.get_hook.return_value.create_agent.call_args.kwargs["output_type"] - assert {m.value for m in output_type} == {"billing", "auth", "general"} + request = mock_hook.create_agent.call_args[0][0] + assert {m.value for m in request.output_type} == {"billing", "auth", "general"} def test_execute_raises_on_no_downstream_tasks(self): """ValueError when the operator has no downstream tasks.""" diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py index 6c970e4326383..53df6bc0e6e04 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_file_analysis.py @@ -23,6 +23,8 @@ import pytest from pydantic import BaseModel +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage +from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAgentHandle from airflow.providers.common.ai.operators.llm_file_analysis import LLMFileAnalysisOperator from airflow.providers.common.ai.utils.file_analysis import FileAnalysisRequest @@ -44,19 +46,11 @@ class Summary(BaseModel): def _make_mock_run_result(output): - mock_result = MagicMock(spec=["output", "usage", "response", "all_messages"]) - mock_result.output = output - mock_result.usage.return_value = MagicMock( - spec=["requests", "tool_calls", "input_tokens", "output_tokens", "total_tokens"], - requests=1, - tool_calls=0, - input_tokens=0, - output_tokens=0, - total_tokens=0, + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0), ) - mock_result.response = MagicMock(spec=["model_name"], model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result def _make_context(ti_id=None): @@ -81,7 +75,7 @@ def test_template_fields(self): } assert set(LLMFileAnalysisOperator.template_fields) == expected - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) @patch( "airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request", autospec=True ) @@ -91,9 +85,11 @@ def test_execute_returns_string_output(self, mock_build_request, mock_hook_cls): resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("Analysis complete") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock(spec=PydanticAgentHandle) + mock_hook_cls.get_agent_hook.return_value.create_agent.return_value = mock_agent + mock_hook_cls.get_agent_hook.return_value.run_agent.return_value = _make_mock_run_result( + "Analysis complete" + ) op = LLMFileAnalysisOperator( task_id="test", @@ -115,10 +111,13 @@ def test_execute_returns_string_output(self, mock_build_request, mock_hook_cls): max_text_chars=100_000, sample_rows=10, ) - mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None) + mock_hook = mock_hook_cls.get_agent_hook.return_value + request = mock_hook.create_agent.call_args[0][0] + assert request.prompt == "prepared prompt" + mock_hook.run_agent.assert_called_once_with(mock_agent, request) @requires_typed_xcom - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) @patch( "airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request", autospec=True ) @@ -128,9 +127,11 @@ def test_execute_structured_output_returns_pydantic_instance(self, mock_build_re resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(Summary(findings=["error spike"])) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock(spec=PydanticAgentHandle) + mock_hook_cls.get_agent_hook.return_value.create_agent.return_value = mock_agent + mock_hook_cls.get_agent_hook.return_value.run_agent.return_value = _make_mock_run_result( + Summary(findings=["error spike"]) + ) op = LLMFileAnalysisOperator( task_id="test", @@ -144,7 +145,7 @@ def test_execute_structured_output_returns_pydantic_instance(self, mock_build_re assert isinstance(result, Summary) assert result.findings == ["error spike"] - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) @patch( "airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request", autospec=True ) @@ -155,9 +156,10 @@ def test_execute_serialize_output_returns_dict(self, mock_build_request, mock_ho resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(Summary(findings=["error spike"])) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock(spec=PydanticAgentHandle) + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = mock_agent + mock_hook.run_agent.return_value = _make_mock_run_result(Summary(findings=["error spike"])) op = LLMFileAnalysisOperator( task_id="test", @@ -201,7 +203,7 @@ def test_parameter_validation(self, mock_build_request): class TestLLMFileAnalysisOperatorApproval: @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) @patch( "airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request", autospec=True ) @@ -215,9 +217,11 @@ def test_execute_with_approval_defers( resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("LLM response") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock(spec=PydanticAgentHandle) + mock_hook_cls.get_agent_hook.return_value.create_agent.return_value = mock_agent + mock_hook_cls.get_agent_hook.return_value.run_agent.return_value = _make_mock_run_result( + "LLM response" + ) op = LLMFileAnalysisOperator( task_id="approval_test", @@ -237,7 +241,7 @@ def test_execute_with_approval_defers( @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) @patch( "airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request", autospec=True ) @@ -251,9 +255,11 @@ def test_execute_with_approval_defers_structured_output_as_json( resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(Summary(findings=["error spike"])) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock(spec=PydanticAgentHandle) + mock_hook_cls.get_agent_hook.return_value.create_agent.return_value = mock_agent + mock_hook_cls.get_agent_hook.return_value.run_agent.return_value = _make_mock_run_result( + Summary(findings=["error spike"]) + ) op = LLMFileAnalysisOperator( task_id="approval_structured_test", @@ -311,7 +317,7 @@ def test_execute_complete_with_approval_restores_modified_structured_output(self @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) @patch( "airflow.providers.common.ai.operators.llm_file_analysis.build_file_analysis_request", autospec=True ) @@ -325,9 +331,9 @@ def test_execute_with_approval_timeout( resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result("output") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + mock_agent = MagicMock(spec=PydanticAgentHandle) + mock_hook_cls.get_agent_hook.return_value.create_agent.return_value = mock_agent + mock_hook_cls.get_agent_hook.return_value.run_agent.return_value = _make_mock_run_result("output") timeout = timedelta(hours=1) op = LLMFileAnalysisOperator( diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py index d719162bc00f7..4271eb9a31368 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py @@ -21,6 +21,7 @@ import pytest +from airflow.providers.common.ai.hooks.base import AgentRunRequest, AgentRunResult, AgentUsage from airflow.providers.common.ai.operators.llm_schema_compare import ( LLMSchemaCompareOperator, SchemaCompareResult, @@ -31,16 +32,13 @@ from airflow.providers.common.sql.hooks.sql import DbApiHook -def _make_mock_run_result(output): - """Create a mock AgentRunResult compatible with log_run_summary.""" - mock_result = MagicMock() - mock_result.output = output - mock_result.usage.return_value = MagicMock( - requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0 +def _make_run_result(output): + """Create an AgentRunResult compatible with log_run_summary.""" + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0), ) - mock_result.response = MagicMock(model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result _BASE_KWARGS = dict(task_id="test_task", prompt="test prompt", llm_conn_id="llm_conn") @@ -260,22 +258,22 @@ def test_execute(self, mock_build_system_prompt, mock_build_schema_context): mock_llm_hook = mock.Mock() mock_agent = mock.Mock() - mock_agent.run_sync.return_value = _make_mock_run_result( + mock_llm_hook.create_agent.return_value = mock_agent + mock_llm_hook.run_agent.return_value = _make_run_result( SchemaCompareResult(compatible=True, mismatches=[], summary="All good") ) - mock_llm_hook.create_agent.return_value = mock_agent op.llm_hook = mock_llm_hook result = op.execute(context={}) mock_build_schema_context.assert_called_once() mock_build_system_prompt.assert_called_once_with("schema_context") - mock_llm_hook.create_agent.assert_called_once_with( - output_type=SchemaCompareResult, - instructions="system_prompt", - param="value", - ) - mock_agent.run_sync.assert_called_once_with("user_prompt", usage_limits=None) + request = mock_llm_hook.create_agent.call_args[0][0] + assert isinstance(request, AgentRunRequest) + assert request.output_type is SchemaCompareResult + assert request.instructions == "system_prompt" + assert request.agent_params == {"param": "value"} + mock_llm_hook.run_agent.assert_called_once_with(mock_agent, request) assert result == {"compatible": True, "mismatches": [], "summary": "All good"} @mock.patch( @@ -336,26 +334,23 @@ def test_execute_schema_comparison_mixed_conn(self, mock_get_db_hook, db_hook): mock_llm_hook = mock.Mock() mock_agent = mock.Mock() - mock_agent.run_sync.return_value = _make_mock_run_result( + mock_llm_hook.create_agent.return_value = mock_agent + mock_llm_hook.run_agent.return_value = _make_run_result( SchemaCompareResult( compatible=True, mismatches=[], summary="S3 and Postgres schemas are compatible" ) ) - mock_llm_hook.create_agent.return_value = mock_agent op.llm_hook = mock_llm_hook with mock.patch.object(op, "_build_schema_context", return_value=schema_context): result = op.execute(context={}) - instructions = mock_llm_hook.create_agent.call_args[1]["instructions"] - assert "schema comparison expert" in instructions - assert "postgresql" in instructions - assert "aws_default" in instructions + request = mock_llm_hook.create_agent.call_args[0][0] + assert "schema comparison expert" in request.instructions + assert "postgresql" in request.instructions + assert "aws_default" in request.instructions - mock_agent.run_sync.assert_called_once_with( - "Compare S3 Parquet schema against the Postgres table and flag breaking changes", - usage_limits=None, - ) + mock_llm_hook.run_agent.assert_called_once_with(mock_agent, request) assert result["compatible"] is True assert result["summary"] == "S3 and Postgres schemas are compatible" @@ -415,18 +410,18 @@ def test_execute_schema_comparison_db_conn_ids_only(self, mock_get_db_hook): mock_llm_hook = mock.Mock() mock_agent = mock.Mock() - mock_agent.run_sync.return_value = _make_mock_run_result( + mock_llm_hook.create_agent.return_value = mock_agent + mock_llm_hook.run_agent.return_value = _make_run_result( SchemaCompareResult(compatible=True, mismatches=[], summary="Schemas are compatible") ) - mock_llm_hook.create_agent.return_value = mock_agent op.llm_hook = mock_llm_hook with mock.patch.object(op, "_build_schema_context", return_value=schema_context): result = op.execute(context={}) - instructions = mock_llm_hook.create_agent.call_args[1]["instructions"] - assert "postgresql" in instructions - assert "snowflake" in instructions + request = mock_llm_hook.create_agent.call_args[0][0] + assert "postgresql" in request.instructions + assert "snowflake" in request.instructions assert result["compatible"] is True def test_execute_schema_comparison_datasources_only(self): @@ -464,22 +459,22 @@ def test_execute_schema_comparison_datasources_only(self): mock_llm_hook = mock.Mock() mock_agent = mock.Mock() - mock_agent.run_sync.return_value = _make_mock_run_result( + mock_llm_hook.create_agent.return_value = mock_agent + mock_llm_hook.run_agent.return_value = _make_run_result( SchemaCompareResult( compatible=False, mismatches=[], summary="Timestamp column type differs between Parquet and CSV", ) ) - mock_llm_hook.create_agent.return_value = mock_agent op.llm_hook = mock_llm_hook with mock.patch.object(op, "_build_schema_context", return_value=schema_context): result = op.execute(context={}) - instructions = mock_llm_hook.create_agent.call_args[1]["instructions"] - assert "aws_lake" in instructions - assert "aws_staging" in instructions + request = mock_llm_hook.create_agent.call_args[0][0] + assert "aws_lake" in request.instructions + assert "aws_staging" in request.instructions assert result["compatible"] is False def test_introspect_full_schema(self, db_hook): diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py index e8e31c6f5de88..b2dffb2d821b9 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py @@ -22,6 +22,7 @@ import pytest +from airflow.providers.common.ai.hooks.base import AgentRunRequest, AgentRunResult, AgentUsage from airflow.providers.common.ai.mixins.approval import ( LLMApprovalMixin, ) @@ -33,23 +34,20 @@ from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS -def _make_mock_run_result(output): - """Create a mock AgentRunResult compatible with log_run_summary.""" - mock_result = MagicMock() - mock_result.output = output - mock_result.usage.return_value = MagicMock( - requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0 +def _make_run_result(output): + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1, tool_calls=0, input_tokens=0, output_tokens=0, total_tokens=0), ) - mock_result.response = MagicMock(model_name="test-model") - mock_result.all_messages.return_value = [] - return mock_result -def _make_mock_agent(output: str): - """Create a mock agent that returns the given output string.""" - mock_agent = MagicMock(spec=["run_sync"]) - mock_agent.run_sync.return_value = _make_mock_run_result(output) - return mock_agent +def _setup_mock_hook(mock_hook_cls, output: str): + """Configure BaseAIHook mock to return the given SQL output from run_agent.""" + mock_hook = mock_hook_cls.get_agent_hook.return_value + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = _make_run_result(output) + return mock_hook class TestStripLLMOutput: @@ -96,11 +94,10 @@ def test_template_fields_include_parent_and_sql_specific(self): } assert set(LLMSQLQueryOperator.template_fields) == expected - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_with_schema_context(self, mock_hook_cls): """Operator uses schema_context and returns generated SQL.""" - mock_agent = _make_mock_agent("SELECT id, name FROM users WHERE active = true") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + _setup_mock_hook(mock_hook_cls, "SELECT id, name FROM users WHERE active = true") op = LLMSQLQueryOperator( task_id="test", @@ -111,32 +108,35 @@ def test_execute_with_schema_context(self, mock_hook_cls): result = op.execute(context=MagicMock()) assert result == "SELECT id, name FROM users WHERE active = true" - mock_agent.run_sync.assert_called_once_with("Get active users", usage_limits=None) + mock_hook = mock_hook_cls.get_agent_hook.return_value + request = mock_hook.create_agent.call_args[0][0] + assert isinstance(request, AgentRunRequest) + assert request.prompt == "Get active users" - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_validation_blocks_unsafe_sql(self, mock_hook_cls): """Validation catches unsafe SQL generated by the LLM.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("DROP TABLE users") + _setup_mock_hook(mock_hook_cls, "DROP TABLE users") op = LLMSQLQueryOperator(task_id="test", prompt="Delete everything", llm_conn_id="my_llm") with pytest.raises(SQLSafetyError, match="not allowed"): op.execute(context=MagicMock()) - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_validation_disabled(self, mock_hook_cls): """When validate_sql=False, unsafe SQL is returned without checks.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("DROP TABLE users") + _setup_mock_hook(mock_hook_cls, "DROP TABLE users") op = LLMSQLQueryOperator(task_id="test", prompt="Drop it", llm_conn_id="my_llm", validate_sql=False) result = op.execute(context=MagicMock()) assert result == "DROP TABLE users" - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_passes_agent_params(self, mock_hook_cls): - """agent_params inherited from LLMOperator are unpacked into create_agent.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("SELECT 1") + """agent_params are forwarded in the AgentRunRequest.""" + _setup_mock_hook(mock_hook_cls, "SELECT 1") op = LLMSQLQueryOperator( task_id="test", @@ -146,14 +146,13 @@ def test_execute_passes_agent_params(self, mock_hook_cls): ) op.execute(context=MagicMock()) - create_agent_call = mock_hook_cls.get_hook.return_value.create_agent.call_args - assert create_agent_call[1]["retries"] == 3 - assert create_agent_call[1]["model_settings"] == {"temperature": 0} + request = mock_hook_cls.get_agent_hook.return_value.create_agent.call_args[0][0] + assert request.agent_params == {"retries": 3, "model_settings": {"temperature": 0}} - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_system_prompt_appended_to_sql_instructions(self, mock_hook_cls): """User-provided system_prompt is appended to built-in SQL safety prompt.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("SELECT 1") + _setup_mock_hook(mock_hook_cls, "SELECT 1") op = LLMSQLQueryOperator( task_id="test", @@ -163,19 +162,17 @@ def test_system_prompt_appended_to_sql_instructions(self, mock_hook_cls): ) op.execute(context=MagicMock()) - instructions = mock_hook_cls.get_hook.return_value.create_agent.call_args[1]["instructions"] - assert "Always use LEFT JOINs." in instructions - # Built-in SQL safety prompt should still be present - assert "Generate only SELECT queries" in instructions - assert "Never generate data modification" in instructions + request = mock_hook_cls.get_agent_hook.return_value.create_agent.call_args[0][0] + assert "Always use LEFT JOINs." in request.instructions + assert "Generate only SELECT queries" in request.instructions + assert "Never generate data modification" in request.instructions class TestLLMSQLQueryOperatorSchemaIntrospection: - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_introspect_schemas_via_db_hook(self, mock_hook_cls): """db_conn_id + table_names triggers schema introspection.""" - mock_agent = _make_mock_agent("SELECT id FROM users") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + _setup_mock_hook(mock_hook_cls, "SELECT id FROM users") mock_db_hook = MagicMock(spec=["get_table_schema", "dialect_name"]) mock_db_hook.get_table_schema.return_value = [ @@ -198,10 +195,9 @@ def test_introspect_schemas_via_db_hook(self, mock_hook_cls): assert result == "SELECT id FROM users" mock_db_hook.get_table_schema.assert_called_once_with("users") - # Verify the system prompt contains the schema info - instructions = mock_hook_cls.get_hook.return_value.create_agent.call_args[1]["instructions"] - assert "users" in instructions - assert "id INTEGER" in instructions + request = mock_hook_cls.get_agent_hook.return_value.create_agent.call_args[0][0] + assert "users" in request.instructions + assert "id INTEGER" in request.instructions def test_introspect_raises_when_no_tables_found(self): """Raise ValueError when all requested tables return empty columns.""" @@ -351,7 +347,7 @@ def test_introspect_schemas_raises_when_no_tables_and_no_datasource(self, mock_e with pytest.raises(ValueError, match="None of the requested tables"): op._introspect_schemas() - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) @patch( "airflow.providers.common.ai.operators.llm_sql.DataFusionEngine", autospec=True, @@ -361,8 +357,7 @@ def test_execute_with_datasource_config_and_db_tables(self, mock_engine_cls, moc mock_engine = mock_engine_cls.return_value mock_engine.get_schema.return_value = "event: TEXT\nts: TIMESTAMP" - mock_agent = _make_mock_agent("SELECT u.id, e.event FROM users u JOIN events e ON u.id = e.user_id") - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent + _setup_mock_hook(mock_hook_cls, "SELECT u.id, e.event FROM users u JOIN events e ON u.id = e.user_id") ds_config = DataSourceConfig( conn_id="aws_default", @@ -390,10 +385,10 @@ def test_execute_with_datasource_config_and_db_tables(self, mock_engine_cls, moc result = op.execute(context=MagicMock()) assert "SELECT" in result - instructions = mock_hook_cls.get_hook.return_value.create_agent.call_args[1]["instructions"] - assert "users" in instructions - assert "events" in instructions - assert "event: TEXT\nts: TIMESTAMP" in instructions + request = mock_hook_cls.get_agent_hook.return_value.create_agent.call_args[0][0] + assert "users" in request.instructions + assert "events" in request.instructions + assert "event: TEXT\nts: TIMESTAMP" in request.instructions class TestLLMSQLQueryOperatorDialect: @@ -460,13 +455,10 @@ def test_approval_flags_default_values(self): @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_with_approval_defers(self, mock_hook_cls, mock_upsert, mock_trigger_cls): """When require_approval=True, execute() defers after generating and validating SQL.""" - - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent( - "SELECT id FROM users WHERE active" - ) + _setup_mock_hook(mock_hook_cls, "SELECT id FROM users WHERE active") op = LLMSQLQueryOperator( task_id="sql_approval", @@ -486,12 +478,12 @@ def test_execute_with_approval_defers(self, mock_hook_cls, mock_upsert, mock_tri @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_with_approval_validates_before_deferring( self, mock_hook_cls, mock_upsert, mock_trigger_cls ): """SQL validation runs before defer_for_approval; unsafe SQL is blocked.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("DROP TABLE users") + _setup_mock_hook(mock_hook_cls, "DROP TABLE users") op = LLMSQLQueryOperator( task_id="sql_unsafe", @@ -508,11 +500,10 @@ def test_execute_with_approval_validates_before_deferring( @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_with_approval_and_modifications(self, mock_hook_cls, mock_upsert, mock_trigger_cls): """allow_modifications=True passes editable params.""" - - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("SELECT 1") + _setup_mock_hook(mock_hook_cls, "SELECT 1") op = LLMSQLQueryOperator( task_id="sql_mod", @@ -531,11 +522,10 @@ def test_execute_with_approval_and_modifications(self, mock_hook_cls, mock_upser @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_with_approval_and_timeout(self, mock_hook_cls, mock_upsert, mock_trigger_cls): """approval_timeout is propagated to the trigger.""" - - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("SELECT 1") + _setup_mock_hook(mock_hook_cls, "SELECT 1") timeout = timedelta(minutes=30) op = LLMSQLQueryOperator( @@ -552,10 +542,10 @@ def test_execute_with_approval_and_timeout(self, mock_hook_cls, mock_upsert, moc assert exc_info.value.timeout == timeout - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_without_approval_returns_sql(self, mock_hook_cls): """When require_approval=False, execute() returns the SQL directly.""" - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent("SELECT 1") + _setup_mock_hook(mock_hook_cls, "SELECT 1") op = LLMSQLQueryOperator( task_id="no_approval", @@ -569,13 +559,10 @@ def test_execute_without_approval_returns_sql(self, mock_hook_cls): @patch("airflow.providers.standard.triggers.hitl.HITLTrigger", autospec=True) @patch("airflow.sdk.execution_time.hitl.upsert_hitl_detail") - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) def test_execute_strips_code_fences_before_deferring(self, mock_hook_cls, mock_upsert, mock_trigger_cls): """Markdown code fences are stripped from LLM output before deferring.""" - - mock_hook_cls.get_hook.return_value.create_agent.return_value = _make_mock_agent( - "```sql\nSELECT 1\n```" - ) + _setup_mock_hook(mock_hook_cls, "```sql\nSELECT 1\n```") op = LLMSQLQueryOperator( task_id="strip_test", @@ -664,11 +651,8 @@ class TestLLMSQLQueryOperatorMultimodalPromptGuard: """LLMSQLQueryOperator.execute raises before agent.run_sync when require_approval=True and self.prompt is not a string.""" - @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) - def test_execute_rejects_sequence_prompt_with_require_approval(self, mock_hook_cls): - mock_agent = MagicMock(spec=["run_sync"]) - mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent - + @patch("airflow.providers.common.ai.operators.llm.BaseAIHook.get_agent_hook", autospec=True) + def test_execute_rejects_sequence_prompt_with_require_approval(self, mock_get_agent_hook): op = LLMSQLQueryOperator( task_id="t", prompt="placeholder", @@ -680,4 +664,4 @@ def test_execute_rejects_sequence_prompt_with_require_approval(self, mock_hook_c with pytest.raises(TypeError, match="require_approval=True"): op.execute(context=_make_context()) - mock_agent.run_sync.assert_not_called() + mock_get_agent_hook.assert_not_called() diff --git a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py index 6f9d976d6f1aa..abf2ffb862a96 100644 --- a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py +++ b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py @@ -25,6 +25,7 @@ # Skip the entire test module on older Airflow versions tested in compat CI. pytest.importorskip("airflow.sdk.definitions.retry_policy", reason="RetryPolicy requires Airflow 3.3+") +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.policies.retry import ( ErrorClassification, LLMRetryPolicy, @@ -32,27 +33,30 @@ from airflow.sdk.definitions.retry_policy import RetryAction, RetryRule -def _make_mock_agent(category, should_retry, delay=0, reasoning="test"): - """Create a mock agent that returns a canned ErrorClassification.""" - mock_result = MagicMock() - mock_result.output = ErrorClassification( - category=category, - should_retry=should_retry, - suggested_delay_seconds=delay, - reasoning=reasoning, +def _make_run_result(output): + return AgentRunResult( + output=output, + model_name="test-model", + usage=AgentUsage(requests=1), ) - mock_agent = MagicMock() - mock_agent.run_sync.return_value = mock_result - return mock_agent + + +def _make_mock_hook(run_result): + mock_hook = MagicMock() + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.return_value = run_result + return mock_hook class TestLLMClassifyDecisions: """Test that _classify maps LLM classification to correct RetryDecisions.""" - @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook", autospec=True) - def test_auth_error_returns_fail(self, mock_hook_cls): - mock_hook_cls.return_value.create_agent.return_value = _make_mock_agent( - "auth", should_retry=False, reasoning="API key expired" + @patch.object(BaseAIHook, "get_agent_hook") + def test_auth_error_returns_fail(self, mock_get_hook): + mock_get_hook.return_value = _make_mock_hook( + _make_run_result( + ErrorClassification(category="auth", should_retry=False, reasoning="API key expired") + ) ) policy = LLMRetryPolicy(llm_conn_id="test") decision = policy.evaluate(PermissionError("403"), try_number=1, max_tries=3) @@ -61,10 +65,14 @@ def test_auth_error_returns_fail(self, mock_hook_cls): assert "auth" in decision.reason assert "API key expired" in decision.reason - @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook", autospec=True) - def test_rate_limit_returns_retry_with_delay(self, mock_hook_cls): - mock_hook_cls.return_value.create_agent.return_value = _make_mock_agent( - "rate_limit", should_retry=True, delay=60, reasoning="429" + @patch.object(BaseAIHook, "get_agent_hook") + def test_rate_limit_returns_retry_with_delay(self, mock_get_hook): + mock_get_hook.return_value = _make_mock_hook( + _make_run_result( + ErrorClassification( + category="rate_limit", should_retry=True, suggested_delay_seconds=60, reasoning="429" + ) + ) ) policy = LLMRetryPolicy(llm_conn_id="test") decision = policy.evaluate(RuntimeError("429"), try_number=1, max_tries=3) @@ -72,11 +80,15 @@ def test_rate_limit_returns_retry_with_delay(self, mock_hook_cls): assert decision.action == RetryAction.RETRY assert decision.retry_delay == timedelta(seconds=60) - @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook", autospec=True) - def test_transient_retry_with_zero_delay_uses_default(self, mock_hook_cls): + @patch.object(BaseAIHook, "get_agent_hook") + def test_transient_retry_with_zero_delay_uses_default(self, mock_get_hook): """suggested_delay_seconds=0 means use the task's default delay, not override.""" - mock_hook_cls.return_value.create_agent.return_value = _make_mock_agent( - "transient", should_retry=True, delay=0 + mock_get_hook.return_value = _make_mock_hook( + _make_run_result( + ErrorClassification( + category="transient", should_retry=True, suggested_delay_seconds=0, reasoning="glitch" + ) + ) ) policy = LLMRetryPolicy(llm_conn_id="test") decision = policy.evaluate(RuntimeError("glitch"), try_number=1, max_tries=3) @@ -84,11 +96,15 @@ def test_transient_retry_with_zero_delay_uses_default(self, mock_hook_cls): assert decision.action == RetryAction.RETRY assert decision.retry_delay is None # None = use task's default - @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook", autospec=True) - def test_negative_delay_treated_as_no_override(self, mock_hook_cls): + @patch.object(BaseAIHook, "get_agent_hook") + def test_negative_delay_treated_as_no_override(self, mock_get_hook): """Negative delay from LLM should not produce a negative timedelta.""" - mock_hook_cls.return_value.create_agent.return_value = _make_mock_agent( - "transient", should_retry=True, delay=-5 + mock_get_hook.return_value = _make_mock_hook( + _make_run_result( + ErrorClassification( + category="transient", should_retry=True, suggested_delay_seconds=-5, reasoning="x" + ) + ) ) policy = LLMRetryPolicy(llm_conn_id="test") decision = policy.evaluate(RuntimeError("x"), try_number=1, max_tries=3) @@ -96,39 +112,46 @@ def test_negative_delay_treated_as_no_override(self, mock_hook_cls): assert decision.action == RetryAction.RETRY assert decision.retry_delay is None - @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook", autospec=True) - def test_prompt_includes_exception_type_and_message(self, mock_hook_cls): - mock_agent = _make_mock_agent("data", should_retry=False) - mock_hook_cls.return_value.create_agent.return_value = mock_agent + @patch.object(BaseAIHook, "get_agent_hook") + def test_prompt_includes_exception_type_and_message(self, mock_get_hook): + mock_hook = _make_mock_hook( + _make_run_result(ErrorClassification(category="data", should_retry=False, reasoning="test")) + ) + mock_get_hook.return_value = mock_hook policy = LLMRetryPolicy(llm_conn_id="test") policy.evaluate(ValueError("bad column type"), try_number=2, max_tries=5) - prompt = mock_agent.run_sync.call_args[0][0] - assert "ValueError: bad column type" in prompt - assert "attempt 2 of 5" in prompt + request = mock_hook.create_agent.call_args[0][0] + assert "ValueError: bad column type" in request.prompt + assert "attempt 2 of 5" in request.prompt - @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook", autospec=True) - def test_custom_instructions_forwarded_to_agent(self, mock_hook_cls): - mock_hook_cls.return_value.create_agent.return_value = _make_mock_agent("x", False) + @patch.object(BaseAIHook, "get_agent_hook") + def test_custom_instructions_forwarded_to_agent(self, mock_get_hook): + mock_hook = _make_mock_hook( + _make_run_result(ErrorClassification(category="x", should_retry=False, reasoning="test")) + ) + mock_get_hook.return_value = mock_hook policy = LLMRetryPolicy(llm_conn_id="test", instructions="My custom prompt") policy.evaluate(ValueError("x"), try_number=1, max_tries=3) - mock_hook_cls.return_value.create_agent.assert_called_once_with( - output_type=ErrorClassification, - instructions="My custom prompt", - ) + request = mock_hook.create_agent.call_args[0][0] + assert request.instructions == "My custom prompt" + assert request.output_type is ErrorClassification - @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook", autospec=True) - def test_timeout_passed_via_model_settings(self, mock_hook_cls): - mock_agent = _make_mock_agent("auth", False) - mock_hook_cls.return_value.create_agent.return_value = mock_agent + @patch.object(BaseAIHook, "get_agent_hook") + def test_timeout_passed_via_model_settings(self, mock_get_hook): + mock_hook = _make_mock_hook( + _make_run_result(ErrorClassification(category="auth", should_retry=False, reasoning="test")) + ) + mock_get_hook.return_value = mock_hook policy = LLMRetryPolicy(llm_conn_id="test", timeout=15.0) policy.evaluate(ValueError("x"), try_number=1, max_tries=3) - model_settings = mock_agent.run_sync.call_args.kwargs["model_settings"] + request = mock_hook.create_agent.call_args[0][0] + model_settings = request.agent_params["model_settings"] assert model_settings["timeout"] == 15.0 @@ -169,12 +192,13 @@ def test_fallback_rules_no_match_returns_default(self): d = policy.evaluate(ValueError("bad"), try_number=1, max_tries=3) assert d.action == RetryAction.DEFAULT - @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook", autospec=True) - def test_agent_run_sync_failure_triggers_fallback(self, mock_hook_cls): - """Failure during run_sync (not hook creation) still triggers fallback.""" - mock_agent = MagicMock() - mock_agent.run_sync.side_effect = RuntimeError("network error mid-call") - mock_hook_cls.return_value.create_agent.return_value = mock_agent + @patch.object(BaseAIHook, "get_agent_hook") + def test_run_agent_failure_triggers_fallback(self, mock_get_hook): + """Failure during run_agent (not hook creation) still triggers fallback.""" + mock_hook = MagicMock() + mock_hook.create_agent.return_value = MagicMock() + mock_hook.run_agent.side_effect = RuntimeError("network error mid-call") + mock_get_hook.return_value = mock_hook policy = LLMRetryPolicy( llm_conn_id="test", @@ -184,10 +208,12 @@ def test_agent_run_sync_failure_triggers_fallback(self, mock_hook_cls): assert d.action == RetryAction.FAIL assert d.reason == "fallback" - @patch("airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook", autospec=True) - def test_hook_creation_failure_triggers_fallback(self, mock_hook_cls): + @patch.object(BaseAIHook, "get_agent_hook") + def test_hook_creation_failure_triggers_fallback(self, mock_get_hook): """Failure during hook.create_agent still triggers fallback.""" - mock_hook_cls.return_value.create_agent.side_effect = RuntimeError("unexpected") + mock_hook = MagicMock() + mock_hook.create_agent.side_effect = RuntimeError("unexpected") + mock_get_hook.return_value = mock_hook policy = LLMRetryPolicy( llm_conn_id="test", diff --git a/providers/common/ai/tests/unit/common/ai/test_observability.py b/providers/common/ai/tests/unit/common/ai/test_observability.py index bfb6f0fbe6aa3..dac4ee49c3e52 100644 --- a/providers/common/ai/tests/unit/common/ai/test_observability.py +++ b/providers/common/ai/tests/unit/common/ai/test_observability.py @@ -25,6 +25,7 @@ from pydantic_ai.models.test import TestModel from airflow.providers.common.ai import observability +from airflow.providers.common.ai.hooks.base import AgentRunRequest from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook @@ -106,13 +107,17 @@ def _run(*, capture: bool): with ( patch.object(observability, "conf", _conf(enabled=True, capture=capture)), patch.object(observability, "_live_tracer_provider", return_value=provider), - patch.object(hook, "get_conn", return_value=TestModel()), + patch.object(hook, "get_model", return_value=TestModel()), ): - agent = hook.create_agent(instructions="be helpful") + request = AgentRunRequest( + prompt=TestEndToEndSpanEmission._PROMPT, + instructions="be helpful", + ) + agent = hook.create_agent(request) # Stand in for the worker's task span: open a parent and run inside it. with provider.get_tracer("test").start_as_current_span("worker.task") as parent: parent_trace_id = parent.get_span_context().trace_id - agent.run_sync(TestEndToEndSpanEmission._PROMPT) + hook.run_agent(agent, request) spans = exporter.get_finished_spans() genai = [s for s in spans if s.attributes and any(k.startswith("gen_ai.") for k in s.attributes)] diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_langchain_bridge.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_langchain_bridge.py index 6187f4e9d4c13..590962e08792b 100644 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_langchain_bridge.py +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_langchain_bridge.py @@ -29,6 +29,7 @@ from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool from pydantic_core import SchemaValidator, core_schema +from airflow.providers.common.ai.hooks.base import BaseToolset, ToolSpec from airflow.providers.common.ai.toolsets.langchain_bridge import airflow_toolset_to_langchain_tools _PASSTHROUGH = SchemaValidator(core_schema.any_schema()) @@ -182,6 +183,95 @@ def test_missing_langchain_raises_optional_feature_exception(self, monkeypatch): airflow_toolset_to_langchain_tools(FakeToolset()) +_TEXT_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], +} + + +class FakeSyncBaseToolset(BaseToolset): + def as_tools(self) -> list[ToolSpec]: + return [ + ToolSpec( + name="greet", + description="Greet someone.", + parameters=_TEXT_SCHEMA, + fn=lambda text: f"hello {text}", + ), + ToolSpec( + name="retry_me", + description="Always retries.", + parameters=_TEXT_SCHEMA, + fn=self._boom, + ), + ] + + @staticmethod + def _boom(text: str) -> str: + raise ModelRetry("please try again") + + +class FakeAsyncBaseToolset(BaseToolset): + def as_tools(self) -> list[ToolSpec]: + return [ + ToolSpec( + name="greet_async", + description="Async greet.", + parameters=_TEXT_SCHEMA, + fn=self._greet, + ), + ToolSpec( + name="retry_async", + description="Always retries (async).", + parameters=_TEXT_SCHEMA, + fn=self._boom, + ), + ] + + @staticmethod + async def _greet(text: str) -> str: + return f"async hello {text}" + + @staticmethod + async def _boom(text: str) -> str: + raise ModelRetry("async please try again") + + +class TestBaseToolsetConversion: + def test_sync_fn_sync_invoke(self): + tools = {t.name: t for t in airflow_toolset_to_langchain_tools(FakeSyncBaseToolset())} + assert tools["greet"].invoke({"text": "world"}) == "hello world" + + def test_sync_fn_async_invoke(self): + tools = {t.name: t for t in airflow_toolset_to_langchain_tools(FakeSyncBaseToolset())} + assert asyncio.run(tools["greet"].ainvoke({"text": "world"})) == "hello world" + + def test_async_fn_sync_invoke(self): + tools = {t.name: t for t in airflow_toolset_to_langchain_tools(FakeAsyncBaseToolset())} + assert tools["greet_async"].invoke({"text": "world"}) == "async hello world" + + def test_async_fn_async_invoke(self): + tools = {t.name: t for t in airflow_toolset_to_langchain_tools(FakeAsyncBaseToolset())} + assert asyncio.run(tools["greet_async"].ainvoke({"text": "world"})) == "async hello world" + + def test_sync_fn_model_retry_returned_as_output_sync(self): + tools = {t.name: t for t in airflow_toolset_to_langchain_tools(FakeSyncBaseToolset())} + assert tools["retry_me"].invoke({"text": "x"}) == "please try again" + + def test_sync_fn_model_retry_returned_as_output_async(self): + tools = {t.name: t for t in airflow_toolset_to_langchain_tools(FakeSyncBaseToolset())} + assert asyncio.run(tools["retry_me"].ainvoke({"text": "x"})) == "please try again" + + def test_async_fn_model_retry_returned_as_output_sync(self): + tools = {t.name: t for t in airflow_toolset_to_langchain_tools(FakeAsyncBaseToolset())} + assert tools["retry_async"].invoke({"text": "x"}) == "async please try again" + + def test_async_fn_model_retry_returned_as_output_async(self): + tools = {t.name: t for t in airflow_toolset_to_langchain_tools(FakeAsyncBaseToolset())} + assert asyncio.run(tools["retry_async"].ainvoke({"text": "x"})) == "async please try again" + + class TestSQLToolsetConversion: def test_sql_toolset_exposes_its_four_tools(self): # get_tools / construction do not touch the database, so no connection diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py index 5e425597a32e1..94b34565f0544 100644 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py @@ -16,7 +16,6 @@ # under the License. from __future__ import annotations -import asyncio import importlib.util import json import sqlite3 @@ -25,6 +24,7 @@ import pytest from pydantic_ai.exceptions import ModelRetry +from airflow.providers.common.ai.hooks.base import BaseToolset, ToolSpec from airflow.providers.common.ai.toolsets.sql import SQLToolset from airflow.providers.common.ai.utils.sql_validation import SQLSafetyError from airflow.providers.common.sql.hooks.sql import DbApiHook @@ -50,65 +50,88 @@ def _make_mock_db_hook( class TestSQLToolsetInit: - def test_id_includes_conn_id(self): + def test_is_base_toolset(self): + assert issubclass(SQLToolset, BaseToolset) + + def test_defaults(self): ts = SQLToolset("my_pg") - assert ts.id == "sql-my_pg" + assert ts._db_conn_id == "my_pg" + assert ts._allowed_tables is None + assert ts._schema is None + assert ts._allow_writes is False + assert ts._max_rows == 50 + + +class TestSQLToolsetAsTools: + def test_returns_four_tool_specs(self): + ts = SQLToolset("pg_default") + tools = ts.as_tools() + assert len(tools) == 4 + assert all(isinstance(t, ToolSpec) for t in tools) + def test_tool_names(self): + ts = SQLToolset("pg_default") + names = [t.name for t in ts.as_tools()] + assert names == ["list_tables", "get_schema", "query", "check_query"] + + def test_tool_descriptions_non_empty(self): + ts = SQLToolset("pg_default") + for spec in ts.as_tools(): + assert spec.description -class TestSQLToolsetGetTools: - def test_returns_four_tools(self): + def test_tool_callables_are_bound_methods(self): + ts = SQLToolset("pg_default") + specs = ts.as_tools() + fns = {s.name: s.fn for s in specs} + assert fns["list_tables"] == ts._list_tables + assert fns["get_schema"] == ts._get_schema + assert fns["query"] == ts._query + assert fns["check_query"] == ts._check_query + + def test_tool_parameters_match_schemas(self): ts = SQLToolset("pg_default") - tools = asyncio.run(ts.get_tools(ctx=MagicMock())) - assert set(tools.keys()) == {"list_tables", "get_schema", "query", "check_query"} + specs = {s.name: s for s in ts.as_tools()} + assert specs["list_tables"].parameters["type"] == "object" + assert "table_name" in specs["get_schema"].parameters["properties"] + assert "sql" in specs["query"].parameters["properties"] + assert "sql" in specs["check_query"].parameters["properties"] - def test_tool_definitions_have_descriptions(self): + def test_tools_are_sequential(self): ts = SQLToolset("pg_default") - tools = asyncio.run(ts.get_tools(ctx=MagicMock())) - for tool in tools.values(): - assert tool.tool_def.description + for spec in ts.as_tools(): + assert spec.sequential is True class TestSQLToolsetListTables: def test_returns_all_tables(self): ts = SQLToolset("pg_default") - mock_hook = _make_mock_db_hook(table_names=["users", "orders", "products"]) - ts._hook = mock_hook + ts._hook = _make_mock_db_hook(table_names=["users", "orders", "products"]) - result = asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock())) - tables = json.loads(result) + tables = json.loads(ts._list_tables()) assert tables == ["users", "orders", "products"] def test_filters_by_allowed_tables(self): ts = SQLToolset("pg_default", allowed_tables=["orders"]) - mock_hook = _make_mock_db_hook(table_names=["users", "orders", "products"]) - ts._hook = mock_hook + ts._hook = _make_mock_db_hook(table_names=["users", "orders", "products"]) - result = asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock())) - tables = json.loads(result) + tables = json.loads(ts._list_tables()) assert tables == ["orders"] class TestSQLToolsetGetSchema: def test_returns_column_info(self): ts = SQLToolset("pg_default") - mock_hook = _make_mock_db_hook() - ts._hook = mock_hook + ts._hook = _make_mock_db_hook() - result = asyncio.run( - ts.call_tool("get_schema", {"table_name": "users"}, ctx=MagicMock(), tool=MagicMock()) - ) - columns = json.loads(result) + columns = json.loads(ts._get_schema("users")) assert columns == [{"name": "id", "type": "INTEGER"}, {"name": "name", "type": "VARCHAR"}] - mock_hook.get_table_schema.assert_called_once_with("users", schema=None) + ts._hook.get_table_schema.assert_called_once_with("users", schema=None) def test_blocks_table_not_in_allowed_list(self): ts = SQLToolset("pg_default", allowed_tables=["orders"]) ts._hook = _make_mock_db_hook() - result = asyncio.run( - ts.call_tool("get_schema", {"table_name": "secrets"}, ctx=MagicMock(), tool=MagicMock()) - ) - data = json.loads(result) + data = json.loads(ts._get_schema("secrets")) assert "error" in data assert "secrets" in data["error"] @@ -121,10 +144,7 @@ def test_returns_rows_as_json(self): last_description=[("id",), ("name",)], ) - result = asyncio.run( - ts.call_tool("query", {"sql": "SELECT id, name FROM users"}, ctx=MagicMock(), tool=MagicMock()) - ) - data = json.loads(result) + data = json.loads(ts._query("SELECT id, name FROM users")) assert data["rows"] == [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}] assert data["count"] == 2 @@ -135,10 +155,7 @@ def test_truncates_at_max_rows(self): last_description=[("id",), ("name",)], ) - result = asyncio.run( - ts.call_tool("query", {"sql": "SELECT id, name FROM users"}, ctx=MagicMock(), tool=MagicMock()) - ) - data = json.loads(result) + data = json.loads(ts._query("SELECT id, name FROM users")) assert len(data["rows"]) == 1 assert data["truncated"] is True assert data["count"] == 3 @@ -148,58 +165,39 @@ def test_blocks_unsafe_sql_by_default(self): ts._hook = _make_mock_db_hook() with pytest.raises(SQLSafetyError, match="not allowed"): - asyncio.run(ts.call_tool("query", {"sql": "DROP TABLE users"}, ctx=MagicMock(), tool=MagicMock())) + ts._query("DROP TABLE users") def test_allows_writes_when_enabled(self): ts = SQLToolset("pg_default", allow_writes=True) - ts._hook = _make_mock_db_hook( - records=[(1,)], - last_description=[("count",)], - ) + ts._hook = _make_mock_db_hook(records=[(1,)], last_description=[("count",)]) - # Should not raise even with INSERT - result = asyncio.run( - ts.call_tool( - "query", {"sql": "INSERT INTO users VALUES (3, 'Eve')"}, ctx=MagicMock(), tool=MagicMock() - ) - ) - # The mock doesn't actually execute, just returns mocked records - data = json.loads(result) + data = json.loads(ts._query("INSERT INTO users VALUES (3, 'Eve')")) assert "rows" in data def test_raises_model_retry_when_query_fails_with_retryable_error(self): - """When the query fails with a retryable error, raise ModelRetry so the model retries.""" ts = SQLToolset("pg_default") ts._hook = _make_mock_db_hook() ts._hook.conn_type = "sqlite" ts._hook.get_records.side_effect = sqlite3.OperationalError("no such column: nonexistent") with pytest.raises(ModelRetry) as exc_info: - asyncio.run( - ts.call_tool( - "query", - {"sql": "SELECT id, nonexistent FROM users"}, - ctx=MagicMock(), - tool=MagicMock(), - ) - ) - assert "nonexistent" in exc_info.value.message - assert "get_schema" in exc_info.value.message - assert "list_tables" in exc_info.value.message + ts._query("SELECT id, nonexistent FROM users") + + assert "nonexistent" in exc_info.value.args[0] + assert "get_schema" in exc_info.value.args[0] + assert "list_tables" in exc_info.value.args[0] def test_model_retry_message_includes_schema_hint(self): - """ModelRetry message tells the model to use get_schema and list_tables for more details.""" ts = SQLToolset("pg_default") ts._hook = _make_mock_db_hook() ts._hook.conn_type = "sqlite" ts._hook.get_records.side_effect = sqlite3.OperationalError("no such table: missing_table") with pytest.raises(ModelRetry) as exc_info: - asyncio.run( - ts.call_tool("query", {"sql": "SELECT foo FROM x"}, ctx=MagicMock(), tool=MagicMock()) - ) - assert "get_schema" in exc_info.value.message - assert "list_tables" in exc_info.value.message + ts._query("SELECT foo FROM x") + + assert "get_schema" in exc_info.value.args[0] + assert "list_tables" in exc_info.value.args[0] def test_non_retryable_error_is_propagated(self): ts = SQLToolset("pg_default") @@ -208,7 +206,7 @@ def test_non_retryable_error_is_propagated(self): ts._hook.get_records.side_effect = sqlite3.OperationalError("database is locked") with pytest.raises(sqlite3.OperationalError, match="database is locked"): - asyncio.run(ts.call_tool("query", {"sql": "SELECT 1"}, ctx=MagicMock(), tool=MagicMock())) + ts._query("SELECT 1") def test_error_propagates_when_hook_conn_type_not_supported(self): ts = SQLToolset("pg_default") @@ -217,7 +215,7 @@ def test_error_propagates_when_hook_conn_type_not_supported(self): ts._hook.get_records.side_effect = RuntimeError("unexpected db error") with pytest.raises(RuntimeError, match="unexpected db error"): - asyncio.run(ts.call_tool("query", {"sql": "SELECT 1"}, ctx=MagicMock(), tool=MagicMock())) + ts._query("SELECT 1") def test_error_propagates_when_hook_has_no_conn_type(self): ts = SQLToolset("pg_default") @@ -227,7 +225,7 @@ def test_error_propagates_when_hook_has_no_conn_type(self): ts._hook = mock_hook with pytest.raises(RuntimeError, match="hook error"): - asyncio.run(ts.call_tool("query", {"sql": "SELECT 1"}, ctx=MagicMock(), tool=MagicMock())) + ts._query("SELECT 1") @pytest.mark.skipif( importlib.util.find_spec("psycopg2") is None, @@ -259,14 +257,7 @@ def test_sqlalchemy_programming_error_with_psycopg2_undefined_column_orig_raises ), pytest.raises(ModelRetry), ): - asyncio.run( - ts.call_tool( - "query", - {"sql": "SELECT id, missing FROM users"}, - ctx=MagicMock(), - tool=MagicMock(), - ) - ) + ts._query("SELECT id, missing FROM users") @pytest.mark.skipif( importlib.util.find_spec("psycopg2") is None, @@ -298,35 +289,18 @@ def test_sqlalchemy_programming_error_with_psycopg2_insufficient_privilege_orig_ ), pytest.raises(ProgrammingError), ): - asyncio.run( - ts.call_tool( - "query", - {"sql": "SELECT id FROM users"}, - ctx=MagicMock(), - tool=MagicMock(), - ) - ) + ts._query("SELECT id FROM users") class TestSQLToolsetCheckQuery: def test_valid_select(self): ts = SQLToolset("pg_default") - ts._hook = _make_mock_db_hook() - - result = asyncio.run( - ts.call_tool("check_query", {"sql": "SELECT 1"}, ctx=MagicMock(), tool=MagicMock()) - ) - data = json.loads(result) + data = json.loads(ts._check_query("SELECT 1")) assert data["valid"] is True def test_invalid_sql(self): ts = SQLToolset("pg_default") - ts._hook = _make_mock_db_hook() - - result = asyncio.run( - ts.call_tool("check_query", {"sql": "DROP TABLE users"}, ctx=MagicMock(), tool=MagicMock()) - ) - data = json.loads(result) + data = json.loads(ts._check_query("DROP TABLE users")) assert data["valid"] is False assert "error" in data @@ -348,7 +322,7 @@ def test_lazy_resolves_db_hook(self, mock_base_hook): @patch("airflow.providers.common.ai.toolsets.sql.BaseHook", autospec=True) def test_raises_for_non_dbapi_hook(self, mock_base_hook): mock_conn = MagicMock(spec=["get_hook"]) - mock_conn.get_hook.return_value = MagicMock() # Not a DbApiHook + mock_conn.get_hook.return_value = MagicMock() mock_base_hook.get_connection.return_value = mock_conn ts = SQLToolset("bad_conn") @@ -367,7 +341,6 @@ def test_caches_hook_after_first_resolution(self, mock_base_hook): ts._get_db_hook() ts._get_db_hook() - # Only called once because result is cached. mock_base_hook.get_connection.assert_called_once() @@ -394,7 +367,7 @@ def test_list_tables_spans_multiple_schemas(self): } ) - result = json.loads(asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock()))) + result = json.loads(ts._list_tables()) assert result == ["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS", "MODEL_CRM.SF_ASTRO_ORGS"] def test_list_tables_never_introspects_none_schema_when_all_qualified(self): @@ -402,7 +375,7 @@ def test_list_tables_never_introspects_none_schema_when_all_qualified(self): ts = SQLToolset("sf", allowed_tables=["MODEL_ASTRO.X", "MODEL_CRM.Y"]) ts._hook = self._schema_aware_hook({"MODEL_ASTRO": ["X"], "MODEL_CRM": ["Y"]}) - asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock())) + ts._list_tables() called_schemas = {c.kwargs.get("schema") for c in ts._hook.inspector.get_table_names.call_args_list} assert called_schemas == {"MODEL_ASTRO", "MODEL_CRM"} @@ -412,7 +385,7 @@ def test_list_tables_mixed_qualified_and_default(self): ts = SQLToolset("pg", allowed_tables=["users", "MODEL_ASTRO.X"], schema="public") ts._hook = self._schema_aware_hook({"public": ["users", "orders"], "MODEL_ASTRO": ["X", "Z"]}) - result = json.loads(asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock()))) + result = json.loads(ts._list_tables()) # Qualified schemas listed first (sorted), then the default schema. assert result == ["MODEL_ASTRO.X", "users"] @@ -420,16 +393,7 @@ def test_get_schema_routes_to_qualified_schema(self): ts = SQLToolset("sf", allowed_tables=["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS"]) ts._hook = self._schema_aware_hook({"MODEL_ASTRO": ["DEPLOYMENT_IMAGE_DETAILS"]}) - result = json.loads( - asyncio.run( - ts.call_tool( - "get_schema", - {"table_name": "MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS"}, - ctx=MagicMock(), - tool=MagicMock(), - ) - ) - ) + result = json.loads(ts._get_schema("MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS")) assert result == [{"name": "id", "type": "INTEGER"}] ts._hook.get_table_schema.assert_called_once_with("DEPLOYMENT_IMAGE_DETAILS", schema="MODEL_ASTRO") @@ -437,13 +401,7 @@ def test_get_schema_blocks_table_outside_allowed_schema(self): ts = SQLToolset("sf", allowed_tables=["MODEL_ASTRO.X"]) ts._hook = self._schema_aware_hook({"MODEL_ASTRO": ["X"]}) - result = json.loads( - asyncio.run( - ts.call_tool( - "get_schema", {"table_name": "SECRETS.PASSWORDS"}, ctx=MagicMock(), tool=MagicMock() - ) - ) - ) + result = json.loads(ts._get_schema("SECRETS.PASSWORDS")) assert "error" in result ts._hook.get_table_schema.assert_not_called() @@ -451,7 +409,7 @@ def test_get_schema_unqualified_uses_default_schema(self): ts = SQLToolset("pg", schema="public") ts._hook = self._schema_aware_hook({"public": ["users"]}) - asyncio.run(ts.call_tool("get_schema", {"table_name": "users"}, ctx=MagicMock(), tool=MagicMock())) + ts._get_schema("users") ts._hook.get_table_schema.assert_called_once_with("users", schema="public") def test_list_tables_matches_case_insensitively(self): @@ -467,23 +425,14 @@ def test_list_tables_matches_case_insensitively(self): } ) - result = json.loads(asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock()))) + result = json.loads(ts._list_tables()) assert result == ["MODEL_ASTRO.deployment_image_details", "MODEL_CRM.sf_astro_orgs"] def test_get_schema_matches_case_insensitively(self): ts = SQLToolset("sf", allowed_tables=["MODEL_ASTRO.DEPLOYMENT_IMAGE_DETAILS"]) ts._hook = self._schema_aware_hook({"MODEL_ASTRO": ["deployment_image_details"]}) - result = json.loads( - asyncio.run( - ts.call_tool( - "get_schema", - {"table_name": "MODEL_ASTRO.deployment_image_details"}, - ctx=MagicMock(), - tool=MagicMock(), - ) - ) - ) + result = json.loads(ts._get_schema("MODEL_ASTRO.deployment_image_details")) assert "error" not in result ts._hook.get_table_schema.assert_called_once_with("deployment_image_details", schema="MODEL_ASTRO") @@ -492,7 +441,7 @@ def test_list_tables_deduplicates_same_table(self): ts = SQLToolset("pg", allowed_tables=["public.users", "users"], schema="public") ts._hook = self._schema_aware_hook({"public": ["users"]}) - result = json.loads(asyncio.run(ts.call_tool("list_tables", {}, ctx=MagicMock(), tool=MagicMock()))) + result = json.loads(ts._list_tables()) assert result == ["public.users"] @@ -507,10 +456,7 @@ def test_describe_allowed_through_query(self): last_description=[("column_name",), ("data_type",)], ) - result = asyncio.run( - ts.call_tool("query", {"sql": "DESCRIBE TABLE users"}, ctx=MagicMock(), tool=MagicMock()) - ) - data = json.loads(result) + data = json.loads(ts._query("DESCRIBE TABLE users")) assert "rows" in data ts._hook.get_records.assert_called_once_with("DESCRIBE TABLE users") @@ -520,8 +466,7 @@ def test_show_allowed_with_snowflake_dialect(self): ts._hook = _make_mock_db_hook(records=[("USERS",)], last_description=[("name",)]) ts._hook.dialect_name = "snowflake" - result = asyncio.run(ts.call_tool("query", {"sql": "SHOW TABLES"}, ctx=MagicMock(), tool=MagicMock())) - data = json.loads(result) + data = json.loads(ts._query("SHOW TABLES")) assert "rows" in data ts._hook.get_records.assert_called_once_with("SHOW TABLES") @@ -537,22 +482,18 @@ def test_query_blocks_disallowed_statements(self, sql): ts._hook.dialect_name = "postgresql" with pytest.raises(SQLSafetyError, match="not allowed"): - asyncio.run(ts.call_tool("query", {"sql": sql}, ctx=MagicMock(), tool=MagicMock())) + ts._query(sql) def test_check_query_accepts_describe(self): ts = SQLToolset("pg_default") ts._hook = _make_mock_db_hook() - result = asyncio.run( - ts.call_tool("check_query", {"sql": "DESCRIBE TABLE users"}, ctx=MagicMock(), tool=MagicMock()) - ) + result = ts._check_query("DESCRIBE TABLE users") assert json.loads(result)["valid"] is True def test_check_query_handles_unresolvable_connection(self): """check_query stays usable (dialect-agnostic) when the connection can't be resolved.""" ts = SQLToolset("missing_conn") with patch.object(ts, "_get_db_hook", side_effect=RuntimeError("no such connection")): - result = asyncio.run( - ts.call_tool("check_query", {"sql": "SELECT 1"}, ctx=MagicMock(), tool=MagicMock()) - ) + result = ts._check_query("SELECT 1") assert json.loads(result)["valid"] is True diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_callables.py b/providers/common/ai/tests/unit/common/ai/utils/test_callables.py new file mode 100644 index 0000000000000..11a8cf0182111 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/utils/test_callables.py @@ -0,0 +1,67 @@ +# 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. +from __future__ import annotations + +import functools + +import pytest + +from airflow.providers.common.ai.utils.callables import is_async_callable + + +async def async_fn(value): + return value + + +def sync_fn(value): + return value + + +class TestIsAsyncCallable: + @pytest.mark.parametrize( + "fn", + [ + async_fn, + functools.partial(async_fn, object()), + ], + ) + def test_detects_async_functions_and_partials(self, fn): + assert is_async_callable(fn) is True + + def test_detects_async_callable_object(self): + class AsyncCallable: + async def __call__(self): + return "ok" + + assert is_async_callable(AsyncCallable()) is True + + @pytest.mark.parametrize( + "fn", + [ + sync_fn, + functools.partial(sync_fn, object()), + ], + ) + def test_rejects_sync_functions_and_partials(self, fn): + assert is_async_callable(fn) is False + + def test_rejects_sync_callable_object(self): + class SyncCallable: + def __call__(self): + return "ok" + + assert is_async_callable(SyncCallable()) is False diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_function_schema.py b/providers/common/ai/tests/unit/common/ai/utils/test_function_schema.py new file mode 100644 index 0000000000000..3f285c4a49b4b --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/utils/test_function_schema.py @@ -0,0 +1,399 @@ +# 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. +from __future__ import annotations + +import functools +import inspect +from typing import Annotated, Any + +import pytest + +from airflow.providers.common.ai.hooks.base import ToolSpec +from airflow.providers.common.ai.utils.function_schema import ( + _EMPTY_OBJECT_SCHEMA, + build_function_json_schema, + callable_to_tool_spec, + extract_function_description, +) + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +def _plain(x: int, y: str = "hi") -> str: # type: ignore[empty-body] + """Do something useful. + + Args: + x: a number. + y: a string. + + Returns: + A string result. + """ + + +def _no_doc(x: int) -> str: # type: ignore[empty-body] + pass + + +def _no_params() -> None: + """No parameters at all.""" + + +def _annotated(q: Annotated[str, "The search query"], limit: Annotated[int, "Max results"] = 10) -> list: # type: ignore[empty-body] + """Search.""" + + +class _CallableObj: + """Callable object used to test non-function callables.""" + + def __call__(self, value: str) -> str: # type: ignore[empty-body] + """Process value.""" + + +# --------------------------------------------------------------------------- +# extract_function_description +# --------------------------------------------------------------------------- + + +class TestExtractFunctionDescription: + def test_returns_first_paragraph(self): + assert extract_function_description(_plain) == "Do something useful." + + def test_no_docstring_falls_back_to_name(self): + assert extract_function_description(_no_doc) == "_no_doc" + + def test_empty_docstring_falls_back_to_name(self): + def fn(): + """""" + + assert extract_function_description(fn) == "fn" + + def test_lambda_falls_back_to_lambda_name(self): + f = lambda x: x + assert extract_function_description(f) == "" + + def test_callable_object_uses_call_docstring(self): + obj = _CallableObj() + # prefers __call__ docstring over class docstring over class name + assert extract_function_description(obj) == "Process value." + + def test_callable_object_falls_back_to_class_docstring(self): + class _NoCallDoc: + """Describes the class.""" + + def __call__(self, x: int) -> int: ... # type: ignore[empty-body] + + assert extract_function_description(_NoCallDoc()) == "Describes the class." + + def test_callable_object_falls_back_to_class_name(self): + class _NoDocs: + def __call__(self, x: int) -> int: ... # type: ignore[empty-body] + + assert extract_function_description(_NoDocs()) == "_NoDocs" + + @pytest.mark.parametrize( + "header", + [ + "Args:", + "Arguments:", + "Parameters:", + "Params:", + "Returns:", + "Return:", + "Yields:", + "Yield:", + "Raises:", + "Raise:", + "Except:", + "Exceptions:", + "Example:", + "Examples:", + "Note:", + "Notes:", + "See also:", + "References:", + ], + ) + def test_stops_before_section_headers(self, header): + def fn(): + pass + + fn.__doc__ = f"First paragraph.\n\n{header}\n detail" + assert extract_function_description(fn) == "First paragraph." + + def test_section_header_case_insensitive(self): + def fn(): + """Summary line. + + ARGS: + x: something. + """ + + assert extract_function_description(fn) == "Summary line." + + def test_multiline_first_paragraph_preserved(self): + def fn(): + """Line one. + Line two. + + Args: + x: something. + """ + + desc = extract_function_description(fn) + assert "Line one." in desc + assert "Line two." in desc + assert "Args" not in desc + + def test_partial_uses_underlying_function_docstring(self): + p = functools.partial(_plain, 1) + assert extract_function_description(p) == "Do something useful." + + +# --------------------------------------------------------------------------- +# build_function_json_schema +# --------------------------------------------------------------------------- + + +class TestBuildFunctionJsonSchema: + def test_no_params_returns_empty_schema(self): + assert build_function_json_schema(_no_params) == _EMPTY_OBJECT_SCHEMA + + def test_required_and_optional_params(self): + schema = build_function_json_schema(_plain) + props = schema["properties"] + assert "x" in props + assert "y" in props + assert schema["required"] == ["x"] + assert props["y"]["default"] == "hi" + + def test_int_type(self): + def fn(n: int): ... + + schema = build_function_json_schema(fn) + assert schema["properties"]["n"]["type"] == "integer" + + def test_str_type(self): + def fn(s: str): ... + + schema = build_function_json_schema(fn) + assert schema["properties"]["s"]["type"] == "string" + + def test_float_type(self): + def fn(f: float): ... + + schema = build_function_json_schema(fn) + assert schema["properties"]["f"]["type"] == "number" + + def test_bool_type(self): + def fn(flag: bool): ... + + schema = build_function_json_schema(fn) + assert schema["properties"]["flag"]["type"] == "boolean" + + def test_list_type(self): + def fn(items: list[str]): ... + + schema = build_function_json_schema(fn) + assert schema["properties"]["items"]["type"] == "array" + + def test_annotated_description_used(self): + schema = build_function_json_schema(_annotated) + assert schema["properties"]["q"]["description"] == "The search query" + assert schema["properties"]["limit"]["description"] == "Max results" + + def test_annotated_default_preserved(self): + schema = build_function_json_schema(_annotated) + assert schema["properties"]["limit"]["default"] == 10 + assert "limit" not in schema.get("required", []) + + def test_self_excluded(self): + class MyClass: + def method(self, x: int): ... + + schema = build_function_json_schema(MyClass.method) + assert "self" not in schema.get("properties", {}) + assert "x" in schema["properties"] + + def test_cls_excluded(self): + class MyClass: + @classmethod + def create(cls, x: int): ... + + schema = build_function_json_schema(MyClass.create) + assert "cls" not in schema.get("properties", {}) + + def test_var_positional_excluded(self): + def fn(*args: int): ... + + schema = build_function_json_schema(fn) + assert "args" not in schema.get("properties", {}) + + def test_var_keyword_excluded(self): + def fn(**kwargs: str): ... + + schema = build_function_json_schema(fn) + assert "kwargs" not in schema.get("properties", {}) + + def test_positional_only_rejected(self): + def fn(x: int, /, y: str): ... + + with pytest.raises(ValueError, match="parameter 'x' is positional-only"): + build_function_json_schema(fn) + + def test_unannotated_param_included_as_any(self): + def fn(x): ... + + schema = build_function_json_schema(fn) + assert "x" in schema["properties"] + + def test_title_stripped_from_schema(self): + schema = build_function_json_schema(_plain) + assert "title" not in schema + + def test_title_stripped_from_properties(self): + schema = build_function_json_schema(_plain) + for prop in schema["properties"].values(): + assert "title" not in prop + + def test_additional_properties_stripped(self): + schema = build_function_json_schema(_plain) + assert "additionalProperties" not in schema + + def test_partial_positional_bind_removes_param(self): + def add(a: int, b: int) -> int: ... # type: ignore[empty-body] + + p = functools.partial(add, 1) + schema = build_function_json_schema(p) + assert "a" not in schema.get("properties", {}) + assert "b" in schema["properties"] + + def test_partial_keyword_bind_keeps_param_as_optional(self): + def add(a: int, b: int) -> int: ... # type: ignore[empty-body] + + p = functools.partial(add, a=1) + schema = build_function_json_schema(p) + assert "a" in schema["properties"] + assert schema["properties"]["a"].get("default") == 1 + assert "b" in schema["required"] + + def test_nested_partial_unwraps_hint_source(self): + def fn(x: int, y: str) -> str: ... # type: ignore[empty-body] + + p = functools.partial(functools.partial(fn, 1), "hello") + schema = build_function_json_schema(p) + # both args bound positionally — empty schema + assert schema == _EMPTY_OBJECT_SCHEMA + + def test_optional_type(self): + def fn(x: int | None = None): ... + + schema = build_function_json_schema(fn) + assert "x" in schema["properties"] + assert "x" not in schema.get("required", []) + + def test_signature_failure_returns_empty_schema(self, monkeypatch): + def fn(x: int): ... + + def raise_value_error(_): + raise ValueError("boom") + + monkeypatch.setattr(inspect, "signature", raise_value_error) + + schema = build_function_json_schema(fn) + assert schema == _EMPTY_OBJECT_SCHEMA + + def test_callable_object_schema_from_call(self): + obj = _CallableObj() + schema = build_function_json_schema(obj) + assert "value" in schema["properties"] + assert "self" not in schema.get("properties", {}) + + @pytest.mark.parametrize("default", [0, "", False, 0.0, [], {}]) + def test_falsy_defaults_preserved(self, default): + def fn(x: Any = None): ... + + fn.__defaults__ = (default,) + import inspect + + fn.__signature__ = inspect.signature(fn) + schema = build_function_json_schema(fn) + assert schema["properties"]["x"].get("default") == default + + +# --------------------------------------------------------------------------- +# callable_to_tool_spec +# --------------------------------------------------------------------------- + + +class TestCallableToToolSpec: + def test_returns_tool_spec_instance(self): + spec = callable_to_tool_spec(_plain) + assert isinstance(spec, ToolSpec) + + def test_name_from_function_name(self): + spec = callable_to_tool_spec(_plain) + assert spec.name == "_plain" + + def test_description_from_docstring(self): + spec = callable_to_tool_spec(_plain) + assert spec.description == "Do something useful." + + def test_parameters_schema_populated(self): + spec = callable_to_tool_spec(_plain) + assert "x" in spec.parameters["properties"] + assert "y" in spec.parameters["properties"] + + def test_fn_is_original_callable(self): + spec = callable_to_tool_spec(_plain) + assert spec.fn is _plain + + def test_sequential_defaults_false(self): + spec = callable_to_tool_spec(_plain) + assert spec.sequential is False + + def test_no_docstring_uses_function_name_as_description(self): + spec = callable_to_tool_spec(_no_doc) + assert spec.description == "_no_doc" + + def test_partial_name_from_underlying_function(self): + p = functools.partial(_plain, 1) + spec = callable_to_tool_spec(p) + assert spec.name == "_plain" + + def test_partial_fn_is_the_partial_not_inner(self): + p = functools.partial(_plain, 1) + spec = callable_to_tool_spec(p) + assert spec.fn is p + + def test_partial_schema_reflects_remaining_params(self): + p = functools.partial(_plain, 1) + spec = callable_to_tool_spec(p) + assert "x" not in spec.parameters.get("properties", {}) + assert "y" in spec.parameters["properties"] + + def test_callable_object_name_from_class(self): + obj = _CallableObj() + spec = callable_to_tool_spec(obj) + assert spec.name == "_CallableObj" + + def test_callable_object_fn_is_the_object(self): + obj = _CallableObj() + spec = callable_to_tool_spec(obj) + assert spec.fn is obj diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_logging.py b/providers/common/ai/tests/unit/common/ai/utils/test_logging.py index 230335a0e027a..767b09df85abc 100644 --- a/providers/common/ai/tests/unit/common/ai/utils/test_logging.py +++ b/providers/common/ai/tests/unit/common/ai/utils/test_logging.py @@ -17,25 +17,17 @@ from __future__ import annotations import logging -from unittest.mock import MagicMock from pydantic import BaseModel -from pydantic_ai.messages import ( - ModelResponse, - ModelResponsePart, - ToolCallPart, -) -from airflow.providers.common.ai.toolsets.logging import LoggingToolset +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage from airflow.providers.common.ai.utils.logging import ( _log_output_debug, log_run_summary, - wrap_toolsets_for_logging, ) -def _make_mock_result(model_name="gpt-5", tool_names=None, usage_kwargs=None): - """Build a mock AgentRunResult with usage, response, and messages.""" +def _make_result(model_name="gpt-5", tool_names=None, usage_kwargs=None): usage_kwargs = usage_kwargs or { "requests": 4, "tool_calls": 3, @@ -43,22 +35,18 @@ def _make_mock_result(model_name="gpt-5", tool_names=None, usage_kwargs=None): "output_tokens": 512, "total_tokens": 3359, } - result = MagicMock() - result.usage.return_value = MagicMock(**usage_kwargs) - result.response = MagicMock(model_name=model_name) - - messages: list = [] - if tool_names: - parts: list[ModelResponsePart] = [ToolCallPart(tool_name=name, args="{}") for name in tool_names] - messages.append(ModelResponse(parts=parts)) - result.all_messages.return_value = messages - return result + return AgentRunResult( + output="test output", + model_name=model_name, + usage=AgentUsage(**usage_kwargs), + tool_names=tool_names, + ) class TestLogRunSummary: def test_logs_usage(self, caplog): logger = logging.getLogger("test.log_run_summary") - result = _make_mock_result() + result = _make_result() with caplog.at_level(logging.INFO, logger="test.log_run_summary"): log_run_summary(logger, result) @@ -76,7 +64,7 @@ def test_logs_usage(self, caplog): def test_logs_tool_sequence(self, caplog): logger = logging.getLogger("test.log_run_summary") - result = _make_mock_result(tool_names=["list_tables", "get_schema", "query"]) + result = _make_result(tool_names=["list_tables", "get_schema", "query"]) with caplog.at_level(logging.INFO, logger="test.log_run_summary"): log_run_summary(logger, result) @@ -88,7 +76,7 @@ def test_logs_tool_sequence(self, caplog): def test_no_tools_skips_sequence_line(self, caplog): logger = logging.getLogger("test.log_run_summary") - result = _make_mock_result(tool_names=None) + result = _make_result(tool_names=None) with caplog.at_level(logging.INFO, logger="test.log_run_summary"): log_run_summary(logger, result) @@ -97,6 +85,17 @@ def test_no_tools_skips_sequence_line(self, caplog): assert len(records) == 2 # summary line + endgroup (no tool sequence) assert records[-1].message == "::endgroup::" + def test_logs_without_usage(self, caplog): + logger = logging.getLogger("test.log_run_summary") + result = AgentRunResult(output="something", model_name="my-model", usage=None) + + with caplog.at_level(logging.INFO, logger="test.log_run_summary"): + log_run_summary(logger, result) + + records = [r for r in caplog.records if r.name == "test.log_run_summary"] + assert "model=my-model" in records[0].message + assert records[-1].message == "::endgroup::" + class TestLogOutputDebug: def test_logs_string_output(self, caplog): @@ -134,18 +133,3 @@ def test_skipped_when_debug_disabled(self, caplog): debug_records = [r for r in caplog.records if r.levelno == logging.DEBUG] assert len(debug_records) == 0 - - -class TestWrapToolsetsForLogging: - def test_wraps_each_toolset(self): - ts_a = MagicMock() - ts_b = MagicMock() - logger = logging.getLogger("test.wrap") - - wrapped = wrap_toolsets_for_logging([ts_a, ts_b], logger) - - assert len(wrapped) == 2 - assert all(isinstance(w, LoggingToolset) for w in wrapped) - assert wrapped[0].wrapped is ts_a - assert wrapped[1].wrapped is ts_b - assert wrapped[0].logger is logger diff --git a/uv.lock b/uv.lock index 796981385893b..49fcaac02e55d 100644 --- a/uv.lock +++ b/uv.lock @@ -4373,7 +4373,7 @@ requires-dist = [ { name = "pyarrow", marker = "python_full_version >= '3.14' and extra == 'parquet'", specifier = ">=22.0.0" }, { name = "pyarrow", marker = "python_full_version < '3.14' and extra == 'parquet'", specifier = ">=18.0.0" }, { name = "pydantic-ai-skills", marker = "extra == 'skills'", specifier = ">=0.11.0" }, - { name = "pydantic-ai-slim", specifier = ">=1.71.0" }, + { name = "pydantic-ai-slim", specifier = ">=1.96.0" }, { name = "pydantic-ai-slim", extras = ["anthropic"], marker = "extra == 'anthropic'" }, { name = "pydantic-ai-slim", extras = ["bedrock"], marker = "extra == 'bedrock'" }, { name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'" }, @@ -18750,7 +18750,7 @@ wheels = [ [[package]] name = "pydantic-ai-slim" -version = "1.93.0" +version = "1.96.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, @@ -18762,9 +18762,9 @@ dependencies = [ { name = "pydantic-graph" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bb/44/438dd99c7d044094037e767dab969d704232aab73e4fffd9f9a1f69bded9/pydantic_ai_slim-1.93.0.tar.gz", hash = "sha256:977364ecd3b6a2201e25d917f4efe80895210e44e66cb6983e1fc0477c78910b", size = 639585, upload-time = "2026-05-09T00:23:25.604Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f1/b0/26299238be57ddbc1ce5b4fc019338dc4856a549d5b636276bf3743a1008/pydantic_ai_slim-1.96.0.tar.gz", hash = "sha256:44ff8bb5cf81023076e82174de1d6a089515277ce508906263d17880e3fcb2f4", size = 699284, upload-time = "2026-05-14T01:09:07.519Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c6/67/bbfa2eda2493de6027e3322bccc676c8a7062dc5fe89d68cacf9e50889ce/pydantic_ai_slim-1.93.0-py3-none-any.whl", hash = "sha256:6733e3f19c83f4121fb9fe42aee918bd8f402ce670a519ecc898f108378fadb7", size = 804814, upload-time = "2026-05-09T00:23:17.122Z" }, + { url = "https://files.pythonhosted.org/packages/32/a0/4ceb29871244a398fed43009b8d6577ee60b8562437e47027e371ef2d22c/pydantic_ai_slim-1.96.0-py3-none-any.whl", hash = "sha256:58044dee3e5429499938a5ef1a44ef44d5e179f3f68e80600bbddea1f6081967", size = 870756, upload-time = "2026-05-14T01:08:58.945Z" }, ] [package.optional-dependencies] @@ -18916,7 +18916,7 @@ wheels = [ [[package]] name = "pydantic-graph" -version = "1.93.0" +version = "1.96.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx" }, @@ -18924,9 +18924,9 @@ dependencies = [ { name = "pydantic" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/63/40/4addcd3c9d06fbdf6c0776026d8a5b87e7ebb17c9896ac2452e714bb17c6/pydantic_graph-1.93.0.tar.gz", hash = "sha256:17effd900200aa7b72ec0509a79f36d3e161c2a6ef02dda6285a381e867ab195", size = 59250, upload-time = "2026-05-09T00:23:28.081Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5b/b3/7e279ee3e8d1db7ff29fb4c6cf21c2b30a002e0900eae94bcc829a66f0e2/pydantic_graph-1.96.0.tar.gz", hash = "sha256:299a2b1e47e232a78b8038779c1ff5b387d6f02d79aebae217806c5d53607f9e", size = 59294, upload-time = "2026-05-14T01:09:10.11Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2b/bc/5f9eb611c6b9315fb0c6ae58bb82a63d9d9f17340b6c8778319261593fbb/pydantic_graph-1.93.0-py3-none-any.whl", hash = "sha256:ef1b0dcd55a6b5a3544de53a9594216ee8f51ac26bf4be79f7ef310747be598a", size = 73066, upload-time = "2026-05-09T00:23:20.847Z" }, + { url = "https://files.pythonhosted.org/packages/a3/58/918e641e1d94b95315a174bf318a78c0c127191333ffe021a92d417f6159/pydantic_graph-1.96.0-py3-none-any.whl", hash = "sha256:5904661751c4f19cba726e4e16a878f2f83722432236c231c88dba2bd887b43d", size = 73047, upload-time = "2026-05-14T01:09:02.476Z" }, ] [[package]]