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..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 @@ -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,13 +152,15 @@ 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'.") + # 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: 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..111cf7738e59a 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={}) @@ -69,6 +69,14 @@ def test_empty_bytes_without_file_type_raises(self): 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: def test_txt_file(self, tmp_path):