diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst b/providers/common/ai/docs/operators/llm_file_analysis.rst index 8b2733ab05acb..920206334a189 100644 --- a/providers/common/ai/docs/operators/llm_file_analysis.rst +++ b/providers/common/ai/docs/operators/llm_file_analysis.rst @@ -153,11 +153,20 @@ This operator also inherits ``LLMOperator``'s HITL review parameters -- Supported Formats ----------------- -- Text-like: ``.log``, ``.json``, ``.csv``, ``.parquet``, ``.avro`` +- Text-like: ``.log``, ``.txt``, ``.json``, ``.csv``, ``.parquet``, ``.avro`` - Multimodal: ``.png``, ``.jpg``, ``.jpeg``, ``.pdf`` when ``multi_modal=True`` -- Gzip-compressed text inputs are supported for ``.log.gz``, ``.json.gz``, and - ``.csv.gz``. -- Gzip is not supported for ``.parquet``, ``.avro``, image, or PDF inputs. +- ``gzip``, ``bzip2``, and ``xz`` compressed text inputs are supported for + ``.log``, ``.txt``, ``.json``, and ``.csv`` (``.log.gz``, ``.csv.bz2``, ``.json.xz``, ...). + ``bzip2`` and ``xz`` need a Python interpreter built with the ``bz2`` and + ``lzma`` modules; otherwise those inputs raise + ``AirflowOptionalProviderFeatureException``. +- Compression is not supported for ``.parquet``, ``.avro``, image, or PDF + inputs. +- For ``bzip2`` and ``xz``, concatenated streams are read in full, but any + data after a point that does not start a valid stream (for example ``xz`` + Stream Padding between streams, or trailing garbage) is silently ignored, + so only the content up to that point is analyzed. ``gzip`` rejects such + files instead (``BadGzipFile``). Parquet and Avro readers require their corresponding optional extras: diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py b/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py index 38c3d6f149293..4f8f994bd1d82 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py +++ b/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py @@ -28,6 +28,17 @@ from pathlib import PurePosixPath from typing import TYPE_CHECKING, Any +# bz2/lzma are optional CPython extensions and may be missing from some interpreter builds +try: + import bz2 +except ImportError: + bz2 = None # type: ignore[assignment] + +try: + import lzma +except ImportError: + lzma = None # type: ignore[assignment] + from pydantic_ai.messages import BinaryContent from airflow.providers.common.ai.exceptions import ( @@ -38,7 +49,7 @@ from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException, ObjectStoragePath if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Callable, Sequence from pydantic_ai.messages import UserContent @@ -64,7 +75,14 @@ "xz": "xz", "zst": "zstd", } -_GZIP_SUPPORTED_FORMATS = frozenset({"csv", "json", "log", "txt"}) +_CODEC_MODULES = {"bzip2": "bz2", "xz": "lzma"} +_KNOWN_CODECS = frozenset({"gzip", *_CODEC_MODULES}) +_DECOMPRESSORS: dict[str, Callable[..., io.BufferedIOBase]] = {"gzip": gzip.open} +if bz2 is not None: + _DECOMPRESSORS["bzip2"] = bz2.open +if lzma is not None: + _DECOMPRESSORS["xz"] = lzma.open +_COMPRESSION_SUPPORTED_FORMATS = frozenset({"csv", "json", "log", "txt"}) _TEXT_SAMPLE_HEAD_CHARS = 8_000 _TEXT_SAMPLE_TAIL_CHARS = 2_000 _MEDIA_TYPES = { @@ -370,15 +388,20 @@ def detect_file_format(path: ObjectStoragePath) -> tuple[str, str | None]: raise LLMFileAnalysisUnsupportedFormatError( f"Unsupported file format {detected!r} for {path}. Supported formats: {', '.join(SUPPORTED_FILE_FORMATS)}." ) - if compression and compression != "gzip": + if compression and compression not in _KNOWN_CODECS: log.info("Rejecting file %s because compression=%s is not supported.", path, compression) raise LLMFileAnalysisUnsupportedFormatError( f"Compression {compression!r} is not supported for file analysis." ) - if compression == "gzip" and detected not in _GZIP_SUPPORTED_FORMATS: + if compression and detected not in _COMPRESSION_SUPPORTED_FORMATS: raise LLMFileAnalysisUnsupportedFormatError( f"Compression {compression!r} is not supported for {detected!r} file analysis." ) + if compression and compression not in _DECOMPRESSORS: + raise AirflowOptionalProviderFeatureException( + f"Compression {compression!r} requires the {_CODEC_MODULES[compression]!r} module, " + "which is missing from this Python build." + ) return detected, compression @@ -536,10 +559,10 @@ def _render_avro(path: ObjectStoragePath, *, sample_rows: int, max_content_bytes def _read_raw_bytes(path: ObjectStoragePath, *, compression: str | None, max_bytes: int) -> bytes: with path.open("rb") as handle: - if compression == "gzip": - with gzip.GzipFile(fileobj=handle) as gzip_handle: - return _read_limited_bytes(gzip_handle, path=path, max_bytes=max_bytes) - return _read_limited_bytes(handle, path=path, max_bytes=max_bytes) + if compression is None: + return _read_limited_bytes(handle, path=path, max_bytes=max_bytes) + with _DECOMPRESSORS[compression](handle) as decompressed: + return _read_limited_bytes(decompressed, path=path, max_bytes=max_bytes) def _read_limited_bytes(handle: io.BufferedIOBase, *, path: ObjectStoragePath, max_bytes: int) -> bytes: diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py b/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py index 2b524998b60e1..d68563f1f0bae 100644 --- a/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py +++ b/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py @@ -30,6 +30,7 @@ LLMFileAnalysisUnsupportedFormatError, ) from airflow.providers.common.ai.utils.file_analysis import ( + _DECOMPRESSORS, FileAnalysisRequest, _infer_partitions, _read_raw_bytes, @@ -256,9 +257,14 @@ def test_total_size_limit_failure(self, tmp_path): mock_prepare.assert_not_called() - def test_gzip_expansion_respects_processed_content_limit(self, tmp_path): - path = tmp_path / "big.log.gz" - path.write_bytes(gzip.compress(b"A" * 20_000)) + @pytest.mark.parametrize( + ("suffix", "module_name"), + [("gz", "gzip"), ("bz2", "bz2"), ("xz", "lzma")], + ) + def test_compressed_expansion_respects_processed_content_limit(self, tmp_path, suffix, module_name): + codec = pytest.importorskip(module_name) + path = tmp_path / f"big.log.{suffix}" + path.write_bytes(codec.compress(b"A" * 20_000)) with pytest.raises(LLMFileAnalysisLimitExceededError, match="processed-content limit"): build_file_analysis_request( @@ -355,6 +361,8 @@ class TestFileAnalysisHelpers: [ ("events.csv", "csv", None), ("events.csv.gz", "csv", "gzip"), + ("events.csv.bz2", "csv", "bzip2"), + ("events.json.xz", "json", "xz"), ("dashboard.jpg", "jpg", None), ("report.pdf", "pdf", None), ("app", "log", None), @@ -378,22 +386,77 @@ def test_detect_file_format_rejects_unsupported_compression(self, tmp_path): with pytest.raises(LLMFileAnalysisUnsupportedFormatError, match="Compression"): detect_file_format(ObjectStoragePath(str(path))) - @pytest.mark.parametrize("filename", ["sample.parquet.gz", "sample.avro.gz", "sample.png.gz"]) - def test_detect_file_format_rejects_unsupported_gzip_format_combinations(self, tmp_path, filename): + @pytest.mark.parametrize( + ("filename", "codec", "module_name"), + [("events.csv.bz2", "bzip2", "bz2"), ("events.json.xz", "xz", "lzma")], + ) + def test_detect_file_format_without_codec_module(self, tmp_path, filename, codec, module_name): + path = tmp_path / filename + path.write_bytes(b"content") + gz_path = tmp_path / "events.csv.gz" + gz_path.write_bytes(b"content") + plain_path = tmp_path / "events.csv" + plain_path.write_bytes(b"content") + + with patch.dict(_DECOMPRESSORS, {"gzip": gzip.open}, clear=True): + with pytest.raises( + AirflowOptionalProviderFeatureException, + match=f"Compression '{codec}' requires the '{module_name}' module", + ): + detect_file_format(ObjectStoragePath(str(path))) + assert detect_file_format(ObjectStoragePath(str(gz_path))) == ("csv", "gzip") + assert detect_file_format(ObjectStoragePath(str(plain_path))) == ("csv", None) + + @pytest.mark.parametrize( + "filename", ["sample.parquet.gz", "sample.avro.bz2", "sample.png.xz", "sample.pdf.gz"] + ) + def test_detect_file_format_rejects_unsupported_compression_format_combinations(self, tmp_path, filename): path = tmp_path / filename path.write_bytes(b"content") - with pytest.raises(LLMFileAnalysisUnsupportedFormatError, match="not supported for"): + with pytest.raises( + LLMFileAnalysisUnsupportedFormatError, match=r"not supported for '\w+' file analysis" + ): detect_file_format(ObjectStoragePath(str(path))) - def test_read_raw_bytes_decompresses_gzip(self, tmp_path): - path = tmp_path / "events.log.gz" - path.write_bytes(gzip.compress(b"line one\nline two\n")) + @pytest.mark.parametrize( + ("suffix", "compression", "module_name"), + [ + ("gz", "gzip", "gzip"), + ("bz2", "bzip2", "bz2"), + ("xz", "xz", "lzma"), + ], + ) + def test_read_raw_bytes_decompresses(self, tmp_path, suffix, compression, module_name): + codec = pytest.importorskip(module_name) + path = tmp_path / f"events.log.{suffix}" + path.write_bytes(codec.compress(b"line one\nline two\n")) - content = _read_raw_bytes(ObjectStoragePath(str(path)), compression="gzip", max_bytes=1_024) + content = _read_raw_bytes(ObjectStoragePath(str(path)), compression=compression, max_bytes=1_024) assert content == b"line one\nline two\n" + @pytest.mark.parametrize( + ("suffix", "compression", "module_name", "separator", "expected"), + [ + ("bz2", "bzip2", "bz2", b"", b"first\nsecond\n"), + ("bz2", "bzip2", "bz2", b"GARBAGE", b"first\n"), + ("xz", "xz", "lzma", b"", b"first\nsecond\n"), + ("xz", "xz", "lzma", b"\x00\x00\x00\x00", b"first\n"), + ], + ids=["bz2-concatenated", "bz2-trailing-garbage", "xz-concatenated", "xz-stream-padding"], + ) + def test_read_raw_bytes_multi_stream_behavior( + self, tmp_path, suffix, compression, module_name, separator, expected + ): + codec = pytest.importorskip(module_name) + path = tmp_path / f"events.log.{suffix}" + path.write_bytes(codec.compress(b"first\n") + separator + codec.compress(b"second\n")) + + content = _read_raw_bytes(ObjectStoragePath(str(path)), compression=compression, max_bytes=1_024) + + assert content == expected + def test_truncate_text_preserves_head_and_tail(self): text = "A" * 9_000 + "B" * 3_000