This is an automated email from the ASF dual-hosted git repository.
Lee-W 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 f24e8fe6551 Validate template fields after rendering in common.ai
operators (#70338)
f24e8fe6551 is described below
commit f24e8fe655166c8258ec4ec51e1361635a152ba6
Author: Stefan Wang <[email protected]>
AuthorDate: Fri Jul 24 00:39:42 2026 -0700
Validate template fields after rendering in common.ai operators (#70338)
Signed-off-by: 1fanwang <[email protected]>
---
.../src/airflow/providers/common/ai/operators/agent.py | 13 ++++++-------
.../providers/common/ai/operators/document_loader.py | 18 +++++++++---------
.../ai/tests/unit/common/ai/operators/test_agent.py | 18 +++++++++++-------
.../unit/common/ai/operators/test_document_loader.py | 12 ++++++++----
scripts/ci/prek/validate_operators_init_exemptions.txt | 2 --
5 files changed, 34 insertions(+), 29 deletions(-)
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 f83a58f0e2a..5afbe3909b5 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
@@ -293,13 +293,6 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
# replay. Reject the combination rather than silently
mis-replaying.
raise ValueError("durable=True and code_mode=True cannot be used
together.")
- if message_history is not None and enable_hitl_review:
- # The post-review transcript is not recoverable today
(run_hitl_review
- # returns only the final string), so emitting the pre-review
transcript
- # would silently drop the human-approved turns. Block until HITL
can
- # surface the final message history.
- raise ValueError("message_history and enable_hitl_review=True
cannot be used together.")
-
self.enable_hitl_review = enable_hitl_review
self.max_hitl_iterations = max_hitl_iterations
self.hitl_timeout = hitl_timeout
@@ -424,6 +417,12 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
)
def execute(self, context: Context) -> Any:
+ # message_history is a template field; validate the combination after
rendering.
+ if self.message_history is not None and self.enable_hitl_review:
+ # run_hitl_review returns only the final string, so the pre-review
transcript would drop
+ # the human-approved turns. Block until HITL can surface the final
message history.
+ raise ValueError("message_history and enable_hitl_review=True
cannot be used together.")
+
if self.enable_hitl_review and not isinstance(self.prompt, str):
raise TypeError(
f"{type(self).__name__}: enable_hitl_review=True is not
supported "
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py
index 1c8c6a935b3..c59848c144f 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py
@@ -136,13 +136,6 @@ class DocumentLoaderOperator(BaseOperator):
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
- if source_path is not None and source_bytes is not None:
- raise ValueError("Provide exactly one of 'source_path' or
'source_bytes', not both.")
- if source_path is None and source_bytes is None:
- raise ValueError("Provide exactly one of 'source_path' or
'source_bytes'.")
- if source_bytes is not None and file_type is None:
- raise ValueError("'file_type' is required when using
'source_bytes' (e.g. '.pdf').")
-
self.source_path = source_path
self.source_conn_id = source_conn_id
self.source_bytes = source_bytes
@@ -155,12 +148,19 @@ class DocumentLoaderOperator(BaseOperator):
self.json_text_field = json_text_field
def execute(self, context: Context) -> list[dict[str, Any]]:
+ # source_path/file_type are template fields; validate after rendering,
not in __init__.
+ if self.source_path is not None and self.source_bytes is not None:
+ raise ValueError("Provide exactly one of 'source_path' or
'source_bytes', not both.")
+ if self.source_path is None and self.source_bytes is None:
+ raise ValueError("Provide exactly one of 'source_path' or
'source_bytes'.")
+ if self.source_bytes is not None and self.file_type is None:
+ raise ValueError("'file_type' is required when using
'source_bytes' (e.g. '.pdf').")
if self.source_bytes is not None:
- assert self.file_type is not None # noqa: S101 -- enforced in
__init__
+ assert self.file_type is not None # noqa: S101 -- enforced above
documents = self._parse_bytes(self.source_bytes, self.file_type)
file_count = 1
else:
- assert self.source_path is not None # noqa: S101 -- enforced in
__init__
+ assert self.source_path is not None # noqa: S101 -- enforced above
files = self._resolve_files(self.source_path)
if not files:
raise FileNotFoundError(f"No files found matching
'{self.source_path}'.")
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 2631c544138..085f8512df8 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
@@ -874,16 +874,20 @@ class TestAgentOperatorMessageHistory:
assert kwargs["usage_limits"] is limits
assert kwargs["message_history"] == []
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible
with Airflow >= 3.1.0"
+ )
def test_message_history_with_hitl_review_raises(self):
"""message_history cannot be combined with HITL review (post-review
transcript is lost)."""
+ op = AgentOperator(
+ task_id="t",
+ prompt="run",
+ llm_conn_id="c",
+ message_history=[],
+ enable_hitl_review=True,
+ )
with pytest.raises(ValueError, match="message_history and
enable_hitl_review"):
- AgentOperator(
- task_id="t",
- prompt="run",
- llm_conn_id="c",
- message_history=[],
- enable_hitl_review=True,
- )
+ op.execute(context={})
@patch("pydantic_ai.models.wrapper.infer_model", side_effect=lambda m: m)
@patch("pydantic_ai.models.infer_model", autospec=True)
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_document_loader.py
b/providers/common/ai/tests/unit/common/ai/operators/test_document_loader.py
index 0a949c19010..575dbaed0a4 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_document_loader.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_document_loader.py
@@ -50,20 +50,24 @@ class TestDocumentLoaderInit:
assert "source_bytes" not in op.template_fields
def test_both_sources_raises(self):
+ op = DocumentLoaderOperator(task_id="test",
source_path="/tmp/file.txt", source_bytes=b"hello")
with pytest.raises(ValueError, match="not both"):
- DocumentLoaderOperator(task_id="test",
source_path="/tmp/file.txt", source_bytes=b"hello")
+ op.execute(context={})
def test_neither_source_raises(self):
+ op = DocumentLoaderOperator(task_id="test")
with pytest.raises(ValueError, match="Provide exactly one"):
- DocumentLoaderOperator(task_id="test")
+ op.execute(context={})
def test_source_bytes_without_file_type_raises(self):
+ op = DocumentLoaderOperator(task_id="test", source_bytes=b"hello")
with pytest.raises(ValueError, match="file_type"):
- DocumentLoaderOperator(task_id="test", source_bytes=b"hello")
+ op.execute(context={})
def test_empty_bytes_without_file_type_raises(self):
+ op = DocumentLoaderOperator(task_id="test", source_bytes=b"")
with pytest.raises(ValueError, match="file_type"):
- DocumentLoaderOperator(task_id="test", source_bytes=b"")
+ op.execute(context={})
class TestTextParser:
diff --git a/scripts/ci/prek/validate_operators_init_exemptions.txt
b/scripts/ci/prek/validate_operators_init_exemptions.txt
index 0f5c1e5e263..ae99a1b3369 100644
--- a/scripts/ci/prek/validate_operators_init_exemptions.txt
+++ b/scripts/ci/prek/validate_operators_init_exemptions.txt
@@ -30,8 +30,6 @@
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/kueue.
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py::KubernetesPodOperator
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/resource.py::KubernetesResourceBaseOperator
providers/cohere/src/airflow/providers/cohere/operators/embedding.py::CohereEmbeddingOperator
-providers/common/ai/src/airflow/providers/common/ai/operators/agent.py::AgentOperator
-providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py::DocumentLoaderOperator
providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposCreateOperator
providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposDeleteOperator
providers/databricks/src/airflow/providers/databricks/operators/databricks_repos.py::DatabricksReposUpdateOperator
