From 792d1c1e0f7ecf1f6867e37d9d4df9c544e82f40 Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Fri, 31 Jul 2026 07:33:22 +0200 Subject: [PATCH 1/6] feat: add extended audio file formats, with new lib, and new input format Signed-off-by: Tomasz Guzik --- ai4rag/components/data/constants.py | 6 ++++ ai4rag/components/data/text_extraction.py | 13 +++++--- dev_utils/file_store.py | 40 +++++++++++++---------- pyproject.toml | 3 +- uv.lock | 23 ++++++++++++- 5 files changed, 62 insertions(+), 23 deletions(-) diff --git a/ai4rag/components/data/constants.py b/ai4rag/components/data/constants.py index 87c7bc2..4763628 100644 --- a/ai4rag/components/data/constants.py +++ b/ai4rag/components/data/constants.py @@ -19,5 +19,11 @@ ".qmd", ".rmd", ".xhtml", + ".wav", + ".mp3", + ".m4a", + ".aac", + ".ogg", + ".flac" } ) diff --git a/ai4rag/components/data/text_extraction.py b/ai4rag/components/data/text_extraction.py index 09da247..6e56489 100644 --- a/ai4rag/components/data/text_extraction.py +++ b/ai4rag/components/data/text_extraction.py @@ -12,9 +12,14 @@ from pathlib import Path from typing import Any +from ai4rag import handler +from ai4rag.components.data.constants import SUPPORTED_EXTENSIONS from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import InputFormat -from docling.datamodel.pipeline_options import PaginatedPipelineOptions, ThreadedPdfPipelineOptions +from docling.datamodel.pipeline_options import ( + PaginatedPipelineOptions, + ThreadedPdfPipelineOptions, +) from docling.document_converter import ( AsciiDocFormatOption, DocumentConverter, @@ -28,10 +33,9 @@ PdfFormatOption, PowerpointFormatOption, WordFormatOption, + AudioFormatOption, ) - -from ai4rag import handler -from ai4rag.components.data.constants import SUPPORTED_EXTENSIONS +from docling.pipeline.asr_pipeline import AsrPipeline _logger = logging.getLogger("text-extraction") _logger.addHandler(handler) @@ -379,6 +383,7 @@ def _build_docling_format_options(do_table_structure: bool = False) -> dict: InputFormat.LATEX: LatexFormatOption(), InputFormat.EPUB: EpubFormatOption(), InputFormat.EMAIL: EmailFormatOption(), + InputFormat.AUDIO: AudioFormatOption(pipeline_cls=AsrPipeline, pipeline_options=AsrPipeline), } diff --git a/dev_utils/file_store.py b/dev_utils/file_store.py index 3c184af..6cd0d03 100644 --- a/dev_utils/file_store.py +++ b/dev_utils/file_store.py @@ -17,23 +17,29 @@ logger = logging.getLogger(__name__) -SUPPORTED_EXTENSIONS = { - ".pdf", - ".docx", - ".pptx", - ".md", - ".html", - ".txt", - ".odt", - ".odp", - ".adoc", - ".tex", - ".epub", - ".eml", - ".qmd", - ".rmd", - ".xhtml", - } +SUPPORTED_EXTENSIONS = { + ".pdf", + ".docx", + ".pptx", + ".md", + ".html", + ".txt", + ".odt", + ".odp", + ".adoc", + ".tex", + ".epub", + ".eml", + ".qmd", + ".rmd", + ".xhtml", + ".wav", + ".mp3", + ".m4a", + ".aac", + ".ogg", + ".flac", +} _DEFAULT_CACHE_DIR = Path(__file__).parent / "local" / "docling_cache" diff --git a/pyproject.toml b/pyproject.toml index a5b4f56..adfe263 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,8 @@ dynamic = ["version"] dependencies = [ "boto3>=1.28", - "docling[feat-chunking,format-opendocument]~=2.107.0", + "docling~=2.107.0", + "docling-slim[feat-chunking,format-opendocument,asr]~=2.107.0", "langchain_chroma~=1.1.0", "langchain-text-splitters~=1.1.0", "multiprocess>=0.70", diff --git a/uv.lock b/uv.lock index 1f65715..0dcfc71 100644 --- a/uv.lock +++ b/uv.lock @@ -32,6 +32,7 @@ source = { editable = "." } dependencies = [ { name = "boto3" }, { name = "docling" }, + { name = "docling-slim", extra = ["feat-chunking", "format-opendocument"] }, { name = "langchain-chroma" }, { name = "langchain-text-splitters" }, { name = "multiprocess" }, @@ -91,7 +92,8 @@ requires-dist = [ { name = "beautifulsoup4", marker = "extra == 'dev'" }, { name = "black", marker = "extra == 'code-check'" }, { name = "boto3", specifier = ">=1.28" }, - { name = "docling", extras = ["feat-chunking", "format-opendocument"], specifier = "~=2.107.0" }, + { name = "docling", specifier = "~=2.107.0" }, + { name = "docling-slim", extras = ["asr", "feat-chunking", "format-opendocument"], specifier = "~=2.107.0" }, { name = "dotenv", marker = "extra == 'dev'" }, { name = "ipykernel", marker = "extra == 'dev'" }, { name = "isort", marker = "extra == 'code-check'" }, @@ -922,6 +924,12 @@ wheels = [ ] [package.optional-dependencies] +feat-chunking = [ + { name = "docling-core", extra = ["chunking"] }, +] +format-opendocument = [ + { name = "odfdo" }, +] standard = [ { name = "accelerate" }, { name = "beautifulsoup4" }, @@ -2385,6 +2393,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/be/9c/92789c596b8df838baa98fa71844d84283302f7604ed565dafe5a6b5041a/oauthlib-3.3.1-py3-none-any.whl", hash = "sha256:88119c938d2b8fb88561af5f6ee0eec8cc8d552b7bb1f712743136eb7523b7a1", size = 160065, upload-time = "2025-06-19T22:48:06.508Z" }, ] +[[package]] +name = "odfdo" +version = "3.23.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "lxml" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/32/7c/77b3ef7ac4074eebd5fea4edd839c0637f2738ac323ee5ba60459b0c4f8e/odfdo-3.23.1.tar.gz", hash = "sha256:86414364d1184f94f7d1788822d9002d402e5b5737fdaa463e9c637a1428baf9", size = 326049, upload-time = "2026-07-22T16:54:06.327Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/86/57/8a1713d09b3366be82adf51571922afe5ac245e481b9ffbdbf5a0de47509/odfdo-3.23.1-py3-none-any.whl", hash = "sha256:06e4136dcffb5407041a8342b15b15bf874906c190d546551d15bbf20fa3062a", size = 428137, upload-time = "2026-07-22T16:54:04.957Z" }, +] + [[package]] name = "ogx-client" version = "1.1.3" From 6e6c4e546cd4eb206cc5036af22881635f31cbb9 Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Fri, 31 Jul 2026 12:31:07 +0200 Subject: [PATCH 2/6] feat: add the audio file extraction Signed-off-by: Tomasz Guzik --- ai4rag/components/data/text_extraction.py | 9 ++++++- dev_utils/file_store.py | 32 ++++++++++++++++++----- 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/ai4rag/components/data/text_extraction.py b/ai4rag/components/data/text_extraction.py index 6e56489..a37f087 100644 --- a/ai4rag/components/data/text_extraction.py +++ b/ai4rag/components/data/text_extraction.py @@ -14,11 +14,13 @@ from ai4rag import handler from ai4rag.components.data.constants import SUPPORTED_EXTENSIONS +from docling.datamodel import asr_model_specs from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import InputFormat from docling.datamodel.pipeline_options import ( PaginatedPipelineOptions, ThreadedPdfPipelineOptions, + AsrPipelineOptions, ) from docling.document_converter import ( AsciiDocFormatOption, @@ -365,6 +367,11 @@ def _build_docling_format_options(do_table_structure: bool = False) -> dict: do_table_structure=do_table_structure, accelerator_options=accel, ) + + asr_pipeline_options = AsrPipelineOptions( + asr_options=asr_model_specs.WHISPER_TURBO_MLX + ) + paginated_pipeline_options = PaginatedPipelineOptions( artifacts_path=ap, generate_page_images=False, @@ -383,7 +390,7 @@ def _build_docling_format_options(do_table_structure: bool = False) -> dict: InputFormat.LATEX: LatexFormatOption(), InputFormat.EPUB: EpubFormatOption(), InputFormat.EMAIL: EmailFormatOption(), - InputFormat.AUDIO: AudioFormatOption(pipeline_cls=AsrPipeline, pipeline_options=AsrPipeline), + InputFormat.AUDIO: AudioFormatOption(pipeline_cls=AsrPipeline, pipeline_options=asr_pipeline_options), } diff --git a/dev_utils/file_store.py b/dev_utils/file_store.py index 6cd0d03..ecd697b 100644 --- a/dev_utils/file_store.py +++ b/dev_utils/file_store.py @@ -8,10 +8,17 @@ from pathlib import Path from typing import Sequence +from docling.datamodel import asr_model_specs from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import InputFormat -from docling.datamodel.pipeline_options import PdfPipelineOptions -from docling.document_converter import DocumentConverter, PdfFormatOption, settings +from docling.datamodel.pipeline_options import PdfPipelineOptions, AsrPipelineOptions +from docling.document_converter import ( + DocumentConverter, + PdfFormatOption, + settings, + AudioFormatOption, +) +from docling.pipeline.asr_pipeline import AsrPipeline from docling_core.types.doc import DoclingDocument from docling_core.types.doc.labels import DocItemLabel @@ -77,17 +84,28 @@ def __init__( self.cache_dir = Path(cache_dir) if cache_dir is not None else None self.files = {} - pipeline_options = PdfPipelineOptions() - pipeline_options.do_ocr = False - pipeline_options.do_table_structure = True - pipeline_options.accelerator_options = AcceleratorOptions(device="auto") + pdf_options = PdfPipelineOptions() + pdf_options.do_ocr = False + pdf_options.do_table_structure = True + pdf_options.accelerator_options = AcceleratorOptions(device="auto") + + asr_options = AsrPipelineOptions() + asr_options.asr_options = asr_model_specs.WHISPER_TURBO + asr_options.asr_options.language = None #if we stay with default the multilang has problems with non-eng texts num_workers = os.cpu_count() or 1 settings.perf.doc_batch_size = num_workers settings.perf.doc_batch_concurrency = num_workers self._converter = DocumentConverter( - format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)} + format_options={ + InputFormat.PDF: PdfFormatOption( + pipeline_options=pdf_options), + InputFormat.AUDIO: AudioFormatOption( + pipeline_cls=AsrPipeline, + pipeline_options=asr_options + ), + } ) def __repr__(self) -> str: From a392526124035d0392865c3b1f175955fa28398d Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Fri, 31 Jul 2026 12:49:08 +0200 Subject: [PATCH 3/6] feat: add the audio file extraction Signed-off-by: Tomasz Guzik --- ai4rag/components/data/text_extraction.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ai4rag/components/data/text_extraction.py b/ai4rag/components/data/text_extraction.py index a37f087..79c34a5 100644 --- a/ai4rag/components/data/text_extraction.py +++ b/ai4rag/components/data/text_extraction.py @@ -369,8 +369,9 @@ def _build_docling_format_options(do_table_structure: bool = False) -> dict: ) asr_pipeline_options = AsrPipelineOptions( - asr_options=asr_model_specs.WHISPER_TURBO_MLX + asr_options=asr_model_specs.WHISPER_TURBO, ) + asr_pipeline_options.asr_options.language = None paginated_pipeline_options = PaginatedPipelineOptions( artifacts_path=ap, From 735a4058b8bc7d5fbaebc81dd613430268010919 Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Tue, 4 Aug 2026 10:08:07 +0200 Subject: [PATCH 4/6] feat: add standard lib Signed-off-by: Tomasz Guzik --- pyproject.toml | 2 +- uv.lock | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index adfe263..f66c78a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ dynamic = ["version"] dependencies = [ "boto3>=1.28", "docling~=2.107.0", - "docling-slim[feat-chunking,format-opendocument,asr]~=2.107.0", + "docling-slim[standard,feat-chunking,format-opendocument,asr]~=2.107.0", "langchain_chroma~=1.1.0", "langchain-text-splitters~=1.1.0", "multiprocess>=0.70", diff --git a/uv.lock b/uv.lock index 0dcfc71..008b5d0 100644 --- a/uv.lock +++ b/uv.lock @@ -32,7 +32,7 @@ source = { editable = "." } dependencies = [ { name = "boto3" }, { name = "docling" }, - { name = "docling-slim", extra = ["feat-chunking", "format-opendocument"] }, + { name = "docling-slim", extra = ["feat-chunking", "format-opendocument", "standard"] }, { name = "langchain-chroma" }, { name = "langchain-text-splitters" }, { name = "multiprocess" }, @@ -93,7 +93,7 @@ requires-dist = [ { name = "black", marker = "extra == 'code-check'" }, { name = "boto3", specifier = ">=1.28" }, { name = "docling", specifier = "~=2.107.0" }, - { name = "docling-slim", extras = ["asr", "feat-chunking", "format-opendocument"], specifier = "~=2.107.0" }, + { name = "docling-slim", extras = ["asr", "feat-chunking", "format-opendocument", "standard"], specifier = "~=2.107.0" }, { name = "dotenv", marker = "extra == 'dev'" }, { name = "ipykernel", marker = "extra == 'dev'" }, { name = "isort", marker = "extra == 'code-check'" }, From 1b1d60279a3718a60144173512596b94482a0665 Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Tue, 4 Aug 2026 10:41:48 +0200 Subject: [PATCH 5/6] tests: add unit tests for the new feature Signed-off-by: Tomasz Guzik --- ai4rag/components/data/constants.py | 2 +- .../ai4rag/components/data/test_discovery.py | 54 +++++++++++++++++++ .../ai4rag/components/data/test_extraction.py | 45 ++++++++++++++++ 3 files changed, 100 insertions(+), 1 deletion(-) diff --git a/ai4rag/components/data/constants.py b/ai4rag/components/data/constants.py index 4763628..0525f71 100644 --- a/ai4rag/components/data/constants.py +++ b/ai4rag/components/data/constants.py @@ -24,6 +24,6 @@ ".m4a", ".aac", ".ogg", - ".flac" + ".flac", } ) diff --git a/tests/unit/ai4rag/components/data/test_discovery.py b/tests/unit/ai4rag/components/data/test_discovery.py index 4583a93..173b3a1 100644 --- a/tests/unit/ai4rag/components/data/test_discovery.py +++ b/tests/unit/ai4rag/components/data/test_discovery.py @@ -326,3 +326,57 @@ def test_list_objects_called_correctly(self, mocker): ) mock_client.list_objects_v2.assert_called_once_with(Bucket="my-bucket", Prefix="prefix/") + + def test_audio_extensions_discovered(self, mocker): + """Audio files with supported extensions must be discovered.""" + contents = [ + _s3_object("audio/meeting.wav", 1000), + _s3_object("audio/podcast.mp3", 2000), + _s3_object("audio/recording.m4a", 1500), + _s3_object("audio/clip.aac", 800), + _s3_object("audio/voice.ogg", 600), + _s3_object("audio/sample.flac", 3000), + ] + mock_client = _make_mock_s3_client(mocker, contents) + + result = discover_documents( + bucket_name="bucket", + prefix="audio/", + sampling_enabled=False, + s3_client=mock_client, + ) + + assert result.count == 6 + keys = [d.key for d in result.documents] + assert "audio/meeting.wav" in keys + assert "audio/podcast.mp3" in keys + assert "audio/recording.m4a" in keys + assert "audio/clip.aac" in keys + assert "audio/voice.ogg" in keys + assert "audio/sample.flac" in keys + + def test_mixed_audio_and_document_extensions(self, mocker): + """Audio and document files should both be discovered together.""" + contents = [ + _s3_object("data/report.pdf", 500), + _s3_object("data/meeting.mp3", 2000), + _s3_object("data/notes.md", 100), + _s3_object("data/recording.wav", 3000), + _s3_object("data/image.png", 400), + ] + mock_client = _make_mock_s3_client(mocker, contents) + + result = discover_documents( + bucket_name="bucket", + prefix="data/", + sampling_enabled=False, + s3_client=mock_client, + ) + + assert result.count == 4 + keys = [d.key for d in result.documents] + assert "data/report.pdf" in keys + assert "data/meeting.mp3" in keys + assert "data/notes.md" in keys + assert "data/recording.wav" in keys + assert "data/image.png" not in keys diff --git a/tests/unit/ai4rag/components/data/test_extraction.py b/tests/unit/ai4rag/components/data/test_extraction.py index e3a2698..3e60c9e 100644 --- a/tests/unit/ai4rag/components/data/test_extraction.py +++ b/tests/unit/ai4rag/components/data/test_extraction.py @@ -371,6 +371,7 @@ def test_formats_registered_have_entries_in_options(self): InputFormat.LATEX, InputFormat.EPUB, InputFormat.EMAIL, + InputFormat.AUDIO, ): assert fmt in options, f"{fmt} missing from format options" @@ -389,3 +390,47 @@ def test_do_table_structure_false(self): options = _build_docling_format_options(do_table_structure=False) pdf_option = options[InputFormat.PDF] assert pdf_option.pipeline_options.do_table_structure is False + + def test_audio_format_uses_asr_pipeline(self): + """Audio format option must use the AsrPipeline class.""" + from docling.datamodel.base_models import InputFormat + from docling.pipeline.asr_pipeline import AsrPipeline + + options = _build_docling_format_options() + audio_option = options[InputFormat.AUDIO] + assert audio_option.pipeline_cls is AsrPipeline + + def test_audio_format_language_is_auto_detect(self): + """Audio ASR options must use language=None for auto-detection.""" + from docling.datamodel.base_models import InputFormat + + options = _build_docling_format_options() + audio_option = options[InputFormat.AUDIO] + assert audio_option.pipeline_options.asr_options.language is None + + +# --------------------------------------------------------------------------- +# SUPPORTED_EXTENSIONS — audio formats +# --------------------------------------------------------------------------- + + +class TestSupportedExtensionsAudio: + """Tests that audio formats are included in SUPPORTED_EXTENSIONS.""" + + AUDIO_EXTENSIONS = {".wav", ".mp3", ".m4a", ".aac", ".ogg", ".flac"} + + def test_audio_extensions_present(self): + """All audio extensions must be in SUPPORTED_EXTENSIONS.""" + from ai4rag.components.data.constants import SUPPORTED_EXTENSIONS + + for ext in self.AUDIO_EXTENSIONS: + assert ext in SUPPORTED_EXTENSIONS, f"{ext} missing from SUPPORTED_EXTENSIONS" + + def test_original_extensions_still_present(self): + """Adding audio extensions must not remove existing document formats.""" + from ai4rag.components.data.constants import SUPPORTED_EXTENSIONS + + original = {".pdf", ".docx", ".pptx", ".md", ".html", ".txt", ".odt", ".odp", + ".adoc", ".tex", ".epub", ".eml", ".qmd", ".rmd", ".xhtml"} + for ext in original: + assert ext in SUPPORTED_EXTENSIONS, f"{ext} missing from SUPPORTED_EXTENSIONS" From fd6dd68c28d466a54c28c3dc2cb532fe6140a3ee Mon Sep 17 00:00:00 2001 From: Tomasz Guzik Date: Tue, 4 Aug 2026 15:05:55 +0200 Subject: [PATCH 6/6] tests: change model to base Signed-off-by: Tomasz Guzik --- ai4rag/components/data/text_extraction.py | 11 ++++++----- dev_utils/file_store.py | 2 +- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/ai4rag/components/data/text_extraction.py b/ai4rag/components/data/text_extraction.py index 79c34a5..37a59b4 100644 --- a/ai4rag/components/data/text_extraction.py +++ b/ai4rag/components/data/text_extraction.py @@ -12,18 +12,17 @@ from pathlib import Path from typing import Any -from ai4rag import handler -from ai4rag.components.data.constants import SUPPORTED_EXTENSIONS from docling.datamodel import asr_model_specs from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import InputFormat from docling.datamodel.pipeline_options import ( + AsrPipelineOptions, PaginatedPipelineOptions, ThreadedPdfPipelineOptions, - AsrPipelineOptions, ) from docling.document_converter import ( AsciiDocFormatOption, + AudioFormatOption, DocumentConverter, EmailFormatOption, EpubFormatOption, @@ -35,10 +34,12 @@ PdfFormatOption, PowerpointFormatOption, WordFormatOption, - AudioFormatOption, ) from docling.pipeline.asr_pipeline import AsrPipeline +from ai4rag import handler +from ai4rag.components.data.constants import SUPPORTED_EXTENSIONS + _logger = logging.getLogger("text-extraction") _logger.addHandler(handler) @@ -369,7 +370,7 @@ def _build_docling_format_options(do_table_structure: bool = False) -> dict: ) asr_pipeline_options = AsrPipelineOptions( - asr_options=asr_model_specs.WHISPER_TURBO, + asr_options=asr_model_specs.WHISPER_BASE, ) asr_pipeline_options.asr_options.language = None diff --git a/dev_utils/file_store.py b/dev_utils/file_store.py index ecd697b..1a98c3d 100644 --- a/dev_utils/file_store.py +++ b/dev_utils/file_store.py @@ -90,7 +90,7 @@ def __init__( pdf_options.accelerator_options = AcceleratorOptions(device="auto") asr_options = AsrPipelineOptions() - asr_options.asr_options = asr_model_specs.WHISPER_TURBO + asr_options.asr_options = asr_model_specs.WHISPER_BASE asr_options.asr_options.language = None #if we stay with default the multilang has problems with non-eng texts num_workers = os.cpu_count() or 1