This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1442b522d20 Add UsageLimits support to common.ai operators (#66248)
1442b522d20 is described below
commit 1442b522d20d09a143689ee406410f8033f3938a
Author: Kaxil Naik <[email protected]>
AuthorDate: Sat May 2 04:55:39 2026 +0100
Add UsageLimits support to common.ai operators (#66248)
Threads pydantic-ai UsageLimits through every common.ai operator
(LLMOperator, AgentOperator, LLMBranchOperator, LLMSQLQueryOperator,
LLMFileAnalysisOperator, LLMSchemaCompareOperator) and their @task.*
decorators by forwarding the new keyword to agent.run_sync. Use it to
fail tasks that exceed token, request, or tool-call budgets -- the
check happens inside pydantic-ai's run loop so the limit applies even
when retries trigger multiple model calls within a single task.
Docs call out the implicit request_limit=50 default in pydantic-ai's
UsageLimits() so users who only want a token cap know to pass
request_limit=None explicitly.
---
providers/common/ai/docs/operators/agent.rst | 5 +++
providers/common/ai/docs/operators/llm.rst | 32 ++++++++++++++++
.../common/ai/example_dags/example_llm.py | 25 ++++++++++++
.../airflow/providers/common/ai/operators/agent.py | 15 ++++++--
.../airflow/providers/common/ai/operators/llm.py | 10 ++++-
.../providers/common/ai/operators/llm_branch.py | 2 +-
.../common/ai/operators/llm_file_analysis.py | 2 +-
.../common/ai/operators/llm_schema_compare.py | 2 +-
.../providers/common/ai/operators/llm_sql.py | 2 +-
.../tests/unit/common/ai/decorators/test_agent.py | 4 +-
.../ai/tests/unit/common/ai/decorators/test_llm.py | 4 +-
.../unit/common/ai/decorators/test_llm_branch.py | 2 +-
.../common/ai/decorators/test_llm_file_analysis.py | 4 +-
.../unit/common/ai/decorators/test_llm_sql.py | 2 +-
.../tests/unit/common/ai/operators/test_agent.py | 44 +++++++++++++++++++++-
.../ai/tests/unit/common/ai/operators/test_llm.py | 21 ++++++++++-
.../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 | 5 ++-
.../tests/unit/common/ai/operators/test_llm_sql.py | 2 +-
20 files changed, 163 insertions(+), 24 deletions(-)
diff --git a/providers/common/ai/docs/operators/agent.rst
b/providers/common/ai/docs/operators/agent.rst
index 71d55617ec0..fc5efa27780 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 219965578bf..a69a592eff8 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
<https://ai.pydantic.dev/api/usage/#pydantic_ai.usage.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 9972d95f840..3ff8509285e 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 __future__ import annotations
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 example_llm_decorator_structured():
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 68da378a994..2660bb408a1 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 @@ from airflow.providers.common.compat.version_compat import
AIRFLOW_V_3_1_PLUS
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 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
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 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
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 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
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 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
"""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 73000eb50a3..95eb6b1442c 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 @@ from airflow.providers.common.compat.sdk import BaseOperator
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 <https://ai.pydantic.dev/api/agent/>`__
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 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
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 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
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 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
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 7a2cda2a671..0395040852f 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 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
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 bb10b66680d..e488aa99d0f 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 @@ class LLMFileAnalysisOperator(LLMOperator):
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 b07413d136e..b51ea7c1f7c 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 @@ class LLMSchemaCompareOperator(LLMOperator):
**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 53c1424fb0f..344b27d122d 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 @@ class LLMSQLQueryOperator(LLMOperator):
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 6865db9a1dd..cfe62a38f27 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 @@ class TestAgentDecoratedOperator:
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 @@ class TestAgentDecoratedOperator:
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 21f663e177a..05768bd7452 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 @@ class TestLLMDecoratedOperator:
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 @@ class TestLLMDecoratedOperator:
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 6620e505db1..c00ed3935e3 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 @@ class TestLLMBranchDecoratedOperator:
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 5b32f061862..41fc750d505 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 @@ class TestLLMFileAnalysisDecoratedOperator:
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 @@ class TestLLMFileAnalysisDecoratedOperator:
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 0f3f943926a..244db8d3150 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 @@ class TestLLMSQLDecoratedOperator:
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 9d2af0d2958..c817e12bd9b 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 @@ from unittest.mock import MagicMock, patch
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 @@ class TestAgentOperatorTemplateFields:
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 @@ class TestAgentOperatorExecute:
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 @@ class TestAgentOperatorRegenerateWithFeedback:
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 @@ class TestAgentOperatorDurable:
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 3a686d014f0..165a2c30b7a 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 @@ from uuid import uuid4
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 @@ class TestLLMOperator:
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 78e6c231080..82f710cc10a 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 @@ class TestLLMBranchOperator:
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 abd98d92326..a2b223c60e4 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 @@ class TestLLMFileAnalysisOperator:
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 5515cce57f1..d719162bc00 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 @@ class TestLLMSchemaCompareOperator:
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 @@ class TestLLMSchemaCompareOperator:
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 4ff3158a356..96719e2662e 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 @@ class TestLLMSQLQueryOperator:
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):
