From ed33ae19fe617f6e5cd1d33475ef69d4e0104767 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Thu, 23 Jul 2026 15:02:55 -0700 Subject: [PATCH 1/2] Validate template fields after rendering in common.ai operators message_history (AgentOperator) and source_path/file_type (DocumentLoaderOperator) are template fields, rendered after __init__ runs. Both constructors validated them there, acting on the un-rendered Jinja expressions: AgentOperator rejected message_history + enable_hitl_review, and DocumentLoaderOperator enforced the source_path/source_bytes/file_type combination. Move those checks into execute(), which runs after rendering. Constructor-only checks that read no template field (durable/code_mode) stay in __init__. related: #70296 Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../providers/common/ai/operators/agent.py | 15 ++++++++------- .../common/ai/operators/document_loader.py | 19 ++++++++++--------- .../unit/common/ai/operators/test_agent.py | 18 +++++++++++------- .../ai/operators/test_document_loader.py | 12 ++++++++---- .../validate_operators_init_exemptions.txt | 2 -- 5 files changed, 37 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 f83a58f0e2aed..5e9d6cdc2a4f0 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 @@ def __init__( # 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,14 @@ def _build_durable_storage(self, context: Context) -> DurableStorageProtocol: ) def execute(self, context: Context) -> Any: + # message_history is a template field; validate its interaction with enable_hitl_review + # here, after rendering, rather than in __init__. + if self.message_history is not None and self.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.") + 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 b452cab72edf5..88001a26e9279 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 @@ def __init__( **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,20 @@ def __init__( 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__ + # where they are still the un-rendered Jinja expressions. + 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 2631c544138f5..085f8512df884 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 @@ def test_usage_limits_still_forwarded_with_history(self, mock_hook_cls): 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 6681c0f463dee..3f4969cd87473 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 @@ def test_template_fields_render_source_path_and_metadata(self): 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 01fd5bc56dbb7..bb530308591ee 100644 --- a/scripts/ci/prek/validate_operators_init_exemptions.txt +++ b/scripts/ci/prek/validate_operators_init_exemptions.txt @@ -34,8 +34,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 From 148c92e3c57925fd1d4591bd86e87b3fe6d63ff1 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Thu, 23 Jul 2026 23:15:23 -0700 Subject: [PATCH 2/2] Tighten common.ai operator validation comments Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../ai/src/airflow/providers/common/ai/operators/agent.py | 8 +++----- .../providers/common/ai/operators/document_loader.py | 3 +-- 2 files changed, 4 insertions(+), 7 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 5e9d6cdc2a4f0..5afbe3909b5e1 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 @@ -417,12 +417,10 @@ def _build_durable_storage(self, context: Context) -> DurableStorageProtocol: ) def execute(self, context: Context) -> Any: - # message_history is a template field; validate its interaction with enable_hitl_review - # here, after rendering, rather than in __init__. + # message_history is a template field; validate the combination after rendering. if self.message_history is not None and self.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. + # 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): 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 88001a26e9279..c6c7db9a789dc 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 @@ -148,8 +148,7 @@ def __init__( 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__ - # where they are still the un-rendered Jinja expressions. + # 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: