diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index 71d55617ec07e..fc5efa27780cf 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -230,6 +230,11 @@ Parameters every tool call is logged in real time. Default ``True``. - ``agent_params``: Additional keyword arguments passed to the pydantic-ai ``Agent`` constructor (e.g. ``retries``, ``model_settings``). +- ``usage_limits``: Optional pydantic-ai ``UsageLimits`` enforced on every + agent run (initial run, durable replay, and HITL regeneration). Use it to + cap requests, tokens, or tool calls per task -- agents are particularly + prone to runaway tool loops, so ``tool_calls_limit`` is a useful guardrail. + See :ref:`howto/operator:llm` for an example. Default ``None``. - ``durable``: When ``True``, enables step-level caching of model responses and tool results via ObjectStorage. On retry, cached steps are replayed instead of re-executing expensive LLM calls. Requires the ``[common.ai] durable_cache_path`` diff --git a/providers/common/ai/docs/operators/llm.rst b/providers/common/ai/docs/operators/llm.rst index 219965578bf65..a69a592eff8b4 100644 --- a/providers/common/ai/docs/operators/llm.rst +++ b/providers/common/ai/docs/operators/llm.rst @@ -66,6 +66,36 @@ the full list of supported parameters. :start-after: [START howto_operator_llm_agent_params] :end-before: [END howto_operator_llm_agent_params] +Usage Limits +------------ + +Set ``usage_limits`` to a +`pydantic-ai UsageLimits `__ +to fail the task when the run exceeds a configured budget — request count, +input/output tokens, or tool calls. The check happens inside pydantic-ai's +run loop, so the limit applies even when ``retries`` triggers multiple model +calls within a single task. + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py + :language: python + :start-after: [START howto_operator_llm_usage_limits] + :end-before: [END howto_operator_llm_usage_limits] + +Common knobs on ``UsageLimits``: + +- ``request_limit`` — max model requests per run (caps retry/tool-loop blow-ups). + pydantic-ai applies a default of ``50`` when ``UsageLimits()`` is constructed + without an explicit value, so passing ``UsageLimits(input_tokens_limit=4_000)`` + silently inherits that 50-request cap. Set ``request_limit=None`` to disable + it explicitly when you only want a token cap. +- ``input_tokens_limit`` / ``output_tokens_limit`` — per-run token caps. +- ``total_tokens_limit`` — combined input + output cap. +- ``tool_calls_limit`` — max tool invocations (``AgentOperator`` only). + +When the limit is hit pydantic-ai raises ``UsageLimitExceeded``, which +propagates to Airflow as a task failure — Airflow's standard retry policy +applies on top. + TaskFlow Decorator ------------------ @@ -132,6 +162,8 @@ Parameters for structured output. - ``agent_params``: Additional keyword arguments passed to the pydantic-ai ``Agent`` constructor (e.g. ``retries``, ``model_settings``, ``tools``). Supports Jinja templating. +- ``usage_limits``: Optional pydantic-ai ``UsageLimits`` enforced on the run. Fails + the task when token / request / tool-call budgets are exceeded. Default ``None``. - ``require_approval``: If ``True``, the task defers after generating output and waits for human review. Default ``False``. - ``approval_timeout``: Maximum time to wait for a review (``timedelta``). ``None`` diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py index 9972d95f84080..3ff8509285e95 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py +++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm.py @@ -21,6 +21,7 @@ from datetime import timedelta from pydantic import BaseModel +from pydantic_ai.usage import UsageLimits from airflow.providers.common.ai.operators.llm import LLMOperator from airflow.providers.common.compat.sdk import dag, task @@ -118,6 +119,30 @@ def extract(text: str): example_llm_decorator_structured() +# [START howto_operator_llm_usage_limits] +@dag +def example_llm_operator_usage_limits(): + LLMOperator( + task_id="capped_summary", + prompt="Summarize the attached design doc in three bullet points.", + llm_conn_id="pydanticai_default", + system_prompt="You are a concise technical reviewer.", + # Fail the task if the run exceeds 5 model requests, 4_000 input + # tokens, or 1_000 output tokens. Useful for guardrails on shared + # connections or untrusted prompts. + usage_limits=UsageLimits( + request_limit=5, + input_tokens_limit=4_000, + output_tokens_limit=1_000, + ), + ) + + +# [END howto_operator_llm_usage_limits] + +example_llm_operator_usage_limits() + + # [START howto_operator_llm_approval] @dag def example_llm_operator_approval(): 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 68da378a99471..2660bb408a187 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 @@ -40,6 +40,7 @@ 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 @@ -102,6 +103,12 @@ class AgentOperator(BaseOperator, HITLReviewMixin): 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 + :class:`~pydantic_ai.usage.UsageLimits` enforced on every agent run + (initial run, durable replay, and HITL regeneration). Pass + ``UsageLimits(request_limit=..., total_tokens_limit=..., tool_calls_limit=..., ...)`` + to fail the task when the agent exceeds the configured token, request, + or tool budget. ``None`` (default) means no enforcement. :param durable: When ``True``, enables step-level caching of model responses and tool results for durable execution. On retry, cached steps are replayed instead of re-executing. Default ``False``. @@ -147,6 +154,7 @@ def __init__( toolsets: list[AbstractToolset] | None = None, enable_tool_logging: bool = True, agent_params: dict[str, Any] | None = None, + usage_limits: UsageLimits | None = None, durable: bool = False, # Agent feedback parameters enable_hitl_review: bool = False, @@ -165,6 +173,7 @@ def __init__( self.toolsets = toolsets self.enable_tool_logging = enable_tool_logging self.agent_params = agent_params or {} + self.usage_limits = usage_limits self.durable = durable @@ -246,9 +255,9 @@ def execute(self, context: Context) -> Any: 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) + result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) else: - result = agent.run_sync(self.prompt) + result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) log_run_summary(self.log, result) @@ -293,7 +302,7 @@ def regenerate_with_feedback(self, *, feedback: str, message_history: Any) -> tu """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) + result = agent.run_sync(feedback, message_history=messages, usage_limits=self.usage_limits) log_run_summary(self.log, result) output = result.output 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 73000eb50a357..95eb6b1442c19 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 @@ -32,6 +32,7 @@ if TYPE_CHECKING: from pydantic_ai import Agent + from pydantic_ai.usage import UsageLimits from airflow.sdk import Context @@ -56,6 +57,11 @@ class LLMOperator(BaseOperator, LLMApprovalMixin): ``Agent`` constructor (e.g. ``retries``, ``model_settings``, ``tools``). See `pydantic-ai Agent docs `__ for the full list. + :param usage_limits: Optional pydantic-ai + :class:`~pydantic_ai.usage.UsageLimits` enforced on the run. Pass + ``UsageLimits(request_limit=..., total_tokens_limit=..., ...)`` to fail + the task when the agent exceeds the configured token, request, or tool + budget. ``None`` (default) means no enforcement. :param require_approval: If ``True``, the task defers after generating output and waits for a human reviewer to approve or reject via the HITL interface. Default ``False``. @@ -83,6 +89,7 @@ def __init__( system_prompt: str = "", output_type: type = str, agent_params: dict[str, Any] | None = None, + usage_limits: UsageLimits | None = None, require_approval: bool = False, approval_timeout: timedelta | None = None, allow_modifications: bool = False, @@ -95,6 +102,7 @@ def __init__( self.system_prompt = system_prompt self.output_type = output_type self.agent_params = agent_params or {} + self.usage_limits = usage_limits self.require_approval = require_approval self.approval_timeout = approval_timeout self.allow_modifications = allow_modifications @@ -118,7 +126,7 @@ def execute(self, context: Context) -> Any: agent: Agent[None, Any] = self.llm_hook.create_agent( output_type=self.output_type, instructions=self.system_prompt, **self.agent_params ) - result = agent.run_sync(self.prompt) + result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) 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 7a2cda2a671e2..0395040852f53 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 @@ -81,7 +81,7 @@ def execute(self, context: Context) -> str | Iterable[str] | None: instructions=self.system_prompt, **self.agent_params, ) - result = agent.run_sync(self.prompt) + result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) 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 bb10b66680def..e488aa99d0f23 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 @@ -134,7 +134,7 @@ def execute(self, context: Context) -> Any: instructions=self._build_system_prompt(), **self.agent_params, ) - result = agent.run_sync(request.user_content) + result = agent.run_sync(request.user_content, usage_limits=self.usage_limits) 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 b07413d136e9d..b51ea7c1f7c3b 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 @@ -309,7 +309,7 @@ def execute(self, context: Context) -> dict[str, Any]: **self.agent_params, ) self.log.info("Running LLM schema comparison...") - result = agent.run_sync(self.prompt) + result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) 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 53c1424fb0f10..344b27d122dfc 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 @@ -147,7 +147,7 @@ def execute(self, context: Context) -> str: agent = self.llm_hook.create_agent( output_type=str, instructions=full_system_prompt, **self.agent_params ) - result = agent.run_sync(self.prompt) + result = agent.run_sync(self.prompt, usage_limits=self.usage_limits) log_run_summary(self.log, result) sql = self._strip_llm_output(result.output) 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 6865db9a1dd77..cfe62a38f27a9 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 @@ -56,7 +56,7 @@ def my_prompt(): 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?") + mock_agent.run_sync.assert_called_once_with("Who is our top customer?", usage_limits=None) @pytest.mark.parametrize( "return_value", @@ -92,7 +92,7 @@ def my_prompt(topic): op.execute(context={"task_instance": MagicMock()}) assert op.prompt == "Analyze revenue trends" - mock_agent.run_sync.assert_called_once_with("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): 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 21f663e177adb..05768bd74522b 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 @@ -54,7 +54,7 @@ def my_prompt(): assert result == "This is a summary." assert op.prompt == "Summarize this text" - mock_agent.run_sync.assert_called_once_with("Summarize this text") + mock_agent.run_sync.assert_called_once_with("Summarize this text", usage_limits=None) @pytest.mark.parametrize( "return_value", @@ -90,4 +90,4 @@ def my_prompt(topic): op.execute(context={"task_instance": MagicMock()}) assert op.prompt == "Summarize quantum computing" - mock_agent.run_sync.assert_called_once_with("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 6620e505db148..c00ed3935e32d 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 @@ -66,7 +66,7 @@ def my_prompt(): assert result == "positive" assert op.prompt == "Route this review" - mock_agent.run_sync.assert_called_once_with("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( 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 5b32f0618627f..41fc750d5051c 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 @@ -71,7 +71,7 @@ def my_prompt(): assert result == "This is a summary." assert op.prompt == "Summarize this text" - mock_agent.run_sync.assert_called_once_with("prepared prompt") + mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None) @pytest.mark.parametrize( "return_value", @@ -115,4 +115,4 @@ def my_prompt(topic): 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") + 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_sql.py b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py index 0f3f943926a70..244db8d315078 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 @@ -54,7 +54,7 @@ def my_prompt_fn(): assert result == "SELECT 1" assert op.prompt == "Get all users" - mock_agent.run_sync.assert_called_once_with("Get all users") + mock_agent.run_sync.assert_called_once_with("Get all users", usage_limits=None) @pytest.mark.parametrize( "return_value", 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 9d2af0d29581a..c817e12bd9b2c 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 @@ -21,6 +21,7 @@ import pytest from pydantic import BaseModel +from pydantic_ai.usage import UsageLimits from airflow.providers.common.ai.operators.agent import AgentOperator, HITLReviewLink from airflow.providers.common.ai.toolsets.logging import LoggingToolset @@ -79,6 +80,44 @@ def test_template_fields(self): class TestAgentOperatorExecute: + @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", 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 + + limits = UsageLimits(request_limit=3, tool_calls_limit=5) + op = AgentOperator( + task_id="test", + prompt="run", + llm_conn_id="my_llm", + usage_limits=limits, + ) + op.execute(context=MagicMock()) + + mock_agent.run_sync.assert_called_once_with("run", usage_limits=limits) + + @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", 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 + + limits = UsageLimits(request_limit=1) + op = AgentOperator( + task_id="test", + prompt="run", + llm_conn_id="my_llm", + usage_limits=limits, + ) + op.regenerate_with_feedback(feedback="Add detail", message_history=[]) + + mock_agent.run_sync.assert_called_once_with( + "Add detail", + message_history=[], + usage_limits=limits, + ) + @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) def test_execute_creates_agent_from_hook(self, mock_hook_cls): mock_agent = _make_mock_agent("The answer is 42.") @@ -97,7 +136,7 @@ def test_execute_creates_agent_from_hook(self, mock_hook_cls): 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?") + mock_agent.run_sync.assert_called_once_with("What is the answer?", usage_limits=None) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) def test_execute_passes_toolsets_in_agent_kwargs(self, mock_hook_cls): @@ -378,6 +417,7 @@ def test_regenerate_with_feedback_calls_agent_with_feedback_and_history(self, mo mock_agent.run_sync.assert_called_once_with( "Add more detail", message_history=msg_history, + usage_limits=None, ) @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True) @@ -462,4 +502,4 @@ def test_execute_non_durable_does_not_wrap(self, mock_hook_cls): op.execute(context=MagicMock()) # run_sync called directly, no override - mock_agent.run_sync.assert_called_once_with("test") + mock_agent.run_sync.assert_called_once_with("test", usage_limits=None) 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 3a686d014f0d0..165a2c30b7a27 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 @@ -22,6 +22,7 @@ import pytest from pydantic import BaseModel +from pydantic_ai.usage import UsageLimits from airflow.providers.common.ai.mixins.approval import ( LLMApprovalMixin, @@ -59,12 +60,30 @@ def test_execute_returns_string_output(self, mock_hook_cls): 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?") + 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 + + limits = UsageLimits(request_limit=2, output_tokens_limit=100) + op = LLMOperator( + task_id="test", + prompt="Summarize", + llm_conn_id="my_llm", + usage_limits=limits, + ) + op.execute(context=MagicMock()) + + mock_agent.run_sync.assert_called_once_with("Summarize", usage_limits=limits) + @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) def test_execute_structured_output_with_all_params(self, mock_hook_cls): """Structured output via model_dump(), with model_id, system_prompt, and agent_params.""" 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 78e6c23108068..82f710cc10a0b 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 @@ -79,7 +79,7 @@ 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") + 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) 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 abd98d9232646..a2b223c60e4a5 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 @@ -101,7 +101,7 @@ 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") + mock_agent.run_sync.assert_called_once_with("prepared prompt", usage_limits=None) @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) @patch( 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 5515cce57f1d6..d719162bc00f7 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 @@ -275,7 +275,7 @@ def test_execute(self, mock_build_system_prompt, mock_build_schema_context): instructions="system_prompt", param="value", ) - mock_agent.run_sync.assert_called_once_with("user_prompt") + mock_agent.run_sync.assert_called_once_with("user_prompt", usage_limits=None) assert result == {"compatible": True, "mismatches": [], "summary": "All good"} @mock.patch( @@ -353,7 +353,8 @@ def test_execute_schema_comparison_mixed_conn(self, mock_get_db_hook, db_hook): assert "aws_default" in instructions mock_agent.run_sync.assert_called_once_with( - "Compare S3 Parquet schema against the Postgres table and flag breaking changes" + "Compare S3 Parquet schema against the Postgres table and flag breaking changes", + usage_limits=None, ) assert result["compatible"] is True assert result["summary"] == "S3 and Postgres schemas are compatible" 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 4ff3158a3562d..96719e2662ec0 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 @@ -111,7 +111,7 @@ 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") + mock_agent.run_sync.assert_called_once_with("Get active users", usage_limits=None) @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook", autospec=True) def test_execute_validation_blocks_unsafe_sql(self, mock_hook_cls):