This is an automated email from the ASF dual-hosted git repository.
shahar1 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 66c31fa7f08 Document usage_limits on all common.ai LLM operators
(#73283)
66c31fa7f08 is described below
commit 66c31fa7f080732b6fc865397f4481e6faa6e254
Author: Jyun-An Chen <[email protected]>
AuthorDate: Thu Sep 17 21:57:32 2026 +0800
Document usage_limits on all common.ai LLM operators (#73283)
---
providers/common/ai/docs/operators/llm.rst | 2 ++
providers/common/ai/docs/operators/llm_branch.rst | 3 +++
providers/common/ai/docs/operators/llm_file_analysis.rst | 3 +++
providers/common/ai/docs/operators/llm_schema_compare.rst | 3 +++
providers/common/ai/docs/operators/llm_sql.rst | 7 +++++++
.../ai/src/airflow/providers/common/ai/operators/llm_branch.py | 3 +++
.../src/airflow/providers/common/ai/operators/llm_file_analysis.py | 3 +++
.../airflow/providers/common/ai/operators/llm_schema_compare.py | 3 +++
.../common/ai/src/airflow/providers/common/ai/operators/llm_sql.py | 3 +++
9 files changed, 30 insertions(+)
diff --git a/providers/common/ai/docs/operators/llm.rst
b/providers/common/ai/docs/operators/llm.rst
index ce0b2e18784..19bb702496c 100644
--- a/providers/common/ai/docs/operators/llm.rst
+++ b/providers/common/ai/docs/operators/llm.rst
@@ -104,6 +104,8 @@ the full list of supported parameters.
:start-after: [START howto_operator_llm_agent_params]
:end-before: [END howto_operator_llm_agent_params]
+.. _howto/operator:llm_usage_limits:
+
Usage Limits
------------
diff --git a/providers/common/ai/docs/operators/llm_branch.rst
b/providers/common/ai/docs/operators/llm_branch.rst
index c8999bfa003..621c78efdc2 100644
--- a/providers/common/ai/docs/operators/llm_branch.rst
+++ b/providers/common/ai/docs/operators/llm_branch.rst
@@ -138,6 +138,9 @@ Parameters
task ID. When ``True`` the LLM may return one or more task IDs.
- ``agent_params``: Additional keyword arguments passed to the pydantic-ai
``Agent``
constructor (e.g. ``retries``, ``model_settings``). Supports Jinja
templating.
+- ``usage_limits``: Optional pydantic-ai ``UsageLimits`` (or a templated
``dict`` of
+ the same fields) enforced on the run; the task fails when a budget is
exceeded.
+ Default ``None``. See :ref:`Usage Limits <howto/operator:llm_usage_limits>`.
- ``require_approval``: If ``True``, the task pauses after the LLM chooses the
branch(es) and waits for human review before branching. Default ``False``.
- ``approval_timeout``: Maximum time to wait for a review (``timedelta``).
``None``
diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst
b/providers/common/ai/docs/operators/llm_file_analysis.rst
index fc9ce6ac29f..ad9da9dcc12 100644
--- a/providers/common/ai/docs/operators/llm_file_analysis.rst
+++ b/providers/common/ai/docs/operators/llm_file_analysis.rst
@@ -152,6 +152,9 @@ Parameters
``BaseModel`` for structured output.
- ``agent_params``: Additional keyword arguments passed to the pydantic-ai
``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+- ``usage_limits``: Optional pydantic-ai ``UsageLimits`` (or a templated
``dict`` of
+ the same fields) enforced on the run; the task fails when a budget is
exceeded.
+ Default ``None``. See :ref:`Usage Limits <howto/operator:llm_usage_limits>`.
- ``serialize_output``: If ``True`` and ``output_type`` is a Pydantic
``BaseModel`` subclass, the model instance is dumped to a ``dict`` via
``model_dump()`` before being pushed to XCom. Default ``False`` -- the
diff --git a/providers/common/ai/docs/operators/llm_schema_compare.rst
b/providers/common/ai/docs/operators/llm_schema_compare.rst
index 426fe5ea228..1132081862e 100644
--- a/providers/common/ai/docs/operators/llm_schema_compare.rst
+++ b/providers/common/ai/docs/operators/llm_schema_compare.rst
@@ -181,6 +181,9 @@ Parameters
:ref:`Customizing the System Prompt <howto/operator:llm_schema_compare>`
above).
- ``agent_params``: Additional keyword arguments passed to the pydantic-ai
``Agent`` constructor.
+- ``usage_limits``: Optional pydantic-ai ``UsageLimits`` (or a templated
``dict`` of
+ the same fields) enforced on the run; the task fails when a budget is
exceeded.
+ Default ``None``. See :ref:`Usage Limits <howto/operator:llm_usage_limits>`.
- ``db_conn_ids``: List of database connection IDs to compare. Each must
resolve
to a ``DbApiHook``.
- ``table_names``: Tables to introspect from each ``db_conn_id``.
diff --git a/providers/common/ai/docs/operators/llm_sql.rst
b/providers/common/ai/docs/operators/llm_sql.rst
index c205fa307e2..45b64c68b67 100644
--- a/providers/common/ai/docs/operators/llm_sql.rst
+++ b/providers/common/ai/docs/operators/llm_sql.rst
@@ -135,6 +135,13 @@ Generate SQL for multiple prompts in parallel using
``expand()``:
:start-after: [START howto_operator_llm_sql_expand]
:end-before: [END howto_operator_llm_sql_expand]
+Usage Limits
+------------
+
+``usage_limits`` caps the request count, token usage, tool calls, or cost of
the
+SQL generation run, and fails the task when a budget is exceeded. It is
inherited
+from ``LLMOperator`` -- see :ref:`Usage Limits
<howto/operator:llm_usage_limits>`.
+
Human-in-the-Loop Approval
--------------------------
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 ddaa39996dc..75291ebe9f8 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
@@ -59,6 +59,9 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
:param agent_params: Additional keyword arguments passed to the pydantic-ai
``Agent`` constructor (e.g. ``retries``, ``model_settings``,
``tools``).
+ ``usage_limits`` is inherited from
+ :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+
Human-in-the-Loop approval parameters are inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
(``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
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 66d31d2eb59..8245312018e 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
@@ -70,6 +70,9 @@ class LLMFileAnalysisOperator(LLMOperator):
read from storage and ``max_text_chars`` limits the final prompt text
budget. Default ``10``.
+ ``usage_limits`` is inherited from
+ :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+
Human-in-the-Loop approval parameters are inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
(``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
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 aa25bf7229d..f1f494bf263 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
@@ -109,6 +109,9 @@ class LLMSchemaCompareOperator(LLMOperator):
``"full"`` to include primary keys, foreign keys, and indexes.
Default ``"full"``.
+ ``usage_limits`` is inherited from
+ :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+
Human-in-the-Loop approval parameters are inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
(``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
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 4e22fdc845e..e4415c2f8a5 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
@@ -82,6 +82,9 @@ class LLMSQLQueryOperator(LLMOperator):
:param dialect: SQL dialect for parsing (``postgres``, ``mysql``, etc.).
Auto-detected from the database hook if not set.
+ ``usage_limits`` is inherited from
+ :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+
Human-in-the-Loop approval parameters are inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
(``require_approval``, ``approval_timeout``, ``on_approval_timeout``,
