feat: add audio support with docling pipelines

ai4rag/components/data/constants.py

@@ -19,5 +19,11 @@
         ".qmd",
         ".rmd",
         ".xhtml",
+        ".wav",
+        ".mp3",
+        ".m4a",
+        ".aac",
+        ".ogg",
+        ".flac",
     }
 )

ai4rag/components/data/text_extraction.py

@@ -12,11 +12,17 @@
 from pathlib import Path
 from typing import Any
 
+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
+from docling.datamodel.pipeline_options import (
+    AsrPipelineOptions,
+    PaginatedPipelineOptions,
+    ThreadedPdfPipelineOptions,
+)
 from docling.document_converter import (
     AsciiDocFormatOption,
+    AudioFormatOption,
     DocumentConverter,
     EmailFormatOption,
     EpubFormatOption,
@@ -29,6 +35,7 @@
     PowerpointFormatOption,
     WordFormatOption,
 )
+from docling.pipeline.asr_pipeline import AsrPipeline
 
 from ai4rag import handler
 from ai4rag.components.data.constants import SUPPORTED_EXTENSIONS
@@ -361,6 +368,12 @@ 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_BASE,
+    )
+    asr_pipeline_options.asr_options.language = None
+
     paginated_pipeline_options = PaginatedPipelineOptions(
         artifacts_path=ap,
         generate_page_images=False,
@@ -379,6 +392,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=asr_pipeline_options),
     }
 
 

dev_utils/file_store.py

@@ -8,32 +8,45 @@
 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
 
 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"
 
@@ -71,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_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
         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:

pyproject.toml

@@ -29,7 +29,7 @@ dynamic = ["version"]
 
 dependencies = [
   "boto3>=1.28",
-  "docling-slim[standard,feat-chunking,format-opendocument]~=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",

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

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"