This is an automated email from the ASF dual-hosted git repository.
potiuk 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 bbe441e75dc Render only the returned value in TaskFlow decorators that
template it (#76244)
bbe441e75dc is described below
commit bbe441e75dca028d632cedf12f80ce5a2fd7efff
Author: Jarek Potiuk <[email protected]>
AuthorDate: Mon Nov 2 14:55:59 2026 +0100
Render only the returned value in TaskFlow decorators that template it
(#76244)
* Render only the returned value in TaskFlow decorators that template it
The decorators that build their main field from the function's return value
called render_template_fields() inside execute() to render it, which
rendered
every other template field a second time after the task runner had already
rendered them. Render only the returned value, as #71402 did, and document
how
to reference function arguments in the returned @task.llm prompt.
Generated-by: Claude Opus 5
* Make positional arguments available to the returned common.ai prompt
Only op_kwargs were merged into the context, so a prompt that referenced an
argument by name
failed with an undefined variable when the task was called positionally, as
in the docs'
summarize("...") example. The six common.ai decorators now merge the bound
call arguments,
without defaults so that a ti=None default cannot shadow the context's ti,
just before rendering
the returned prompt. The docs, the examples the page includes and the Usage
docstrings reference
the argument by name instead of formatting it in with an f-string.
The returned-command test of @task.bash runs on Airflow 3 only: on Airflow
2 the decorator is the
one in Airflow core.
Generated-by: Claude Opus 5
---
providers/cncf/kubernetes/docs/changelog.rst | 5 ++
.../cncf/kubernetes/decorators/kubernetes_cmd.py | 7 +--
.../kubernetes/decorators/test_kubernetes_cmd.py | 23 +++++++
.../decorators/test_kubernetes_commons.py | 5 +-
providers/common/ai/docs/changelog.rst | 9 +++
providers/common/ai/docs/operators/llm.rst | 15 +++++
.../providers/common/ai/decorators/agent.py | 9 ++-
.../airflow/providers/common/ai/decorators/llm.py | 11 +++-
.../providers/common/ai/decorators/llm_branch.py | 11 +++-
.../common/ai/decorators/llm_file_analysis.py | 7 ++-
.../common/ai/decorators/llm_schema_compare.py | 7 ++-
.../providers/common/ai/decorators/llm_sql.py | 7 ++-
.../common/ai/example_dags/example_llm.py | 4 +-
.../example_dags/example_llm_analysis_pipeline.py | 7 ++-
.../airflow/providers/common/ai/operators/agent.py | 3 +-
.../airflow/providers/common/ai/utils/prompt.py | 20 +++++-
.../tests/unit/common/ai/decorators/test_agent.py | 20 ++++++
.../ai/tests/unit/common/ai/decorators/test_llm.py | 72 ++++++++++++++++++++++
.../unit/common/ai/decorators/test_llm_branch.py | 17 +++++
.../common/ai/decorators/test_llm_file_analysis.py | 21 +++++++
.../ai/decorators/test_llm_schema_compare.py | 22 +++++++
.../unit/common/ai/decorators/test_llm_sql.py | 17 +++++
.../tests/unit/common/ai/operators/test_agent.py | 2 +-
providers/common/sql/docs/changelog.rst | 5 ++
.../providers/common/sql/decorators/analytics.py | 4 +-
.../airflow/providers/common/sql/decorators/sql.py | 4 +-
.../unit/common/sql/decorators/test_analytics.py | 20 ++++++
.../tests/unit/common/sql/decorators/test_sql.py | 19 ++++++
providers/standard/docs/changelog.rst | 6 ++
.../airflow/providers/standard/decorators/bash.py | 4 +-
.../tests/unit/standard/decorators/test_bash.py | 21 +++++++
31 files changed, 378 insertions(+), 26 deletions(-)
diff --git a/providers/cncf/kubernetes/docs/changelog.rst
b/providers/cncf/kubernetes/docs/changelog.rst
index d431c2a808c..a1d3c140a60 100644
--- a/providers/cncf/kubernetes/docs/changelog.rst
+++ b/providers/cncf/kubernetes/docs/changelog.rst
@@ -27,6 +27,11 @@
Changelog
---------
+.. note::
+ ``@task.kubernetes_cmd`` now renders only the command returned by the
decorated function inside
+ ``execute()``. Its other template fields are rendered once before
``execute()`` and are no
+ longer rendered a second time.
+
.. note::
In multi-team mode, a team's ``KubernetesExecutor`` now reads
``in_cluster``, ``config_file``,
``cluster_context``, ``verify_ssl``, ``ssl_ca_cert`` and
``api_client_retry_configuration`` from the
diff --git
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/decorators/kubernetes_cmd.py
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/decorators/kubernetes_cmd.py
index fb15f5097bc..f97528c62cc 100644
---
a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/decorators/kubernetes_cmd.py
+++
b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/decorators/kubernetes_cmd.py
@@ -70,16 +70,15 @@ class _KubernetesCmdDecoratedOperator(DecoratedOperator,
KubernetesPodOperator):
)
def execute(self, context: Context):
- self.render_template_fields(context)
-
- generated = self._generate_cmds(context)
+ # Render only the command the callable returned. The other template
fields were already
+ # rendered before execute().
+ generated = self.render_template(self._generate_cmds(context), context)
if self.args_only:
self.cmds = []
self.arguments = generated
else:
self.cmds = generated
self.arguments = []
- self.render_template_fields(context)
return super().execute(context)
def _generate_cmds(self, context: Context) -> list[str]:
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_cmd.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_cmd.py
index a9832240345..0b7feb1efab 100644
---
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_cmd.py
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_cmd.py
@@ -377,6 +377,29 @@ class
TestKubernetesCmdDecorator(TestKubernetesDecoratorsBase):
assert containers[0].image == f"python:{DAG_ID}"
+ def test_execute_renders_only_the_returned_command(self):
+ """Fields rendered before execute() are not rendered again; only the
returned value is."""
+ with self.dag_maker(params={"message": "{{ dag.dag_id }}"}):
+
+ @task.kubernetes_cmd(
+ image="python:3.10-slim-buster",
+ in_cluster=False,
+ cluster_context="default",
+ config_file="/tmp/fake_file",
+ namespace="default",
+ env_vars={"MESSAGE": "{{ params.message }}"},
+ )
+ def hello():
+ return ["echo", "{{ ti.task_id }}"]
+
+ hello_task = hello()
+
+ self.execute_task(hello_task)
+
+ containers =
self.mock_create_pod.call_args.kwargs["pod"].spec.containers
+ assert containers[0].command == ["echo", "hello"]
+ assert [(env.name, env.value) for env in containers[0].env] ==
[("MESSAGE", "{{ dag.dag_id }}")]
+
def test_airflow_skip(self):
"""Test that the operator is skipped if the task is skipped"""
with self.dag_maker:
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_commons.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_commons.py
index 16c4ed65c2e..37f70ab210c 100644
---
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_commons.py
+++
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/decorators/test_kubernetes_commons.py
@@ -150,7 +150,10 @@ class TestKubernetesDecoratorsBase:
session = self.dag_maker.session
dag_run =
self.dag_maker.create_dagrun(run_id=f"k8s_decorator_test_{DEFAULT_DATE.date()}")
ti = dag_run.get_task_instance(task.operator.task_id, session=session)
- return_val = task.operator.execute(context=get_template_context(ti,
task.operator, session=session))
+ context = get_template_context(ti, task.operator, session=session)
+ # The task runner renders the template fields before it calls
execute().
+ task.operator.render_template_fields(context)
+ return_val = task.operator.execute(context=context)
return ti, return_val
diff --git a/providers/common/ai/docs/changelog.rst
b/providers/common/ai/docs/changelog.rst
index d94f9540cde..fdeb728db53 100644
--- a/providers/common/ai/docs/changelog.rst
+++ b/providers/common/ai/docs/changelog.rst
@@ -25,6 +25,15 @@
Changelog
---------
+.. note::
+ ``@task.llm``, ``@task.agent``, ``@task.llm_sql``, ``@task.llm_branch``,
+ ``@task.llm_file_analysis`` and ``@task.llm_schema_compare`` now render only
the prompt returned
+ by the decorated function inside ``execute()``. Their other template fields
(``system_prompt``,
+ ``op_args``, ``op_kwargs``, ...) are rendered once before ``execute()`` and
are no longer
+ rendered a second time. The returned prompt is still rendered as a Jinja
template, and every
+ argument the task is called with, positional ones included, can be
referenced in it by name;
+ see the ``@task.llm`` documentation for how to include text from function
arguments in it.
+
1.1.0
.....
diff --git a/providers/common/ai/docs/operators/llm.rst
b/providers/common/ai/docs/operators/llm.rst
index dccdbf04c65..a4cbcd62b01 100644
--- a/providers/common/ai/docs/operators/llm.rst
+++ b/providers/common/ai/docs/operators/llm.rst
@@ -206,6 +206,21 @@ With structured output:
:start-after: [START howto_decorator_llm_structured]
:end-before: [END howto_decorator_llm_structured]
+The returned prompt is rendered as a Jinja template, like the operator's
``prompt``
+argument, so ``{{ ds }}`` and other template variables work in it. Text you
insert into
+the returned string yourself, for example with an f-string, is rendered as
part of the
+template too. To include text that comes from a function argument, such as a
document or
+an upstream task's output, reference the argument by name in the template
instead of
+formatting it into the string, as the examples above do with ``{{ text }}``.
Every argument
+the task is called with, positionally or by keyword, is available when the
prompt is
+rendered, and its value is inserted as it is. A parameter left at its default
value is not
+available; pass it explicitly to reference it:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_llm_analysis_pipeline.py
+ :language: python
+ :start-after: [START howto_decorator_llm_prompt_argument]
+ :end-before: [END howto_decorator_llm_prompt_argument]
+
Multimodal prompts
^^^^^^^^^^^^^^^^^^
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/decorators/agent.py
b/providers/common/ai/src/airflow/providers/common/ai/decorators/agent.py
index d2a92e07668..39070e90cb1 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/decorators/agent.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/decorators/agent.py
@@ -28,6 +28,7 @@ from collections.abc import Callable, Collection, Mapping,
Sequence
from typing import TYPE_CHECKING, Any, ClassVar
from airflow.providers.common.ai.operators.agent import AgentOperator
+from airflow.providers.common.ai.utils.prompt import call_arguments
from airflow.providers.common.ai.utils.validation import (
reject_sequence_with_unsupported_feature,
validate_prompt,
@@ -98,7 +99,11 @@ class _AgentDecoratedOperator(DecoratedOperator,
AgentOperator):
feature_enabled=self.enable_hitl_review,
)
- self.render_template_fields(context)
+ # Render only the prompt the callable returned. The other template
fields were already
+ # rendered before execute().
+ # The arguments of the call, positional ones too, can be referenced by
name in the prompt.
+ context_merge(context, call_arguments(self.python_callable,
self.op_args, self.op_kwargs))
+ self.prompt = self.render_template(self.prompt, context)
return AgentOperator.execute(self, context)
@@ -121,7 +126,7 @@ def agent_task(
toolsets=[SQLToolset(db_conn_id="postgres_default")],
)
def analyze(question: str):
- return f"Answer: {question}"
+ return "Answer: {{ question }}"
:param python_callable: Function to decorate.
"""
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm.py
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm.py
index ca18d5225b5..32514c62513 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm.py
@@ -30,6 +30,7 @@ from collections.abc import Callable, Collection, Mapping,
Sequence
from typing import TYPE_CHECKING, Any, ClassVar
from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.ai.utils.prompt import call_arguments
from airflow.providers.common.ai.utils.validation import validate_prompt
from airflow.providers.common.compat.sdk import (
SET_DURING_EXECUTION,
@@ -91,7 +92,11 @@ class _LLMDecoratedOperator(DecoratedOperator, LLMOperator):
validate_prompt(self.prompt, decorator_name="@task.llm")
- self.render_template_fields(context)
+ # Render only the prompt the callable returned. The other template
fields were already
+ # rendered before execute().
+ # The arguments of the call, positional ones too, can be referenced by
name in the prompt.
+ context_merge(context, call_arguments(self.python_callable,
self.op_args, self.op_kwargs))
+ self.prompt = self.render_template(self.prompt, context)
return LLMOperator.execute(self, context)
@@ -113,7 +118,7 @@ def llm_task(
system_prompt="Summarize concisely.",
)
def summarize(text: str):
- return f"Summarize this article: {text}"
+ return "Summarize this article: {{ text }}"
With structured output::
@@ -123,7 +128,7 @@ def llm_task(
output_type=Entities,
)
def extract(text: str):
- return f"Extract entities from: {text}"
+ return "Extract entities from: {{ text }}"
:param python_callable: Function to decorate.
"""
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_branch.py
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_branch.py
index 4686a9c68fe..9ca53eacdcc 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_branch.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_branch.py
@@ -28,6 +28,7 @@ from collections.abc import Callable, Collection, Mapping,
Sequence
from typing import TYPE_CHECKING, Any, ClassVar
from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator
+from airflow.providers.common.ai.utils.prompt import call_arguments
from airflow.providers.common.ai.utils.validation import validate_prompt
from airflow.providers.common.compat.sdk import (
SET_DURING_EXECUTION,
@@ -90,7 +91,11 @@ class _LLMBranchDecoratedOperator(DecoratedOperator,
LLMBranchOperator):
validate_prompt(self.prompt, decorator_name="@task.llm_branch")
- self.render_template_fields(context)
+ # Render only the prompt the callable returned. The other template
fields were already
+ # rendered before execute().
+ # The arguments of the call, positional ones too, can be referenced by
name in the prompt.
+ context_merge(context, call_arguments(self.python_callable,
self.op_args, self.op_kwargs))
+ self.prompt = self.render_template(self.prompt, context)
return LLMBranchOperator.execute(self, context)
@@ -112,7 +117,7 @@ def llm_branch_task(
system_prompt="Route support tickets to the right team.",
)
def route_ticket(message: str):
- return f"Route this ticket: {message}"
+ return "Route this ticket: {{ message }}"
With multiple branches::
@@ -122,7 +127,7 @@ def llm_branch_task(
allow_multiple_branches=True,
)
def classify(text: str):
- return f"Classify this text: {text}"
+ return "Classify this text: {{ text }}"
:param python_callable: Function to decorate.
"""
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_file_analysis.py
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_file_analysis.py
index 1e3f6974826..3c78ad8782e 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_file_analysis.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_file_analysis.py
@@ -22,6 +22,7 @@ from collections.abc import Callable, Collection, Mapping,
Sequence
from typing import TYPE_CHECKING, Any, ClassVar
from airflow.providers.common.ai.operators.llm_file_analysis import
LLMFileAnalysisOperator
+from airflow.providers.common.ai.utils.prompt import call_arguments
from airflow.providers.common.compat.sdk import (
SET_DURING_EXECUTION,
DecoratedOperator,
@@ -90,7 +91,11 @@ class _LLMFileAnalysisDecoratedOperator(DecoratedOperator,
LLMFileAnalysisOperat
"The returned value from the @task.llm_file_analysis callable
must be a non-empty string."
)
- self.render_template_fields(context)
+ # Render only the prompt the callable returned. The other template
fields were already
+ # rendered before execute().
+ # The arguments of the call, positional ones too, can be referenced by
name in the prompt.
+ context_merge(context, call_arguments(self.python_callable,
self.op_args, self.op_kwargs))
+ self.prompt = self.render_template(self.prompt, context)
return LLMFileAnalysisOperator.execute(self, context)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_schema_compare.py
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_schema_compare.py
index 5440aa4b5b0..09a1ba29c3e 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_schema_compare.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_schema_compare.py
@@ -28,6 +28,7 @@ from collections.abc import Callable, Collection, Mapping,
Sequence
from typing import TYPE_CHECKING, Any, ClassVar
from airflow.providers.common.ai.operators.llm_schema_compare import
LLMSchemaCompareOperator
+from airflow.providers.common.ai.utils.prompt import call_arguments
from airflow.providers.common.ai.utils.validation import validate_prompt
from airflow.providers.common.compat.sdk import (
SET_DURING_EXECUTION,
@@ -90,7 +91,11 @@ class _LLMSchemaCompareDecoratedOperator(DecoratedOperator,
LLMSchemaCompareOper
validate_prompt(self.prompt, decorator_name="@task.llm_schema_compare")
- self.render_template_fields(context)
+ # Render only the prompt the callable returned. The other template
fields were already
+ # rendered before execute().
+ # The arguments of the call, positional ones too, can be referenced by
name in the prompt.
+ context_merge(context, call_arguments(self.python_callable,
self.op_args, self.op_kwargs))
+ self.prompt = self.render_template(self.prompt, context)
return LLMSchemaCompareOperator.execute(self, context)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_sql.py
b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_sql.py
index 7c5082a9f44..9ebb1b6a0b6 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_sql.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/decorators/llm_sql.py
@@ -28,6 +28,7 @@ from collections.abc import Callable, Collection, Mapping,
Sequence
from typing import TYPE_CHECKING, Any, ClassVar
from airflow.providers.common.ai.operators.llm_sql import LLMSQLQueryOperator
+from airflow.providers.common.ai.utils.prompt import call_arguments
from airflow.providers.common.ai.utils.validation import validate_prompt
from airflow.providers.common.compat.sdk import (
SET_DURING_EXECUTION,
@@ -89,7 +90,11 @@ class _LLMSQLDecoratedOperator(DecoratedOperator,
LLMSQLQueryOperator):
validate_prompt(self.prompt, decorator_name="@task.llm_sql")
- self.render_template_fields(context)
+ # Render only the prompt the callable returned. The other template
fields were already
+ # rendered before execute().
+ # The arguments of the call, positional ones too, can be referenced by
name in the prompt.
+ context_merge(context, call_arguments(self.python_callable,
self.op_args, self.op_kwargs))
+ self.prompt = self.render_template(self.prompt, context)
# Call LLMSQLQueryOperator.execute directly, not super().execute(),
# because we need to skip DecoratedOperator.execute — the callable
# invocation is already handled above.
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 25139db825c..4e2f764dbc7 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
@@ -98,7 +98,7 @@ example_llm_operator_model_settings()
def example_llm_decorator():
@task.llm(llm_conn_id="pydanticai_default", system_prompt="Summarize
concisely.")
def summarize(text: str):
- return f"Summarize this article: {text}"
+ return "Summarize this article: {{ text }}"
summarize("Apache Airflow is a platform for programmatically authoring...")
@@ -117,7 +117,7 @@ def example_llm_decorator_structured():
output_type=Entities,
)
def extract(text: str):
- return f"Extract entities from: {text}"
+ return "Extract entities from: {{ text }}"
extract("Alice visited Paris and met Bob in London.")
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_analysis_pipeline.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_analysis_pipeline.py
index b13dc811ae9..dec163533ff 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_analysis_pipeline.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_llm_analysis_pipeline.py
@@ -69,8 +69,13 @@ def example_llm_analysis_pipeline():
),
output_type=TicketAnalysis,
)
+ # [START howto_decorator_llm_prompt_argument]
def analyze_ticket(ticket: str):
- return f"Analyze this support ticket:\n\n{ticket}"
+ # The returned prompt is rendered as a Jinja template. Reference the
argument by name so
+ # the ticket text is inserted as it is, rather than formatting it into
the template.
+ return "Analyze this support ticket:\n\n{{ ticket }}"
+
+ # [END howto_decorator_llm_prompt_argument]
@task
def store_results(analyses: list[TicketAnalysis]):
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 37068585f70..590ee28301b 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
@@ -783,8 +783,7 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator,
HITLReviewMixin, Usa
rendered = copy.copy(toolset)
self._do_render_template_fields(rendered, fields, context,
jinja_env, seen_oids)
# The rendered connection is recorded nowhere else, so this line
is the audit trail
- # of which connection this task instance's agent was given.
@task.agent renders a
- # second time, when the id no longer changes, so this logs once
per task instance.
+ # of which connection this task instance's agent was given.
if rendered.id != toolset.id:
self.log.info("Rendered toolset %s", rendered.id)
return rendered
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/prompt.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/prompt.py
index a3266b6814a..6409addd635 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/prompt.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/prompt.py
@@ -16,7 +16,9 @@
# under the License.
from __future__ import annotations
-from collections.abc import Sequence
+import inspect
+from collections.abc import Callable, Collection, Mapping, Sequence
+from typing import Any
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from pydantic_ai.messages import (
@@ -32,6 +34,22 @@ from pydantic_ai.messages import (
)
+def call_arguments(
+ python_callable: Callable, op_args: Collection[Any], op_kwargs:
Mapping[str, Any]
+) -> dict[str, Any]:
+ """
+ Return the arguments a decorated function was called with, by parameter
name.
+
+ Positional arguments are included, so a returned prompt can reference ``{{
text }}`` for
+ ``summarize("...")`` as well as for ``summarize(text="...")``. Parameters
left at their default
+ are not: a default such as ``ti=None`` must not shadow the ``ti`` of the
task context.
+ """
+ try:
+ return dict(inspect.signature(python_callable).bind_partial(*op_args,
**op_kwargs).arguments)
+ except (TypeError, ValueError):
+ return dict(op_kwargs)
+
+
def describe_prompt(prompt: str | Sequence[UserContent]) -> str:
if isinstance(prompt, str):
return prompt
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 a1c4a5330fa..00e017f068c 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
@@ -26,6 +26,7 @@ from pydantic_ai.toolsets.function import FunctionToolset
from airflow.providers.common.ai.decorators.agent import
_AgentDecoratedOperator
from airflow.providers.common.ai.toolsets.logging import ToolLoggingCapability
+from airflow.providers.common.ai.toolsets.sql import SQLToolset
from airflow.providers.common.ai.utils.prompt_cache import PromptCaching
from airflow.providers.common.ai.utils.toolset_base import MaskingToolset
@@ -158,6 +159,25 @@ class TestAgentDecoratedOperator:
"Analyze revenue trends", usage_limits=None, run_id="ti-1",
cancellation_token=ANY, usage=ANY
)
+
@patch("airflow.providers.common.ai.operators.agent.AgentOperator.execute",
autospec=True)
+ def test_execute_renders_only_the_returned_prompt(self, mock_execute):
+ """Fields rendered before execute() are not rendered again; only the
returned value is."""
+ op = _AgentDecoratedOperator(
+ task_id="test",
+ python_callable=lambda: "Summarize {{ topic }}",
+ llm_conn_id="my_llm",
+ system_prompt="{{ note }}",
+ toolsets=[SQLToolset(db_conn_id="{{ note }}")],
+ )
+ context = {"topic": "logs", "note": "{{ other }}", "other": "rendered
twice"}
+
+ op.render_template_fields(context)
+ op.execute(context=context)
+
+ assert op.prompt == "Summarize logs"
+ assert op.system_prompt == "{{ other }}"
+ assert op.toolsets[0].id == "sql-{{ other }}"
+
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
def test_execute_passes_toolsets_through(self, mock_hook_cls,
make_mock_run_result):
"""Toolsets passed to the decorator are forwarded to the agent."""
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 a067981dc7f..4c7db1a248d 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
@@ -144,6 +144,78 @@ class TestLLMDecoratedOperator:
"Summarize quantum computing", usage=ANY, usage_limits=None,
run_id=ANY, cancellation_token=ANY
)
+ @patch("airflow.providers.common.ai.operators.llm.LLMOperator.execute",
autospec=True)
+ def test_execute_renders_only_the_returned_prompt(self, mock_execute):
+ """Fields rendered before execute() are not rendered again; only the
returned prompt is."""
+ op = _LLMDecoratedOperator(
+ task_id="test",
+ python_callable=lambda: "Summarize {{ topic }}",
+ llm_conn_id="my_llm",
+ system_prompt="{{ note }}",
+ )
+ context = {"topic": "logs", "note": "{{ other }}", "other": "rendered
twice"}
+
+ op.render_template_fields(context)
+ op.execute(context=context)
+
+ assert op.prompt == "Summarize logs"
+ assert op.system_prompt == "{{ other }}"
+
+ @patch("airflow.providers.common.ai.operators.llm.LLMOperator.execute",
autospec=True)
+ def test_argument_referenced_in_returned_prompt_is_inserted_as_is(self,
mock_execute):
+ """An argument referenced by name in the returned prompt is inserted
without being rendered."""
+
+ def analyze_ticket(ticket):
+ return "Analyze this support ticket:\n\n{{ ticket }}"
+
+ op = _LLMDecoratedOperator(
+ task_id="test",
+ python_callable=analyze_ticket,
+ llm_conn_id="my_llm",
+ # As an upstream task's output, already resolved when execute()
runs.
+ op_kwargs={"ticket": "Total is {{ 7 * 7 }}"},
+ )
+
+ op.execute(context={"task_instance": MagicMock()})
+
+ assert op.prompt == "Analyze this support ticket:\n\nTotal is {{ 7 * 7
}}"
+
+ @patch("airflow.providers.common.ai.operators.llm.LLMOperator.execute",
autospec=True)
+ def test_positional_argument_referenced_in_returned_prompt(self,
mock_execute):
+ """A positional call, such as ``summarize("...")``, makes the argument
available by name too."""
+
+ def summarize(text):
+ return "Summarize this article: {{ text }}"
+
+ op = _LLMDecoratedOperator(
+ task_id="test",
+ python_callable=summarize,
+ llm_conn_id="my_llm",
+ op_args=["Total is {{ 7 * 7 }}"],
+ )
+
+ op.execute(context={"task_instance": MagicMock()})
+
+ assert op.prompt == "Summarize this article: Total is {{ 7 * 7 }}"
+
+ @patch("airflow.providers.common.ai.operators.llm.LLMOperator.execute",
autospec=True)
+ def test_default_argument_does_not_shadow_the_context(self, mock_execute):
+ """A parameter left at its default is not merged, so ``ti=None`` keeps
the context's ``ti``."""
+
+ def summarize(text, ti=None):
+ return "{{ text }} from {{ ti }}"
+
+ op = _LLMDecoratedOperator(
+ task_id="test",
+ python_callable=summarize,
+ llm_conn_id="my_llm",
+ op_args=["logs"],
+ )
+
+ op.execute(context={"task_instance": MagicMock(), "ti": "the real ti"})
+
+ assert op.prompt == "logs from the real ti"
+
@pytest.mark.skipif(not AIRFLOW_V_3_1_PLUS, reason="require_approval needs
Airflow >= 3.1.0")
@pytest.mark.parametrize(
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 747bbd46717..8d0ccf48618 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
@@ -131,3 +131,20 @@ class TestLLMBranchDecoratedOperator:
op.execute(context={"task_instance": MagicMock()})
assert op.prompt == "Route this billing ticket"
+
+
@patch("airflow.providers.common.ai.operators.llm_branch.LLMBranchOperator.execute",
autospec=True)
+ def test_execute_renders_only_the_returned_prompt(self, mock_execute):
+ """Fields rendered before execute() are not rendered again; only the
returned value is."""
+ op = _LLMBranchDecoratedOperator(
+ task_id="test",
+ python_callable=lambda: "Summarize {{ topic }}",
+ llm_conn_id="my_llm",
+ system_prompt="{{ note }}",
+ )
+ context = {"topic": "logs", "note": "{{ other }}", "other": "rendered
twice"}
+
+ op.render_template_fields(context)
+ op.execute(context=context)
+
+ assert op.prompt == "Summarize logs"
+ assert op.system_prompt == "{{ other }}"
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 5903138f817..d7e2dda3fba 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
@@ -124,3 +124,24 @@ class TestLLMFileAnalysisDecoratedOperator:
mock_agent.run_sync.assert_called_once_with(
"prepared prompt", usage=ANY, usage_limits=None,
cancellation_token=ANY
)
+
+ @patch(
+
"airflow.providers.common.ai.operators.llm_file_analysis.LLMFileAnalysisOperator.execute",
+ autospec=True,
+ )
+ def test_execute_renders_only_the_returned_prompt(self, mock_execute):
+ """Fields rendered before execute() are not rendered again; only the
returned prompt is."""
+ op = _LLMFileAnalysisDecoratedOperator(
+ task_id="test",
+ python_callable=lambda: "Summarize {{ topic }}",
+ llm_conn_id="my_llm",
+ system_prompt="{{ note }}",
+ file_path="/tmp/app.log",
+ )
+ context = {"topic": "logs", "note": "{{ other }}", "other": "rendered
twice"}
+
+ op.render_template_fields(context)
+ op.execute(context=context)
+
+ assert op.prompt == "Summarize logs"
+ assert op.system_prompt == "{{ other }}"
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
index 73adba3f2d3..35c47fd5ffa 100644
---
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
+++
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
@@ -152,3 +152,25 @@ class TestLLMSchemaCompareDecoratedOperator:
op.execute(context={"task_instance": MagicMock()})
assert op.prompt == "Compare schemas for production environment"
+
+ @patch(
+
"airflow.providers.common.ai.operators.llm_schema_compare.LLMSchemaCompareOperator.execute",
+ autospec=True,
+ )
+ def test_execute_renders_only_the_returned_prompt(self, mock_execute):
+ """Fields rendered before execute() are not rendered again; only the
returned value is."""
+ op = _LLMSchemaCompareDecoratedOperator(
+ task_id="test",
+ python_callable=lambda: "Summarize {{ topic }}",
+ llm_conn_id="my_llm",
+ system_prompt="{{ note }}",
+ db_conn_ids=["postgres_default", "snowflake_default"],
+ table_names=["test_table"],
+ )
+ context = {"topic": "logs", "note": "{{ other }}", "other": "rendered
twice"}
+
+ op.render_template_fields(context)
+ op.execute(context=context)
+
+ assert op.prompt == "Summarize logs"
+ assert op.system_prompt == "{{ other }}"
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 aed2f559586..b15b08bf89d 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
@@ -114,3 +114,20 @@ class TestLLMSQLDecoratedOperator:
op.execute(context={"task_instance": MagicMock()})
assert op.prompt == "Get all rows from users"
+
+
@patch("airflow.providers.common.ai.operators.llm_sql.LLMSQLQueryOperator.execute",
autospec=True)
+ def test_execute_renders_only_the_returned_prompt(self, mock_execute):
+ """Fields rendered before execute() are not rendered again; only the
returned value is."""
+ op = _LLMSQLDecoratedOperator(
+ task_id="test",
+ python_callable=lambda: "Summarize {{ topic }}",
+ llm_conn_id="my_llm",
+ system_prompt="{{ note }}",
+ )
+ context = {"topic": "logs", "note": "{{ other }}", "other": "rendered
twice"}
+
+ op.render_template_fields(context)
+ op.execute(context=context)
+
+ assert op.prompt == "Summarize logs"
+ assert op.system_prompt == "{{ other }}"
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 74c97a553c9..877edcaeded 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
@@ -481,7 +481,7 @@ class TestAgentOperatorToolsetTemplating:
assert "Rendered toolset sql-tenant_acme" in caplog.messages
def test_rendering_twice_logs_once(self, caplog):
- """@task.agent renders a second time; by then the id no longer
changes."""
+ """On a second render the id no longer changes."""
op = AgentOperator(
task_id="t",
prompt="p",
diff --git a/providers/common/sql/docs/changelog.rst
b/providers/common/sql/docs/changelog.rst
index e4e0b4be499..9de154a1e82 100644
--- a/providers/common/sql/docs/changelog.rst
+++ b/providers/common/sql/docs/changelog.rst
@@ -25,6 +25,11 @@
Changelog
---------
+.. note::
+ ``@task.sql`` and ``@task.analytics`` now render only the SQL returned by
the decorated function
+ inside ``execute()``. Their other template fields (``parameters``,
``op_args``, ``op_kwargs``, ...)
+ are rendered once before ``execute()`` and are no longer rendered a second
time.
+
.. warning::
The deprecated ``DbApiHook.get_pandas_df`` and
``DbApiHook.get_pandas_df_by_chunks`` methods have been removed. Use
``DbApiHook.get_df(..., df_type="pandas")`` and
``DbApiHook.get_df_by_chunks(..., df_type="pandas")`` instead.
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/decorators/analytics.py
b/providers/common/sql/src/airflow/providers/common/sql/decorators/analytics.py
index 1c297ce8254..4923ab6b94f 100644
---
a/providers/common/sql/src/airflow/providers/common/sql/decorators/analytics.py
+++
b/providers/common/sql/src/airflow/providers/common/sql/decorators/analytics.py
@@ -117,7 +117,9 @@ class _AnalyticsDecoratedOperator(DecoratedOperator,
AnalyticsOperator):
# AnalyticsOperator expects queries as a list of strings
self.queries = [result] if isinstance(result, str) else result
- self.render_template_fields(context)
+ # Render only the queries the callable returned. The other template
fields were already
+ # rendered before execute().
+ self.queries = self.render_template(self.queries, context)
return AnalyticsOperator.execute(self, context)
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/decorators/sql.py
b/providers/common/sql/src/airflow/providers/common/sql/decorators/sql.py
index 72f9b48da7f..7f800a80d44 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/decorators/sql.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/decorators/sql.py
@@ -110,7 +110,9 @@ class _SQLDecoratedOperator(DecoratedOperator,
SQLExecuteQueryOperator):
"string or a non-empty list of non-empty strings."
)
- self.render_template_fields(context)
+ # Render only the SQL the callable returned. The other template fields
were already
+ # rendered before execute().
+ self.sql = self.render_template(self.sql, context)
return super().execute(context)
diff --git
a/providers/common/sql/tests/unit/common/sql/decorators/test_analytics.py
b/providers/common/sql/tests/unit/common/sql/decorators/test_analytics.py
index 36e964bbece..c11e019d288 100644
--- a/providers/common/sql/tests/unit/common/sql/decorators/test_analytics.py
+++ b/providers/common/sql/tests/unit/common/sql/decorators/test_analytics.py
@@ -100,6 +100,26 @@ class TestAnalyticsDecoratedOperator:
with pytest.raises(TypeError, match="non-empty string"):
op.execute(context={})
+ @patch(
+
"airflow.providers.common.sql.operators.analytics.AnalyticsOperator.execute",
+ autospec=True,
+ )
+ def test_execute_renders_only_the_returned_queries(self, mock_execute):
+ """Fields rendered before execute() are not rendered again; only the
returned value is."""
+ op = _AnalyticsDecoratedOperator(
+ task_id="test",
+ python_callable=lambda: "SELECT * FROM {{ table }}",
+ datasource_configs=DATASOURCE_CONFIGS,
+ result_output_format="{{ output_format }}",
+ )
+ context = {"table": "users_data", "output_format": "{{ other }}",
"other": "rendered twice"}
+
+ op.render_template_fields(context)
+ op.execute(context=context)
+
+ assert op.queries == ["SELECT * FROM users_data"]
+ assert op.result_output_format == "{{ other }}"
+
@patch(
"airflow.providers.common.sql.operators.analytics.AnalyticsOperator.execute",
autospec=True,
diff --git a/providers/common/sql/tests/unit/common/sql/decorators/test_sql.py
b/providers/common/sql/tests/unit/common/sql/decorators/test_sql.py
index 1c56a833f23..37bbbbd2686 100644
--- a/providers/common/sql/tests/unit/common/sql/decorators/test_sql.py
+++ b/providers/common/sql/tests/unit/common/sql/decorators/test_sql.py
@@ -120,6 +120,25 @@ class TestSqlDecorator:
assert sql_task.operator.sql == "SELECT 1;"
+
@mock.patch("airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator.get_db_hook",
MagicMock())
+ def test_execute_renders_only_the_returned_sql(self):
+ """Fields rendered before execute() are not rendered again; only the
returned value is."""
+ with self.dag_maker:
+
+ @task.sql(conn_id=conn_id, parameters={"name": "{{ name }}"})
+ def sql():
+ return "SELECT :name FROM {{ table }}"
+
+ sql_task = sql()
+
+ op = sql_task.operator
+ context = {"table": "users", "name": "{{ other }}", "other": "rendered
twice"}
+ op.render_template_fields(context)
+ op.execute(context=context)
+
+ assert op.sql == "SELECT :name FROM users"
+ assert op.parameters == {"name": "{{ other }}"}
+
@mock.patch("airflow.providers.common.sql.operators.sql.SQLExecuteQueryOperator.get_db_hook",
MagicMock())
@pytest.mark.parametrize("sql_query", ["", 1, None, [""], ["SELECT 1;",
""]])
def test_invalid_sql_query(self, sql_query):
diff --git a/providers/standard/docs/changelog.rst
b/providers/standard/docs/changelog.rst
index 678e7ae4657..febfca0589f 100644
--- a/providers/standard/docs/changelog.rst
+++ b/providers/standard/docs/changelog.rst
@@ -35,6 +35,12 @@
Changelog
---------
+.. note::
+ ``@task.bash`` now renders only the command returned by the decorated
function inside
+ ``execute()``. Its other template fields (``env``, ``cwd``, ``op_args``,
``op_kwargs``, ...) are
+ rendered once before ``execute()``, as for ``BashOperator``, and are no
longer rendered a
+ second time.
+
.. warning::
The deprecated ``TimeDeltaSensorAsync`` class has been removed. Use
``TimeDeltaSensor(..., deferrable=True)`` instead.
diff --git
a/providers/standard/src/airflow/providers/standard/decorators/bash.py
b/providers/standard/src/airflow/providers/standard/decorators/bash.py
index 77fb9861aac..c92d7dc6e9a 100644
--- a/providers/standard/src/airflow/providers/standard/decorators/bash.py
+++ b/providers/standard/src/airflow/providers/standard/decorators/bash.py
@@ -82,7 +82,9 @@ class _BashDecoratedOperator(DecoratedOperator, BashOperator):
if not isinstance(self.bash_command, str) or self.bash_command.strip()
== "":
raise TypeError("The returned value from the TaskFlow callable
must be a non-empty string.")
- self.render_template_fields(context)
+ # Render only the command the callable returned. The other template
fields were already
+ # rendered before execute().
+ self.bash_command = self.render_template(self.bash_command, context)
return super().execute(context)
diff --git a/providers/standard/tests/unit/standard/decorators/test_bash.py
b/providers/standard/tests/unit/standard/decorators/test_bash.py
index 19e3f9cae14..48ae0cc47f3 100644
--- a/providers/standard/tests/unit/standard/decorators/test_bash.py
+++ b/providers/standard/tests/unit/standard/decorators/test_bash.py
@@ -581,3 +581,24 @@ class TestBashDecorator:
op = render_template_fields(ti, task_arg.operator, context=context)
result = op.execute(context=context)
assert result == "test_templated_fields_task"
+
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_0_PLUS, reason="On Airflow 2, @task.bash is the
decorator in Airflow core"
+ )
+ def test_execute_renders_only_the_returned_command(self):
+ """Fields rendered before execute() are not rendered again; only the
returned value is."""
+ with self.dag_maker:
+
+ @task.bash(env={"message": "{{ message }}"})
+ def bash():
+ return 'echo "{{ greeting }} $message"'
+
+ bash_task = bash()
+
+ op = bash_task.operator
+ context = {"greeting": "hello", "message": "{{ other }}", "other":
"rendered twice"}
+ op.render_template_fields(context)
+
+ assert op.execute(context=context) == "hello {{ other }}"
+ assert op.bash_command == 'echo "hello $message"'
+ assert op.env == {"message": "{{ other }}"}
