From 89c36915f12233caf947f64882fbfde6057bde32 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sat, 23 May 2026 12:30:43 +0100 Subject: [PATCH 01/28] Add BaseAIHook and Update usage --- providers/common/ai/AGENTS.md | 15 +- providers/common/ai/docs/changelog.rst | 9 + providers/common/ai/docs/hooks/index.rst | 5 +- providers/common/ai/docs/operators/agent.rst | 28 +- providers/common/ai/docs/toolsets.rst | 56 ++- .../example_dags/example_pydantic_ai_hook.py | 27 +- .../providers/common/ai/hooks/base_ai.py | 353 +++++++++++++++++ .../providers/common/ai/hooks/pydantic_ai.py | 179 +++++++-- .../providers/common/ai/operators/agent.py | 170 ++++---- .../providers/common/ai/operators/llm.py | 32 +- .../common/ai/operators/llm_branch.py | 10 +- .../common/ai/operators/llm_file_analysis.py | 12 +- .../common/ai/operators/llm_schema_compare.py | 10 +- .../providers/common/ai/operators/llm_sql.py | 12 +- .../providers/common/ai/policies/retry.py | 24 +- .../providers/common/ai/toolsets/sql.py | 90 ++--- .../providers/common/ai/utils/logging.py | 62 +-- .../unit/common/ai/decorators/test_agent.py | 109 +++--- .../unit/common/ai/decorators/test_llm.py | 70 ++-- .../common/ai/decorators/test_llm_branch.py | 65 ++- .../ai/decorators/test_llm_file_analysis.py | 50 ++- .../ai/decorators/test_llm_schema_compare.py | 64 ++- .../unit/common/ai/decorators/test_llm_sql.py | 69 ++-- .../unit/common/ai/hooks/test_base_ai.py | 343 ++++++++++++++++ .../unit/common/ai/hooks/test_pydantic_ai.py | 369 +++++++++++++++--- .../unit/common/ai/operators/test_agent.py | 308 ++++++++------- .../unit/common/ai/operators/test_llm.py | 120 +++--- .../common/ai/operators/test_llm_branch.py | 60 ++- .../ai/operators/test_llm_file_analysis.py | 70 ++-- .../ai/operators/test_llm_schema_compare.py | 69 ++-- .../unit/common/ai/operators/test_llm_sql.py | 125 +++--- .../unit/common/ai/policies/test_retry.py | 136 ++++--- .../tests/unit/common/ai/toolsets/test_sql.py | 179 ++++----- .../unit/common/ai/utils/test_logging.py | 60 ++- 34 files changed, 2179 insertions(+), 1181 deletions(-) create mode 100644 providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py create mode 100644 providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py diff --git a/providers/common/ai/AGENTS.md b/providers/common/ai/AGENTS.md index 8e8711c8ab78f..d59f1208313c4 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_ai.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/docs/changelog.rst b/providers/common/ai/docs/changelog.rst index 9a9e4cfc81276..956e1a3a563fd 100644 --- a/providers/common/ai/docs/changelog.rst +++ b/providers/common/ai/docs/changelog.rst @@ -47,6 +47,15 @@ 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. +Next release +............ + +Features +^^^^^^^^ + +* Add ``BaseAIHook`` contract with framework-agnostic ``create_agent`` / ``run_agent`` / + ``get_model`` interface so ``AgentOperator`` selects the agent backend via connection type. + 0.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/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index 8e7c8ad5f3983..7feb1f4661a9c 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,18 @@ 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_ai.BaseToolset` subclasses + (``SQLToolset``), pydantic-ai ``AbstractToolset`` implementations + (``HookToolset``, ``MCPToolset``, ``DataFusionToolset``, + ``AgentSkillsToolset`` for :ref:`agent-skills`, third-party toolsets), + plain Python callables, or native pydantic-ai ``Tool`` objects. Mixed lists + are supported. +- ``enable_tool_logging``: When ``True`` (default), wraps each tool call with + real-time logging. For pydantic-ai ``AbstractToolset`` items this is done via + :class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset`; for + plain callables and :class:`~airflow.providers.common.ai.hooks.base_ai.BaseToolset` + items it is applied at the callable level. - ``agent_params``: Additional keyword arguments passed to the pydantic-ai ``Agent`` constructor (e.g. ``retries``, ``model_settings``, ``capabilities``). See :ref:`capabilities-passthrough` for how to enable pydantic-ai capabilities @@ -327,9 +337,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..8b75a3deecf5e 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -34,29 +34,41 @@ 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_ai.BaseToolset` interface. +All three can be passed to +:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`, which routes each +toolset to the correct agent parameter automatically. .. note:: - ``AgentOperator`` accepts **any** ``AbstractToolset`` implementation — not - just the Airflow-native toolsets above. PydanticAI's own MCP server - classes (``MCPServerStreamableHTTP``, ``MCPServerSSE``, ``MCPServerStdio``) - and third-party toolsets work too. The Airflow-native toolsets add - connection management, secret backend integration, and the connection UI, - but you are not locked in. + ``AgentOperator`` accepts a mixed ``toolsets`` list containing any + combination of: + - pydantic-ai ``AbstractToolset`` implementations (``HookToolset``, + ``MCPToolset``, ``DataFusionToolset``). + - Any third-party ``AbstractToolset``, including PydanticAI's own MCP + server classes (``MCPServerStreamableHTTP``, ``MCPServerSSE``, + ``MCPServerStdio``). + - :class:`~airflow.providers.common.ai.hooks.base_ai.BaseToolset` + subclasses (``SQLToolset``). + - Plain Python callables (``def my_tool(...): ...``). + - Native pydantic-ai ``Tool`` objects. -Using Toolsets Directly with PydanticAI ---------------------------------------- + The hook routes each item to the correct agent parameter automatically. -Toolsets are standard pydantic-ai ``AbstractToolset`` implementations with no -dependency on ``AgentOperator`` or ``@task.agent``. You can use them anywhere -you can run Python within Airflow -- ``@task`` functions, ``PythonOperator`` -callables, or any custom operator's ``execute()`` method -- by creating a -``pydantic_ai.Agent`` yourself: + +Using Toolsets Directly +----------------------- + +Toolsets can be used anywhere you can run Python within Airflow — ``@task`` +functions, ``PythonOperator`` callables, or any custom operator's +``execute()`` method — without needing ``AgentOperator`` or ``@task.agent``. +Pass toolsets via :class:`~airflow.providers.common.ai.hooks.base_ai.AgentRunRequest` +and call the hook yourself: .. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py :language: python @@ -67,11 +79,13 @@ This works because toolsets resolve Airflow connections lazily via ``BaseHook.get_connection()``, which is available in any task execution context. -This approach gives you full control over the agent lifecycle -- you can call -``agent.run_sync()`` multiple times, swap models at runtime, or combine -results from several agents in a single task. The tradeoff is that you lose +This approach gives you direct control over the agent lifecycle — you can +build and run multiple agents in a single task, or combine results from +several runs. The tradeoff is that you lose the durable execution (step-level caching with retry replay), HITL review -integration, and automatic tool call logging that ``AgentOperator`` provides. +integration, and the automatic tool call logging and routing that +``AgentOperator`` provides via +:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset`. ``HookToolset`` 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..a58514bc3cc9c 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_ai import AgentRunRequest, BaseAIHook from airflow.providers.common.compat.sdk import dag, task @@ -29,9 +29,10 @@ def example_pydantic_ai_hook(): @task def generate_summary(text: str) -> str: - hook = PydanticAIHook(llm_conn_id="pydanticai_default") - agent = hook.create_agent(output_type=str, instructions="Summarize concisely.") - result = agent.run_sync(text) + hook = BaseAIHook.get_agent_hook("pydanticai_default") + request = AgentRunRequest(prompt=text, output_type=str, instructions="Summarize concisely.") + agent = hook.create_agent(request) + result = hook.run_agent(agent, request) return result.output generate_summary("Apache Airflow is a platform for programmatically authoring...") @@ -51,12 +52,14 @@ class SQLResult(BaseModel): query: str explanation: str - hook = PydanticAIHook(llm_conn_id="pydanticai_default") - agent = hook.create_agent( + hook = BaseAIHook.get_agent_hook("pydanticai_default") + request = AgentRunRequest( + prompt=prompt, output_type=SQLResult, instructions="Generate a SQL query and explain it.", ) - result = agent.run_sync(prompt) + agent = hook.create_agent(request) + result = hook.run_agent(agent, request) return result.output.model_dump() generate_sql("Find the top 10 customers by revenue") @@ -76,8 +79,9 @@ def example_task_with_toolsets(): def analyze_revenue() -> str: from airflow.providers.common.ai.toolsets.sql import SQLToolset - hook = PydanticAIHook(llm_conn_id="pydanticai_default") - agent = hook.create_agent( + hook = BaseAIHook.get_agent_hook("pydanticai_default") + request = AgentRunRequest( + prompt="Which customers have spent the most? Show the top 5.", output_type=str, instructions=( "You are a sales analytics assistant. " @@ -91,7 +95,8 @@ def analyze_revenue() -> str: ), ], ) - result = agent.run_sync("Which customers have spent the most? Show the top 5.") + agent = hook.create_agent(request) + result = hook.run_agent(agent, request) return result.output analyze_revenue() diff --git a/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py b/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py new file mode 100644 index 0000000000000..119752345161b --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py @@ -0,0 +1,353 @@ +# 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 inspect +import json +import time +from abc import ABCMeta, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any, ClassVar + +from airflow.providers.common.compat.sdk import BaseHook + + +@dataclass +class AgentUsage: + """Token and request usage from an agent run, when the backend exposes it.""" + + requests: int = 0 + tool_calls: int = 0 + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 + + +@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*. + """ + + name: str + description: str + parameters: dict[str, Any] + fn: Callable[..., Any] + + +@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. + + :param prompt: User prompt for this invocation. + :param output_type: Expected structured output type (default: ``str``). + :param instructions: System-level instructions for the agent. + :param toolsets: List of :class:`BaseToolset` instances the agent may call. + :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 each tool callable with a logging shim. + :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 + output_type: type[Any] = 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, 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``). + + Subclasses implement :meth:`get_model`, :meth:`create_agent`, :meth:`run_agent`, and + :meth:`_tool_spec_to_native`. + + Shared helpers :meth:`_init_durable`, :meth:`_resolve_tools`, :meth:`_logged_callable`, and + :meth:`_cached_callable` are provided for all hooks. + """ + + conn_name_attr = "llm_conn_id" + + supports_toolsets: ClassVar[bool] = False + supports_durable: ClassVar[bool] = False + supports_usage_limits: ClassVar[bool] = False + + @classmethod + def get_agent_hook(cls, conn_id: str, *, hook_params: dict[str, Any] | None = None) -> BaseAIHook: + """ + 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. Delegates to :meth:`get_model`.""" + return self.get_model() + + @abstractmethod + def create_agent(self, request: AgentRunRequest) -> Any: + """ + 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*. + + When :attr:`AgentRunRequest.durable_context` is set, implementations + should call :meth:`_init_durable` and store the returned storage/counter + on the instance so that :meth:`run_agent` can use them. + + :param request: All parameters needed to configure the agent. + :returns: Framework-native agent object, ready to be passed to :meth:`run_agent`. + """ + + @abstractmethod + def run_agent(self, agent: Any, request: AgentRunRequest) -> AgentRunResult: + """ + Execute *agent* for *request* and return a normalized :class:`AgentRunResult`. + + Implementations that store durable state on the instance (set during + :meth:`create_agent`) should apply it here and clean up after the run. + + :param agent: Framework-native agent 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 _init_durable(self, ctx: DurableContext) -> tuple[Any, Any]: + """ + Create and return a ``DurableStorage`` / ``DurableStepCounter`` pair for *ctx*. + + Hooks call this inside :meth:`create_agent` when + :attr:`AgentRunRequest.durable_context` is set. + """ + from airflow.providers.common.ai.durable.step_counter import DurableStepCounter + from airflow.providers.common.ai.durable.storage import DurableStorage + + storage = DurableStorage( + dag_id=ctx.dag_id, + task_id=ctx.task_id, + run_id=ctx.run_id, + map_index=ctx.map_index, + ) + counter = DurableStepCounter() + return storage, counter + + def _resolve_tools( + self, + toolsets: list[Any], + enable_logging: bool, + storage: Any, + counter: Any, + ) -> 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`. + * Plain Python function (``def`` / ``lambda``) — auto-wraps into a :class:`ToolSpec` + using ``__name__`` and ``__doc__``, 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: + *fn* → optional cache wrap → optional log wrap → :meth:`_tool_spec_to_native`. + + :param toolsets: Mix of :class:`BaseToolset` instances, plain callables, and native tool objects. + :param enable_logging: When ``True``, wrap each callable with :meth:`_logged_callable`. + :param storage: ``DurableStorage`` instance, or ``None`` to skip caching. + :param counter: ``DurableStepCounter`` instance, or ``None`` to skip caching. + """ + native: list[Any] = [] + for ts in toolsets: + if isinstance(ts, BaseToolset): + specs = ts.as_tools() + elif inspect.isfunction(ts): + specs = [ToolSpec(name=ts.__name__, description=ts.__doc__ or "", parameters={}, fn=ts)] + else: + native.append(ts) + continue + for spec in specs: + fn = spec.fn + if storage is not None and counter is not None: + fn = self._cached_callable(fn, storage, counter) + if enable_logging: + fn = self._logged_callable(fn, self.log) + adapted = ToolSpec( + name=spec.name, + description=spec.description, + parameters=spec.parameters, + fn=fn, + ) + native.append(self._tool_spec_to_native(adapted)) + return native + + @staticmethod + def _logged_callable(fn: Callable[..., Any], logger: Any) -> Callable[..., Any]: + """Wrap *fn* to log tool name, args, timing, and exceptions.""" + + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + name = fn.__name__ + logger.info("::group::Tool call: %s", 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", name, elapsed) + logger.info("::endgroup::") + return result + except Exception: + elapsed = time.monotonic() - start + logger.exception("Tool %s failed after %.2fs", name, elapsed) + logger.info("::endgroup::") + raise + + return wrapper + + @staticmethod + def _cached_callable( + fn: Callable[..., Any], + storage: Any, + counter: Any, + ) -> Callable[..., Any]: + """Wrap *fn* to cache its result in *storage* using a monotonic step counter.""" + + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + 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 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..4bbee2e930baf 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,30 @@ # under the License. from __future__ import annotations -from typing import TYPE_CHECKING, Any, TypeVar, overload +from typing import TYPE_CHECKING, Any from pydantic_ai import Agent from pydantic_ai.models import infer_model from pydantic_ai.providers import infer_provider, infer_provider_class +from airflow.providers.common.ai.hooks.base_ai import ( + AgentRunRequest, + AgentRunResult, + AgentUsage, + BaseAIHook, + DurableStats, + ToolSpec, +) from airflow.providers.common.ai.observability import genai_instrumentation_settings -from airflow.providers.common.compat.sdk import BaseHook - -OutputT = TypeVar("OutputT") if TYPE_CHECKING: from pydantic_ai.models import KnownModelName, Model + from airflow.providers.common.ai.durable.step_counter import DurableStepCounter + from airflow.providers.common.ai.durable.storage import DurableStorage + -class PydanticAIHook(BaseHook): +class PydanticAIHook(BaseAIHook): """ Hook for LLM access via pydantic-ai. @@ -57,6 +65,10 @@ class PydanticAIHook(BaseHook): conn_type = "pydanticai" hook_name = "Pydantic AI" + supports_toolsets = True + supports_durable = True + supports_usage_limits = True + def __init__( self, llm_conn_id: str | None = None, @@ -71,6 +83,8 @@ def __init__( self.llm_conn_id = llm_conn_id if llm_conn_id is not None else self.default_conn_name self.model_id = model_id self._model: Model | None = None + self._durable_storage: DurableStorage | None = None + self._durable_counter: DurableStepCounter | None = None @staticmethod def get_ui_field_behaviour() -> dict[str, Any]: @@ -117,7 +131,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``. @@ -172,42 +186,147 @@ 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 + # ------------------------------------------------------------------ + + def _tool_spec_to_native(self, spec: ToolSpec) -> Any: + """Convert a :class:`~airflow.providers.common.ai.hooks.base_ai.ToolSpec` to a pydantic-ai ``Tool``.""" + from pydantic_ai.tools import Tool - @overload - def create_agent(self, *, instructions: str, **agent_kwargs) -> Agent[None, str]: ... + return Tool(spec.fn, name=spec.name, description=spec.description) - def create_agent( - self, output_type: type[Any] = str, *, instructions: str, **agent_kwargs - ) -> Agent[None, Any]: + def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: """ - Create a pydantic-ai Agent configured with this hook's model. + Build a pydantic-ai ``Agent`` 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 stores them on the instance 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. """ - 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 + if request.durable_context is not None: + storage, counter = self._init_durable(request.durable_context) + self._durable_storage = storage + self._durable_counter = counter + else: + self._durable_storage = None + self._durable_counter = None + + extra_kwargs = dict(request.agent_params or {}) + if request.toolsets: + from pydantic_ai.toolsets.abstract import AbstractToolset + + 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 = self._resolve_tools( + pipeline_items, + request.enable_tool_logging, + self._durable_storage, + self._durable_counter, + ) + if resolved: + extra_kwargs["tools"] = resolved + + if abstract_items: + processed: list[Any] = list(abstract_items) + if self._durable_storage is not None and self._durable_counter is not None: + from airflow.providers.common.ai.durable.caching_toolset import CachingToolset + + processed = [ + CachingToolset( + wrapped=ts, + storage=self._durable_storage, + counter=self._durable_counter, + ) + for ts in processed + ] + if request.enable_tool_logging: + from airflow.providers.common.ai.toolsets.logging import LoggingToolset + + processed = [LoggingToolset(wrapped=ts, logger=self.log) for ts in processed] + extra_kwargs["toolsets"] = processed + + agent = Agent( + self.get_model(), + output_type=request.output_type, + instructions=request.instructions, + **extra_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 + def run_agent(self, agent: Agent[None, Any], request: AgentRunRequest) -> AgentRunResult: + """Run *agent* synchronously for *request* and return a normalized :class:`~airflow.providers.common.ai.hooks.base_ai.AgentRunResult`.""" + from pydantic_ai.messages import ToolCallPart + + 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 + + if self._durable_storage is not None and self._durable_counter is not None: + from airflow.providers.common.ai.durable.caching_model import CachingModel + + resolved_model = infer_model(agent.model) + caching_model = CachingModel( + resolved_model, + storage=self._durable_storage, + counter=self._durable_counter, + ) + with agent.override(model=caching_model): + result = agent.run_sync(request.prompt, **run_kwargs) + else: + result = 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 self._durable_storage is not None: + run_result.durable_stats = DurableStats( + replayed_model=self._durable_counter.replayed_model, + replayed_tool=self._durable_counter.replayed_tool, + cached_model=self._durable_counter.cached_model, + cached_tool=self._durable_counter.cached_tool, + ) + self._durable_storage.cleanup() + self._durable_storage = None + self._durable_counter = None + + return run_result + def test_connection(self) -> tuple[bool, str]: """ Test connection by resolving the model. @@ -218,7 +337,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/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index b41a0c54d8b2f..5c679e26c118c 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,13 @@ from pydantic import BaseModel -from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook +from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, BaseAIHook, 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.output_type import rehydrate_pydantic_output +from airflow.providers.common.ai.utils.output_type import ( + iter_base_model_classes, + rehydrate_pydantic_output, +) from airflow.providers.common.compat.sdk import ( AirflowOptionalProviderFeatureException, BaseOperator, @@ -47,12 +50,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 +91,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,14 +110,14 @@ 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_ai.BaseToolset` + instances the agent can use. + :param enable_tool_logging: When ``True`` (default), wraps each tool callable with a + logging shim that logs 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 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=..., ...)`` @@ -170,7 +172,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 +217,69 @@ 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." + return BaseAIHook.get_agent_hook(self.llm_conn_id, hook_params=hook_params) + + def _validate_hook_capabilities(self) -> None: + """Raise if operator options are incompatible with the resolved agent hook.""" + hook = self.llm_hook + if self.toolsets and not hook.supports_toolsets: + raise ValueError( + f"toolsets are not supported for connection {self.llm_conn_id!r} " + f"(conn_type resolves to {type(hook).__name__}). " + "Use a connection type with toolset support (e.g. pydanticai, pydanticai-bedrock)." + ) + if self.usage_limits is not None and not hook.supports_usage_limits: + raise ValueError( + f"usage_limits are not supported for connection {self.llm_conn_id!r} " + f"(conn_type resolves to {type(hook).__name__})." + ) + if self.durable and not hook.supports_durable: + raise ValueError( + f"durable=True requires a hook that supports durable execution; got {type(hook).__name__} " + f"for connection {self.llm_conn_id!r}." ) - 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( + def _build_request(self, *, prompt: str, message_history: Any = None) -> AgentRunRequest: + """Build an :class:`~airflow.providers.common.ai.hooks.base_ai.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() + 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), + ) - 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 + def execute(self, context: Context) -> Any: + self._validate_hook_capabilities() - from airflow.providers.common.ai.durable.caching_model import CachingModel + 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 +294,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 +319,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..8672869fb72d4 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_ai 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_ai.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..b23c932848612 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_ai 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..974eb7e25440b 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_ai 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..3033a676f52f6 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_ai 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..abb25fb8514bc 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_ai.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_ai 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/sql.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py index ee3128705a1f9..05a3c863e0bd5 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 ( @@ -34,19 +34,9 @@ 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_ai 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 +92,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. @@ -191,6 +181,7 @@ def _is_table_allowed(self, name: str) -> bool: def id(self) -> str: return f"sql-{self._db_conn_id}" + # ------------------------------------------------------------------ # Lazy hook resolution # ------------------------------------------------------------------ @@ -208,50 +199,36 @@ 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, - 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}") + 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, + ), + ToolSpec( + name="get_schema", + description="Get column names and types for a table.", + parameters=_GET_SCHEMA_SCHEMA, + fn=self._get_schema, + ), + ToolSpec( + name="query", + description="Execute a SQL query and return rows as JSON.", + parameters=_QUERY_SCHEMA, + fn=self._query, + ), + ToolSpec( + name="check_query", + description="Validate SQL syntax without executing it.", + parameters=_CHECK_QUERY_SCHEMA, + fn=self._check_query, + ), + ] # ------------------------------------------------------------------ # Tool implementations @@ -321,11 +298,10 @@ def _query(self, sql: str) -> str: rows = hook.get_records(sql) except Exception as e: if self._is_retryable_query_error(hook, e): - raise ModelRetry( + raise ValueError( 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/logging.py b/providers/common/ai/src/airflow/providers/common/ai/utils/logging.py index 47cabf7ce6ccc..2b2f4e7ecce31 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_ai 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..4cbc131cb3063 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_ai import AgentRunResult, AgentUsage, BaseAIHook try: from airflow.sdk.serde import SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as _CORE_WALKER @@ -40,38 +40,39 @@ 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.supports_toolsets = True + mock_hook.supports_durable = False + mock_hook.supports_usage_limits = True + 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,35 +89,33 @@ 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.""" + from pydantic_ai.messages import ImageUrl 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 + from pydantic_ai.messages import ImageUrl op = _AgentDecoratedOperator( task_id="test", @@ -127,14 +126,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 +138,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 +156,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 +173,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 +189,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..fb5939b174992 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_ai 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..dc2e3da1ff160 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_ai 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..7a1f5904a2939 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_ai 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..6746c1e88d09a 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_ai 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..d5db22e512232 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_ai 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_ai.py b/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py new file mode 100644 index 0000000000000..244c4ad2e2d0a --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py @@ -0,0 +1,343 @@ +# 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 + +from unittest.mock import MagicMock, patch + +import pytest + +from airflow.providers.common.ai.hooks.base_ai import ( + AgentRunRequest, + AgentRunResult, + AgentUsage, + BaseAIHook, + BaseToolset, + DurableContext, + DurableStats, + ToolSpec, +) +from airflow.providers.common.compat.sdk import BaseHook + + +class TestBaseAIHookGetAgentHook: + @patch("airflow.providers.common.ai.hooks.base_ai.BaseHook.get_hook", autospec=True) + def test_returns_hook_when_instance_is_base_ai_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_ai.BaseHook.get_hook", autospec=True) + def test_raises_when_hook_is_not_base_ai_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 TestAgentRunResult: + 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 create_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, storage=None, counter=None) + + 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 create_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, storage=None, counter=None) + wrapped_fn() + + assert calls == ["original"] + mock_log.info.assert_called() + + def test_resolve_tools_wraps_plain_callable(self): + """A bare Python function is auto-wrapped using __name__ and __doc__.""" + + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + + def get_model(self): + return None + + def create_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} + + 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, storage=None, counter=None) + + assert len(result) == 1 + assert result[0]["name"] == "roll_dice" + assert result[0]["description"] == "Roll a six-sided die and return the result." + + 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 create_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, storage=None, counter=None) + + 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 create_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, storage=None, counter=None + ) + + assert result == ["converted:greet", native_tool] + + +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() + + +class TestBaseAIHookCachedCallable: + def test_cached_callable_saves_and_returns(self): + storage = MagicMock() + counter = MagicMock() + counter.next_step.return_value = 1 + storage.load_tool_result.return_value = (False, None) + + calls = [] + + def fn(): + calls.append(1) + return "computed" + + wrapped = BaseAIHook._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() + counter = MagicMock() + counter.next_step.return_value = 1 + storage.load_tool_result.return_value = (True, "cached_value") + + calls = [] + + def fn(): + calls.append(1) + return "computed" + + wrapped = BaseAIHook._cached_callable(fn, storage, counter) + result = wrapped() + + assert result == "cached_value" + assert calls == [] + counter.replayed_tool += 1 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..f2baad0423ea1 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 @@ -25,6 +25,7 @@ from pydantic_ai.models.test import TestModel from airflow.models.connection import Connection +from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, BaseAIHook from airflow.providers.common.ai.hooks.pydantic_ai import ( PydanticAIAzureHook, PydanticAIBedrockHook, @@ -33,6 +34,16 @@ ) +class TestPydanticAIHookBaseContract: + def test_is_base_ai_hook(self): + assert issubclass(PydanticAIHook, BaseAIHook) + + def test_capability_flags(self): + assert PydanticAIHook.supports_toolsets is True + assert PydanticAIHook.supports_durable is True + assert PydanticAIHook.supports_usage_limits is True + + class TestPydanticAIHookInit: def test_default_conn_id(self): hook = PydanticAIHook() @@ -57,11 +68,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_initialised_to_none(self): + hook = PydanticAIHook() + assert hook._durable_storage is None + assert hook._durable_counter is None + -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 +92,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 +109,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 +122,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 +141,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 +154,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 +168,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,28 +186,35 @@ 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: @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) @@ -205,12 +224,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 +237,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 +259,260 @@ 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_inits_durable_when_context_set(self, mock_agent_cls, mock_infer_model): + from airflow.providers.common.ai.hooks.base_ai 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() + 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=(mock_storage, mock_counter)), + ): + hook.create_agent(request) + + assert hook._durable_storage is mock_storage + assert hook._durable_counter is mock_counter + + @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_clears_durable_when_no_context(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") + hook._durable_storage = MagicMock() + hook._durable_counter = MagicMock() + + request = AgentRunRequest(prompt="hi") + conn = Connection(conn_id="test_conn", conn_type="pydanticai") + with patch.object(hook, "get_connection", return_value=conn): + hook.create_agent(request) + + assert hook._durable_storage is None + assert hook._durable_counter 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 the BaseToolset pipeline.""" + 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 resolved; native Tool objects are passed through unchanged.""" + from pydantic_ai.tools import Tool + + from airflow.providers.common.ai.hooks.base_ai import BaseToolset, ToolSpec + + mock_model = MagicMock(spec=Model) + mock_infer_model.return_value = mock_model + + def my_fn() -> str: + return "result" + + class MyToolset(BaseToolset): + def as_tools(self): + return [ToolSpec(name="my_fn", description="desc", parameters={}, fn=my_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) == 2 + assert any(t is native_tool for t in 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_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_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 are wrapped with CachingToolset (inner) and LoggingToolset (outer) 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_ai 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() + + 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=(mock_storage, mock_counter)), + ): + 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, LoggingToolset) + assert isinstance(outer.wrapped, CachingToolset) + assert outer.wrapped.wrapped is abstract_ts + + +class TestPydanticAIHookRunAgent: + def test_run_agent_returns_agent_run_result(self): + hook = PydanticAIHook() + mock_agent = MagicMock() + mock_usage = MagicMock(requests=1, tool_calls=0, input_tokens=5, output_tokens=10, total_tokens=15) + mock_result = MagicMock() + mock_result.output = "done" + mock_result.usage.return_value = mock_usage + mock_result.response = MagicMock(model_name="openai:gpt-5") + mock_result.all_messages.return_value = [] + mock_agent.run_sync.return_value = mock_result + + request = AgentRunRequest(prompt="hello") + run_result = hook.run_agent(mock_agent, request) + + assert isinstance(run_result, AgentRunResult) + assert run_result.output == "done" + assert run_result.model_name == "openai:gpt-5" + mock_agent.run_sync.assert_called_once_with("hello") + + def test_run_agent_forwards_message_history_and_usage_limits(self): + hook = PydanticAIHook() + mock_agent = MagicMock() + mock_result = MagicMock() + mock_result.output = "ok" + 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="m") + mock_result.all_messages.return_value = ["history"] + mock_agent.run_sync.return_value = mock_result + limits = MagicMock() + history = ["prior"] + + request = AgentRunRequest(prompt="more", message_history=history, usage_limits=limits) + hook.run_agent(mock_agent, request) + + mock_agent.run_sync.assert_called_once_with("more", message_history=history, usage_limits=limits) + + def test_run_agent_durable_applies_caching_model(self): + """When durable state is set, run_agent wraps model with CachingModel.""" + hook = PydanticAIHook() + mock_agent = MagicMock() + mock_agent.model = MagicMock() + mock_result = MagicMock() + mock_result.output = "ok" + 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="m") + mock_result.all_messages.return_value = [] + mock_agent.run_sync.return_value = mock_result + mock_agent.override.return_value.__enter__ = MagicMock(return_value=None) + mock_agent.override.return_value.__exit__ = MagicMock(return_value=False) + + 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 + hook._durable_storage = mock_storage + hook._durable_counter = mock_counter + + with patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", return_value=MagicMock()): + with patch("airflow.providers.common.ai.durable.caching_model.CachingModel"): + request = AgentRunRequest(prompt="hi") + run_result = hook.run_agent(mock_agent, request) + + mock_agent.override.assert_called_once() + assert run_result.durable_stats is not None + assert hook._durable_storage is None + assert hook._durable_counter is None + mock_storage.cleanup.assert_called_once() + class TestPydanticAIHookCreateAgentInstrumentation: """create_agent() wires OpenTelemetry instrumentation from observability.""" @@ -387,12 +657,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 +675,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 +686,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 +696,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 +740,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 +749,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 +774,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 +821,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 +925,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 +934,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 +958,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/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 4a0b08cf1483e..9cba27de5c162 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,8 @@ from pydantic import BaseModel from pydantic_ai.usage import UsageLimits +from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, AgentUsage, BaseAIHook 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 +44,44 @@ 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.supports_toolsets = True + mock_hook.supports_durable = True + mock_hook.supports_usage_limits = True + 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.supports_toolsets = False + mock_hook.supports_durable = False + mock_hook.supports_usage_limits = False + 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 are not supported"): + op.execute(context=MagicMock()) class TestAgentOperatorValidation: @@ -95,11 +116,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 +131,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 +150,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 +170,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 +198,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 +218,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 +236,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 +262,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 +276,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 +310,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 +336,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 +360,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 +430,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 + + 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 - @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 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 +490,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 +547,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..a1200d6d8f993 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,7 @@ from pydantic import BaseModel from pydantic_ai.usage import UsageLimits +from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, AgentUsage from airflow.providers.common.ai.mixins.approval import ( LLMApprovalMixin, ) @@ -54,15 +55,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 +68,34 @@ 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="" - ) - mock_hook_cls.get_hook.assert_called_once_with("my_llm", hook_params={"model_id": None}) - - @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 + 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_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 +106,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 +131,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. @@ -186,14 +191,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 +218,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 +245,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 +272,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 +296,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={}) 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..0edc93746d98f 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_ai 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..45e4cd7c67302 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,7 @@ import pytest from pydantic import BaseModel +from airflow.providers.common.ai.hooks.base_ai import AgentRunResult, AgentUsage from airflow.providers.common.ai.operators.llm_file_analysis import LLMFileAnalysisOperator from airflow.providers.common.ai.utils.file_analysis import FileAnalysisRequest @@ -44,19 +45,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 +74,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 +84,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() + 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 +110,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 +126,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() + 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", @@ -201,7 +201,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 +215,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() + 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 +239,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 +253,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() + 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 +315,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 +329,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() + 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..0e3892d6acd34 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_ai 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..61e6c09444aea 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_ai 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", 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..c358347b953b3 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_ai 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/toolsets/test_sql.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py index 5e425597a32e1..e1bb0815c6d04 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,15 +16,14 @@ # under the License. from __future__ import annotations -import asyncio import importlib.util import json import sqlite3 from unittest.mock import MagicMock, PropertyMock, patch import pytest -from pydantic_ai.exceptions import ModelRetry +from airflow.providers.common.ai.hooks.base_ai 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 +49,83 @@ 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"] -class TestSQLToolsetGetTools: - def test_returns_four_tools(self): + def test_tool_descriptions_non_empty(self): ts = SQLToolset("pg_default") - tools = asyncio.run(ts.get_tools(ctx=MagicMock())) - assert set(tools.keys()) == {"list_tables", "get_schema", "query", "check_query"} + for spec in ts.as_tools(): + assert spec.description - def test_tool_definitions_have_descriptions(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())) - for tool in tools.values(): - assert tool.tool_def.description + 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"] 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 +138,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 +149,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 +159,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 + with pytest.raises(ValueError) as exc_info: + 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 + with pytest.raises(ValueError) as exc_info: + 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 +200,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 +209,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 +219,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, @@ -257,16 +249,9 @@ def test_sqlalchemy_programming_error_with_psycopg2_undefined_column_orig_raises "airflow.providers.common.ai.toolsets.sql._SQLALCHEMY_RETRYABLE_EXCEPTIONS", (ProgrammingError,), ), - pytest.raises(ModelRetry), + pytest.raises(ValueError), ): - 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 +283,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 +316,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 +335,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() 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..f943365218291 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_ai 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 From 07de78a3d7252aa51dfabf248cfa582538ffe8b0 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sun, 24 May 2026 23:17:02 +0100 Subject: [PATCH 02/28] Update usage and resolve conflicts --- .../ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 4bbee2e930baf..592a5391bee0a 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 @@ -293,7 +293,7 @@ def run_agent(self, agent: Agent[None, Any], request: AgentRunRequest) -> AgentR else: result = agent.run_sync(request.prompt, **run_kwargs) - usage = result.usage() + usage = result.usage tool_names: list[str] = [] for message in result.all_messages(): for part in getattr(message, "parts", []): From 1a48f4bb2e1110f18e5586013f7867a7e9d3a40a Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Mon, 25 May 2026 18:59:24 +0100 Subject: [PATCH 03/28] Resolve comments --- providers/common/ai/docs/toolsets.rst | 15 ++- .../providers/common/ai/hooks/base_ai.py | 55 +++++++- .../providers/common/ai/hooks/pydantic_ai.py | 124 +++++++++--------- .../providers/common/ai/operators/agent.py | 7 + .../providers/common/ai/toolsets/sql.py | 8 +- .../unit/common/ai/hooks/test_base_ai.py | 39 +++++- .../unit/common/ai/hooks/test_pydantic_ai.py | 99 +++++++++++--- .../unit/common/ai/operators/test_llm.py | 9 +- .../unit/common/ai/operators/test_llm_sql.py | 9 +- .../tests/unit/common/ai/toolsets/test_sql.py | 12 +- 10 files changed, 274 insertions(+), 103 deletions(-) diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index 8b75a3deecf5e..4ffc848373664 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -81,11 +81,16 @@ context. 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 the automatic tool call logging and routing that -``AgentOperator`` provides via -:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset`. +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_ai.BaseAIHook`): +callable-level logging and caching for +:class:`~airflow.providers.common.ai.hooks.base_ai.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/src/airflow/providers/common/ai/hooks/base_ai.py b/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py index 119752345161b..9770b57fb4c31 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py @@ -29,6 +29,14 @@ from airflow.providers.common.compat.sdk import BaseHook +_EMPTY_OBJECT_SCHEMA: dict[str, Any] = {"type": "object", "properties": {}} + +# Durable storage/counter pairs keyed by ``id(agent)``. +# pydantic-ai ``Agent`` is not hashable, so ``WeakKeyDictionary`` cannot be used. +# ``create_agent`` and ``run_agent`` run synchronously in the same task, so ``id()`` +# is stable until ``_pop_agent_durable`` removes the entry. +_AGENT_DURABLE: dict[int, tuple[Any, Any]] = {} + @dataclass class AgentUsage: @@ -85,12 +93,15 @@ class ToolSpec: :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 @@ -156,6 +167,9 @@ class BaseAIHook(BaseHook, metaclass=ABCMeta): :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:`create_agent`, :meth:`run_agent`, and :meth:`_tool_spec_to_native`. @@ -169,6 +183,16 @@ class BaseAIHook(BaseHook, metaclass=ABCMeta): supports_durable: ClassVar[bool] = False supports_usage_limits: ClassVar[bool] = False + 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: """ @@ -203,8 +227,9 @@ def create_agent(self, request: AgentRunRequest) -> Any: with the model, tools, instructions, and output type from *request*. When :attr:`AgentRunRequest.durable_context` is set, implementations - should call :meth:`_init_durable` and store the returned storage/counter - on the instance so that :meth:`run_agent` can use them. + should call :meth:`_init_durable` and bind the returned storage/counter + to the agent via :meth:`_bind_agent_durable` so that :meth:`run_agent` + can retrieve and clean them up. :param request: All parameters needed to configure the agent. :returns: Framework-native agent object, ready to be passed to :meth:`run_agent`. @@ -215,8 +240,10 @@ def run_agent(self, agent: Any, request: AgentRunRequest) -> AgentRunResult: """ Execute *agent* for *request* and return a normalized :class:`AgentRunResult`. - Implementations that store durable state on the instance (set during - :meth:`create_agent`) should apply it here and clean up after the run. + Implementations with durable execution should pop durable state via + :meth:`_pop_agent_durable`, 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 produced by :meth:`create_agent`. :param request: The same request used to create the agent (prompt, usage @@ -254,6 +281,16 @@ def _init_durable(self, ctx: DurableContext) -> tuple[Any, Any]: counter = DurableStepCounter() return storage, counter + @staticmethod + def _bind_agent_durable(agent: Any, storage: Any, counter: Any) -> None: + """Associate *storage* and *counter* with *agent* until :meth:`run_agent` completes.""" + _AGENT_DURABLE[id(agent)] = (storage, counter) + + @staticmethod + def _pop_agent_durable(agent: Any) -> tuple[Any, Any] | None: + """Remove and return durable state bound to *agent*, if any.""" + return _AGENT_DURABLE.pop(id(agent), None) + def _resolve_tools( self, toolsets: list[Any], @@ -285,7 +322,14 @@ def _resolve_tools( if isinstance(ts, BaseToolset): specs = ts.as_tools() elif inspect.isfunction(ts): - specs = [ToolSpec(name=ts.__name__, description=ts.__doc__ or "", parameters={}, fn=ts)] + specs = [ + ToolSpec( + name=ts.__name__, + description=ts.__doc__ or "", + parameters=_EMPTY_OBJECT_SCHEMA, + fn=ts, + ) + ] else: native.append(ts) continue @@ -300,6 +344,7 @@ def _resolve_tools( description=spec.description, parameters=spec.parameters, fn=fn, + sequential=spec.sequential, ) native.append(self._tool_spec_to_native(adapted)) return native 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 592a5391bee0a..9e5dd38d3fcb0 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 @@ -35,9 +35,6 @@ if TYPE_CHECKING: from pydantic_ai.models import KnownModelName, Model - from airflow.providers.common.ai.durable.step_counter import DurableStepCounter - from airflow.providers.common.ai.durable.storage import DurableStorage - class PydanticAIHook(BaseAIHook): """ @@ -75,16 +72,13 @@ def __init__( 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 - self._durable_storage: DurableStorage | None = None - self._durable_counter: DurableStepCounter | None = None @staticmethod def get_ui_field_behaviour() -> dict[str, Any]: @@ -194,14 +188,19 @@ def _tool_spec_to_native(self, spec: ToolSpec) -> Any: """Convert a :class:`~airflow.providers.common.ai.hooks.base_ai.ToolSpec` to a pydantic-ai ``Tool``.""" from pydantic_ai.tools import Tool - return Tool(spec.fn, name=spec.name, description=spec.description) + return Tool( + spec.fn, + name=spec.name, + description=spec.description, + sequential=spec.sequential, + ) def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: """ Build a pydantic-ai ``Agent`` from *request*. When :attr:`~AgentRunRequest.durable_context` is set, initialises durable - storage and step counter and stores them on the instance for use by + storage and step counter and binds them to the returned 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 @@ -209,13 +208,9 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: :param request: Agent configuration — output type, instructions, toolsets, extra params. """ + storage = counter = None if request.durable_context is not None: storage, counter = self._init_durable(request.durable_context) - self._durable_storage = storage - self._durable_counter = counter - else: - self._durable_storage = None - self._durable_counter = None extra_kwargs = dict(request.agent_params or {}) if request.toolsets: @@ -228,22 +223,22 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: resolved = self._resolve_tools( pipeline_items, request.enable_tool_logging, - self._durable_storage, - self._durable_counter, + storage, + counter, ) if resolved: extra_kwargs["tools"] = resolved if abstract_items: processed: list[Any] = list(abstract_items) - if self._durable_storage is not None and self._durable_counter is not None: + if storage is not None and counter is not None: from airflow.providers.common.ai.durable.caching_toolset import CachingToolset processed = [ CachingToolset( wrapped=ts, - storage=self._durable_storage, - counter=self._durable_counter, + storage=storage, + counter=counter, ) for ts in processed ] @@ -259,6 +254,8 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: instructions=request.instructions, **extra_kwargs, ) + if storage is not None and counter is not None: + self._bind_agent_durable(agent, storage, counter) if "instrument" not in extra_kwargs: # Set the public ``agent.instrument`` property rather than the # ``Agent(instrument=...)`` constructor kwarg, which is deprecated in @@ -279,52 +276,57 @@ def run_agent(self, agent: Agent[None, Any], request: AgentRunRequest) -> AgentR if request.usage_limits is not None: run_kwargs["usage_limits"] = request.usage_limits - if self._durable_storage is not None and self._durable_counter is not None: - from airflow.providers.common.ai.durable.caching_model import CachingModel + durable = self._pop_agent_durable(agent) + storage, counter = durable if durable else (None, None) - resolved_model = infer_model(agent.model) - caching_model = CachingModel( - resolved_model, - storage=self._durable_storage, - counter=self._durable_counter, - ) - with agent.override(model=caching_model): + try: + if storage is not None and counter is not None: + from airflow.providers.common.ai.durable.caching_model import CachingModel + + 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(request.prompt, **run_kwargs) + else: result = agent.run_sync(request.prompt, **run_kwargs) - else: - result = 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 self._durable_storage is not None: - run_result.durable_stats = DurableStats( - replayed_model=self._durable_counter.replayed_model, - replayed_tool=self._durable_counter.replayed_tool, - cached_model=self._durable_counter.cached_model, - cached_tool=self._durable_counter.cached_tool, + 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, ) - self._durable_storage.cleanup() - self._durable_storage = None - self._durable_counter = 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, + ) + except BaseException: + raise + else: + if storage is not None: + storage.cleanup() return run_result def test_connection(self) -> tuple[bool, str]: 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 5c679e26c118c..5796de3307524 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 @@ -33,6 +33,7 @@ iter_base_model_classes, 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, @@ -268,6 +269,12 @@ def _build_request(self, *, prompt: str, message_history: Any = None) -> AgentRu ) 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._validate_hook_capabilities() self._durable_ti = context["task_instance"] if self.durable else None 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 05a3c863e0bd5..6d50f3d9d7129 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 @@ -34,6 +34,8 @@ raise AirflowOptionalProviderFeatureException(e) +from pydantic_ai.exceptions import ModelRetry + from airflow.providers.common.ai.hooks.base_ai import BaseToolset, ToolSpec from airflow.providers.common.compat.sdk import BaseHook @@ -209,24 +211,28 @@ def as_tools(self) -> list[ToolSpec]: description="List available table names in the database.", parameters=_LIST_TABLES_SCHEMA, fn=self._list_tables, + sequential=True, ), 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, ), ] @@ -298,7 +304,7 @@ def _query(self, sql: str) -> str: rows = hook.get_records(sql) except Exception as e: if self._is_retryable_query_error(hook, e): - raise ValueError( + raise ModelRetry( f"error: {e!s}, Use get_schema and list_tables tools for more details." ) from e raise diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py b/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py index 244c4ad2e2d0a..fc6f3e64cbaef 100644 --- a/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py @@ -52,6 +52,41 @@ def test_raises_when_hook_is_not_base_ai_hook(self, mock_get_hook): 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 create_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 TestBaseAIHookAgentDurable: + def test_bind_pop_round_trip(self): + # Real pydantic-ai Agent is unhashable; id()-keyed storage must still work. + agent = object() + storage = MagicMock() + counter = MagicMock() + + BaseAIHook._bind_agent_durable(agent, storage, counter) + assert BaseAIHook._pop_agent_durable(agent) == (storage, counter) + assert BaseAIHook._pop_agent_durable(agent) is None + + class TestAgentRunResult: def test_dataclass_fields(self): usage = AgentUsage(requests=1, tool_calls=2, total_tokens=10) @@ -326,6 +361,7 @@ def fn(): def test_cached_callable_replays_on_hit(self): storage = MagicMock() counter = MagicMock() + counter.replayed_tool = 0 counter.next_step.return_value = 1 storage.load_tool_result.return_value = (True, "cached_value") @@ -340,4 +376,5 @@ def fn(): assert result == "cached_value" assert calls == [] - counter.replayed_tool += 1 + assert counter.replayed_tool == 1 + storage.save_tool_result.assert_not_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 f2baad0423ea1..0b118ae85a302 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 @@ -25,7 +25,7 @@ from pydantic_ai.models.test import TestModel from airflow.models.connection import Connection -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, BaseAIHook +from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, BaseAIHook, ToolSpec from airflow.providers.common.ai.hooks.pydantic_ai import ( PydanticAIAzureHook, PydanticAIBedrockHook, @@ -68,10 +68,10 @@ def test_vertex_hook_uses_own_default_conn_name(self): hook = PydanticAIVertexHook() assert hook.llm_conn_id == "pydanticai_vertex_default" - def test_durable_state_initialised_to_none(self): + def test_durable_state_not_stored_on_hook_instance(self): hook = PydanticAIHook() - assert hook._durable_storage is None - assert hook._durable_counter is None + assert not hasattr(hook, "_durable_storage") + assert not hasattr(hook, "_durable_counter") class TestPydanticAIHookGetModel: @@ -280,26 +280,24 @@ def test_create_agent_inits_durable_when_context_set(self, mock_agent_cls, mock_ ): hook.create_agent(request) - assert hook._durable_storage is mock_storage - assert hook._durable_counter is mock_counter + mock_agent = mock_agent_cls.return_value + assert BaseAIHook._pop_agent_durable(mock_agent) == (mock_storage, mock_counter) @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_clears_durable_when_no_context(self, mock_agent_cls, mock_infer_model): + def test_create_agent_does_not_bind_durable_when_no_context(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") - hook._durable_storage = MagicMock() - hook._durable_counter = MagicMock() request = AgentRunRequest(prompt="hi") conn = Connection(conn_id="test_conn", conn_type="pydanticai") with patch.object(hook, "get_connection", return_value=conn): hook.create_agent(request) - assert hook._durable_storage is None - assert hook._durable_counter is None + mock_agent = mock_agent_cls.return_value + assert BaseAIHook._pop_agent_durable(mock_agent) 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) @@ -437,6 +435,37 @@ def test_create_agent_wraps_abstract_toolset_with_caching_when_durable( assert isinstance(outer.wrapped, CachingToolset) assert outer.wrapped.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_binds_durable_per_agent_not_on_hook(self, mock_agent_cls, mock_infer_model): + """Second create_agent must not overwrite durable state for the first agent.""" + from airflow.providers.common.ai.hooks.base_ai 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() + mock_agent_a = MagicMock() + mock_agent_b = MagicMock() + mock_agent_cls.side_effect = [mock_agent_a, mock_agent_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=[(storage_a, counter_a), (storage_b, counter_b)] + ): + agent_a = hook.create_agent(AgentRunRequest(prompt="a", durable_context=ctx_a)) + agent_b = hook.create_agent(AgentRunRequest(prompt="b", durable_context=ctx_b)) + + assert agent_a is mock_agent_a + assert agent_b is mock_agent_b + assert BaseAIHook._pop_agent_durable(agent_a) == (storage_a, counter_a) + assert BaseAIHook._pop_agent_durable(agent_b) == (storage_b, counter_b) + class TestPydanticAIHookRunAgent: def test_run_agent_returns_agent_run_result(self): @@ -499,8 +528,7 @@ def test_run_agent_durable_applies_caching_model(self): mock_counter.replayed_tool = 0 mock_counter.cached_model = 0 mock_counter.cached_tool = 0 - hook._durable_storage = mock_storage - hook._durable_counter = mock_counter + BaseAIHook._bind_agent_durable(mock_agent, mock_storage, mock_counter) with patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", return_value=MagicMock()): with patch("airflow.providers.common.ai.durable.caching_model.CachingModel"): @@ -509,10 +537,51 @@ def test_run_agent_durable_applies_caching_model(self): mock_agent.override.assert_called_once() assert run_result.durable_stats is not None - assert hook._durable_storage is None - assert hook._durable_counter is None + assert BaseAIHook._pop_agent_durable(mock_agent) is None mock_storage.cleanup.assert_called_once() + def test_run_agent_preserves_durable_cache_on_exception(self): + hook = PydanticAIHook() + mock_agent = MagicMock() + mock_agent.model = MagicMock() + mock_agent.run_sync.side_effect = RuntimeError("boom") + mock_agent.override.return_value.__enter__ = MagicMock(return_value=None) + mock_agent.override.return_value.__exit__ = MagicMock(return_value=False) + + mock_storage = MagicMock() + mock_counter = MagicMock() + BaseAIHook._bind_agent_durable(mock_agent, mock_storage, mock_counter) + + with patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", return_value=MagicMock()): + with patch("airflow.providers.common.ai.durable.caching_model.CachingModel"): + with pytest.raises(RuntimeError, match="boom"): + hook.run_agent(mock_agent, AgentRunRequest(prompt="hi")) + + mock_storage.cleanup.assert_not_called() + assert BaseAIHook._pop_agent_durable(mock_agent) is None + + def test_tool_spec_to_native_forwards_sequential(self): + hook = PydanticAIHook() + + def fn() -> str: + return "ok" + + spec = ToolSpec( + name="fn", + description="desc", + parameters={"type": "object", "properties": {}}, + fn=fn, + sequential=True, + ) + with patch("pydantic_ai.tools.Tool") as mock_tool_cls: + hook._tool_spec_to_native(spec) + mock_tool_cls.assert_called_once_with( + fn, + name="fn", + description="desc", + sequential=True, + ) + class TestPydanticAIHookCreateAgentInstrumentation: """create_agent() wires OpenTelemetry instrumentation from observability.""" 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 a1200d6d8f993..fabf21bfe5e0d 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 @@ -371,11 +371,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", @@ -387,4 +384,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_sql.py b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py index 61e6c09444aea..e140b0a5ed71d 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 @@ -651,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", @@ -667,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/toolsets/test_sql.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_sql.py index e1bb0815c6d04..51fc76f6e7b54 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 @@ -22,6 +22,7 @@ from unittest.mock import MagicMock, PropertyMock, patch import pytest +from pydantic_ai.exceptions import ModelRetry from airflow.providers.common.ai.hooks.base_ai import BaseToolset, ToolSpec from airflow.providers.common.ai.toolsets.sql import SQLToolset @@ -95,6 +96,11 @@ def test_tool_parameters_match_schemas(self): assert "sql" in specs["query"].parameters["properties"] assert "sql" in specs["check_query"].parameters["properties"] + def test_tools_are_sequential(self): + ts = SQLToolset("pg_default") + for spec in ts.as_tools(): + assert spec.sequential is True + class TestSQLToolsetListTables: def test_returns_all_tables(self): @@ -174,7 +180,7 @@ def test_raises_model_retry_when_query_fails_with_retryable_error(self): ts._hook.conn_type = "sqlite" ts._hook.get_records.side_effect = sqlite3.OperationalError("no such column: nonexistent") - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ModelRetry) as exc_info: ts._query("SELECT id, nonexistent FROM users") assert "nonexistent" in exc_info.value.args[0] @@ -187,7 +193,7 @@ def test_model_retry_message_includes_schema_hint(self): ts._hook.conn_type = "sqlite" ts._hook.get_records.side_effect = sqlite3.OperationalError("no such table: missing_table") - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ModelRetry) as exc_info: ts._query("SELECT foo FROM x") assert "get_schema" in exc_info.value.args[0] @@ -249,7 +255,7 @@ def test_sqlalchemy_programming_error_with_psycopg2_undefined_column_orig_raises "airflow.providers.common.ai.toolsets.sql._SQLALCHEMY_RETRYABLE_EXCEPTIONS", (ProgrammingError,), ), - pytest.raises(ValueError), + pytest.raises(ModelRetry), ): ts._query("SELECT id, missing FROM users") From 7a51429c6a308317f1d596d5e59b10b599bc174f Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Mon, 25 May 2026 20:12:36 +0100 Subject: [PATCH 04/28] Resolve ci failures --- docs/spelling_wordlist.txt | 1 + providers/common/ai/provider.yaml | 3 +++ .../src/airflow/providers/common/ai/get_provider_info.py | 4 ++++ .../ai/src/airflow/providers/common/ai/hooks/base_ai.py | 7 ++++--- .../src/airflow/providers/common/ai/hooks/pydantic_ai.py | 7 +++++-- .../ai/tests/unit/common/ai/decorators/test_agent.py | 3 --- 6 files changed, 17 insertions(+), 8 deletions(-) diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 00c557d1eb2a0..b0f6fd8d93927 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -831,6 +831,7 @@ init initdb initialisation initialiser +initialises initialising initializer inout diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index 695a88ec2c103..14842636d8494 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -62,6 +62,9 @@ integrations: tags: [ai] hooks: + - integration-name: Base Hook + python-modules: + - airflow.providers.common.ai.hooks.base_ai - integration-name: Pydantic AI python-modules: - airflow.providers.common.ai.hooks.pydantic_ai 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..d95d22efa084d 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,10 @@ def get_provider_info(): }, ], "hooks": [ + { + "integration-name": "Base Hook", + "python-modules": ["airflow.providers.common.ai.hooks.base_ai"], + }, { "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_ai.py b/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py index 9770b57fb4c31..24c4c332c7ce9 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py @@ -23,7 +23,7 @@ import json import time from abc import ABCMeta, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass, field from typing import Any, ClassVar @@ -123,7 +123,8 @@ class AgentRunRequest: framework-neutral structure, so that :class:`~airflow.providers.common.ai.operators.agent.AgentOperator` has zero framework-specific imports. - :param prompt: User prompt for this invocation. + :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 :class:`BaseToolset` instances the agent may call. @@ -135,7 +136,7 @@ class AgentRunRequest: Use this escape hatch for framework-specific options. """ - prompt: str + prompt: str | Sequence[Any] output_type: type[Any] = str instructions: str = "" toolsets: list[Any] | None = None 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 9e5dd38d3fcb0..5305a77161376 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 @@ -142,7 +142,7 @@ def get_model(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", "") @@ -283,7 +283,10 @@ def run_agent(self, agent: Agent[None, Any], request: AgentRunRequest) -> AgentR if storage is not None and counter is not None: from airflow.providers.common.ai.durable.caching_model import CachingModel - resolved_model = infer_model(agent.model) + if agent.model is None: + raise ValueError("Agent model must be set when durable=True") + model = agent.model + resolved_model = infer_model(model) if isinstance(model, str) else model caching_model = CachingModel( resolved_model, storage=storage, 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 4cbc131cb3063..c1876139c06f0 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 @@ -91,7 +91,6 @@ def test_execute_raises_on_invalid_prompt(self, return_value): def test_execute_accepts_sequence_prompt(self): """A non-empty Sequence[UserContent] return value is forwarded as-is.""" - from pydantic_ai.messages import ImageUrl image = ImageUrl(url="https://example.com/x.png") prompt = ["Describe this:", image] @@ -115,8 +114,6 @@ def test_sequence_prompt_with_hitl_review_raises(self): if not AIRFLOW_V_3_1_PLUS: pytest.skip("enable_hitl_review requires Airflow >= 3.1.0") - from pydantic_ai.messages import ImageUrl - op = _AgentDecoratedOperator( task_id="test", python_callable=lambda: ["x", ImageUrl(url="https://example.com/x.png")], From 586e412378d67595dd14ffa937d07e2277bb1edf Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Wed, 27 May 2026 04:32:28 +0100 Subject: [PATCH 05/28] Resolve round 2 comments --- providers/common/ai/docs/changelog.rst | 6 + .../providers/common/ai/hooks/base_ai.py | 41 +++- .../providers/common/ai/hooks/pydantic_ai.py | 117 ++++----- .../providers/common/ai/operators/agent.py | 22 -- .../unit/common/ai/hooks/test_base_ai.py | 59 ++++- .../unit/common/ai/hooks/test_pydantic_ai.py | 226 ++++++++++++------ .../unit/common/ai/operators/test_agent.py | 8 + .../unit/common/ai/operators/test_llm.py | 18 +- 8 files changed, 334 insertions(+), 163 deletions(-) diff --git a/providers/common/ai/docs/changelog.rst b/providers/common/ai/docs/changelog.rst index 956e1a3a563fd..8b53791d5aa3a 100644 --- a/providers/common/ai/docs/changelog.rst +++ b/providers/common/ai/docs/changelog.rst @@ -56,6 +56,12 @@ Features * Add ``BaseAIHook`` contract with framework-agnostic ``create_agent`` / ``run_agent`` / ``get_model`` interface so ``AgentOperator`` selects the agent backend via connection type. +.. For the next release: 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_ai.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. + 0.4.0 ..... diff --git a/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py b/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py index 24c4c332c7ce9..3e4526116c086 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py @@ -31,11 +31,8 @@ _EMPTY_OBJECT_SCHEMA: dict[str, Any] = {"type": "object", "properties": {}} -# Durable storage/counter pairs keyed by ``id(agent)``. -# pydantic-ai ``Agent`` is not hashable, so ``WeakKeyDictionary`` cannot be used. -# ``create_agent`` and ``run_agent`` run synchronously in the same task, so ``id()`` -# is stable until ``_pop_agent_durable`` removes the entry. -_AGENT_DURABLE: dict[int, tuple[Any, Any]] = {} +# Attribute name for durable storage/counter bound to a framework agent instance. +_AIRFLOW_DURABLE_ATTR = "_airflow_durable_state" @dataclass @@ -232,6 +229,9 @@ def create_agent(self, request: AgentRunRequest) -> Any: to the agent via :meth:`_bind_agent_durable` so that :meth:`run_agent` can retrieve and clean them up. + Implementations must call :meth:`validate_run_request` at the start of + this method before any agent construction or durable initialisation. + :param request: All parameters needed to configure the agent. :returns: Framework-native agent object, ready to be passed to :meth:`run_agent`. """ @@ -263,6 +263,29 @@ def _tool_spec_to_native(self, spec: ToolSpec) -> Any: 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. + + Hook implementations call this at the start of :meth:`create_agent`. + """ + hook_name = type(self).__name__ + conn_id = self.llm_conn_id or "unknown" + if request.toolsets and not self.supports_toolsets: + raise ValueError( + f"toolsets are not supported for connection {conn_id!r} (conn_type resolves to {hook_name}). " + ) + if request.usage_limits is not None and not self.supports_usage_limits: + raise ValueError( + f"usage_limits are not supported for connection {conn_id!r} " + f"(conn_type resolves to {hook_name})." + ) + if request.durable_context is not None and not self.supports_durable: + raise ValueError( + f"durable execution requires a hook that supports durable caching; " + f"got {hook_name} for connection {conn_id!r}." + ) + def _init_durable(self, ctx: DurableContext) -> tuple[Any, Any]: """ Create and return a ``DurableStorage`` / ``DurableStepCounter`` pair for *ctx*. @@ -285,12 +308,16 @@ def _init_durable(self, ctx: DurableContext) -> tuple[Any, Any]: @staticmethod def _bind_agent_durable(agent: Any, storage: Any, counter: Any) -> None: """Associate *storage* and *counter* with *agent* until :meth:`run_agent` completes.""" - _AGENT_DURABLE[id(agent)] = (storage, counter) + setattr(agent, _AIRFLOW_DURABLE_ATTR, (storage, counter)) @staticmethod def _pop_agent_durable(agent: Any) -> tuple[Any, Any] | None: """Remove and return durable state bound to *agent*, if any.""" - return _AGENT_DURABLE.pop(id(agent), None) + state = getattr(agent, _AIRFLOW_DURABLE_ATTR, None) + if state is None: + return None + delattr(agent, _AIRFLOW_DURABLE_ATTR) + return state def _resolve_tools( self, 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 5305a77161376..fc8155e55184d 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 @@ -19,9 +19,14 @@ 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_ai import ( AgentRunRequest, AgentRunResult, @@ -31,6 +36,7 @@ ToolSpec, ) from airflow.providers.common.ai.observability import genai_instrumentation_settings +from airflow.providers.common.ai.toolsets.logging import LoggingToolset if TYPE_CHECKING: from pydantic_ai.models import KnownModelName, Model @@ -186,8 +192,6 @@ def _provider_factory(pname: str) -> Any: def _tool_spec_to_native(self, spec: ToolSpec) -> Any: """Convert a :class:`~airflow.providers.common.ai.hooks.base_ai.ToolSpec` to a pydantic-ai ``Tool``.""" - from pydantic_ai.tools import Tool - return Tool( spec.fn, name=spec.name, @@ -208,13 +212,22 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: :param request: Agent configuration — output type, instructions, toolsets, extra params. """ + self.validate_run_request(request) + storage = counter = None if request.durable_context is not None: storage, counter = self._init_durable(request.durable_context) extra_kwargs = dict(request.agent_params or {}) if request.toolsets: - from pydantic_ai.toolsets.abstract import AbstractToolset + 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)] @@ -226,14 +239,11 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: storage, counter, ) - if resolved: - extra_kwargs["tools"] = resolved + extra_kwargs["tools"] = resolved if abstract_items: processed: list[Any] = list(abstract_items) if storage is not None and counter is not None: - from airflow.providers.common.ai.durable.caching_toolset import CachingToolset - processed = [ CachingToolset( wrapped=ts, @@ -243,8 +253,6 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: for ts in processed ] if request.enable_tool_logging: - from airflow.providers.common.ai.toolsets.logging import LoggingToolset - processed = [LoggingToolset(wrapped=ts, logger=self.log) for ts in processed] extra_kwargs["toolsets"] = processed @@ -268,8 +276,6 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: def run_agent(self, agent: Agent[None, Any], request: AgentRunRequest) -> AgentRunResult: """Run *agent* synchronously for *request* and return a normalized :class:`~airflow.providers.common.ai.hooks.base_ai.AgentRunResult`.""" - from pydantic_ai.messages import ToolCallPart - run_kwargs: dict[str, Any] = {} if request.message_history is not None: run_kwargs["message_history"] = request.message_history @@ -279,57 +285,52 @@ def run_agent(self, agent: Agent[None, Any], request: AgentRunRequest) -> AgentR durable = self._pop_agent_durable(agent) storage, counter = durable if durable else (None, None) - try: - if storage is not None and counter is not None: - from airflow.providers.common.ai.durable.caching_model import CachingModel - - if agent.model is None: - raise ValueError("Agent model must be set when durable=True") - model = agent.model - resolved_model = infer_model(model) if isinstance(model, str) else model - caching_model = CachingModel( - resolved_model, - storage=storage, - counter=counter, - ) - with agent.override(model=caching_model): - result = agent.run_sync(request.prompt, **run_kwargs) - else: + if storage is not None and counter is not None: + if agent.model is None: + raise ValueError("Agent model must be set when durable=True") + model = agent.model + resolved_model = infer_model(model) if isinstance(model, str) else model + caching_model = CachingModel( + resolved_model, + storage=storage, + counter=counter, + ) + with agent.override(model=caching_model): result = agent.run_sync(request.prompt, **run_kwargs) + else: + result = 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, + ) - 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 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, - ) - except BaseException: - raise - else: - if storage is not None: - storage.cleanup() + if storage is not None: + storage.cleanup() return run_result def test_connection(self) -> tuple[bool, str]: 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 5796de3307524..23ace9d96283f 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 @@ -225,26 +225,6 @@ def llm_hook(self) -> BaseAIHook: } return BaseAIHook.get_agent_hook(self.llm_conn_id, hook_params=hook_params) - def _validate_hook_capabilities(self) -> None: - """Raise if operator options are incompatible with the resolved agent hook.""" - hook = self.llm_hook - if self.toolsets and not hook.supports_toolsets: - raise ValueError( - f"toolsets are not supported for connection {self.llm_conn_id!r} " - f"(conn_type resolves to {type(hook).__name__}). " - "Use a connection type with toolset support (e.g. pydanticai, pydanticai-bedrock)." - ) - if self.usage_limits is not None and not hook.supports_usage_limits: - raise ValueError( - f"usage_limits are not supported for connection {self.llm_conn_id!r} " - f"(conn_type resolves to {type(hook).__name__})." - ) - if self.durable and not hook.supports_durable: - raise ValueError( - f"durable=True requires a hook that supports durable execution; got {type(hook).__name__} " - f"for connection {self.llm_conn_id!r}." - ) - def _build_request(self, *, prompt: str, message_history: Any = None) -> AgentRunRequest: """Build an :class:`~airflow.providers.common.ai.hooks.base_ai.AgentRunRequest` from operator config.""" durable_context: DurableContext | None = None @@ -275,8 +255,6 @@ def execute(self, context: Context) -> Any: feature_name="enable_hitl_review", feature_enabled=self.enable_hitl_review, ) - self._validate_hook_capabilities() - self._durable_ti = context["task_instance"] if self.durable else None request = self._build_request(prompt=self.prompt) diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py b/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py index fc6f3e64cbaef..c8a82b5fc2fbc 100644 --- a/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py @@ -19,6 +19,8 @@ from unittest.mock import MagicMock, patch import pytest +from pydantic_ai import Agent +from pydantic_ai.models.test import TestModel from airflow.providers.common.ai.hooks.base_ai import ( AgentRunRequest, @@ -77,14 +79,67 @@ def _tool_spec_to_native(self, spec): class TestBaseAIHookAgentDurable: def test_bind_pop_round_trip(self): - # Real pydantic-ai Agent is unhashable; id()-keyed storage must still work. - agent = object() + agent = Agent(TestModel()) storage = MagicMock() counter = MagicMock() BaseAIHook._bind_agent_durable(agent, storage, counter) + assert agent._airflow_durable_state == (storage, counter) assert BaseAIHook._pop_agent_durable(agent) == (storage, counter) assert BaseAIHook._pop_agent_durable(agent) is None + assert not hasattr(agent, "_airflow_durable_state") + + +class TestValidateRunRequest: + def test_rejects_toolsets_when_unsupported(self): + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + supports_toolsets = False + supports_usage_limits = True + supports_durable = True + + def get_model(self): + return None + + def create_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 are not supported"): + hook.validate_run_request(request) + + def test_rejects_usage_limits_when_unsupported(self): + class ConcreteHook(BaseAIHook): + conn_type = "test" + hook_name = "Test" + supports_toolsets = True + supports_usage_limits = False + supports_durable = True + + def get_model(self): + return None + + def create_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 are not supported"): + hook.validate_run_request(request) class TestAgentRunResult: 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 0b118ae85a302..9b0385fb5ae4e 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 @@ -21,8 +21,12 @@ from unittest.mock import MagicMock, patch import pytest +from pydantic_ai import Agent +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_ai import AgentRunRequest, AgentRunResult, BaseAIHook, ToolSpec @@ -34,6 +38,43 @@ ) +def _test_agent() -> Agent[None, str]: + return Agent(TestModel()) + + +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 TestPydanticAIHookBaseContract: def test_is_base_ai_hook(self): assert issubclass(PydanticAIHook, BaseAIHook) @@ -261,7 +302,46 @@ def test_create_agent_with_agent_params(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_inits_durable_when_context_set(self, mock_agent_cls, mock_infer_model): + 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_ai import DurableContext mock_model = MagicMock(spec=Model) @@ -278,14 +358,12 @@ def test_create_agent_inits_durable_when_context_set(self, mock_agent_cls, mock_ patch.object(hook, "get_connection", return_value=conn), patch.object(hook, "_init_durable", return_value=(mock_storage, mock_counter)), ): - hook.create_agent(request) + agent = hook.create_agent(request) - mock_agent = mock_agent_cls.return_value - assert BaseAIHook._pop_agent_durable(mock_agent) == (mock_storage, mock_counter) + assert BaseAIHook._pop_agent_durable(agent) == (mock_storage, mock_counter) @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_does_not_bind_durable_when_no_context(self, mock_agent_cls, mock_infer_model): + def test_create_agent_does_not_bind_durable_when_no_context(self, mock_infer_model): mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model @@ -294,10 +372,9 @@ def test_create_agent_does_not_bind_durable_when_no_context(self, mock_agent_cls request = AgentRunRequest(prompt="hi") conn = Connection(conn_id="test_conn", conn_type="pydanticai") with patch.object(hook, "get_connection", return_value=conn): - hook.create_agent(request) + agent = hook.create_agent(request) - mock_agent = mock_agent_cls.return_value - assert BaseAIHook._pop_agent_durable(mock_agent) is None + assert BaseAIHook._pop_agent_durable(agent) 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) @@ -436,8 +513,7 @@ def test_create_agent_wraps_abstract_toolset_with_caching_when_durable( assert outer.wrapped.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_binds_durable_per_agent_not_on_hook(self, mock_agent_cls, mock_infer_model): + def test_create_agent_binds_durable_per_agent_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_ai import DurableContext @@ -449,9 +525,6 @@ def test_create_agent_binds_durable_per_agent_not_on_hook(self, mock_agent_cls, ctx_b = DurableContext(dag_id="d", task_id="t", run_id="r2") storage_a, counter_a = MagicMock(), MagicMock() storage_b, counter_b = MagicMock(), MagicMock() - mock_agent_a = MagicMock() - mock_agent_b = MagicMock() - mock_agent_cls.side_effect = [mock_agent_a, mock_agent_b] conn = Connection(conn_id="test_conn", conn_type="pydanticai") with patch.object(hook, "get_connection", return_value=conn): @@ -461,8 +534,7 @@ def test_create_agent_binds_durable_per_agent_not_on_hook(self, mock_agent_cls, agent_a = hook.create_agent(AgentRunRequest(prompt="a", durable_context=ctx_a)) agent_b = hook.create_agent(AgentRunRequest(prompt="b", durable_context=ctx_b)) - assert agent_a is mock_agent_a - assert agent_b is mock_agent_b + assert agent_a is not agent_b assert BaseAIHook._pop_agent_durable(agent_a) == (storage_a, counter_a) assert BaseAIHook._pop_agent_durable(agent_b) == (storage_b, counter_b) @@ -470,57 +542,60 @@ def test_create_agent_binds_durable_per_agent_not_on_hook(self, mock_agent_cls, class TestPydanticAIHookRunAgent: def test_run_agent_returns_agent_run_result(self): hook = PydanticAIHook() - mock_agent = MagicMock() - mock_usage = MagicMock(requests=1, tool_calls=0, input_tokens=5, output_tokens=10, total_tokens=15) - mock_result = MagicMock() - mock_result.output = "done" - mock_result.usage.return_value = mock_usage - mock_result.response = MagicMock(model_name="openai:gpt-5") - mock_result.all_messages.return_value = [] - mock_agent.run_sync.return_value = mock_result + agent = _test_agent() + mock_result = _pydantic_run_result( + "done", + model_name="openai:gpt-5", + input_tokens=5, + output_tokens=10, + ) request = AgentRunRequest(prompt="hello") - run_result = hook.run_agent(mock_agent, request) + with patch.object(agent, "run_sync", return_value=mock_result) as mock_run_sync: + run_result = hook.run_agent(agent, request) assert isinstance(run_result, AgentRunResult) assert run_result.output == "done" assert run_result.model_name == "openai:gpt-5" - mock_agent.run_sync.assert_called_once_with("hello") + 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.supports_usage_limits = False + with pytest.raises(ValueError, match="usage_limits are not supported"): + hook.create_agent(AgentRunRequest(prompt="hi", usage_limits=UsageLimits())) def test_run_agent_forwards_message_history_and_usage_limits(self): hook = PydanticAIHook() - mock_agent = MagicMock() - mock_result = MagicMock() - mock_result.output = "ok" - 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="m") - mock_result.all_messages.return_value = ["history"] - mock_agent.run_sync.return_value = mock_result - limits = MagicMock() + agent = _test_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) - hook.run_agent(mock_agent, request) - - mock_agent.run_sync.assert_called_once_with("more", message_history=history, usage_limits=limits) - - def test_run_agent_durable_applies_caching_model(self): + with patch.object(agent, "run_sync", return_value=mock_result) as mock_run_sync: + hook.run_agent(agent, 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.""" hook = PydanticAIHook() - mock_agent = MagicMock() - mock_agent.model = MagicMock() - mock_result = MagicMock() - mock_result.output = "ok" - 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="m") - mock_result.all_messages.return_value = [] - mock_agent.run_sync.return_value = mock_result - mock_agent.override.return_value.__enter__ = MagicMock(return_value=None) - mock_agent.override.return_value.__exit__ = MagicMock(return_value=False) + 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() @@ -528,37 +603,42 @@ def test_run_agent_durable_applies_caching_model(self): mock_counter.replayed_tool = 0 mock_counter.cached_model = 0 mock_counter.cached_tool = 0 - BaseAIHook._bind_agent_durable(mock_agent, mock_storage, mock_counter) + BaseAIHook._bind_agent_durable(agent, mock_storage, mock_counter) - with patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", return_value=MagicMock()): - with patch("airflow.providers.common.ai.durable.caching_model.CachingModel"): - request = AgentRunRequest(prompt="hi") - run_result = hook.run_agent(mock_agent, request) + request = AgentRunRequest(prompt="hi") + run_result = hook.run_agent(agent, request) - mock_agent.override.assert_called_once() + 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 - assert BaseAIHook._pop_agent_durable(mock_agent) is None + assert BaseAIHook._pop_agent_durable(agent) is None mock_storage.cleanup.assert_called_once() - def test_run_agent_preserves_durable_cache_on_exception(self): + @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, + ): hook = PydanticAIHook() - mock_agent = MagicMock() - mock_agent.model = MagicMock() - mock_agent.run_sync.side_effect = RuntimeError("boom") - mock_agent.override.return_value.__enter__ = MagicMock(return_value=None) - mock_agent.override.return_value.__exit__ = MagicMock(return_value=False) + agent = _test_agent() + mock_override.return_value = _noop_override_context() mock_storage = MagicMock() mock_counter = MagicMock() - BaseAIHook._bind_agent_durable(mock_agent, mock_storage, mock_counter) + BaseAIHook._bind_agent_durable(agent, mock_storage, mock_counter) - with patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", return_value=MagicMock()): - with patch("airflow.providers.common.ai.durable.caching_model.CachingModel"): - with pytest.raises(RuntimeError, match="boom"): - hook.run_agent(mock_agent, AgentRunRequest(prompt="hi")) + with pytest.raises(RuntimeError, match="boom"): + hook.run_agent(agent, AgentRunRequest(prompt="hi")) mock_storage.cleanup.assert_not_called() - assert BaseAIHook._pop_agent_durable(mock_agent) is None + assert BaseAIHook._pop_agent_durable(agent) is None def test_tool_spec_to_native_forwards_sequential(self): hook = PydanticAIHook() @@ -573,7 +653,7 @@ def fn() -> str: fn=fn, sequential=True, ) - with patch("pydantic_ai.tools.Tool") as mock_tool_cls: + with patch("airflow.providers.common.ai.hooks.pydantic_ai.Tool") as mock_tool_cls: hook._tool_spec_to_native(spec) mock_tool_cls.assert_called_once_with( fn, 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 9cba27de5c162..41a7146042a1a 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 @@ -56,6 +56,7 @@ def _make_agent_run_result(output, *, message_history=None) -> AgentRunResult: 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.supports_toolsets = True mock_hook.supports_durable = True mock_hook.supports_usage_limits = True @@ -69,9 +70,16 @@ 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.supports_toolsets = False mock_hook.supports_durable = False mock_hook.supports_usage_limits = False + + 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( 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 fabf21bfe5e0d..ea4ec9e45cd66 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,7 +24,7 @@ from pydantic import BaseModel from pydantic_ai.usage import UsageLimits -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, AgentUsage +from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.mixins.approval import ( LLMApprovalMixin, ) @@ -89,6 +89,22 @@ def test_execute_returns_string_output(self, mock_hook_cls): 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.supports_usage_limits = False + 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), + ) + with pytest.raises(ValueError, match="usage_limits are not supported"): + op.execute(context=MagicMock()) + @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.""" From 037f457908bf1451e11895f2b5db2aa119978599 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Wed, 27 May 2026 09:24:06 +0100 Subject: [PATCH 06/28] Fix static checks --- providers/common/ai/provider.yaml | 2 +- .../ai/src/airflow/providers/common/ai/get_provider_info.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index 14842636d8494..dfec53bfbf90d 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -62,7 +62,7 @@ integrations: tags: [ai] hooks: - - integration-name: Base Hook + - integration-name: Common AI python-modules: - airflow.providers.common.ai.hooks.base_ai - integration-name: Pydantic AI 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 d95d22efa084d..ba001ae19ed0d 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 @@ -68,7 +68,7 @@ def get_provider_info(): ], "hooks": [ { - "integration-name": "Base Hook", + "integration-name": "Common AI", "python-modules": ["airflow.providers.common.ai.hooks.base_ai"], }, { From afc7800e1751fd2261270aab9799e88c33bee9db Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sat, 30 May 2026 10:26:15 +0100 Subject: [PATCH 07/28] Resolve round 3 comments --- .../example_dags/example_callable_toolsets.py | 231 ++++++++++++++++++ .../providers/common/ai/hooks/base_ai.py | 82 ++++++- .../providers/common/ai/hooks/pydantic_ai.py | 37 ++- .../providers/common/ai/operators/agent.py | 13 +- .../unit/common/ai/hooks/test_base_ai.py | 138 ++++++++++- .../unit/common/ai/hooks/test_pydantic_ai.py | 93 ++++++- 6 files changed, 564 insertions(+), 30 deletions(-) create mode 100644 providers/common/ai/src/airflow/providers/common/ai/example_dags/example_callable_toolsets.py 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/hooks/base_ai.py b/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py index 3e4526116c086..4c93be99959fd 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py @@ -27,6 +27,8 @@ from dataclasses import dataclass, field from typing import Any, ClassVar +from typing_extensions import get_type_hints + from airflow.providers.common.compat.sdk import BaseHook _EMPTY_OBJECT_SCHEMA: dict[str, Any] = {"type": "object", "properties": {}} @@ -127,7 +129,9 @@ class AgentRunRequest: :param toolsets: List of :class:`BaseToolset` instances the agent may call. :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 each tool callable with a logging shim. + :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. @@ -332,15 +336,18 @@ def _resolve_tools( Three cases per item: * :class:`BaseToolset` — calls ``as_tools()`` and processes each :class:`ToolSpec`. - * Plain Python function (``def`` / ``lambda``) — auto-wraps into a :class:`ToolSpec` - using ``__name__`` and ``__doc__``, then processes it the same way. + * 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: *fn* → optional cache wrap → optional log wrap → :meth:`_tool_spec_to_native`. - :param toolsets: Mix of :class:`BaseToolset` instances, plain callables, and native tool objects. + :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 storage: ``DurableStorage`` instance, or ``None`` to skip caching. :param counter: ``DurableStepCounter`` instance, or ``None`` to skip caching. @@ -349,11 +356,17 @@ def _resolve_tools( for ts in toolsets: if isinstance(ts, BaseToolset): specs = ts.as_tools() - elif inspect.isfunction(ts): + elif callable(ts): + if isinstance(ts, functools.partial): + name = getattr(ts.func, "__name__", type(ts.func).__name__) + doc = ts.func.__doc__ or "" + else: + name = getattr(ts, "__name__", type(ts).__name__) + doc = ts.__doc__ or "" specs = [ ToolSpec( - name=ts.__name__, - description=ts.__doc__ or "", + name=name, + description=doc, parameters=_EMPTY_OBJECT_SCHEMA, fn=ts, ) @@ -382,8 +395,8 @@ def _logged_callable(fn: Callable[..., Any], logger: Any) -> Callable[..., Any]: """Wrap *fn* to log tool name, args, timing, and exceptions.""" @functools.wraps(fn) - def wrapper(*args: Any, **kwargs: Any) -> Any: - name = fn.__name__ + def wrapper(*args, **kwargs): + name = getattr(fn, "__name__", type(fn).__name__) logger.info("::group::Tool call: %s", name) if kwargs: logger.debug("Tool args: %s", json.dumps(kwargs, default=str)) @@ -400,6 +413,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: logger.info("::endgroup::") raise + BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) return wrapper @staticmethod @@ -411,7 +425,7 @@ def _cached_callable( """Wrap *fn* to cache its result in *storage* using a monotonic step counter.""" @functools.wraps(fn) - def wrapper(*args: Any, **kwargs: Any) -> Any: + def wrapper(*args, **kwargs): step = counter.next_step() key = f"tool_step_{step}" found, cached = storage.load_tool_result(key) @@ -423,4 +437,52 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: counter.cached_tool += 1 return result + BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) return wrapper + + @staticmethod + def _copy_wrapper_introspection_metadata( + fn: Callable[..., Any], + wrapper: Callable[..., Any], + ) -> None: + """ + Keep *wrapper* introspection aligned with *fn*. + + Generic logging/caching wrappers use ``*args`` and ``**kwargs``, which can + lose or mismatch the original callable's signature and annotations. This + helper copies consistent metadata onto the wrapper so schema/introspection + code sees the same callable shape as the original tool. + """ + signature_source, annotation_source = BaseAIHook._get_wrapper_metadata_sources(fn) + wrapper.__signature__ = inspect.signature(signature_source) + wrapper.__module__ = getattr(annotation_source, "__module__", __name__) + wrapper.__annotations__ = getattr(annotation_source, "__annotations__", {}).copy() + + try: + hints = get_type_hints(annotation_source, include_extras=True) + except (NameError, TypeError): + return + + annotations = {name: hints[name] for name in wrapper.__signature__.parameters if name in hints} + if "return" in hints: + annotations["return"] = hints["return"] + wrapper.__annotations__ = annotations + + @staticmethod + def _get_wrapper_metadata_sources( + fn: Callable[..., Any], + ) -> tuple[Callable[..., Any], Callable[..., Any]]: + """ + Return the best signature and annotation sources for *fn*. + + Most callables can use the same object for both. Partials need the bound + signature from the partial itself but annotations from the underlying + function, and callable objects expose their useful metadata on + ``obj.__call__``. + """ + if isinstance(fn, functools.partial): + return fn, fn.func + if inspect.ismethod(fn) or inspect.isfunction(fn): + return fn, fn + + return fn.__call__, fn.__call__ 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 fc8155e55184d..1b7fe28f32f0d 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 @@ -211,6 +211,11 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: :mod:`airflow.providers.common.ai.observability`. :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. """ self.validate_run_request(request) @@ -233,13 +238,28 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: pipeline_items = [ts for ts in request.toolsets if not isinstance(ts, AbstractToolset)] if pipeline_items: - resolved = self._resolve_tools( - pipeline_items, - request.enable_tool_logging, - storage, - counter, - ) + resolved: list[Any] = [] + for item in pipeline_items: + if isinstance(item, Tool): + # Native pydantic-ai Tool objects are preserved as-is. They keep their + # original schema/configuration, but do not receive Airflow's callable + # logging or durable tool-result caching wrappers. + resolved.append(item) + else: + resolved.extend( + self._resolve_tools( + [item], + request.enable_tool_logging, + storage, + counter, + ) + ) 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) @@ -255,6 +275,11 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: if request.enable_tool_logging: processed = [LoggingToolset(wrapped=ts, logger=self.log) 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], + ) agent = Agent( self.get_model(), 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 23ace9d96283f..9897883681b3a 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 @@ -113,9 +113,10 @@ class AgentOperator(BaseOperator, HITLReviewMixin): cannot be deserialized from XCom. :param toolsets: List of :class:`~airflow.providers.common.ai.hooks.base_ai.BaseToolset` instances the agent can use. - :param enable_tool_logging: When ``True`` (default), wraps each tool callable with a - logging shim that logs calls with timing at INFO level and arguments at DEBUG level. - Set to ``False`` to disable. + :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 @@ -125,8 +126,10 @@ class AgentOperator(BaseOperator, HITLReviewMixin): 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): diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py b/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py index c8a82b5fc2fbc..4b9406f5c1c7e 100644 --- a/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py @@ -16,6 +16,9 @@ # under the License. from __future__ import annotations +import functools +import inspect +from typing import get_type_hints from unittest.mock import MagicMock, patch import pytest @@ -274,7 +277,7 @@ def as_tools(self): mock_log.info.assert_called() def test_resolve_tools_wraps_plain_callable(self): - """A bare Python function is auto-wrapped using __name__ and __doc__.""" + """A plain function is auto-wrapped using __name__ and __doc__.""" class ConcreteHook(BaseAIHook): conn_type = "test" @@ -290,7 +293,7 @@ def run_agent(self, agent, request): return AgentRunResult(output="") def _tool_spec_to_native(self, spec): - return {"name": spec.name, "description": spec.description} + return {"name": spec.name, "description": spec.description, "fn": spec.fn} hook = ConcreteHook.__new__(ConcreteHook) @@ -303,6 +306,107 @@ def roll_dice() -> str: 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 create_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, storage=None, counter=None) + + 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 create_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, storage=None, counter=None) + + 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 create_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, storage=None, counter=None) + + 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.""" @@ -392,6 +496,36 @@ def failing(): logger.exception.assert_called_once() + 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")) + assert get_type_hints(wrapped) == { + "metric_name": str, + "return": float, + } + + 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",) + assert get_type_hints(wrapped) == { + "customer_id": str, + "return": dict[str, str], + } + class TestBaseAIHookCachedCallable: def test_cached_callable_saves_and_returns(self): 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 9b0385fb5ae4e..21e4a53141d60 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,6 +16,7 @@ # under the License. from __future__ import annotations +import functools import json import sys from unittest.mock import MagicMock, patch @@ -75,6 +76,17 @@ def _noop_override_context() -> MagicMock: 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_ai_hook(self): assert issubclass(PydanticAIHook, BaseAIHook) @@ -258,6 +270,66 @@ def test_get_conn_delegates_to_get_model(self): 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"] + @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): @@ -379,7 +451,7 @@ def test_create_agent_does_not_bind_durable_when_no_context(self, mock_infer_mod @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 the BaseToolset pipeline.""" + """Native pydantic-ai Tool objects bypass Airflow callable wrappers.""" from pydantic_ai.tools import Tool mock_model = MagicMock(spec=Model) @@ -398,7 +470,7 @@ def test_create_agent_passes_native_tools_through_directly(self, mock_agent_cls, @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 resolved; native Tool objects are passed through unchanged.""" + """BaseToolset items are expanded; native Tool objects are passed through unchanged.""" from pydantic_ai.tools import Tool from airflow.providers.common.ai.hooks.base_ai import BaseToolset, ToolSpec @@ -406,12 +478,18 @@ def test_create_agent_mixes_base_toolset_and_native_tool(self, mock_agent_cls, m mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model - def my_fn() -> str: - return "result" + def first_fn() -> str: + return "first" + + def second_fn() -> str: + return "second" class MyToolset(BaseToolset): def as_tools(self): - return [ToolSpec(name="my_fn", description="desc", parameters={}, fn=my_fn)] + 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) @@ -423,8 +501,9 @@ def as_tools(self): call_kwargs = mock_agent_cls.call_args[1] tools = call_kwargs["tools"] - assert len(tools) == 2 - assert any(t is native_tool for t in 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) From 077218f089b7f64e08698c91ea6a5bec00a3dd25 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sat, 30 May 2026 10:39:56 +0100 Subject: [PATCH 08/28] Rename base_ai to base --- providers/common/ai/AGENTS.md | 2 +- providers/common/ai/docs/changelog.rst | 2 +- providers/common/ai/docs/operators/agent.rst | 4 ++-- providers/common/ai/docs/toolsets.rst | 10 +++++----- providers/common/ai/provider.yaml | 2 +- .../ai/example_dags/example_pydantic_ai_hook.py | 2 +- .../airflow/providers/common/ai/get_provider_info.py | 5 +---- .../common/ai/hooks/{base_ai.py => base.py} | 0 .../airflow/providers/common/ai/hooks/pydantic_ai.py | 6 +++--- .../airflow/providers/common/ai/operators/agent.py | 6 +++--- .../src/airflow/providers/common/ai/operators/llm.py | 4 ++-- .../providers/common/ai/operators/llm_branch.py | 2 +- .../common/ai/operators/llm_file_analysis.py | 1 - .../common/ai/operators/llm_schema_compare.py | 2 +- .../airflow/providers/common/ai/operators/llm_sql.py | 2 +- .../airflow/providers/common/ai/policies/retry.py | 4 ++-- .../src/airflow/providers/common/ai/toolsets/sql.py | 2 +- .../src/airflow/providers/common/ai/utils/logging.py | 2 +- .../ai/tests/unit/common/ai/decorators/test_agent.py | 2 +- .../ai/tests/unit/common/ai/decorators/test_llm.py | 2 +- .../unit/common/ai/decorators/test_llm_branch.py | 2 +- .../common/ai/decorators/test_llm_file_analysis.py | 2 +- .../common/ai/decorators/test_llm_schema_compare.py | 2 +- .../tests/unit/common/ai/decorators/test_llm_sql.py | 2 +- .../ai/hooks/{test_base_ai.py => test_base.py} | 10 +++++----- .../tests/unit/common/ai/hooks/test_pydantic_ai.py | 12 ++++++------ .../ai/tests/unit/common/ai/operators/test_agent.py | 2 +- .../ai/tests/unit/common/ai/operators/test_llm.py | 2 +- .../unit/common/ai/operators/test_llm_branch.py | 2 +- .../common/ai/operators/test_llm_file_analysis.py | 2 +- .../common/ai/operators/test_llm_schema_compare.py | 2 +- .../tests/unit/common/ai/operators/test_llm_sql.py | 2 +- .../ai/tests/unit/common/ai/policies/test_retry.py | 2 +- .../ai/tests/unit/common/ai/toolsets/test_sql.py | 2 +- .../ai/tests/unit/common/ai/utils/test_logging.py | 2 +- 35 files changed, 53 insertions(+), 57 deletions(-) rename providers/common/ai/src/airflow/providers/common/ai/hooks/{base_ai.py => base.py} (100%) rename providers/common/ai/tests/unit/common/ai/hooks/{test_base_ai.py => test_base.py} (97%) diff --git a/providers/common/ai/AGENTS.md b/providers/common/ai/AGENTS.md index d59f1208313c4..15dbd055a95e2 100644 --- a/providers/common/ai/AGENTS.md +++ b/providers/common/ai/AGENTS.md @@ -70,7 +70,7 @@ building a wrapper here. ## Key Paths - Hooks: `src/airflow/providers/common/ai/hooks/pydantic_ai.py` (pydantic-ai) -- Base hook contract: `src/airflow/providers/common/ai/hooks/base_ai.py` +- 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/docs/changelog.rst b/providers/common/ai/docs/changelog.rst index 8b53791d5aa3a..a4acb696e01f9 100644 --- a/providers/common/ai/docs/changelog.rst +++ b/providers/common/ai/docs/changelog.rst @@ -58,7 +58,7 @@ Features .. For the next release: 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_ai.AgentRunRequest` instead of keyword arguments. + :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. diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index 7feb1f4661a9c..73d1f154eb4af 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -308,7 +308,7 @@ Parameters - ``output_type``: Expected output type (default: ``str``). Set to a Pydantic ``BaseModel`` for structured output. - ``toolsets``: List of toolsets the agent can use. Accepts - :class:`~airflow.providers.common.ai.hooks.base_ai.BaseToolset` subclasses + :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), @@ -317,7 +317,7 @@ Parameters - ``enable_tool_logging``: When ``True`` (default), wraps each tool call with real-time logging. For pydantic-ai ``AbstractToolset`` items this is done via :class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset`; for - plain callables and :class:`~airflow.providers.common.ai.hooks.base_ai.BaseToolset` + 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``). diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index 4ffc848373664..8b91db324f4a1 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -38,7 +38,7 @@ Three toolsets are included: :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_ai.BaseToolset` interface. +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. @@ -53,7 +53,7 @@ toolset to the correct agent parameter automatically. - Any third-party ``AbstractToolset``, including PydanticAI's own MCP server classes (``MCPServerStreamableHTTP``, ``MCPServerSSE``, ``MCPServerStdio``). - - :class:`~airflow.providers.common.ai.hooks.base_ai.BaseToolset` + - :class:`~airflow.providers.common.ai.hooks.base.BaseToolset` subclasses (``SQLToolset``). - Plain Python callables (``def my_tool(...): ...``). - Native pydantic-ai ``Tool`` objects. @@ -67,7 +67,7 @@ Using Toolsets Directly Toolsets can be used anywhere you can run Python within Airflow — ``@task`` functions, ``PythonOperator`` callables, or any custom operator's ``execute()`` method — without needing ``AgentOperator`` or ``@task.agent``. -Pass toolsets via :class:`~airflow.providers.common.ai.hooks.base_ai.AgentRunRequest` +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 @@ -85,9 +85,9 @@ 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_ai.BaseAIHook`): +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_ai.BaseToolset` tools and plain +: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. diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index dfec53bfbf90d..78da663821091 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -64,7 +64,7 @@ integrations: hooks: - integration-name: Common AI python-modules: - - airflow.providers.common.ai.hooks.base_ai + - 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/example_dags/example_pydantic_ai_hook.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_pydantic_ai_hook.py index a58514bc3cc9c..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 @@ -20,7 +20,7 @@ from pydantic import BaseModel -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, BaseAIHook +from airflow.providers.common.ai.hooks.base import AgentRunRequest, BaseAIHook from airflow.providers.common.compat.sdk import dag, task 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 ba001ae19ed0d..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,10 +67,7 @@ def get_provider_info(): }, ], "hooks": [ - { - "integration-name": "Common AI", - "python-modules": ["airflow.providers.common.ai.hooks.base_ai"], - }, + {"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_ai.py b/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py similarity index 100% rename from providers/common/ai/src/airflow/providers/common/ai/hooks/base_ai.py rename to providers/common/ai/src/airflow/providers/common/ai/hooks/base.py 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 1b7fe28f32f0d..fdabc116c3b0e 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 @@ -27,7 +27,7 @@ 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_ai import ( +from airflow.providers.common.ai.hooks.base import ( AgentRunRequest, AgentRunResult, AgentUsage, @@ -191,7 +191,7 @@ def _provider_factory(pname: str) -> Any: # ------------------------------------------------------------------ def _tool_spec_to_native(self, spec: ToolSpec) -> Any: - """Convert a :class:`~airflow.providers.common.ai.hooks.base_ai.ToolSpec` to a pydantic-ai ``Tool``.""" + """Convert a :class:`~airflow.providers.common.ai.hooks.base.ToolSpec` to a pydantic-ai ``Tool``.""" return Tool( spec.fn, name=spec.name, @@ -300,7 +300,7 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: return agent def run_agent(self, agent: Agent[None, Any], request: AgentRunRequest) -> AgentRunResult: - """Run *agent* synchronously for *request* and return a normalized :class:`~airflow.providers.common.ai.hooks.base_ai.AgentRunResult`.""" + """Run *agent* synchronously for *request* and return a normalized :class:`~airflow.providers.common.ai.hooks.base.AgentRunResult`.""" run_kwargs: dict[str, Any] = {} if request.message_history is not None: run_kwargs["message_history"] = request.message_history 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 9897883681b3a..f640338b264d9 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 @@ -26,7 +26,7 @@ from pydantic import BaseModel -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, BaseAIHook, DurableContext +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.output_type import ( @@ -111,7 +111,7 @@ 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 :class:`~airflow.providers.common.ai.hooks.base_ai.BaseToolset` + :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. @@ -229,7 +229,7 @@ def llm_hook(self) -> BaseAIHook: 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_ai.AgentRunRequest` from operator config.""" + """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 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 8672869fb72d4..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.base_ai import AgentRunRequest, BaseAIHook +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 @@ -50,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.base_ai.BaseAIHook` + 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 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 b23c932848612..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,7 +22,7 @@ from enum import Enum from typing import TYPE_CHECKING, Any -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest +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 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 2fc82b043cdd6..8ae086831aaca 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 @@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, Any 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 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 974eb7e25440b..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,7 +25,7 @@ from pydantic import BaseModel, Field -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest +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 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 3033a676f52f6..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,7 +34,7 @@ raise AirflowOptionalProviderFeatureException(e) -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest +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 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 abb25fb8514bc..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.base_ai.BaseAIHook` + 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. @@ -141,7 +141,7 @@ def _classify( ) -> RetryDecision: from pydantic_ai.settings import ModelSettings - from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, BaseAIHook + 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}) 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 6d50f3d9d7129..521016ab3bce1 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 @@ -36,7 +36,7 @@ from pydantic_ai.exceptions import ModelRetry -from airflow.providers.common.ai.hooks.base_ai import BaseToolset, ToolSpec +from airflow.providers.common.ai.hooks.base import BaseToolset, ToolSpec from airflow.providers.common.compat.sdk import BaseHook _LIST_TABLES_SCHEMA: dict[str, Any] = { 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 2b2f4e7ecce31..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 @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from airflow.providers.common.ai.hooks.base_ai import AgentRunResult + from airflow.providers.common.ai.hooks.base import AgentRunResult from airflow.sdk.types import Logger _MAX_OUTPUT_LEN = 500 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 c1876139c06f0..9a1a14169a8c0 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.hooks.base_ai import AgentRunResult, AgentUsage, BaseAIHook +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook try: from airflow.sdk.serde import SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as _CORE_WALKER 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 fb5939b174992..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,7 +22,7 @@ from pydantic_ai.messages import ImageUrl from airflow.providers.common.ai.decorators.llm import _LLMDecoratedOperator -from airflow.providers.common.ai.hooks.base_ai import AgentRunResult, AgentUsage, BaseAIHook +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook def _make_run_result(output): 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 dc2e3da1ff160..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,7 +23,7 @@ from pydantic_ai.messages import ImageUrl from airflow.providers.common.ai.decorators.llm_branch import _LLMBranchDecoratedOperator -from airflow.providers.common.ai.hooks.base_ai import AgentRunResult, AgentUsage, BaseAIHook +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator 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 7a1f5904a2939..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,7 +21,7 @@ import pytest from airflow.providers.common.ai.decorators.llm_file_analysis import _LLMFileAnalysisDecoratedOperator -from airflow.providers.common.ai.hooks.base_ai import AgentRunResult, AgentUsage, BaseAIHook +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.utils.file_analysis import FileAnalysisRequest 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 6746c1e88d09a..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,7 +22,7 @@ from pydantic_ai.messages import ImageUrl from airflow.providers.common.ai.decorators.llm_schema_compare import _LLMSchemaCompareDecoratedOperator -from airflow.providers.common.ai.hooks.base_ai import AgentRunResult, AgentUsage, BaseAIHook +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.operators.llm_schema_compare import ( LLMSchemaCompareOperator, SchemaCompareResult, 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 d5db22e512232..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,7 +22,7 @@ from pydantic_ai.messages import ImageUrl from airflow.providers.common.ai.decorators.llm_sql import _LLMSQLDecoratedOperator -from airflow.providers.common.ai.hooks.base_ai import AgentRunResult, AgentUsage, BaseAIHook +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook def _make_run_result(output): diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py b/providers/common/ai/tests/unit/common/ai/hooks/test_base.py similarity index 97% rename from providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py rename to providers/common/ai/tests/unit/common/ai/hooks/test_base.py index 4b9406f5c1c7e..aed65112fa8cb 100644 --- a/providers/common/ai/tests/unit/common/ai/hooks/test_base_ai.py +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_base.py @@ -25,7 +25,7 @@ from pydantic_ai import Agent from pydantic_ai.models.test import TestModel -from airflow.providers.common.ai.hooks.base_ai import ( +from airflow.providers.common.ai.hooks.base import ( AgentRunRequest, AgentRunResult, AgentUsage, @@ -39,8 +39,8 @@ class TestBaseAIHookGetAgentHook: - @patch("airflow.providers.common.ai.hooks.base_ai.BaseHook.get_hook", autospec=True) - def test_returns_hook_when_instance_is_base_ai_hook(self, mock_get_hook): + @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 @@ -49,8 +49,8 @@ def test_returns_hook_when_instance_is_base_ai_hook(self, mock_get_hook): assert result is mock_hook mock_get_hook.assert_called_once_with("my_conn", hook_params=None) - @patch("airflow.providers.common.ai.hooks.base_ai.BaseHook.get_hook", autospec=True) - def test_raises_when_hook_is_not_base_ai_hook(self, mock_get_hook): + @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"): 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 21e4a53141d60..2a20c69ea87c3 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 @@ -30,7 +30,7 @@ from pydantic_ai.usage import RunUsage, UsageLimits from airflow.models.connection import Connection -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, BaseAIHook, ToolSpec +from airflow.providers.common.ai.hooks.base import AgentRunRequest, AgentRunResult, BaseAIHook, ToolSpec from airflow.providers.common.ai.hooks.pydantic_ai import ( PydanticAIAzureHook, PydanticAIBedrockHook, @@ -88,7 +88,7 @@ def get_model(self) -> TestModel: class TestPydanticAIHookBaseContract: - def test_is_base_ai_hook(self): + def test_is_base_hook(self): assert issubclass(PydanticAIHook, BaseAIHook) def test_capability_flags(self): @@ -414,7 +414,7 @@ def test_create_agent_rejects_toolsets_in_agent_params_with_toolsets( @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_ai import DurableContext + from airflow.providers.common.ai.hooks.base import DurableContext mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model @@ -473,7 +473,7 @@ def test_create_agent_mixes_base_toolset_and_native_tool(self, mock_agent_cls, m """BaseToolset items are expanded; native Tool objects are passed through unchanged.""" from pydantic_ai.tools import Tool - from airflow.providers.common.ai.hooks.base_ai import BaseToolset, ToolSpec + from airflow.providers.common.ai.hooks.base import BaseToolset, ToolSpec mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model @@ -561,7 +561,7 @@ def test_create_agent_wraps_abstract_toolset_with_caching_when_durable( from pydantic_ai.toolsets.abstract import AbstractToolset from airflow.providers.common.ai.durable.caching_toolset import CachingToolset - from airflow.providers.common.ai.hooks.base_ai import DurableContext + from airflow.providers.common.ai.hooks.base import DurableContext from airflow.providers.common.ai.toolsets.logging import LoggingToolset mock_model = MagicMock(spec=Model) @@ -594,7 +594,7 @@ def test_create_agent_wraps_abstract_toolset_with_caching_when_durable( @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) def test_create_agent_binds_durable_per_agent_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_ai import DurableContext + from airflow.providers.common.ai.hooks.base import DurableContext mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model 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 41a7146042a1a..fb3756a2d6304 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,7 +23,7 @@ from pydantic import BaseModel from pydantic_ai.usage import UsageLimits -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, AgentUsage, BaseAIHook +from airflow.providers.common.ai.hooks.base import AgentRunRequest, AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.operators.agent import AgentOperator, HITLReviewLink from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS 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 ea4ec9e45cd66..662aa1f7c7c3c 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,7 +24,7 @@ from pydantic import BaseModel from pydantic_ai.usage import UsageLimits -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, AgentUsage, BaseAIHook +from airflow.providers.common.ai.hooks.base import AgentRunRequest, AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.mixins.approval import ( LLMApprovalMixin, ) 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 0edc93746d98f..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,7 +21,7 @@ import pytest -from airflow.providers.common.ai.hooks.base_ai import AgentRunResult, AgentUsage +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 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 45e4cd7c67302..19fb67bf853a5 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,7 +23,7 @@ import pytest from pydantic import BaseModel -from airflow.providers.common.ai.hooks.base_ai import AgentRunResult, AgentUsage +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage from airflow.providers.common.ai.operators.llm_file_analysis import LLMFileAnalysisOperator from airflow.providers.common.ai.utils.file_analysis import FileAnalysisRequest 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 0e3892d6acd34..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,7 +21,7 @@ import pytest -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, AgentUsage +from airflow.providers.common.ai.hooks.base import AgentRunRequest, AgentRunResult, AgentUsage from airflow.providers.common.ai.operators.llm_schema_compare import ( LLMSchemaCompareOperator, SchemaCompareResult, 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 e140b0a5ed71d..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,7 +22,7 @@ import pytest -from airflow.providers.common.ai.hooks.base_ai import AgentRunRequest, AgentRunResult, AgentUsage +from airflow.providers.common.ai.hooks.base import AgentRunRequest, AgentRunResult, AgentUsage from airflow.providers.common.ai.mixins.approval import ( LLMApprovalMixin, ) 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 c358347b953b3..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,7 +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_ai import AgentRunResult, AgentUsage, BaseAIHook +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage, BaseAIHook from airflow.providers.common.ai.policies.retry import ( ErrorClassification, LLMRetryPolicy, 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 51fc76f6e7b54..067d0fe05bd20 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 @@ -24,7 +24,7 @@ import pytest from pydantic_ai.exceptions import ModelRetry -from airflow.providers.common.ai.hooks.base_ai import BaseToolset, ToolSpec +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 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 f943365218291..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 @@ -20,7 +20,7 @@ from pydantic import BaseModel -from airflow.providers.common.ai.hooks.base_ai import AgentRunResult, AgentUsage +from airflow.providers.common.ai.hooks.base import AgentRunResult, AgentUsage from airflow.providers.common.ai.utils.logging import ( _log_output_debug, log_run_summary, From 7169214de45ccdf8e9e59337e19a9a04c328dbe7 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sat, 30 May 2026 11:05:58 +0100 Subject: [PATCH 09/28] Rebase --- providers/common/ai/docs/changelog.rst | 19 +++++-------------- .../unit/common/ai/operators/test_llm.py | 9 +++++---- 2 files changed, 10 insertions(+), 18 deletions(-) diff --git a/providers/common/ai/docs/changelog.rst b/providers/common/ai/docs/changelog.rst index a4acb696e01f9..7b21a1b73eaad 100644 --- a/providers/common/ai/docs/changelog.rst +++ b/providers/common/ai/docs/changelog.rst @@ -47,20 +47,11 @@ 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. -Next release -............ - -Features -^^^^^^^^ - -* Add ``BaseAIHook`` contract with framework-agnostic ``create_agent`` / ``run_agent`` / - ``get_model`` interface so ``AgentOperator`` selects the agent backend via connection type. - -.. For the next release: 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. +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. 0.4.0 ..... 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 662aa1f7c7c3c..565865775fb44 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 @@ -163,12 +163,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", From 881ec59cfaf7f821bf4bf7e5a02afae2f991ef51 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sat, 30 May 2026 14:32:55 +0100 Subject: [PATCH 10/28] Fix mypy and refactor introspection --- docs/spelling_wordlist.txt | 1 + .../airflow/providers/common/ai/hooks/base.py | 33 ++-- .../tests/unit/common/ai/hooks/test_base.py | 147 ++++++++++++++++++ 3 files changed, 170 insertions(+), 11 deletions(-) diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index b0f6fd8d93927..b6c1c022d01f7 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -677,6 +677,7 @@ fsGroup fsspec fullname func +functools ga Gantt gantt 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 index 4c93be99959fd..17a22b4fc3e5e 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py @@ -25,7 +25,7 @@ from abc import ABCMeta, abstractmethod from collections.abc import Callable, Sequence from dataclasses import dataclass, field -from typing import Any, ClassVar +from typing import Any, ClassVar, cast from typing_extensions import get_type_hints @@ -454,19 +454,26 @@ def _copy_wrapper_introspection_metadata( code sees the same callable shape as the original tool. """ signature_source, annotation_source = BaseAIHook._get_wrapper_metadata_sources(fn) - wrapper.__signature__ = inspect.signature(signature_source) - wrapper.__module__ = getattr(annotation_source, "__module__", __name__) - wrapper.__annotations__ = getattr(annotation_source, "__annotations__", {}).copy() + try: + signature = inspect.signature(signature_source) + except (ValueError, TypeError): + return + setattr(wrapper, "__signature__", signature) + wrapper.__module__ = cast( + "str", + getattr(annotation_source, "__module__", None) or getattr(fn, "__module__", __name__), + ) try: hints = get_type_hints(annotation_source, include_extras=True) except (NameError, TypeError): + wrapper.__annotations__ = getattr(annotation_source, "__annotations__", {}).copy() return - annotations = {name: hints[name] for name in wrapper.__signature__.parameters if name in hints} + resolved = {name: hints[name] for name in signature.parameters if name in hints} if "return" in hints: - annotations["return"] = hints["return"] - wrapper.__annotations__ = annotations + resolved["return"] = hints["return"] + wrapper.__annotations__ = resolved @staticmethod def _get_wrapper_metadata_sources( @@ -477,12 +484,16 @@ def _get_wrapper_metadata_sources( Most callables can use the same object for both. Partials need the bound signature from the partial itself but annotations from the underlying - function, and callable objects expose their useful metadata on - ``obj.__call__``. + function (unwrapped through any nesting), and callable objects expose + their useful metadata on ``obj.__call__``. """ if isinstance(fn, functools.partial): - return fn, fn.func + inner: Callable[..., Any] = fn.func + while isinstance(inner, functools.partial): + inner = inner.func + return fn, inner if inspect.ismethod(fn) or inspect.isfunction(fn): return fn, fn - return fn.__call__, fn.__call__ + call = cast("Callable[..., Any]", object.__getattribute__(fn, "__call__")) + return call, call 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 index aed65112fa8cb..f04b961a1515d 100644 --- a/providers/common/ai/tests/unit/common/ai/hooks/test_base.py +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_base.py @@ -468,6 +468,153 @@ def as_tools(self): assert result == ["converted:greet", native_tool] +class TestGetWrapperMetadataSources: + def test_plain_function_returns_self_for_both(self): + def fn(x: int) -> str: + return str(x) + + sig_src, ann_src = BaseAIHook._get_wrapper_metadata_sources(fn) + assert sig_src is fn + assert ann_src is fn + + def test_bound_method_returns_self_for_both(self): + class MyClass: + def method(self, x: int) -> str: + return str(x) + + obj = MyClass() + # Store once — each attribute access on an instance creates a new bound-method object. + bound = obj.method + sig_src, ann_src = BaseAIHook._get_wrapper_metadata_sources(bound) + assert sig_src is bound + assert ann_src is bound + + def test_simple_partial_returns_partial_and_underlying_func(self): + def fn(x: int, y: str) -> float: + return 1.0 + + p = functools.partial(fn, x=1) + sig_src, ann_src = BaseAIHook._get_wrapper_metadata_sources(p) + assert sig_src is p + assert ann_src is fn + + def test_nested_partial_unwraps_to_original_function(self): + def fn(url: str, method: str, timeout: int) -> str: + return "" + + p1 = functools.partial(fn, url="https://example.com") + p2 = functools.partial(p1, method="POST") + + sig_src, ann_src = BaseAIHook._get_wrapper_metadata_sources(p2) + assert sig_src is p2 + assert ann_src is fn + + def test_callable_object_returns_call_method(self): + class Searcher: + def __call__(self, query: str) -> str: + return query + + obj = Searcher() + sig_src, ann_src = BaseAIHook._get_wrapper_metadata_sources(obj) + assert sig_src is ann_src + assert tuple(inspect.signature(sig_src).parameters) == ("query",) + + +class TestCopyWrapperIntrospectionMetadata: + def test_plain_function_copies_signature_and_annotations(self): + def fn(x: int, y: str) -> float: + return 1.0 + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + return fn(*args, **kwargs) + + BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) + + assert inspect.signature(wrapper) == inspect.signature(fn) + assert get_type_hints(wrapper) == {"x": int, "y": str, "return": float} + + def test_nested_partial_annotations_resolved_from_underlying_function(self): + def fn(url: str, method: str, timeout: int) -> str: + return "" + + p1 = functools.partial(fn, url="https://example.com") + p2 = functools.partial(p1, method="POST") + + @functools.wraps(p2) + def wrapper(*args, **kwargs): + return p2(*args, **kwargs) + + BaseAIHook._copy_wrapper_introspection_metadata(p2, wrapper) + + sig = inspect.signature(wrapper) + # inspect.signature keeps bound params with defaults rather than removing them. + assert set(sig.parameters) == {"url", "method", "timeout"} + # Annotations are resolved from the unwrapped underlying function, covering all params. + assert get_type_hints(wrapper) == {"url": str, "method": str, "timeout": int, "return": str} + + def test_uninspectable_signature_returns_early_without_raising(self): + def fn(x: int) -> str: + return str(x) + + # Plain wrapper with no __signature__ pre-set. + def wrapper(*args, **kwargs): + return fn(*args, **kwargs) + + with patch( + "airflow.providers.common.ai.hooks.base.inspect.signature", + side_effect=ValueError("no introspectable signature"), + ): + BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) + + # Must not raise, and __signature__ must not be set on the wrapper. + assert not hasattr(wrapper, "__signature__") + + def test_get_type_hints_failure_falls_back_to_raw_annotations(self): + def fn(x: int) -> str: + return str(x) + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + return fn(*args, **kwargs) + + with patch( + "airflow.providers.common.ai.hooks.base.get_type_hints", + side_effect=NameError("NonExistentType"), + ): + BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) + + # Falls back to raw __annotations__ from functools.wraps + assert "x" in wrapper.__annotations__ + + def test_module_falls_back_to_fn_module_when_annotation_source_lacks_it(self): + def fn(x: int) -> int: + return x + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + return fn(*args, **kwargs) + + BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) + assert wrapper.__module__ == fn.__module__ + + def test_callable_object_copies_call_signature(self): + class Lookup: + def __call__(self, customer_id: str) -> dict: + return {} + + obj = Lookup() + + @functools.wraps(obj) + def wrapper(*args, **kwargs): + return obj(*args, **kwargs) + + BaseAIHook._copy_wrapper_introspection_metadata(obj, wrapper) + + sig = inspect.signature(wrapper) + assert tuple(sig.parameters) == ("customer_id",) + + class TestBaseAIHookLoggedCallable: def test_logged_callable_logs_and_returns(self): logger = MagicMock() From abd834f947c56d49d3c5c40335a94a0b4ffcf71a Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sat, 30 May 2026 14:51:04 +0100 Subject: [PATCH 11/28] Fix tool name in logs --- .../airflow/providers/common/ai/hooks/base.py | 17 +++++--- .../tests/unit/common/ai/hooks/test_base.py | 39 +++++++++++++++++++ 2 files changed, 50 insertions(+), 6 deletions(-) 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 index 17a22b4fc3e5e..70420d16fc8bb 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py @@ -379,7 +379,7 @@ def _resolve_tools( if storage is not None and counter is not None: fn = self._cached_callable(fn, storage, counter) if enable_logging: - fn = self._logged_callable(fn, self.log) + fn = self._logged_callable(fn, self.log, name=spec.name) adapted = ToolSpec( name=spec.name, description=spec.description, @@ -391,25 +391,30 @@ def _resolve_tools( return native @staticmethod - def _logged_callable(fn: Callable[..., Any], logger: Any) -> Callable[..., Any]: + 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__) @functools.wraps(fn) def wrapper(*args, **kwargs): - name = getattr(fn, "__name__", type(fn).__name__) - logger.info("::group::Tool call: %s", name) + 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", name, elapsed) + logger.info("Tool %s returned in %.2fs", _tool_name, elapsed) logger.info("::endgroup::") return result except Exception: elapsed = time.monotonic() - start - logger.exception("Tool %s failed after %.2fs", name, elapsed) + logger.exception("Tool %s failed after %.2fs", _tool_name, elapsed) logger.info("::endgroup::") raise 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 index f04b961a1515d..2042b23402dee 100644 --- a/providers/common/ai/tests/unit/common/ai/hooks/test_base.py +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_base.py @@ -643,6 +643,45 @@ def failing(): 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() From 5fe951d5fa6cd186276f628d9560f34a1eb447c6 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sat, 30 May 2026 22:38:04 +0100 Subject: [PATCH 12/28] Bump pydantic-ai-slim version --- providers/common/ai/README.rst | 2 +- providers/common/ai/docs/index.rst | 2 +- providers/common/ai/pyproject.toml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/providers/common/ai/README.rst b/providers/common/ai/README.rst index 8be522e21cf19..af4b697da3f42 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.102.0`` ========================================== ================== Cross provider package dependencies diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index 4b3014ca1ffbb..868f849e2620d 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.102.0`` ========================================== ================== Cross provider package dependencies diff --git a/providers/common/ai/pyproject.toml b/providers/common/ai/pyproject.toml index 910ef1401c978..d7d15fedc1df0 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.102.0", ] # The optional dependencies should be modified in place in the generated file From 76a8e258b70b6e8e9c9c94390dd3562efcb3e59c Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sun, 31 May 2026 04:26:48 +0100 Subject: [PATCH 13/28] Simplify tool discovery schema and assignment --- providers/common/ai/README.rst | 2 +- providers/common/ai/docs/index.rst | 2 +- providers/common/ai/pyproject.toml | 2 +- .../airflow/providers/common/ai/hooks/base.py | 83 +--- .../providers/common/ai/hooks/pydantic_ai.py | 5 +- .../common/ai/utils/function_schema.py | 163 ++++++++ .../tests/unit/common/ai/hooks/test_base.py | 156 -------- .../unit/common/ai/hooks/test_pydantic_ai.py | 27 +- .../common/ai/utils/test_function_schema.py | 373 ++++++++++++++++++ 9 files changed, 558 insertions(+), 255 deletions(-) create mode 100644 providers/common/ai/src/airflow/providers/common/ai/utils/function_schema.py create mode 100644 providers/common/ai/tests/unit/common/ai/utils/test_function_schema.py diff --git a/providers/common/ai/README.rst b/providers/common/ai/README.rst index af4b697da3f42..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.102.0`` +``pydantic-ai-slim`` ``>=1.96.0`` ========================================== ================== Cross provider package dependencies diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index 868f849e2620d..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.102.0`` +``pydantic-ai-slim`` ``>=1.96.0`` ========================================== ================== Cross provider package dependencies diff --git a/providers/common/ai/pyproject.toml b/providers/common/ai/pyproject.toml index d7d15fedc1df0..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.102.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/hooks/base.py b/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py index 70420d16fc8bb..d8d8589649c59 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py @@ -19,20 +19,16 @@ from __future__ import annotations import functools -import inspect import json import time from abc import ABCMeta, abstractmethod from collections.abc import Callable, Sequence from dataclasses import dataclass, field -from typing import Any, ClassVar, cast - -from typing_extensions import get_type_hints +from typing import Any, ClassVar +from airflow.providers.common.ai.utils.function_schema import callable_to_tool_spec from airflow.providers.common.compat.sdk import BaseHook -_EMPTY_OBJECT_SCHEMA: dict[str, Any] = {"type": "object", "properties": {}} - # Attribute name for durable storage/counter bound to a framework agent instance. _AIRFLOW_DURABLE_ATTR = "_airflow_durable_state" @@ -357,20 +353,7 @@ def _resolve_tools( if isinstance(ts, BaseToolset): specs = ts.as_tools() elif callable(ts): - if isinstance(ts, functools.partial): - name = getattr(ts.func, "__name__", type(ts.func).__name__) - doc = ts.func.__doc__ or "" - else: - name = getattr(ts, "__name__", type(ts).__name__) - doc = ts.__doc__ or "" - specs = [ - ToolSpec( - name=name, - description=doc, - parameters=_EMPTY_OBJECT_SCHEMA, - fn=ts, - ) - ] + specs = [callable_to_tool_spec(ts)] else: native.append(ts) continue @@ -418,7 +401,6 @@ def wrapper(*args, **kwargs): logger.info("::endgroup::") raise - BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) return wrapper @staticmethod @@ -442,63 +424,4 @@ def wrapper(*args, **kwargs): counter.cached_tool += 1 return result - BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) return wrapper - - @staticmethod - def _copy_wrapper_introspection_metadata( - fn: Callable[..., Any], - wrapper: Callable[..., Any], - ) -> None: - """ - Keep *wrapper* introspection aligned with *fn*. - - Generic logging/caching wrappers use ``*args`` and ``**kwargs``, which can - lose or mismatch the original callable's signature and annotations. This - helper copies consistent metadata onto the wrapper so schema/introspection - code sees the same callable shape as the original tool. - """ - signature_source, annotation_source = BaseAIHook._get_wrapper_metadata_sources(fn) - try: - signature = inspect.signature(signature_source) - except (ValueError, TypeError): - return - setattr(wrapper, "__signature__", signature) - wrapper.__module__ = cast( - "str", - getattr(annotation_source, "__module__", None) or getattr(fn, "__module__", __name__), - ) - - try: - hints = get_type_hints(annotation_source, include_extras=True) - except (NameError, TypeError): - wrapper.__annotations__ = getattr(annotation_source, "__annotations__", {}).copy() - return - - resolved = {name: hints[name] for name in signature.parameters if name in hints} - if "return" in hints: - resolved["return"] = hints["return"] - wrapper.__annotations__ = resolved - - @staticmethod - def _get_wrapper_metadata_sources( - fn: Callable[..., Any], - ) -> tuple[Callable[..., Any], Callable[..., Any]]: - """ - Return the best signature and annotation sources for *fn*. - - Most callables can use the same object for both. Partials need the bound - signature from the partial itself but annotations from the underlying - function (unwrapped through any nesting), and callable objects expose - their useful metadata on ``obj.__call__``. - """ - if isinstance(fn, functools.partial): - inner: Callable[..., Any] = fn.func - while isinstance(inner, functools.partial): - inner = inner.func - return fn, inner - if inspect.ismethod(fn) or inspect.isfunction(fn): - return fn, fn - - call = cast("Callable[..., Any]", object.__getattribute__(fn, "__call__")) - return call, call 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 fdabc116c3b0e..6537e3a65beff 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 @@ -192,11 +192,12 @@ def _provider_factory(pname: str) -> Any: 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( + return Tool.from_schema( spec.fn, name=spec.name, description=spec.description, sequential=spec.sequential, + json_schema=spec.parameters, ) def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: @@ -324,7 +325,7 @@ def run_agent(self, agent: Agent[None, Any], request: AgentRunRequest) -> AgentR result = agent.run_sync(request.prompt, **run_kwargs) else: result = agent.run_sync(request.prompt, **run_kwargs) - + print(result.all_messages()) usage = result.usage tool_names: list[str] = [] for message in result.all_messages(): 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..6a06f8b4a5238 --- /dev/null +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/function_schema.py @@ -0,0 +1,163 @@ +# 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_function_description(fn: Callable[..., Any]) -> str: + """Return the first paragraph of *fn*'s docstring, stopping before Args/Returns sections.""" + doc = inspect.getdoc(fn) + if not doc: + return getattr(fn, "__name__", type(fn).__name__) + 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() or getattr(fn, "__name__", type(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. + 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 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 desc else Field(default=default), + ) + + 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 + + if isinstance(fn, functools.partial): + inner: Callable[..., Any] = fn.func + while isinstance(inner, functools.partial): + inner = inner.func + name = getattr(inner, "__name__", type(inner).__name__) + source = inner + else: + name = getattr(fn, "__name__", type(fn).__name__) + source = fn + + return ToolSpec( + name=name, + description=extract_function_description(source), + parameters=build_function_json_schema(fn), # fn not source: sig already reflects bound args + fn=fn, + ) 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 index 2042b23402dee..98787e0c6a36d 100644 --- a/providers/common/ai/tests/unit/common/ai/hooks/test_base.py +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_base.py @@ -18,7 +18,6 @@ import functools import inspect -from typing import get_type_hints from unittest.mock import MagicMock, patch import pytest @@ -468,153 +467,6 @@ def as_tools(self): assert result == ["converted:greet", native_tool] -class TestGetWrapperMetadataSources: - def test_plain_function_returns_self_for_both(self): - def fn(x: int) -> str: - return str(x) - - sig_src, ann_src = BaseAIHook._get_wrapper_metadata_sources(fn) - assert sig_src is fn - assert ann_src is fn - - def test_bound_method_returns_self_for_both(self): - class MyClass: - def method(self, x: int) -> str: - return str(x) - - obj = MyClass() - # Store once — each attribute access on an instance creates a new bound-method object. - bound = obj.method - sig_src, ann_src = BaseAIHook._get_wrapper_metadata_sources(bound) - assert sig_src is bound - assert ann_src is bound - - def test_simple_partial_returns_partial_and_underlying_func(self): - def fn(x: int, y: str) -> float: - return 1.0 - - p = functools.partial(fn, x=1) - sig_src, ann_src = BaseAIHook._get_wrapper_metadata_sources(p) - assert sig_src is p - assert ann_src is fn - - def test_nested_partial_unwraps_to_original_function(self): - def fn(url: str, method: str, timeout: int) -> str: - return "" - - p1 = functools.partial(fn, url="https://example.com") - p2 = functools.partial(p1, method="POST") - - sig_src, ann_src = BaseAIHook._get_wrapper_metadata_sources(p2) - assert sig_src is p2 - assert ann_src is fn - - def test_callable_object_returns_call_method(self): - class Searcher: - def __call__(self, query: str) -> str: - return query - - obj = Searcher() - sig_src, ann_src = BaseAIHook._get_wrapper_metadata_sources(obj) - assert sig_src is ann_src - assert tuple(inspect.signature(sig_src).parameters) == ("query",) - - -class TestCopyWrapperIntrospectionMetadata: - def test_plain_function_copies_signature_and_annotations(self): - def fn(x: int, y: str) -> float: - return 1.0 - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - return fn(*args, **kwargs) - - BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) - - assert inspect.signature(wrapper) == inspect.signature(fn) - assert get_type_hints(wrapper) == {"x": int, "y": str, "return": float} - - def test_nested_partial_annotations_resolved_from_underlying_function(self): - def fn(url: str, method: str, timeout: int) -> str: - return "" - - p1 = functools.partial(fn, url="https://example.com") - p2 = functools.partial(p1, method="POST") - - @functools.wraps(p2) - def wrapper(*args, **kwargs): - return p2(*args, **kwargs) - - BaseAIHook._copy_wrapper_introspection_metadata(p2, wrapper) - - sig = inspect.signature(wrapper) - # inspect.signature keeps bound params with defaults rather than removing them. - assert set(sig.parameters) == {"url", "method", "timeout"} - # Annotations are resolved from the unwrapped underlying function, covering all params. - assert get_type_hints(wrapper) == {"url": str, "method": str, "timeout": int, "return": str} - - def test_uninspectable_signature_returns_early_without_raising(self): - def fn(x: int) -> str: - return str(x) - - # Plain wrapper with no __signature__ pre-set. - def wrapper(*args, **kwargs): - return fn(*args, **kwargs) - - with patch( - "airflow.providers.common.ai.hooks.base.inspect.signature", - side_effect=ValueError("no introspectable signature"), - ): - BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) - - # Must not raise, and __signature__ must not be set on the wrapper. - assert not hasattr(wrapper, "__signature__") - - def test_get_type_hints_failure_falls_back_to_raw_annotations(self): - def fn(x: int) -> str: - return str(x) - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - return fn(*args, **kwargs) - - with patch( - "airflow.providers.common.ai.hooks.base.get_type_hints", - side_effect=NameError("NonExistentType"), - ): - BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) - - # Falls back to raw __annotations__ from functools.wraps - assert "x" in wrapper.__annotations__ - - def test_module_falls_back_to_fn_module_when_annotation_source_lacks_it(self): - def fn(x: int) -> int: - return x - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - return fn(*args, **kwargs) - - BaseAIHook._copy_wrapper_introspection_metadata(fn, wrapper) - assert wrapper.__module__ == fn.__module__ - - def test_callable_object_copies_call_signature(self): - class Lookup: - def __call__(self, customer_id: str) -> dict: - return {} - - obj = Lookup() - - @functools.wraps(obj) - def wrapper(*args, **kwargs): - return obj(*args, **kwargs) - - BaseAIHook._copy_wrapper_introspection_metadata(obj, wrapper) - - sig = inspect.signature(wrapper) - assert tuple(sig.parameters) == ("customer_id",) - - class TestBaseAIHookLoggedCallable: def test_logged_callable_logs_and_returns(self): logger = MagicMock() @@ -691,10 +543,6 @@ def fetch_metric(environment: str, metric_name: str) -> float: wrapped = BaseAIHook._logged_callable(functools.partial(fetch_metric, "prod"), logger) assert inspect.signature(wrapped) == inspect.signature(functools.partial(fetch_metric, "prod")) - assert get_type_hints(wrapped) == { - "metric_name": str, - "return": float, - } def test_logged_callable_preserves_callable_object_introspection(self): logger = MagicMock() @@ -707,10 +555,6 @@ def __call__(self, customer_id: str) -> dict[str, str]: signature = inspect.signature(wrapped) assert tuple(signature.parameters) == ("customer_id",) - assert get_type_hints(wrapped) == { - "customer_id": str, - "return": dict[str, str], - } class TestBaseAIHookCachedCallable: 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 2a20c69ea87c3..2f154126744f2 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 @@ -473,7 +473,7 @@ def test_create_agent_mixes_base_toolset_and_native_tool(self, mock_agent_cls, m """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, ToolSpec + from airflow.providers.common.ai.hooks.base import BaseToolset mock_model = MagicMock(spec=Model) mock_infer_model.return_value = mock_model @@ -719,26 +719,25 @@ def test_run_agent_preserves_durable_cache_on_exception( mock_storage.cleanup.assert_not_called() assert BaseAIHook._pop_agent_durable(agent) is None - def test_tool_spec_to_native_forwards_sequential(self): + def test_tool_spec_to_native_tools_called(self): hook = PydanticAIHook() - def fn() -> str: + def fn(customer_id: int) -> str: + """test function""" return "ok" - spec = ToolSpec( - name="fn", - description="desc", - parameters={"type": "object", "properties": {}}, - fn=fn, - sequential=True, - ) with patch("airflow.providers.common.ai.hooks.pydantic_ai.Tool") as mock_tool_cls: - hook._tool_spec_to_native(spec) - mock_tool_cls.assert_called_once_with( + hook._resolve_tools(toolsets=[fn], enable_logging=False, storage=None, counter=None) + mock_tool_cls.from_schema.assert_called_once_with( fn, name="fn", - description="desc", - sequential=True, + description="test function", + sequential=False, + json_schema={ + "properties": {"customer_id": {"type": "integer"}}, + "required": ["customer_id"], + "type": "object", + }, ) 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..965710d3853e5 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/utils/test_function_schema.py @@ -0,0 +1,373 @@ +# 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 +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: + """Do something useful. + + Args: + x: a number. + y: a string. + + Returns: + A string result. + """ + + +def _no_doc(x: int) -> str: + pass + + +def _no_params() -> None: + """No parameters at all.""" + + +def _annotated(q: Annotated[str, "The search query"], limit: Annotated[int, "Max results"] = 10) -> list: + """Search.""" + + +class _CallableObj: + """Callable object used to test non-function callables.""" + + def __call__(self, value: str) -> str: + """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_falls_back_to_class_name(self): + obj = _CallableObj() + # callable objects have no __name__; falls back to type name + result = extract_function_description(obj) + assert result == "_CallableObj" + + @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_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: ... + + 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: ... + + 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: ... + + 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_introspection_failure_returns_empty_schema(self): + # Built-in functions have no inspectable signature. + schema = build_function_json_schema(len) + 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 From 0c6c14e4d24010dd0c38bc1010d96708cf9c6457 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sun, 31 May 2026 07:07:20 +0100 Subject: [PATCH 14/28] Fixup tests --- .../providers/common/ai/hooks/pydantic_ai.py | 1 - .../common/ai/utils/function_schema.py | 38 ++++++++++--------- 2 files changed, 21 insertions(+), 18 deletions(-) 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 6537e3a65beff..809c2038fc22e 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 @@ -325,7 +325,6 @@ def run_agent(self, agent: Agent[None, Any], request: AgentRunRequest) -> AgentR result = agent.run_sync(request.prompt, **run_kwargs) else: result = agent.run_sync(request.prompt, **run_kwargs) - print(result.all_messages()) usage = result.usage tool_names: list[str] = [] for message in result.all_messages(): 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 index 6a06f8b4a5238..a4464826648e6 100644 --- 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 @@ -55,15 +55,24 @@ def extract_function_description(fn: Callable[..., Any]) -> str: """Return the first paragraph of *fn*'s docstring, stopping before Args/Returns sections.""" + # 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__; use the class name. + if not hasattr(fn, "__name__"): + return type(fn).__name__ + doc = inspect.getdoc(fn) + name: str = fn.__name__ # type: ignore[assignment] if not doc: - return getattr(fn, "__name__", type(fn).__name__) + return name 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() or getattr(fn, "__name__", type(fn).__name__) + return "\n".join(result).strip() or name def build_function_json_schema(fn: Callable[..., Any]) -> dict[str, Any]: @@ -77,6 +86,9 @@ def build_function_json_schema(fn: Callable[..., Any]) -> dict[str, Any]: ``self``, ``cls``, ``*args``, and ``**kwargs`` are excluded. For ``functools.partial``, only the remaining free parameters appear. """ + if inspect.isbuiltin(fn): + return _EMPTY_OBJECT_SCHEMA + # Partials: sig from partial (bound args already removed), hints from inner fn. hint_source: Callable[..., Any] = fn if isinstance(fn, functools.partial): @@ -114,10 +126,7 @@ def build_function_json_schema(fn: Callable[..., Any]) -> dict[str, Any]: actual_type = annotation desc = None - field_defs[param_name] = ( - actual_type, - Field(default=default, description=desc) if desc else Field(default=default), - ) + field_defs[param_name] = (actual_type, Field(default=default, description=desc)) if not field_defs: return _EMPTY_OBJECT_SCHEMA @@ -145,19 +154,14 @@ def callable_to_tool_spec(fn: Callable[..., Any]) -> ToolSpec: # this module imports ToolSpec from base. from airflow.providers.common.ai.hooks.base import ToolSpec - if isinstance(fn, functools.partial): - inner: Callable[..., Any] = fn.func - while isinstance(inner, functools.partial): - inner = inner.func - name = getattr(inner, "__name__", type(inner).__name__) - source = inner - else: - name = getattr(fn, "__name__", type(fn).__name__) - source = fn + 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(source), - parameters=build_function_json_schema(fn), # fn not source: sig already reflects bound args + description=extract_function_description(fn), + parameters=build_function_json_schema(fn), fn=fn, ) From 42a67c753cffb969b8e1fc9e6b3bc870f1f4604a Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sun, 31 May 2026 07:35:52 +0100 Subject: [PATCH 15/28] Fixump tests --- .../common/ai/utils/function_schema.py | 32 ++++++++++++------- .../common/ai/utils/test_function_schema.py | 21 +++++++++--- 2 files changed, 37 insertions(+), 16 deletions(-) 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 index a4464826648e6..d027f01ce1dfd 100644 --- 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 @@ -53,26 +53,34 @@ ) +def _first_docstring_paragraph(obj: Any) -> str: + 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 the first paragraph of *fn*'s docstring, stopping before Args/Returns sections.""" # 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__; use the class name. + # 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 type(fn).__name__ + return ( + _first_docstring_paragraph(type(fn).__call__) + or _first_docstring_paragraph(fn) + or type(fn).__name__ + ) - doc = inspect.getdoc(fn) - name: str = fn.__name__ # type: ignore[assignment] - if not doc: - return name - 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() or name + return _first_docstring_paragraph(fn) or fn.__name__ # type: ignore[return-value] def build_function_json_schema(fn: Callable[..., Any]) -> dict[str, Any]: 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 index 965710d3853e5..6927edf053892 100644 --- 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 @@ -87,11 +87,24 @@ def test_lambda_falls_back_to_lambda_name(self): f = lambda x: x assert extract_function_description(f) == "" - def test_callable_object_falls_back_to_class_name(self): + def test_callable_object_uses_call_docstring(self): obj = _CallableObj() - # callable objects have no __name__; falls back to type name - result = extract_function_description(obj) - assert result == "_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: ... + + 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: ... + + assert extract_function_description(_NoDocs()) == "_NoDocs" @pytest.mark.parametrize( "header", From c367f9bff32a5659005db77c2792535f1c71ded1 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sun, 31 May 2026 12:14:51 +0100 Subject: [PATCH 16/28] Fix mypy --- .../common/ai/utils/function_schema.py | 2 +- .../common/ai/utils/test_function_schema.py | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) 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 index d027f01ce1dfd..a9eb78a17a7f1 100644 --- 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 @@ -80,7 +80,7 @@ def extract_function_description(fn: Callable[..., Any]) -> str: or type(fn).__name__ ) - return _first_docstring_paragraph(fn) or fn.__name__ # type: ignore[return-value] + return _first_docstring_paragraph(fn) or fn.__name__ def build_function_json_schema(fn: Callable[..., Any]) -> dict[str, Any]: 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 index 6927edf053892..9e5ec71530c91 100644 --- 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 @@ -34,7 +34,7 @@ # --------------------------------------------------------------------------- -def _plain(x: int, y: str = "hi") -> str: +def _plain(x: int, y: str = "hi") -> str: # type: ignore[empty-body] """Do something useful. Args: @@ -46,7 +46,7 @@ def _plain(x: int, y: str = "hi") -> str: """ -def _no_doc(x: int) -> str: +def _no_doc(x: int) -> str: # type: ignore[empty-body] pass @@ -54,14 +54,14 @@ def _no_params() -> None: """No parameters at all.""" -def _annotated(q: Annotated[str, "The search query"], limit: Annotated[int, "Max results"] = 10) -> list: +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: + def __call__(self, value: str) -> str: # type: ignore[empty-body] """Process value.""" @@ -96,13 +96,13 @@ def test_callable_object_falls_back_to_class_docstring(self): class _NoCallDoc: """Describes the class.""" - def __call__(self, x: int) -> int: ... + 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: ... + def __call__(self, x: int) -> int: ... # type: ignore[empty-body] assert extract_function_description(_NoDocs()) == "_NoDocs" @@ -270,7 +270,7 @@ def test_additional_properties_stripped(self): assert "additionalProperties" not in schema def test_partial_positional_bind_removes_param(self): - def add(a: int, b: int) -> int: ... + def add(a: int, b: int) -> int: ... # type: ignore[empty-body] p = functools.partial(add, 1) schema = build_function_json_schema(p) @@ -278,7 +278,7 @@ def add(a: int, b: int) -> int: ... assert "b" in schema["properties"] def test_partial_keyword_bind_keeps_param_as_optional(self): - def add(a: int, b: int) -> int: ... + def add(a: int, b: int) -> int: ... # type: ignore[empty-body] p = functools.partial(add, a=1) schema = build_function_json_schema(p) @@ -287,7 +287,7 @@ def add(a: int, b: int) -> int: ... assert "b" in schema["required"] def test_nested_partial_unwraps_hint_source(self): - def fn(x: int, y: str) -> str: ... + 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) From 60e9aa3210432aeb57ec3937e71e5f34fedec392 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sun, 31 May 2026 16:15:33 +0100 Subject: [PATCH 17/28] Update functions schema --- .../common/ai/utils/function_schema.py | 27 +++++++++++++------ .../common/ai/utils/test_function_schema.py | 19 ++++++++++--- 2 files changed, 35 insertions(+), 11 deletions(-) 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 index a9eb78a17a7f1..d7836dc786a66 100644 --- 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 @@ -53,7 +53,8 @@ ) -def _first_docstring_paragraph(obj: Any) -> str: +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 "" @@ -66,7 +67,7 @@ def _first_docstring_paragraph(obj: Any) -> str: def extract_function_description(fn: Callable[..., Any]) -> str: - """Return the first paragraph of *fn*'s docstring, stopping before Args/Returns sections.""" + """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) @@ -75,12 +76,12 @@ def extract_function_description(fn: Callable[..., Any]) -> str: # Prefer __call__ docstring (what calling does), then class docstring, then class name. if not hasattr(fn, "__name__"): return ( - _first_docstring_paragraph(type(fn).__call__) - or _first_docstring_paragraph(fn) + _extract_docstring_summary(type(fn).__call__) + or _extract_docstring_summary(fn) or type(fn).__name__ ) - return _first_docstring_paragraph(fn) or fn.__name__ + return _extract_docstring_summary(fn) or fn.__name__ def build_function_json_schema(fn: Callable[..., Any]) -> dict[str, Any]: @@ -92,11 +93,10 @@ def build_function_json_schema(fn: Callable[..., Any]) -> dict[str, Any]: 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. """ - if inspect.isbuiltin(fn): - return _EMPTY_OBJECT_SCHEMA - # Partials: sig from partial (bound args already removed), hints from inner fn. hint_source: Callable[..., Any] = fn if isinstance(fn, functools.partial): @@ -118,6 +118,17 @@ def build_function_json_schema(fn: Callable[..., Any]) -> 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 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 index 9e5ec71530c91..3f285c4a49b4b 100644 --- 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 @@ -17,6 +17,7 @@ from __future__ import annotations import functools +import inspect from typing import Annotated, Any import pytest @@ -250,6 +251,12 @@ 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): ... @@ -301,9 +308,15 @@ def fn(x: int | None = None): ... assert "x" in schema["properties"] assert "x" not in schema.get("required", []) - def test_introspection_failure_returns_empty_schema(self): - # Built-in functions have no inspectable signature. - schema = build_function_json_schema(len) + 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): From da76d7ec58d6ac9373e746d951a44c0bc16536db Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Mon, 1 Jun 2026 17:11:51 +0100 Subject: [PATCH 18/28] Fixup tests --- .../unit/common/ai/hooks/test_pydantic_ai.py | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) 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 2f154126744f2..73d22267f9d33 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 @@ -753,8 +753,8 @@ 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()): + agent = hook.create_agent(AgentRunRequest(prompt="test", instructions="hi")) assert agent.instrument is sentinel @@ -762,22 +762,21 @@ def test_instrument_set_when_settings_returned(self, mock_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()): + agent = hook.create_agent(AgentRunRequest(prompt="test", instructions="hi")) mock_settings.assert_called_once() assert 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() From 344ee003c30cedb44ade5ee4c9e58f155e2bded0 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Tue, 2 Jun 2026 22:54:38 +0100 Subject: [PATCH 19/28] Resolve comments --- docs/spelling_wordlist.txt | 1 + .../airflow/providers/common/ai/hooks/base.py | 218 +++++++------- .../providers/common/ai/hooks/pydantic_ai.py | 92 ++++-- .../providers/common/ai/mixins/durable.py | 97 ++++++ .../providers/common/ai/utils/callables.py | 30 ++ .../unit/common/ai/decorators/test_agent.py | 6 +- .../tests/unit/common/ai/hooks/test_base.py | 284 +++++++++++++----- .../unit/common/ai/hooks/test_durable.py | 121 ++++++++ .../unit/common/ai/hooks/test_pydantic_ai.py | 249 ++++++++++++--- .../unit/common/ai/operators/test_agent.py | 18 +- .../unit/common/ai/operators/test_llm.py | 11 +- .../unit/common/ai/test_observability.py | 11 +- .../unit/common/ai/utils/test_callables.py | 67 +++++ uv.lock | 2 +- 14 files changed, 926 insertions(+), 281 deletions(-) create mode 100644 providers/common/ai/src/airflow/providers/common/ai/mixins/durable.py create mode 100644 providers/common/ai/src/airflow/providers/common/ai/utils/callables.py create mode 100644 providers/common/ai/tests/unit/common/ai/hooks/test_durable.py create mode 100644 providers/common/ai/tests/unit/common/ai/utils/test_callables.py diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index b6c1c022d01f7..8e21d1e37c3cf 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -25,6 +25,7 @@ adsinsights afterall agentcore agentic +AgentT ai aio aiobotocore 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 index d8d8589649c59..0433c9ea2bb1e 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py @@ -24,24 +24,40 @@ from abc import ABCMeta, abstractmethod from collections.abc import Callable, Sequence from dataclasses import dataclass, field -from typing import Any, ClassVar +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 -# Attribute name for durable storage/counter bound to a framework agent instance. -_AIRFLOW_DURABLE_ATTR = "_airflow_durable_state" +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 = 0 - tool_calls: int = 0 - input_tokens: int = 0 - output_tokens: int = 0 - total_tokens: int = 0 + 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 @@ -116,11 +132,13 @@ class AgentRunRequest: 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. + 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 output_type: Expected structured output type or backend-specific JSON schema + mapping (default: ``str``). :param instructions: System-level instructions for the agent. :param toolsets: List of :class:`BaseToolset` instances the agent may call. :param usage_limits: Backend-specific usage limits; ignored if the hook does not support them. @@ -134,7 +152,7 @@ class AgentRunRequest: """ prompt: str | Sequence[Any] - output_type: type[Any] = str + output_type: type[Any] | dict[str, Any] | None = str instructions: str = "" toolsets: list[Any] | None = None usage_limits: Any = None @@ -158,7 +176,7 @@ def as_tools(self) -> list[ToolSpec]: """Return the list of tools this toolset exposes.""" -class BaseAIHook(BaseHook, metaclass=ABCMeta): +class BaseAIHook(BaseHook, Generic[AgentT], metaclass=ABCMeta): """ Abstract hook for multi-turn LLM agents. @@ -168,18 +186,17 @@ class BaseAIHook(BaseHook, metaclass=ABCMeta): :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:`create_agent`, :meth:`run_agent`, and + Subclasses implement :meth:`get_model`, :meth:`_build_agent`, :meth:`run_agent`, and :meth:`_tool_spec_to_native`. - Shared helpers :meth:`_init_durable`, :meth:`_resolve_tools`, :meth:`_logged_callable`, and - :meth:`_cached_callable` are provided for all hooks. + 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" - supports_toolsets: ClassVar[bool] = False - supports_durable: ClassVar[bool] = False - supports_usage_limits: ClassVar[bool] = False + capabilities: ClassVar[frozenset[Capability]] = frozenset() def __init__( self, @@ -192,7 +209,7 @@ def __init__( self.model_id = model_id @classmethod - def get_agent_hook(cls, conn_id: str, *, hook_params: dict[str, Any] | None = None) -> BaseAIHook: + 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. @@ -212,11 +229,15 @@ def get_model(self) -> Any: """Return the backend model/client used to construct agents.""" def get_conn(self) -> Any: - """Return the backend model/client. Delegates to :meth:`get_model`.""" + """ + 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() - @abstractmethod - def create_agent(self, request: AgentRunRequest) -> Any: + def create_agent(self, request: AgentRunRequest) -> AgentT: """ Build (but do not run) the agent described by *request*. @@ -224,29 +245,27 @@ def create_agent(self, request: AgentRunRequest) -> Any: :meth:`_resolve_tools` and constructing the framework-native agent object with the model, tools, instructions, and output type from *request*. - When :attr:`AgentRunRequest.durable_context` is set, implementations - should call :meth:`_init_durable` and bind the returned storage/counter - to the agent via :meth:`_bind_agent_durable` so that :meth:`run_agent` - can retrieve and clean them up. - - Implementations must call :meth:`validate_run_request` at the start of - this method before any agent construction or durable initialisation. - :param request: All parameters needed to configure the agent. - :returns: Framework-native agent object, ready to be passed to :meth:`run_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: Any, request: AgentRunRequest) -> AgentRunResult: + def run_agent(self, agent: AgentT, request: AgentRunRequest) -> AgentRunResult: """ Execute *agent* for *request* and return a normalized :class:`AgentRunResult`. - Implementations with durable execution should pop durable state via - :meth:`_pop_agent_durable`, 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). + 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 produced by :meth:`create_agent`. + :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.). """ @@ -267,64 +286,30 @@ def validate_run_request(self, request: AgentRunRequest) -> None: """ Raise if *request* uses features this hook implementation does not support. - Hook implementations call this at the start of :meth:`create_agent`. + :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 not self.supports_toolsets: + if request.toolsets and Capability.TOOLSETS not in self.capabilities: raise ValueError( - f"toolsets are not supported for connection {conn_id!r} (conn_type resolves to {hook_name}). " + f"toolsets not supported for connection {conn_id!r} (conn_type resolves to {hook_name})." ) - if request.usage_limits is not None and not self.supports_usage_limits: + if request.usage_limits is not None and Capability.USAGE_LIMITS not in self.capabilities: raise ValueError( - f"usage_limits are not supported for connection {conn_id!r} " - f"(conn_type resolves to {hook_name})." + f"usage_limits not supported for connection {conn_id!r} (conn_type resolves to {hook_name})." ) - if request.durable_context is not None and not self.supports_durable: + if request.durable_context is not None and Capability.DURABLE not in self.capabilities: raise ValueError( - f"durable execution requires a hook that supports durable caching; " - f"got {hook_name} for connection {conn_id!r}." + f"durable execution not supported for connection {conn_id!r} (conn_type resolves to {hook_name})." ) - def _init_durable(self, ctx: DurableContext) -> tuple[Any, Any]: - """ - Create and return a ``DurableStorage`` / ``DurableStepCounter`` pair for *ctx*. - - Hooks call this inside :meth:`create_agent` when - :attr:`AgentRunRequest.durable_context` is set. - """ - from airflow.providers.common.ai.durable.step_counter import DurableStepCounter - from airflow.providers.common.ai.durable.storage import DurableStorage - - storage = DurableStorage( - dag_id=ctx.dag_id, - task_id=ctx.task_id, - run_id=ctx.run_id, - map_index=ctx.map_index, - ) - counter = DurableStepCounter() - return storage, counter - - @staticmethod - def _bind_agent_durable(agent: Any, storage: Any, counter: Any) -> None: - """Associate *storage* and *counter* with *agent* until :meth:`run_agent` completes.""" - setattr(agent, _AIRFLOW_DURABLE_ATTR, (storage, counter)) - - @staticmethod - def _pop_agent_durable(agent: Any) -> tuple[Any, Any] | None: - """Remove and return durable state bound to *agent*, if any.""" - state = getattr(agent, _AIRFLOW_DURABLE_ATTR, None) - if state is None: - return None - delattr(agent, _AIRFLOW_DURABLE_ATTR) - return state - def _resolve_tools( self, toolsets: list[Any], enable_logging: bool, - storage: Any, - counter: Any, + 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. @@ -339,14 +324,18 @@ def _resolve_tools( * 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: - *fn* → optional cache wrap → optional log wrap → :meth:`_tool_spec_to_native`. + 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 storage: ``DurableStorage`` instance, or ``None`` to skip caching. - :param counter: ``DurableStepCounter`` instance, or ``None`` to skip caching. + :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: @@ -359,16 +348,16 @@ def _resolve_tools( continue for spec in specs: fn = spec.fn - if storage is not None and counter is not None: - fn = self._cached_callable(fn, storage, counter) 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, + sequential=spec.sequential or force_sequential, ) native.append(self._tool_spec_to_native(adapted)) return native @@ -383,8 +372,30 @@ def _logged_callable( """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 wrapper(*args, **kwargs): + 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)) @@ -393,35 +404,12 @@ def wrapper(*args, **kwargs): result = fn(*args, **kwargs) elapsed = time.monotonic() - start logger.info("Tool %s returned in %.2fs", _tool_name, elapsed) - logger.info("::endgroup::") return result except Exception: elapsed = time.monotonic() - start logger.exception("Tool %s failed after %.2fs", _tool_name, elapsed) - logger.info("::endgroup::") raise + finally: + logger.info("::endgroup::") - return wrapper - - @staticmethod - def _cached_callable( - fn: Callable[..., Any], - storage: Any, - counter: Any, - ) -> Callable[..., Any]: - """Wrap *fn* to cache its result in *storage* using a monotonic step counter.""" - - @functools.wraps(fn) - def 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 wrapper + 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 809c2038fc22e..4283e12e6e614 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,6 +16,8 @@ # under the License. from __future__ import annotations +import functools +from dataclasses import dataclass from typing import TYPE_CHECKING, Any from pydantic_ai import Agent @@ -32,9 +34,11 @@ 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.ai.toolsets.logging import LoggingToolset @@ -42,7 +46,15 @@ from pydantic_ai.models import KnownModelName, Model -class PydanticAIHook(BaseAIHook): +@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. @@ -68,9 +80,7 @@ class PydanticAIHook(BaseAIHook): conn_type = "pydanticai" hook_name = "Pydantic AI" - supports_toolsets = True - supports_durable = True - supports_usage_limits = True + capabilities = frozenset({Capability.TOOLSETS, Capability.USAGE_LIMITS, Capability.DURABLE}) def __init__( self, @@ -200,13 +210,13 @@ def _tool_spec_to_native(self, spec: ToolSpec) -> Any: json_schema=spec.parameters, ) - def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: + def _build_agent(self, request: AgentRunRequest) -> PydanticAgentHandle: """ - Build a pydantic-ai ``Agent`` from *request*. + Build a pydantic-ai ``Agent`` handle from *request*. When :attr:`~AgentRunRequest.durable_context` is set, initialises durable - storage and step counter and binds them to the returned agent for use by - :meth:`run_agent`. When ``[common.ai] otel_export_enabled`` is set and the + 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`. @@ -218,11 +228,9 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: framework-neutral callables / ``BaseToolset`` specs and pydantic-ai ``AbstractToolset`` instances, but not to native ``Tool`` instances. """ - self.validate_run_request(request) - - storage = counter = None + durable_state = None if request.durable_context is not None: - storage, counter = self._init_durable(request.durable_context) + durable_state = self._init_durable(request.durable_context) extra_kwargs = dict(request.agent_params or {}) if request.toolsets: @@ -242,17 +250,25 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: resolved: list[Any] = [] for item in pipeline_items: if isinstance(item, Tool): - # Native pydantic-ai Tool objects are preserved as-is. They keep their - # original schema/configuration, but do not receive Airflow's callable - # logging or durable tool-result caching wrappers. + # 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, - storage, - counter, + cache_wrapper, + force_sequential=durable_state is not None, ) ) extra_kwargs["tools"] = resolved @@ -264,17 +280,17 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: if abstract_items: processed: list[Any] = list(abstract_items) - if storage is not None and counter is not None: + 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=storage, - counter=counter, + storage=durable_state.storage, + counter=durable_state.counter, ) for ts in processed ] - if request.enable_tool_logging: - processed = [LoggingToolset(wrapped=ts, logger=self.log) for ts in processed] extra_kwargs["toolsets"] = processed self.log.info( "Agent abstract toolsets configured: count=%d types=%s", @@ -288,8 +304,6 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: instructions=request.instructions, **extra_kwargs, ) - if storage is not None and counter is not None: - self._bind_agent_durable(agent, storage, counter) if "instrument" not in extra_kwargs: # Set the public ``agent.instrument`` property rather than the # ``Agent(instrument=...)`` constructor kwarg, which is deprecated in @@ -298,33 +312,45 @@ def create_agent(self, request: AgentRunRequest) -> Agent[None, Any]: 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: Agent[None, Any], request: AgentRunRequest) -> AgentRunResult: + 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 - durable = self._pop_agent_durable(agent) - storage, counter = durable if durable else (None, None) + 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 agent.model is None: + if native_agent.model is None: raise ValueError("Agent model must be set when durable=True") - model = agent.model + 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 agent.override(model=caching_model): - result = agent.run_sync(request.prompt, **run_kwargs) + with native_agent.override(model=caching_model): + result = native_agent.run_sync(request.prompt, **run_kwargs) else: - result = agent.run_sync(request.prompt, **run_kwargs) + result = native_agent.run_sync(request.prompt, **run_kwargs) usage = result.usage tool_names: list[str] = [] for message in result.all_messages(): 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/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/tests/unit/common/ai/decorators/test_agent.py b/providers/common/ai/tests/unit/common/ai/decorators/test_agent.py index 9a1a14169a8c0..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.hooks.base import AgentRunResult, AgentUsage, BaseAIHook +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 @@ -46,9 +46,7 @@ def _make_run_result(output): def _make_mock_hook(run_result): mock_hook = MagicMock(spec=BaseAIHook) - mock_hook.supports_toolsets = True - mock_hook.supports_durable = False - mock_hook.supports_usage_limits = True + 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 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 index 98787e0c6a36d..94831ff89809e 100644 --- a/providers/common/ai/tests/unit/common/ai/hooks/test_base.py +++ b/providers/common/ai/tests/unit/common/ai/hooks/test_base.py @@ -16,13 +16,12 @@ # under the License. from __future__ import annotations +import asyncio import functools import inspect from unittest.mock import MagicMock, patch import pytest -from pydantic_ai import Agent -from pydantic_ai.models.test import TestModel from airflow.providers.common.ai.hooks.base import ( AgentRunRequest, @@ -30,6 +29,7 @@ AgentUsage, BaseAIHook, BaseToolset, + Capability, DurableContext, DurableStats, ToolSpec, @@ -65,7 +65,7 @@ class ConcreteHook(BaseAIHook): def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -79,32 +79,17 @@ def _tool_spec_to_native(self, spec): assert hook.model_id == "openai:gpt-5" -class TestBaseAIHookAgentDurable: - def test_bind_pop_round_trip(self): - agent = Agent(TestModel()) - storage = MagicMock() - counter = MagicMock() - - BaseAIHook._bind_agent_durable(agent, storage, counter) - assert agent._airflow_durable_state == (storage, counter) - assert BaseAIHook._pop_agent_durable(agent) == (storage, counter) - assert BaseAIHook._pop_agent_durable(agent) is None - assert not hasattr(agent, "_airflow_durable_state") - - class TestValidateRunRequest: def test_rejects_toolsets_when_unsupported(self): class ConcreteHook(BaseAIHook): conn_type = "test" hook_name = "Test" - supports_toolsets = False - supports_usage_limits = True - supports_durable = True + capabilities = frozenset({Capability.USAGE_LIMITS, Capability.DURABLE}) def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -115,21 +100,19 @@ def _tool_spec_to_native(self, spec): hook = ConcreteHook(llm_conn_id="test_conn") request = AgentRunRequest(prompt="hi", toolsets=[MagicMock()]) - with pytest.raises(ValueError, match="toolsets are not supported"): + 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" - supports_toolsets = True - supports_usage_limits = False - supports_durable = True + capabilities = frozenset({Capability.TOOLSETS, Capability.DURABLE}) def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -140,11 +123,77 @@ def _tool_spec_to_native(self, spec): hook = ConcreteHook(llm_conn_id="test_conn") request = AgentRunRequest(prompt="hi", usage_limits=MagicMock()) - with pytest.raises(ValueError, match="usage_limits are not supported"): + 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( @@ -210,7 +259,7 @@ class ConcreteHook(BaseAIHook): def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -228,7 +277,7 @@ 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, storage=None, counter=None) + result = hook._resolve_tools([MyToolset()], enable_logging=False) assert len(result) == 1 assert result[0]["name"] == "my_tool" @@ -248,7 +297,7 @@ def log(self): def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -269,7 +318,7 @@ class SimpleToolset(BaseToolset): def as_tools(self): return [ToolSpec(name="original", description="", parameters={}, fn=original)] - [wrapped_fn] = hook._resolve_tools([SimpleToolset()], enable_logging=True, storage=None, counter=None) + [wrapped_fn] = hook._resolve_tools([SimpleToolset()], enable_logging=True) wrapped_fn() assert calls == ["original"] @@ -285,7 +334,7 @@ class ConcreteHook(BaseAIHook): def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -300,7 +349,7 @@ def roll_dice() -> str: """Roll a six-sided die and return the result.""" return "4" - result = hook._resolve_tools([roll_dice], enable_logging=False, storage=None, counter=None) + result = hook._resolve_tools([roll_dice], enable_logging=False) assert len(result) == 1 assert result[0]["name"] == "roll_dice" @@ -317,7 +366,7 @@ class ConcreteHook(BaseAIHook): def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -335,7 +384,7 @@ def search(self, query: str) -> str: helper = MyHelper() bound_method = helper.search - result = hook._resolve_tools([bound_method], enable_logging=False, storage=None, counter=None) + result = hook._resolve_tools([bound_method], enable_logging=False) assert len(result) == 1 assert result[0]["name"] == "search" @@ -352,7 +401,7 @@ class ConcreteHook(BaseAIHook): def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -368,7 +417,7 @@ def query_db(db: str, query: str) -> str: return f"{db}: {query}" partial_tool = functools.partial(query_db, db="prod") - result = hook._resolve_tools([partial_tool], enable_logging=False, storage=None, counter=None) + result = hook._resolve_tools([partial_tool], enable_logging=False) assert len(result) == 1 assert result[0]["name"] == "query_db" @@ -385,7 +434,7 @@ class ConcreteHook(BaseAIHook): def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -401,7 +450,7 @@ def __call__(self, query: str) -> str: return query searcher = Searcher() - result = hook._resolve_tools([searcher], enable_logging=False, storage=None, counter=None) + result = hook._resolve_tools([searcher], enable_logging=False) assert len(result) == 1 assert result[0]["name"] == "Searcher" @@ -417,7 +466,7 @@ class ConcreteHook(BaseAIHook): def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -429,7 +478,7 @@ def _tool_spec_to_native(self, spec): hook = ConcreteHook.__new__(ConcreteHook) native_tool_obj = object() # not a function, not a BaseToolset - result = hook._resolve_tools([native_tool_obj], enable_logging=True, storage=None, counter=None) + result = hook._resolve_tools([native_tool_obj], enable_logging=True) assert result == [native_tool_obj] @@ -443,7 +492,7 @@ class ConcreteHook(BaseAIHook): def get_model(self): return None - def create_agent(self, request): + def _build_agent(self, request): return None def run_agent(self, agent, request): @@ -460,12 +509,102 @@ 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, storage=None, counter=None - ) + 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): @@ -556,44 +695,39 @@ def __call__(self, customer_id: str) -> dict[str, str]: signature = inspect.signature(wrapped) assert tuple(signature.parameters) == ("customer_id",) + def test_logged_callable_preserves_async_function_behavior(self): + logger = MagicMock() -class TestBaseAIHookCachedCallable: - def test_cached_callable_saves_and_returns(self): - storage = MagicMock() - counter = MagicMock() - counter.next_step.return_value = 1 - storage.load_tool_result.return_value = (False, None) + async def fn(value): + return value * 2 - calls = [] + wrapped = BaseAIHook._logged_callable(fn, logger) - def fn(): - calls.append(1) - return "computed" + assert inspect.iscoroutinefunction(wrapped) + assert asyncio.run(wrapped(3)) == 6 + logger.info.assert_called() - wrapped = BaseAIHook._cached_callable(fn, storage, counter) - result = wrapped() + def test_logged_callable_preserves_async_partial_behavior(self): + logger = MagicMock() - assert result == "computed" - assert calls == [1] - storage.save_tool_result.assert_called_once_with("tool_step_1", "computed") + async def fn(prefix, value): + return f"{prefix}:{value}" - def test_cached_callable_replays_on_hit(self): - storage = MagicMock() - counter = MagicMock() - counter.replayed_tool = 0 - counter.next_step.return_value = 1 - storage.load_tool_result.return_value = (True, "cached_value") + wrapped = BaseAIHook._logged_callable(functools.partial(fn, "prod"), logger) - calls = [] + assert inspect.iscoroutinefunction(wrapped) + assert asyncio.run(wrapped("cpu")) == "prod:cpu" + logger.info.assert_called() - def fn(): - calls.append(1) - return "computed" + def test_logged_callable_preserves_async_callable_object_behavior(self): + logger = MagicMock() - wrapped = BaseAIHook._cached_callable(fn, storage, counter) - result = wrapped() + class Lookup: + async def __call__(self, value): + return value.upper() - assert result == "cached_value" - assert calls == [] - assert counter.replayed_tool == 1 - storage.save_tool_result.assert_not_called() + 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_durable.py b/providers/common/ai/tests/unit/common/ai/hooks/test_durable.py new file mode 100644 index 0000000000000..7213933d7aae2 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/hooks/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/hooks/test_pydantic_ai.py b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py index 73d22267f9d33..d03c649057981 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 @@ -30,19 +30,31 @@ from pydantic_ai.usage import RunUsage, UsageLimits from airflow.models.connection import Connection -from airflow.providers.common.ai.hooks.base import AgentRunRequest, AgentRunResult, BaseAIHook, ToolSpec +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 def _test_agent() -> Agent[None, str]: return Agent(TestModel()) +def _test_handle() -> PydanticAgentHandle: + return PydanticAgentHandle(agent=_test_agent()) + + def _pydantic_run_result( output: str, *, @@ -92,9 +104,9 @@ def test_is_base_hook(self): assert issubclass(PydanticAIHook, BaseAIHook) def test_capability_flags(self): - assert PydanticAIHook.supports_toolsets is True - assert PydanticAIHook.supports_durable is True - assert PydanticAIHook.supports_usage_limits is True + assert Capability.TOOLSETS in PydanticAIHook.capabilities + assert Capability.DURABLE in PydanticAIHook.capabilities + assert Capability.USAGE_LIMITS in PydanticAIHook.capabilities class TestPydanticAIHookInit: @@ -330,6 +342,74 @@ def fetch_metric(environment: str, metric_name: str) -> float: 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 + assert calls == ["bound", "prod", "callable"] + @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): @@ -425,17 +505,18 @@ def test_create_agent_inits_durable_when_context_set(self, mock_infer_model): 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=(mock_storage, mock_counter)), + patch.object(hook, "_init_durable", return_value=durable_state), ): - agent = hook.create_agent(request) + handle = hook.create_agent(request) - assert BaseAIHook._pop_agent_durable(agent) == (mock_storage, mock_counter) + assert handle.durable_state is durable_state @patch("airflow.providers.common.ai.hooks.pydantic_ai.infer_model", autospec=True) - def test_create_agent_does_not_bind_durable_when_no_context(self, mock_infer_model): + 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 @@ -444,9 +525,10 @@ def test_create_agent_does_not_bind_durable_when_no_context(self, mock_infer_mod request = AgentRunRequest(prompt="hi") conn = Connection(conn_id="test_conn", conn_type="pydanticai") with patch.object(hook, "get_connection", return_value=conn): - agent = hook.create_agent(request) + handle = hook.create_agent(request) - assert BaseAIHook._pop_agent_durable(agent) is None + 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) @@ -505,6 +587,46 @@ def as_tools(self): 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): @@ -557,7 +679,7 @@ def test_create_agent_wraps_abstract_toolset_with_logging(self, mock_agent_cls, def test_create_agent_wraps_abstract_toolset_with_caching_when_durable( self, mock_agent_cls, mock_infer_model ): - """AbstractToolset items are wrapped with CachingToolset (inner) and LoggingToolset (outer) for durable runs.""" + """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 @@ -570,6 +692,7 @@ def test_create_agent_wraps_abstract_toolset_with_caching_when_durable( 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") @@ -579,7 +702,7 @@ def test_create_agent_wraps_abstract_toolset_with_caching_when_durable( ) with ( patch.object(hook, "get_connection", return_value=conn), - patch.object(hook, "_init_durable", return_value=(mock_storage, mock_counter)), + patch.object(hook, "_init_durable", return_value=durable_state), ): hook.create_agent(request) @@ -587,12 +710,12 @@ def test_create_agent_wraps_abstract_toolset_with_caching_when_durable( toolsets = call_kwargs["toolsets"] assert len(toolsets) == 1 outer = toolsets[0] - assert isinstance(outer, LoggingToolset) - assert isinstance(outer.wrapped, CachingToolset) + 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_binds_durable_per_agent_not_on_hook(self, mock_infer_model): + 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 @@ -604,24 +727,25 @@ def test_create_agent_binds_durable_per_agent_not_on_hook(self, mock_infer_model 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=[(storage_a, counter_a), (storage_b, counter_b)] - ): - agent_a = hook.create_agent(AgentRunRequest(prompt="a", durable_context=ctx_a)) - agent_b = hook.create_agent(AgentRunRequest(prompt="b", durable_context=ctx_b)) + 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 agent_a is not agent_b - assert BaseAIHook._pop_agent_durable(agent_a) == (storage_a, counter_a) - assert BaseAIHook._pop_agent_durable(agent_b) == (storage_b, counter_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", @@ -631,7 +755,7 @@ def test_run_agent_returns_agent_run_result(self): request = AgentRunRequest(prompt="hello") with patch.object(agent, "run_sync", return_value=mock_result) as mock_run_sync: - run_result = hook.run_agent(agent, request) + run_result = hook.run_agent(handle, request) assert isinstance(run_result, AgentRunResult) assert run_result.output == "done" @@ -641,20 +765,21 @@ def test_run_agent_returns_agent_run_result(self): def test_create_agent_rejects_unsupported_usage_limits(self): hook = PydanticAIHook() - hook.supports_usage_limits = False - with pytest.raises(ValueError, match="usage_limits are not supported"): + 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(agent, request) + hook.run_agent(handle, request) mock_run_sync.assert_called_once_with("more", message_history=history, usage_limits=limits) @@ -670,6 +795,8 @@ def test_run_agent_durable_applies_caching_model( 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") @@ -682,16 +809,21 @@ def test_run_agent_durable_applies_caching_model( mock_counter.replayed_tool = 0 mock_counter.cached_model = 0 mock_counter.cached_tool = 0 - BaseAIHook._bind_agent_durable(agent, mock_storage, mock_counter) + handle = PydanticAgentHandle( + agent=agent, + durable_state=DurableState(storage=mock_storage, counter=mock_counter), + ) - request = AgentRunRequest(prompt="hi") - run_result = hook.run_agent(agent, request) + 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 - assert BaseAIHook._pop_agent_durable(agent) is None mock_storage.cleanup.assert_called_once() @patch.object(Agent, "override") @@ -705,19 +837,58 @@ def test_run_agent_preserves_durable_cache_on_exception( 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() - BaseAIHook._bind_agent_durable(agent, mock_storage, mock_counter) + handle = PydanticAgentHandle( + agent=agent, + durable_state=DurableState(storage=mock_storage, counter=mock_counter), + ) with pytest.raises(RuntimeError, match="boom"): - hook.run_agent(agent, AgentRunRequest(prompt="hi")) + 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() - assert BaseAIHook._pop_agent_durable(agent) is None + + 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() @@ -727,7 +898,7 @@ def fn(customer_id: int) -> str: return "ok" with patch("airflow.providers.common.ai.hooks.pydantic_ai.Tool") as mock_tool_cls: - hook._resolve_tools(toolsets=[fn], enable_logging=False, storage=None, counter=None) + hook._resolve_tools(toolsets=[fn], enable_logging=False) mock_tool_cls.from_schema.assert_called_once_with( fn, name="fn", @@ -754,19 +925,19 @@ def test_instrument_set_when_settings_returned(self, mock_settings): mock_settings.return_value = sentinel hook = self._hook() with patch.object(hook, "get_model", return_value=TestModel()): - agent = hook.create_agent(AgentRunRequest(prompt="test", instructions="hi")) + 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_model", return_value=TestModel()): - agent = hook.create_agent(AgentRunRequest(prompt="test", instructions="hi")) + 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") 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 fb3756a2d6304..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,7 +23,13 @@ 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.hooks.base import ( + AgentRunRequest, + AgentRunResult, + AgentUsage, + BaseAIHook, + Capability, +) from airflow.providers.common.ai.operators.agent import AgentOperator, HITLReviewLink from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS @@ -57,9 +63,7 @@ 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.supports_toolsets = True - mock_hook.supports_durable = True - mock_hook.supports_usage_limits = True + 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) @@ -71,9 +75,7 @@ class TestAgentOperatorHookCapabilities: 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.supports_toolsets = False - mock_hook.supports_durable = False - mock_hook.supports_usage_limits = False + mock_hook.capabilities = frozenset() # no capabilities — all features rejected def create_agent(request): BaseAIHook.validate_run_request(mock_hook, request) @@ -88,7 +90,7 @@ def create_agent(request): llm_conn_id="strands_conn", toolsets=[MagicMock()], ) - with pytest.raises(ValueError, match="toolsets are not supported"): + with pytest.raises(ValueError, match="toolsets not supported"): op.execute(context=MagicMock()) 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 565865775fb44..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,7 +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.hooks.base import ( + AgentRunRequest, + AgentRunResult, + AgentUsage, + BaseAIHook, +) from airflow.providers.common.ai.mixins.approval import ( LLMApprovalMixin, ) @@ -93,7 +98,7 @@ def test_execute_returns_string_output(self, mock_hook_cls): 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.supports_usage_limits = False + 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( @@ -102,7 +107,7 @@ def test_execute_rejects_usage_limits_when_hook_unsupported(self, mock_hook_cls) llm_conn_id="my_llm", usage_limits=UsageLimits(request_limit=1), ) - with pytest.raises(ValueError, match="usage_limits are not supported"): + with pytest.raises(ValueError, match="usage_limits not supported"): op.execute(context=MagicMock()) @patch("airflow.providers.common.ai.operators.llm.BaseAIHook", autospec=True) 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/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/uv.lock b/uv.lock index 796981385893b..cb3451910fb8e 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'" }, From 68c2ed9b7ee449d3447aa7987b6a611d2d427cdf Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Tue, 2 Jun 2026 23:19:57 +0100 Subject: [PATCH 20/28] Resolve conflicts --- .../providers/common/ai/operators/agent.py | 7 ++----- .../ai/operators/test_llm_file_analysis.py | 20 ++++++++++--------- 2 files changed, 13 insertions(+), 14 deletions(-) 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 f640338b264d9..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 @@ -28,11 +28,8 @@ 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.output_type import ( - iter_base_model_classes, - rehydrate_pydantic_output, -) +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, 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 19fb67bf853a5..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 @@ -24,6 +24,7 @@ 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 @@ -84,7 +85,7 @@ 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() + 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" @@ -126,7 +127,7 @@ def test_execute_structured_output_returns_pydantic_instance(self, mock_build_re resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock() + 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"]) @@ -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", @@ -215,7 +217,7 @@ def test_execute_with_approval_defers( resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock() + 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" @@ -253,7 +255,7 @@ def test_execute_with_approval_defers_structured_output_as_json( resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock() + 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"]) @@ -329,7 +331,7 @@ def test_execute_with_approval_timeout( resolved_paths=["/tmp/app.log"], total_size_bytes=10, ) - mock_agent = MagicMock() + 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") From d3546c0179b69d6b437b32cb57d70b6befdb497b Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Wed, 3 Jun 2026 00:09:21 +0100 Subject: [PATCH 21/28] move tests --- .../ai/tests/unit/common/ai/{hooks => mixins}/test_durable.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename providers/common/ai/tests/unit/common/ai/{hooks => mixins}/test_durable.py (100%) diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_durable.py b/providers/common/ai/tests/unit/common/ai/mixins/test_durable.py similarity index 100% rename from providers/common/ai/tests/unit/common/ai/hooks/test_durable.py rename to providers/common/ai/tests/unit/common/ai/mixins/test_durable.py From 4ee61618776510daa7a1aaa146b7be7e5e369afc Mon Sep 17 00:00:00 2001 From: GPK Date: Wed, 3 Jun 2026 00:16:36 +0100 Subject: [PATCH 22/28] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../common/ai/src/airflow/providers/common/ai/hooks/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 0433c9ea2bb1e..10b520c460859 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py @@ -140,7 +140,7 @@ class AgentRunRequest: :param output_type: Expected structured output type or backend-specific JSON schema mapping (default: ``str``). :param instructions: System-level instructions for the agent. - :param toolsets: List of :class:`BaseToolset` instances the agent may call. + :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 From fb7ff439c5494f9ee3c115b9bd69bb66e9ecee53 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Wed, 3 Jun 2026 17:53:02 +0100 Subject: [PATCH 23/28] Fix tests and mypy --- .../providers/common/ai/hooks/pydantic_ai.py | 8 +++-- .../common/ai/toolsets/langchain_bridge.py | 36 ++++++++++++++++++- 2 files changed, 40 insertions(+), 4 deletions(-) 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 4283e12e6e614..d09cf1495ba4a 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 @@ -298,11 +298,13 @@ def _build_agent(self, request: AgentRunRequest) -> PydanticAgentHandle: [type(toolset).__name__ for toolset in processed], ) + 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(), - output_type=request.output_type, - instructions=request.instructions, - **extra_kwargs, + **agent_kwargs, ) if "instrument" not in extra_kwargs: # Set the public ``agent.instrument`` property rather than the 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..f064bd997f51a 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,16 @@ from pydantic_ai.models.test import TestModel from pydantic_ai.usage import RunUsage +from airflow.providers.common.ai.hooks.base import BaseToolset + 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 +70,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 +123,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 +177,30 @@ 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: + return spec.fn(**kwargs) + except ModelRetry as e: + return str(e) + + async def _async_call(**kwargs: Any) -> Any: + try: + 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, + ) From c9c8a4dfe67e6fa4fd879cbcab6bc09b79d02bb9 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Wed, 3 Jun 2026 18:17:09 +0100 Subject: [PATCH 24/28] Fix toolset assignment --- .../common/ai/toolsets/langchain_bridge.py | 4 + .../ai/toolsets/test_langchain_bridge.py | 90 +++++++++++++++++++ 2 files changed, 94 insertions(+) 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 f064bd997f51a..514044c5c9632 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 @@ -187,12 +187,16 @@ def _build_structured_tool_from_spec( def _sync_call(**kwargs: Any) -> Any: try: + if asyncio.iscoroutinefunction(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 asyncio.iscoroutinefunction(spec.fn): + return await spec.fn(**kwargs) return spec.fn(**kwargs) except ModelRetry as e: return str(e) 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 From d1622ff4a506a389e5a7c82f33b2d964f32a5b6d Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Wed, 3 Jun 2026 20:28:00 +0100 Subject: [PATCH 25/28] Fixup tests --- .../airflow/providers/common/ai/operators/llm_file_analysis.py | 1 + .../common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) 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 8ae086831aaca..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 @@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Any 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 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 d03c649057981..d4210d56b137f 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 @@ -408,7 +408,8 @@ def __call__(self) -> str: assert run_result.usage is not None assert run_result.usage.tool_calls == 3 - assert calls == ["bound", "prod", "callable"] + # 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) From 12717046243f2a46e798540d62d6347e82173582 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Fri, 5 Jun 2026 22:39:08 +0100 Subject: [PATCH 26/28] Update chanelogs and remove dict output type --- providers/common/ai/docs/changelog.rst | 8 +++ providers/common/ai/docs/operators/agent.rst | 7 ++- providers/common/ai/docs/toolsets.rst | 23 +++++++++ .../airflow/providers/common/ai/hooks/base.py | 5 +- .../providers/common/ai/hooks/pydantic_ai.py | 6 +++ .../common/ai/toolsets/langchain_bridge.py | 5 +- .../unit/common/ai/hooks/test_pydantic_ai.py | 50 ++++++++++++++++++- 7 files changed, 96 insertions(+), 8 deletions(-) diff --git a/providers/common/ai/docs/changelog.rst b/providers/common/ai/docs/changelog.rst index 7b21a1b73eaad..5e8190a1995f1 100644 --- a/providers/common/ai/docs/changelog.rst +++ b/providers/common/ai/docs/changelog.rst @@ -53,6 +53,14 @@ and :meth:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook.run_age 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/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index 73d1f154eb4af..52f2e2a3bf3a4 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -312,8 +312,11 @@ Parameters (``SQLToolset``), pydantic-ai ``AbstractToolset`` implementations (``HookToolset``, ``MCPToolset``, ``DataFusionToolset``, ``AgentSkillsToolset`` for :ref:`agent-skills`, third-party toolsets), - plain Python callables, or native pydantic-ai ``Tool`` objects. Mixed lists - are supported. + 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 diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index 8b91db324f4a1..5355bcf498d59 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -53,12 +53,35 @@ toolset to the correct agent parameter automatically. - 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. 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``: + + .. 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 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 index 10b520c460859..f85ce7d4ae509 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py +++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/base.py @@ -137,8 +137,7 @@ class AgentRunRequest: :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 or backend-specific JSON schema - mapping (default: ``str``). + :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. @@ -152,7 +151,7 @@ class AgentRunRequest: """ prompt: str | Sequence[Any] - output_type: type[Any] | dict[str, Any] | None = str + output_type: type[Any] | None = str instructions: str = "" toolsets: list[Any] | None = None usage_limits: Any = None 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 d09cf1495ba4a..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 @@ -298,6 +298,12 @@ def _build_agent(self, request: AgentRunRequest) -> PydanticAgentHandle: [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 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 514044c5c9632..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 @@ -41,6 +41,7 @@ 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 @@ -187,7 +188,7 @@ def _build_structured_tool_from_spec( def _sync_call(**kwargs: Any) -> Any: try: - if asyncio.iscoroutinefunction(spec.fn): + if is_async_callable(spec.fn): return _run_coro_sync(spec.fn(**kwargs)) return spec.fn(**kwargs) except ModelRetry as e: @@ -195,7 +196,7 @@ def _sync_call(**kwargs: Any) -> Any: async def _async_call(**kwargs: Any) -> Any: try: - if asyncio.iscoroutinefunction(spec.fn): + if is_async_callable(spec.fn): return await spec.fn(**kwargs) return spec.fn(**kwargs) except ModelRetry as e: 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 d4210d56b137f..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 @@ -19,10 +19,12 @@ import functools import json import sys +from typing import TYPE_CHECKING from unittest.mock import MagicMock, patch import pytest -from pydantic_ai import Agent +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 @@ -46,6 +48,9 @@ ) 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()) @@ -453,6 +458,25 @@ def test_create_agent_with_agent_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): @@ -650,6 +674,30 @@ def test_create_agent_routes_abstract_toolset_to_toolsets_kwarg(self, mock_agent 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): From 65824cc04cbe56ac0b988db7b23895f499c674c6 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sat, 6 Jun 2026 12:34:02 +0100 Subject: [PATCH 27/28] fix tests --- generated/provider_dependencies.json | 2 +- .../provider_dependencies.json.sha256sum | 2 +- .../providers/common/ai/toolsets/sql.py | 1 - .../tests/unit/common/ai/toolsets/test_sql.py | 60 +++++-------------- 4 files changed, 16 insertions(+), 49 deletions(-) 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/src/airflow/providers/common/ai/toolsets/sql.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py index 521016ab3bce1..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 @@ -183,7 +183,6 @@ def _is_table_allowed(self, name: str) -> bool: def id(self) -> str: return f"sql-{self._db_conn_id}" - # ------------------------------------------------------------------ # Lazy hook resolution # ------------------------------------------------------------------ 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 067d0fe05bd20..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 @@ -367,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): @@ -375,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"} @@ -385,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"] @@ -393,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") @@ -410,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() @@ -424,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): @@ -440,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") @@ -465,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"] @@ -480,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") @@ -493,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") @@ -510,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 From 8ab47657300977e1681e8ab1a27e05be5871bbe2 Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Sat, 6 Jun 2026 12:59:30 +0100 Subject: [PATCH 28/28] Resolve uv --- uv.lock | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/uv.lock b/uv.lock index cb3451910fb8e..49fcaac02e55d 100644 --- a/uv.lock +++ b/uv.lock @@ -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]]