refactor(phase-5): core domain logic parsers

openrag/components/indexer/loaders/audio/local_whisper.py

@@ -1,121 +1,67 @@
+"""
+Local Whisper-backed audio loader.
+
+The Ray actor + pool that drive ``faster-whisper`` (``WhisperActor``,
+``WhisperPool``) and the services-side :class:`BasePooledParser`
+implementation now live in
+``services/workers/parsers/whisper_workers.py``; this module re-exports
+``WhisperActor`` and ``WhisperPool`` for legacy import paths
+(``components.indexer.loaders.audio.local_whisper.WhisperActor`` is
+still used by the OpenAI audio loader for language detection, and by
+``utils/dependencies.py`` for the actor bootstrap).
+
+``LocalWhisperLoader`` is now a thin :class:`BaseLoader` adapter that
+delegates to
+:class:`core.indexing.parsers.audio.local_whisper.LocalWhisperParser`,
+which itself wraps the services-side pool. New code should call the
+core parser directly; this shim keeps the legacy loader-discovery path
+alive until consumers migrate.
+"""
+
 import asyncio
 from pathlib import Path
 
-import ray
-import torch
-from config import load_config
-from faster_whisper import WhisperModel
+from core.indexing.parsers.audio.local_whisper import LocalWhisperParser
+from core.models.document import Document as CoreDocument
 from langchain_core.documents.base import Document
+from services.workers.parsers.whisper_workers import (  # noqa: F401  (re-exported for legacy import paths)
+    LocalWhisperLoader as _ServicesWhisperPool,
+)
+from services.workers.parsers.whisper_workers import (  # noqa: F401
+    WhisperActor,
+    WhisperPool,
+)
 from utils.logger import get_logger
 
 from ..base import BaseLoader
 
 logger = get_logger()
-config = load_config()
-
-
-if torch.cuda.is_available():
-    WHISPER_NUM_GPUS = config.loader.local_whisper.whisper_num_gpus
-else:  # On CPU
-    WHISPER_NUM_GPUS = 0
-
-WHISPER_CONCURRENCY_PER_WORKER = config.loader.local_whisper.whisper_concurrency_per_worker
-
-
-@ray.remote(
-    num_gpus=WHISPER_NUM_GPUS, max_restarts=5, max_concurrency=WHISPER_CONCURRENCY_PER_WORKER
-)  # Ensure each worker processes one file at a time
-class WhisperActor:
-    def __init__(self):
-        import torch
-        from config import load_config
-        from utils.logger import get_logger
-
-        self.logger = get_logger()
-        self.config = load_config()
-
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        compute_type = "float16" if device == "cuda" else "int8"
-        model_name = self.config.loader.local_whisper.model
-
-        self.logger.info("Loading Whisper model", model_name=model_name, device=device, compute_type=compute_type)
-        self.model = WhisperModel(model_name, device=device, compute_type=compute_type)
-        self.logger.info("Whisper model loaded successfully", model_name=model_name, device=device)
-
-    async def transcribe(self, wav_path: str | Path) -> str:
-        self.logger.info("Transcribing audio file", file_path=Path(wav_path).name)
-
-        def _transcribe_sync() -> str:
-            segments, _ = self.model.transcribe(str(wav_path))
-            return "".join(segment.text for segment in segments)
-
-        return await asyncio.to_thread(_transcribe_sync)
-
-    async def detect_language(self, wav_path: str | Path, fallback_language="en") -> str:
-        try:
-            self.logger.info("Detecting language for audio file", file_path=Path(wav_path).name)
-
-            def _detect_language_sync() -> str:
-                # beam_size=1 + max_new_tokens=1 runs only language detection, no full transcription
-                _, info = self.model.transcribe(str(wav_path), beam_size=1, max_new_tokens=1)
-                return info.language
-
-            return await asyncio.to_thread(_detect_language_sync)
-
-        except Exception as e:
-            self.logger.error("Error detecting language", error=str(e))
-            return fallback_language
-
-
-@ray.remote
-class WhisperPool:
-    def __init__(self):
-        from utils.logger import get_logger
-
-        self.logger = get_logger()
-
-        n_workers = config.loader.local_whisper.whisper_n_workers
-        self.logger.info(f"Starting WhisperPool with {n_workers} workers")
-        self.workers = [WhisperActor.remote() for _ in range(n_workers)]
-        self._pending = [0] * n_workers
-
-    async def transcribe(self, path):
-        from components.ray_utils import call_ray_actor_with_timeout, retry_with_backoff
-
-        timeout = config.loader.local_whisper.whisper_timeout
-
-        async def attempt(i: int):
-            idx = min(range(len(self._pending)), key=lambda j: self._pending[j])
-            self._pending[idx] += 1
-            try:
-                return await call_ray_actor_with_timeout(
-                    self.workers[idx].transcribe.remote(path),
-                    timeout=timeout,
-                    task_description=f"WhisperPool transcribe ({path})",
-                )
-            finally:
-                self._pending[idx] -= 1
-
-        return await retry_with_backoff(
-            attempt,
-            max_retries=config.loader.local_whisper.whisper_max_task_retry,
-            base_delay=config.loader.local_whisper.whisper_retry_base_delay,
-            task_description=f"WhisperPool transcribe ({path})",
-        )
 
 
 class LocalWhisperLoader(BaseLoader):
+    """Adapter shim — delegates to ``LocalWhisperParser`` via the services-side pool."""
+
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.whisper_actor: WhisperPool = ray.get_actor("WhisperPool", namespace="openrag")
+        self._parser = LocalWhisperParser(pool=_ServicesWhisperPool())
 
     async def aload_document(self, file_path, metadata: dict = None, save_markdown=False):
+        path = Path(file_path)
+        raw_bytes = await asyncio.to_thread(path.read_bytes)
+        core_doc = CoreDocument(
+            filename=path.name,
+            content_type=CoreDocument.detect_content_type(path.name),
+            raw_bytes=raw_bytes,
+            metadata=dict(metadata) if metadata else {},
+        )
         try:
-            content = await self.whisper_actor.transcribe.remote(file_path)
-            doc = Document(page_content=content, metadata=metadata)
-            if save_markdown:
-                self.save_content(content, str(file_path))
-            return doc
+            processed = await self._parser.parse(core_doc)
         except Exception as e:
-            self.logger.error("Error loading document", error=str(e))
-            raise e
+            logger.error("Error loading document", error=str(e))
+            raise
+
+        content = "".join(b.text for b in processed.text_blocks)
+        doc = Document(page_content=content, metadata=dict(metadata) if metadata else {})
+        if save_markdown:
+            self.save_content(content, str(file_path))
+        return doc

openrag/components/indexer/loaders/audio/openai.py

@@ -1,11 +1,27 @@
+"""
+OpenAI-compatible audio loader.
+
+The transcription client now lives in
+``services/inference/parsers/openai_audio.py`` as
+:class:`OpenAIAudioClient` (a :class:`BaseClientParser`).
+``OpenAIAudioLoader`` is a thin :class:`BaseLoader` adapter that
+constructs the services-side client (with a Whisper-actor-backed
+language detector when ``transcriber.use_whisper_lang_detector`` is
+enabled) and wraps it in
+:class:`core.indexing.parsers.audio.client_based.ClientAudioParser`.
+New code should call the core parser directly; this shim keeps the
+legacy loader-discovery path alive until consumers migrate.
+"""
+
 import asyncio
 from pathlib import Path
 
-import ray
-from components.utils import get_audio_semaphore
+from core.indexing.parsers.audio.client_based import ClientAudioParser
+from core.models.document import Document as CoreDocument
+from core.models.document import DocumentType
 from langchain_core.documents.base import Document
-from openai import AsyncOpenAI
 from pydub import AudioSegment
+from services.inference.parsers.openai_audio import OpenAIAudioClient
 from utils.logger import get_logger
 
 from ..base import BaseLoader
@@ -17,108 +33,64 @@
 LANG_DETECT_SAMPLE_MS = 30_000  # 30 s
 
 
-class AudioTranscriber:
-    """Transcribes audio in a single request (no chunking).
-
-    Language detection is handled locally by WhisperActor (faster-whisper).
-    vLLM's native language detection fix is not yet merged (PR #34342) missed the v0.16.0 branch
-    cut (Feb 8) — it was merged Feb 21 and will ship in v0.17.0.
-    """
-
-    def __init__(self, config):
-        self.client = AsyncOpenAI(
-            base_url=config.loader.transcriber.base_url,
-            api_key=config.loader.transcriber.api_key,
-            timeout=config.loader.transcriber.timeout,
-        )
-        self.model_name = config.loader.transcriber.model_name
-        self.use_whisper_lang_detector = config.loader.transcriber.use_whisper_lang_detector
-        self.direct_upload_suffixes = config.loader.transcriber.direct_upload_suffixes
-
-    async def transcribe(self, file_path: Path) -> str:
-        # Formats in self.direct_upload_suffixes (configurable via
-        # TRANSCRIBER_DIRECT_UPLOAD_SUFFIXES) are sent as-is to avoid the ~10x
-        # size inflation from WAV conversion (Scaleway cap: 100 MB; OpenAI: 25 MB).
-        # Everything else falls back to WAV for vLLM/libsndfile deployments.
-
-        try:
-            logger.bind(file=file_path.name)
-            suffix = file_path.suffix.lower()
-            if suffix in self.direct_upload_suffixes:
-                wav_path = file_path
-                tmp_wav = None
-                # We still need to load the audio so language detection can
-                # extract its 30-second sample. ``AudioSegment.from_file``
-                # uses ffmpeg under the hood, so it handles every format.
-                sound = await asyncio.to_thread(AudioSegment.from_file, file_path)
-            else:
-                sound = await asyncio.to_thread(AudioSegment.from_file, file_path)
-                logger.info("Converting audio to WAV (unsupported container)", duration_s=f"{len(sound) / 1000:.1f}")
-                tmp_wav = file_path.with_suffix(".wav")
-                await asyncio.to_thread(sound.export, tmp_wav, format="wav")
-                wav_path = tmp_wav
-
-            language = await self._detect_language(sound, wav_path) if self.use_whisper_lang_detector else None
-            logger.info("Transcribing audio as a single request", language=language)
-
-            async with get_audio_semaphore():
-                return await self._transcribe_file(wav_path, language)
-        except Exception as e:
-            logger.exception("Error in transcribe", error=str(e))
-            raise e
-        finally:
-            if tmp_wav:
-                await asyncio.to_thread(tmp_wav.unlink, True)
+def _get_whisper_actor():
+    try:
+        return WhisperActor.options(name="WhisperActor", namespace="openrag", get_if_exists=True).remote()
+    except Exception as e:
+        logger.error("Error getting WhisperActor", error=str(e))
+        raise
 
-    async def _detect_language(self, sound: AudioSegment, wav_path: Path, fallback: str = "en") -> str:
-        """Detect language via local WhisperActor from a short audio sample."""
-        sample = sound[:LANG_DETECT_SAMPLE_MS]
-        tmp_path = wav_path.parent / f"{wav_path.stem}_langdetect.wav"
-        await asyncio.to_thread(sample.export, tmp_path, format="wav")
-        try:
-            whisper_actor = self._get_whisper_actor()
-            return await whisper_actor.detect_language.remote(tmp_path, fallback)
-        except Exception as e:
-            logger.exception("Language detection failed", error=str(e))
-            return fallback
-        finally:
-            await asyncio.to_thread(tmp_path.unlink, True)
 
-    def _get_whisper_actor(self):
-        actor_name = "WhisperActor"
-        try:
-            return ray.get_actor(actor_name, namespace="openrag")
-        except ValueError:
-            return WhisperActor.options(name=actor_name, namespace="openrag").remote()
-        except Exception as e:
-            logger.error("Error getting WhisperActor", error=str(e))
-            raise
-
-    async def _transcribe_file(self, wav_path: Path, language: str = None) -> str:
-        """Send a single file to the transcription endpoint."""
-        try:
-            kwargs = {"model": self.model_name, "file": wav_path}
-            if language:
-                kwargs["language"] = language
-            result = await self.client.audio.transcriptions.create(**kwargs)
-            return result.text
-        except Exception as e:
-            logger.exception("Error transcribing file", file=wav_path.name, error=str(e))
-            raise e
+async def _whisper_language_detector(file_path: Path) -> str | None:
+    """Detect language via the singleton ``WhisperActor`` from a short audio sample."""
+    sound = await asyncio.to_thread(AudioSegment.from_file, file_path)
+    sample = sound[:LANG_DETECT_SAMPLE_MS]
+    tmp_path = file_path.parent / f"{file_path.stem}_langdetect.wav"
+    await asyncio.to_thread(sample.export, tmp_path, format="wav")
+    try:
+        whisper_actor = _get_whisper_actor()
+        return await whisper_actor.detect_language.remote(tmp_path, "en")
+    except Exception as e:
+        logger.exception("Language detection failed", error=str(e))
+        return None
+    finally:
+        await asyncio.to_thread(tmp_path.unlink, True)
 
 
 class OpenAIAudioLoader(BaseLoader):
+    """Adapter shim — delegates to ``OpenAIAudioClient`` via ``ClientAudioParser``."""
+
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.transcriber = AudioTranscriber(config=self.config)
+        cfg = self.config.loader.transcriber
+        _client = OpenAIAudioClient(
+            base_url=cfg.base_url,
+            api_key=cfg.api_key,
+            model=cfg.model_name,
+            timeout=cfg.timeout,
+            direct_upload_suffixes=cfg.direct_upload_suffixes,
+            language_detector=_whisper_language_detector if cfg.use_whisper_lang_detector else None,
+        )
+        self._parser = ClientAudioParser(client=_client)
 
     async def aload_document(self, file_path, metadata: dict = None, save_markdown=False):
+        if metadata is None:
+            metadata = {}
+        path = Path(file_path)
+        raw_bytes = await asyncio.to_thread(path.read_bytes)
+        core_doc = CoreDocument(
+            filename=path.name,
+            content_type=DocumentType.AUDIO,
+            raw_bytes=raw_bytes,
+            metadata=dict(metadata),
+        )
         try:
-            content = await self.transcriber.transcribe(Path(file_path))
-            doc = Document(page_content=content, metadata=metadata)
-            if save_markdown:
-                self.save_content(content, str(file_path))
-            return doc
-        except Exception as e:
-            logger.exception("Error in OpenAIAudioLoader", path=file_path, error=str(e))
-            raise e
+            processed = await self._parser.parse(core_doc)
+        except Exception:
+            logger.exception("Error in OpenAIAudioLoader", path=str(file_path))
+            raise
+        content = "".join(b.text for b in processed.text_blocks)
+        doc = Document(page_content=content, metadata=metadata)
+        if save_markdown:
+            self.save_content(content, str(file_path))
+        return doc