diff --git a/REFACTORING_DECISION_LOG.md b/REFACTORING_DECISION_LOG.md index 8cf2a8003..148fb2e2d 100644 --- a/REFACTORING_DECISION_LOG.md +++ b/REFACTORING_DECISION_LOG.md @@ -68,7 +68,349 @@ a separate error handler mapping (mandragora style).** status code mapping in api/error_handlers.py). Cleaner for hexagonal purity but rejected for backward compatibility. - Follow-up: strip status codes from core exceptions in Phase 10 when - api/error_handlers.py is built. The error handler will own the mapping. + `api/error_handlers.py` is built. The error handler will own the mapping. + +--- + +## Phase 5D — Indexing domain logic + parsers (2026-04-30) + +**1. `core/indexing/validators.py` is fully framework-free.** +- All FastAPI types removed (`Form`, `UploadFile`, `HTTPException`, + `status`, `Depends`); validators are pure functions on `str` / `dict` + / `Iterable[str]`. +- `accepted_formats` / `accepted_mimetypes` are passed as args instead + of read from Hydra `config` at module import (legacy module-level reads + of `ACCEPTED_FILE_FORMATS` / `DICT_MIMETYPES` / + `FORBIDDEN_CHARS_IN_FILE_ID` are gone). +- `ValidationError` accepts a `status_code` (and `code`) kwarg. Phase 1 + hardcoded 422; the original validators raised HTTP 400 (invalid + `file_id` / metadata JSON) and HTTP 415 (unsupported format), so a + status-code override is needed to preserve those codes from a + pure-domain exception. Existing precedent in the same module + (`LLMParsingError` overrides `status_code` after `super().__init__`) + shows the pattern is already accepted. +- HTTP translation flows through the existing global + `openrag_exception_handler` (`@app.exception_handler(OpenRAGError)`) + wired in Phase 1, not local `HTTPException` raises in routers. +- Why: A core module that imports FastAPI or reaches into Hydra is not + framework-free, blocks reuse from non-HTTP entry points, and + re-introduces the boundary violation the refactor exists to fix. + Stripping only `Depends()` — the literal task description — would + leave the boundary half-broken. +- Trade-off: error body becomes `{"detail": "[CODE]: msg", "extra": {}}` + instead of FastAPI's `{"detail": "..."}` — matches every other + `OpenRAGError`. +- Alternatives considered: (a) consolidate everything on 422 — rejected, + observable behaviour change; (b) introduce specific subclasses + (`UnsupportedFileFormatError`, etc.) — rejected as premature, only two + call sites need non-default codes today (Phase 10's API error-handler + layer can re-evaluate); (c) keep module-level Hydra reads — rejected, + embeds the infrastructure config object into core; (d) catch and + re-raise as `HTTPException` in the router wrappers — rejected, + duplicates the global handler. + +**2. Exception shims under `utils/exceptions/` use `core.X`, not `openrag.core.X`.** +The legacy shims imported via `openrag.core.utils.exceptions`. With both +`/app` and `/app/openrag/` reachable, Python loads the same file as two +distinct modules, producing two distinct `OpenRAGError` classes — +`isinstance` failed and the global handler never fired. +- Why: Unifying on the bare `core.X` path matches `pythonpath = ./openrag` + and the relative-imports-within-`core/` convention (commit 4528c71). +- Follow-up: ~20 other `from openrag.X` imports across `core/`, `config/`, + and components are latent dual-import traps and should be migrated + in a separate pass. + +**3. Parser layering: native in core, services-backed in services/workers, type-marker bases without vendor names, DI for pools.** +- Native-bytes parsers (PyMuPDF, html_to_markdown, chardet, image) live + in `core/indexing/parsers/`. Service-/Ray-backed parsers (Marker, + LocalWhisper) live in `services/workers/parsers/`. +- Empty marker subclasses `BasePooledParser` / `BaseClientParser` in + `core/indexing/parsers/document_parser.py` categorize parsers by *how* + they get their work done (actor-pool vs HTTP-client) without naming + the implementation. A core base class called `RayPoolParser` or + `OpenaiClientParser` would leak vendor/infrastructure into the + framework-free layer and foreclose swapping the backend. +- Core facades (`MarkerParser`, `LocalWhisperParser`, `ClientPdfParser`, + `ClientAudioParser`) accept any pool/client of the appropriate marker + type via `__init__`; services own the actor lifecycle. +- Why: `@ray.remote` decoration imports infrastructure at + class-definition time and can't be hidden behind a port. DI keeps + facades testable with in-memory fakes. +- Alternatives considered: (a) all parsers in core with Ray injected + via DI — rejected, class-level decoration can't be deferred to + composition; (b) have core facades resolve the actor by name + themselves — rejected, couples core to Ray's named-actor registry. + +**4. Image preprocessing helpers extracted to `core/indexing/image_preprocessor.py`.** +Pure helpers (`ensure_png_compatible_mode`, `pil_to_png_bytes`, +`pil_to_base64`, `is_http_url`, `is_data_uri`, `HTTP_IMAGE_PATTERN`, +`DATA_URI_IMAGE_PATTERN`, `MIN_IMAGE_PIXELS`). Used by the core image +parser and by Marker captioning in services. +- Why: Both layers need PNG normalization and markdown image-reference + detection. Sharing via core (no VLM, no langchain imports) avoids + services depending on `components/indexer/loaders/base.py`. +- Alternative considered: leave helpers in + `components/indexer/loaders/base.py`. Rejected — + services-importing-components is a layering violation, and `base.py` + drags in langchain. + +**5. `services/workers/ray_utils.py` keeps function and decorator forms together; `description=` is a format-string template.** +- `call_ray_actor_with_timeout` / `@with_timeout` and `retry_with_backoff` + / `@with_retry` (with jitter) live in one module — STRATEGY's proposed + `_retry.py` / `_timeout.py` split for `services/inference/` doesn't + apply here because workers need both forms in practice (decorator at + class-definition for static-param call sites, function form for + callsite-resolved values). The decorators delegate to the function + form internally; splitting across two files would duplicate that + wiring. +- `description=` accepts a **format string** like + `"PDF parse ({file_path})"`; `_resolve_description` binds it via + `inspect.signature.bind` against the wrapped call's args at call + time. **Callables (lambdas) are NOT supported** — they fall through + to `if "{" not in template:` and raise `TypeError: argument of type + 'function' is not iterable`. (One outlier in `marker_workers.py` used + a lambda and was fixed in Phase 5E.) +- Inline `call_ray_actor_with_timeout(worker.X.remote(...))` calls in + workers are extracted into one-line `@with_timeout`-decorated helper + methods (`_transcribe_chunk`, `_check_pool_broken`, + `_reset_worker_pool`, `_run_chunk`, `_convert_pdf`) returning the + `ObjectRef`; the decorator awaits it with timeout. Worker files use + only decorator form — no mixed styles. +- Retry-around-timeout semantics preserved: `@with_retry` outer, + `@with_timeout` inner — `TimeoutError` propagates from the inner + helper and the outer decorator re-runs the whole method body (slot + pick, fresh `.remote()`, fresh timeout). +- Alternatives considered: (a) mirror inference's `_retry.py` / + `_timeout.py` split verbatim — rejected, adds files that just import + from each other; (b) keep description static, drop to function form + when dynamic — rejected, re-introduces the verbose + `call_ray_actor_with_timeout(...)` call sites the decorator was meant + to remove; (c) keep function form for the inline cases — rejected, + leaves a mix of styles in the same file with no clear rule. + +**6. `ray_utils` canonical home moved from `components/` to `services/workers/`.** +`components/ray_utils.py` is now a back-compat shim re-exporting from +`services.workers.ray_utils`. +- Why: Ray-actor concurrency primitives belong in the services layer, + not in `components/` (which is on the deprecation path). Routers and + pipeline still import via the components shim during the transition. +- Follow-up: migrate the remaining `components.ray_utils` imports + (pipeline, search router, indexer router, workspaces router, indexer + utils) and delete the shim in Phase 5E. + +**7. Docling and DoclingV2 PDF backends deferred — not migrated in Phase 5D.** +No `core/indexing/parsers/pdf/docling*` modules will be created in this +pass. Legacy `DoclingLoader` and `DoclingLoader2` stay where they are +for now. +- Why: This is a PDF backend we haven't used or tested recently — + porting it now would pin a stale integration into the new layer. We'll + revisit and re-port it (or drop it) in a later pass once the refactor + has shaken out and we know whether Docling is still wanted. +- Alternative considered: port now alongside Marker / OpenAI / DotsOCR + for completeness. Rejected — moves dead-feeling code into the new + layer without verifying it still works. +- Follow-up: revisit during a later parser-coverage sweep. If the + decision is to drop, the legacy modules get deleted in Phase 5E rather + than shimmed. + +**8. `ImageBlock` is the parser↔caption contract — captioning is a downstream stage's job.** +- Every parser (Image, Markdown, Docx, Pptx, Eml, Marker, + `DotsOCRPdfClient`) emits `ImageBlock` with `caption=None`. The + caption stage fills it in. For VLM-PDF specifically, the picture-bbox + crop becomes an `ImageBlock(image_bytes=…, page_number=N)` — the + parser never issues the second VLM call. One uniform contract beats + per-parser carve-outs; the chunker sees the same `ImageBlock` shape + from every parser, including `DotsOCRPdfClient`. +- `ImageBlock.metadata['markdown_ref']` holds the in-text placeholder + (data-URI, `![](pptx-image-N)`, `![](docx-image-N)`, + `![](marker-key)`); the caption stage `str.replace`s it. No + placeholder ⇒ no `markdown_ref` ⇒ caption stage emits a + free-standing `TextBlock`. Contract is documented on `ImageBlock` + itself. +- `ImageBlock` carries `image_bytes` (default `b""`) AND `source_url`. + Locally-extracted images set bytes; HTTP refs (`![alt](https://…)`) + leave bytes empty and set `source_url`. The `image_url` property + returns `data:{mime};base64,…` when bytes are present, else + `source_url` — consumers read `image_url` regardless of shape. +- Why: Refs are per-image-unique and chunk-stable. Legacy + `MarkdownLoader` captioned HTTP images via langchain `ChatOpenAI` + (which accepts URLs natively). The new VLM ABC takes bytes only, so a + fetch stage has to populate them — but the parser still emits one + `ImageBlock` per in-text image, keeping the contract uniform. +- Alternatives considered: (a) positional matching of refs to images — + rejected as fragile; (b) embedding image bytes inside `TextBlock` — + rejected as a heavier model change. + +**9. Paginated parsers emit `list[TextBlock]` with `page_number`; in-band `[PAGE_N]` markers are gone.** +Marker and PPTX previously concatenated all page content into one +`TextBlock` with `[PAGE_N]` markers between pages. They now emit one +`TextBlock` per page with `page_number` set, matching what PyMuPDF +already does. Parsers without natural pagination +(text/html/md/docx/doc/eml/whisper/image) still emit a single +`page_number=1` block. +- Why: Pagination is metadata, not content. Leaking `[PAGE_N]` markers + into chunk text forced every consumer to know the marker syntax; + `TextBlock.page_number` is the canonical channel and was already + half-used. +- Implication for chunking: the chunker must NOT scan for `[PAGE_N]` + markers. Iterate `ProcessedDocument.text_blocks` and carry + `block.page_number` onto every emitted chunk. Page boundaries are + block boundaries. + +**10. Client-backed parsers: generic `Client*Parser` facades; `BaseOpenAIPdfClient` is scaffolding only.** +- Renamed `OpenAIPdfParser` → `ClientPdfParser` + (`core/indexing/parsers/pdf/openai.py` → `pdf/client_based.py`); added + `ClientAudioParser` at `core/indexing/parsers/audio/client_based.py`. + Both accept any `BaseClientParser` and delegate `parse()`. "OpenAI" + was a leaky model-specific label on a class that takes any + HTTP-client-backed parser; whatever DotsOCR / Whisper-vLLM / + Scaleway-Speech is called next quarter, the facade stays the same — + what varies is the injected `BaseClientParser`. +- `BaseOpenAIPdfClient` provides reusable helpers (PDF page rendering, + semaphore-protected `_ocr_one(page_img, prompt) → str | None`, + JSON-fence stripping, JSON loading, picture-bbox cropping). It does + **NOT** define `parse()`, a `PROMPT` class attribute, or abstract + `_caption_images` / `_result_to_md` / `_parse_ocr_response` hooks. + The file was renamed `_openai.py` → `_base_openai_parser.py` to + match the new role. +- Why: The previous abstract pipeline imposed assumptions ("there's one + OCR response per page", "captioning is a parser concern") that didn't + generalise. Treat the base as a toolbox; let each concrete client + (DotsOCR, future variants) drive its own `parse()` and block-emission + strategy. +- Trade-off: more code per concrete subclass. Accepted — + model-specific variation (response schema, block layout, bbox + handling) lives in the subclass anyway. +- Alternative considered: keep one model-specific facade per backend. + Rejected — duplicates the same isinstance + delegate boilerplate. + +**11. DotsOCR response is validated through Pydantic.** +`DotsOCRElement` / `DotsOCRPage(RootModel[list[DotsOCRElement]])` / +`DotsOCRCategory` (Enum) capture the layout-element shape; +`DotsOCRPdfClient._parse_page` runs `model_validate` and returns `None` +on bad payloads. The `{"items": [...]}` envelope is tolerated alongside +a bare list. +- Why: Replaces dict shuffling (`page_res.get("category") == "Picture"`, + `item.get("bbox")`) with typed access (`element.category is + DotsOCRCategory.PICTURE`, `element.bbox`). Bad payloads fail loudly + via `ValidationError` instead of silently returning empty markdown. + +**12. `OpenAIAudioClient` keeps language detection as an injected callable, not a Ray ref-getter.** +Legacy `AudioTranscriber` looked up a `WhisperActor` Ray actor by name. +The new `OpenAIAudioClient` takes `language_detector: Callable[[Path], +Awaitable[str | None]] | None` in its constructor and skips detection +when `None` (vLLM auto-detects). +- Why: Keep the client free of Ray coupling so it can be instantiated + and tested without a Ray cluster. The wiring layer passes a closure + that calls the Whisper actor when `USE_WHISPER_LANG_DETECTOR=true`. +- Alternative considered: keep the Ray actor lookup inside the client + guarded by a config flag. Rejected — pulls Ray into the + `services/inference` layer where the rest of the file is plain HTTP. + +--- + +## Phase 5E — Loader → Parser shims (2026-05-06) + +**1. Legacy loaders are *adapter* shims, not re-export shims.** +The earlier compat-shim pass (commit `93476a6`) used pure `from X +import Y` re-exports because the symbols moved unchanged +(`ray_utils`, `text_sanitizer`, exceptions). The loader→parser move +can't do that: `BaseLoader.aload_document(file_path) → langchain +Document` and `DocumentParser.parse(document) → ProcessedDocument` +have different names *and* different contracts. Each legacy loader +becomes a `BaseLoader` adapter that reads the file into bytes, builds +a `CoreDocument`, calls the new parser, and maps `ProcessedDocument` +back to a langchain `Document`. +- Why: Preserves dynamic loader-discovery + (`BaseLoader.__subclasses__()` in `loaders/__init__.py`) and the + config-string lookup (`file_loaders.pdf: "MarkerLoader"`) without + forcing every consumer to migrate at once. +- Alternative considered: pure re-exports aliasing `*Parser` as + `*Loader`. Rejected — the discovery walk only finds `BaseLoader` + subclasses, so an aliased `DocumentParser` would silently disappear + from the loader registry. + +**2. Shimmed in this pass: text/markdown, image, docx, doc, pptx, pymupdf, marker, local-whisper, openai-audio.** +Each adapter delegates to its core parser and, when the parser emits +`ImageBlock`s with `markdown_ref` set, layers VLM captioning on top +via the existing `BaseLoader` mixin (`self.image_captioning`, +`self.caption_images`, `self.replace_markdown_images_with_captions`). +- Why: Keeps the legacy contract intact (captioned markdown in + `page_content`) while the canonical home is the parser. The + `markdown_ref` substitution path is the same one the future + caption-stage will use. + +**3. `base.py` Stage 1: re-export the four image_preprocessor symbols already in core, leave the captioning mixin in place.** +`ensure_png_compatible_mode`, `HTTP_IMAGE_PATTERN`, +`DATA_URI_IMAGE_PATTERN`, `MIN_IMAGE_PIXELS` now point at the +canonical `core.indexing.image_preprocessor` symbols (class attrs +hold module-level references for `self.X` access). +`_pil_image_to_base64` rewritten on top of `pil_to_png_bytes`. The +VLM endpoint setup, `get_image_description`, `caption_images`, +`replace_markdown_images_with_captions` stay in `base.py` for now. +- Why: Mechanical, behavior-identical change. Stage 2 (move VLM + captioning to `services/inference/captioning`) needs a design call + (where it lives, how the shim acquires it) and is deferred. + +**4. `PyMuPDFParser`: single dedicated thread + retain empty pages for 1-to-1 pagination.** +- PyMuPDF/pymupdf4llm are not thread-safe; concurrent calls raise + `ValueError: not a textpage of this page`. Upstream maintainer + (`pymupdf/PyMuPDF#3771`, closed wontfix) confirms this is documented + behaviour, not a bug. The parser now uses a module-level + `ThreadPoolExecutor(max_workers=1)` instead of `asyncio.to_thread`; + concurrent `parse()` calls queue on the executor, eliminating the + race against the default thread pool. The rest of the indexing + pipeline still parallelizes — only the pymupdf step is serialized. +- Empty pages now produce a `TextBlock` with empty `text` (was + previously dropped while keeping `page_count` accurate). Reverted so + every page produces a `TextBlock`, keeping a 1-to-1 mapping with the + source PDF's pagination — the legacy `\n[PAGE_N]\n` anchor format + the loader-shim emits aligns exactly with the source. + +**5. `TranscriberConfig.direct_upload_suffixes` got lost in the core/config migration; ported to `core/config/indexation.py`.** +The legacy `config/models.py:TranscriberConfig` had the field + +`|`-separated string validator + a default frozenset of audio +extensions. The active `core/config/indexation.py:TranscriberConfig` +(loaded via `openrag.core.config.loader.load_config`) was missing it, +producing `AttributeError: 'TranscriberConfig' object has no attribute +'direct_upload_suffixes'` when the audio shim accessed it. +- Why: `config/models.py` is now vestigial — kept for legacy imports + but no longer drives `load_config()`. Fields added there but not + mirrored to `core/config` are silently inactive at runtime. + +**6. Skipped: eml, `pdf_loaders/openai.py`, `pdf_loaders/dotsocr.py`.** +- `eml_loader.py`: the new `EmlParser` takes `attachment_parsers: + Mapping[str, DocumentParser]`, but the old loader dispatches + attachments through `BaseLoader`-keyed `get_loader_classes` with a + multi-tier PDF fallback chain (`MarkerLoader` → `PyMuPDFLoader` → + `PyMuPDF4LLMLoader` → `DoclingLoader`). The contract bridge isn't + trivial; deferred until services-side attachment-parser composition + lands. +- `pdf_loaders/openai.py` + `pdf_loaders/dotsocr.py`: services-side + `BaseOpenAIPdfClient` / `DotsOCRPdfClient` exist but require a + concrete `core.vlm.VLM` to instantiate, and `vlm_registry` is empty + (no concrete VLM impl exists yet). Both legacy classes are also dead + code on this branch — not in any Hydra config, no external imports. +- Why: Both gaps need new services-side work (attachment-parser DI, + `LangchainOpenAIVLM`-style concrete) before a meaningful shim is + possible. Re-export-only "shims" would relocate the file without + going through the new architecture, defeating the purpose. + +**7. Stale files flagged for deletion (Phase 12 cleanup).** +- `components/indexer/loaders/CustomHTMLLoader.py` and + `components/indexer/loaders/CustomDocLoader.py` — legacy + `BaseLoader` subclasses, not referenced by any Hydra config or + external import. Discoverable via `BaseLoader.__subclasses__()` but + never instantiated. `CustomDocLoader` uses + `UnstructuredWordDocumentLoader` / `UnstructuredODTLoader` — no + clean parser equivalent in core (`DocxParser` uses MarkItDown). +- `config/models.py` (the whole file, incl. its `TranscriberConfig`) + — superseded by `core/config/*`; kept only so legacy imports don't + break. Drift between the two has already caused one runtime bug + (entry 5). +- Why: Out of scope for the loader-shim pass; flagged here so they + don't get re-shimmed by future passes. Removal coordinates with + Phase 12 ("delete old re-export shims"). --- diff --git a/openrag/components/indexer/loaders/audio/local_whisper.py b/openrag/components/indexer/loaders/audio/local_whisper.py index 7768cfaa0..0005c3aa3 100644 --- a/openrag/components/indexer/loaders/audio/local_whisper.py +++ b/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 diff --git a/openrag/components/indexer/loaders/audio/openai.py b/openrag/components/indexer/loaders/audio/openai.py index 8b7f42173..0d3748ec8 100644 --- a/openrag/components/indexer/loaders/audio/openai.py +++ b/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 diff --git a/openrag/components/indexer/loaders/base.py b/openrag/components/indexer/loaders/base.py index 70f05e394..647e06fd8 100644 --- a/openrag/components/indexer/loaders/base.py +++ b/openrag/components/indexer/loaders/base.py @@ -2,11 +2,23 @@ import base64 import re from abc import ABC, abstractmethod -from io import BytesIO from pathlib import Path from components.prompts import IMAGE_DESCRIBER from components.utils import get_vlm_semaphore, load_config +from core.indexing.image_preprocessor import ( + DATA_URI_IMAGE_PATTERN as _CORE_DATA_URI_IMAGE_PATTERN, +) +from core.indexing.image_preprocessor import ( + HTTP_IMAGE_PATTERN as _CORE_HTTP_IMAGE_PATTERN, +) +from core.indexing.image_preprocessor import ( + MIN_IMAGE_PIXELS as _CORE_MIN_IMAGE_PIXELS, +) +from core.indexing.image_preprocessor import ( + ensure_png_compatible_mode, # noqa: F401 (re-exported for legacy import path) + pil_to_png_bytes, +) from langchain_core.messages import HumanMessage from langchain_openai import ChatOpenAI from openai import BadRequestError @@ -19,20 +31,13 @@ config = load_config() -def ensure_png_compatible_mode(image: Image.Image) -> Image.Image: - """Convert incompatible PIL image modes to PNG-saveable modes.""" - if image.mode in ("CMYK", "YCbCr", "LAB"): - return image.convert("RGB") - if image.mode in ("P", "LA", "PA"): - return image.convert("RGBA") - return image - - class BaseLoader(ABC): - # Class-level compiled regex patterns (shared across all instances) - HTTP_IMAGE_PATTERN = re.compile(r"!\[(.*?)\]\((https?://[^)]+)\)") - DATA_URI_IMAGE_PATTERN = re.compile(r"!\[(.*?)\]\((data:image/[^;]+;base64,[^)]+)\)") - MIN_IMAGE_PIXELS = 784 # Qwen2.5-VL min_pixels threshold + # Class-level compiled regex patterns and constants — single source of truth + # lives in ``core.indexing.image_preprocessor``; pinned here as class attrs so + # subclasses keep working via ``self.X``. + HTTP_IMAGE_PATTERN = _CORE_HTTP_IMAGE_PATTERN + DATA_URI_IMAGE_PATTERN = _CORE_DATA_URI_IMAGE_PATTERN + MIN_IMAGE_PIXELS = _CORE_MIN_IMAGE_PIXELS def __init__(self, **kwargs) -> None: self.page_sep = "[PAGE_SEP]" @@ -68,14 +73,12 @@ def save_content(self, text_content: str, path: str): def _pil_image_to_base64(self, image: Image.Image) -> str: """Convert PIL Image to base64 string.""" - buffered = BytesIO() try: - image = ensure_png_compatible_mode(image) - image.save(buffered, format="PNG") + png_bytes = pil_to_png_bytes(image) except Exception as e: logger.warning("Failed to convert image to PNG", error=str(e), mode=getattr(image, "mode", "unknown")) return "" - return base64.b64encode(buffered.getvalue()).decode() + return base64.b64encode(png_bytes).decode() def _is_http_url(self, data: str) -> bool: """Check if string is an HTTP/HTTPS URL.""" diff --git a/openrag/components/indexer/loaders/doc.py b/openrag/components/indexer/loaders/doc.py index 387b46543..7ca0f909f 100644 --- a/openrag/components/indexer/loaders/doc.py +++ b/openrag/components/indexer/loaders/doc.py @@ -1,46 +1,77 @@ +""" +Legacy ``.doc`` file loader implementation. + +``DocLoader`` is now a thin :class:`BaseLoader` adapter that delegates +to :class:`core.indexing.parsers.doc_parser.DocParser` (which itself +runs Spire.Doc → .docx conversion and then ``DocxParser``) and layers +VLM captioning of embedded images on top. New code should call the +core parser directly; this shim keeps the legacy loader-discovery path +alive until consumers migrate. +""" + +import asyncio import os -import tempfile +from io import BytesIO +from pathlib import Path +from core.indexing.parsers.doc_parser import DocParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document as LCDocument -from spire.doc import Document, FileFormat +from PIL import Image from utils.logger import get_logger from .base import BaseLoader from .docx import DocxLoader -os.environ["DOTNET_SYSTEM_GLOBALIZATION_INVARIANT"] = "1" # Disable Globalization +os.environ["DOTNET_SYSTEM_GLOBALIZATION_INVARIANT"] = "1" logger = get_logger() class DocLoader(BaseLoader): + """Adapter shim — delegates to ``DocParser``; layers image captioning on top.""" + def __init__(self, **kwargs) -> None: super().__init__(**kwargs) self.MDLoader = DocxLoader(**kwargs) + self._parser = DocParser() + async def aload_document(self, file_path, metadata, save_markdown=False): - """Convert .doc to .docx format, then use DocxLoader to convert to markdown. - Falls back to plain text extraction if the .docx conversion fails.""" - temp_path = None - document = Document() - try: - document.LoadFromFile(str(file_path)) - with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as temp_file: - temp_path = temp_file.name - document.SaveToFile(temp_path, FileFormat.Docx2016) - except Exception as e: - logger.bind(file_id=metadata.get("file_id"), partition=metadata.get("partition")).warning( - f"Spire.Doc conversion to .docx failed, falling back to text extraction: {e}" - ) - text = document.GetText() - doc = LCDocument(page_content=text, metadata=metadata) - if save_markdown: - self.save_content(text, str(file_path)) - return doc + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.DOC, + raw_bytes=raw_bytes, + metadata=dict(metadata) if metadata else {}, + ) + processed = await self._parser.parse(core_doc) + result = "\n\n".join(b.text for b in processed.text_blocks).strip() + + if not self.image_captioning: + logger.info("Image captioning disabled. Ignoring images.") else: - result = await self.MDLoader.aload_document(temp_path, metadata, save_markdown) - return result - finally: - document.Close() - if temp_path and os.path.exists(temp_path): - os.remove(temp_path) + if processed.images: + pil_images: list[Image.Image] = [] + for block in processed.images: + img = Image.open(BytesIO(block.image_bytes)) + img.load() + pil_images.append(img) + captions = await self.caption_images(pil_images, desc="Captioning embedded images") + for block, caption in zip(processed.images, captions): + ref = (block.metadata or {}).get("markdown_ref") + if ref: + result = result.replace(ref, caption.replace("\\", "/")) + + result = await self.replace_markdown_images_with_captions( + result, + caption_data_uris=False, + desc="Captioning linked images", + ) + + doc = LCDocument(page_content=result, metadata=dict(metadata) if metadata else {}) + if save_markdown: + self.save_content(result, str(file_path)) + return doc diff --git a/openrag/components/indexer/loaders/docx.py b/openrag/components/indexer/loaders/docx.py index 615f58abf..6692cac20 100644 --- a/openrag/components/indexer/loaders/docx.py +++ b/openrag/components/indexer/loaders/docx.py @@ -1,10 +1,24 @@ -import re +""" +DOCX file loader implementation. + +``DocxLoader`` is now a thin :class:`BaseLoader` adapter that delegates +extraction to :class:`core.indexing.parsers.docx_parser.DocxParser` and +layers VLM captioning of embedded images on top via the ``BaseLoader`` +mixin. The legacy ``convert_to_png_image`` helper and the +``get_images_from_zip`` instance method are preserved for backward +compatibility with existing test consumers; new code should use the +core parser directly. +""" + +import asyncio import zipfile from io import BytesIO +from pathlib import Path -from docx import Document as DocxDocument +from core.indexing.parsers.docx_parser import DocxParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document -from markitdown import MarkItDown from PIL import Image from utils.logger import get_logger @@ -18,68 +32,59 @@ def convert_to_png_image(image: Image.Image) -> Image.Image: with BytesIO() as buffer: image.save(buffer, format="PNG") buffer.seek(0) - # Reload the image from the buffer as a PNG png_image = Image.open(buffer).convert("RGBA") return png_image class DocxLoader(BaseLoader): + """Adapter shim — delegates to ``DocxParser``; layers image captioning on top.""" + def __init__(self, **kwargs): super().__init__(**kwargs) - self.converter = MarkItDown() + self._parser = DocxParser() async def aload_document(self, file_path, metadata, save_markdown=False): - try: - result = self.converter.convert(file_path).text_content - except Exception as markitdown_err: - logger.warning( - "MarkItDown conversion failed, falling back to python-docx plain text extraction", - path=str(file_path), - error=str(markitdown_err), - ) - try: - result = self._fallback_extract_text(file_path) - except Exception as docx_err: - raise RuntimeError( - f"DOCX conversion failed with both MarkItDown ({markitdown_err}) and python-docx ({docx_err})" - ) from docx_err - - if self.image_captioning: - # Handle embedded images (extracted from docx zip) - # images may contain None entries for unsupported formats (e.g. EMF, WMF) - images = self.get_images_from_zip(file_path) - valid_images = [img for img in images if img is not None] - captions = await self.caption_images(valid_images, desc="Captioning embedded images") - - # Rebuild caption list preserving positional alignment with markdown refs - caption_iter = iter(captions) - for img in images: - caption = next(caption_iter) if img is not None else "" - result = re.sub( - r"!\[.*?\]\(data:image/.*?\)", - caption.replace("\\", "/") if caption else "", - string=result, - count=1, - ) - - # Handle linked images (HTTP URLs) using shared method - # Only caption HTTP URLs, data URIs are already handled above + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.DOCX, + raw_bytes=raw_bytes, + metadata=dict(metadata) if metadata else {}, + ) + processed = await self._parser.parse(core_doc) + result = "\n\n".join(b.text for b in processed.text_blocks).strip() + + if not self.image_captioning: + logger.info("Image captioning disabled. Ignoring images.") + for block in processed.images: + ref = (block.metadata or {}).get("markdown_ref") + if ref: + result = result.replace(ref, "") + else: + pil_images: list[Image.Image] = [] + for block in processed.images: + img = Image.open(BytesIO(block.image_bytes)) + img.load() + pil_images.append(img) + captions = await self.caption_images(pil_images, desc="Captioning embedded images") + for block, caption in zip(processed.images, captions): + ref = (block.metadata or {}).get("markdown_ref") + if ref: + result = result.replace(ref, caption.replace("\\", "/")) + result = await self.replace_markdown_images_with_captions( result, caption_data_uris=False, desc="Captioning linked images", ) - else: - logger.info("Image captioning disabled. Ignoring images.") - doc = Document(page_content=result, metadata=metadata) + doc = Document(page_content=result, metadata=dict(metadata) if metadata else {}) if save_markdown: self.save_content(result, str(file_path)) return doc - def _fallback_extract_text(self, file_path) -> str: - doc = DocxDocument(file_path) - return "\n\n".join(p.text for p in doc.paragraphs if p.text.strip()) + # ----- legacy helpers retained for test_docx_loader.py compatibility ----- def get_images_from_zip(self, input_file): try: @@ -89,16 +94,11 @@ def get_images_from_zip(self, input_file): return [] with docx: file_names = docx.namelist() - # word/media/ may also contain non-image files (e.g. oleObject, hdphoto, ink) image_files = [f for f in file_names if f.startswith("word/media/")] if not image_files: return [] images_not_in_order, order = [], [] - - # the images got from the original file is not in the right order - # but the target_ref contains the position of the image in the document - for image_file in image_files: image_data = docx.read(image_file) image_extension = image_file.split(".")[-1].lower() @@ -116,7 +116,6 @@ def get_images_from_zip(self, input_file): if not images_not_in_order: return [] - # Reorder images by their original position in the document max_order = max(order) images = [None] * max_order for i, pos in enumerate(order): diff --git a/openrag/components/indexer/loaders/image.py b/openrag/components/indexer/loaders/image.py index a1356c802..d22bd8f7f 100644 --- a/openrag/components/indexer/loaders/image.py +++ b/openrag/components/indexer/loaders/image.py @@ -1,7 +1,22 @@ +""" +Image file loader implementation. + +``ImageLoader`` is now a thin :class:`BaseLoader` adapter that delegates +decode (raster + SVG) to +:class:`core.indexing.parsers.image_parser.ImageParser` and then layers +VLM captioning on top via the ``BaseLoader`` mixin. The legacy +``ImageLoadError`` contract is preserved on decode failure. New code +should call the core parser directly; this shim keeps the legacy +loader-discovery path alive until consumers migrate. +""" + +import asyncio from io import BytesIO from pathlib import Path -import cairosvg +from core.indexing.parsers.image_parser import ImageParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents import Document from PIL import Image from utils.logger import get_logger @@ -18,27 +33,41 @@ class ImageLoadError(Exception): class ImageLoader(BaseLoader): def __init__(self, **kwargs): super().__init__(**kwargs) + # ``min_pixels=0`` so the parser does not drop small images; the + # size threshold is enforced by ``get_image_description`` (which + # returns the legacy "Image too small for captioning" marker). + self._parser = ImageParser(min_pixels=0) async def aload_document(self, file_path, metadata=None, save_markdown=False): - path = Path(file_path) + if metadata is None: + metadata = {} + path = Path(file_path) try: - # Handle SVG files by converting to PNG first - if path.suffix.lower() == ".svg": - png_data = cairosvg.svg2png(url=str(path)) - img = Image.open(BytesIO(png_data)) - else: - img = Image.open(path) + raw_bytes = await asyncio.to_thread(path.read_bytes) except Exception as e: log.error( - "Failed to load image file", + "Failed to read image file", file_path=str(path), error_type=type(e).__name__, error=str(e), ) raise ImageLoadError(f"Cannot load image '{path.name}': {type(e).__name__}") from e + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.IMAGE, + raw_bytes=raw_bytes, + metadata=metadata, + ) + processed = await self._parser.parse(core_doc) + if not processed.images: + raise ImageLoadError(f"Cannot load image '{path.name}': failed to decode") + + img = Image.open(BytesIO(processed.images[0].image_bytes)) + img.load() description = await self.get_image_description(image_data=img) + doc = Document(page_content=description, metadata=metadata) if save_markdown: self.save_content(description, str(path)) diff --git a/openrag/components/indexer/loaders/pdf_loaders/marker.py b/openrag/components/indexer/loaders/pdf_loaders/marker.py index 62e136dc1..03cb0edbe 100644 --- a/openrag/components/indexer/loaders/pdf_loaders/marker.py +++ b/openrag/components/indexer/loaders/pdf_loaders/marker.py @@ -1,306 +1,52 @@ +""" +Marker-backed PDF loader. + +The Ray actor + pool that drive Marker (``MarkerWorker``, +``MarkerPool``) and the services-side :class:`BasePooledParser` +implementation now live in +``services/workers/parsers/marker_workers.py``; this module re-exports +``MarkerWorker`` and ``MarkerPool`` for legacy import paths +(``utils.dependencies`` constructs the named ``MarkerPool`` actor at +startup via ``get_or_create_actor``). + +``MarkerLoader`` is now a thin :class:`BaseLoader` adapter that +delegates to :class:`core.indexing.parsers.pdf.marker.MarkerParser`, +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 -import gc -import re import time +from io import BytesIO from pathlib import Path -import pypdfium2 -import ray -import torch -from config import load_config +from core.indexing.parsers.pdf.marker import MarkerParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document -from marker.converters.pdf import PdfConverter +from PIL import Image +from services.workers.parsers.marker_workers import ( # noqa: F401 (re-exported for legacy import paths) + MarkerLoader as _ServicesMarkerPool, +) +from services.workers.parsers.marker_workers import ( # noqa: F401 + MarkerPool, + MarkerWorker, +) from utils.logger import get_logger from ..base import BaseLoader logger = get_logger() -config = load_config() - -if torch.cuda.is_available(): - MARKER_NUM_GPUS = config.loader.marker_num_gpus -else: # On CPU - MARKER_NUM_GPUS = 0 - - -@ray.remote(num_gpus=MARKER_NUM_GPUS, max_restarts=5) -class MarkerWorker: - def __init__(self): - import os - - from config import load_config - from utils.logger import get_logger - - self.logger = get_logger() - self.config = load_config() - self.page_sep = "[PAGE_SEP]" - - self._workers = self.config.loader.marker_max_processes - - self.converter_config = { - "output_format": "markdown", - "paginate_output": True, - "page_separator": self.page_sep, - "pdftext_workers": self.config.loader.marker_pdftext_workers, - "disable_multiprocessing": False, - } - os.environ["RAY_ADDRESS"] = "auto" - - self.executor = None - self.init_resources() - - def init_resources(self): - from marker.models import create_model_dict - - self.model_dict = create_model_dict() - for v in self.model_dict.values(): - if hasattr(v.model, "share_memory"): - v.model.share_memory() - - self.setup_mp() - - def setup_mp(self): - """Initialize ProcessPoolExecutor for PDF processing. - - We use ProcessPoolExecutor instead of multiprocessing.Pool because: - - Ray actors run as daemon processes - - Pool workers are daemonic by default and cannot spawn children - - The pdftext library (used by Marker) internally spawns processes - - ProcessPoolExecutor workers are non-daemon, allowing nested process creation - """ - from concurrent.futures import ProcessPoolExecutor - - import torch.multiprocessing as mp - - if self.executor: - self.logger.warning("Resetting ProcessPoolExecutor") - self.executor.shutdown(wait=False, cancel_futures=True) - self.executor = None - - # Ensure spawn method for CUDA compatibility - try: - if mp.get_start_method(allow_none=True) != "spawn": - mp.set_start_method("spawn", force=True) - except RuntimeError: - self.logger.warning("Process start method already set, using existing method") - - self.logger.info(f"Initializing MarkerWorker with {self._workers} workers") - self.executor = ProcessPoolExecutor( - max_workers=self._workers, - initializer=self._worker_init, - initargs=(self.model_dict,), - mp_context=mp.get_context("spawn"), - max_tasks_per_child=self.config.loader.marker_max_tasks_per_child, - ) - self.logger.info("MarkerWorker initialized with ProcessPoolExecutor") - - @staticmethod - def _worker_init(model_dict): - global worker_model_dict - worker_model_dict = model_dict - logger.debug("Worker initialized with model dictionary") - - @staticmethod - def _process_pdf(file_path, config): - global worker_model_dict - - page_range = config.get("page_range") - if page_range is not None: - label = f"[p{page_range[0]}-{page_range[-1]}]" - else: - label = "(all pages)" - - try: - logger.debug("Processing PDF", path=file_path, label=label) - converter = PdfConverter( - artifact_dict=worker_model_dict, - config=config, - ) - render = converter(file_path) - return render - except Exception as e: - logger.exception("Error processing PDF", path=file_path, label=label, error=str(e)) - raise - finally: - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() - - async def process_pdf(self, file_path: str, page_range: list[int] | None = None): - from concurrent.futures import TimeoutError as FuturesTimeoutError - - converter_config = self.converter_config.copy() - if page_range is not None: - converter_config["page_range"] = page_range - - loop = asyncio.get_event_loop() - timeout = self.config.loader.marker_timeout - - def run_with_timeout(): - future = self.executor.submit(self._process_pdf, file_path, converter_config) - try: - result = future.result(timeout=timeout) - return result - except FuturesTimeoutError: - self.logger.exception("MarkerWorker child process timed out", path=file_path) - raise - except Exception: - self.logger.exception("Error processing with MarkerWorker", path=file_path) - raise - - result = await loop.run_in_executor(None, run_with_timeout) - return result.markdown, result.images - - def is_pool_broken(self): - # ProcessPoolExecutor auto-replaces dead/finished workers on next - # submit(), so counting live processes is unreliable and unnecessary. - # Only a None or shut-down executor requires reinitialization. - return self.executor is None or bool(getattr(self.executor, "_broken", False)) - - def __del__(self): - """Clean up ProcessPoolExecutor on actor destruction""" - if self.executor: - try: - self.executor.shutdown(wait=False, cancel_futures=True) - except Exception: - pass # Best effort cleanup - - -@ray.remote(max_restarts=5) -class MarkerPool: - def __init__(self): - from config import load_config - from utils.logger import get_logger - - self.logger = get_logger() - self.config = load_config() - self.max_processes = self.config.loader.marker_max_processes - self.pool_size = self.config.loader.marker_pool_size - self.actors = [MarkerWorker.remote() for _ in range(self.pool_size)] - self._queue: asyncio.Queue[ray.actor.ActorHandle] = asyncio.Queue() - - for _ in range(self.max_processes): - for actor in self.actors: - self._queue.put_nowait(actor) - - self.logger.info( - f"Marker pool: {self.pool_size} actors × {self.max_processes} slots = " - f"{self.pool_size * self.max_processes} PDF concurrency" - ) - - @staticmethod - def _get_page_count(file_path: str) -> int: - pdf = pypdfium2.PdfDocument(file_path) - try: - return len(pdf) - finally: - pdf.close() - - @staticmethod - def _create_chunks(page_count: int, chunk_size: int) -> list[tuple[list[int], str]]: - if page_count <= chunk_size: - return [(list(range(page_count)), f"({page_count}p)")] - chunks = [] - for start in range(0, page_count, chunk_size): - end = min(start + chunk_size, page_count) - page_range = list(range(start, end)) - label = f"[p{start}-{end - 1}]" - chunks.append((page_range, label)) - return chunks - - async def ensure_worker_pool_healthy(self, worker): - from components.ray_utils import call_ray_actor_with_timeout - - timeout = self.config.loader.marker_timeout - broken = await call_ray_actor_with_timeout( - worker.is_pool_broken.remote(), - timeout=timeout, - task_description="MarkerWorker pool health check", - ) - if broken: - self.logger.warning("Worker ProcessPoolExecutor is broken. Reinitializing pool...") - await call_ray_actor_with_timeout( - worker.setup_mp.remote(), - timeout=timeout, - task_description="MarkerWorker pool reset", - ) - - async def _process_chunk(self, file_path: str, page_range: list[int] | None, label: str): - """Acquire a worker slot, process a PDF chunk, and release the slot. - - Retries on failure with exponential backoff up to marker_max_task_retry times. - A fresh worker is acquired per attempt so a flaky worker can be sidestepped - and ensure_worker_pool_healthy re-runs each time. - """ - from components.ray_utils import call_ray_actor_with_timeout, retry_with_backoff - - timeout = self.config.loader.marker_timeout - - async def attempt(i: int): - worker = await self._queue.get() - try: - self.logger.info(f"MarkerWorker allocated for {label} (attempt {i + 1})") - await self.ensure_worker_pool_healthy(worker) - future = worker.process_pdf.remote(file_path, page_range=page_range) - return await call_ray_actor_with_timeout( - future, - timeout=timeout, - task_description=f"MarkerPool PDF {label} ({file_path})", - ) - finally: - await self._queue.put(worker) - self.logger.debug(f"MarkerWorker returned to pool for {label}") - - return await retry_with_backoff( - attempt, - max_retries=self.config.loader.marker_max_task_retry, - base_delay=self.config.loader.marker_retry_base_delay, - task_description=f"MarkerPool PDF {label} ({file_path})", - ) - - async def process_pdf(self, file_path: str): - chunk_size = self.config.loader.marker_chunk_size - - if chunk_size <= 0: - return await self._process_chunk(file_path, page_range=None, label="(all pages)") - - page_count = self._get_page_count(file_path) - chunks = self._create_chunks(page_count, chunk_size) - - if len(chunks) == 1: - page_range, label = chunks[0] - return await self._process_chunk(file_path, page_range=None, label=label) - - self.logger.info( - f"Splitting {page_count}-page PDF into {len(chunks)} chunks of ~{chunk_size} pages for parallel processing" - ) - - tasks = [asyncio.create_task(self._process_chunk(file_path, page_range, label)) for page_range, label in chunks] - try: - results = await asyncio.gather(*tasks) - except Exception: - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - raise - - # Reassemble: concatenate markdown in order, merge image dicts - all_markdown = [] - all_images = {} - for markdown, images in results: - all_markdown.append(markdown) - all_images.update(images) - - combined_markdown = "\n\n".join(all_markdown) - return combined_markdown, all_images class MarkerLoader(BaseLoader): + """Adapter shim — delegates to ``MarkerParser`` via the services-side pool.""" + def __init__(self, **kwargs): super().__init__(**kwargs) self.page_sep = "[PAGE_SEP]" - self.worker = ray.get_actor("MarkerPool", namespace="openrag") + self._parser = MarkerParser(pool=_ServicesMarkerPool()) async def aload_document( self, @@ -308,41 +54,46 @@ async def aload_document( metadata: dict | None = None, save_markdown: bool = False, ) -> Document: - from components.ray_utils import call_ray_actor_with_timeout - if metadata is None: metadata = {} + path = Path(file_path) file_path_str = str(file_path) start = time.time() try: - timeout = self.config.loader.marker_timeout - future = self.worker.process_pdf.remote(file_path_str) - markdown, images = await call_ray_actor_with_timeout( - future, - timeout=timeout, - task_description=f"MarkerLoader PDF loading ({file_path_str})", + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.PDF, + raw_bytes=raw_bytes, + metadata=dict(metadata), ) + processed = await self._parser.parse(core_doc) + markdown = "".join(f"{b.text}\n[PAGE_{b.page_number}]\n" for b in processed.text_blocks) if not markdown: raise RuntimeError(f"Conversion failed for {file_path_str}") - if self.image_captioning: - keys = list(images.keys()) - captions = await self.caption_images(list(images.values())) - for key, caption in zip(keys, captions): - markdown = markdown.replace(f"![]({key})", caption) - - else: + if not self.image_captioning: logger.debug("Image captioning disabled.") - - markdown = markdown.split(self.page_sep, 1)[1] - markdown = re.sub(r"\{(\d+)\}" + re.escape(self.page_sep), r"[PAGE_\1]", markdown) - markdown = markdown.replace("
", "").strip() + for block in processed.images: + ref = (block.metadata or {}).get("markdown_ref") + if ref: + markdown = markdown.replace(ref, "") + elif processed.images: + pil_images: list[Image.Image] = [] + for block in processed.images: + img = Image.open(BytesIO(block.image_bytes)) + img.load() + pil_images.append(img) + captions = await self.caption_images(pil_images) + for block, caption in zip(processed.images, captions): + ref = (block.metadata or {}).get("markdown_ref") + if ref: + markdown = markdown.replace(ref, caption) doc = Document(page_content=markdown, metadata=metadata) - if save_markdown: self.save_content(markdown, file_path_str) diff --git a/openrag/components/indexer/loaders/pdf_loaders/pymupdf.py b/openrag/components/indexer/loaders/pdf_loaders/pymupdf.py index f00018b4f..c856c9523 100644 --- a/openrag/components/indexer/loaders/pdf_loaders/pymupdf.py +++ b/openrag/components/indexer/loaders/pdf_loaders/pymupdf.py @@ -1,25 +1,60 @@ +""" +PyMuPDF-backed PDF loader implementation. + +``PyMuPDFLoader`` and ``PyMuPDF4LLMLoader`` are now thin +:class:`BaseLoader` adapters that delegate to +:class:`core.indexing.parsers.pdf.pymupdf.PyMuPDFParser` (text and +markdown modes respectively). The markdown adapter additionally layers +VLM captioning of embedded images on top via the ``BaseLoader`` mixin. +New code should call the core parser directly; this shim keeps the +legacy loader-discovery path alive until consumers migrate. +""" + +import asyncio +from io import BytesIO from pathlib import Path -import pymupdf4llm -from langchain_community.document_loaders import PyMuPDFLoader as pymupdf_loader +from core.indexing.parsers.pdf.pymupdf import PyMuPDFParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document +from PIL import Image +from utils.logger import get_logger from ..base import BaseLoader +logger = get_logger() + + +def _join_pages_with_anchors(text_blocks) -> str: + """Join one ``TextBlock`` per page with the legacy ``\\n[PAGE_N]\\n`` anchors.""" + return "".join(f"{b.text}\n[PAGE_{b.page_number}]\n" for b in text_blocks) + + +async def _read_pdf_bytes(file_path) -> tuple[Path, bytes]: + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + return path, raw_bytes + class PyMuPDFLoader(BaseLoader): + """Adapter shim — delegates to ``PyMuPDFParser(mode='text')``.""" + def __init__(self, **kwargs): super().__init__(**kwargs) + self._parser = PyMuPDFParser(mode="text") async def aload_document(self, file_path, metadata: dict = None, save_markdown=False): - loader = pymupdf_loader( - file_path=Path(file_path), + metadata = {} if metadata is None else dict(metadata) + path, raw_bytes = await _read_pdf_bytes(file_path) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.PDF, + raw_bytes=raw_bytes, + metadata=metadata, ) - pages = await loader.aload() - - s = "" - for page_num, segment in enumerate(pages, start=1): - s += segment.page_content.strip() + f"\n[PAGE_{page_num}]\n" + processed = await self._parser.parse(core_doc) + s = _join_pages_with_anchors(processed.text_blocks) doc = Document(page_content=s, metadata=metadata) if save_markdown: @@ -28,15 +63,43 @@ async def aload_document(self, file_path, metadata: dict = None, save_markdown=F class PyMuPDF4LLMLoader(BaseLoader): + """Adapter shim — delegates to ``PyMuPDFParser(mode='markdown')``; layers image captioning on top.""" + def __init__(self, **kwargs) -> None: super().__init__(**kwargs) + self._parser = PyMuPDFParser(mode="markdown") async def aload_document(self, file_path, metadata: dict = None, save_markdown=False): - pages = pymupdf4llm.to_markdown(file_path, write_images=False, page_chunks=True) + metadata = {} if metadata is None else dict(metadata) + path, raw_bytes = await _read_pdf_bytes(file_path) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.PDF, + raw_bytes=raw_bytes, + metadata=metadata, + ) + processed = await self._parser.parse(core_doc) + s = _join_pages_with_anchors(processed.text_blocks) - s = "" - for page_num, segment in enumerate(pages, start=1): - s += segment.get("text").strip() + f"\n[PAGE_{page_num}]\n" + if not self.image_captioning: + # Legacy parity: the old loader called ``pymupdf4llm`` with the + # default ``embed_images=False`` and surfaced no images. The new + # parser embeds them as data URIs; strip those refs to match. + for block in processed.images: + ref = (block.metadata or {}).get("markdown_ref") + if ref: + s = s.replace(ref, "") + elif processed.images: + pil_images: list[Image.Image] = [] + for block in processed.images: + img = Image.open(BytesIO(block.image_bytes)) + img.load() + pil_images.append(img) + captions = await self.caption_images(pil_images, desc="Captioning embedded images") + for block, caption in zip(processed.images, captions): + ref = (block.metadata or {}).get("markdown_ref") + if ref: + s = s.replace(ref, caption.replace("\\", "/")) doc = Document(page_content=s, metadata=metadata) if save_markdown: diff --git a/openrag/components/indexer/loaders/pptx_loader.py b/openrag/components/indexer/loaders/pptx_loader.py index 69fa45be5..f91c75832 100644 --- a/openrag/components/indexer/loaders/pptx_loader.py +++ b/openrag/components/indexer/loaders/pptx_loader.py @@ -1,9 +1,20 @@ -import html -import re +""" +PPTX file loader implementation. + +``PPTXLoader`` is now a thin :class:`BaseLoader` adapter that delegates +extraction to :class:`core.indexing.parsers.pptx_parser.PptxParser` and +layers VLM captioning of slide pictures on top via the ``BaseLoader`` +mixin. New code should call the core parser directly; this shim keeps +the legacy loader-discovery path alive until consumers migrate. +""" + +import asyncio from io import BytesIO +from pathlib import Path -import pptx -from html_to_markdown import convert +from core.indexing.parsers.pptx_parser import PptxParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document from PIL import Image from utils.logger import get_logger @@ -13,149 +24,46 @@ logger = get_logger() -class PPTXConverter: - """Implementation based on PPTX converter in MarkItDown library. - - https://github.com/microsoft/markitdown/blob/main/packages/markitdown/src/markitdown/converters/_pptx_converter.py - """ - - def __init__(self, image_placeholder=r"", page_separator: str = "[PAGE_SEP]"): - self.image_placeholder = image_placeholder - self.page_separator = page_separator - - def convert(self, local_path): - md_content = "" - presentation = pptx.Presentation(local_path) - slide_num = 0 - images_list = [] - for slide in presentation.slides: - slide_num += 1 - - title = slide.shapes.title - for shape in slide.shapes: - if self._is_picture(shape): - images_list.append(Image.open(BytesIO(shape.image.blob))) - md_content += self.image_placeholder - - # Tables - if self._is_table(shape): - html_table = "" - first_row = True - for row in shape.table.rows: - html_table += "" - for cell in row.cells: - if first_row: - html_table += "" - else: - html_table += "" - html_table += "" - first_row = False - html_table += "
" + html.escape(cell.text) + "" + html.escape(cell.text) + "
" - md_content += "\n" + convert(html_table).strip() + "\n" - - # Charts - if shape.has_chart: - md_content += self._convert_chart_to_markdown(shape.chart) - - # Text areas - elif shape.has_text_frame: - if shape == title: - md_content += "# " + shape.text.lstrip() + "\n" - else: - md_content += shape.text + "\n" - - md_content = md_content.strip() - - if slide.has_notes_slide: - md_content += "\n\n### Notes:\n" - notes_frame = slide.notes_slide.notes_text_frame - if notes_frame is not None: - md_content += notes_frame.text - md_content = md_content.strip() - - md_content += f"\n[PAGE_{slide_num}]\n" - - return md_content, images_list - - def _is_picture(self, shape): - try: - if shape.shape_type == pptx.enum.shapes.MSO_SHAPE_TYPE.PICTURE: - return True - if shape.shape_type == pptx.enum.shapes.MSO_SHAPE_TYPE.PLACEHOLDER: - if hasattr(shape, "image"): - return True - except NotImplementedError: - # https://python-pptx.readthedocs.io/en/latest/_modules/pptx/shapes/autoshape.html - # Not all shape types are implemented in python-pptx - logger.warning("Encountered an unimplemented shape type.") - - return False - - def _is_table(self, shape): - try: - if shape.shape_type == pptx.enum.shapes.MSO_SHAPE_TYPE.TABLE: - return True - except NotImplementedError: - # # https://python-pptx.readthedocs.io/en/latest/_modules/pptx/shapes/autoshape.html - # Not all shape types are implemented in python-pptx - logger.warning("Encountered an unimplemented shape type.") - return False - - def _convert_chart_to_markdown(self, chart): - try: - md = "\n\n### Chart" - if chart.has_title: - md += f": {chart.chart_title.text_frame.text}" - md += "\n\n" - data = [] - category_names = [c.label for c in chart.plots[0].categories] - series_names = [s.name for s in chart.series] - data.append(["Category"] + series_names) - - for idx, category in enumerate(category_names): - row = [category] - for series in chart.series: - row.append(series.values[idx]) - data.append(row) - - markdown_table = [] - for row in data: - markdown_table.append("| " + " | ".join(map(str, row)) + " |") - header = markdown_table[0] - separator = "|" + "|".join(["---"] * len(data[0])) + "|" - return md + "\n".join([header, separator] + markdown_table[1:]) - except ValueError as e: - # Handle the specific error for unsupported chart types - if "unsupported plot type" in str(e): - return "\n\n[unsupported chart]\n\n" - except Exception: - # Catch any other exceptions that might occur - return "\n\n[unsupported chart]\n\n" - - class PPTXLoader(BaseLoader): + """Adapter shim — delegates to ``PptxParser``; layers image captioning on top.""" + def __init__(self, **kwargs) -> None: super().__init__(**kwargs) - self.image_placeholder = r"" - self.converter = PPTXConverter(image_placeholder=self.image_placeholder, page_separator=self.page_sep) + self._parser = PptxParser() async def aload_document(self, file_path, metadata=None, save_markdown=False): - md_content, imgs = self.converter.convert(local_path=file_path) - - if self.image_captioning: - images_captions = await self.caption_images(imgs, desc="Generating captions") - - for caption in images_captions: - md_content = re.sub( - self.image_placeholder, - caption.replace("\\", "/"), - md_content, - count=1, - ) - else: + metadata = {} if metadata is None else dict(metadata) + path = Path(file_path) + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.PPTX, + raw_bytes=raw_bytes, + metadata=dict(metadata) if metadata else {}, + ) + processed = await self._parser.parse(core_doc) + + # Reconstitute the legacy ``\n[PAGE_N]\n`` page-anchored layout. + slides = [f"{b.text}\n[PAGE_{b.page_number}]" for b in processed.text_blocks] + md_content = ("\n".join(slides) + "\n") if slides else "" + + if not self.image_captioning: logger.info("Image captioning disabled. Ignoring images.") - # Remove image placeholders when captioning is disabled - md_content = md_content.replace(self.image_placeholder, "") + for block in processed.images: + ref = (block.metadata or {}).get("markdown_ref") + if ref: + md_content = md_content.replace(ref, "") + elif processed.images: + pil_images: list[Image.Image] = [] + for block in processed.images: + img = Image.open(BytesIO(block.image_bytes)) + img.load() + pil_images.append(img) + captions = await self.caption_images(pil_images, desc="Generating captions") + for block, caption in zip(processed.images, captions): + ref = (block.metadata or {}).get("markdown_ref") + if ref: + md_content = md_content.replace(ref, caption.replace("\\", "/")) doc = Document(page_content=md_content, metadata=metadata) if save_markdown: diff --git a/openrag/components/indexer/loaders/test_doc_loader.py b/openrag/components/indexer/loaders/test_doc_loader.py index 68d2f2b30..a74e01b24 100644 --- a/openrag/components/indexer/loaders/test_doc_loader.py +++ b/openrag/components/indexer/loaders/test_doc_loader.py @@ -1,21 +1,25 @@ """ -Unit tests for DocLoader .doc to .docx conversion with fallback. - -Mocks spire.doc.Document entirely since it's a native .NET library -that cannot run without real .doc files. +Unit tests for the legacy ``DocLoader`` shim. + +The .doc → .docx → markdown conversion logic itself is tested in +``core/indexing/parsers/test_doc_parser.py``. These tests cover only +shim-level concerns: the langchain ``Document`` round-trip, the +``save_markdown=True`` integration with ``BaseLoader.save_content``, +and that errors raised by the underlying ``DocParser`` propagate +without being swallowed. """ -import os from unittest.mock import AsyncMock, MagicMock, patch import pytest from config.models import LoaderConfig, VLMConfig +from core.models.document import ProcessedDocument, TextBlock from langchain_core.documents.base import Document as LCDocument @pytest.fixture def mock_config(): - """Create a minimal mock config for BaseLoader.""" + """Minimal mock config for BaseLoader.""" config = MagicMock() config.vlm = VLMConfig(model="mock", base_url="http://mock", api_key="mock") config.loader = LoaderConfig(image_captioning=False, image_captioning_url=False) @@ -27,21 +31,17 @@ def metadata(): return {"file_id": "test-file-id", "partition": "test-partition"} -# All patches needed to import and instantiate DocLoader without real dependencies _PATCHES = [ - patch("components.indexer.loaders.doc.Document"), - patch("components.indexer.loaders.doc.DocxLoader"), + patch("components.indexer.loaders.doc.DocParser"), patch("components.indexer.loaders.base.ChatOpenAI"), patch("components.indexer.loaders.base.load_config"), ] def _start_patches(mock_config): - """Start all patches and return (MockSpireDoc, MockDocxLoader).""" mocks = [p.start() for p in _PATCHES] - mock_spire_doc, mock_docx_loader_cls, mock_chat, mock_load_config = mocks + _mock_doc_parser_cls, _mock_chat, mock_load_config = mocks mock_load_config.return_value = mock_config - return mock_spire_doc, mock_docx_loader_cls def _stop_patches(): @@ -58,150 +58,88 @@ def _patch_cleanup(): _stop_patches() -class TestDocLoader: - """Test DocLoader .doc to .docx conversion and fallback logic.""" - - def _make_loader(self, mock_config): - """Create a DocLoader with all dependencies mocked. Patches must be active.""" - from components.indexer.loaders.doc import DocLoader +def _make_loader(mock_config): + from components.indexer.loaders.doc import DocLoader - loader = DocLoader(config=mock_config) - return loader - - @pytest.mark.asyncio - async def test_successful_conversion(self, mock_config, metadata): - """Test happy path: .doc converts to .docx successfully.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) + return DocLoader(config=mock_config) - expected_doc = LCDocument(page_content="converted markdown", metadata=metadata) - loader.MDLoader.aload_document = AsyncMock(return_value=expected_doc) - mock_doc_instance = MagicMock() - mock_spire_doc.return_value = mock_doc_instance +def _processed(text: str = "markdown") -> ProcessedDocument: + return ProcessedDocument( + document_id="test", + text_blocks=[TextBlock(text=text, page_number=1)] if text else [], + metadata={}, + page_count=1 if text else 0, + ) - result = await loader.aload_document("/fake/path.doc", metadata) - mock_doc_instance.LoadFromFile.assert_called_once_with("/fake/path.doc") - mock_doc_instance.SaveToFile.assert_called_once() - mock_doc_instance.Close.assert_called_once() - - # DocxLoader was called with a temp .docx path - loader.MDLoader.aload_document.assert_called_once() - call_args = loader.MDLoader.aload_document.call_args - assert call_args[0][0].endswith(".docx") - - assert result == expected_doc - mock_doc_instance.GetText.assert_not_called() +class TestDocLoaderShim: + """Shim-level integration: ``DocParser`` ↔ langchain ``Document`` ↔ ``BaseLoader``.""" @pytest.mark.asyncio - async def test_fallback_on_spire_exception(self, mock_config, metadata): - """Test fallback to text extraction when SaveToFile crashes.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) - - mock_doc_instance = MagicMock() - mock_doc_instance.SaveToFile.side_effect = Exception("TypeInitialization_Type_NoTypeAvailable") - mock_doc_instance.GetText.return_value = "Plain text content from .doc" - mock_spire_doc.return_value = mock_doc_instance + async def test_happy_path_returns_langchain_document(self, mock_config, metadata, tmp_path): + """Parser output is joined into ``page_content``; ``metadata`` is passed through.""" + _start_patches(mock_config) + loader = _make_loader(mock_config) + loader._parser = MagicMock() + loader._parser.parse = AsyncMock(return_value=_processed("converted markdown")) - result = await loader.aload_document("/fake/path.doc", metadata) + file_path = tmp_path / "x.doc" + file_path.write_bytes(b"\xd0\xcf\x11\xe0fake-doc") - mock_doc_instance.SaveToFile.assert_called_once() - mock_doc_instance.GetText.assert_called_once() - mock_doc_instance.Close.assert_called_once() + result = await loader.aload_document(str(file_path), metadata) - # DocxLoader should NOT have been called - loader.MDLoader.aload_document.assert_not_called() - assert result.page_content == "Plain text content from .doc" + assert isinstance(result, LCDocument) + assert result.page_content == "converted markdown" assert result.metadata == metadata - @pytest.mark.asyncio - async def test_temp_file_cleaned_up_on_success(self, mock_config, metadata): - """Test temp file is removed after successful conversion.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) - - created_temp_files = [] - - mock_doc_instance = MagicMock() - - def capture_temp_path(path, fmt): - created_temp_files.append(path) - - mock_doc_instance.SaveToFile.side_effect = capture_temp_path - mock_spire_doc.return_value = mock_doc_instance - - expected_doc = LCDocument(page_content="content", metadata=metadata) - loader.MDLoader.aload_document = AsyncMock(return_value=expected_doc) - - await loader.aload_document("/fake/path.doc", metadata) - - # The temp file should have been cleaned up by the finally block - for path in created_temp_files: - assert not os.path.exists(path) + loader._parser.parse.assert_awaited_once() + forwarded = loader._parser.parse.await_args.args[0] + assert forwarded.raw_bytes == b"\xd0\xcf\x11\xe0fake-doc" + assert forwarded.filename == "x.doc" @pytest.mark.asyncio - async def test_temp_file_cleaned_up_on_failure(self, mock_config, metadata): - """Test temp file is removed even when conversion fails.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) - - mock_doc_instance = MagicMock() - created_temp_files = [] - - def save_then_fail(path, fmt): - created_temp_files.append(path) - # Create the file so we can verify it's cleaned up - with open(path, "w") as f: - f.write("partial") - raise Exception("Spire crash") - - mock_doc_instance.SaveToFile.side_effect = save_then_fail - mock_doc_instance.GetText.return_value = "fallback text" - mock_spire_doc.return_value = mock_doc_instance - - result = await loader.aload_document("/fake/path.doc", metadata) - - assert result.page_content == "fallback text" - for path in created_temp_files: - assert not os.path.exists(path), f"Temp file was not cleaned up: {path}" + async def test_empty_parser_result_yields_empty_content(self, mock_config, metadata, tmp_path): + """An empty ``ProcessedDocument`` yields an empty langchain document.""" + _start_patches(mock_config) + loader = _make_loader(mock_config) + loader._parser = MagicMock() + loader._parser.parse = AsyncMock(return_value=_processed(text="")) + + file_path = tmp_path / "x.doc" + file_path.write_bytes(b"") + + result = await loader.aload_document(str(file_path), metadata) + assert result.page_content == "" + assert result.metadata == metadata @pytest.mark.asyncio - async def test_fallback_with_save_markdown(self, mock_config, metadata, tmp_path): - """Test fallback path respects save_markdown flag.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) - - mock_doc_instance = MagicMock() - mock_doc_instance.SaveToFile.side_effect = Exception("Spire crash") - mock_doc_instance.GetText.return_value = "Extracted text" - mock_spire_doc.return_value = mock_doc_instance + async def test_save_markdown_writes_extracted_content(self, mock_config, metadata, tmp_path): + """``save_markdown=True`` calls ``BaseLoader.save_content`` with the extracted text and source path.""" + _start_patches(mock_config) + loader = _make_loader(mock_config) + loader._parser = MagicMock() + loader._parser.parse = AsyncMock(return_value=_processed("Extracted text")) - file_path = str(tmp_path / "test.doc") + file_path = tmp_path / "x.doc" + file_path.write_bytes(b"x") with patch.object(loader, "save_content") as mock_save: - result = await loader.aload_document(file_path, metadata, save_markdown=True) - mock_save.assert_called_once_with("Extracted text", file_path) + result = await loader.aload_document(str(file_path), metadata, save_markdown=True) + mock_save.assert_called_once_with("Extracted text", str(file_path)) assert result.page_content == "Extracted text" @pytest.mark.asyncio - async def test_docx_loader_error_propagates(self, mock_config, metadata): - """Test that MDLoader errors are NOT caught by the Spire fallback.""" - mock_spire_doc, _ = _start_patches(mock_config) - loader = self._make_loader(mock_config) - - mock_doc_instance = MagicMock() - mock_spire_doc.return_value = mock_doc_instance - - # Spire conversion succeeds, but DocxLoader fails - loader.MDLoader.aload_document = AsyncMock(side_effect=ValueError("DocxLoader broke")) - - with pytest.raises(ValueError, match="DocxLoader broke"): - await loader.aload_document("/fake/path.doc", metadata) - - # GetText fallback should NOT have been used - mock_doc_instance.GetText.assert_not_called() - # But Close should still be called (via finally) - mock_doc_instance.Close.assert_called_once() + async def test_parser_error_propagates(self, mock_config, metadata, tmp_path): + """Exceptions from the underlying ``DocParser`` are not swallowed.""" + _start_patches(mock_config) + loader = _make_loader(mock_config) + loader._parser = MagicMock() + loader._parser.parse = AsyncMock(side_effect=ValueError("DocParser broke")) + + file_path = tmp_path / "x.doc" + file_path.write_bytes(b"x") + + with pytest.raises(ValueError, match="DocParser broke"): + await loader.aload_document(str(file_path), metadata) diff --git a/openrag/components/indexer/loaders/txt_loader.py b/openrag/components/indexer/loaders/txt_loader.py index 33e272663..775b798c0 100644 --- a/openrag/components/indexer/loaders/txt_loader.py +++ b/openrag/components/indexer/loaders/txt_loader.py @@ -1,11 +1,24 @@ """ Text and Markdown file loader implementation. + +``TextLoader`` and ``MarkdownLoader`` are now thin :class:`BaseLoader` +adapters that delegate extraction to the corresponding core parsers +(:class:`core.indexing.parsers.text_parser.TextParser`, +:class:`core.indexing.parsers.markdown_parser.MarkdownParser`). Image +captioning is layered on top of the markdown adapter via the +``BaseLoader`` mixin. New code should call the core parsers directly; +these shims keep the legacy loader-discovery path alive until consumers +migrate. """ +import asyncio from pathlib import Path from components.indexer.loaders.base import BaseLoader -from langchain_community.document_loaders import TextLoader as LangchainTextLoader +from core.indexing.parsers.markdown_parser import MarkdownParser +from core.indexing.parsers.text_parser import TextParser +from core.models.document import Document as CoreDocument +from core.models.document import DocumentType from langchain_core.documents.base import Document from utils.logger import get_logger @@ -13,12 +26,11 @@ class TextLoader(BaseLoader): - """ - Loader for plain text files (.txt). - """ + """Adapter shim — delegates to ``TextParser`` and returns a LangChain ``Document``.""" def __init__(self, **kwargs) -> None: super().__init__(**kwargs) + self._parser = TextParser() async def aload_document( self, @@ -30,13 +42,15 @@ async def aload_document( metadata = {} path = Path(file_path) - loader = LangchainTextLoader(file_path=str(path), autodetect_encoding=True) - - # Load document segments asynchronously - doc_segments = await loader.aload() - - # Create final document - content = doc_segments[0].page_content.strip() + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.TEXT, + raw_bytes=raw_bytes, + metadata=metadata, + ) + processed = await self._parser.parse(core_doc) + content = "\n\n".join(block.text for block in processed.text_blocks).strip() doc = Document(page_content=content, metadata=metadata) if save_markdown: @@ -46,12 +60,11 @@ async def aload_document( class MarkdownLoader(BaseLoader): - """ - Loader for markdown files (.md). - """ + """Adapter shim — delegates to ``MarkdownParser`` and layers image captioning on top.""" def __init__(self, **kwargs) -> None: super().__init__(**kwargs) + self._parser = MarkdownParser() async def aload_document( self, @@ -63,15 +76,16 @@ async def aload_document( metadata = {} path = Path(file_path) - loader = LangchainTextLoader(file_path=str(path), autodetect_encoding=True) - - # Load document segments asynchronously - doc_segments = await loader.aload() - - # Create final document - content = doc_segments[0].page_content.strip() + raw_bytes = await asyncio.to_thread(path.read_bytes) + core_doc = CoreDocument( + filename=path.name, + content_type=DocumentType.MARKDOWN, + raw_bytes=raw_bytes, + metadata=metadata, + ) + processed = await self._parser.parse(core_doc) + content = "\n\n".join(block.text for block in processed.text_blocks).strip() - # Caption any images in the markdown content = await self.replace_markdown_images_with_captions(content) doc = Document(page_content=content, metadata=metadata) diff --git a/openrag/components/indexer/utils/text_sanitizer.py b/openrag/components/indexer/utils/text_sanitizer.py index 304a094cb..ef14f6f8d 100644 --- a/openrag/components/indexer/utils/text_sanitizer.py +++ b/openrag/components/indexer/utils/text_sanitizer.py @@ -1,149 +1,7 @@ -""" -Text sanitization utilities for cleaning extracted text and improving quality. - -This module provides functions to clean and normalize text extracted from various -document sources (PDFs, Office files, etc.) by removing excessive whitespace, -special characters, and other artifacts that don't add value. -""" - -import re -import unicodedata - - -def sanitize_text( - text: str, - normalize_whitespace: bool = True, - remove_control_chars: bool = True, - remove_zero_width_chars: bool = True, - max_consecutive_newlines: int = 2, - normalize_unicode: bool = True, -) -> str: - """ - Sanitize text by removing useless characters and normalizing whitespace. - - This function performs comprehensive text cleaning including: - - Removing or normalizing control characters - - Removing zero-width spaces and invisible characters - - Normalizing excessive whitespace (spaces, tabs) - - Limiting consecutive newlines - - Unicode normalization - - Args: - text: The input text to sanitize - normalize_whitespace: If True, normalize spaces and tabs to single spaces - remove_control_chars: If True, remove control characters (except \n, \r, \t) - remove_zero_width_chars: If True, remove zero-width spaces and similar chars - max_consecutive_newlines: Maximum number of consecutive newlines to keep (0 = unlimited) - normalize_unicode: If True, normalize unicode to NFC form - - Returns: - Sanitized text string - - Examples: - >>> sanitize_text("Hello world\\n\\n\\n\\nTest") - 'Hello world\\n\\nTest' - >>> sanitize_text("Text with\\t\\ttabs") - 'Text with tabs' - """ - if not text: - return text - - # Normalize unicode to NFC form (composed form) - if normalize_unicode: - text = unicodedata.normalize("NFC", text) - - # Remove zero-width spaces and similar invisible characters - if remove_zero_width_chars: - # Zero-width space (U+200B), zero-width non-joiner (U+200C), - # zero-width joiner (U+200D), word joiner (U+2060), - # zero-width no-break space (U+FEFF) - text = re.sub(r"[\u200B-\u200D\u2060\uFEFF]", "", text) - - # Remove control characters except newline, carriage return, and tab - if remove_control_chars: - # Remove C0 control characters (0x00-0x1F) except \t (0x09), \n (0x0A), \r (0x0D) - # and C1 control characters (0x80-0x9F) - text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F-\x9F]", "", text) - - # Normalize whitespace - if normalize_whitespace: - # Convert multiple spaces to single space - text = re.sub(r" {2,}", " ", text) - - # Convert tabs to single space - text = re.sub(r"\t+", " ", text) - - # Remove spaces at the beginning of lines - text = re.sub(r"(?m)^ +", "", text) - - # Remove spaces at the end of lines - text = re.sub(r"(?m) +$", "", text) - - # Normalize line breaks - # First, normalize different line break styles to \n - text = re.sub(r"\r\n", "\n", text) - text = re.sub(r"\r", "\n", text) - - # Limit consecutive newlines - if max_consecutive_newlines > 0: - pattern = r"\n{" + str(max_consecutive_newlines + 1) + r",}" - replacement = "\n" * max_consecutive_newlines - text = re.sub(pattern, replacement, text) - - # Remove leading/trailing whitespace - text = text.strip() - return text - - -def clean_markdown_table_spacing(markdown_table: str) -> str: - """ - Normalize spacing inside a markdown table: - - trims each cell - - keeps table shape intact - - Args: - markdown_table: Markdown table text to clean - - Returns: - Cleaned markdown table with normalized spacing - """ - cleaned_lines = [] - - for line in markdown_table.strip().split("\n"): - if "|" not in line: - cleaned_lines.append(line.strip()) - continue - - # Split row into cells (preserve leading/trailing pipes) - parts = line.split("|") - - # Strip each cell except the outer empty ones - cleaned_cells = [cell.strip() for cell in parts] - - # Rebuild with a single space around each cell - new_line = "| " + " | ".join(cleaned_cells[1:-1]) + " |" - cleaned_lines.append(new_line) - - return "\n".join(cleaned_lines) - - -def sanitize_extracted_text(text: str) -> str: - """ - Convenience function for sanitizing text extracted from documents. - - This applies a standard set of cleaning operations suitable for - text extraction endpoints and general document processing. - Uses the default sanitization settings which include: - - Normalize whitespace - - Remove control characters - - Remove zero-width characters - - Limit consecutive newlines to 2 - - Normalize Unicode - - Args: - text: The extracted text to sanitize - - Returns: - Sanitized text - """ - return sanitize_text(text) +# Re-export from canonical location for backward compatibility. +# New code should import from `core.utils.text` directly. +from core.utils.text import ( # noqa: F401 + clean_markdown_table_spacing, + sanitize_extracted_text, + sanitize_text, +) diff --git a/openrag/components/ray_utils.py b/openrag/components/ray_utils.py index baecf358c..83ec534b1 100644 --- a/openrag/components/ray_utils.py +++ b/openrag/components/ray_utils.py @@ -1,90 +1,6 @@ -import asyncio -from collections.abc import Callable -from typing import Any - -import ray -from ray.exceptions import RayTaskError, TaskCancelledError -from utils.logger import get_logger - -logger = get_logger() - - -async def call_ray_actor_with_timeout( - future: ray.ObjectRef, - timeout: float, - task_description: str = "Ray task", -) -> Any: - """ - Wait for a Ray actor call with timeout and proper cancellation handling. - - This utility provides consistent error handling for Ray actor calls: - - Timeout with proper task cancellation - - Propagation of asyncio cancellation to Ray tasks - - Proper handling of Ray-specific exceptions - - Args: - future: The Ray ObjectRef from a remote call - timeout: Timeout in seconds - task_description: Description for error messages - - Returns: - The result of the Ray task - - Raises: - TimeoutError: If the task exceeds the timeout - asyncio.CancelledError: If the calling coroutine is cancelled - TaskCancelledError: If the Ray task was cancelled - RuntimeError: If the Ray task failed with an error - """ - try: - result = await asyncio.wait_for(asyncio.gather(future), timeout=timeout) - return result[0] # gather returns a list - - except TimeoutError: - logger.warning(f"{task_description} timed out, cancelling Ray task") - ray.cancel(future, recursive=True) - raise - - except asyncio.CancelledError: - logger.warning(f"{task_description} cancelled, cancelling Ray task") - ray.cancel(future, recursive=True) - raise - - except TaskCancelledError: - logger.warning(f"{task_description} Ray task was cancelled") - raise - - except RayTaskError as e: - raise RuntimeError(f"{task_description} failed") from e - - -async def retry_with_backoff( - attempt_fn: Callable[[int], Any], - max_retries: int, - base_delay: float, - task_description: str = "task", -) -> Any: - """ - Run `attempt_fn(attempt_index)` (an async callable) with exponential-backoff - retries. The callable owns its own resource acquire/release per attempt. - - Backoff: base_delay * 2**attempt seconds. CancelledError is never retried. - """ - last_exc: Exception | None = None - for attempt in range(max_retries + 1): - try: - return await attempt_fn(attempt) - except asyncio.CancelledError: - raise - except Exception as e: - last_exc = e - if attempt >= max_retries: - logger.error(f"{task_description} failed after {attempt + 1} attempts: {e}") - raise - delay = base_delay * (2**attempt) - logger.warning( - f"{task_description} failed (attempt {attempt + 1}/{max_retries + 1}): {e}. Retrying in {delay:.1f}s..." - ) - await asyncio.sleep(delay) - - raise last_exc # unreachable +# Re-export from canonical location for backward compatibility. +# New code should import from `services.workers.ray_utils` directly. +from services.workers.ray_utils import ( # noqa: F401 + call_ray_actor_with_timeout, + retry_with_backoff, +) diff --git a/openrag/consts.py b/openrag/consts.py index 51bc12904..240bba006 100644 --- a/openrag/consts.py +++ b/openrag/consts.py @@ -1,6 +1,6 @@ -PARTITION_PREFIX = "openrag-" -LEGACY_PARTITION_PREFIX = "ragondin-" - -FILE_READ_CHUNK_SIZE = 1024 * 1024 # Read file in blocks of 1MB to preserve RAM - -IMAGE_PLACEHOLDER = """\n\n[Image Placeholder]\n\n""" +from core.utils.conts import ( # noqa: F401,F403 + FILE_READ_CHUNK_SIZE, + IMAGE_PLACEHOLDER, + LEGACY_PARTITION_PREFIX, + PARTITION_PREFIX, +) diff --git a/openrag/core/config/indexation.py b/openrag/core/config/indexation.py index fbebc34cc..39e006206 100644 --- a/openrag/core/config/indexation.py +++ b/openrag/core/config/indexation.py @@ -2,7 +2,9 @@ from __future__ import annotations -from pydantic import Field +from typing import Any + +from pydantic import Field, field_validator from .base import ConfigMixin @@ -10,6 +12,17 @@ # Transcriber (nested under loader) # --------------------------------------------------------------------------- +_DEFAULT_DIRECT_UPLOAD_SUFFIXES = frozenset( + {".wav", ".flac", ".ogg", ".mp3", ".mp4", ".m4a", ".webm", ".mpeg", ".mpga"} +) + + +def _normalize_suffix(s: str) -> str: + s = s.strip().lower() + if not s: + return "" + return s if s.startswith(".") else f".{s}" + class TranscriberConfig(ConfigMixin): base_url: str = "http://transcriber:8000/v1" @@ -18,6 +31,14 @@ class TranscriberConfig(ConfigMixin): timeout: int = 3600 max_concurrent_chunks: int = 20 use_whisper_lang_detector: bool = True + direct_upload_suffixes: set[str] = Field(default_factory=lambda: set(_DEFAULT_DIRECT_UPLOAD_SUFFIXES)) + + @field_validator("direct_upload_suffixes", mode="before") + @classmethod + def _split_suffixes(cls, v: Any) -> Any: + if isinstance(v, str): + return {n for raw in v.split("|") if (n := _normalize_suffix(raw))} + return v # --------------------------------------------------------------------------- diff --git a/openrag/core/indexing/contextualize.py b/openrag/core/indexing/contextualize.py new file mode 100644 index 000000000..25f8843d1 --- /dev/null +++ b/openrag/core/indexing/contextualize.py @@ -0,0 +1,147 @@ +"""Chunk contextualization against the ``LLM`` ABC. + +Framework-free implementation of contextual retrieval: for each chunk, +ask an LLM to write a short situating context based on the document's +opening chunks plus the immediate preceding neighbourhood, then prepend +that context to the chunk text so embeddings capture document-level +meaning. + +Inputs and outputs are :class:`core.models.chunk.Chunk` instances. The +caller supplies the LLM, the system prompt, and any concurrency / timeout +limits — core does not reach into Hydra config or the global VLM +semaphore. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Sequence + +from ..llm import LLM +from ..models.chunk import Chunk + +logger = logging.getLogger(__name__) + + +BASE_CHUNK_FORMAT = "* filename: {filename}\n\n[CHUNK_START]\n\n{content}\n\n[CHUNK_END]" +CHUNK_FORMAT = "[CONTEXT]\n\n{chunk_context}\n\n" + BASE_CHUNK_FORMAT + +DEFAULT_TIMEOUT_SECONDS = 30.0 +DEFAULT_MAX_CONCURRENT = 4 + + +def format_chunk(content: str, filename: str, chunk_context: str | None = None) -> str: + """Render the canonical chunk wrapping (with or without context).""" + if chunk_context: + return CHUNK_FORMAT.format(content=content, filename=filename, chunk_context=chunk_context) + return BASE_CHUNK_FORMAT.format(content=content, filename=filename) + + +class ChunkContextualizer: + """Generate a per-chunk context string and prepend it to the chunk text.""" + + def __init__( + self, + llm: LLM, + system_prompt: str, + *, + timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS, + max_concurrent: int = DEFAULT_MAX_CONCURRENT, + semaphore: asyncio.Semaphore | None = None, + ): + self._llm = llm + self._system_prompt = system_prompt + self._timeout = timeout_seconds + self._batch_size = max(1, max_concurrent) + self._semaphore = semaphore or asyncio.Semaphore(max_concurrent) + + async def _generate_context( + self, + first_chunks: Sequence[Chunk], + prev_chunks: Sequence[Chunk], + current_chunk: Chunk, + filename: str, + lang: str, + ) -> str: + first_block = "\n--\n".join(c.text for c in first_chunks) + prev_block = "\n--\n".join(c.text for c in prev_chunks) + user_msg = ( + "Here is the context to consider for generating the context:\n" + f"- Filename: {filename}\n" + f"- First chunks:\n{first_block}\n\n" + f"- Previous chunks:\n{prev_block}\n\n" + f"Here is the current chunk to contextualize strictly in this {lang} language:\n" + f"- Current chunk:\n\n{current_chunk.text}" + ) + messages = [ + {"role": "system", "content": self._system_prompt}, + {"role": "user", "content": user_msg}, + ] + async with self._semaphore: + try: + return await asyncio.wait_for(self._llm.chat(messages), timeout=self._timeout) + except TimeoutError: + logger.warning("LLM timeout contextualizing chunk (filename=%s)", filename) + return "" + except Exception as exc: + logger.warning("Error contextualizing chunk (filename=%s): %s", filename, exc) + return "" + + async def contextualize( + self, + chunks: Sequence[Chunk], + *, + filename: str = "", + lang: str = "en", + ) -> list[Chunk]: + """Return new chunks with context prepended to ``text``. + + Each returned chunk preserves the input's id, metadata, and other + fields; ``text`` is rewritten to the formatted (context + content) + string used for embedding, ``context`` holds the generated context, + and ``content`` holds the original chunk text. + + Falls back to returning the input chunks unchanged on any + unrecoverable error. + """ + chunks = list(chunks) + if not chunks: + return [] + + try: + first_chunks = chunks[:2] + contexts: list[str] = [] + # Schedule one batch at a time so prompt strings + coroutine + # objects don't all sit in memory upfront on large documents. + for start in range(0, len(chunks), self._batch_size): + end = min(start + self._batch_size, len(chunks)) + batch = [ + self._generate_context( + first_chunks=first_chunks, + prev_chunks=chunks[max(0, i - 2) : i] if i > 0 else [], + current_chunk=chunks[i], + filename=filename, + lang=lang, + ) + for i in range(start, end) + ] + contexts.extend(await asyncio.gather(*batch)) + + return [ + chunk.model_copy( + update={ + "text": format_chunk( + content=chunk.text, + filename=filename, + chunk_context=context, + ), + "context": context, + "content": chunk.text, + } + ) + for chunk, context in zip(chunks, contexts, strict=True) + ] + except Exception as exc: + logger.warning("Error contextualizing chunks from %s: %s", filename, exc) + return chunks diff --git a/openrag/core/indexing/image_preprocessor.py b/openrag/core/indexing/image_preprocessor.py new file mode 100644 index 000000000..50cead047 --- /dev/null +++ b/openrag/core/indexing/image_preprocessor.py @@ -0,0 +1,122 @@ +"""Image preprocessing helpers for the indexing pipeline. + +Pure helpers — no VLM, no langchain, no infrastructure imports. Used by +parsers (core) and Ray-pool adapters (services) that need to: + +- normalize PIL Image modes for PNG encoding +- encode PIL Images as PNG bytes or base64 data URIs +- detect / decode markdown image references (HTTP / data URI) in extracted text + +Extracted from the legacy ``components/indexer/loaders/base.py``; the +legacy module is kept as a back-compat shim until existing imports are +migrated. +""" + +from __future__ import annotations + +import base64 +import logging +import re +from io import BytesIO +from typing import Any + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Markdown image-reference patterns (compile once; shared regex objects) +# --------------------------------------------------------------------------- + +HTTP_IMAGE_PATTERN = re.compile(r"!\[(.*?)\]\((https?://[^)]+)\)") +DATA_URI_IMAGE_PATTERN = re.compile(r"!\[(.*?)\]\((data:image/[^;]+;base64,[^)]+)\)") + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +# Qwen2.5-VL ``min_pixels`` threshold; images below this break the model. +MIN_IMAGE_PIXELS = 784 + + +# --------------------------------------------------------------------------- +# PIL mode normalization & encoding +# --------------------------------------------------------------------------- + + +def ensure_png_compatible_mode(image: Any) -> Any: + """Convert PIL image modes that PNG can't encode directly. + + CMYK/YCbCr/LAB → RGB; P/LA/PA → RGBA. Others returned unchanged. + """ + if image.mode in ("CMYK", "YCbCr", "LAB"): + return image.convert("RGB") + if image.mode in ("P", "LA", "PA"): + return image.convert("RGBA") + return image + + +def pil_to_png_bytes(image: Any) -> bytes: + """Encode a PIL Image as PNG bytes. ``bytes`` input is passed through.""" + if isinstance(image, bytes): + return image + image = ensure_png_compatible_mode(image) + buf = BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + + +# --------------------------------------------------------------------------- +# URL / data URI detection +# --------------------------------------------------------------------------- + + +def decode_data_uri(data_uri: str) -> bytes | None: + """Decode a ``data:image/...;base64,...`` URI into raw bytes. ``None`` on failure.""" + try: + _, b64 = data_uri.split(",", 1) + return base64.b64decode(b64) + except Exception as exc: + logger.warning("Failed to decode data URI: %s", exc) + return None + + +def mime_from_data_uri(data_uri: str) -> str: + """Pull the mime type out of a data URI; fall back to ``image/png``. + + Example: ``data:image/jpeg;base64,xxx`` → ``image/jpeg``. + """ + try: + return data_uri.split(",", 1)[0].split(":", 1)[1].split(";", 1)[0] + except Exception: + return "image/png" + + +def extract_data_uri_image_blocks(text: str, *, page_number: int = 1) -> list[Any]: + """Build ``ImageBlock``s for every ``![alt](data:image/...;base64,...)`` ref. + + The original markdown ref is preserved in ``metadata['markdown_ref']`` + so a downstream caption stage can substitute the wrapped caption back + into the corresponding ``TextBlock`` via ``str.replace``. + + Returns ``list[ImageBlock]`` (declared as ``list[Any]`` only because + importing the model would create a cycle in some build orderings — + the caller side is type-correct). + """ + if not text: + return [] + # Local import to avoid a top-level cycle with ``core.models``. + from ..models.document import ImageBlock + + blocks: list[Any] = [] + for alt, data_uri in DATA_URI_IMAGE_PATTERN.findall(text): + payload = decode_data_uri(data_uri) + if payload is None: + continue + blocks.append( + ImageBlock( + image_bytes=payload, + page_number=page_number, + mime_type=mime_from_data_uri(data_uri), + metadata={"markdown_ref": f"![{alt}]({data_uri})", "alt": alt}, + ) + ) + return blocks diff --git a/openrag/core/indexing/parsers/audio/__init__.py b/openrag/core/indexing/parsers/audio/__init__.py new file mode 100644 index 000000000..7604a3d84 --- /dev/null +++ b/openrag/core/indexing/parsers/audio/__init__.py @@ -0,0 +1,12 @@ +"""Audio parser facades. + +Each backend lives in its own module so its heavy dependencies (Ray +worker pools, cloud SDKs, …) are only pulled in by the concrete impl +in ``services/`` — the core facade just declares the parser type and +delegates ``parse()`` to an injected pool/client. +""" + +from .client_based import ClientAudioParser +from .local_whisper import LocalWhisperParser + +__all__ = ["ClientAudioParser", "LocalWhisperParser"] diff --git a/openrag/core/indexing/parsers/audio/client_based.py b/openrag/core/indexing/parsers/audio/client_based.py new file mode 100644 index 000000000..4718ba5c6 --- /dev/null +++ b/openrag/core/indexing/parsers/audio/client_based.py @@ -0,0 +1,32 @@ +"""Client-backed audio ``DocumentParser`` (thin core facade). + +Holds a ``BaseClientParser`` (the actual HTTP-client / OpenAI-SDK +implementation lives in ``services/`` and is composed in at startup) +and delegates ``parse()`` to it. + +Mirrors the :class:`ClientPdfParser` pattern: core stays free of vendor +SDKs while the facade names "client-backed audio" as a first-class +parser type. +""" + +from __future__ import annotations + +from ....models.document import Document, ProcessedDocument +from ..document_parser import BaseClientParser, DocumentParser +from ..registry import parser_registry + + +@parser_registry.register("audio_client") +class ClientAudioParser(DocumentParser): + """Public audio parser facade backed by an OpenAI-compatible transcription client.""" + + def __init__(self, client: BaseClientParser) -> None: + if not isinstance(client, BaseClientParser): + raise ValueError("ClientAudioParser requires a BaseClientParser instance as client") + self._client = client + + def supported_types(self) -> list[str]: + return self._client.supported_types() + + async def parse(self, document: Document) -> ProcessedDocument: + return await self._client.parse(document) diff --git a/openrag/core/indexing/parsers/audio/local_whisper.py b/openrag/core/indexing/parsers/audio/local_whisper.py new file mode 100644 index 000000000..c880488ee --- /dev/null +++ b/openrag/core/indexing/parsers/audio/local_whisper.py @@ -0,0 +1,34 @@ +"""Local Whisper-backed audio ``DocumentParser`` (thin core facade). + +Holds a reference to a ``BasePooledParser`` (the actual Ray-pool / +GPU-model implementation lives in ``services/`` and is not yet wired +up) and delegates ``parse()`` to it. The split keeps core free of Ray +and GPU lifecycle code while still naming the local-Whisper backend as +a first-class parser type. + +The injected pool is a generic ``BasePooledParser``; if a more specific +``WhisperPoolParser`` ABC emerges in services, this class can tighten +its type without changing call sites. +""" + +from __future__ import annotations + +from ....models.document import Document, DocumentType, ProcessedDocument +from ..document_parser import BasePooledParser, DocumentParser +from ..registry import parser_registry + + +@parser_registry.register("local_whisper") +class LocalWhisperParser(DocumentParser): + """Public audio parser facade backed by a local-Whisper worker pool.""" + + def __init__(self, pool: BasePooledParser) -> None: + if not isinstance(pool, BasePooledParser): + raise ValueError("LocalWhisperParser requires a BasePooledParser instance as pool") + self._pool = pool + + def supported_types(self) -> list[str]: + return [DocumentType.AUDIO.value, DocumentType.VIDEO.value] + + async def parse(self, document: Document) -> ProcessedDocument: + return await self._pool.parse(document) diff --git a/openrag/core/indexing/parsers/doc_parser.py b/openrag/core/indexing/parsers/doc_parser.py new file mode 100644 index 000000000..c7371a593 --- /dev/null +++ b/openrag/core/indexing/parsers/doc_parser.py @@ -0,0 +1,104 @@ +"""Legacy ``.doc`` (binary Word 97-2003) ``DocumentParser``. + +Converts ``.doc`` to ``.docx`` via the ``spire.doc`` library, then +delegates to :class:`DocxParser` for Markdown extraction. Falls back to +plain-text extraction (``Document.GetText()``) if Spire's conversion +fails. + +Spire.Doc requires DOTNET; the constructor sets the invariant-globalization +env var that makes Spire usable without the full ICU data. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import tempfile +from pathlib import Path + +from ...models.document import Document, DocumentType, ProcessedDocument, TextBlock +from .document_parser import DocumentParser +from .docx_parser import DocxParser +from .registry import parser_registry + +logger = logging.getLogger(__name__) + +os.environ.setdefault("DOTNET_SYSTEM_GLOBALIZATION_INVARIANT", "1") + + +@parser_registry.register("doc") +class DocParser(DocumentParser): + """Parse legacy ``.doc`` files via .docx conversion + DocxParser.""" + + def __init__(self, docx_parser: DocxParser | None = None) -> None: + """Pass an explicit ``DocxParser`` to share VLM / semaphore config; + otherwise a captioning-disabled instance is constructed. + """ + self._docx = docx_parser or DocxParser() + + def supported_types(self) -> list[str]: + return [DocumentType.DOC.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as src_path: + docx_bytes, fallback_text = await asyncio.to_thread(self._convert, str(src_path)) + + if docx_bytes: + docx_doc = document.model_copy(update={"raw_bytes": docx_bytes, "content_type": DocumentType.DOCX}) + return await self._docx.parse(docx_doc) + + text = (fallback_text or "").strip() + text_blocks = [TextBlock(text=text, page_number=1)] if text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + metadata=dict(document.metadata), + page_count=1 if text else 0, + ) + + @staticmethod + def _convert(path: str) -> tuple[bytes | None, str | None]: + """Run blocking Spire.Doc conversion. Returns ``(docx_bytes, fallback_text)``. + + Exactly one of the two will be non-None on success; both ``None`` + means total failure (caller emits an empty ProcessedDocument). + """ + try: + from spire.doc import Document as SpireDocument + from spire.doc import FileFormat + except ImportError: + logger.warning("spire.doc not available; cannot parse legacy .doc files") + return None, None + + spire_doc = SpireDocument() + out_path: str | None = None + try: + spire_doc.LoadFromFile(path) + with tempfile.NamedTemporaryFile(suffix=".docx", delete=False) as out: + out_path = out.name + spire_doc.SaveToFile(out_path, FileFormat.Docx2016) + return Path(out_path).read_bytes(), None + except Exception as exc: + logger.warning("Spire.Doc .doc → .docx conversion failed (%s); falling back to plain text", exc) + try: + return None, spire_doc.GetText() + except Exception as fallback_exc: + logger.warning("Spire.Doc fallback text extraction also failed: %s", fallback_exc) + return None, None + finally: + try: + spire_doc.Close() + except Exception: + pass + if out_path and os.path.exists(out_path): + try: + os.remove(out_path) + except OSError: + pass diff --git a/openrag/core/indexing/parsers/document_parser.py b/openrag/core/indexing/parsers/document_parser.py index aedd0b383..d69d12c8b 100644 --- a/openrag/core/indexing/parsers/document_parser.py +++ b/openrag/core/indexing/parsers/document_parser.py @@ -1,10 +1,32 @@ -"""Abstract document parser interface.""" +"""Abstract document parser interface and category markers. + +``DocumentParser`` is the single port every concrete parser implements. + +Two empty subclasses are exposed alongside it as **type markers** — +they categorize a parser by *how* it gets its work done, without adding +any behaviour: + +- ``BasePooledParser`` — a parser whose ``parse()`` is satisfied by a + pool of workers (Ray actors, ProcessPoolExecutor, asyncio task group, + …). Concrete impls live in ``services/``. +- ``BaseClientParser`` — a parser whose ``parse()`` is satisfied by an + external client (HTTP service, gRPC, vendor SDK, …). Concrete impls + live in ``services/``. + +The markers exist so consumers (e.g. ``MarkerParser(pool: BasePooledParser)``) +can constrain the *kind* of parser they accept tighter than the +generic ``DocumentParser``. Concrete subclasses implement +``parse()`` and ``supported_types()`` directly — no extra hook method. + +If a shared pattern (retry, timeout, semaphore, …) ever materialises +across multiple subclasses, lift it into the base then. +""" from __future__ import annotations from abc import ABC, abstractmethod -from openrag.core.models.document import Document, ProcessedDocument +from ...models.document import Document, ProcessedDocument class DocumentParser(ABC): @@ -19,3 +41,11 @@ async def parse(self, document: Document) -> ProcessedDocument: def supported_types(self) -> list[str]: """Return list of DocumentType values this parser handles.""" ... + + +class BasePooledParser(DocumentParser, ABC): + """Marker for parsers backed by a worker pool. Concrete impl in services/.""" + + +class BaseClientParser(DocumentParser, ABC): + """Marker for parsers backed by an external client. Concrete impl in services/.""" diff --git a/openrag/core/indexing/parsers/docx_parser.py b/openrag/core/indexing/parsers/docx_parser.py new file mode 100644 index 000000000..a1702eb7c --- /dev/null +++ b/openrag/core/indexing/parsers/docx_parser.py @@ -0,0 +1,200 @@ +"""DOCX ``DocumentParser`` implementation. + +Conversion to Markdown via the ``markitdown`` library; falls back to +plain-text extraction via ``python-docx`` if MarkItDown fails. + +MarkItDown emits a generic ``![](data:image/png;base64...)`` placeholder +(literal, truncated) for every embedded image, with no per-image +identifier. Actual image bytes are pulled from the DOCX zip +(``word/media/``) and matched to placeholders **positionally**, in +document order. Each placeholder is rewritten to a unique synthetic +``![](docx-image-N)`` ref, and the matching :class:`ImageBlock` stores +that ref in ``metadata['markdown_ref']`` for downstream caption +substitution. + +Captioning is not done here — see :class:`ImageBlock` for the +parser→caption contract. + +Output: +- A single ``TextBlock`` containing the rewritten Markdown. +- One ``ImageBlock`` per embedded zip image; ``caption=None``. + +Failures fall back gracefully: missing libraries or malformed zips +degrade to leaving content untouched rather than raising. Images that +can't be decoded by PIL (e.g. EMF, WMF) are skipped — the matching +placeholder in the markdown is left in place for downstream cleanup. +""" + +from __future__ import annotations + +import asyncio +import logging +import re +import zipfile +from io import BytesIO + +from ...models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from ..image_preprocessor import ensure_png_compatible_mode, pil_to_png_bytes +from .document_parser import DocumentParser +from .registry import parser_registry + +logger = logging.getLogger(__name__) + + +# Match MarkItDown's image refs in the rendered markdown. The current +# version emits a truncated placeholder (``![](data:image/png;base64...)``) +# but older / future versions may emit a full data URI or non-empty alt +# text. The pattern below matches both shapes — same regex the legacy +# loader used (``components/indexer/loaders/docx.py``). +_MARKITDOWN_IMAGE_PLACEHOLDER = re.compile(r"!\[.*?\]\(data:image/[^)]*\)") + + +def _image_ref(index: int) -> str: + """Synthetic markdown image ref used as a placeholder for embedded DOCX images.""" + return f"![](docx-image-{index})" + + +@parser_registry.register("docx") +class DocxParser(DocumentParser): + """Parse DOCX into a Markdown TextBlock + one ImageBlock per embedded image.""" + + def supported_types(self) -> list[str]: + return [DocumentType.DOCX.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as path: + markdown, embedded = await asyncio.to_thread(self._extract, str(path)) + + markdown, images = self._rewrite_placeholders_and_build_blocks(markdown, embedded) + markdown = markdown.strip() + text_blocks = [TextBlock(text=markdown, page_number=1)] if markdown else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=dict(document.metadata), + page_count=1 if markdown else 0, + ) + + # ----- helpers ----- + + @classmethod + def _extract(cls, path: str) -> tuple[str, list[bytes | None]]: + """Run MarkItDown + zip-image extraction in one thread hop.""" + return cls._convert_to_markdown(path), cls._extract_embedded_images(path) + + @staticmethod + def _convert_to_markdown(path: str) -> str: + try: + from markitdown import MarkItDown + except ImportError: + logger.warning("markitdown not available; falling back to python-docx text extraction") + return DocxParser._fallback_extract_text(path) + try: + return MarkItDown().convert(path).text_content + except Exception as exc: + logger.warning("MarkItDown DOCX conversion failed (%s); falling back to plain text", exc) + return DocxParser._fallback_extract_text(path) + + @staticmethod + def _fallback_extract_text(path: str) -> str: + try: + from docx import Document as DocxDocument + except ImportError: + logger.warning("python-docx not available; cannot extract DOCX text") + return "" + try: + doc = DocxDocument(path) + return "\n\n".join(p.text for p in doc.paragraphs if p.text.strip()) + except Exception as exc: + logger.warning("python-docx fallback failed: %s", exc) + return "" + + @staticmethod + def _extract_embedded_images(path: str) -> list[bytes | None]: + """Return PNG bytes for each embedded image in document order. + + ``None`` entries preserve positional alignment with markdown + placeholders for unsupported formats (EMF, WMF, …). + """ + try: + from PIL import Image + except ImportError: + logger.warning("PIL not available; skipping DOCX image extraction") + return [] + try: + with zipfile.ZipFile(path, "r") as zf: + media = [n for n in zf.namelist() if n.startswith("word/media/")] + if not media: + return [] + ordered: dict[int, bytes | None] = {} + for name in media: + raw = zf.read(name) + try: + order_num = int(name.split("media/image")[1].split(".")[0]) + except (IndexError, ValueError): + continue + try: + with Image.open(BytesIO(raw)) as im: + im = ensure_png_compatible_mode(im) + ordered[order_num] = pil_to_png_bytes(im) + except Exception as exc: + logger.warning("Skipping unsupported DOCX media %s: %s", name, exc) + ordered[order_num] = None + if not ordered: + return [] + max_order = max(ordered) + return [ordered.get(i + 1) for i in range(max_order)] + except zipfile.BadZipFile: + logger.warning("DOCX is not a valid zip archive; skipping image extraction") + return [] + except Exception as exc: + logger.warning("DOCX image extraction failed: %s", exc) + return [] + + @staticmethod + def _rewrite_placeholders_and_build_blocks( + markdown: str, embedded: list[bytes | None] + ) -> tuple[str, list[ImageBlock]]: + """Replace each MarkItDown placeholder with a unique synthetic ref + and emit one ``ImageBlock`` per successfully-decoded zip image. + + Positional matching: the i-th placeholder in the markdown maps to + the i-th entry in ``embedded``. ``None`` entries (unsupported + formats) collapse the placeholder to an empty string. + """ + if not markdown or not embedded: + return markdown, [] + + images: list[ImageBlock] = [] + idx = 0 + consumed = 0 # count of `embedded` entries used so far + + def replacer(_match: re.Match[str]) -> str: + nonlocal idx, consumed + if consumed >= len(embedded): + return "" # more placeholders than zip images: drop extras + payload = embedded[consumed] + consumed += 1 + if payload is None: + return "" # zip image was unsupported; drop placeholder + ref = _image_ref(idx) + idx += 1 + images.append( + ImageBlock( + image_bytes=payload, + page_number=1, + mime_type="image/png", + metadata={"markdown_ref": ref}, + ) + ) + return ref + + new_markdown = _MARKITDOWN_IMAGE_PLACEHOLDER.sub(replacer, markdown) + return new_markdown, images diff --git a/openrag/core/indexing/parsers/eml_parser.py b/openrag/core/indexing/parsers/eml_parser.py new file mode 100644 index 000000000..760ff9626 --- /dev/null +++ b/openrag/core/indexing/parsers/eml_parser.py @@ -0,0 +1,216 @@ +"""EML (RFC822 email) ``DocumentParser`` implementation. + +Extracts the message body (``text/plain`` preferred, ``text/html`` +fallback) and dispatches each attachment to a parser supplied via DI. +Email headers (subject, from, to, date, message-id) and an attachment +manifest are merged into the output ``ProcessedDocument.metadata``. + +Attachment dispatch contract: + +- ``attachment_parsers`` maps lowercased extension (``"pdf"``, ``"docx"``, + no leading dot) to a :class:`DocumentParser`. +- Each attachment becomes a synthetic :class:`Document` (raw bytes, the + appropriate ``DocumentType`` if recognised, ``DocumentType.TEXT`` otherwise). +- The dispatched parser's text output is appended after a header block + giving filename, content-type, and size. +- Any ``ImageBlock``s emitted by the dispatched parser are propagated + into the EML's own ``ProcessedDocument.images``. +- Image attachments with no registered parser are emitted directly as + ``ImageBlock``s (no ``markdown_ref`` — there is no in-body placeholder + for them; see :class:`ImageBlock` for the parser→caption contract). +- Unknown non-image attachments include only the manifest header. + +Failures are tolerated: a single attachment that errors does not +propagate; we log and continue. +""" + +from __future__ import annotations + +import email +import logging +from collections.abc import Mapping +from email import policy +from email.utils import parsedate_to_datetime + +from ...models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from .document_parser import DocumentParser +from .registry import parser_registry + +logger = logging.getLogger(__name__) + + +_IMAGE_EXTS = {"png", "jpg", "jpeg", "gif", "webp", "bmp", "svg"} + + +@parser_registry.register("eml") +class EmlParser(DocumentParser): + """Parse ``.eml`` into a single text block plus ImageBlocks; dispatch attachments via DI.""" + + def __init__(self, attachment_parsers: Mapping[str, DocumentParser] | None = None) -> None: + self._attachment_parsers = dict(attachment_parsers or {}) + + def supported_types(self) -> list[str]: + return [DocumentType.EML.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + try: + msg = email.message_from_bytes(document.raw_bytes, policy=policy.default) + except Exception as exc: + logger.warning("Failed to parse EML: %s", exc) + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + headers = self._extract_headers(msg) + body, attachments = self._walk_parts(msg) + + attachments_text, images = await self._render_attachments(attachments) + full_text = (body + attachments_text).strip() + + metadata = dict(document.metadata) + metadata.update( + { + "email_subject": headers["subject"], + "email_from": headers["from"], + "email_to": headers["to"], + "email_date": headers["date"], + "email_message_id": headers["message-id"], + "email_attachment_count": len(attachments), + "email_attachment_filenames": [a["filename"] for a in attachments], + } + ) + if attachments: + metadata["email_attachments"] = [ + {"filename": a["filename"], "content_type": a["content_type"], "size": a["size"]} for a in attachments + ] + + text_blocks = [TextBlock(text=full_text, page_number=1)] if full_text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=metadata, + page_count=1 if full_text else 0, + ) + + # ----- helpers ----- + + @staticmethod + def _extract_headers(msg: email.message.Message) -> dict[str, str]: + # Under policy.default, msg.get(...) returns Header-like objects whose + # str() is the RFC 2047-decoded value; cast eagerly so the metadata is + # plain strings. + headers = { + "subject": str(msg.get("subject", "") or ""), + "from": str(msg.get("from", "") or ""), + "to": str(msg.get("to", "") or ""), + "date": str(msg.get("date", "") or ""), + "message-id": str(msg.get("message-id", "") or ""), + } + if headers["date"]: + try: + headers["date"] = parsedate_to_datetime(headers["date"]).isoformat() + except Exception: + pass + return headers + + @staticmethod + def _walk_parts(msg: email.message.Message) -> tuple[str, list[dict]]: + body = "" + attachments: list[dict] = [] + + for part in msg.walk(): + content_type = part.get_content_type() + disposition = part.get_content_disposition() + + if disposition in ("attachment", "inline"): + filename = part.get_filename() + payload = part.get_payload(decode=True) + if filename and payload: + attachments.append( + { + "filename": filename, + "content_type": content_type, + "size": len(payload), + "raw": payload, + } + ) + continue + + if content_type in ("text/plain", "text/html"): + payload = part.get_payload(decode=True) + if not payload: + continue + try: + text = payload.decode("utf-8") if isinstance(payload, bytes) else str(payload) + except UnicodeDecodeError: + text = payload.decode("latin-1", errors="ignore") if isinstance(payload, bytes) else str(payload) + # text/plain wins; only use text/html if we have nothing yet + if content_type == "text/plain" or not body: + body = text + + return body.strip(), attachments + + async def _render_attachments(self, attachments: list[dict]) -> tuple[str, list[ImageBlock]]: + """Render the attachment-section text and collect any ImageBlocks. + + Returns ``("", [])`` when there are no attachments. + """ + if not attachments: + return "", [] + + rendered: list[str] = ["\n\n--- ATTACHMENTS ---\n"] + images: list[ImageBlock] = [] + for att in attachments: + ext = self._extension(att["filename"]) + header = ( + f"\nAttachment: {att['filename']}\nContent-Type: {att['content_type']}\nSize: {att['size']} bytes\n" + ) + content, att_images = await self._render_one(att, ext) + rendered.append(header + content + "---\n") + images.extend(att_images) + return "".join(rendered), images + + async def _render_one(self, attachment: dict, ext: str) -> tuple[str, list[ImageBlock]]: + """Dispatch one attachment. Returns ``(text_to_inline, image_blocks)``.""" + parser = self._attachment_parsers.get(ext) + if parser is not None: + try: + doc = Document( + filename=attachment["filename"], + raw_bytes=attachment["raw"], + content_type=Document.detect_content_type(attachment["filename"]), + metadata={"source": f"attachment:{attachment['filename']}"}, + ) + processed = await parser.parse(doc) + content = "\n\n".join(b.text for b in processed.text_blocks if b.text) + inline = f"Content:\n{content}\n" if content else "" + return inline, list(processed.images) + except Exception as exc: + logger.warning("Attachment parser failed for %s: %s", attachment["filename"], exc) + + if ext in _IMAGE_EXTS: + # No parser registered — emit the image as an ImageBlock so a + # downstream caption stage can describe it. No ``markdown_ref`` + # because there is no in-body placeholder pointing to it. + return "", [ + ImageBlock( + image_bytes=attachment["raw"], + page_number=1, + mime_type=attachment["content_type"] or "image/png", + metadata={"source": f"attachment:{attachment['filename']}"}, + ) + ] + + return "", [] + + @staticmethod + def _extension(filename: str) -> str: + return filename.rsplit(".", 1)[-1].lower() if "." in filename else "" diff --git a/openrag/core/indexing/parsers/html_parser.py b/openrag/core/indexing/parsers/html_parser.py new file mode 100644 index 000000000..d51d2abb1 --- /dev/null +++ b/openrag/core/indexing/parsers/html_parser.py @@ -0,0 +1,58 @@ +"""HTML ``DocumentParser`` implementation. + +Converts HTML to Markdown via the project's ``html_to_markdown`` dep +(already used by the websearch and pptx pipelines), then emits a single +text block. No file I/O, no JavaScript execution, no image fetching — +purely structural conversion. +""" + +from __future__ import annotations + +import asyncio + +from ...models.document import Document, DocumentType, ProcessedDocument, TextBlock +from ..text_preprocessor import decode_bytes +from .document_parser import DocumentParser +from .registry import parser_registry + + +@parser_registry.register("html") +class HtmlParser(DocumentParser): + """Parse HTML documents into a single Markdown text block.""" + + def __init__(self, *, encoding: str | None = None) -> None: + """``encoding`` forces a specific decode of ``raw_bytes``; ``None`` + auto-detects (UTF-8 first, then chardet). + """ + self._encoding = encoding + + def supported_types(self) -> list[str]: + return [DocumentType.HTML.value] + + async def parse(self, document: Document) -> ProcessedDocument: + markdown = (await asyncio.to_thread(self._html_to_markdown, document)).strip() + text_blocks = [TextBlock(text=markdown, page_number=1)] if markdown else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + metadata=dict(document.metadata), + page_count=1 if markdown else 0, + ) + + def _html_to_markdown(self, document: Document) -> str: + """Decode + HTML→Markdown in one shot. Sync; runs in a thread.""" + html = self._extract_html(document) + return self._to_markdown(html) if html else "" + + def _extract_html(self, document: Document) -> str: + if document.text is not None: + return document.text + if document.raw_bytes: + return decode_bytes(document.raw_bytes, encoding=self._encoding) + return "" + + @staticmethod + def _to_markdown(html: str) -> str: + from html_to_markdown import convert + + return convert(html) diff --git a/openrag/core/indexing/parsers/image_parser.py b/openrag/core/indexing/parsers/image_parser.py new file mode 100644 index 000000000..4e6d5b767 --- /dev/null +++ b/openrag/core/indexing/parsers/image_parser.py @@ -0,0 +1,141 @@ +"""Image ``DocumentParser`` implementation. + +Decodes an image into normalized PNG bytes and emits a single +:class:`ImageBlock`. Supports raster formats (PNG/JPEG/etc. — anything +PIL opens) and SVG (rasterized to PNG via cairosvg). + +Captioning is not done here — see :class:`ImageBlock` for the +parser→caption contract. + +Output: +- A single ``ImageBlock`` with the normalized PNG bytes and no caption. +- No ``TextBlock`` is emitted; downstream stages produce text from the image. + +Failures (decode errors, undersized images) emit an empty +``ProcessedDocument`` rather than raising — RAG pipelines should not die +on a single bad image. +""" + +from __future__ import annotations + +import asyncio +import logging + +from ...models.document import Document, DocumentType, ImageBlock, ProcessedDocument +from ..image_preprocessor import MIN_IMAGE_PIXELS, ensure_png_compatible_mode +from .document_parser import DocumentParser +from .registry import parser_registry + +logger = logging.getLogger(__name__) + + +@parser_registry.register("image") +class ImageParser(DocumentParser): + """Decode an image and emit it as a single ``ImageBlock``.""" + + def __init__(self, *, min_pixels: int = MIN_IMAGE_PIXELS) -> None: + self._min_pixels = max(0, min_pixels) + + def supported_types(self) -> list[str]: + return [DocumentType.IMAGE.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + png_bytes = await asyncio.to_thread(self._normalize_to_png, document) + if png_bytes is None: + logger.warning("ImageParser: failed to decode image (id=%s)", document.id) + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + if await asyncio.to_thread(self._below_min_pixels, png_bytes): + logger.warning("ImageParser: image below min_pixels threshold (id=%s)", document.id) + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + return ProcessedDocument( + document_id=document.id, + images=[ + ImageBlock( + image_bytes=png_bytes, + page_number=1, + mime_type="image/png", + ) + ], + metadata=dict(document.metadata), + page_count=1, + ) + + def _normalize_to_png(self, document: Document) -> bytes | None: + """Return PNG bytes for any supported image input, or None on failure.""" + raw = document.raw_bytes + if not raw: + return None + if self._is_svg(raw, document.filename): + return self._svg_to_png(raw) + return self._raster_to_png(raw) + + @staticmethod + def _is_svg(raw: bytes, filename: str) -> bool: + if filename.lower().endswith(".svg"): + return True + head = raw[:200].lstrip().lower() + return head.startswith((b" bytes | None: + try: + import cairosvg + + return cairosvg.svg2png(bytestring=raw) + except Exception as exc: + logger.warning("Failed to rasterize SVG: %s", exc) + return None + + @staticmethod + def _raster_to_png(raw: bytes) -> bytes | None: + """Decode raw bytes through PIL and re-encode as PNG. + + Re-encoding normalizes the format so downstream consumers + (caption stage, vector-store image fields) only need to handle + one mime type, and validates the image is decodable. + """ + try: + from io import BytesIO + + from PIL import Image + except ImportError: + logger.warning("PIL not available; cannot decode raster image") + return None + try: + with Image.open(BytesIO(raw)) as image: + image = ensure_png_compatible_mode(image) + buf = BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + except Exception as exc: + logger.warning("Failed to decode image: %s", exc) + return None + + def _below_min_pixels(self, png_bytes: bytes) -> bool: + if self._min_pixels <= 0: + return False + try: + from io import BytesIO + + from PIL import Image + except ImportError: + return False + try: + with Image.open(BytesIO(png_bytes)) as image: + return image.width * image.height < self._min_pixels + except Exception: + return False diff --git a/openrag/core/indexing/parsers/markdown_parser.py b/openrag/core/indexing/parsers/markdown_parser.py new file mode 100644 index 000000000..566ea6de0 --- /dev/null +++ b/openrag/core/indexing/parsers/markdown_parser.py @@ -0,0 +1,69 @@ +"""Markdown ``DocumentParser`` implementation. + +Decodes a Markdown document into a single :class:`TextBlock` and emits +one :class:`ImageBlock` per image reference in the source: + +- Data-URI refs (``![alt](data:image/...;base64,...)``) are decoded and + the bytes stored on the block. +- HTTP/HTTPS refs (``![alt](https://...)``) yield an :class:`ImageBlock` + with empty ``image_bytes`` and ``source_url`` set; a downstream fetch + stage can populate the bytes later. The :attr:`ImageBlock.image_url` + property gives a uniform VLM-friendly URL in either case. + +Captioning is not done here — see :class:`ImageBlock` for the +parser→caption contract. +""" + +from __future__ import annotations + +from ...models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from ..image_preprocessor import HTTP_IMAGE_PATTERN, extract_data_uri_image_blocks +from ..text_preprocessor import decode_bytes +from .document_parser import DocumentParser +from .registry import parser_registry + + +@parser_registry.register("markdown") +class MarkdownParser(DocumentParser): + """Parse Markdown documents and emit ImageBlocks for every image ref.""" + + def __init__(self, *, encoding: str | None = None) -> None: + self._encoding = encoding + + def supported_types(self) -> list[str]: + return [DocumentType.MARKDOWN.value] + + async def parse(self, document: Document) -> ProcessedDocument: + text = self._extract_text(document).strip() + images = self._extract_image_blocks(text) + + text_blocks = [TextBlock(text=text, page_number=1)] if text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=dict(document.metadata), + page_count=1 if text else 0, + ) + + def _extract_text(self, document: Document) -> str: + if document.text is not None: + return document.text + if document.raw_bytes: + return decode_bytes(document.raw_bytes, encoding=self._encoding) + return "" + + @staticmethod + def _extract_image_blocks(text: str) -> list[ImageBlock]: + if not text: + return [] + blocks: list[ImageBlock] = list(extract_data_uri_image_blocks(text, page_number=1)) + for alt, url in HTTP_IMAGE_PATTERN.findall(text): + blocks.append( + ImageBlock( + source_url=url, + page_number=1, + metadata={"markdown_ref": f"![{alt}]({url})", "alt": alt}, + ) + ) + return blocks diff --git a/openrag/core/indexing/parsers/pdf/__init__.py b/openrag/core/indexing/parsers/pdf/__init__.py new file mode 100644 index 000000000..d48cea8b6 --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/__init__.py @@ -0,0 +1,9 @@ +"""PDF parser backends. + +Each backend lives in its own module so its heavy dependencies (Marker, +Docling, DotsOCR, …) are only imported when the backend's submodule is +itself imported. Consumers do +``from core.indexing.parsers.pdf.marker import MarkerParser`` rather +than going through this package, so importing ``pdf`` does not pull in +any backend. +""" diff --git a/openrag/core/indexing/parsers/pdf/client_based.py b/openrag/core/indexing/parsers/pdf/client_based.py new file mode 100644 index 000000000..5e355b945 --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/client_based.py @@ -0,0 +1,33 @@ +"""OpenAI-VLM-backed PDF ``DocumentParser`` (thin core facade). + +Holds a ``BaseClientParser`` (the actual HTTP-client / OpenAI-SDK +implementation lives in ``services/`` and is composed in at startup) and +delegates ``parse()`` to it. + +Mirrors the :class:`MarkerParser` pattern: core stays free of vendor +SDKs while the facade names "OpenAI-VLM PDF" as a first-class parser +type. Concrete subclasses of the services-side base (e.g. DotsOCR) can +be swapped in without changing this facade. +""" + +from __future__ import annotations + +from ....models.document import Document, ProcessedDocument +from ..document_parser import BaseClientParser, DocumentParser +from ..registry import parser_registry + + +@parser_registry.register("pdf_client") +class ClientPdfParser(DocumentParser): + """Public PDF parser facade backed by an OpenAI-compatible VLM client.""" + + def __init__(self, client: BaseClientParser) -> None: + if not isinstance(client, BaseClientParser): + raise ValueError("ClientPdfParser requires a BaseClientParser instance as client") + self._client = client + + def supported_types(self) -> list[str]: + return self._client.supported_types() + + async def parse(self, document: Document) -> ProcessedDocument: + return await self._client.parse(document) diff --git a/openrag/core/indexing/parsers/pdf/marker.py b/openrag/core/indexing/parsers/pdf/marker.py new file mode 100644 index 000000000..db17f0ce8 --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/marker.py @@ -0,0 +1,36 @@ +"""Marker-backed PDF ``DocumentParser`` (thin core facade). + +Holds a reference to a ``BasePooledParser`` (the actual Ray-pool / +GPU-model / process-pool implementation lives in ``services/`` and is +not yet wired up) and delegates ``parse()`` to it. The split keeps core +free of Ray and GPU lifecycle code while still naming the Marker +backend as a first-class parser type. + +The injected pool is a generic ``BasePooledParser``; if a more specific +``MarkerPoolParser`` ABC emerges in services, this class can tighten +its type without changing call sites. +""" + +from __future__ import annotations + +from ....models.document import Document, ProcessedDocument +from ..document_parser import BasePooledParser, DocumentParser +from ..registry import parser_registry + + +@parser_registry.register("marker") +class MarkerParser(DocumentParser): + """Public PDF parser facade backed by a Marker worker pool.""" + + def __init__(self, pool: BasePooledParser) -> None: + # check pool is a BasePooledParser? and not empty + if not isinstance(pool, BasePooledParser) or pool is None: + raise ValueError("MarkerParser requires a BasePooledParser instance as pool") + + self._pool = pool + + def supported_types(self) -> list[str]: + return self._pool.supported_types() + + async def parse(self, document: Document) -> ProcessedDocument: + return await self._pool.parse(document) diff --git a/openrag/core/indexing/parsers/pdf/pymupdf.py b/openrag/core/indexing/parsers/pdf/pymupdf.py new file mode 100644 index 000000000..f42b33755 --- /dev/null +++ b/openrag/core/indexing/parsers/pdf/pymupdf.py @@ -0,0 +1,114 @@ +"""PyMuPDF-backed PDF ``DocumentParser``. + +The lightweight, no-VLM, no-GPU PDF backend. Uses ``pymupdf`` (a.k.a. +``fitz``) for plain-text extraction and ``pymupdf4llm`` for Markdown +extraction. Operates on ``Document.raw_bytes`` — file I/O is upstream. + +In ``mode="markdown"``, embedded images are surfaced as ``ImageBlock``s +via ``pymupdf4llm``'s ``embed_images=True`` (each image becomes a +``data:image/png;base64,…`` ref in the markdown, which we decode into +an :class:`ImageBlock` with ``markdown_ref`` set so a downstream caption +stage can substitute a description back in). ``mode="text"`` does not +extract images. + +Threading note: PyMuPDF is **not** thread-safe — concurrent calls to +``page.get_text`` / ``pymupdf4llm.to_markdown`` from different threads +can raise ``ValueError: not a textpage of this page`` (upstream +maintainer position: documented limitation, won't fix). We therefore +serialize all pymupdf work onto a single dedicated worker thread via +``_PYMUPDF_EXECUTOR``. The async ``parse`` method stays concurrent — +multiple callers will queue on the executor, but only one pymupdf +operation runs at a time. +""" + +from __future__ import annotations + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from typing import Literal + +import pymupdf +import pymupdf4llm + +from ....models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from ...image_preprocessor import extract_data_uri_image_blocks +from ..document_parser import DocumentParser +from ..registry import parser_registry + +ParseMode = Literal["markdown", "text"] + +# Single dedicated worker for pymupdf — see "Threading note" in module docstring. +_PYMUPDF_EXECUTOR = ThreadPoolExecutor(max_workers=1, thread_name_prefix="pymupdf") + + +def _extract_text(raw: bytes) -> tuple[list[str], list[ImageBlock]]: + """Return one stripped plain-text string per page; no images.""" + with pymupdf.open(stream=raw, filetype="pdf") as doc: + return [page.get_text().strip() for page in doc], [] + + +def _extract_markdown(raw: bytes) -> tuple[list[str], list[ImageBlock]]: + """Return Markdown per page + ``ImageBlock``s built from embedded data URIs. + + ``embed_images=True`` makes ``pymupdf4llm`` write images as base64 + data URIs in-line. We decode each ref into an ``ImageBlock`` and + leave the ref in the page text untouched so the caption stage can + substitute later via ``ImageBlock.metadata['markdown_ref']``. + """ + with pymupdf.open(stream=raw, filetype="pdf") as doc: + chunks = pymupdf4llm.to_markdown( + doc, + page_chunks=True, + embed_images=True, + write_images=False, + dpi=300, + ) + pages: list[str] = [] + images: list[ImageBlock] = [] + for i, chunk in enumerate(chunks, start=1): + text = (chunk.get("text") or "").strip() + pages.append(text) + if text: + images.extend(extract_data_uri_image_blocks(text, page_number=i)) + return pages, images + + +@parser_registry.register("pymupdf") +class PyMuPDFParser(DocumentParser): + """Extract text from a PDF as one ``TextBlock`` per page (+ ImageBlocks in markdown mode). + + ``mode="markdown"`` (default) uses ``pymupdf4llm`` for layout-preserving + Markdown — better for downstream embedding and chunking, and surfaces + embedded images. ``mode="text"`` uses raw ``pymupdf`` for plain text — + slightly faster, no formatting, no images. + """ + + def __init__(self, *, mode: ParseMode = "markdown") -> None: + if mode not in ("markdown", "text"): + raise ValueError(f"PyMuPDFParser: unsupported mode {mode!r}") + self._mode = mode + self._extract = _extract_text if mode == "text" else _extract_markdown + + def supported_types(self) -> list[str]: + return [DocumentType.PDF.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + pages, images = await asyncio.get_running_loop().run_in_executor( + _PYMUPDF_EXECUTOR, self._extract, document.raw_bytes + ) + # Keep one TextBlock per source page (including empties) so callers + # can preserve a 1-to-1 mapping with the original PDF's pagination. + text_blocks = [TextBlock(text=text, page_number=i) for i, text in enumerate(pages, start=1)] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=dict(document.metadata), + page_count=len(pages), + ) diff --git a/openrag/core/indexing/parsers/pptx_parser.py b/openrag/core/indexing/parsers/pptx_parser.py new file mode 100644 index 000000000..0ccd1b8d6 --- /dev/null +++ b/openrag/core/indexing/parsers/pptx_parser.py @@ -0,0 +1,195 @@ +"""PPTX ``DocumentParser`` implementation. + +Walks slides via ``python-pptx``, converting each slide to Markdown: +title → ``#`` heading, text frames → paragraphs, tables → HTML→Markdown, +charts → Markdown tables, pictures → ``![](pptx-image-N)`` synthetic +markdown image refs. Speaker notes are appended as ``### Notes:``. + +Captioning is not done here — see :class:`ImageBlock` for the +parser→caption contract. + +Output is one :class:`TextBlock` per slide (1-indexed ``page_number``) +plus one :class:`ImageBlock` per slide picture. ``page_number`` on each +``ImageBlock`` is the slide number it came from. + +Implementation derived from the legacy ``PPTXConverter`` (which itself +mirrored the MarkItDown PPTX converter). +""" + +from __future__ import annotations + +import asyncio +import html +import logging +from io import BytesIO +from typing import Any + +from ...models.document import Document, DocumentType, ImageBlock, ProcessedDocument, TextBlock +from ..image_preprocessor import ensure_png_compatible_mode, pil_to_png_bytes +from .document_parser import DocumentParser +from .registry import parser_registry + +logger = logging.getLogger(__name__) + + +def _image_ref(index: int) -> str: + """Synthetic markdown image ref used as a placeholder for slide pictures.""" + return f"![](pptx-image-{index})" + + +@parser_registry.register("pptx") +class PptxParser(DocumentParser): + """Parse PPTX into one TextBlock per slide plus one ImageBlock per picture.""" + + def supported_types(self) -> list[str]: + return [DocumentType.PPTX.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as path: + slide_count, slides, images = await asyncio.to_thread(self._convert, str(path)) + + text_blocks = [TextBlock(text=text, page_number=slide_num) for slide_num, text in slides] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=dict(document.metadata), + page_count=slide_count, + ) + + # ----- conversion ----- + + def _convert(self, path: str) -> tuple[int, list[tuple[int, str]], list[ImageBlock]]: + try: + import pptx + from PIL import Image + except ImportError: + logger.warning("python-pptx or PIL not available; cannot parse PPTX") + return 0, [], [] + + try: + presentation = pptx.Presentation(path) + except Exception as exc: + logger.warning("Failed to open PPTX: %s", exc) + return 0, [], [] + + slides: list[tuple[int, str]] = [] + images: list[ImageBlock] = [] + + for slide_num, slide in enumerate(presentation.slides, start=1): + md = "" + title = slide.shapes.title + + for shape in slide.shapes: + if self._is_picture(shape): + try: + with Image.open(BytesIO(shape.image.blob)) as im: + im = ensure_png_compatible_mode(im) + png_bytes = pil_to_png_bytes(im) + ref = _image_ref(len(images)) + images.append( + ImageBlock( + image_bytes=png_bytes, + page_number=slide_num, + mime_type="image/png", + metadata={"markdown_ref": ref}, + ) + ) + md += ref + except Exception as exc: + logger.warning("Skipping unreadable PPTX picture: %s", exc) + elif self._is_table(shape): + md += "\n" + self._table_to_markdown(shape.table) + "\n" + elif getattr(shape, "has_chart", False): + md += self._chart_to_markdown(shape.chart) + elif getattr(shape, "has_text_frame", False): + if shape == title: + md += "# " + shape.text.lstrip() + "\n" + else: + md += shape.text + "\n" + + md = md.strip() + if slide.has_notes_slide: + notes_frame = slide.notes_slide.notes_text_frame + if notes_frame is not None: + md += "\n\n### Notes:\n" + notes_frame.text + md = md.strip() + + if md: + slides.append((slide_num, md)) + + return len(presentation.slides), slides, images + + @staticmethod + def _is_picture(shape: Any) -> bool: + try: + from pptx.enum.shapes import MSO_SHAPE_TYPE + + if shape.shape_type == MSO_SHAPE_TYPE.PICTURE: + return True + if shape.shape_type == MSO_SHAPE_TYPE.PLACEHOLDER and hasattr(shape, "image"): + return True + except NotImplementedError: + logger.debug("Encountered an unimplemented shape type") + except Exception: + return False + return False + + @staticmethod + def _is_table(shape: Any) -> bool: + try: + from pptx.enum.shapes import MSO_SHAPE_TYPE + + return shape.shape_type == MSO_SHAPE_TYPE.TABLE + except NotImplementedError: + logger.debug("Encountered an unimplemented shape type") + return False + except Exception: + return False + + @staticmethod + def _table_to_markdown(table: Any) -> str: + from html_to_markdown import convert + + html_rows = [""] + first_row = True + for row in table.rows: + html_rows.append("") + for cell in row.cells: + tag = "th" if first_row else "td" + html_rows.append(f"<{tag}>{html.escape(cell.text)}") + html_rows.append("") + first_row = False + html_rows.append("
") + return convert("".join(html_rows)).strip() + + @staticmethod + def _chart_to_markdown(chart: Any) -> str: + try: + md = "\n\n### Chart" + if chart.has_title: + md += f": {chart.chart_title.text_frame.text}" + md += "\n\n" + category_names = [c.label for c in chart.plots[0].categories] + series_names = [s.name for s in chart.series] + data: list[list[str]] = [["Category"] + series_names] + for idx, category in enumerate(category_names): + row = [category] + for series in chart.series: + row.append(series.values[idx]) + data.append(row) + rows = ["| " + " | ".join(map(str, r)) + " |" for r in data] + separator = "|" + "|".join(["---"] * len(data[0])) + "|" + return md + "\n".join([rows[0], separator] + rows[1:]) + except ValueError as exc: + if "unsupported plot type" in str(exc): + return "\n\n[unsupported chart]\n\n" + return "\n\n[unsupported chart]\n\n" + except Exception: + return "\n\n[unsupported chart]\n\n" diff --git a/openrag/core/indexing/parsers/test_doc_parser.py b/openrag/core/indexing/parsers/test_doc_parser.py new file mode 100644 index 000000000..454de181f --- /dev/null +++ b/openrag/core/indexing/parsers/test_doc_parser.py @@ -0,0 +1,118 @@ +"""Unit tests for :class:`DocParser` (.doc → DocxParser delegation + fallback).""" + +from __future__ import annotations + +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from ...models.document import Document, DocumentType, ProcessedDocument, TextBlock +from .doc_parser import DocParser + + +@pytest.fixture +def fake_spire(): + """Inject a fake ``spire.doc`` into ``sys.modules`` for the duration of a test. + + Returns the ``Document`` mock class so tests can configure the + instance returned by ``Document()``. + """ + spire = MagicMock() + spire_doc = MagicMock() + spire.doc = spire_doc + saved = {k: sys.modules.get(k) for k in ("spire", "spire.doc")} + sys.modules["spire"] = spire + sys.modules["spire.doc"] = spire_doc + try: + yield spire_doc.Document + finally: + for k, v in saved.items(): + if v is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = v + + +def _doc_document(raw: bytes = b"\xd0\xcf\x11\xe0fake-doc") -> Document: + return Document(filename="x.doc", content_type=DocumentType.DOC, raw_bytes=raw) + + +class TestParse: + @pytest.mark.asyncio + async def test_empty_raw_bytes_returns_empty(self): + doc = _doc_document(raw=b"") + result = await DocParser().parse(doc) + assert result.text_blocks == [] and result.images == [] and result.page_count == 0 + + @pytest.mark.asyncio + async def test_successful_conversion_delegates_to_docx(self, fake_spire, tmp_path): + # Spire writes a real .docx file at the path given to SaveToFile. + dummy_docx = b"DOCX-CONTENT" + + instance = MagicMock() + + def save_to_file(path: str, _fmt) -> None: + with open(path, "wb") as fh: + fh.write(dummy_docx) + + instance.SaveToFile.side_effect = save_to_file + fake_spire.return_value = instance + + docx_parser = MagicMock() + expected = ProcessedDocument( + document_id="test", + text_blocks=[TextBlock(text="from-docx", page_number=1)], + page_count=1, + ) + docx_parser.parse = AsyncMock(return_value=expected) + + parser = DocParser(docx_parser=docx_parser) + result = await parser.parse(_doc_document()) + + assert result is expected + docx_parser.parse.assert_awaited_once() + forwarded = docx_parser.parse.await_args.args[0] + assert forwarded.raw_bytes == dummy_docx + assert forwarded.content_type is DocumentType.DOCX + instance.LoadFromFile.assert_called_once() + instance.Close.assert_called_once() + + @pytest.mark.asyncio + async def test_save_failure_falls_back_to_get_text(self, fake_spire): + instance = MagicMock() + instance.SaveToFile.side_effect = RuntimeError("Spire crashed") + instance.GetText.return_value = " plain text content " + fake_spire.return_value = instance + + docx_parser = MagicMock() + docx_parser.parse = AsyncMock() + result = await DocParser(docx_parser=docx_parser).parse(_doc_document()) + + assert result.text_blocks == [TextBlock(text="plain text content", page_number=1)] + assert result.page_count == 1 + instance.GetText.assert_called_once() + instance.Close.assert_called_once() + docx_parser.parse.assert_not_awaited() # never delegated + + @pytest.mark.asyncio + async def test_total_failure_returns_empty(self, fake_spire): + instance = MagicMock() + instance.SaveToFile.side_effect = RuntimeError("Spire crashed") + instance.GetText.side_effect = RuntimeError("GetText crashed") + fake_spire.return_value = instance + + result = await DocParser().parse(_doc_document()) + assert result.text_blocks == [] and result.page_count == 0 + instance.Close.assert_called_once() + + @pytest.mark.asyncio + async def test_missing_spire_returns_empty(self, monkeypatch): + # spire-doc is in the runtime deps, so just omitting fake_spire would + # actually drive a real Spire instance against malformed bytes. + # Pin the import to None so ``_convert``'s ``import spire.doc`` raises + # ImportError deterministically. + monkeypatch.setitem(sys.modules, "spire", None) + monkeypatch.setitem(sys.modules, "spire.doc", None) + result = await DocParser().parse(_doc_document()) + assert result.text_blocks == [] and result.page_count == 0 diff --git a/openrag/core/indexing/parsers/test_docx_parser.py b/openrag/core/indexing/parsers/test_docx_parser.py new file mode 100644 index 000000000..cf612b069 --- /dev/null +++ b/openrag/core/indexing/parsers/test_docx_parser.py @@ -0,0 +1,141 @@ +"""Unit tests for :class:`DocxParser`.""" + +from __future__ import annotations + +import tempfile +import zipfile +from io import BytesIO +from pathlib import Path + +import pytest +from PIL import Image + +from ...models.document import Document, DocumentType +from .docx_parser import DocxParser, _image_ref + + +def _png_bytes(color: str = "red") -> bytes: + img = Image.new("RGBA", (10, 10), color) + buf = BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + +def _fake_docx(media_files: dict[str, bytes]) -> Path: + """Build a minimal .docx zip with given ``word/media/`` entries.""" + tmp = tempfile.NamedTemporaryFile(suffix=".docx", delete=False) + with zipfile.ZipFile(tmp, "w") as zf: + for name, data in media_files.items(): + zf.writestr(f"word/media/{name}", data) + return Path(tmp.name) + + +class TestExtractEmbeddedImages: + """Mirrors legacy ``TestGetImagesFromZip`` against the new staticmethod.""" + + def test_valid_images_kept_in_order(self): + docx = _fake_docx({"image2.png": _png_bytes("blue"), "image1.png": _png_bytes("red")}) + result = DocxParser._extract_embedded_images(str(docx)) + assert len(result) == 2 + assert result[0] is not None and result[1] is not None + + def test_unsupported_format_collapses_to_none_at_position(self): + docx = _fake_docx( + { + "image1.png": _png_bytes(), + "image2.emf": b"\x01\x00\x00\x00garbage", + "image3.png": _png_bytes(), + } + ) + result = DocxParser._extract_embedded_images(str(docx)) + assert len(result) == 3 + assert result[0] is not None + assert result[1] is None + assert result[2] is not None + + def test_non_image_media_skipped(self): + docx = _fake_docx( + { + "image1.png": _png_bytes(), + "oleObject1.bin": b"OLE", + "hdphoto1.wdp": b"WDP", + } + ) + result = DocxParser._extract_embedded_images(str(docx)) + assert sum(1 for x in result if x is not None) == 1 + + def test_all_unsupported_returns_empty(self): + docx = _fake_docx({"image1.emf": b"EMF", "image2.wmf": b"WMF"}) + # Two unsupported entries: positional list still has length 2 with None slots. + result = DocxParser._extract_embedded_images(str(docx)) + assert all(x is None for x in result) + + def test_no_media_returns_empty(self): + tmp = tempfile.NamedTemporaryFile(suffix=".docx", delete=False) + with zipfile.ZipFile(tmp, "w") as zf: + zf.writestr("word/document.xml", "") + result = DocxParser._extract_embedded_images(tmp.name) + assert result == [] + + def test_invalid_zip_returns_empty(self): + tmp = tempfile.NamedTemporaryFile(suffix=".docx", delete=False) + tmp.write(b"not a zip") + tmp.flush() + assert DocxParser._extract_embedded_images(tmp.name) == [] + + +class TestRewritePlaceholdersAndBuildBlocks: + """The parser→caption contract: synthetic refs + ImageBlock metadata.""" + + def test_assigns_unique_refs_in_order(self): + md = "before ![](data:image/png;base64,trunc...) middle ![](data:image/png;base64,trunc...) end" + embedded = [_png_bytes("red"), _png_bytes("blue")] + new_md, blocks = DocxParser._rewrite_placeholders_and_build_blocks(md, embedded) + + assert _image_ref(0) in new_md + assert _image_ref(1) in new_md + assert len(blocks) == 2 + assert blocks[0].metadata["markdown_ref"] == _image_ref(0) + assert blocks[1].metadata["markdown_ref"] == _image_ref(1) + assert blocks[0].image_bytes == embedded[0] + assert blocks[1].image_bytes == embedded[1] + assert all(b.page_number == 1 and b.mime_type == "image/png" for b in blocks) + + def test_none_entry_drops_placeholder(self): + md = "![](data:image/png;base64,a) ![](data:image/png;base64,b)" + embedded = [None, _png_bytes()] + new_md, blocks = DocxParser._rewrite_placeholders_and_build_blocks(md, embedded) + + # First placeholder collapses to empty; second becomes ref-0 (only one block emitted). + assert _image_ref(0) in new_md + assert _image_ref(1) not in new_md + assert len(blocks) == 1 + + def test_more_placeholders_than_zip_entries_drops_extras(self): + md = "![](data:image/png;base64,a) ![](data:image/png;base64,b)" + embedded = [_png_bytes()] + new_md, blocks = DocxParser._rewrite_placeholders_and_build_blocks(md, embedded) + + assert _image_ref(0) in new_md + # The extra placeholder is dropped — the regex match collapses to "". + assert "data:image" not in new_md + assert len(blocks) == 1 + + def test_no_placeholders_passthrough(self): + new_md, blocks = DocxParser._rewrite_placeholders_and_build_blocks("plain text", [_png_bytes()]) + assert new_md == "plain text" + assert blocks == [] + + def test_empty_inputs(self): + assert DocxParser._rewrite_placeholders_and_build_blocks("", []) == ("", []) + assert DocxParser._rewrite_placeholders_and_build_blocks("text", []) == ("text", []) + + +class TestParse: + @pytest.mark.asyncio + async def test_empty_raw_bytes_returns_empty(self): + doc = Document(filename="x.docx", content_type=DocumentType.DOCX, raw_bytes=b"") + result = await DocxParser().parse(doc) + assert result.text_blocks == [] + assert result.images == [] + assert result.page_count == 0 diff --git a/openrag/core/indexing/parsers/text_parser.py b/openrag/core/indexing/parsers/text_parser.py new file mode 100644 index 000000000..b23b1a2b4 --- /dev/null +++ b/openrag/core/indexing/parsers/text_parser.py @@ -0,0 +1,49 @@ +"""Plain-text ``DocumentParser`` implementation. + +Decodes ``Document.raw_bytes`` (or uses ``Document.text`` if already +populated) into a single :class:`TextBlock`. Performs no image captioning, +no markdown-image extraction, and no file I/O — those are upstream +concerns. Handles the ``TEXT`` content type; Markdown (with image +captioning) lives in :class:`MarkdownParser`; HTML, PDF, and richer +formats live in their own parsers. +""" + +from __future__ import annotations + +import asyncio + +from ...models.document import Document, DocumentType, ProcessedDocument, TextBlock +from ..text_preprocessor import decode_bytes +from .document_parser import DocumentParser +from .registry import parser_registry + + +@parser_registry.register("text") +class TextParser(DocumentParser): + """Parse plain-text documents into a single text block.""" + + def __init__(self, *, encoding: str | None = None) -> None: + """If ``encoding`` is ``None``, raw bytes are auto-detected via + :func:`core.indexing.text_preprocessor.decode_bytes`. + """ + self._encoding = encoding + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + + async def parse(self, document: Document) -> ProcessedDocument: + text = (await asyncio.to_thread(self._extract_text, document)).strip() + text_blocks = [TextBlock(text=text, page_number=1)] if text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + metadata=dict(document.metadata), + page_count=1 if text else 0, + ) + + def _extract_text(self, document: Document) -> str: + if document.text is not None: + return document.text + if document.raw_bytes: + return decode_bytes(document.raw_bytes, encoding=self._encoding) + return "" diff --git a/openrag/core/indexing/test_image_preprocessor.py b/openrag/core/indexing/test_image_preprocessor.py new file mode 100644 index 000000000..29ae9ec5d --- /dev/null +++ b/openrag/core/indexing/test_image_preprocessor.py @@ -0,0 +1,98 @@ +"""Unit tests for ``core.indexing.image_preprocessor``.""" + +from __future__ import annotations + +import base64 + +from PIL import Image + +from .image_preprocessor import ( + MIN_IMAGE_PIXELS, + decode_data_uri, + ensure_png_compatible_mode, + extract_data_uri_image_blocks, + mime_from_data_uri, + pil_to_png_bytes, +) + + +class TestEnsurePngCompatibleMode: + def test_cmyk_to_rgb(self): + assert ensure_png_compatible_mode(Image.new("CMYK", (10, 10))).mode == "RGB" + + def test_palette_to_rgba(self): + assert ensure_png_compatible_mode(Image.new("P", (10, 10))).mode == "RGBA" + + def test_la_to_rgba(self): + assert ensure_png_compatible_mode(Image.new("LA", (10, 10))).mode == "RGBA" + + def test_rgb_unchanged(self): + assert ensure_png_compatible_mode(Image.new("RGB", (10, 10))).mode == "RGB" + + def test_rgba_unchanged(self): + assert ensure_png_compatible_mode(Image.new("RGBA", (10, 10))).mode == "RGBA" + + +class TestPilToPngBytes: + def test_rgb_round_trip(self): + img = Image.new("RGB", (32, 32), "red") + png = pil_to_png_bytes(img) + assert png[:8] == b"\x89PNG\r\n\x1a\n" + + def test_cmyk_normalised_then_encoded(self): + png = pil_to_png_bytes(Image.new("CMYK", (16, 16))) + assert png[:8] == b"\x89PNG\r\n\x1a\n" + + def test_bytes_passthrough(self): + raw = b"already-bytes" + assert pil_to_png_bytes(raw) is raw + + +class TestDecodeDataUri: + def test_round_trip(self): + payload = b"hello" + uri = f"data:image/png;base64,{base64.b64encode(payload).decode()}" + assert decode_data_uri(uri) == payload + + def test_malformed_returns_none(self): + assert decode_data_uri("not-a-data-uri") is None + assert decode_data_uri("data:image/png;base64,!!!not-base64") is None + + +class TestMimeFromDataUri: + def test_jpeg(self): + assert mime_from_data_uri("data:image/jpeg;base64,xxx") == "image/jpeg" + + def test_png(self): + assert mime_from_data_uri("data:image/png;base64,xxx") == "image/png" + + def test_malformed_falls_back_to_png(self): + assert mime_from_data_uri("garbage") == "image/png" + + +class TestExtractDataUriImageBlocks: + def _data_uri(self, payload: bytes = b"x", mime: str = "image/png") -> str: + return f"data:{mime};base64,{base64.b64encode(payload).decode()}" + + def test_emits_one_block_per_match(self): + uri = self._data_uri(b"hello") + text = f"intro ![alt-1]({uri}) middle ![alt-2]({uri}) end" + blocks = extract_data_uri_image_blocks(text, page_number=3) + + assert len(blocks) == 2 + assert all(b.image_bytes == b"hello" for b in blocks) + assert all(b.page_number == 3 for b in blocks) + assert blocks[0].metadata["alt"] == "alt-1" + assert blocks[0].metadata["markdown_ref"] == f"![alt-1]({uri})" + + def test_no_matches_returns_empty(self): + assert extract_data_uri_image_blocks("plain text") == [] + assert extract_data_uri_image_blocks("") == [] + + def test_skips_undecodable(self): + text = "![](data:image/png;base64,!!!not-base64)" + assert extract_data_uri_image_blocks(text) == [] + + +def test_min_image_pixels_constant(): + assert MIN_IMAGE_PIXELS == 784 diff --git a/openrag/core/indexing/test_validators.py b/openrag/core/indexing/test_validators.py new file mode 100644 index 000000000..0ffae81ac --- /dev/null +++ b/openrag/core/indexing/test_validators.py @@ -0,0 +1,69 @@ +"""Unit tests for ``core.indexing.validators``.""" + +from __future__ import annotations + +import pytest + +from ..utils.exceptions import ValidationError +from .validators import parse_metadata, validate_file_format, validate_file_id + + +class TestParseMetadata: + def test_none_returns_empty(self): + assert parse_metadata(None) == {} + + def test_empty_string_returns_empty(self): + assert parse_metadata("") == {} + + def test_dict_passthrough(self): + d = {"a": 1, "b": [1, 2]} + assert parse_metadata(d) is d + + def test_valid_json_string(self): + assert parse_metadata('{"k": "v"}') == {"k": "v"} + + def test_invalid_json_raises_400(self): + with pytest.raises(ValidationError) as exc: + parse_metadata("{not-json") + assert exc.value.status_code == 400 + + def test_non_object_json_raises_400(self): + with pytest.raises(ValidationError) as exc: + parse_metadata('["a", "b"]') + assert exc.value.status_code == 400 + + +class TestValidateFileId: + def test_valid(self): + assert validate_file_id("abc-123") == "abc-123" + + def test_default_forbidden_slash(self): + with pytest.raises(ValidationError) as exc: + validate_file_id("a/b") + assert exc.value.status_code == 400 + + def test_empty_or_whitespace_raises(self): + for bad in ("", " "): + with pytest.raises(ValidationError): + validate_file_id(bad) + + def test_custom_forbidden_chars(self): + with pytest.raises(ValidationError): + validate_file_id("hello?world", forbidden_chars="?") + assert validate_file_id("hello/world", forbidden_chars="?") == "hello/world" + + +class TestValidateFileFormat: + formats = ("pdf", "docx") + mimetypes = ("application/pdf", "application/vnd.openxmlformats-officedocument.wordprocessingml.document") + + def test_extension_match(self): + assert validate_file_format("doc.PDF", self.formats, self.mimetypes) == "pdf" + + def test_mimetype_match_when_no_extension(self): + assert validate_file_format("noext", self.formats, self.mimetypes, mimetype="application/pdf") == "" + + def test_unsupported_raises_415(self): + with pytest.raises(ValidationError) as exc: + validate_file_format("img.exe", self.formats, self.mimetypes, mimetype="application/x-msdownload") + assert exc.value.status_code == 415 diff --git a/openrag/core/indexing/text_preprocessor.py b/openrag/core/indexing/text_preprocessor.py new file mode 100644 index 000000000..8f31137ef --- /dev/null +++ b/openrag/core/indexing/text_preprocessor.py @@ -0,0 +1,16 @@ +"""Text preprocessing utilities for the indexing pipeline. + +Re-exports the canonical implementations from `core.utils.text`. Kept as a +named entry point under `core.indexing` so callers can import preprocessing +helpers alongside parsers, validators, and contextualization without +reaching into the generic utils package. +""" + +from ..utils.text import clean_markdown_table_spacing, decode_bytes, sanitize_extracted_text, sanitize_text + +__all__ = [ + "clean_markdown_table_spacing", + "sanitize_extracted_text", + "sanitize_text", + "decode_bytes", +] diff --git a/openrag/core/indexing/validators.py b/openrag/core/indexing/validators.py new file mode 100644 index 000000000..0dbc33848 --- /dev/null +++ b/openrag/core/indexing/validators.py @@ -0,0 +1,75 @@ +"""Framework-free validators for indexing inputs. + +Pure functions on plain types — no FastAPI, no Hydra. Routers translate +incoming HTTP requests into these inputs and let the global ``OpenRAGError`` +handler convert raised ``ValidationError`` instances into HTTP responses. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable +from typing import Any + +from ..utils.exceptions import ValidationError + +DEFAULT_FORBIDDEN_CHARS_IN_FILE_ID: frozenset[str] = frozenset("/") + + +def parse_metadata(raw: Any | None) -> dict: + """Parse JSON-encoded metadata into a dict. + + Accepts ``None`` / empty string (returns ``{}``), an existing dict + (returned as-is), or a JSON string that decodes to a dict. + """ + if raw is None or raw == "": + return {} + if isinstance(raw, dict): + return raw + try: + decoded = json.loads(raw) + except (json.JSONDecodeError, TypeError) as exc: + raise ValidationError("Invalid JSON in metadata", status_code=400) from exc + if not isinstance(decoded, dict): + raise ValidationError("Metadata must be a JSON object", status_code=400) + return decoded + + +def validate_file_id( + file_id: str, + forbidden_chars: Iterable[str] = DEFAULT_FORBIDDEN_CHARS_IN_FILE_ID, +) -> str: + """Return ``file_id`` if valid, else raise ``ValidationError`` (HTTP 400).""" + forbidden = frozenset(forbidden_chars) + if any(c in file_id for c in forbidden): + raise ValidationError( + f"File ID contains forbidden characters: {', '.join(sorted(forbidden))}", + status_code=400, + ) + if not file_id.strip(): + raise ValidationError("File ID cannot be empty.", status_code=400) + return file_id + + +def validate_file_format( + filename: str, + accepted_formats: Iterable[str], + accepted_mimetypes: Iterable[str], + mimetype: str | None = None, +) -> str: + """Validate the file by extension or mimetype. + + Returns the lowercased file extension (without the leading dot, possibly + empty). Raises ``ValidationError`` (HTTP 415) on rejection. + """ + file_extension = filename.rsplit(".", 1)[-1].lower() if "." in filename else "" + formats = set(accepted_formats) + mimetypes = set(accepted_mimetypes) + if file_extension not in formats and mimetype not in mimetypes: + details = ( + f"Unsupported file format: {file_extension} or file mimetype.\n" + f"Supported formats: {', '.join(sorted(formats))}\n" + f"Supported mimetypes: {', '.join(sorted(mimetypes))}" + ) + raise ValidationError(details, status_code=415) + return file_extension diff --git a/openrag/core/models/document.py b/openrag/core/models/document.py index 3652844fd..9dbfd3a69 100644 --- a/openrag/core/models/document.py +++ b/openrag/core/models/document.py @@ -2,9 +2,16 @@ from __future__ import annotations +import asyncio +import base64 +import os +import tempfile import uuid +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from datetime import UTC, datetime from enum import Enum +from pathlib import Path from typing import Any from pydantic import BaseModel, Field @@ -34,14 +41,56 @@ class TextBlock(BaseModel): class ImageBlock(BaseModel): - """An image extracted from a document.""" - - image_bytes: bytes = Field(exclude=True, repr=False) + """An image extracted from a document. + + Parser→caption contract: + - Parsers emit ``ImageBlock`` with ``caption=None``. A downstream + caption stage fills it in via a VLM. + - When the source text contains a placeholder for the image + (``![alt](data:image/...)``, ``![](pptx-image-3)``, ``![](marker-key-7)``, + …), the parser stores the exact placeholder string in + ``metadata["markdown_ref"]``. The caption stage substitutes the + wrapped caption back into the corresponding ``TextBlock`` via + ``str.replace`` on that ref. + - When there is no in-text placeholder (standalone image uploads, + EML image attachments), ``metadata["markdown_ref"]`` is omitted; + the caption stage produces a free-standing captioned ``TextBlock`` + instead. + + Bytes vs. URL: + - Locally-extracted images set ``image_bytes`` (raw PNG / JPEG bytes) + and leave ``source_url`` as ``None``. + - Remote images parsed from a markdown ``![](http://…)`` ref leave + ``image_bytes`` empty and set ``source_url`` to the URL. A + downstream fetch stage may populate ``image_bytes`` later. + - The :attr:`image_url` property is the unified VLM-friendly form: + a ``data:`` URI built from the bytes when present, otherwise the + ``source_url`` as-is. + """ + + image_bytes: bytes = Field(default=b"", exclude=True, repr=False) + source_url: str | None = None page_number: int | None = None caption: str | None = None mime_type: str = "image/png" metadata: dict[str, Any] = Field(default_factory=dict) + @property + def image_url(self) -> str: + """A VLM-friendly URL for this image. + + - Bytes present → ``data:{mime_type};base64,{...}`` URI. + - Otherwise → ``source_url`` if set, else empty string. + - On any encoding failure → falls back to ``source_url`` (or ""). + """ + if self.image_bytes: + try: + b64 = base64.b64encode(self.image_bytes).decode() + return f"data:{self.mime_type};base64,{b64}" + except Exception: + pass + return self.source_url or "" + class Document(BaseModel): """A document before or during indexing.""" @@ -101,6 +150,66 @@ def to_langchain(self) -> Any: } return LCDocument(page_content=self.text or "", metadata=metadata) + @asynccontextmanager + async def as_temporary_file(self, *, suffix: str | None = None) -> AsyncIterator[Path]: + """Materialize ``raw_bytes`` to a temporary file and yield its ``Path``. + + Parsers wrapping a sync library that requires a path on disk + (Marker, Whisper, MarkItDown, python-pptx, Spire.Doc, …) use this + helper instead of rolling their own ``NamedTemporaryFile`` dance. + The file is removed on context exit even if the body raises. + + ``suffix`` defaults to ``filename``'s extension, falling back to + a content-type-appropriate default. + """ + if self.raw_bytes is None: + raise ValueError("Document.as_temporary_file requires raw_bytes") + + if suffix is None: + suffix = Path(self.filename).suffix or _DEFAULT_TEMPFILE_SUFFIX.get(self.content_type, "") + + raw = self.raw_bytes + + def _write_temp() -> str: + # Close before yielding so sync callers (Marker/Whisper/MarkItDown/ + # python-pptx/Spire.Doc) can reopen the path on Windows, where + # NamedTemporaryFile(delete=True) holds an exclusive handle. + tf = tempfile.NamedTemporaryFile(suffix=suffix, delete=False) + try: + tf.write(raw) + finally: + tf.close() + return tf.name + + path = await asyncio.to_thread(_write_temp) + try: + yield Path(path) + finally: + await asyncio.to_thread(_safe_unlink, path) + + +def _safe_unlink(path: str) -> None: + """``os.unlink`` that swallows missing-file errors (sync callers may have already removed it).""" + try: + os.unlink(path) + except FileNotFoundError: + pass + + +_DEFAULT_TEMPFILE_SUFFIX: dict[DocumentType, str] = { + DocumentType.PDF: ".pdf", + DocumentType.DOCX: ".docx", + DocumentType.PPTX: ".pptx", + DocumentType.DOC: ".doc", + DocumentType.AUDIO: ".wav", + DocumentType.VIDEO: ".mp4", + DocumentType.EML: ".eml", + DocumentType.IMAGE: ".png", + DocumentType.HTML: ".html", + DocumentType.MARKDOWN: ".md", + DocumentType.TEXT: ".txt", +} + class ProcessedDocument(BaseModel): """Document after parsing/extraction — contains text blocks and images.""" diff --git a/openrag/core/utils/conts.py b/openrag/core/utils/conts.py new file mode 100644 index 000000000..9f457374d --- /dev/null +++ b/openrag/core/utils/conts.py @@ -0,0 +1,10 @@ +PARTITION_PREFIX = "openrag-" +LEGACY_PARTITION_PREFIX = "ragondin-" + +FILE_READ_CHUNK_SIZE = 1024 * 1024 # Read file in blocks of 1MB to preserve RAM + + +IMG_WRAPPER_OPEN = "\n\n" +IMG_WRAPPER_CLOSE = "\n\n" + +IMAGE_PLACEHOLDER = f"""{IMG_WRAPPER_OPEN}[Image Placeholder]{IMG_WRAPPER_CLOSE}""" diff --git a/openrag/core/utils/exceptions.py b/openrag/core/utils/exceptions.py index 25df37303..7ccebdb25 100644 --- a/openrag/core/utils/exceptions.py +++ b/openrag/core/utils/exceptions.py @@ -131,10 +131,15 @@ def __init__(self, message: str, **kwargs): class ValidationError(OpenRAGError): - """Input validation or business rule violation. Maps to HTTP 422.""" + """Input validation or business rule violation. Maps to HTTP 422 by default. - def __init__(self, message: str, **kwargs): - super().__init__(message, code="VALIDATION_ERROR", status_code=422, **kwargs) + Accepts a custom ``status_code`` so callers can preserve more specific + semantics (e.g. 400 Bad Request for malformed input, 415 Unsupported + Media Type for rejected file formats). + """ + + def __init__(self, message: str, *, status_code: int = 422, code: str = "VALIDATION_ERROR", **kwargs): + super().__init__(message, code=code, status_code=status_code, **kwargs) # --------------------------------------------------------------------------- diff --git a/openrag/core/utils/text.py b/openrag/core/utils/text.py index b4ddf4540..b992191f6 100644 --- a/openrag/core/utils/text.py +++ b/openrag/core/utils/text.py @@ -9,6 +9,38 @@ import re import unicodedata +DEFAULT_FALLBACK_ENCODING = "utf-8" + + +def decode_bytes(raw: bytes, encoding: str | None = None) -> str: + """Decode ``raw`` to ``str`` with a UTF-8-first detection strategy. + + chardet alone misclassifies short ASCII-heavy UTF-8 as Latin-1, which + produces mojibake on common short inputs. Trying strict UTF-8 first + catches the common case; chardet handles genuinely non-UTF-8 inputs. + Falls back to UTF-8 with ``errors="replace"`` so this never raises. + """ + if encoding: + try: + return raw.decode(encoding, errors="replace") + except LookupError: + # Invalid codec name — fall through to detection. + pass + try: + return raw.decode("utf-8") + except UnicodeDecodeError: + pass + try: + import chardet + except ImportError: + return raw.decode(DEFAULT_FALLBACK_ENCODING, errors="replace") + guess = chardet.detect(raw) + detected = guess.get("encoding") or DEFAULT_FALLBACK_ENCODING + try: + return raw.decode(detected, errors="replace") + except LookupError: + return raw.decode(DEFAULT_FALLBACK_ENCODING, errors="replace") + def sanitize_text( text: str, diff --git a/openrag/routers/utils.py b/openrag/routers/utils.py index bfd8a908c..656251d92 100644 --- a/openrag/routers/utils.py +++ b/openrag/routers/utils.py @@ -1,4 +1,3 @@ -import json import os from pathlib import Path from typing import Any @@ -6,6 +5,7 @@ import consts import openai from config import load_config +from core.indexing import validators as core_validators from fastapi import Depends, Form, HTTPException, Request, UploadFile, status from openai import AsyncOpenAI from utils.dependencies import get_task_state_manager, get_vectordb @@ -272,48 +272,24 @@ async def check_user_file_quota( return user -def is_file_id_valid(file_id: str) -> bool: - return not any(c in file_id for c in FORBIDDEN_CHARS_IN_FILE_ID) - - async def validate_file_id(file_id: str): - if not is_file_id_valid(file_id): - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"File ID contains forbidden characters: {', '.join(FORBIDDEN_CHARS_IN_FILE_ID)}", - ) - if not file_id.strip(): - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="File ID cannot be empty.") - return file_id + return core_validators.validate_file_id(file_id, FORBIDDEN_CHARS_IN_FILE_ID) async def validate_metadata(metadata: Any | None = Form(None)): - try: - processed_metadata = metadata or "{}" - processed_metadata = json.loads(processed_metadata) - return processed_metadata - except json.JSONDecodeError: - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid JSON in metadata") + return core_validators.parse_metadata(metadata) async def validate_file_format( file: UploadFile, metadata: dict = Depends(validate_metadata), ): - file_extension = file.filename.split(".")[-1].lower() if "." in file.filename else "" - mimetype = metadata.get("mimetype", None) - - if file_extension not in ACCEPTED_FILE_FORMATS and mimetype not in DICT_MIMETYPES.keys(): - details = ( - f"Unsupported file format: {file_extension} or file mimetype.\n" - f"Supported formats: {', '.join(ACCEPTED_FILE_FORMATS)}\n" - f"Supported mimetypes: {', '.join(DICT_MIMETYPES.keys())}" - ) - raise HTTPException( - status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, - detail=details, - ) - + core_validators.validate_file_format( + filename=file.filename, + accepted_formats=ACCEPTED_FILE_FORMATS, + accepted_mimetypes=DICT_MIMETYPES.keys(), + mimetype=metadata.get("mimetype"), + ) return file diff --git a/openrag/services/inference/parsers/__init__.py b/openrag/services/inference/parsers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/services/inference/parsers/_base_openai_parser.py b/openrag/services/inference/parsers/_base_openai_parser.py new file mode 100644 index 000000000..3b46320e0 --- /dev/null +++ b/openrag/services/inference/parsers/_base_openai_parser.py @@ -0,0 +1,99 @@ +"""Common scaffolding for OpenAI-VLM-backed PDF parsers. + +Provides reusable helpers — PDF rendering, single-page VLM calls under a +semaphore, JSON-fence stripping, picture-bbox cropping — but takes no +opinion on response shape or block layout. Concrete subclasses +implement ``parse()`` and stitch blocks together however suits the +model they target. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from abc import ABC +from io import BytesIO +from typing import Any + +from core.indexing.image_preprocessor import pil_to_png_bytes +from core.indexing.parsers.document_parser import BaseClientParser +from core.models.document import DocumentType +from core.vlm import VLM + +logger = logging.getLogger(__name__) + + +class BaseOpenAIPdfClient(BaseClientParser, ABC): + """OpenAI-compatible VLM-backed PDF parser scaffolding.""" + + def __init__( + self, + vlm: VLM, + *, + scale: float = 1.0, + concurrency_limit: int = 4, + ) -> None: + self._vlm = vlm + self._scale = scale + self._semaphore = asyncio.Semaphore(max(1, concurrency_limit)) + + def supported_types(self) -> list[str]: + return [DocumentType.PDF.value] + + # ----- helpers ----- + + @staticmethod + def _render_pdf_pages(raw_bytes: bytes, scale: float) -> list[Any]: + """Render every PDF page into a PIL Image. Pure-CPU; runs in a thread.""" + import pypdfium2 as pdfium + + pdf = pdfium.PdfDocument(raw_bytes) + try: + return [page.render(scale=scale).to_pil() for page in pdf] + finally: + pdf.close() + + async def _ocr_one(self, page_img: Any, prompt: str) -> str | None: + """Send one page image through the VLM with ``prompt``; return raw text.""" + async with self._semaphore: + try: + png_bytes = pil_to_png_bytes(page_img) + return await self._vlm.caption_image(png_bytes, prompt=prompt) + except Exception as exc: + logger.warning("OpenAI VLM OCR call failed: %s", exc) + return None + + @staticmethod + def _strip_json_fences(raw: str) -> str: + """Strip ```json ... ``` fences and surrounding whitespace from a VLM response.""" + text = raw.strip() + if text.startswith("```"): + text = text.strip("`") + if text.lower().startswith("json"): + text = text[4:].lstrip() + return text + + @staticmethod + def _load_json(raw: str | None) -> Any | None: + """Decode a JSON payload from a VLM response, tolerating fences and whitespace.""" + if not raw: + return None + text = BaseOpenAIPdfClient._strip_json_fences(raw) + try: + return json.loads(text) + except json.JSONDecodeError as exc: + logger.warning("OCR response was not valid JSON: %s", exc) + return None + + @staticmethod + def _crop_to_png_bytes(page_img: Any, bbox: Any) -> bytes | None: + """Crop a region from a PIL page image and return PNG bytes.""" + try: + cropped = page_img.crop(tuple(bbox)) + buf = BytesIO() + cropped.save(buf, format="PNG") + return buf.getvalue() + except Exception as exc: + logger.warning("Failed to crop bbox %s: %s", bbox, exc) + return None diff --git a/openrag/services/inference/parsers/dotsocr.py b/openrag/services/inference/parsers/dotsocr.py new file mode 100644 index 000000000..d65904540 --- /dev/null +++ b/openrag/services/inference/parsers/dotsocr.py @@ -0,0 +1,143 @@ +"""DotsOCR PDF parser — concrete :class:`BaseOpenAIPdfClient` subclass. + +DotsOCR returns a JSON list of layout elements (``Picture``, ``Table``, +``Text``, ``Title`` …) with bounding boxes and text content, sorted by +reading order. + +Block emission: + +- One :class:`TextBlock` per page (1-indexed ``page_number``), holding + every non-``Picture`` element's text joined in reading order. +- One :class:`ImageBlock` per ``Picture`` element, carrying the cropped + PNG bytes. Captioning is left to a downstream stage — the parser + does **not** call the VLM for captions. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from enum import Enum + +from core.models.document import Document, ImageBlock, ProcessedDocument, TextBlock +from pydantic import BaseModel, RootModel, ValidationError + +from ._base_openai_parser import BaseOpenAIPdfClient + +logger = logging.getLogger(__name__) + + +_DOTSOCR_PROMPT = """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox. + +1. Bbox format: [x1, y1, x2, y2] + +2. Layout Categories: The possible categories are ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title']. + +3. Text Extraction & Formatting Rules: + - Picture: For the 'Picture' category, the text field should be omitted. + - Formula: Format its text as LaTeX. + - Table: Format its text as HTML. + - All Others (Text, Title, etc.): Format their text as Markdown. + +4. Constraints: + - The output text must be the original text from the image, with no translation. + - All layout elements must be sorted according to human reading order. + +5. Final Output: The entire output must be a single JSON object. +""" + + +class DotsOCRCategory(str, Enum): + CAPTION = "Caption" + FOOTNOTE = "Footnote" + FORMULA = "Formula" + LIST_ITEM = "List-item" + PAGE_FOOTER = "Page-footer" + PAGE_HEADER = "Page-header" + PICTURE = "Picture" + SECTION_HEADER = "Section-header" + TABLE = "Table" + TEXT = "Text" + TITLE = "Title" + + +class DotsOCRElement(BaseModel): + """One layout element on a page.""" + + bbox: tuple[float, float, float, float] + category: DotsOCRCategory + text: str = "" + + +class DotsOCRPage(RootModel[list[DotsOCRElement]]): + """One page's DotsOCR output: layout elements in reading order.""" + + def pictures(self) -> list[DotsOCRElement]: + return [e for e in self.root if e.category is DotsOCRCategory.PICTURE] + + def text(self) -> str: + """Join every non-``Picture`` element's text in reading order.""" + return "\n".join( + e.text.strip() for e in self.root if e.category is not DotsOCRCategory.PICTURE and e.text and e.text.strip() + ) + + +class DotsOCRPdfClient(BaseOpenAIPdfClient): + """OpenAI-compatible PDF parser using the DotsOCR layout-aware prompt.""" + + PROMPT: str = _DOTSOCR_PROMPT + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + start = time.time() + try: + page_imgs = await asyncio.to_thread(self._render_pdf_pages, document.raw_bytes, self._scale) + raw_responses = await asyncio.gather(*(self._ocr_one(img, self.PROMPT) for img in page_imgs)) + except Exception: + logger.exception("DotsOCR PDF parse failed (id=%s)", document.id) + raise + + text_blocks: list[TextBlock] = [] + images: list[ImageBlock] = [] + for page_number, (page_img, raw) in enumerate(zip(page_imgs, raw_responses, strict=True), start=1): + page = self._parse_page(raw) + if page is None: + continue + page_text = page.text() + if page_text: + text_blocks.append(TextBlock(text=page_text, page_number=page_number)) + for element in page.pictures(): + png = self._crop_to_png_bytes(page_img, element.bbox) + if png is not None: + images.append(ImageBlock(image_bytes=png, page_number=page_number)) + + logger.info("DotsOCR PDF parsed (id=%s) in %.2fs", document.id, time.time() - start) + + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=images, + metadata=dict(document.metadata), + page_count=len(page_imgs), + ) + + @classmethod + def _parse_page(cls, raw: str | None) -> DotsOCRPage | None: + """Validate one page's raw VLM response into a :class:`DotsOCRPage`.""" + payload = cls._load_json(raw) + if payload is None: + return None + # Tolerate ``{"items": [...]}`` envelope as well as a bare list. + if isinstance(payload, dict) and "items" in payload: + payload = payload["items"] + try: + return DotsOCRPage.model_validate(payload) + except ValidationError as exc: + logger.warning("DotsOCR response did not match expected schema: %s", exc) + return None diff --git a/openrag/services/inference/parsers/openai_audio.py b/openrag/services/inference/parsers/openai_audio.py new file mode 100644 index 000000000..54058d0ad --- /dev/null +++ b/openrag/services/inference/parsers/openai_audio.py @@ -0,0 +1,130 @@ +"""OpenAI-compatible audio transcription client. + +Pipeline: + +1. Materialize ``Document.raw_bytes`` to a temporary file via + :meth:`Document.as_temporary_file`. +2. If the file's suffix is in ``direct_upload_suffixes``, send it to the + transcription endpoint as-is. Otherwise, decode through + ``pydub.AudioSegment`` and re-encode as WAV (libsndfile-friendly). +3. Optionally run a caller-provided language detector against the + prepared file (its result is forwarded to the OpenAI ``language`` + param). The detector is a plain async callable so this client stays + free of Ray / model-loader coupling — the wiring layer can plug in a + Whisper actor or any other implementation. +4. Send the file to ``audio.transcriptions.create`` and emit a single + :class:`TextBlock` with the resulting transcript. + +Adapted from the legacy +``components/indexer/loaders/audio/openai.py`` ``AudioTranscriber``; +the new version drops the in-memory ``components.utils`` semaphore (now +per-instance via ``concurrency_limit``) and the embedded WhisperActor +ref-getter (now an injected callable). +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from collections.abc import Awaitable, Callable, Iterable +from pathlib import Path + +from core.indexing.parsers.document_parser import BaseClientParser +from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock +from openai import AsyncOpenAI +from pydub import AudioSegment + +logger = logging.getLogger(__name__) + + +# Suffixes the transcription backend can ingest as-is, avoiding the ~10x +# size inflation from WAV conversion (Scaleway cap: 100 MB; OpenAI: 25 MB). +_DEFAULT_DIRECT_UPLOAD_SUFFIXES: tuple[str, ...] = (".mp3", ".m4a", ".ogg", ".webm", ".wav") + +LanguageDetector = Callable[[Path], Awaitable[str | None]] + + +class OpenAIAudioClient(BaseClientParser): + """OpenAI-compatible audio transcription client.""" + + def __init__( + self, + *, + base_url: str, + api_key: str, + model: str, + timeout: float = 120.0, + direct_upload_suffixes: Iterable[str] = _DEFAULT_DIRECT_UPLOAD_SUFFIXES, + language_detector: LanguageDetector | None = None, + concurrency_limit: int = 1, + ) -> None: + self._client = AsyncOpenAI(base_url=base_url, api_key=api_key, timeout=timeout) + self._model = model + self._direct_upload_suffixes = {s.lower() for s in direct_upload_suffixes} + self._language_detector = language_detector + self._semaphore = asyncio.Semaphore(max(1, concurrency_limit)) + + def supported_types(self) -> list[str]: + return [DocumentType.AUDIO.value, DocumentType.VIDEO.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + start = time.time() + try: + async with document.as_temporary_file() as input_path: + async with self._semaphore: + upload_path, cleanup = await self._prepare_upload(input_path) + try: + language: str | None = None + if self._language_detector is not None: + try: + language = await self._language_detector(upload_path) + except Exception as exc: + logger.warning("Language detection failed: %s", exc) + text = await self._transcribe(upload_path, language=language) + finally: + if cleanup: + await asyncio.to_thread(upload_path.unlink, True) + except Exception: + logger.exception("OpenAI audio transcription failed (id=%s)", document.id) + raise + + logger.info("OpenAI audio transcribed (id=%s) in %.2fs", document.id, time.time() - start) + + text = text.strip() + text_blocks = [TextBlock(text=text, page_number=1)] if text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + metadata=dict(document.metadata), + page_count=1 if text else 0, + ) + + async def _prepare_upload(self, input_path: Path) -> tuple[Path, bool]: + """Return ``(path_to_upload, needs_cleanup)``. + + Files in :attr:`_direct_upload_suffixes` are sent as-is; others + are decoded by ``pydub`` (ffmpeg) and re-exported as WAV next to + the input — the caller unlinks that temporary on the way out. + """ + if input_path.suffix.lower() in self._direct_upload_suffixes: + return input_path, False + + sound = await asyncio.to_thread(AudioSegment.from_file, input_path) + logger.info("Converting audio to WAV (duration=%.1fs)", len(sound) / 1000) + wav_path = input_path.with_suffix(".wav") + await asyncio.to_thread(sound.export, wav_path, format="wav") + return wav_path, True + + async def _transcribe(self, path: Path, *, language: str | None) -> str: + kwargs: dict[str, object] = {"model": self._model, "file": path} + if language: + kwargs["language"] = language + response = await self._client.audio.transcriptions.create(**kwargs) + return response.text or "" diff --git a/openrag/services/inference/parsers/test_openai_audio.py b/openrag/services/inference/parsers/test_openai_audio.py new file mode 100644 index 000000000..e6c665e2c --- /dev/null +++ b/openrag/services/inference/parsers/test_openai_audio.py @@ -0,0 +1,149 @@ +"""Unit tests for :class:`OpenAIAudioClient`. + +``pydub`` is shimmed at import time via ``sys.modules`` so the test +runs on Python 3.13 (where ``audioop`` was dropped from stdlib and +plain ``import pydub`` fails). The mock is good enough for the control +flow we exercise — neither real audio decoding nor a real OpenAI +client is needed. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import pytest + +# ---- shim pydub before importing openai_audio ------------------------------ + +if "pydub" not in sys.modules: + pydub = types.ModuleType("pydub") + pydub.AudioSegment = MagicMock() # type: ignore[attr-defined] + sys.modules["pydub"] = pydub + +from core.models.document import Document, DocumentType # noqa: E402 + +from .openai_audio import OpenAIAudioClient # noqa: E402 + +# ---- shared fixtures ------------------------------------------------------- + + +@pytest.fixture +def mock_openai_client(): + """Build an ``AsyncOpenAI``-shaped mock with an awaitable ``audio.transcriptions.create``.""" + fake = MagicMock() + fake.audio = MagicMock() + fake.audio.transcriptions = MagicMock() + fake.audio.transcriptions.create = AsyncMock() + return fake + + +def _client(mock_openai_client, **overrides) -> OpenAIAudioClient: + defaults = {"base_url": "http://x", "api_key": "k", "model": "whisper-mock"} + client = OpenAIAudioClient(**{**defaults, **overrides}) + # Constructor stores config only; swap in our mock before any call. + client._client = mock_openai_client + return client + + +def _audio_doc(raw: bytes = b"audio-bytes", filename: str = "x.mp3") -> Document: + return Document(filename=filename, content_type=DocumentType.AUDIO, raw_bytes=raw) + + +# ---- _prepare_upload ------------------------------------------------------- + + +class TestPrepareUpload: + @pytest.mark.asyncio + async def test_direct_upload_skips_conversion(self, mock_openai_client): + client = _client(mock_openai_client) + path = Path("/tmp/audio.mp3") + upload, cleanup = await client._prepare_upload(path) + assert upload == path + assert cleanup is False + + @pytest.mark.asyncio + async def test_unsupported_suffix_falls_back_to_wav(self, mock_openai_client, monkeypatch): + from services.inference.parsers import openai_audio as mod + + sound = MagicMock() + sound.__len__ = MagicMock(return_value=1500) + sound.export = MagicMock() + from_file = MagicMock(return_value=sound) + monkeypatch.setattr(mod.AudioSegment, "from_file", from_file) + + client = _client(mock_openai_client) + path = Path("/tmp/audio.flac") + upload, cleanup = await client._prepare_upload(path) + + assert upload == path.with_suffix(".wav") + assert cleanup is True + from_file.assert_called_once_with(path) + sound.export.assert_called_once() + assert sound.export.call_args.kwargs == {"format": "wav"} + + +# ---- parse() --------------------------------------------------------------- + + +class TestParse: + @pytest.mark.asyncio + async def test_empty_raw_bytes_returns_empty(self, mock_openai_client): + result = await _client(mock_openai_client).parse(_audio_doc(raw=b"")) + assert result.text_blocks == [] and result.page_count == 0 + + @pytest.mark.asyncio + async def test_returns_text_block_on_success(self, mock_openai_client): + mock_openai_client.audio.transcriptions.create.return_value = MagicMock(text=" hello world ") + result = await _client(mock_openai_client).parse(_audio_doc()) + + assert len(result.text_blocks) == 1 + assert result.text_blocks[0].text == "hello world" + assert result.text_blocks[0].page_number == 1 + assert result.page_count == 1 + mock_openai_client.audio.transcriptions.create.assert_awaited_once() + kwargs = mock_openai_client.audio.transcriptions.create.await_args.kwargs + assert kwargs["model"] == "whisper-mock" + assert "language" not in kwargs + + @pytest.mark.asyncio + async def test_empty_transcript_yields_no_text_block(self, mock_openai_client): + mock_openai_client.audio.transcriptions.create.return_value = MagicMock(text=" ") + result = await _client(mock_openai_client).parse(_audio_doc()) + assert result.text_blocks == [] and result.page_count == 0 + + @pytest.mark.asyncio + async def test_language_detector_result_forwarded(self, mock_openai_client): + mock_openai_client.audio.transcriptions.create.return_value = MagicMock(text="bonjour") + detector = AsyncMock(return_value="fr") + result = await _client(mock_openai_client, language_detector=detector).parse(_audio_doc()) + + detector.assert_awaited_once() + kwargs = mock_openai_client.audio.transcriptions.create.await_args.kwargs + assert kwargs["language"] == "fr" + assert result.text_blocks[0].text == "bonjour" + + @pytest.mark.asyncio + async def test_language_detector_failure_is_swallowed(self, mock_openai_client): + mock_openai_client.audio.transcriptions.create.return_value = MagicMock(text="ok") + detector = AsyncMock(side_effect=RuntimeError("detector down")) + result = await _client(mock_openai_client, language_detector=detector).parse(_audio_doc()) + + # Transcription proceeds without ``language`` and the call still succeeds. + kwargs = mock_openai_client.audio.transcriptions.create.await_args.kwargs + assert "language" not in kwargs + assert result.text_blocks[0].text == "ok" + + @pytest.mark.asyncio + async def test_transcribe_exception_propagates(self, mock_openai_client): + mock_openai_client.audio.transcriptions.create.side_effect = RuntimeError("api down") + with pytest.raises(RuntimeError, match="api down"): + await _client(mock_openai_client).parse(_audio_doc()) + + +def test_supported_types(mock_openai_client): + types_ = _client(mock_openai_client).supported_types() + assert DocumentType.AUDIO.value in types_ + assert DocumentType.VIDEO.value in types_ diff --git a/openrag/services/workers/parsers/__init__.py b/openrag/services/workers/parsers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openrag/services/workers/parsers/marker_workers.py b/openrag/services/workers/parsers/marker_workers.py new file mode 100644 index 000000000..61f22bb76 --- /dev/null +++ b/openrag/services/workers/parsers/marker_workers.py @@ -0,0 +1,455 @@ +import asyncio +import gc +import re +import time + +import pypdfium2 +import ray +import torch +from config import load_config +from core.indexing.image_preprocessor import pil_to_png_bytes +from core.indexing.parsers.document_parser import BasePooledParser +from core.models.document import ( + Document, + DocumentType, + ImageBlock, + ProcessedDocument, + TextBlock, +) +from marker.converters.pdf import PdfConverter +from utils.logger import get_logger + +from ..ray_utils import with_retry, with_timeout + +logger = get_logger() +config = load_config() + +if torch.cuda.is_available(): + MARKER_NUM_GPUS = config.loader.marker_num_gpus +else: # On CPU + MARKER_NUM_GPUS = 0 + + +@ray.remote(num_gpus=MARKER_NUM_GPUS, max_restarts=5) +class MarkerWorker: + def __init__(self): + import os + + from config import load_config + from utils.logger import get_logger + + self.logger = get_logger() + self.config = load_config() + self.page_sep = "[PAGE_SEP]" + + self._workers = self.config.loader.marker_max_processes + + self.converter_config = { + "output_format": "markdown", + "paginate_output": True, + "page_separator": self.page_sep, + "pdftext_workers": self.config.loader.marker_pdftext_workers, + "disable_multiprocessing": False, + } + os.environ["RAY_ADDRESS"] = "auto" + + self.executor = None + self.init_resources() + + def init_resources(self): + from marker.models import create_model_dict + + self.model_dict = create_model_dict() + for v in self.model_dict.values(): + if hasattr(v.model, "share_memory"): + v.model.share_memory() + + self.setup_mp() + + def setup_mp(self): + """Initialize ProcessPoolExecutor for PDF processing. + + We use ProcessPoolExecutor instead of multiprocessing.Pool because: + - Ray actors run as daemon processes + - Pool workers are daemonic by default and cannot spawn children + - The pdftext library (used by Marker) internally spawns processes + - ProcessPoolExecutor workers are non-daemon, allowing nested process creation + """ + from concurrent.futures import ProcessPoolExecutor + + import torch.multiprocessing as mp + + if self.executor: + self.logger.warning("Resetting ProcessPoolExecutor") + self.executor.shutdown(wait=False, cancel_futures=True) + self.executor = None + + # Ensure spawn method for CUDA compatibility + try: + if mp.get_start_method(allow_none=True) != "spawn": + mp.set_start_method("spawn", force=True) + except RuntimeError: + self.logger.warning("Process start method already set, using existing method") + + self.logger.info(f"Initializing MarkerWorker with {self._workers} workers") + self.executor = ProcessPoolExecutor( + max_workers=self._workers, + initializer=self._worker_init, + initargs=(self.model_dict,), + mp_context=mp.get_context("spawn"), + max_tasks_per_child=self.config.loader.marker_max_tasks_per_child, + ) + self.logger.info("MarkerWorker initialized with ProcessPoolExecutor") + + @staticmethod + def _worker_init(model_dict): + global worker_model_dict + worker_model_dict = model_dict + logger.debug("Worker initialized with model dictionary") + + @staticmethod + def _process_pdf(file_path, config): + global worker_model_dict + + page_range = config.get("page_range") + if page_range is not None: + label = f"[p{page_range[0]}-{page_range[-1]}]" + else: + label = "(all pages)" + + try: + logger.debug("Processing PDF", path=file_path, label=label) + converter = PdfConverter( + artifact_dict=worker_model_dict, + config=config, + ) + render = converter(file_path) + return render + except Exception as e: + logger.exception("Error processing PDF", path=file_path, label=label, error=str(e)) + raise + finally: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + async def process_pdf(self, file_path: str, page_range: list[int] | None = None): + from concurrent.futures import TimeoutError as FuturesTimeoutError + + converter_config = self.converter_config.copy() + if page_range is not None: + converter_config["page_range"] = page_range + + loop = asyncio.get_event_loop() + timeout = self.config.loader.marker_timeout + + def run_with_timeout(): + future = self.executor.submit(self._process_pdf, file_path, converter_config) + try: + result = future.result(timeout=timeout) + return result + except FuturesTimeoutError: + self.logger.exception("MarkerWorker child process timed out", path=file_path) + raise + except Exception: + self.logger.exception("Error processing with MarkerWorker", path=file_path) + raise + + result = await loop.run_in_executor(None, run_with_timeout) + return result.markdown, result.images + + def is_pool_broken(self): + # ProcessPoolExecutor auto-replaces dead/finished workers on next + # submit(), so counting live processes is unreliable and unnecessary. + # Only a None or shut-down executor requires reinitialization. + return self.executor is None or bool(getattr(self.executor, "_broken", False)) + + def __del__(self): + """Clean up ProcessPoolExecutor on actor destruction""" + if self.executor: + try: + self.executor.shutdown(wait=False, cancel_futures=True) + except Exception: + pass # Best effort cleanup + + +@ray.remote(max_restarts=5) +class MarkerPool: + def __init__(self): + from config import load_config + from utils.logger import get_logger + + self.logger = get_logger() + self.config = load_config() + self.max_processes = self.config.loader.marker_max_processes + self.pool_size = self.config.loader.marker_pool_size + self.actors = [MarkerWorker.remote() for _ in range(self.pool_size)] + self._queue: asyncio.Queue[ray.actor.ActorHandle] = asyncio.Queue() + + for _ in range(self.max_processes): + for actor in self.actors: + self._queue.put_nowait(actor) + + self.logger.info( + f"Marker pool: {self.pool_size} actors × {self.max_processes} slots = " + f"{self.pool_size * self.max_processes} PDF concurrency" + ) + + @staticmethod + def _get_page_count(file_path: str) -> int: + pdf = pypdfium2.PdfDocument(file_path) + try: + return len(pdf) + finally: + pdf.close() + + @staticmethod + def _create_chunks(page_count: int, chunk_size: int) -> list[tuple[list[int], str]]: + if page_count <= chunk_size: + return [(list(range(page_count)), f"({page_count}p)")] + chunks = [] + for start in range(0, page_count, chunk_size): + end = min(start + chunk_size, page_count) + page_range = list(range(start, end)) + label = f"[p{start}-{end - 1}]" + chunks.append((page_range, label)) + return chunks + + @with_timeout( + seconds=config.loader.marker_timeout, + description="MarkerWorker pool health check", + ) + async def _check_pool_broken(self, worker): + return worker.is_pool_broken.remote() + + @with_timeout( + seconds=config.loader.marker_timeout, + description="MarkerWorker pool reset", + ) + async def _reset_worker_pool(self, worker): + return worker.setup_mp.remote() + + async def ensure_worker_pool_healthy(self, worker): + if await self._check_pool_broken(worker): + self.logger.warning("Worker ProcessPoolExecutor is broken. Reinitializing pool...") + await self._reset_worker_pool(worker) + + @with_timeout( + seconds=config.loader.marker_timeout, + description="MarkerPool PDF {label} ({file_path})", + ) + async def _run_chunk(self, worker, file_path: str, page_range: list[int] | None, label: str): + return worker.process_pdf.remote(file_path, page_range=page_range) + + @with_retry( + max_retries=config.loader.marker_max_task_retry, + base_delay=config.loader.marker_retry_base_delay, + description="MarkerPool PDF {label} ({file_path})", + ) + async def _process_chunk(self, file_path: str, page_range: list[int] | None, label: str): + """Acquire a worker slot, process a PDF chunk, and release the slot. + + A fresh worker is acquired per attempt so a flaky worker can be + sidestepped and ``ensure_worker_pool_healthy`` re-runs each time. + Retries are handled by ``@with_retry``. + """ + worker = await self._queue.get() + try: + self.logger.info(f"MarkerWorker allocated for {label}") + await self.ensure_worker_pool_healthy(worker) + return await self._run_chunk(worker, file_path, page_range, label) + finally: + await self._queue.put(worker) + self.logger.debug(f"MarkerWorker returned to pool for {label}") + + async def process_pdf(self, file_path: str): + chunk_size = self.config.loader.marker_chunk_size + + if chunk_size <= 0: + return await self._process_chunk(file_path, page_range=None, label="(all pages)") + + page_count = self._get_page_count(file_path) + chunks = self._create_chunks(page_count, chunk_size) + + if len(chunks) == 1: + page_range, label = chunks[0] + return await self._process_chunk(file_path, page_range=None, label=label) + + self.logger.info( + f"Splitting {page_count}-page PDF into {len(chunks)} chunks of ~{chunk_size} pages for parallel processing" + ) + + tasks = [asyncio.create_task(self._process_chunk(file_path, page_range, label)) for page_range, label in chunks] + try: + results = await asyncio.gather(*tasks) + except Exception: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + + # Reassemble: concatenate markdown in order, merge image dicts + all_markdown = [] + all_images = {} + for markdown, images in results: + all_markdown.append(markdown) + all_images.update(images) + + combined_markdown = "\n\n".join(all_markdown) + return combined_markdown, all_images + + +_MARKER_KEY_PAGE_RE = re.compile(r"_page_(\d+)_") + + +def _marker_key_to_page(key: str) -> int | None: + """Extract the 1-indexed page number from a Marker image key. + + Marker emits keys like ``_page_0_Picture_1.jpeg`` (0-indexed). We + return ``N + 1`` so callers see 1-indexed pages aligned with the + ``[PAGE_N]`` markers produced by the post-processing step. + Returns ``None`` if the key doesn't match the expected pattern. + """ + match = _MARKER_KEY_PAGE_RE.search(key) + if match is None: + return None + try: + return int(match.group(1)) + 1 + except (TypeError, ValueError): + return None + + +class MarkerLoader(BasePooledParser): + """Public ``BasePooledParser`` facade for the Marker Ray pool. + + Holds a handle to the named ``MarkerPool`` Ray actor and dispatches + each ``parse()`` call to it. Marker requires a file path on disk, so + ``Document.raw_bytes`` is materialized to a temporary file (via + ``Document.as_temporary_file``) before handoff. + + Output: one ``TextBlock`` per page (1-indexed ``page_number``) plus + one ``ImageBlock`` per Marker image. Each ``ImageBlock`` carries the + ``![](key)`` markdown ref in ``metadata['markdown_ref']`` so a + downstream caption stage can substitute the wrapped caption back + into the markdown by string match. Captioning is not done here — + see :class:`ImageBlock` for the parser→caption contract. + """ + + PAGE_SEP = "[PAGE_SEP]" + _PAGE_MARKER_RE = re.compile(r"\{(\d+)\}" + re.escape(PAGE_SEP)) + + def __init__(self) -> None: + self.worker = ray.get_actor("MarkerPool", namespace="openrag") + + def supported_types(self) -> list[str]: + return [DocumentType.PDF.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as path: + markdown, images = await self._dispatch(str(path)) + + pages = self._split_pages(markdown) + image_blocks = self._build_image_blocks(images) + text_blocks = [TextBlock(text=text, page_number=page) for page, text in pages] + + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + images=image_blocks, + metadata=dict(document.metadata), + page_count=pages[-1][0] if pages else 0, + ) + + # ----- helpers ----- + + @with_timeout( + seconds=config.loader.marker_timeout, + description="MarkerLoader PDF loading ({file_path})", + ) + async def _convert_pdf(self, file_path: str): + return self.worker.process_pdf.remote(file_path) + + async def _dispatch(self, file_path: str) -> tuple[str, dict]: + start = time.time() + try: + markdown, images = await self._convert_pdf(file_path) + if not markdown: + raise RuntimeError(f"Conversion failed for {file_path}") + duration = time.time() - start + logger.info(f"Processed {file_path} in {duration:.2f}s") + return markdown, images or {} + except Exception: + logger.exception("Error in MarkerLoader.parse", path=file_path) + raise + + @staticmethod + def _build_image_blocks(images: dict) -> list[ImageBlock]: + """Convert Marker's ``{key: PIL_image}`` dict into ``ImageBlock``s. + + Each block records the ``![](key)`` markdown ref in + ``metadata['markdown_ref']`` so a downstream caption stage can + substitute the wrapped caption back into the text. The page + number is parsed from Marker's key format + (``_page_{N}_Picture_{i}.{ext}``) and stored 1-indexed to match + the ``[PAGE_N]`` markers in the post-processed markdown. + """ + blocks: list[ImageBlock] = [] + for key, pil_image in images.items(): + try: + png_bytes = pil_to_png_bytes(pil_image) + except Exception as exc: + logger.warning("Failed to encode Marker image %s: %s", key, exc) + continue + blocks.append( + ImageBlock( + image_bytes=png_bytes, + page_number=_marker_key_to_page(str(key)), + mime_type="image/png", + metadata={"markdown_ref": f"![]({key})", "marker_key": str(key)}, + ) + ) + return blocks + + @classmethod + def _split_pages(cls, markdown: str) -> list[tuple[int, str]]: + """Clean Marker output and split it into ``[(page_number, text), …]``. + + Marker emits ``{1}[PAGE_SEP]{2}[PAGE_SEP]…``. We + drop the leading ``[PAGE_SEP]`` segment (Marker prefixes one), + strip ``
``, then split on each ``{N}[PAGE_SEP]`` marker — + the captured ``N`` is the 1-indexed page that just ended. + + Blank pages are preserved (text=``""``) so ``page_number`` and + ``page_count`` reflect the source document, not just the + non-empty subset. Trailing text after the last marker (rare) is + assigned to ``last_page + 1``. Markdown with no markers collapses + to a single page-1 entry. + """ + if markdown is None: + return [] + if cls.PAGE_SEP in markdown: + markdown = markdown.split(cls.PAGE_SEP, 1)[1] + markdown = markdown.replace("
", "") + + pairs: list[tuple[int, str]] = [] + cursor = 0 + last_page = 0 + for match in cls._PAGE_MARKER_RE.finditer(markdown): + page = int(match.group(1)) + text = markdown[cursor : match.start()].strip() + pairs.append((page, text)) + cursor = match.end() + last_page = page + tail = markdown[cursor:].strip() + if tail: + pairs.append((last_page + 1, tail)) + elif not pairs and markdown.strip(): + pairs.append((1, markdown.strip())) + return pairs diff --git a/openrag/services/workers/parsers/whisper_workers.py b/openrag/services/workers/parsers/whisper_workers.py new file mode 100644 index 000000000..c0aeb5075 --- /dev/null +++ b/openrag/services/workers/parsers/whisper_workers.py @@ -0,0 +1,152 @@ +import asyncio +from pathlib import Path + +import ray +import torch +from config import load_config +from core.indexing.parsers.document_parser import BasePooledParser +from core.models.document import ( + Document, + DocumentType, + ProcessedDocument, + TextBlock, +) +from faster_whisper import WhisperModel +from utils.logger import get_logger + +from ..ray_utils import call_ray_actor_with_timeout, with_retry, with_timeout + +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: + """Ray-actor pool of ``WhisperActor``s. Internal — the public + ``BasePooledParser`` face is ``LocalWhisperLoader``. + """ + + 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 + + @with_timeout( + seconds=config.loader.local_whisper.whisper_timeout, + description="WhisperPool transcribe ({path})", + ) + async def _transcribe_chunk(self, idx: int, path): + return self.workers[idx].transcribe.remote(path) + + @with_retry( + max_retries=config.loader.local_whisper.whisper_max_task_retry, + base_delay=config.loader.local_whisper.whisper_retry_base_delay, + description="WhisperPool transcribe ({path})", + ) + async def transcribe(self, path): + idx = min(range(len(self._pending)), key=lambda j: self._pending[j]) + self._pending[idx] += 1 + try: + return await self._transcribe_chunk(idx, path) + finally: + self._pending[idx] -= 1 + + +class LocalWhisperLoader(BasePooledParser): + """Public ``BasePooledParser`` facade for the local-Whisper Ray pool. + + Holds a handle to the named ``WhisperPool`` Ray actor and dispatches + each ``parse()`` call to it. Whisper requires a file path on disk, + so ``Document.raw_bytes`` is written to a NamedTemporaryFile before + handoff. + """ + + def __init__(self): + self.whisper_actor: WhisperPool = ray.get_actor("WhisperPool", namespace="openrag") + + def supported_types(self) -> list[str]: + return [DocumentType.AUDIO.value, DocumentType.VIDEO.value] + + async def parse(self, document: Document) -> ProcessedDocument: + if not document.raw_bytes: + return ProcessedDocument( + document_id=document.id, + metadata=dict(document.metadata), + ) + + async with document.as_temporary_file() as path: + try: + text = await call_ray_actor_with_timeout( + self.whisper_actor.transcribe.remote(str(path)), + timeout=config.loader.local_whisper.whisper_timeout, + task_description=f"WhisperPool transcribe ({path})", + ) + except Exception as e: + logger.error("Error transcribing audio", error=str(e)) + raise + + text_blocks = [TextBlock(text=text, page_number=1)] if text else [] + return ProcessedDocument( + document_id=document.id, + text_blocks=text_blocks, + metadata=dict(document.metadata), + page_count=1 if text else 0, + ) diff --git a/openrag/services/workers/ray_utils.py b/openrag/services/workers/ray_utils.py new file mode 100644 index 000000000..1a63b7614 --- /dev/null +++ b/openrag/services/workers/ray_utils.py @@ -0,0 +1,222 @@ +"""Ray-actor concurrency helpers. + +Two pairs of utilities, each in function and decorator form: + +- timeout: ``call_ray_actor_with_timeout`` / ``@with_timeout``. Awaits a + ``ray.ObjectRef`` with proper cancel-on-timeout semantics. +- retry: ``retry_with_backoff`` / ``@with_retry``. Exponential backoff + + jitter; ``CancelledError`` is never retried. + +Use the decorator form when params are static (or pulled from a +module-level config); use the function form when params are dynamic per +call. Cancellation paths are translated into a predictable shape: + +- caller-side timeout → ``ray.cancel(future)`` then re-raise ``TimeoutError`` +- caller-side ``asyncio.CancelledError`` → ``ray.cancel(future)`` then re-raise +- worker-side ``TaskCancelledError`` → re-raise as-is +- worker-side ``RayTaskError`` → re-raise as ``RuntimeError`` (cause preserved) +""" + +from __future__ import annotations + +import asyncio +import functools +import inspect +import random +from collections.abc import Callable +from typing import Any + +import ray +from ray.exceptions import RayTaskError, TaskCancelledError +from utils.logger import get_logger + +logger = get_logger() + + +def _resolve_description(template: str, fn: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]) -> str: + """Format ``template`` with the wrapped call's bound arguments. + + A description like ``"PDF parse ({file_path})"`` gets ``{file_path}`` + substituted with the value passed to ``fn`` for that parameter. If + ``template`` contains no ``{`` it is returned unchanged — no inspect + cost in the hot path for plain-string descriptions. + + ``KeyError`` from a missing placeholder is caught and the raw + template is returned, so a typo in a placeholder name degrades to a + log-line oddity rather than a runtime crash on the wrapped call. + """ + if "{" not in template: + return template + try: + bound = inspect.signature(fn).bind(*args, **kwargs).arguments + return template.format(**bound) + except (KeyError, TypeError) as exc: + logger.warning(f"description template missing placeholder for {exc}") + return template + + +# --------------------------------------------------------------------------- +# Timeout +# --------------------------------------------------------------------------- + + +async def call_ray_actor_with_timeout( + future: ray.ObjectRef, + timeout: float, + task_description: str = "Ray task", +) -> Any: + """Await a Ray ``ObjectRef`` with a timeout, propagating cancellation. + + Raises: + TimeoutError: If the task exceeds ``timeout``. + asyncio.CancelledError: If the calling coroutine is cancelled. + TaskCancelledError: If the Ray task was cancelled by the worker. + RuntimeError: If the Ray task failed (original exception chained). + """ + try: + result = await asyncio.wait_for(asyncio.gather(future), timeout=timeout) + return result[0] + + except TimeoutError: + logger.warning(f"{task_description} timed out, cancelling Ray task") + ray.cancel(future, recursive=True) + raise + + except asyncio.CancelledError: + logger.warning(f"{task_description} cancelled, cancelling Ray task") + ray.cancel(future, recursive=True) + raise + + except TaskCancelledError: + logger.warning(f"{task_description} Ray task was cancelled") + raise + + except RayTaskError as e: + raise RuntimeError(f"{task_description} failed") from e + + +def with_timeout( + *, + seconds: float, + description: str = "Ray task", +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Decorator: wrap an async function returning a ``ray.ObjectRef``. + + The wrapped function is called normally; its return value (an + ``ObjectRef``) is then awaited via ``call_ray_actor_with_timeout``. + + ``description`` may embed any of the wrapped function's parameter + names as ``str.format``-style placeholders; they are substituted + with the per-call argument values for log lines. + + Example:: + + @with_timeout( + seconds=30.0, + description="caption_image ({path})", + ) + async def caption(self, path): + return self.actor.caption.remote(path) + """ + + def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: + @functools.wraps(fn) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + desc = _resolve_description(description, fn, args, kwargs) + future = fn(*args, **kwargs) + if asyncio.iscoroutine(future): + future = await future + return await call_ray_actor_with_timeout(future, seconds, desc) + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------- +# Retry +# --------------------------------------------------------------------------- + + +async def retry_with_backoff( + attempt_fn: Callable[[int], Any], + max_retries: int, + base_delay: float, + task_description: str = "task", + jitter: bool = True, +) -> Any: + """Run ``attempt_fn(attempt_index)`` with exponential backoff. + + Backoff is ``base_delay * 2**attempt`` seconds, plus uniform jitter + in ``[0, base_delay)`` when ``jitter=True``. ``attempt_fn`` is an + async callable; it owns acquire/release of any per-attempt resources + so a flaky resource can be sidestepped on retry. + """ + last_exc: Exception | None = None + for attempt in range(max_retries + 1): + try: + return await attempt_fn(attempt) + except (asyncio.CancelledError, TaskCancelledError): + raise + except Exception as e: + last_exc = e + if attempt >= max_retries: + logger.error(f"{task_description} failed after {attempt + 1} attempts: {e}") + raise + delay = base_delay * (2**attempt) + if jitter: + delay += random.uniform(0, base_delay) + logger.warning( + f"{task_description} failed (attempt {attempt + 1}/{max_retries + 1}): {e}. Retrying in {delay:.1f}s..." + ) + await asyncio.sleep(delay) + + raise last_exc # unreachable + + +def with_retry( + *, + max_retries: int, + base_delay: float, + description: str = "task", + jitter: bool = True, +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Decorator: retry an async function with exponential backoff + jitter. + + Each invocation counts as one attempt. ``CancelledError`` is never + retried. + + ``description`` may embed any of the wrapped function's parameter + names as ``str.format``-style placeholders; they are substituted + with the per-call argument values for log lines. + + Example:: + + @with_retry( + max_retries=3, + base_delay=0.5, + description="transcribe ({path})", + ) + async def transcribe(self, path): + return await self.actor.transcribe.remote(path) + """ + + def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: + @functools.wraps(fn) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + desc = _resolve_description(description, fn, args, kwargs) + + async def attempt(_i: int) -> Any: + return await fn(*args, **kwargs) + + return await retry_with_backoff( + attempt, + max_retries=max_retries, + base_delay=base_delay, + task_description=desc, + jitter=jitter, + ) + + return wrapper + + return decorator diff --git a/openrag/utils/exceptions/__init__.py b/openrag/utils/exceptions/__init__.py index 857408bdd..7ed1e2a97 100644 --- a/openrag/utils/exceptions/__init__.py +++ b/openrag/utils/exceptions/__init__.py @@ -1,3 +1,3 @@ # Re-export from canonical location for backward compatibility. -# New code should import from openrag.core.utils.exceptions directly. -from openrag.core.utils.exceptions import * # noqa: F401,F403 +# New code should import from `core.utils.exceptions` directly. +from core.utils.exceptions import * # noqa: F401,F403 diff --git a/openrag/utils/exceptions/base.py b/openrag/utils/exceptions/base.py index 927bb617f..c0fed1294 100644 --- a/openrag/utils/exceptions/base.py +++ b/openrag/utils/exceptions/base.py @@ -1,6 +1,6 @@ # Re-export from canonical location for backward compatibility. -# New code should import from openrag.core.utils.exceptions directly. -from openrag.core.utils.exceptions import ( # noqa: F401 +# New code should import from `core.utils.exceptions` directly. +from core.utils.exceptions import ( # noqa: F401 EmbeddingError, OpenRAGError, VDBError, diff --git a/openrag/utils/exceptions/embeddings.py b/openrag/utils/exceptions/embeddings.py index ae5fee12c..84be79cef 100644 --- a/openrag/utils/exceptions/embeddings.py +++ b/openrag/utils/exceptions/embeddings.py @@ -1,6 +1,6 @@ # Re-export from canonical location for backward compatibility. -# New code should import from openrag.core.utils.exceptions directly. -from openrag.core.utils.exceptions import ( # noqa: F401 +# New code should import from `core.utils.exceptions` directly. +from core.utils.exceptions import ( # noqa: F401 EmbeddingAPIError, EmbeddingResponseError, UnexpectedEmbeddingError, diff --git a/openrag/utils/exceptions/vectordb.py b/openrag/utils/exceptions/vectordb.py index 2ec5c1a84..54dd12df1 100644 --- a/openrag/utils/exceptions/vectordb.py +++ b/openrag/utils/exceptions/vectordb.py @@ -1,6 +1,6 @@ # Re-export from canonical location for backward compatibility. -# New code should import from openrag.core.utils.exceptions directly. -from openrag.core.utils.exceptions import ( # noqa: F401 +# New code should import from `core.utils.exceptions` directly. +from core.utils.exceptions import ( # noqa: F401 UnexpectedVDBError, VDBConnectionError, VDBCreateOrLoadCollectionError,