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..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 @@ -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,12 @@ 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 b452cab72edf5..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 @@ -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,19 @@ 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'.") + 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