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 7993ee19955 Remove `agent_params` from LLM operators so they cannot
run tools (#74721)
7993ee19955 is described below
commit 7993ee19955bc8f71aa23a4b22b86c5913025b32
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Oct 14 12:39:02 2026 +0100
Remove `agent_params` from LLM operators so they cannot run tools (#74721)
LLMOperator and the operators built on it pass agent_params straight to
pydantic-ai's Agent, so tools, toolsets or capabilities turned a single
model call into an agent loop without AgentOperator's tool approval,
durable replay or tool logging. Replace agent_params with typed
model_settings and output_retries arguments; AgentOperator keeps
agent_params.
---
providers/common/ai/docs/changelog.rst | 10 ++++
providers/common/ai/docs/operators/llm.rst | 43 ++++++++++----
providers/common/ai/docs/operators/llm_branch.rst | 6 +-
.../common/ai/docs/operators/llm_file_analysis.rst | 6 +-
.../ai/docs/operators/llm_schema_compare.rst | 6 +-
.../example_dags/example_aip_progress_tracker.py | 6 +-
.../common/ai/example_dags/example_llm.py | 11 ++--
.../airflow/providers/common/ai/operators/llm.py | 43 ++++++++++----
.../providers/common/ai/operators/llm_branch.py | 9 +--
.../common/ai/operators/llm_file_analysis.py | 11 ++--
.../common/ai/operators/llm_schema_compare.py | 8 ++-
.../providers/common/ai/operators/llm_sql.py | 10 ++--
.../ai/tests/unit/common/ai/operators/test_llm.py | 66 +++++++++++++++++++---
.../unit/common/ai/operators/test_llm_branch.py | 25 ++++++++
.../common/ai/operators/test_llm_file_analysis.py | 31 +++++++++-
.../common/ai/operators/test_llm_schema_compare.py | 9 +--
.../tests/unit/common/ai/operators/test_llm_sql.py | 13 +++--
17 files changed, 241 insertions(+), 72 deletions(-)
diff --git a/providers/common/ai/docs/changelog.rst
b/providers/common/ai/docs/changelog.rst
index 82706b87d39..4095637fe5c 100644
--- a/providers/common/ai/docs/changelog.rst
+++ b/providers/common/ai/docs/changelog.rst
@@ -26,6 +26,16 @@ Changelog
---------
.. note::
+ ``LLMOperator``, the operators built on it (``LLMBranchOperator``,
``LLMSQLQueryOperator``,
+ ``LLMSchemaCompareOperator``, ``LLMFileAnalysisOperator``) and their
``@task.llm*``
+ decorators no longer accept ``agent_params``. These operators make one model
call and run
+ no tools in the task; an agent with function tools, toolsets or capabilities
belongs in
+ ``AgentOperator``, which keeps ``agent_params``. Passing ``agent_params`` to
an LLM operator now fails when the Dag is
+ parsed. ``agent_params={"model_settings": {...}}`` becomes
``model_settings={...}``,
+ ``agent_params={"retries": n}`` becomes ``output_retries=n``, and ``tools``,
``toolsets``,
+ ``capabilities`` or ``spec_file`` move to ``AgentOperator`` or
``@task.agent``. See
+ :ref:`howto/operator:llm`.
+
On Airflow >= 3.3, the ``usage_limits`` of the ``LLM*`` operators
(``LLMOperator``,
``LLMBranchOperator``, ``LLMSQLQueryOperator``, ``LLMSchemaCompareOperator``,
``LLMFileAnalysisOperator``) counts usage across every attempt of the task
instance
diff --git a/providers/common/ai/docs/operators/llm.rst
b/providers/common/ai/docs/operators/llm.rst
index fa0ef52a493..b789fa0472a 100644
--- a/providers/common/ai/docs/operators/llm.rst
+++ b/providers/common/ai/docs/operators/llm.rst
@@ -59,18 +59,35 @@ XCom unchanged, so downstream tasks can type-hint the class
and use attribute ac
:start-after: [START howto_operator_llm_structured]
:end-before: [END howto_operator_llm_structured]
-Agent parameters
-----------------
+Model settings and output retries
+---------------------------------
-Pass additional keyword arguments to the pydantic-ai ``Agent`` constructor
-via ``agent_params``: for example ``retries``, ``model_settings``, or
``tools``.
-See the `pydantic-ai Agent docs <https://ai.pydantic.dev/api/agent/>`__ for
-the full list of supported parameters.
+Set ``model_settings`` to pass Pydantic AI
+`ModelSettings <https://ai.pydantic.dev/api/settings/>`__ to the model call,
such as
+``temperature`` or ``max_tokens``. Set ``output_retries`` to choose how many
times the
+model is asked again when its output fails validation against ``output_type``.
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
:language: python
- :start-after: [START howto_operator_llm_agent_params]
- :end-before: [END howto_operator_llm_agent_params]
+ :start-after: [START howto_operator_llm_model_settings]
+ :end-before: [END howto_operator_llm_model_settings]
+
+No tool runs in the task: ``LLMOperator`` accepts no function tools, toolsets
or
+capabilities, so ``model_settings`` cannot register one. A provider-hosted
tool that a
+provider's own setting turns on, such as ``openai_native_tools`` for OpenAI's
web search,
+runs on the provider's side. For a model that calls your tools, use
+:ref:`AgentOperator <howto/operator:agent>`.
+
+Moving from ``agent_params``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Common AI 1.0 removes ``agent_params`` from ``LLMOperator``, the operators
built on it and
+their ``@task.llm*`` decorators. Passing it fails when the Dag is parsed. Move
each key:
+
+- ``agent_params={"model_settings": {...}}`` becomes ``model_settings={...}``.
+- ``agent_params={"retries": n}`` becomes ``output_retries=n``.
+- ``tools``, ``toolsets``, ``capabilities`` or ``spec_file`` in
``agent_params``: use
+ ``AgentOperator`` or ``@task.agent``, which keep ``agent_params``.
.. _howto/operator:llm_usage_limits:
@@ -81,8 +98,8 @@ 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.
+run loop, so the limit applies even when ``output_retries`` asks the model
again
+within a single attempt.
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm.py
:language: python
@@ -263,8 +280,10 @@ Parameters
the operator to return the output without a person, from 0 to 1; no reported
confidence counts as uncertain. ``on_uncertain`` is ``"review"`` (default) or
``"fail"``. Default ``None``: no gate.
-- ``agent_params``: Additional keyword arguments passed to the pydantic-ai
``Agent``
- constructor (e.g. ``retries``, ``model_settings``, ``tools``). Supports
Jinja templating.
+- ``model_settings``: Pydantic AI ``ModelSettings`` for the model call, such as
+ ``temperature`` or ``max_tokens``. Supports Jinja templating.
+- ``output_retries``: How many times the model is asked again when its output
fails
+ validation. Default ``None`` (Pydantic AI's default).
- ``usage_limits``: Optional pydantic-ai ``UsageLimits`` enforced on the run,
or a
``dict`` of the same fields (templated via Jinja, then coerced per field
type).
Fails the task when token / request / tool-call budgets are exceeded, or
when a
diff --git a/providers/common/ai/docs/operators/llm_branch.rst
b/providers/common/ai/docs/operators/llm_branch.rst
index 0e113aed050..04abfe903f2 100644
--- a/providers/common/ai/docs/operators/llm_branch.rst
+++ b/providers/common/ai/docs/operators/llm_branch.rst
@@ -313,8 +313,10 @@ Parameters
Jinja templating. Default ``None``.
- ``allow_multiple_branches``: When ``False`` (default) the LLM returns a
single
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.
+- ``model_settings``: Pydantic AI ``ModelSettings`` for the model call, such as
+ ``temperature`` or ``max_tokens``. Supports Jinja templating.
+- ``output_retries``: How many times the model is asked again when its output
fails
+ validation. Default ``None`` (Pydantic AI's default).
- ``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>`.
diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst
b/providers/common/ai/docs/operators/llm_file_analysis.rst
index 0f9cdbc6279..1a666b16893 100644
--- a/providers/common/ai/docs/operators/llm_file_analysis.rst
+++ b/providers/common/ai/docs/operators/llm_file_analysis.rst
@@ -150,8 +150,10 @@ Parameters
built-in read-only guidance.
- ``output_type``: Expected output type (default: ``str``). Set to a Pydantic
``BaseModel`` for structured output.
-- ``agent_params``: Additional keyword arguments passed to the pydantic-ai
- ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+- ``model_settings``: Pydantic AI ``ModelSettings`` for the model call, such as
+ ``temperature`` or ``max_tokens``. Supports Jinja templating.
+- ``output_retries``: How many times the model is asked again when its output
fails
+ validation. Default ``None`` (Pydantic AI's default).
- ``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>`.
diff --git a/providers/common/ai/docs/operators/llm_schema_compare.rst
b/providers/common/ai/docs/operators/llm_schema_compare.rst
index 39be5fefa9d..aaaad00d60a 100644
--- a/providers/common/ai/docs/operators/llm_schema_compare.rst
+++ b/providers/common/ai/docs/operators/llm_schema_compare.rst
@@ -189,8 +189,10 @@ Parameters
severity definitions. Passing a value **replaces** the default, so
concatenate
with ``DEFAULT_SYSTEM_PROMPT`` to extend it (see
:ref:`Customizing the System Prompt <howto/operator:llm_schema_compare>`
above).
-- ``agent_params``: Additional keyword arguments passed to the pydantic-ai
- ``Agent`` constructor.
+- ``model_settings``: Pydantic AI ``ModelSettings`` for the model call, such as
+ ``temperature`` or ``max_tokens``. Supports Jinja templating.
+- ``output_retries``: How many times the model is asked again when its output
fails
+ validation. Default ``None`` (Pydantic AI's default).
- ``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>`.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aip_progress_tracker.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aip_progress_tracker.py
index 191935e42aa..3d322fcd616 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aip_progress_tracker.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aip_progress_tracker.py
@@ -662,7 +662,7 @@ def example_aip_progress_tracker():
system_prompt=ANALYSIS_SYSTEM_PROMPT,
output_type=AIPStatus,
serialize_output=True,
- agent_params={"model_settings": {"temperature": 0}},
+ model_settings={"temperature": 0},
).expand(prompt=prompts)
# [END aip_tracker_dtm_analysis]
@@ -750,7 +750,7 @@ Prioritize AIPs that are close to completion or have shared
dependencies.
Use only the data below -- do not add external information.
{{ ti.xcom_pull(task_ids='format_report') }}""",
- agent_params={"model_settings": {"temperature": 0}},
+ model_settings={"temperature": 0},
usage_limits=UsageLimits(
request_limit=5,
input_tokens_limit=20_000,
@@ -787,7 +787,7 @@ Flag any claims not grounded in the evidence.
input_tokens_limit=30_000,
output_tokens_limit=8_000,
),
- agent_params={"model_settings": {"temperature": 0}},
+ model_settings={"temperature": 0},
)
# [END aip_tracker_validation]
synthesize >> validate
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 b8d2c48b80b..25139db825c 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
@@ -75,21 +75,22 @@ def example_llm_operator_structured():
example_llm_operator_structured()
-# [START howto_operator_llm_agent_params]
+# [START howto_operator_llm_model_settings]
@dag(tags=["example"])
-def example_llm_operator_agent_params():
+def example_llm_operator_model_settings():
LLMOperator(
task_id="creative_writing",
prompt="Write a haiku about data pipelines.",
llm_conn_id="pydanticai_default",
system_prompt="You are a creative writer.",
- agent_params={"model_settings": {"temperature": 0.9}, "retries": 3},
+ model_settings={"temperature": 0.9},
+ output_retries=3,
)
-# [END howto_operator_llm_agent_params]
+# [END howto_operator_llm_model_settings]
-example_llm_operator_agent_params()
+example_llm_operator_model_settings()
# [START howto_decorator_llm]
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 8d06c7764e4..ac894d4903d 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
@@ -65,6 +65,7 @@ if TYPE_CHECKING:
from pydantic_ai import Agent
from pydantic_ai.agent import AgentRunResult
+ from pydantic_ai.settings import ModelSettings
from pydantic_ai.usage import UsageLimits
from airflow.sdk import Context
@@ -91,6 +92,11 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
The Pydantic class must be defined at module scope: classes nested inside
a function or ``@dag``-decorated body cannot be deserialized from XCom.
+ No tool runs in the task: the operator accepts no function tools, toolsets
or
+ capabilities, makes one model call (repeated only when the output fails
validation)
+ and returns its output. For an agent that calls tools, use
+ :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`.
+
Once ``_clear_usage_budget`` has run, nothing may fail the attempt (raising
``DownstreamTasksSkipped`` to signal success is fine): a failure would be
retried
with a fresh budget. This applies to subclasses that override ``execute``
or
@@ -114,10 +120,12 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
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 agent_params: Additional keyword arguments passed to the pydantic-ai
- ``Agent`` constructor (e.g. ``retries``, ``model_settings``,
``tools``).
- See `pydantic-ai Agent docs <https://ai.pydantic.dev/api/agent/>`__
- for the full list.
+ :param model_settings: Pydantic AI
:class:`~pydantic_ai.settings.ModelSettings`
+ for the model call, such as ``temperature`` or ``max_tokens``,
including
+ provider-specific settings. ``None`` (default) uses the model's
defaults.
+ :param output_retries: How many times the model is asked again when its
output
+ fails validation against ``output_type``. ``None`` (default) uses
Pydantic AI's
+ default.
:param usage_limits: Optional pydantic-ai
:class:`~pydantic_ai.usage.UsageLimits` enforced on the run, or a dict
of the same fields (e.g.
@@ -198,7 +206,7 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
"model_id",
"fallback_conn_ids",
"system_prompt",
- "agent_params",
+ "model_settings",
"usage_limits",
)
@@ -211,7 +219,8 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
fallback_conn_ids: list[str] | None = None,
system_prompt: str = "",
output_type: type | UnionType = str,
- agent_params: dict[str, Any] | None = None,
+ model_settings: ModelSettings | None = None,
+ output_retries: int | None = None,
usage_limits: UsageLimits | dict[str, Any] | None = None,
require_approval: bool = False,
approval_timeout: timedelta | None = None,
@@ -241,7 +250,8 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
# for deserialization (its worker-side DAG walk); otherwise, or when
the
# user opts in, dump to a dict so the value is deserializable anywhere.
self._serialize_model_output = serialize_output or not _CORE_WALKER
- self.agent_params = agent_params or {}
+ self.model_settings = model_settings
+ self.output_retries = output_retries
# No validation here -- see coerce_usage_limits() docstring for why.
self.usage_limits = usage_limits
if on_approval_timeout not in ("fail",
*LLMApprovalMixin.TIMEOUT_DEFAULTS):
@@ -305,6 +315,21 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
}
return PydanticAIHook.get_hook(self.llm_conn_id,
hook_params=hook_params)
+ def _create_agent(self, *, output_type: Any, instructions: str) ->
Agent[object, Any]:
+ """
+ Build the agent for this operator's model call.
+
+ The only way any LLM operator builds its agent. Only
``model_settings`` and
+ ``output_retries`` reach it, so no function tool, toolset or
capability is ever
+ registered; an agent with tools is what ``AgentOperator`` is for.
+ """
+ return self.llm_hook.create_agent(
+ output_type=output_type,
+ instructions=instructions,
+ model_settings=self.model_settings,
+ retries=self.output_retries,
+ )
+
def execute(self, context: Context) -> Any:
if self._may_review:
self.validate_approval_prompt() # type: ignore[misc]
@@ -312,9 +337,7 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin, Usag
# Coerced first so a bad rendered value fails before the expensive
setup below.
usage_limits = coerce_usage_limits(self.usage_limits)
- agent: Agent[object, Any] = self.llm_hook.create_agent(
- output_type=self.output_type, instructions=self.system_prompt,
**self.agent_params
- )
+ agent = self._create_agent(output_type=self.output_type,
instructions=self.system_prompt)
result = self._run_llm(context, agent, self.prompt,
usage_limits=usage_limits)
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 f5d35cb1543..ec2a4511e96 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
@@ -122,8 +122,10 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
every downstream task rather than only the direct ones, so a task whose
trigger rule would still run it is skipped too. Only takes effect when
a
review is opened. Default ``False``.
- :param agent_params: Additional keyword arguments passed to the pydantic-ai
- ``Agent`` constructor (e.g. ``retries``, ``model_settings``,
``tools``).
+ :param model_settings: Pydantic AI ``ModelSettings`` for the model call,
such as
+ ``temperature`` or ``max_tokens``. See
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+ :param output_retries: How many times the model is asked again when its
output
+ fails validation. See
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
``usage_limits`` is inherited from
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
@@ -229,10 +231,9 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
# Coerced first so a bad rendered value fails before the expensive
setup below.
usage_limits = coerce_usage_limits(self.usage_limits)
- agent = self.llm_hook.create_agent(
+ agent = self._create_agent(
output_type=output_type,
instructions=self.system_prompt,
- **self.agent_params,
)
result = self._run_llm(context, agent, self.prompt,
usage_limits=usage_limits)
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 cf176992eeb..9e61e34a3e8 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
@@ -28,8 +28,6 @@ from airflow.providers.common.ai.utils.file_analysis import
build_file_analysis_
from airflow.providers.common.ai.utils.usage import coerce_usage_limits
if TYPE_CHECKING:
- from pydantic_ai import Agent
-
from airflow.sdk import Context
@@ -55,8 +53,10 @@ class LLMFileAnalysisOperator(LLMOperator):
for how blank entries in the list are dropped.
:param system_prompt: Additional instructions appended to the built-in
file-analysis system prompt.
- :param agent_params: Additional keyword arguments passed to the pydantic-ai
- ``Agent`` constructor (e.g. ``retries``, ``model_settings``,
``tools``).
+ :param model_settings: Pydantic AI ``ModelSettings`` for the model call,
such as
+ ``temperature`` or ``max_tokens``. See
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+ :param output_retries: How many times the model is asked again when its
output
+ fails validation. See
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
:param file_path: File or prefix to analyze.
:param file_conn_id: Optional connection ID for the storage backend.
Overrides a connection embedded in ``file_path``.
@@ -163,10 +163,9 @@ class LLMFileAnalysisOperator(LLMOperator):
self.sample_rows,
)
self.log.debug("Resolved file analysis paths: %s",
request.resolved_paths)
- agent: Agent[object, Any] = self.llm_hook.create_agent(
+ agent = self._create_agent(
output_type=self.output_type,
instructions=self._build_system_prompt(),
- **self.agent_params,
)
result = self._run_llm(context, agent, request.user_content,
usage_limits=usage_limits)
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 91723596627..2e2f8d11502 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
@@ -106,7 +106,10 @@ class LLMSchemaCompareOperator(LLMOperator):
:param system_prompt: Instructions included in the LLM system prompt.
Defaults to
``DEFAULT_SYSTEM_PROMPT`` which contains cross-system type
equivalences and
severity definitions. Passing a value **replaces** the default system
prompt
- :param agent_params: Extra keyword arguments for the pydantic-ai ``Agent``.
+ :param model_settings: Pydantic AI ``ModelSettings`` for the model call,
such as
+ ``temperature`` or ``max_tokens``. See
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+ :param output_retries: How many times the model is asked again when its
output
+ fails validation. See
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
:param data_sources: List of DataSourceConfig objects, one per system.
:param db_conn_ids: Connection IDs for databases to compare (used with
``table_names``).
@@ -337,10 +340,9 @@ class LLMSchemaCompareOperator(LLMOperator):
full_system_prompt = self._build_system_prompt(schema_context)
- agent = self.llm_hook.create_agent(
+ agent = self._create_agent(
output_type=SchemaCompareResult,
instructions=full_system_prompt,
- **self.agent_params,
)
self.log.info("Running LLM schema comparison...")
result = self._run_llm(context, agent, self.prompt,
usage_limits=usage_limits)
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 28247e3eae4..931ad82f0a1 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
@@ -73,8 +73,10 @@ class LLMSQLQueryOperator(LLMOperator):
for how blank entries in the list are dropped.
:param system_prompt: Additional instructions appended to the built-in SQL
safety prompt. Use for domain-specific guidance.
- :param agent_params: Additional keyword arguments passed to the pydantic-ai
- ``Agent`` constructor (e.g. ``retries``, ``model_settings``).
+ :param model_settings: Pydantic AI ``ModelSettings`` for the model call,
such as
+ ``temperature`` or ``max_tokens``. See
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
+ :param output_retries: How many times the model is asked again when its
output
+ fails validation. See
:class:`~airflow.providers.common.ai.operators.llm.LLMOperator`.
:param db_conn_id: Connection ID for database schema introspection.
The connection must resolve to a ``DbApiHook``.
:param table_names: Tables to include in the LLM's schema context.
@@ -161,9 +163,7 @@ class LLMSQLQueryOperator(LLMOperator):
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
- )
+ agent = self._create_agent(output_type=str,
instructions=full_system_prompt)
result = self._run_llm(context, agent, self.prompt,
usage_limits=usage_limits)
sql = self._strip_llm_output(result.output,
dialect=self._resolved_dialect)
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 4d679cf7e24..6d57211e248 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
@@ -25,7 +25,7 @@ from uuid import uuid4
import pytest
from pydantic import BaseModel
from pydantic_ai import Agent
-from pydantic_ai.exceptions import UsageLimitExceeded
+from pydantic_ai.exceptions import UnexpectedModelBehavior, UsageLimitExceeded
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits
@@ -35,10 +35,14 @@ from airflow.providers.common.ai.mixins.approval import (
)
from airflow.providers.common.ai.operators import llm as llm_module
from airflow.providers.common.ai.operators.llm import DecisionPolicy,
LLMOperator
+from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator
+from airflow.providers.common.ai.operators.llm_file_analysis import
LLMFileAnalysisOperator
+from airflow.providers.common.ai.operators.llm_schema_compare import
LLMSchemaCompareOperator
+from airflow.providers.common.ai.operators.llm_sql import LLMSQLQueryOperator
from airflow.providers.common.ai.utils.logging import MODEL_NAME_XCOM_KEY
from airflow.providers.common.ai.utils.usage_budget import USAGE_BUDGET_KEY
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_3_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
try:
from airflow.sdk.serde import SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as
_CORE_WALKER
@@ -46,7 +50,11 @@ except ImportError:
_CORE_WALKER = False
from airflow.providers.common.compat.notifier import BaseNotifier
-from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, TaskDeferred
+from airflow.providers.common.compat.sdk import (
+ AirflowException,
+ AirflowOptionalProviderFeatureException,
+ TaskDeferred,
+)
from airflow.providers.standard.exceptions import HITLRejectException,
HITLTimeoutError
if AIRFLOW_V_3_3_PLUS:
@@ -99,11 +107,54 @@ class TestLLMOperator:
"model_id",
"fallback_conn_ids",
"system_prompt",
- "agent_params",
+ "model_settings",
"usage_limits",
}
assert set(LLMOperator.template_fields) == expected
+ @pytest.mark.parametrize(
+ ("operator_cls", "operator_kwargs"),
+ [
+ pytest.param(LLMOperator, {}, id="llm"),
+ pytest.param(LLMBranchOperator, {}, id="llm_branch"),
+ pytest.param(LLMSQLQueryOperator, {}, id="llm_sql"),
+ pytest.param(LLMFileAnalysisOperator, {"file_path":
"/tmp/app.log"}, id="llm_file_analysis"),
+ pytest.param(
+ LLMSchemaCompareOperator,
+ {"db_conn_ids": ["db_a", "db_b"], "table_names": ["orders"]},
+ id="llm_schema_compare",
+ ),
+ ],
+ )
+ def test_llm_operators_reject_agent_params(self, operator_cls,
operator_kwargs):
+ """No LLM operator accepts agent_params, the one argument that reached
Agent(tools=...)."""
+ # Airflow 2's BaseOperator reports an unexpected argument as
AirflowException.
+ with pytest.raises(TypeError if AIRFLOW_V_3_0_PLUS else
AirflowException, match="agent_params"):
+ operator_cls(
+ task_id="test",
+ prompt="p",
+ llm_conn_id="my_llm",
+ agent_params={"tools": []},
+ **operator_kwargs,
+ )
+
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_output_retries_bound_how_often_invalid_output_is_retried(self,
mock_hook_cls):
+ calls = []
+
+ def respond(messages, info: AgentInfo):
+ calls.append(messages)
+ return ModelResponse(parts=[TextPart(content="not a number")])
+
+ mock_hook_cls.get_hook.return_value.create_agent.side_effect = lambda
**kw: Agent(
+ FunctionModel(respond), **kw
+ )
+ op = LLMOperator(task_id="test", prompt="p", llm_conn_id="my_llm",
output_type=int, output_retries=2)
+
+ with pytest.raises(UnexpectedModelBehavior):
+ op.execute(context=MagicMock())
+ assert len(calls) == 3
+
@pytest.mark.parametrize(
"bad",
[pytest.param([], id="list"), pytest.param("0.5", id="str"),
pytest.param(5, id="int")],
@@ -139,7 +190,7 @@ class TestLLMOperator:
"What is the capital of France?", usage_limits=None,
cancellation_token=ANY
)
mock_hook_cls.get_hook.return_value.create_agent.assert_called_once_with(
- output_type=str, instructions=""
+ output_type=str, instructions="", model_settings=None, retries=None
)
mock_hook_cls.get_hook.assert_called_once_with(
"my_llm", hook_params={"model_id": None, "fallback_conn_ids": None}
@@ -294,7 +345,8 @@ class TestLLMOperator:
model_id="openai:gpt-5",
system_prompt="You are an extractor.",
output_type=Entities,
- agent_params={"retries": 3, "model_settings": {"temperature":
0.9}},
+ model_settings={"temperature": 0.9},
+ output_retries=3,
)
result = op.execute(context=MagicMock())
@@ -306,8 +358,8 @@ class TestLLMOperator:
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},
+ retries=3,
)
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
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 ac0f46870b4..8960add157e 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
@@ -67,6 +67,31 @@ class TestLLMBranchOperator:
output_type=int,
)
+ @patch.object(LLMBranchOperator, "do_branch")
+ @patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
+ def test_execute_builds_agent_with_model_settings_and_output_retries_only(
+ self, mock_hook_cls, mock_do_branch, make_mock_run_result
+ ):
+ mock_agent = MagicMock(spec=["run_sync"])
+ mock_agent.run_sync.return_value = make_mock_run_result("task_a")
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+ mock_do_branch.return_value = "task_a"
+
+ op = LLMBranchOperator(
+ task_id="test",
+ prompt="Pick a branch",
+ llm_conn_id="my_llm",
+ model_settings={"temperature": 0},
+ output_retries=2,
+ )
+ op.downstream_task_ids = {"task_a", "task_b"}
+ op.execute(MagicMock())
+
+ create_agent_kwargs =
mock_hook_cls.get_hook.return_value.create_agent.call_args.kwargs
+ assert set(create_agent_kwargs) == {"output_type", "instructions",
"model_settings", "retries"}
+ assert create_agent_kwargs["model_settings"] == {"temperature": 0}
+ assert create_agent_kwargs["retries"] == 2
+
@patch.object(LLMBranchOperator, "do_branch")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_single_branch(self, mock_hook_cls, mock_do_branch,
make_mock_run_result):
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 2a118e08764..74e2f8f2b59 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
@@ -95,13 +95,42 @@ class TestLLMFileAnalysisOperator:
"model_id",
"fallback_conn_ids",
"system_prompt",
- "agent_params",
+ "model_settings",
"usage_limits",
"file_path",
"file_conn_id",
}
assert set(LLMFileAnalysisOperator.template_fields) == expected
+ @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_builds_agent_with_model_settings_and_output_retries_only(
+ self, mock_build_request, mock_hook_cls
+ ):
+ 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("Analysis
complete")
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
+
+ op = LLMFileAnalysisOperator(
+ task_id="test",
+ prompt="Summarize the file",
+ llm_conn_id="my_llm",
+ file_path="/tmp/app.log",
+ model_settings={"temperature": 0},
+ output_retries=2,
+ )
+ op.execute(context={})
+
+ create_agent_kwargs =
mock_hook_cls.get_hook.return_value.create_agent.call_args.kwargs
+ assert set(create_agent_kwargs) == {"output_type", "instructions",
"model_settings", "retries"}
+ assert create_agent_kwargs["model_settings"] == {"temperature": 0}
+ assert create_agent_kwargs["retries"] == 2
+
@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
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 5a0f84e82c3..c8cc658bb93 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
@@ -166,11 +166,11 @@ class TestLLMSchemaCompareOperator:
data_sources=[ds, _make_ds_config(conn_id="ds2")],
context_strategy="basic",
system_prompt="additional instructions",
- agent_params={"temperature": 0.5},
+ model_settings={"temperature": 0.5},
)
assert op.context_strategy == "basic"
assert op.system_prompt == "additional instructions"
- assert op.agent_params == {"temperature": 0.5}
+ assert op.model_settings == {"temperature": 0.5}
@mock.patch("airflow.providers.common.ai.operators.llm_schema_compare.BaseHook")
def test_get_db_hook_success(self, mock_base_hook):
@@ -272,7 +272,7 @@ class TestLLMSchemaCompareOperator:
llm_conn_id="llm_conn",
db_conn_ids=["postgres_default", "snowflake_default"],
table_names=["orders"],
- agent_params={"param": "value"},
+ model_settings={"temperature": 0},
)
mock_llm_hook = mock.Mock()
@@ -290,7 +290,8 @@ class TestLLMSchemaCompareOperator:
mock_llm_hook.create_agent.assert_called_once_with(
output_type=SchemaCompareResult,
instructions="system_prompt",
- param="value",
+ model_settings={"temperature": 0},
+ retries=None,
)
mock_agent.run_sync.assert_called_once_with(
"user_prompt", usage_limits=None, cancellation_token=mock.ANY
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 d70853f6a76..3d5a696e1a5 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
@@ -209,7 +209,7 @@ class TestLLMSQLQueryOperator:
"model_id",
"fallback_conn_ids",
"system_prompt",
- "agent_params",
+ "model_settings",
"usage_limits",
"db_conn_id",
"table_names",
@@ -321,8 +321,8 @@ class TestLLMSQLQueryOperator:
assert result == "DROP TABLE users"
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
- def test_execute_passes_agent_params(self, mock_hook_cls,
make_mock_run_result):
- """agent_params inherited from LLMOperator are unpacked into
create_agent."""
+ def test_execute_passes_model_settings_and_output_retries(self,
mock_hook_cls, make_mock_run_result):
+ """model_settings and output_retries inherited from LLMOperator reach
create_agent."""
mock_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
"SELECT 1", make_mock_run_result
)
@@ -331,13 +331,14 @@ class TestLLMSQLQueryOperator:
task_id="test",
prompt="test",
llm_conn_id="my_llm",
- agent_params={"retries": 3, "model_settings": {"temperature": 0}},
+ model_settings={"temperature": 0},
+ output_retries=3,
)
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}
+ assert create_agent_call.kwargs["retries"] == 3
+ assert create_agent_call.kwargs["model_settings"] == {"temperature": 0}
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_system_prompt_appended_to_sql_instructions(self, mock_hook_cls,
make_mock_run_result):
