From 71ce9e2bcb47249a0b6b9e683e701c48717613fb Mon Sep 17 00:00:00 2001 From: baha-bouali Date: Tue, 28 Jul 2026 15:06:42 +0100 Subject: [PATCH 1/6] Restore __init__ validation for common AI operators --- .../providers/common/ai/operators/agent.py | 13 +++++++------ .../common/ai/operators/document_loader.py | 13 ++++++++----- .../unit/common/ai/operators/test_agent.py | 18 +++++++----------- .../ai/operators/test_document_loader.py | 8 ++++---- 4 files changed, 26 insertions(+), 26 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 83a03bfda02f7..dac5e5b283077 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 @@ -297,6 +297,13 @@ 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 @@ -421,12 +428,6 @@ def _build_durable_storage(self, context: Context) -> DurableStorageProtocol: ) 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 217917f8112b2..f0c36adcbe06d 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,6 +136,10 @@ 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'.") self.source_path = source_path self.source_conn_id = source_conn_id self.source_bytes = source_bytes @@ -148,11 +152,10 @@ 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__. - 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'.") + # source_path/source_bytes provision is checked in __init__ (that's just "was an + # argument passed"). file_type is different: it backs the assert below, and since + # file_type is itself a template field, whether it was *actually* supplied is only + # knowable after rendering -- so this one check has to stay here, not in __init__. 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').") 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 760bacb8cc7f6..41470724d6008 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 @@ -877,20 +877,16 @@ 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"): - op.execute(context={}) + AgentOperator( + task_id="t", + prompt="run", + llm_conn_id="c", + message_history=[], + enable_hitl_review=True, + ) @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 575dbaed0a405..97f767c42fa5f 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,16 +50,16 @@ 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") + # source_path/source_bytes provision is a constructor-time check now. with pytest.raises(ValueError, match="not both"): - op.execute(context={}) + DocumentLoaderOperator(task_id="test", source_path="/tmp/file.txt", source_bytes=b"hello") def test_neither_source_raises(self): - op = DocumentLoaderOperator(task_id="test") with pytest.raises(ValueError, match="Provide exactly one"): - op.execute(context={}) + DocumentLoaderOperator(task_id="test") def test_source_bytes_without_file_type_raises(self): + # file_type is a template field, so this check only fires at execute() time. op = DocumentLoaderOperator(task_id="test", source_bytes=b"hello") with pytest.raises(ValueError, match="file_type"): op.execute(context={}) From 47531bb6cebdcbe4b20d76e9ac0e009ac7896136 Mon Sep 17 00:00:00 2001 From: baha-bouali Date: Wed, 29 Jul 2026 00:27:15 +0100 Subject: [PATCH 2/6] Guard rendered source_path against None; fix misleading comment --- .../common/ai/operators/document_loader.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) 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 f0c36adcbe06d..209281fc4784b 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 @@ -152,12 +152,17 @@ def __init__( self.json_text_field = json_text_field def execute(self, context: Context) -> list[dict[str, Any]]: - # source_path/source_bytes provision is checked in __init__ (that's just "was an - # argument passed"). file_type is different: it backs the assert below, and since - # file_type is itself a template field, whether it was *actually* supplied is only - # knowable after rendering -- so this one check has to stay here, not in __init__. + # file_type and source_path can each be *supplied* (a non-None argument) yet still + # render to None -- e.g. under render_template_as_native_obj=True, a Jinja + # expression can legitimately evaluate to None. These aren't provision checks (that + # already happened in __init__); they guard the rendered value itself, since + # _parse_bytes/_resolve_files need a real value to work with. Checking this in + # __init__ would validate the unrendered template string instead of the value + # actually used here. 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 None and self.source_path is None: + raise ValueError("Provide exactly one of 'source_path' or 'source_bytes'.") if self.source_bytes is not None: if TYPE_CHECKING: From af4d38f63cf5246d94aebda67320f75cb56b0be3 Mon Sep 17 00:00:00 2001 From: Baha Bouali <74876479+baha-bouali@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:34:05 +0100 Subject: [PATCH 3/6] Guard rendered source_path against None; fix misleading comment Clarify comments regarding file_type and source_path checks in execute method. --- .../providers/common/ai/operators/document_loader.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) 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 f0c36adcbe06d..b74fcf8841cb8 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 @@ -152,12 +152,15 @@ def __init__( self.json_text_field = json_text_field def execute(self, context: Context) -> list[dict[str, Any]]: - # source_path/source_bytes provision is checked in __init__ (that's just "was an - # argument passed"). file_type is different: it backs the assert below, and since - # file_type is itself a template field, whether it was *actually* supplied is only - # knowable after rendering -- so this one check has to stay here, not in __init__. + # file_type and source_path can each be *supplied* (as non-None argument) yet still + # render to None. These aren't provision checks (that already happened in __init__); + # they guard the rendered value itself, since _parse_bytes/_resolve_files need a real + # value to work with. Checking this in __init__ would validate the unrendered template + # string instead of the value actually used here. 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 None and self.source_path is None: + raise ValueError("Provide exactly one of 'source_path' or 'source_bytes'.") if self.source_bytes is not None: if TYPE_CHECKING: From c52aabd650b3367a272209d51fc1e1c1319b5032 Mon Sep 17 00:00:00 2001 From: Baha Bouali <74876479+baha-bouali@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:28:40 +0100 Subject: [PATCH 4/6] Add test for None source_path in DocumentLoaderOperator --- .../tests/unit/common/ai/operators/test_document_loader.py | 7 +++++++ 1 file changed, 7 insertions(+) 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 97f767c42fa5f..37567524665ff 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 @@ -68,6 +68,13 @@ def test_empty_bytes_without_file_type_raises(self): op = DocumentLoaderOperator(task_id="test", source_bytes=b"") with pytest.raises(ValueError, match="file_type"): op.execute(context={}) + def test_source_path_none_after_render_raises(self): + # source_path can render to None even when supplied -- must raise ValueError, + # not a TypeError from _resolve_files. + op = DocumentLoaderOperator(task_id="test", source_path="{{ none }}") + op.source_path = None # simulate the rendered value, bypassing real templating + with pytest.raises(ValueError, match="Provide exactly one"): + op.execute(context={}) class TestTextParser: From a4ec1331aab79093a79abc92fa076cb2fb742cf2 Mon Sep 17 00:00:00 2001 From: Baha Bouali <74876479+baha-bouali@users.noreply.github.com> Date: Sun, 2 Aug 2026 12:00:12 +0100 Subject: [PATCH 5/6] static checks: delete trailing spaces in a comment Updated comments for clarity and consistency in the execute method. --- .../providers/common/ai/operators/document_loader.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 b74fcf8841cb8..f38d9c75ab091 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 @@ -153,9 +153,9 @@ def __init__( def execute(self, context: Context) -> list[dict[str, Any]]: # file_type and source_path can each be *supplied* (as non-None argument) yet still - # render to None. These aren't provision checks (that already happened in __init__); - # they guard the rendered value itself, since _parse_bytes/_resolve_files need a real - # value to work with. Checking this in __init__ would validate the unrendered template + # render to None. These aren't provision checks (that already happened in __init__); + # they guard the rendered value itself, since _parse_bytes/_resolve_files need a real + # value to work with. Checking this in __init__ would validate the unrendered template # string instead of the value actually used here. 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').") From d04ba43624492670e03cf0888cfa950718cf1e43 Mon Sep 17 00:00:00 2001 From: Baha Bouali <74876479+baha-bouali@users.noreply.github.com> Date: Sun, 2 Aug 2026 12:02:09 +0100 Subject: [PATCH 6/6] static checks: add line break before a function --- .../ai/tests/unit/common/ai/operators/test_document_loader.py | 1 + 1 file changed, 1 insertion(+) 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 37567524665ff..92da323386e43 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 @@ -68,6 +68,7 @@ def test_empty_bytes_without_file_type_raises(self): op = DocumentLoaderOperator(task_id="test", source_bytes=b"") with pytest.raises(ValueError, match="file_type"): op.execute(context={}) + def test_source_path_none_after_render_raises(self): # source_path can render to None even when supplied -- must raise ValueError, # not a TypeError from _resolve_files.