diff --git a/README.md b/README.md index 4335720a8..97875e949 100644 --- a/README.md +++ b/README.md @@ -148,7 +148,7 @@ All supported file format parsers are pre-configured. For PDF processing, **[Mar
For more PDF options -For CPU-only deployments or lightweight testing scenarios, you can consider switching to **`PyMuPDF4LLMLoader`** or **`PyMuPDFLoader`**. To change the loader, set the **`PDFLoader`** variable like this `PDFLoader=PyMuPDF4LLMLoader`. +For CPU-only deployments or lightweight testing scenarios, you can consider switching to **`PyMuPDFLoader`**. To change the loader, set the **`PDFLoader`** variable like this `PDFLoader=PyMuPDFLoader`. > ⚠️ **Important**: These alternative loaders have limitations - they cannot process non-searchable (image-based) PDFs and do not extract or handle embedded images.
diff --git a/docs/content/docs/documentation/env_vars.md b/docs/content/docs/documentation/env_vars.md index 398c89644..b0ce83a76 100644 --- a/docs/content/docs/documentation/env_vars.md +++ b/docs/content/docs/documentation/env_vars.md @@ -22,10 +22,10 @@ Openrag loads all files into a pivot markdown file format before proceeding to c | `IMAGE_CAPTIONING_URL` | `bool` | `true` | If `true`, HTTP/HTTPS image URLs in markdown files are fetched and described by the VLM. | | `SAVE_MARKDOWN` | `bool` | `false` | If `true`, the pivot-format markdown produced during parsing is saved. Useful for debugging and verifying the correctness of the generated markdown. | |`SAVE_UPLOADED_FILES`|`bool`|`false`| When `true`, uploaded files are stored on disk. You must enable this option if you want Chainlit to show sources while chatting.| -| `PDFLoader` | `str` | `MarkerLoader` | Specifies the PDF parsing engine to use. Available options: `PyMuPDFLoader`, `PyMuPDF4LLMLoader`, `MarkerLoader` and `DotsOCRLoader`.| +| `PDFLoader` | `str` | `MarkerLoader` | Specifies the PDF parsing engine to use. Available options: `PyMuPDFLoader`, `MarkerLoader` and `DotsOCRLoader`.| :::caution -`PyMuPDFLoader` and `PyMuPDF4LLMLoader` are lightweight pdf loaders that cannot process non-searchable (image-based) PDFs and do not extract or handle embedded images. +`PyMuPDFLoader` is a lightweight pdf loader that cannot process non-searchable (image-based) PDFs and does not extract or handle embedded images. ::: #### PDF Loader diff --git a/docs/content/docs/getting_started/quickstart.mdx b/docs/content/docs/getting_started/quickstart.mdx index efd57a483..9952043bc 100644 --- a/docs/content/docs/getting_started/quickstart.mdx +++ b/docs/content/docs/getting_started/quickstart.mdx @@ -11,7 +11,7 @@ OpenRAG is an open source Retrieval Augmented Generation (RAG) solution. This gu ### Prerequisites - [Docker](https://www.docker.com/get-started) and **Docker Compose** - Your hardware should meet these specifications: - - **CPU deployment**: Minimum **13 GiB** RAM for light PDF parsers (**`PyMuPDF4LLMLoader`, `PyMuPDFLoader`**), or **23 GiB** RAM for heavier parsers like **`MarkerLoader`** (refer to [this section](/openrag/getting_started/quickstart/#3-file-parser-configuration) for details) + - **CPU deployment**: Minimum **13 GiB** RAM for light PDF parsers (**`PyMuPDFLoader`**), or **23 GiB** RAM for heavier parsers like **`MarkerLoader`** (refer to [this section](/openrag/getting_started/quickstart/#3-file-parser-configuration) for details) - **GPU deployment**: **16 GB** GPU memory recommended (for systems with separate CPU and GPU memory) ### Installation and Configuration @@ -37,7 +37,7 @@ Here is a brief overview of key environment variables to configure: All supported file format parsers are pre-configured. For PDF processing, **[MarkerLoader](https://github.com/datalab-to/marker)** serves as the default parser, offering comprehensive support for OCR-scanned documents, complex layouts, tables, and embedded images. MarkerLoader operates efficiently on both GPU and CPU environments. :::note -For **`CPU-only deployments`** or lightweight testing scenarios, you can consider switching to **`PyMuPDF4LLMLoader`** or **`PyMuPDFLoader`**. To change the loader, set the **`PDFLoader`** variable like this `PDFLoader=PyMuPDF4LLMLoader`. +For **`CPU-only deployments`** or lightweight testing scenarios, you can consider switching to **`PyMuPDFLoader`**. To change the loader, set the **`PDFLoader`** variable like this `PDFLoader=PyMuPDFLoader`. :::caution[Important] These alternative loaders have limitations - they cannot process non-searchable (image-based) PDFs and do not extract or handle embedded images. ::: diff --git a/openrag/api/main.py b/openrag/api/main.py index e03279adb..2da2858e3 100644 --- a/openrag/api/main.py +++ b/openrag/api/main.py @@ -155,7 +155,7 @@ async def lifespan(app: FastAPI): # ``ensure_worker_bootstrap`` imports ``services.workers.bootstrap`` # for its side effect: creating the long-lived detached worker - # actors (TaskStateManager, DocSerializer, MarkerPool, semaphores). + # actors (TaskStateManager, MarkerPool, semaphores). # The indirection through :mod:`di.workers` keeps API code free of # direct ``services.workers`` imports. logger.info("Startup: initializing worker bootstrap") diff --git a/openrag/api/routers/admin/tools.py b/openrag/api/routers/admin/tools.py index ede394a6b..257123120 100644 --- a/openrag/api/routers/admin/tools.py +++ b/openrag/api/routers/admin/tools.py @@ -1,8 +1,8 @@ """Tools routes — thin HTTP layer over :class:`ConversionService`. -Phase 8E: the ``extractText`` serialization moved to -``services.orchestrators.conversion_service.ConversionService`` (the Ray -``DocSerializer`` actor now sits behind the ``FileSerializer`` port). +The ``extractText`` serialization lives in +``services.orchestrators.conversion_service.ConversionService``, which calls +the ``FileSerializer`` port (an in-process parser-dispatcher serializer). This module keeps HTTP transport only: the saved-file IO + cleanup, tool validation/dispatch, and the 4xx/5xx error mapping whose exact ``{"detail": ...}`` body the legacy endpoint returned via diff --git a/openrag/core/indexing/serializer.py b/openrag/core/indexing/serializer.py index a2b8da941..4046a06de 100644 --- a/openrag/core/indexing/serializer.py +++ b/openrag/core/indexing/serializer.py @@ -1,15 +1,15 @@ -"""Transitional port for the document-serialization operation. +"""Port for the document-serialization operation (file → raw text). -``ConversionService`` (Phase 8E) exposes the ``extractText`` tool — -serialize an uploaded file to raw text. The work runs in the -``DocSerializer`` Ray actor; defining it on a dedicated port keeps the -orchestrator Ray-free (8H: no Ray import / remote call under -``services/orchestrators/``). A small shim in ``services/storage/`` -adapts the actor to this interface during the shim period; Phase 9 -swaps it for a direct serializer call and deletes the shim. +``ConversionService`` exposes the ``extractText`` tool / ``/extract`` route — +serialize an uploaded file to raw text. Defining the operation on a dedicated +port keeps the orchestrator decoupled from the parser/Ray infrastructure +(no Ray import / remote call under ``services/orchestrators/``); the +composition root injects the concrete implementation +(``services/workers/parsers/file_serializer.py::ParserFileSerializer``, which +runs the parser dispatcher in-process). -No Ray / LangChain types leak across this boundary — the serialized -document is returned as its plain text content. +No Ray / LangChain types leak across this boundary — the serialized document +is returned as its plain text content. """ from __future__ import annotations diff --git a/openrag/core/models/document.py b/openrag/core/models/document.py index 459a8e637..f0d0e21bb 100644 --- a/openrag/core/models/document.py +++ b/openrag/core/models/document.py @@ -118,6 +118,10 @@ def detect_content_type(filename: str) -> DocumentType: "png": DocumentType.IMAGE, "jpg": DocumentType.IMAGE, "jpeg": DocumentType.IMAGE, + "svg": DocumentType.IMAGE, + "gif": DocumentType.IMAGE, + "webp": DocumentType.IMAGE, + "bmp": DocumentType.IMAGE, "mp3": DocumentType.AUDIO, "wav": DocumentType.AUDIO, "flac": DocumentType.AUDIO, diff --git a/openrag/di/container.py b/openrag/di/container.py index b391a98d3..7fd99080c 100644 --- a/openrag/di/container.py +++ b/openrag/di/container.py @@ -589,17 +589,18 @@ def job_service(self) -> JobService: def conversion_service(self) -> ConversionService: """ConversionService — lazily built, cached for the container's lifetime. - The serializer is the Ray-backed ``SerializerRayShim`` during the - Phase-8 shim period (Ray cleanup is Phase 9); the DocSerializer - actor is resolved lazily per call inside the shim. + The serializer is the in-process ``ParserFileSerializer`` — it runs the + parser dispatcher directly (GPU backends still dispatch to their pool + actors) and implements the ``FileSerializer`` port, so the orchestrator + stays decoupled from the parser/Ray infrastructure. """ if self._conversion_service is None: from services.orchestrators.conversion_service import ConversionService - from services.workers.parsers.doc_serializer_adapter import from_ray_namespace + from services.workers.parsers.file_serializer import build_file_serializer settings = self._require_settings() self._conversion_service = ConversionService( - serializer=from_ray_namespace(), + serializer=build_file_serializer(), vector_store=self.vector_store, collection=settings.vectordb.collection_name, ) diff --git a/openrag/services/orchestrators/conversion_service.py b/openrag/services/orchestrators/conversion_service.py index a77631f41..68684e0d5 100644 --- a/openrag/services/orchestrators/conversion_service.py +++ b/openrag/services/orchestrators/conversion_service.py @@ -4,9 +4,9 @@ tool) and ``routers/extract.py`` (chunk-by-id lookup). Both were thin wrappers; this service keeps them Ray-free: -- serialization runs in the ``DocSerializer`` Ray actor, reached through - the :class:`~core.indexing.serializer.FileSerializer` port (the - container injects the ``SerializerRayShim`` during the shim period); +- serialization goes through the :class:`~core.indexing.serializer.FileSerializer` + port (the container injects the in-process ``ParserFileSerializer``, which runs + the parser dispatcher directly — GPU backends still dispatch to their pools); - chunk lookup goes through the clean :class:`VectorStore` port (``query_chunks_by_filter`` on the Milvus ``_id``), mirroring how PartitionService reads chunks — no LangChain ``Document`` leaks out. diff --git a/openrag/services/workers/bootstrap.py b/openrag/services/workers/bootstrap.py index 89a81ef33..a6e7278bf 100644 --- a/openrag/services/workers/bootstrap.py +++ b/openrag/services/workers/bootstrap.py @@ -4,7 +4,6 @@ detached actors the request path looks up by name: * TaskStateManager — shared task-state actor -* DocSerializer — loader dispatcher * MarkerPool / DoclingPool / WhisperPool / WhisperActor — GPU parsers * llmSemaphore / vlmSemaphore / audioSemaphore — cluster-wide rate limiters @@ -65,12 +64,6 @@ def get_task_state_manager(): return get_or_create_actor("TaskStateManager", TaskStateManager, lifetime="detached") -def get_serializer(): - from services.workers.parsers.doc_serializer import DocSerializer - - return get_or_create_actor("DocSerializer", DocSerializer, lifetime="detached") - - def get_marker_pool(): from services.workers.parsers.docling_workers import DoclingPool from services.workers.parsers.marker_workers import MarkerPool @@ -78,7 +71,7 @@ def get_marker_pool(): config = _require_settings() pdf_loader = config.loader.file_loaders.pdf match pdf_loader: - case "DoclingLoader2": + case "DoclingLoader": return get_or_create_actor("DoclingPool", DoclingPool, lifetime="detached") case "MarkerLoader": return get_or_create_actor("MarkerPool", MarkerPool, lifetime="detached") @@ -146,4 +139,3 @@ def initialize_worker_bootstrap(settings: "Settings") -> None: init_audio_actor() get_marker_pool() get_task_state_manager() - get_serializer() diff --git a/openrag/services/workers/indexer_actor.py b/openrag/services/workers/indexer_actor.py index 736ce9669..b93d8faa0 100644 --- a/openrag/services/workers/indexer_actor.py +++ b/openrag/services/workers/indexer_actor.py @@ -1,12 +1,12 @@ from __future__ import annotations +import asyncio import traceback from datetime import datetime from pathlib import Path from typing import Any from core.models.document import Document -from services.workers.parsers.doc_serializer_bridge import INDEXATION_CONFIG_METADATA_KEY from services.workers.pipeline_builder import IndexingPipeline @@ -61,13 +61,13 @@ async def process_file( """ await self._tsm.set_state.remote(task_id, "SERIALIZING") try: - document = _load_document(path, metadata, partition, indexation_config=indexation_config) + document = await _load_document(path, metadata, partition) # One indexation timestamp for this file, shared by the Milvus chunks # (via the store stage) and the Postgres catalog row, so they agree. row: dict[str, Any] = { "document": document, "partition": partition, - "filename": Path(path).name, + "filename": document.filename, "language": metadata.get("language", "en"), "replace": replace, "user": user, @@ -178,24 +178,47 @@ async def _replace_topic_tags_if_needed( ) -def _load_document( +async def _load_document( path: str, metadata: dict[str, Any], partition: str, - *, - indexation_config: dict[str, Any] | None = None, ) -> Document: p = Path(path) - document_metadata = dict(metadata) - if indexation_config is not None: - document_metadata[INDEXATION_CONFIG_METADATA_KEY] = dict(indexation_config) + file_id = metadata.get("file_id") + if not file_id: + # file_id is a required route path param, force-set by + # IndexingService._build_metadata. Missing here means a broken upstream + # contract — fail loudly rather than silently persisting chunks under a + # non-queryable id (e.g. the temp upload's basename). + raise ValueError("_load_document requires metadata['file_id']") + # ``Document.id`` is the file's identity, not a random uuid: parsers set + # ``ProcessedDocument.document_id = document.id`` and the chunker uses that as + # ``Chunk.document_id`` / ``file_id``. If this defaulted to uuid4, chunks would + # persist under an id the ``/partition/{partition}/file/{file_id}`` lookup + # never queries by (zero chunks found). + # + # Per-partition indexation_config reaches the pipeline via ``row["indexation_config"]`` + # (see IndexerWorker.process_file); it is intentionally not stamped into the + # document metadata so it never leaks into chunk metadata. + filename = _display_filename(path, metadata) + raw_bytes = await asyncio.to_thread(p.read_bytes) return Document( - filename=metadata.get("file_id") or p.name, - raw_bytes=p.read_bytes(), - content_type=Document.detect_content_type(p.name), + id=file_id, + filename=filename, + raw_bytes=raw_bytes, + content_type=Document.detect_content_type(filename), partition=partition, - metadata=document_metadata, + metadata=dict(metadata), ) +def _display_filename(path: str, metadata: dict[str, Any]) -> str: + """Return the user-facing filename while falling back to the stored path.""" + + filename = metadata.get("original_filename") or metadata.get("filename") + if filename: + return str(filename) + return Path(path).name + + __all__ = ["IndexerWorker"] diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py index badd52dad..486fafb01 100644 --- a/openrag/services/workers/indexer_pool.py +++ b/openrag/services/workers/indexer_pool.py @@ -22,12 +22,13 @@ def __init__(self) -> None: from core.embeddings import embedder_registry from services.storage.milvus_store import MilvusVectorStore from services.storage.postgres_store import PostgresStore - from services.workers.parsers.doc_serializer_bridge import DocSerializerBridgeParser + from services.workers.parsers.parser_dispatcher import build_caption_vlm, build_parser_dispatcher from services.workers.pipeline_builder import build_indexing_pipeline cfg = load_config() - parser = DocSerializerBridgeParser(config=cfg) + parser = build_parser_dispatcher(cfg) + vlm = build_caption_vlm(cfg) chunker = _build_chunker(cfg) embedder_factory = _build_embedder_factory(cfg) contextualizer_factory = _build_contextualizer_factory(cfg) @@ -51,6 +52,8 @@ def __init__(self) -> None: chunker=chunker, embedder=embedder, vector_store=self._vector_store, + vlm=vlm, + image_captioning=cfg.loader.image_captioning, chunker_factory=_build_chunker_from_config, embedder_factory=embedder_factory, contextualizer_factory=contextualizer_factory, diff --git a/openrag/services/workers/parsers/doc_serializer.py b/openrag/services/workers/parsers/doc_serializer.py deleted file mode 100644 index 5fe3e51d9..000000000 --- a/openrag/services/workers/parsers/doc_serializer.py +++ /dev/null @@ -1,83 +0,0 @@ -"""DocSerializer Ray actor. - -Moved from ``components/indexer/loaders/serializer.py``; the old module -re-exports this class for backward compatibility. -""" - -from __future__ import annotations - -import gc -from pathlib import Path - -import ray -import torch -from langchain_core.documents.base import Document -from services.workers.parsers.legacy_loaders import get_loader_classes - - -@ray.remote(max_restarts=5) -class DocSerializer: - def __init__(self, data_dir=None, **kwargs) -> None: - from core.config import load_config - from core.utils.logging import get_logger - - self.logger = get_logger() - self.config = load_config() - self.data_dir = data_dir - self.kwargs = kwargs - self.kwargs["config"] = self.config - self.save_markdown = self.config.loader.save_markdown - - self.loader_classes = get_loader_classes(config=self.config) - self.logger.info("DocSerializer initialized.") - - async def serialize_document( - self, - task_id: str, - path: str | Path, - metadata: dict | None = None, - ) -> Document: - metadata = metadata or {} - log = self.logger.bind( - file_id=metadata.get("file_id"), - partition=metadata.get("partition"), - task_id=task_id, - ) - task_state_manager = ray.get_actor("TaskStateManager", namespace="openrag") - await task_state_manager.set_state.remote(task_id, "SERIALIZING") - - log.info("Starting document serialization") - - p = Path(path) - file_ext = p.suffix.lower() - mimetype = metadata.get("mimetype", None) - mimetypes = self.config.loader.mimetypes.to_dict() - if mimetype is None: - loader_cls = self.loader_classes.get(file_ext) - else: - loader_cls = self.loader_classes.get(mimetypes.get(mimetype)) - - if loader_cls is None: - log.warning(f"No loader available for {p.name}") - raise ValueError(f"No loader available for file type {file_ext}.") - - log.debug(f"Loading document: {p.name} with loader {loader_cls.__name__}") - loader = loader_cls(**self.kwargs) - - try: - doc: Document = await loader.aload_document( - file_path=path, metadata=metadata, save_markdown=self.save_markdown - ) - del loader - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() - log.info("Document serialized successfully") - return doc - except Exception as e: - log.exception("Failed to serialize document", error=str(e)) - raise - - -__all__ = ["DocSerializer"] diff --git a/openrag/services/workers/parsers/doc_serializer_adapter.py b/openrag/services/workers/parsers/doc_serializer_adapter.py deleted file mode 100644 index 525c62dc7..000000000 --- a/openrag/services/workers/parsers/doc_serializer_adapter.py +++ /dev/null @@ -1,37 +0,0 @@ -"""FileSerializer adapter over the DocSerializer Ray actor. - -Replaces ``services/storage/serializer_ray_shim.py`` (Phase 9E). The adapter -lives in the workers layer because it wraps a worker Ray actor; the storage -layer no longer references Ray directly. -""" - -from __future__ import annotations - -from core.indexing.serializer import FileSerializer - -_FALLBACK_TASK_ID = "tools-extract" - - -class DocSerializerAdapter(FileSerializer): - """Implements FileSerializer by delegating to the DocSerializer Ray actor.""" - - async def serialize(self, path: str, metadata: dict) -> str: - import ray - from core.config import load_config - from services.workers.ray_utils import call_ray_actor_with_timeout - - cfg = load_config() - timeout = cfg.ray.indexer.serialize_timeout - task_id = ray.get_runtime_context().get_task_id() or _FALLBACK_TASK_ID - serializer = ray.get_actor("DocSerializer", namespace="openrag") - doc = await call_ray_actor_with_timeout( - future=serializer.serialize_document.remote(task_id, path, metadata=metadata or {}), - timeout=timeout, - task_description=f"Serialization task {task_id}", - ) - return doc.page_content - - -def from_ray_namespace() -> DocSerializerAdapter: - """Build the adapter. Convenience for the composition root.""" - return DocSerializerAdapter() diff --git a/openrag/services/workers/parsers/doc_serializer_bridge.py b/openrag/services/workers/parsers/doc_serializer_bridge.py deleted file mode 100644 index cfe3727e4..000000000 --- a/openrag/services/workers/parsers/doc_serializer_bridge.py +++ /dev/null @@ -1,100 +0,0 @@ -from __future__ import annotations - -import tempfile -from pathlib import Path -from typing import Any - -from core.indexing.parsers.document_parser import DocumentParser -from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock - -INDEXATION_CONFIG_METADATA_KEY = "_openrag_indexation_config" - - -class DocSerializerBridgeParser(DocumentParser): - """Transitional parser backed by the legacy loader registry.""" - - def __init__(self, config: Any) -> None: - from services.workers.parsers.legacy_loaders import get_loader_classes - - self._config = config - self._loader_classes = get_loader_classes(config=config) - self._save_markdown = getattr(config.loader, "save_markdown", False) - - def supported_types(self) -> list[str]: - return [doc_type.value for doc_type in DocumentType] - - async def parse(self, document: Document) -> ProcessedDocument: - suffix = _suffix_from_document(document) - with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as handle: - handle.write(document.raw_bytes or b"") - temp_path = handle.name - try: - return await self._load_via_legacy(temp_path, document) - finally: - Path(temp_path).unlink(missing_ok=True) - - async def _load_via_legacy(self, path: str, document: Document) -> ProcessedDocument: - metadata = dict(document.metadata or {}) - indexation_config = metadata.pop(INDEXATION_CONFIG_METADATA_KEY, None) - loader_cls = self._loader_for(path, metadata) - if loader_cls is None: - raise ValueError(f"No loader registered for file extension {Path(path).suffix.lower()!r}") - - loader = loader_cls(config=_legacy_loader_config(self._config, indexation_config)) - lang_doc = await loader.aload_document( - file_path=path, - metadata=metadata, - save_markdown=self._save_markdown, - ) - - return ProcessedDocument( - document_id=document.filename or "unknown", - text_blocks=[TextBlock(text=lang_doc.page_content or "")], - metadata=lang_doc.metadata or {}, - ) - - def _loader_for(self, path: str, metadata: dict[str, Any]) -> Any | None: - mimetype = metadata.get("mimetype") - if mimetype: - try: - from services.workers.parsers.doc_serializer import DICT_MIMETYPES - - loader_cls = self._loader_classes.get(DICT_MIMETYPES.get(mimetype)) - except Exception: - loader_cls = None - if loader_cls is not None: - return loader_cls - - return self._loader_classes.get(Path(path).suffix.lower()) - - -def _suffix_from_document(document: Document) -> str: - source = (document.metadata or {}).get("source") - if source: - suffix = Path(str(source)).suffix - if suffix: - return suffix - if document.filename: - suffix = Path(document.filename).suffix - if suffix: - return suffix - return f".{document.content_type.value}" if document.content_type else "" - - -def _legacy_loader_config(config: Any, indexation_config: Any) -> Any: - """Apply per-file indexation overrides to legacy loader config.""" - if not isinstance(indexation_config, dict): - return config - if indexation_config.get("enable_image_captioning", True): - return config - - loader = config.loader.model_copy( - update={ - "image_captioning": False, - "image_captioning_url": False, - } - ) - return config.model_copy(update={"loader": loader}) - - -__all__ = ["DocSerializerBridgeParser"] diff --git a/openrag/services/workers/parsers/file_serializer.py b/openrag/services/workers/parsers/file_serializer.py new file mode 100644 index 000000000..20ea61048 --- /dev/null +++ b/openrag/services/workers/parsers/file_serializer.py @@ -0,0 +1,70 @@ +"""In-process ``FileSerializer`` over the new parser stack. + +``ConversionService`` (the ``extractText`` tool + ``/extract`` route) needs +file → text. This implements the :class:`~core.indexing.serializer.FileSerializer` +port by running :class:`ParserDispatcher` directly: parse the file into a +``ProcessedDocument``, caption images when a VLM is configured, and join the +text blocks. + +No Ray-actor hop — the GPU backends (marker/docling/whisper) still dispatch to +their pool actors from inside the dispatcher, so only the lightweight dispatch +and CPU parsers run in-process. Replaces the former ``DocSerializer`` Ray actor ++ ``DocSerializerAdapter`` (the Phase-9 shim that indexing no longer shares). +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +from core.indexing.serializer import FileSerializer + + +class ParserFileSerializer(FileSerializer): + """Implements ``FileSerializer`` by running the parser dispatcher in-process.""" + + def __init__(self) -> None: + from core.config import load_config + from services.workers.parsers.parser_dispatcher import build_caption_vlm, build_parser_dispatcher + + config = load_config() + self._dispatcher = build_parser_dispatcher(config) + self._vlm = build_caption_vlm(config) + # Global gate for captioning images embedded in other documents. + # Standalone image files are always captioned (see ``serialize``). + self._image_captioning = bool(config.loader.image_captioning) + + async def serialize(self, path: str, metadata: dict) -> str: + from core.models.document import Document, DocumentType + from services.workers.stages.caption import _caption_document + + metadata = dict(metadata or {}) + p = Path(path) + name_for_type = metadata.get("filename") or p.name + raw_bytes = await asyncio.to_thread(p.read_bytes) + + document = Document( + filename=name_for_type, + raw_bytes=raw_bytes, + content_type=Document.detect_content_type(name_for_type), + metadata=metadata, + ) + + processed = await self._dispatcher.parse(document) + should_caption = ( + self._vlm is not None + and processed.images + and (document.content_type is DocumentType.IMAGE or self._image_captioning) + ) + if should_caption: + processed = await _caption_document(processed, self._vlm, None) + + return "\n\n".join(block.text for block in processed.text_blocks if block.text) + + +def build_file_serializer() -> ParserFileSerializer: + """Build the in-process file serializer. Convenience for the composition root.""" + return ParserFileSerializer() + + +__all__ = ["ParserFileSerializer", "build_file_serializer"] diff --git a/openrag/services/workers/parsers/legacy_loaders/CustomDocLoader.py b/openrag/services/workers/parsers/legacy_loaders/CustomDocLoader.py deleted file mode 100644 index 3238d2c26..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/CustomDocLoader.py +++ /dev/null @@ -1,39 +0,0 @@ -from pathlib import Path - -from langchain_community.document_loaders import ( - UnstructuredODTLoader, - UnstructuredWordDocumentLoader, -) -from langchain_core.documents.base import Document - -from .base import BaseLoader - - -class CustomDocLoader(BaseLoader): - doc_loaders = { - ".docx": UnstructuredWordDocumentLoader, - ".doc": UnstructuredWordDocumentLoader, - ".odt": UnstructuredODTLoader, - } - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - async def aload_document(self, file_path, metadata: dict = None): - path = Path(file_path) - cls_loader = CustomDocLoader.doc_loaders.get(path.suffix, None) - - if cls_loader is None: - raise ValueError(f"This loader only supports {CustomDocLoader.doc_loaders.keys()} format") - - loader = cls_loader( - file_path=str(file_path), - mode="single", - ) - pages = await loader.aload() - - s = "" - for page_num, p in enumerate(pages, start=1): - s += p.page_content.strip() + f"\n[PAGE_{page_num}]\n" - - return Document(page_content=s, metadata=metadata) diff --git a/openrag/services/workers/parsers/legacy_loaders/CustomHTMLLoader.py b/openrag/services/workers/parsers/legacy_loaders/CustomHTMLLoader.py deleted file mode 100644 index 680b376a3..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/CustomHTMLLoader.py +++ /dev/null @@ -1,22 +0,0 @@ -from pathlib import Path - -from langchain_community.document_loaders import UnstructuredHTMLLoader -from langchain_core.documents.base import Document - -from .base import BaseLoader - - -class CustomHTMLLoader(BaseLoader): - def __init__(self, **kwargs): - super().__init__(**kwargs) - - async def aload_document(self, file_path, metadata: dict = None): - path = Path(file_path) - loader = UnstructuredHTMLLoader(file_path=str(path), autodetect_encoding=True) - doc = await loader.aload() - - s = "" - for page_num, segment in enumerate(doc, start=1): - s += segment.page_content.strip() + f"\n[PAGE_{page_num}]\n" - - return Document(page_content=s, metadata=metadata) diff --git a/openrag/services/workers/parsers/legacy_loaders/__init__.py b/openrag/services/workers/parsers/legacy_loaders/__init__.py deleted file mode 100644 index b30b10575..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/__init__.py +++ /dev/null @@ -1,64 +0,0 @@ -""" -Loader registry and initialization module. - -This module handles the dynamic loading and registration of all document loaders. -""" - -import importlib -import pkgutil -from pathlib import Path - -from core.utils.logging import get_logger - -from .base import BaseLoader - -logger = get_logger() - - -def get_loader_classes(config) -> dict[str, type[BaseLoader]]: - # 1. Discover all subclasses - root_pkg = "services.workers.parsers.legacy_loaders" - root_path = Path(__file__).parent - - discovered: dict[str, type[BaseLoader]] = {} - - for finder, module_name, is_pkg in pkgutil.walk_packages(path=[str(root_path)], prefix=f"{root_pkg}."): - try: - module = importlib.import_module(module_name) - except ImportError as e: - logger.warning(f"Could not import module {module_name}: {e}") - continue - - for attr in vars(module).values(): - if isinstance(attr, type) and issubclass(attr, BaseLoader) and attr is not BaseLoader: - discovered[attr.__name__] = attr - - # logger.debug(f"Discovered loaders: {discovered}") - - # 2. Read your config map of extensions → class names - loader_classes: dict[str, type[BaseLoader]] = {} - file_loaders = config.loader.file_loaders.model_dump() - - for ext, cls_name in file_loaders.items(): - cls = discovered.get(cls_name) - if cls is None: - logger.error(f"Configured loader '{cls_name}' for '.{ext}' not found") - continue - loader_classes[f".{ext}"] = cls - logger.debug(f"Registered {cls_name} for .{ext}") - - logger.debug(f"Final loader map: {loader_classes.keys()}") - return loader_classes - - -def get_supported_extensions(loader_classes: dict[str, type[BaseLoader]]) -> set[str]: - """ - Get the set of supported file extensions from the loaded classes. - - Args: - loader_classes: Dictionary mapping file extensions to loader classes - - Returns: - Set of supported file extensions - """ - return set(loader_classes.keys()) diff --git a/openrag/services/workers/parsers/legacy_loaders/audio/__init__.py b/openrag/services/workers/parsers/legacy_loaders/audio/__init__.py deleted file mode 100644 index 07ca00839..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/audio/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .local_whisper import * -from .openai import * diff --git a/openrag/services/workers/parsers/legacy_loaders/audio/local_whisper.py b/openrag/services/workers/parsers/legacy_loaders/audio/local_whisper.py deleted file mode 100644 index 51f9c8bc0..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/audio/local_whisper.py +++ /dev/null @@ -1,67 +0,0 @@ -""" -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 -(``services.workers.parsers.legacy_loaders.audio.local_whisper.WhisperActor`` is -still used by the OpenAI audio loader for language detection, and by -``services/workers/bootstrap.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 - -from core.indexing.parsers.audio.local_whisper import LocalWhisperParser -from core.models.document import Document as CoreDocument -from core.utils.logging import get_logger -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 ..base import BaseLoader - -logger = get_logger() - - -class LocalWhisperLoader(BaseLoader): - """Adapter shim — delegates to ``LocalWhisperParser`` via the services-side pool.""" - - def __init__(self, **kwargs): - super().__init__(**kwargs) - 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: - processed = await self._parser.parse(core_doc) - except Exception as 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/services/workers/parsers/legacy_loaders/audio/openai.py b/openrag/services/workers/parsers/legacy_loaders/audio/openai.py deleted file mode 100644 index 4e8f976f0..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/audio/openai.py +++ /dev/null @@ -1,73 +0,0 @@ -""" -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 - -from core.indexing.parsers.audio.client_based import ClientAudioParser -from core.models.document import Document as CoreDocument -from core.models.document import DocumentType -from core.utils.logging import get_logger -from langchain_core.documents.base import Document -from services.inference.parsers.openai_audio import OpenAIAudioClient -from services.workers.parsers.whisper_workers import detect_language_via_actor - -from ..base import BaseLoader - -logger = get_logger() - - -async def _whisper_language_detector(file_path: Path) -> str | None: - """Detect language via the singleton ``WhisperActor`` (worker-side helper).""" - return await detect_language_via_actor(file_path) - - -class OpenAIAudioLoader(BaseLoader): - """Adapter shim — delegates to ``OpenAIAudioClient`` via ``ClientAudioParser``.""" - - def __init__(self, **kwargs): - super().__init__(**kwargs) - 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: - processed = await self._parser.parse(core_doc) - except Exception: - logger.exception("Error in OpenAIAudioLoader", path=str(file_path)) - raise - content = "\n\n".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/services/workers/parsers/legacy_loaders/base.py b/openrag/services/workers/parsers/legacy_loaders/base.py deleted file mode 100644 index e7321567a..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/base.py +++ /dev/null @@ -1,285 +0,0 @@ -import asyncio -import base64 -import re -from abc import ABC, abstractmethod -from pathlib import Path - -from core.config import 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 core.prompts import load_template_by_key -from core.utils.external_errors import is_external_resource_error -from core.utils.logging import get_logger -from langchain_core.messages import HumanMessage -from langchain_openai import ChatOpenAI -from openai import BadRequestError -from PIL import Image -from services.inference.runtime import get_vlm_semaphore -from tqdm.asyncio import tqdm - -logger = get_logger() - - -class BaseLoader(ABC): - # 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]" - self.config = kwargs.get("config") or load_config() - settings: dict = self.config.vlm.model_dump() - model_settings = { - "temperature": 0.2, - "max_retries": 3, - "timeout": 60, - # "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, - } - settings.update(model_settings) - - self.image_captioning = self.config.loader.image_captioning - self.image_captioning_url = self.config.loader.image_captioning_url - self.image_describer_prompt = "" - if self.image_captioning or self.image_captioning_url: - self.image_describer_prompt = load_template_by_key( - self.config.paths.prompts_dir, - self.config.prompts, - "image_describer", - ) - - self.vlm_endpoint = ChatOpenAI(**settings).with_retry(stop_after_attempt=2) - - @abstractmethod - async def aload_document( - self, - file_path: str | Path, - metadata: dict | None = None, - save_markdown: bool = False, - ): - pass - - def save_content(self, text_content: str, path: str): - path = re.sub(r"\..*", ".md", path) - with open(path, "w", encoding="utf-8") as f: - f.write(text_content) - logger.debug(f"Document saved to {path}") - - def _pil_image_to_base64(self, image: Image.Image) -> str: - """Convert PIL Image to base64 string.""" - try: - 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(png_bytes).decode() - - def _is_http_url(self, data: str) -> bool: - """Check if string is an HTTP/HTTPS URL.""" - return isinstance(data, str) and data.startswith(("http://", "https://")) - - def _is_data_uri(self, data: str) -> bool: - """Check if string is a data URI.""" - return isinstance(data, str) and data.startswith("data:image/") - - async def get_image_description( - self, - image_data: Image.Image | str, - ) -> str: - """ - Creates a description for an image using the LLM model. - - Args: - image_data: Can be one of: - - PIL.Image object - - str: HTTP/HTTPS URL - - str: data URI (data:image/...;base64,...) - - Returns: - str: Description of the image wrapped in XML tags - """ - # Early exit for small PIL images (below VLM min_pixels threshold) - if isinstance(image_data, Image.Image): - w, h = image_data.size - if w * h < self.MIN_IMAGE_PIXELS: - logger.debug("Skipping image below minimum size", size=f"{w}x{h}") - return "\n\nImage too small for captioning\n\n" - - async with get_vlm_semaphore(): - try: - # Determine the type of image data and create appropriate message content - if isinstance(image_data, Image.Image): - # Convert PIL Image to base64 - img_b64 = self._pil_image_to_base64(image_data) - if not img_b64: - return "\n\nFailed to convert image\n\n" - image_url = f"data:image/png;base64,{img_b64}" - - elif self._is_http_url(image_data): - # Handle HTTP/HTTPS URL - image_url = image_data - logger.debug(f"Processing HTTP URL: {image_data}") - - elif self._is_data_uri(image_data): - # Handle data URI - use as-is - image_url = image_data - logger.debug(f"Processing data URI: {image_data[:50]}...") - - else: - # Handle raw base64 string (assume it's base64 encoded image) - if isinstance(image_data, str): - try: - # Try to decode to verify it's valid base64 - base64.b64decode(image_data) - image_url = f"data:image/png;base64,{image_data}" - logger.debug("Processing raw base64 string") - except Exception: - logger.error(f"Invalid image data type or format: {type(image_data)}") - return """\n\nInvalid image data format\n\n""" - else: - logger.error(f"Unsupported image data type: {type(image_data)}") - return """\n\nUnsupported image data type\n\n""" - - # Create message for LLM - message = HumanMessage( - content=[ - { - "type": "image_url", - "image_url": {"url": image_url}, - }, - {"type": "text", "text": self.image_describer_prompt}, - ] - ) - - # Get description from LLM - response = await self.vlm_endpoint.ainvoke([message]) - image_description = response.content - - except BadRequestError as e: - # VLM returned 400 - log as warning without stack trace - logger.warning("VLM rejected image captioning request", error=str(e)[:300]) - image_description = "" - - except Exception as e: - is_external, status_code, url = is_external_resource_error(e) - if is_external: - # Log external resource errors as warnings, not exceptions - # These are expected when VLM cannot fetch external URLs - log_msg = "Failed to fetch external image resource" - log_extra = {"error": str(e)[:200]} - if status_code: - log_extra["http_status"] = status_code - if url: - log_extra["url"] = url - elif self._is_http_url(str(image_data)): - log_extra["url"] = str(image_data) - logger.warning(log_msg, **log_extra) - else: - logger.exception("Error while generating image description", error=str(e)) - image_description = "" - - return f"""\n\n{image_description}\n\n""" - - async def caption_images(self, images: list[Image.Image], desc: str = "Captioning images") -> list[str]: - """Generate captions for a list of PIL images concurrently. - - Args: - images: List of PIL Image objects to caption. - desc: Description for the progress bar. - - Returns: - List of captions in the same order as input images. - """ - if not images: - return [] - - tasks = [self.get_image_description(image_data=img) for img in images] - try: - results = await tqdm.gather(*tasks, desc=desc) - except asyncio.CancelledError: - for task in tasks: - if hasattr(task, "cancel"): - task.cancel() - raise - return results - - async def replace_markdown_images_with_captions( - self, - content: str, - caption_http_urls: bool | None = None, - caption_data_uris: bool = True, - desc: str = "Captioning images", - ) -> str: - """Find markdown image references and replace with VLM-generated captions. - - Args: - content: Markdown content containing ![alt](url) image references. - caption_http_urls: Whether to caption HTTP/HTTPS URLs. - If None, uses config value `loader.image_captioning_url`. - caption_data_uris: Whether to caption data URI images. - desc: Description for the progress bar. - - Returns: - Content with image references replaced by captions. - """ - if not self.image_captioning: - return content - - # Determine URL captioning setting - if caption_http_urls is None: - caption_http_urls = self.image_captioning_url - - # Find all images - http_matches = self.HTTP_IMAGE_PATTERN.findall(content) - data_uri_matches = self.DATA_URI_IMAGE_PATTERN.findall(content) - - logger.debug( - "Found images in markdown", - http_images=len(http_matches), - data_uri_images=len(data_uri_matches), - ) - - # Build tasks dict mapping markdown syntax to coroutine - tasks = {} - - if caption_http_urls: - for alt, url in http_matches: - markdown_syntax = f"![{alt}]({url})" - tasks[markdown_syntax] = self.get_image_description(url) - - if caption_data_uris: - for alt, data_uri in data_uri_matches: - markdown_syntax = f"![{alt}]({data_uri})" - tasks[markdown_syntax] = self.get_image_description(data_uri) - - if not tasks: - return content - - # Execute all captioning tasks concurrently - try: - captions = await tqdm.gather(*tasks.values(), desc=desc) - image_to_caption = dict(zip(tasks.keys(), captions)) - - # Replace images with captions - logger.debug("Replacing image references", image_count=len(image_to_caption)) - for md_syntax, caption in image_to_caption.items(): - content = content.replace(md_syntax, caption) - - except asyncio.CancelledError: - logger.warning("Image captioning cancelled") - raise - - return content diff --git a/openrag/services/workers/parsers/legacy_loaders/doc.py b/openrag/services/workers/parsers/legacy_loaders/doc.py deleted file mode 100644 index cf0c7d882..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/doc.py +++ /dev/null @@ -1,78 +0,0 @@ -""" -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 -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 core.utils.logging import get_logger -from langchain_core.documents.base import Document as LCDocument -from PIL import Image - -from .base import BaseLoader - -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._parser = DocParser() - - async def aload_document(self, file_path, metadata, save_markdown=False): - 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.") - for block in processed.images: - ref = (block.metadata or {}).get("markdown_ref") - if ref: - result = result.replace(ref, "") - else: - 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/services/workers/parsers/legacy_loaders/docx.py b/openrag/services/workers/parsers/legacy_loaders/docx.py deleted file mode 100644 index 7688f87bc..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/docx.py +++ /dev/null @@ -1,123 +0,0 @@ -""" -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 core.indexing.parsers.docx_parser import DocxParser -from core.models.document import Document as CoreDocument -from core.models.document import DocumentType -from core.utils.logging import get_logger -from langchain_core.documents.base import Document -from PIL import Image - -from .base import BaseLoader, ensure_png_compatible_mode - -logger = get_logger() - - -def convert_to_png_image(image: Image.Image) -> Image.Image: - image = ensure_png_compatible_mode(image) - with BytesIO() as buffer: - image.save(buffer, format="PNG") - buffer.seek(0) - 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._parser = DocxParser() - - async def aload_document(self, file_path, metadata, save_markdown=False): - 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) - - result = await self.replace_markdown_images_with_captions( - result, - caption_data_uris=False, - desc="Captioning linked images", - ) - - doc = Document(page_content=result, metadata=dict(metadata) if metadata else {}) - if save_markdown: - self.save_content(result, str(file_path)) - return doc - - # ----- legacy helpers retained for test_docx_loader.py compatibility ----- - - def get_images_from_zip(self, input_file): - try: - docx = zipfile.ZipFile(input_file, "r") - except zipfile.BadZipFile: - logger.warning("File is not a valid zip archive; skipping image extraction.", path=str(input_file)) - return [] - with docx: - file_names = docx.namelist() - image_files = [f for f in file_names if f.startswith("word/media/")] - if not image_files: - return [] - - images_not_in_order, order = [], [] - for image_file in image_files: - image_data = docx.read(image_file) - image_extension = image_file.split(".")[-1].lower() - try: - image = Image.open(BytesIO(image_data)) - image = convert_to_png_image(image) - order_num = int(image_file.split("media/image")[1].split(f".{image_extension}")[0]) - except Exception as e: - logger.warning(f"Skipping unsupported media file {image_file}: {e}") - continue - - images_not_in_order.append(image) - order.append(order_num) - - if not images_not_in_order: - return [] - - max_order = max(order) - images = [None] * max_order - for i, pos in enumerate(order): - images[pos - 1] = images_not_in_order[i] - return images diff --git a/openrag/services/workers/parsers/legacy_loaders/eml_loader.py b/openrag/services/workers/parsers/legacy_loaders/eml_loader.py deleted file mode 100644 index 8ef94d1f7..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/eml_loader.py +++ /dev/null @@ -1,334 +0,0 @@ -import datetime -import email -import io -import os -import tempfile -from email.utils import parsedate_to_datetime -from pathlib import Path - -from langchain_core.documents.base import Document -from PIL import Image - -from . import get_loader_classes -from .base import BaseLoader - - -def json_serial(obj): - if isinstance(obj, datetime.datetime): - serial = obj.isoformat() - return serial - - -class EmlLoader(BaseLoader): - def __init__(self, **kwargs): - super().__init__(**kwargs) - # Store all kwargs for passing to sub-loaders - self.kwargs = kwargs - # Get available loaders for processing attachments - self.loader_classes = get_loader_classes(config=self.config) - # Cap how deeply nested .eml attachments may be processed; operator- - # tunable via loader config. Bounds recursion when .eml files are - # nested inside one another. - self.max_eml_recursion_depth = self.config.loader.eml_max_recursion_depth - - async def aload_document( - self, - file_path, - metadata: dict | None = None, - save_markdown: bool = False, - _eml_recursion_depth: int = 0, - ): - try: - with open(file_path, "rb") as fhdl: - raw_email = fhdl.read() - - # Parse email using standard email library - email_msg = email.message_from_bytes(raw_email) - - # Extract email metadata - email_data = { - "header": { - "subject": email_msg.get("subject", ""), - "from": email_msg.get("from", ""), - "to": email_msg.get("to", ""), - "date": email_msg.get("date", ""), - "message-id": email_msg.get("message-id", ""), - }, - "body": [], - "attachment": [], - } - - # Parse date if available - if email_data["header"]["date"]: - try: - email_data["header"]["date"] = parsedate_to_datetime(email_data["header"]["date"]).isoformat() - except Exception: - pass - - # Extract body content and attachments - body_content = "" - - for part in email_msg.walk(): - content_type = part.get_content_type() - content_disposition = part.get_content_disposition() - - if content_disposition == "attachment" or content_disposition == "inline": - # Handle attachments - filename = part.get_filename() - if filename: - payload = part.get_payload(decode=True) - if payload: - attachment_info = { - "filename": filename, - "content_type": content_type, - "size": len(payload), - "raw": payload, - } - email_data["attachment"].append(attachment_info) - - elif content_type.startswith("text/"): - # Handle text content - if content_type == "text/plain" or content_type == "text/html": - text_content = part.get_payload(decode=True) - if text_content: - try: - # Try to decode as UTF-8, fallback to latin-1 - if isinstance(text_content, bytes): - try: - text_content = text_content.decode("utf-8") - except UnicodeDecodeError: - text_content = text_content.decode("latin-1", errors="ignore") - - body_info = { - "content": text_content, - "content_type": content_type, - } - email_data["body"].append(body_info) - - # Use plain text as primary body content - if content_type == "text/plain" or not body_content: - body_content = text_content - except Exception as e: - print(f"Failed to decode text content: {e}") - - # Extract body content - content_body = body_content.strip() if body_content else "" - - # Process attachments using appropriate loaders - attachments_text = "" - if email_data["attachment"]: - attachments_text = "\n\n--- ATTACHMENTS ---\n" - for attachment in email_data["attachment"]: - filename = attachment.get("filename", "unknown") - content_type = attachment.get("content_type", "unknown") - size = attachment.get("size", "unknown") - - attachments_text += f"\nAttachment: {filename}\n" - attachments_text += f"Content-Type: {content_type}\n" - attachments_text += f"Size: {size} bytes\n" - - # Try to process attachment using appropriate loader - if "raw" in attachment: - try: - # Get file extension from filename - file_ext = Path(filename).suffix.lower() - - # Check if we have a loader for this file type - loader_cls = self.loader_classes.get(file_ext) - - # Stop the recursion before we descend into another - # .eml — past max_eml_recursion_depth we annotate - # the chain and skip the load. - if loader_cls is EmlLoader and _eml_recursion_depth + 1 >= self.max_eml_recursion_depth: - attachments_text += ( - f"Skipped nested .eml attachment '{filename}': " - f"recursion depth limit ({self.max_eml_recursion_depth}) reached.\n" - ) - attachments_text += "---\n" - continue - - if loader_cls: - # Save attachment to temporary file - with tempfile.NamedTemporaryFile(suffix=file_ext, delete=False) as temp_file: - temp_file.write(attachment["raw"]) - temp_file_path = temp_file.name - - try: - # Use appropriate loader to process attachment - loader = loader_cls(**self.kwargs) - sub_kwargs: dict = {} - if loader_cls is EmlLoader: - sub_kwargs["_eml_recursion_depth"] = _eml_recursion_depth + 1 - attachment_doc = await loader.aload_document( - temp_file_path, - metadata={"source": f"attachment:{filename}"}, - **sub_kwargs, - ) - attachments_text += f"Content:\n{attachment_doc.page_content}\n" - except Exception as e: - attachments_text += f"Failed to process attachment with loader ({loader_cls.__name__}): {str(e)[:200]}...\n" - - # Special fallback handling for PDFs with alternative loaders - if file_ext == ".pdf": - pdf_fallback_loaders = [ - "PyMuPDFLoader", - "PyMuPDF4LLMLoader", - "DoclingLoader", - ] - fallback_success = False - - for fallback_loader_name in pdf_fallback_loaders: - if ( - fallback_loader_name != loader_cls.__name__ - ): # Don't try the same loader again - try: - # Try to get the fallback loader class - fallback_loader_cls = None - for ( - ext, - cls, - ) in self.loader_classes.items(): - if cls.__name__ == fallback_loader_name: - fallback_loader_cls = cls - break - - if fallback_loader_cls: - attachments_text += ( - f"Trying fallback PDF loader: {fallback_loader_name}\n" - ) - fallback_loader = fallback_loader_cls(**self.kwargs) - attachment_doc = await fallback_loader.aload_document( - temp_file_path, - metadata={"source": f"attachment:{filename}"}, - ) - attachments_text += f"Content (via {fallback_loader_name}):\n{attachment_doc.page_content}\n" - fallback_success = True - break - except Exception as fallback_e: - attachments_text += f"Fallback {fallback_loader_name} also failed: {str(fallback_e)[:100]}...\n" - - if not fallback_success: - attachments_text += f"All PDF loaders failed for {filename}\n" - - # Try fallback processing for images - if file_ext in [".png", ".jpg", ".jpeg", ".svg"]: - try: - if self.image_captioning: - # Try to load image directly from bytes as fallback - image = Image.open(io.BytesIO(attachment["raw"])) - caption = await self.get_image_description(image_data=image) - attachments_text += f"Fallback Image Description:\n{caption}\n" - else: - attachments_text += ( - "Image attachment present but image captioning disabled\n" - ) - except Exception as img_e: - attachments_text += f"Image fallback also failed: {str(img_e)[:100]}...\n" - - # Try text fallback for other text-based formats - elif file_ext in [".txt", ".docx", ".doc"] or ( - file_ext == ".pdf" and not fallback_success - ): - try: - # Try to extract any readable text directly - text_content = attachment["raw"].decode("utf-8", errors="ignore") - if text_content.strip(): - attachments_text += ( - f"Fallback text extraction:\n{text_content[:1000]}...\n" - ) - else: - attachments_text += "No readable text found in attachment\n" - except Exception as text_e: - attachments_text += f"Text fallback failed: {str(text_e)[:100]}...\n" - finally: - # Clean up temporary file - if os.path.exists(temp_file_path): - os.unlink(temp_file_path) - - # Special handling for images with captioning if no specific loader or captioning is enabled - elif ( - file_ext - in [ - ".png", - ".jpg", - ".jpeg", - ".svg", - ] - and self.image_captioning - ): - try: - # Load image from raw bytes - image = Image.open(io.BytesIO(attachment["raw"])) - # Verify image can be processed - image.verify() - # Reopen image since verify() closes it - image = Image.open(io.BytesIO(attachment["raw"])) - # Generate caption using the base loader's method - caption = await self.get_image_description(image_data=image) - attachments_text += f"Image Description:\n{caption}\n" - except Exception as e: - attachments_text += f"Failed to generate image caption: {str(e)[:200]}...\n" - # Try to show basic image info if available - try: - size_info = f"Image size: {len(attachment['raw'])} bytes" - attachments_text += ( - f"Image attachment present but corrupted or unreadable. {size_info}\n" - ) - except Exception: - attachments_text += "Image attachment present but corrupted or unreadable\n" - - elif content_type.startswith("text/"): - # For text attachments, decode directly - attachment_content = attachment["raw"].decode("utf-8", errors="ignore") - attachments_text += f"Content:\n{attachment_content}\n" - else: - # For other binary content, just show metadata - attachments_text += f"Binary content (size: {len(attachment['raw'])} bytes)\n" - except Exception as e: - attachments_text += f"Content could not be processed: {e}\n" - attachments_text += "---\n" - - # Combine body and attachments - content_body = content_body + attachments_text - - # Prepare metadata - if metadata is None: - metadata = {} - - # Add email metadata to document metadata - metadata.update( - { - "email_subject": email_data["header"]["subject"], - "email_from": email_data["header"]["from"], - "email_to": email_data["header"]["to"], - "email_date": email_data["header"]["date"], - "email_message_id": email_data["header"]["message-id"], - "email_attachment_count": len(email_data["attachment"]), - "email_attachment_filenames": [att["filename"] for att in email_data["attachment"]], - } - ) - - # Add attachment metadata if there are attachments - if email_data["attachment"]: - attachment_metadata = [] - for att in email_data["attachment"]: - attachment_metadata.append( - { - "filename": att["filename"], - "content_type": att["content_type"], - "size": att["size"], - } - ) - metadata["email_attachments"] = attachment_metadata - - # Save content body to a file if requested - if save_markdown: - markdown_path = Path(file_path).with_suffix(".md") - with open(markdown_path, "w", encoding="utf-8") as md_file: - md_file.write(content_body) - metadata["markdown_path"] = str(markdown_path) - except Exception as e: - raise ValueError(f"Failed to parse the EML file {file_path}: {e}") - - document = Document(page_content=content_body, metadata=metadata) - return document diff --git a/openrag/services/workers/parsers/legacy_loaders/image.py b/openrag/services/workers/parsers/legacy_loaders/image.py deleted file mode 100644 index 9aa418b8c..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/image.py +++ /dev/null @@ -1,84 +0,0 @@ -""" -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 - -from core.indexing.parsers.image_parser import ImageParser -from core.models.document import Document as CoreDocument -from core.models.document import DocumentType -from core.utils.exceptions import OpenRAGError -from core.utils.logging import get_logger -from langchain_core.documents import Document -from PIL import Image - -from .base import BaseLoader - -log = get_logger() - - -class ImageLoadError(OpenRAGError): - """Raised when an image file cannot be loaded or converted.""" - - def __init__(self, message: str, **kwargs): - super().__init__(message, code="IMAGE_LOAD_ERROR", status_code=500, **kwargs) - - -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): - if metadata is None: - metadata = {} - - path = Path(file_path) - try: - raw_bytes = await asyncio.to_thread(path.read_bytes) - except Exception as e: - log.error( - "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, - ) - try: - processed = await self._parser.parse(core_doc) - if not processed.images or not processed.images[0].image_bytes: - raise ImageLoadError(f"Cannot load image '{path.name}': failed to decode") - - img = Image.open(BytesIO(processed.images[0].image_bytes)) - img.load() - except ImageLoadError: - raise - except Exception as e: - log.error("Failed to decode 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 - 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)) - return doc diff --git a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/__init__.py b/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/__init__.py deleted file mode 100644 index 8b1378917..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/docling.py b/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/docling.py deleted file mode 100644 index dd4a7b046..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/docling.py +++ /dev/null @@ -1,84 +0,0 @@ -import asyncio - -import torch -from core.utils.logging import get_logger -from core.utils.singleton import SingletonMeta -from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend -from docling.datamodel.base_models import InputFormat -from docling.datamodel.document import ConversionResult -from docling.datamodel.pipeline_options import ( - AcceleratorDevice, - AcceleratorOptions, - PdfPipelineOptions, - TableFormerMode, - TableStructureOptions, -) -from docling.document_converter import DocumentConverter, PdfFormatOption -from docling_core.types.doc.document import PictureItem -from langchain_core.documents.base import Document - -from ..base import BaseLoader - -logger = get_logger() - - -class DoclingConverter(metaclass=SingletonMeta): - def __init__(self): - img_scale = 1 - pipeline_options = PdfPipelineOptions( - do_ocr=True, - do_table_structure=True, - generate_picture_images=True, - images_scale=img_scale, - generate_table_images=True, - # generate_page_images=True - ) - pipeline_options.table_structure_options = TableStructureOptions( - do_cell_matching=True, mode=TableFormerMode.ACCURATE - ) - - pipeline_options.accelerator_options = AcceleratorOptions(device=AcceleratorDevice.AUTO) - self.converter = DocumentConverter( - format_options={ - InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options, backend=PyPdfiumDocumentBackend) - } - ) - - async def convert_to_md(self, file_path) -> ConversionResult: - o = await asyncio.to_thread(self.converter.convert, str(file_path)) - return o - - -class DoclingLoader(BaseLoader): - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.converter = DoclingConverter() - - async def convert_to_md(self, file_path) -> ConversionResult: - return await asyncio.to_thread(self.converter.convert, str(file_path)) - - async def aload_document(self, file_path, metadata, save_markdown=False): - with torch.no_grad(): - result = await self.converter.convert_to_md(file_path) - - n_pages = len(result.pages) - - s = "" - for i in range(1, n_pages + 1): - s += result.document.export_to_markdown(page_no=i) - s += f"\n[PAGE_{i}]\n" - - enriched_content = s - if self.image_captioning: - pictures: list[PictureItem] = result.document.pictures - images = [p.image.pil_image for p in pictures] - descriptions = await self.caption_images(images, desc="Captioning imgs") - for description in descriptions: - enriched_content = enriched_content.replace("", description, 1) - else: - logger.debug("Image captioning disabled. Ignoring images.") - - doc = Document(page_content=enriched_content, metadata=metadata) - if save_markdown: - self.save_content(enriched_content, str(file_path)) - return doc diff --git a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/docling2.py b/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/docling2.py deleted file mode 100644 index 86a19b182..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/docling2.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Docling-backed PDF loader. - -The Ray actor + pool (``DoclingWorker``, ``DoclingPool``) and the -services-side :class:`BasePooledParser` implementation now live in -``services/workers/parsers/docling_workers.py``; this module re-exports -them for legacy import paths (``services.workers.bootstrap`` constructs the -named ``DoclingPool`` actor at startup via ``get_or_create_actor``). - -``DoclingLoader2`` is a thin :class:`BaseLoader` adapter that delegates -to :class:`core.indexing.parsers.pdf.docling.DoclingParser`, which -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. -""" - -from __future__ import annotations - -from core.utils.logging import get_logger -from langchain_core.documents.base import Document -from services.workers.parsers.docling_workers import ( # noqa: F401 (re-exported for legacy paths) - DoclingLoader, - DoclingPool, - DoclingWorker, -) - -from ..base import BaseLoader - -logger = get_logger() - - -class DoclingLoader2(BaseLoader): - """Adapter shim — delegates to ``DoclingParser`` via the services-side pool.""" - - def __init__(self, **kwargs): - super().__init__(**kwargs) - from core.indexing.parsers.pdf.docling import DoclingParser - from services.workers.parsers.docling_workers import DoclingLoader as _DoclingLoader - - self._parser = DoclingParser(pool=_DoclingLoader()) - - async def aload_document(self, file_path, metadata, save_markdown=False): - import asyncio - from pathlib import Path - - from core.models.document import Document as CoreDocument - from core.models.document import DocumentType - - path = Path(file_path) - 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 or {}), - ) - processed = await self._parser.parse(core_doc) - - markdown = "" - for block in processed.text_blocks: - markdown += block.text + f"\n[PAGE_{block.page_number}]\n" - - if self.image_captioning and processed.images: - import io - - from PIL import Image - - pil_images = [] - for img_block in processed.images: - if img_block.image_bytes: - pil_images.append(Image.open(io.BytesIO(img_block.image_bytes))) - - if pil_images: - captions = await self.caption_images(pil_images) - for img_block, caption in zip(processed.images, captions): - ref = (img_block.metadata or {}).get("markdown_ref") - if ref: - markdown = markdown.replace("", caption, 1) - else: - logger.debug("Image captioning disabled. Ignoring images.") - - doc = Document(page_content=markdown, metadata=metadata) - if save_markdown: - self.save_document(Document(page_content=markdown), str(file_path)) - return doc diff --git a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/dotsocr.py b/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/dotsocr.py deleted file mode 100644 index 7c1e7a6af..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/dotsocr.py +++ /dev/null @@ -1,62 +0,0 @@ -from core.utils.logging import get_logger -from PIL import Image -from tqdm.asyncio import tqdm - -from .openai import OpenAILoader - -logger = get_logger() - - -class DotsOCRLoader(OpenAILoader): - """PDF loader using 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. -""" - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - async def _caption_images(self, page_img: Image.Image, page_res: list): - """Extract picture elements and caption them.""" - picture_items = [item for item in page_res if item.get("category") == "Picture"] - if not picture_items: - return - - picture_crops = [] - for item in picture_items: - bbox = item.get("bbox") - if bbox and len(bbox) == 4: - try: - cropped = page_img.crop(bbox) - picture_crops.append((item, cropped)) - except Exception as e: - logger.warning(f"Failed to crop image bbox {bbox}: {e}") - - if picture_crops: - desc_tasks = [self._get_caption(crop) for _, crop in picture_crops] - desc_results = await tqdm.gather( - *desc_tasks, - desc="Captioning images", - total=len(desc_tasks), - ) - for (item, _), desc in zip(picture_crops, desc_results): - item["text"] = desc.strip() if isinstance(desc, str) else "" - - def _result_to_md(self, result: list[dict]) -> str: - return "\n".join(item.get("text", "").strip() for item in result if item.get("text")) diff --git a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/marker.py b/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/marker.py deleted file mode 100644 index 111db223e..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/marker.py +++ /dev/null @@ -1,106 +0,0 @@ -""" -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 -(``services.workers.bootstrap`` 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 time -from io import BytesIO -from pathlib import Path - -from core.indexing.parsers.pdf.marker import MarkerParser -from core.models.document import Document as CoreDocument -from core.models.document import DocumentType -from core.utils.logging import get_logger -from langchain_core.documents.base import Document -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 ..base import BaseLoader - -logger = get_logger() - - -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._parser = MarkerParser(pool=_ServicesMarkerPool()) - - async def aload_document( - self, - file_path: str | Path, - metadata: dict | None = None, - save_markdown: bool = False, - ) -> Document: - if metadata is None: - metadata = {} - - path = Path(file_path) - file_path_str = str(file_path) - start = time.time() - - try: - 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 not self.image_captioning: - logger.debug("Image captioning disabled.") - 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) - - duration = time.time() - start - logger.info(f"Processed {file_path_str} in {duration:.2f}s") - return doc - - except Exception: - logger.exception("Error in aload_document", path=file_path_str) - raise diff --git a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/openai.py b/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/openai.py deleted file mode 100644 index 2ad4461c8..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/openai.py +++ /dev/null @@ -1,132 +0,0 @@ -import asyncio -import base64 -import io -import json -import time -from abc import ABC, abstractmethod -from pathlib import Path - -import pypdfium2 as pdfium -from core.utils.logging import get_logger -from langchain.schema import Document -from langchain_openai import ChatOpenAI -from PIL import Image - -from ..base import BaseLoader - -logger = get_logger() - - -async def pdf_to_images(pdf_path: str, scale: float = 1.0) -> list[Image.Image]: - pdf: pdfium.PdfDocument = await asyncio.to_thread(pdfium.PdfDocument, pdf_path) - return [p.render(scale=scale).to_pil() for p in pdf] - - -class OpenAILoader(BaseLoader, ABC): - """Generic OpenAI-compatible loader for multimodal OCR-style models.""" - - PROMPT: str = "" # To be defined by subclasses - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - self.llm = ChatOpenAI( - base_url=self.config.loader.openai.base_url, - api_key=self.config.loader.openai.api_key, - model=self.config.loader.openai.model, - temperature=self.config.loader.openai.temperature, - timeout=self.config.loader.openai.timeout, - max_retries=self.config.loader.openai.max_retries, - top_p=self.config.loader.openai.top_p, - ) - self.llm_semaphore = asyncio.Semaphore(self.config.loader.openai.concurrency_limit) - - async def aload_document( - self, - file_path: str | Path, - metadata: dict | None = None, - save_markdown: bool = False, - ) -> Document: - """Main pipeline: PDF → OCR → Caption → Markdown.""" - if metadata is None: - metadata = {} - - start_time = time.time() - file_path = str(file_path) - - try: - pages = await pdf_to_images(file_path) - ocr_results = await self._run_ocr_on_pages(pages) - markdown = await self._assemble_markdown(pages, ocr_results) - - if save_markdown: - self.save_content(markdown, file_path) - - duration = time.time() - start_time - logger.info(f"Processed {file_path} in {duration:.2f}s") - return Document(page_content=markdown, metadata=metadata) - - except Exception: - logger.exception("Error in OpenAILoader.aload_document", path=file_path) - raise - - async def _run_ocr_on_pages(self, pages: list[Image.Image]) -> list[dict]: - tasks = [self._img2result(img) for img in pages] - return await asyncio.gather(*tasks) - - async def _assemble_markdown(self, pages: list[Image.Image], results: list[dict]) -> str: - markdown_parts = [] - for page_img, page_res in zip(pages, results): - if not page_res: - continue - if self.image_captioning: - await self._caption_images(page_img, page_res) - markdown_parts.append(self._result_to_md(page_res)) - return "\n\n".join(markdown_parts).strip() - - async def _get_caption(self, img: Image.Image) -> str: - try: - return await self.get_image_description(image_data=img) - except Exception as e: - logger.warning(f"Captioning failed: {e}") - return "" - - async def _img2result(self, img: Image.Image, format: str = "PNG") -> dict: - """Send an image to the OpenAI-compatible OCR model.""" - async with self.llm_semaphore: - try: - buffer = io.BytesIO() - img.save(buffer, format=format) - img_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8") - - messages = [ - { - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": {"url": f"data:image/{format.lower()};base64,{img_b64}"}, - }, - { - "type": "text", - "text": f"<|img|><|imgpad|><|endofimg|>{self.PROMPT}", - }, - ], - } - ] - - response = await self.llm.ainvoke(messages) - data = json.loads(response.content) - return data - - except Exception as e: - logger.error("Error in _img2result", error=str(e)) - return {} - - @abstractmethod - def _result_to_md(self, result: list[dict]) -> str: - """Convert structured OCR + caption results to markdown format.""" - - @abstractmethod - async def _caption_images(self, page_img: Image.Image, page_res: list): - """Extract picture elements and caption them.""" diff --git a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/pymupdf.py b/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/pymupdf.py deleted file mode 100644 index c46dd0d89..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/pdf_loaders/pymupdf.py +++ /dev/null @@ -1,107 +0,0 @@ -""" -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 - -from core.indexing.parsers.pdf.pymupdf import PyMuPDFParser -from core.models.document import Document as CoreDocument -from core.models.document import DocumentType -from core.utils.logging import get_logger -from langchain_core.documents.base import Document -from PIL import Image - -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): - 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) - - doc = Document(page_content=s, metadata=metadata) - if save_markdown: - self.save_content(s, str(file_path)) - return doc - - -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): - 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) - - 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: - self.save_content(s, str(file_path)) - return doc diff --git a/openrag/services/workers/parsers/legacy_loaders/pptx_loader.py b/openrag/services/workers/parsers/legacy_loaders/pptx_loader.py deleted file mode 100644 index 18b14f177..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/pptx_loader.py +++ /dev/null @@ -1,71 +0,0 @@ -""" -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 - -from core.indexing.parsers.pptx_parser import PptxParser -from core.models.document import Document as CoreDocument -from core.models.document import DocumentType -from core.utils.logging import get_logger -from langchain_core.documents.base import Document -from PIL import Image - -from .base import BaseLoader - -logger = get_logger() - - -class PPTXLoader(BaseLoader): - """Adapter shim — delegates to ``PptxParser``; layers image captioning on top.""" - - def __init__(self, **kwargs) -> None: - super().__init__(**kwargs) - self._parser = PptxParser() - - async def aload_document(self, file_path, metadata=None, save_markdown=False): - 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.") - 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) - - doc = Document(page_content=md_content, metadata=metadata) - if save_markdown: - self.save_content(md_content, str(file_path)) - return doc diff --git a/openrag/services/workers/parsers/legacy_loaders/txt_loader.py b/openrag/services/workers/parsers/legacy_loaders/txt_loader.py deleted file mode 100644 index 82c4062c7..000000000 --- a/openrag/services/workers/parsers/legacy_loaders/txt_loader.py +++ /dev/null @@ -1,94 +0,0 @@ -""" -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 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 core.utils.logging import get_logger -from langchain_core.documents.base import Document -from services.workers.parsers.legacy_loaders.base import BaseLoader - -logger = get_logger() - - -class TextLoader(BaseLoader): - """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, - file_path: str | Path, - metadata: dict | None = None, - save_markdown: bool = False, - ) -> Document: - 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.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: - self.save_content(content, str(path)) - - return doc - - -class MarkdownLoader(BaseLoader): - """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, - file_path: str | Path, - metadata: dict | None = None, - save_markdown: bool = False, - ) -> Document: - 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.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() - - content = await self.replace_markdown_images_with_captions(content) - - doc = Document(page_content=content, metadata=metadata) - if save_markdown: - self.save_content(text_content=content, path=str(path)) - return doc diff --git a/openrag/services/workers/parsers/parser_dispatcher.py b/openrag/services/workers/parsers/parser_dispatcher.py new file mode 100644 index 000000000..d79ea076f --- /dev/null +++ b/openrag/services/workers/parsers/parser_dispatcher.py @@ -0,0 +1,250 @@ +"""Content-type → parser dispatch over the ``core/indexing/parsers`` stack. + +``ParserDispatcher`` is the single ``DocumentParser`` the indexing pipeline +and the extract path use. It routes each ``Document`` to the right concrete +parser by ``Document.content_type`` (``DocumentType``), resolving the PDF and +audio backends from the existing ``config.loader.file_loaders`` map so behavior +matches the GPU pools ``bootstrap.py`` provisions. + +This replaces the transitional ``DocSerializerBridgeParser`` (which delegated +to the legacy ``BaseLoader`` registry). Backends are built lazily on first use +and cached, so importing this module is cheap and a deploy never needs a +pool/library for a backend it doesn't exercise. +""" + +from __future__ import annotations + +import importlib +from typing import Any + +from core.indexing.parsers.document_parser import DocumentParser +from core.models.document import Document, DocumentType, ProcessedDocument +from core.utils.logging import get_logger + +logger = get_logger() + +# Translate the legacy ``file_loaders`` class-name values into new registry +# backend names. Mirrors ``bootstrap.get_marker_pool`` / ``init_audio_actor`` so +# the dispatcher always resolves to a backend whose pool actually exists. +_PDF_BACKENDS: dict[str, str] = { + "MarkerLoader": "marker", + "DoclingLoader": "docling", + "PyMuPDFLoader": "pymupdf", + "DotsOCRLoader": "pdf_client", + "OpenAILoader": "pdf_client", +} +_AUDIO_BACKENDS: dict[str, str] = { + "LocalWhisperLoader": "local_whisper", + "OpenAIAudioLoader": "audio_client", +} + +# Attachment ext (lowercased, no dot) → DocumentType, used to wire the EML +# parser's per-attachment sub-parsers. +_MAX_EML_ATTACHMENT_DEPTH = 3 +_EML_ATTACHMENT_TYPES: dict[str, DocumentType] = { + "txt": DocumentType.TEXT, + "md": DocumentType.MARKDOWN, + "html": DocumentType.HTML, + "htm": DocumentType.HTML, + "eml": DocumentType.EML, + "docx": DocumentType.DOCX, + "doc": DocumentType.DOC, + "pptx": DocumentType.PPTX, + "pdf": DocumentType.PDF, + "png": DocumentType.IMAGE, + "jpg": DocumentType.IMAGE, + "jpeg": DocumentType.IMAGE, + "gif": DocumentType.IMAGE, + "webp": DocumentType.IMAGE, + "bmp": DocumentType.IMAGE, + "svg": DocumentType.IMAGE, +} + + +def _create(module_path: str, name: str, **kwargs: Any) -> DocumentParser: + """Import a parser module (triggering registration) then build it by name.""" + importlib.import_module(module_path) + from core.indexing.parsers.registry import parser_registry + + return parser_registry.create(name, **kwargs) + + +class ParserDispatcher(DocumentParser): + """Route a document to the configured concrete parser by content type.""" + + def __init__(self, config: Any) -> None: + self._config = config + self._by_name: dict[str, DocumentParser] = {} + + def supported_types(self) -> list[str]: + return [doc_type.value for doc_type in DocumentType] + + async def parse(self, document: Document) -> ProcessedDocument: + backend = self._resolve_backend(document.content_type, _suffix(document.filename)) + parser = self._get(backend) + return await parser.parse(document) + + # ----- backend resolution ----- + + def _resolve_backend(self, content_type: DocumentType, ext: str) -> str: + if content_type is DocumentType.PDF: + return self._resolve_pdf_backend() + if content_type in (DocumentType.AUDIO, DocumentType.VIDEO): + return self._resolve_audio_backend(ext) + # Every other type maps 1:1 to a registered parser whose name is the + # type value (TEXT->"text", DOCX->"docx", EML->"eml", ...). + return content_type.value + + def _resolve_pdf_backend(self) -> str: + configured = self._config.loader.file_loaders.pdf + backend = _PDF_BACKENDS.get(configured) + if backend is None: + raise ValueError( + f"Unsupported PDF loader configuration {configured!r}; expected one of {sorted(_PDF_BACKENDS)}" + ) + return backend + + def _resolve_audio_backend(self, ext: str) -> str: + file_loaders = self._config.loader.file_loaders + configured = ( + getattr(file_loaders, ext, None) or getattr(file_loaders, "mp3", None) or getattr(file_loaders, "wav", None) + ) + backend = _AUDIO_BACKENDS.get(configured) + if backend is None: + raise ValueError( + f"Unsupported audio loader configuration {configured!r}; expected one of {sorted(_AUDIO_BACKENDS)}" + ) + return backend + + # ----- lazy backend construction ----- + + def _get(self, name: str) -> DocumentParser: + parser = self._by_name.get(name) + if parser is None: + parser = self._build(name) + self._by_name[name] = parser + return parser + + def _build(self, name: str) -> DocumentParser: + logger.debug(f"Building parser backend: {name}") + builder = _BUILDERS.get(name) + if builder is not None: + return builder(self) + # Convention for simple, dependency-free parsers: registry name ``X`` + # lives in ``core.indexing.parsers.X_parser`` and registers as ``X``. + return _create(f"core.indexing.parsers.{name}_parser", name) + + # ----- backend builders (lazy heavy imports live inside these) ----- + + def _build_eml(self, attachment_depth: int = 0) -> DocumentParser: + attachment_parsers: dict[str, DocumentParser] = {} + for ext, dtype in _EML_ATTACHMENT_TYPES.items(): + try: + if dtype is DocumentType.EML: + if attachment_depth >= _MAX_EML_ATTACHMENT_DEPTH: + continue + attachment_parsers[ext] = self._build_eml(attachment_depth + 1) + continue + attachment_parsers[ext] = self._get(self._resolve_backend(dtype, ext)) + except Exception as exc: # a missing backend must not break .eml parsing + logger.warning(f"EML attachment parser for '.{ext}' unavailable: {exc}") + return _create("core.indexing.parsers.eml_parser", "eml", attachment_parsers=attachment_parsers) + + def _build_marker(self) -> DocumentParser: + from services.workers.parsers.marker_workers import MarkerLoader + + return _create("core.indexing.parsers.pdf.marker", "marker", pool=MarkerLoader()) + + def _build_docling(self) -> DocumentParser: + from services.workers.parsers.docling_workers import DoclingLoader + + return _create("core.indexing.parsers.pdf.docling", "docling", pool=DoclingLoader()) + + def _build_local_whisper(self) -> DocumentParser: + from services.workers.parsers.whisper_workers import LocalWhisperLoader + + return _create("core.indexing.parsers.audio.local_whisper", "local_whisper", pool=LocalWhisperLoader()) + + def _build_pdf_client(self) -> DocumentParser: + from services.inference.parsers.dotsocr import DotsOCRPdfClient + + ocfg = self._config.loader.openai + vlm = _build_vlm(ocfg.base_url, ocfg.model, ocfg.api_key, ocfg.timeout) + client = DotsOCRPdfClient(vlm, concurrency_limit=ocfg.concurrency_limit) + return _create("core.indexing.parsers.pdf.client_based", "pdf_client", client=client) + + def _build_audio_client(self) -> DocumentParser: + from services.inference.parsers.openai_audio import OpenAIAudioClient + + tcfg = self._config.loader.transcriber + language_detector = None + if tcfg.use_whisper_lang_detector: + from services.workers.parsers.whisper_workers import detect_language_via_actor + + async def language_detector(path): # noqa: E731 - small adapter to the (path) -> str|None contract + return await detect_language_via_actor(path) + + client = OpenAIAudioClient( + base_url=tcfg.base_url, + api_key=tcfg.api_key, + model=tcfg.model_name, + timeout=tcfg.timeout, + direct_upload_suffixes=tcfg.direct_upload_suffixes, + language_detector=language_detector, + concurrency_limit=tcfg.max_concurrent_chunks, + ) + return _create("core.indexing.parsers.audio.client_based", "audio_client", client=client) + + +# Only backends that need a non-conventional module path (``pymupdf`` lives +# under ``pdf/``) or injected dependencies (pooled/client/eml). Everything else +# is built by convention in ``ParserDispatcher._build``. +_BUILDERS: dict[str, Any] = { + "pymupdf": lambda d: _create("core.indexing.parsers.pdf.pymupdf", "pymupdf"), + "eml": lambda d: d._build_eml(), + "marker": lambda d: d._build_marker(), + "docling": lambda d: d._build_docling(), + "local_whisper": lambda d: d._build_local_whisper(), + "pdf_client": lambda d: d._build_pdf_client(), + "audio_client": lambda d: d._build_audio_client(), +} + + +def _suffix(filename: str) -> str: + """Lowercased extension without the dot (``"report.PDF"`` → ``"pdf"``).""" + return filename.rsplit(".", 1)[-1].lower() if "." in filename else "" + + +def _build_vlm(base_url: str, model: str, api_key: str, timeout: float) -> Any: + """Construct a vLLM-backed VLM client for VLM-OCR / captioning.""" + import services.inference.vllm_client # noqa: F401 - registers "vllm" + from core.vlm import vlm_registry + + return vlm_registry.create("vllm", endpoint=base_url, model_name=model, api_key=api_key, timeout=timeout) + + +def build_parser_dispatcher(config: Any) -> ParserDispatcher: + """Build the content-type dispatcher over the new parser stack.""" + return ParserDispatcher(config) + + +def build_caption_vlm(config: Any) -> Any | None: + """Build the captioning VLM, or ``None`` when no VLM endpoint is configured. + + This only decides VLM *availability* (an endpoint must be set), not the + captioning *policy*, which is applied downstream: + + - Standalone image files are always captioned when a VLM is available — + their caption is the only text content (legacy ``ImageLoader`` parity, + which never consulted ``image_captioning``). + - Images embedded in other documents are gated by the global + ``config.loader.image_captioning`` flag and the per-partition + ``enable_image_captioning`` setting (see ``IndexingPipeline``). + """ + vlm_cfg = config.vlm + if not getattr(vlm_cfg, "base_url", ""): + return None + return _build_vlm(vlm_cfg.base_url, vlm_cfg.model, vlm_cfg.api_key, vlm_cfg.timeout) + + +__all__ = ["ParserDispatcher", "build_parser_dispatcher", "build_caption_vlm"] diff --git a/openrag/services/workers/pipeline_builder.py b/openrag/services/workers/pipeline_builder.py index f4d8b8b13..52db8aac7 100644 --- a/openrag/services/workers/pipeline_builder.py +++ b/openrag/services/workers/pipeline_builder.py @@ -10,6 +10,7 @@ from core.indexing.contextualize import ChunkContextualizer from core.indexing.parsers.document_parser import DocumentParser from core.indexing.topic_tags import TopicTagger +from core.models.document import Document, DocumentType from core.vector_stores.vector_store import VectorStore from core.vlm.vlm import VLM from services.workers.stages.caption import caption_stage @@ -47,6 +48,10 @@ class IndexingPipeline: embedder: Embedder vector_store: VectorStore vlm: VLM | None = None + # Global gate for captioning images *embedded* in other documents + # (mirrors ``config.loader.image_captioning``). Standalone image files are + # always captioned when a VLM is available, regardless of this flag. + image_captioning: bool = True contextualizer: ChunkContextualizer | None = None topic_tagger: TopicTagger | None = None timeouts: PipelineTimeouts = PipelineTimeouts() @@ -65,18 +70,19 @@ async def run(self, row: MutableMapping[str, Any]) -> MutableMapping[str, Any]: parser = self._select_parser(config) chunker = self._select_chunker(config) embedder = self._select_embedder(row) - vlm = self._select_vlm(config) contextualizer = self._select_contextualizer(config) topic_tagger = self._select_topic_tagger(config) await parse_stage(row, parser, timeout=self.timeouts.parse) - if vlm is not None: - await caption_stage( - row, - vlm, - timeout=self.timeouts.caption, - per_image_timeout=self.timeouts.caption_per_image, - ) + if self._should_caption(row, config): + vlm = self._select_vlm(config) + if vlm is not None: + await caption_stage( + row, + vlm, + timeout=self.timeouts.caption, + per_image_timeout=self.timeouts.caption_per_image, + ) await chunk_stage(row, chunker, timeout=self.timeouts.chunk) if contextualizer is not None: await contextualize_stage( @@ -134,13 +140,26 @@ def _select_embedder(self, row: MutableMapping[str, Any]) -> Embedder: return self.embedder def _select_vlm(self, config: IndexationPipelineConfig | None) -> VLM | None: - if config is not None: - if not config.enable_image_captioning: - return None - if self.vlm_factory is not None: - return self.vlm_factory(config.vlm or "default") + """Pick the captioning VLM instance (availability only — policy is in + ``_should_caption``).""" + if config is not None and self.vlm_factory is not None: + return self.vlm_factory(config.vlm or "default") return self.vlm + def _should_caption(self, row: MutableMapping[str, Any], config: IndexationPipelineConfig | None) -> bool: + """Decide whether to caption this document's images. + + A standalone image file's caption is its only text content, so it is + always captioned when a VLM is available (legacy ``ImageLoader`` + parity). Images embedded in other documents are gated by the global + ``image_captioning`` flag and the per-partition setting. + """ + document = row.get("document") + if isinstance(document, Document) and document.content_type is DocumentType.IMAGE: + return True + per_partition = config.enable_image_captioning if config is not None else True + return self.image_captioning and per_partition + def _select_contextualizer(self, config: IndexationPipelineConfig | None) -> ChunkContextualizer | None: if config is not None: if not config.enable_contextualization: @@ -165,6 +184,7 @@ def build_indexing_pipeline( embedder: Embedder, vector_store: VectorStore, vlm: VLM | None = None, + image_captioning: bool = True, contextualizer: ChunkContextualizer | None = None, topic_tagger: TopicTagger | None = None, timeouts: PipelineTimeouts | None = None, @@ -184,6 +204,7 @@ def build_indexing_pipeline( embedder=embedder, vector_store=vector_store, vlm=vlm, + image_captioning=image_captioning, contextualizer=contextualizer, topic_tagger=topic_tagger, timeouts=timeouts or PipelineTimeouts(), diff --git a/tests/unit/services/workers/parsers/legacy_loaders/audio/test_openai.py b/tests/unit/services/workers/parsers/legacy_loaders/audio/test_openai.py deleted file mode 100644 index fdd809423..000000000 --- a/tests/unit/services/workers/parsers/legacy_loaders/audio/test_openai.py +++ /dev/null @@ -1,313 +0,0 @@ -""" -Unit tests for openai audio processing functionality (./openai.py). - -These tests validate the pydub operations used in openai.py without -importing the full module (which has complex dependencies). -""" - -import warnings - -# Filter pydub warnings before importing -warnings.filterwarnings("ignore", category=SyntaxWarning, module="pydub") -warnings.filterwarnings("ignore", category=DeprecationWarning, module="pydub") - -import tempfile # noqa: E402 -from pathlib import Path # noqa: E402 - -import pytest # noqa: E402 - -# These tests exercise real pydub audio operations. pydub depends on the stdlib -# ``audioop`` module, which was removed in Python 3.13 — guard on it directly so -# the module skips cleanly (other tests register a mock ``pydub`` in sys.modules, -# so importorskip("pydub") would not catch the missing dependency). -pytest.importorskip("audioop") - -from pydub import AudioSegment # noqa: E402 -from pydub.generators import Sine # noqa: E402 - - -class TestAudioSegmentOperations: - """Test pydub AudioSegment operations used in openai.""" - - def test_create_audio_segment(self): - """Test creating an audio segment.""" - # Create a 1 second sine wave at 440Hz - audio = Sine(440).to_audio_segment(duration=1000) - assert len(audio) == 1000 # 1000ms - - def test_audio_duration(self): - """Test getting audio duration in milliseconds.""" - audio = AudioSegment.silent(duration=500) # 500ms of silence - assert len(audio) == 500 - - def test_audio_slicing(self): - """Test slicing audio by milliseconds.""" - audio = Sine(440).to_audio_segment(duration=2000) # 2 seconds - - # Slice first second - first_half = audio[:1000] - assert len(first_half) == 1000 - - # Slice second half - second_half = audio[1000:] - assert len(second_half) == 1000 - - # Slice middle portion - middle = audio[500:1500] - assert len(middle) == 1000 - - def test_set_channels_mono(self): - """Test converting to mono.""" - # Create stereo audio - stereo = Sine(440).to_audio_segment(duration=1000) - stereo = stereo.set_channels(2) - assert stereo.channels == 2 - - # Convert to mono - mono = stereo.set_channels(1) - assert mono.channels == 1 - - def test_set_frame_rate(self): - """Test changing sample rate.""" - audio = Sine(440).to_audio_segment(duration=1000) - - # Downsample to 16kHz - downsampled = audio.set_frame_rate(16000) - assert downsampled.frame_rate == 16000 - - def test_export_and_load_wav(self): - """Test exporting and loading WAV files.""" - audio = Sine(440).to_audio_segment(duration=1000) - - with tempfile.TemporaryDirectory() as tmpdir: - wav_path = Path(tmpdir) / "test.wav" - - # Export - audio.export(wav_path, format="wav") - assert wav_path.exists() - - # Load back - loaded = AudioSegment.from_wav(wav_path) - assert len(loaded) == 1000 - - def test_export_and_load_mp3(self): - """Test exporting and loading MP3 files (if ffmpeg available).""" - audio = Sine(440).to_audio_segment(duration=1000) - - with tempfile.TemporaryDirectory() as tmpdir: - mp3_path = Path(tmpdir) / "test.mp3" - - try: - # Export as MP3 - audio.export(mp3_path, format="mp3") - assert mp3_path.exists() - - # Load back - loaded = AudioSegment.from_file(mp3_path, format="mp3") - # MP3 may have slight duration differences due to encoding - assert abs(len(loaded) - 1000) < 100 - except (FileNotFoundError, OSError): - # Skip if ffmpeg not available - pytest.skip("ffmpeg not available for MP3 encoding") - - -class TestSilenceDetection: - """Test pydub silence detection used in openai.""" - - def test_detect_silence_all_silent(self): - """Test detection when entire audio is silent.""" - from pydub import silence - - audio = AudioSegment.silent(duration=1000) - silences = silence.detect_silence(audio, min_silence_len=100, silence_thresh=-40) - - # Should detect one long silence - assert len(silences) >= 1 - assert silences[0][0] == 0 - - def test_detect_silence_no_silence(self): - """Test detection when there is no silence.""" - from pydub import silence - - # Loud sine wave - audio = Sine(440).to_audio_segment(duration=1000) - - silences = silence.detect_silence(audio, min_silence_len=100, silence_thresh=-40) - - # Should not detect any silence - assert len(silences) == 0 - - def test_detect_silence_in_middle(self): - """Test detection of silence in the middle of audio.""" - from pydub import silence - - # Loud - Silent - Loud pattern - loud = Sine(440).to_audio_segment(duration=500) - silent = AudioSegment.silent(duration=300) - - audio = loud + silent + loud - - silences = silence.detect_silence(audio, min_silence_len=100, silence_thresh=-40) - - # Should detect one silence segment - assert len(silences) >= 1 - # Silence should start around 500ms - assert 400 <= silences[0][0] <= 600 - - def test_detect_multiple_silences(self): - """Test detection of multiple silence segments.""" - from pydub import silence - - # Pattern: loud-silent-loud-silent-loud - loud = Sine(440).to_audio_segment(duration=300) - silent = AudioSegment.silent(duration=200) - - audio = loud + silent + loud + silent + loud - - silences = silence.detect_silence(audio, min_silence_len=100, silence_thresh=-40) - - # Should detect two silence segments - assert len(silences) >= 2 - - def test_short_silence_ignored(self): - """Test that silences shorter than min_silence_len are ignored.""" - from pydub import silence - - # Loud with very short silence (50ms) - loud = Sine(440).to_audio_segment(duration=500) - short_silent = AudioSegment.silent(duration=50) - - audio = loud + short_silent + loud - - # Request minimum 200ms silence - silences = silence.detect_silence(audio, min_silence_len=200, silence_thresh=-40) - - # Should not detect the short silence - assert len(silences) == 0 - - -class TestAudioChunking: - """Test chunking logic similar to openai._get_audio_chunks.""" - - def get_audio_chunks( - self, - sound: AudioSegment, - max_chunk_ms: int, - min_silence_len_ms: int, - silence_thresh_db: int, - ) -> list: - """ - Reproduce chunking logic from openai. - """ - from pydub import silence - - total_ms = len(sound) - if total_ms <= max_chunk_ms: - return [(0, total_ms)] - - downsampled_sound = sound.set_channels(1).set_frame_rate(16000) - silences = silence.detect_silence( - downsampled_sound, - min_silence_len=min_silence_len_ms, - silence_thresh=silence_thresh_db, - ) - - chunks = [] - start = 0 - while start < total_ms: - target_end = start + max_chunk_ms - if target_end >= total_ms: - end = total_ms - else: - valid_silences = [s for s in silences if start < s[0] < target_end] - if valid_silences: - end = valid_silences[-1][0] - else: - end = target_end - chunks.append((start, end)) - start = end - - return chunks - - def test_short_audio_single_chunk(self): - """Test that short audio results in single chunk.""" - audio = Sine(440).to_audio_segment(duration=5000) # 5 seconds - - chunks = self.get_audio_chunks( - audio, - max_chunk_ms=10000, # 10 seconds max - min_silence_len_ms=500, - silence_thresh_db=-40, - ) - - assert len(chunks) == 1 - assert chunks[0] == (0, 5000) - - def test_long_audio_multiple_chunks(self): - """Test that long audio is split into multiple chunks.""" - # Create 10 seconds of audio with silences - segment = Sine(440).to_audio_segment(duration=2000) + AudioSegment.silent(duration=500) - audio = segment * 3 # Repeat pattern - - chunks = self.get_audio_chunks( - audio, - max_chunk_ms=3000, # 3 seconds max per chunk - min_silence_len_ms=200, - silence_thresh_db=-40, - ) - - # Should have multiple chunks - assert len(chunks) >= 2 - - # All chunks should be within max size (with some tolerance) - for start, end in chunks[:-1]: # Except last chunk - assert end - start <= 3500 # Allow some tolerance - - def test_chunks_cover_entire_audio(self): - """Test that chunks cover the entire audio without gaps.""" - audio = Sine(440).to_audio_segment(duration=10000) # 10 seconds - - chunks = self.get_audio_chunks( - audio, - max_chunk_ms=3000, - min_silence_len_ms=500, - silence_thresh_db=-40, - ) - - # First chunk should start at 0 - assert chunks[0][0] == 0 - - # Last chunk should end at audio length - assert chunks[-1][1] == 10000 - - # Chunks should be contiguous - for i in range(len(chunks) - 1): - assert chunks[i][1] == chunks[i + 1][0] - - -class TestTranscribeFinallyCleanup: - @staticmethod - def _buggy_transcribe(): - try: - raise RuntimeError("audio decode failed") - tmp_wav = None # noqa: F841 (unreachable, matches old code) - finally: - if tmp_wav: # noqa: F821 (deliberately exercises the bug) - pass - - @staticmethod - def _fixed_transcribe(): - tmp_wav = None - try: - raise RuntimeError("audio decode failed") - finally: - if tmp_wav: - pass - - def test_buggy_pattern_masks_original_error(self): - with pytest.raises(UnboundLocalError): - self._buggy_transcribe() - - def test_fixed_pattern_propagates_original_error(self): - with pytest.raises(RuntimeError, match="audio decode failed"): - self._fixed_transcribe() diff --git a/tests/unit/services/workers/parsers/legacy_loaders/test_base_loader.py b/tests/unit/services/workers/parsers/legacy_loaders/test_base_loader.py deleted file mode 100644 index 3ab2c342e..000000000 --- a/tests/unit/services/workers/parsers/legacy_loaders/test_base_loader.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Unit tests for BaseLoader image conversion utilities.""" - -from unittest.mock import MagicMock - -import pytest -from PIL import Image -from services.workers.parsers.legacy_loaders.base import BaseLoader, ensure_png_compatible_mode - - -class ConcreteLoader(BaseLoader): - """Minimal concrete subclass for testing.""" - - def __init__(self): - # Skip BaseLoader.__init__ which needs config/VLM - self.image_captioning = True - - async def aload_document(self, file_path, metadata=None, save_markdown=False): - pass - - -class TestPilImageToBase64: - def setup_method(self): - self.loader = ConcreteLoader() - - def test_rgb_image(self): - img = Image.new("RGB", (100, 100), "red") - result = self.loader._pil_image_to_base64(img) - assert isinstance(result, str) - assert len(result) > 0 - - def test_cmyk_image_converted(self): - img = Image.new("CMYK", (100, 100), (0, 0, 0, 0)) - result = self.loader._pil_image_to_base64(img) - assert isinstance(result, str) - assert len(result) > 0 - - def test_rgba_image(self): - img = Image.new("RGBA", (100, 100), (255, 0, 0, 128)) - result = self.loader._pil_image_to_base64(img) - assert isinstance(result, str) - - def test_palette_image(self): - img = Image.new("P", (100, 100)) - result = self.loader._pil_image_to_base64(img) - assert isinstance(result, str) - - def test_unsaveable_image_returns_empty(self): - """Images that can't be converted at all return empty string.""" - img = MagicMock(spec=Image.Image) - img.mode = "UNKNOWN" - img.convert.side_effect = Exception("Cannot convert") - img.save.side_effect = Exception("Cannot save") - result = self.loader._pil_image_to_base64(img) - assert result == "" - - -class TestGetImageDescription: - def setup_method(self): - self.loader = ConcreteLoader() - - @pytest.mark.asyncio - async def test_small_image_skipped(self): - """Images below minimum pixel threshold should be skipped without calling VLM.""" - small_img = Image.new("RGB", (10, 10), "red") - result = await self.loader.get_image_description(small_img) - assert "Image too small for captioning" in result - - def test_min_image_pixels_threshold(self): - """Verify the threshold constant is set correctly.""" - assert BaseLoader.MIN_IMAGE_PIXELS == 784 - - -class TestEnsurePngCompatibleMode: - def test_cmyk_to_rgb(self): - img = Image.new("CMYK", (10, 10)) - result = ensure_png_compatible_mode(img) - assert result.mode == "RGB" - - def test_palette_to_rgba(self): - img = Image.new("P", (10, 10)) - result = ensure_png_compatible_mode(img) - assert result.mode == "RGBA" - - def test_rgb_unchanged(self): - img = Image.new("RGB", (10, 10)) - result = ensure_png_compatible_mode(img) - assert result.mode == "RGB" - - def test_rgba_unchanged(self): - img = Image.new("RGBA", (10, 10)) - result = ensure_png_compatible_mode(img) - assert result.mode == "RGBA" diff --git a/tests/unit/services/workers/parsers/legacy_loaders/test_customdocloader.py b/tests/unit/services/workers/parsers/legacy_loaders/test_customdocloader.py deleted file mode 100644 index 2bd6f82f9..000000000 --- a/tests/unit/services/workers/parsers/legacy_loaders/test_customdocloader.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Regression test for CustomDocLoader page accumulation (#376). - -The previous loop body used ``s = ...`` instead of ``s += ...``, so only -the final page survived. This test confirms every page's content is now -in the returned ``Document``. -""" - -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest -from langchain_core.documents.base import Document as LCDocument - - -@pytest.mark.asyncio -async def test_customdocloader_accumulates_all_pages(tmp_path): - from services.workers.parsers.legacy_loaders.CustomDocLoader import CustomDocLoader - - fake_pages = [ - LCDocument(page_content="page-one"), - LCDocument(page_content="page-two"), - LCDocument(page_content="page-three"), - ] - fake_loader_instance = MagicMock() - fake_loader_instance.aload = AsyncMock(return_value=fake_pages) - fake_loader_cls = MagicMock(return_value=fake_loader_instance) - - file_path = tmp_path / "stub.docx" - file_path.write_text("ignored") - - with patch.dict(CustomDocLoader.doc_loaders, {".docx": fake_loader_cls}, clear=True): - # BaseLoader.__init__ pulls a config; we bypass it with object.__new__ - loader = object.__new__(CustomDocLoader) - result = await loader.aload_document(str(file_path), metadata={"src": "x"}) - - assert "page-one" in result.page_content - assert "page-two" in result.page_content - assert "page-three" in result.page_content - assert "[PAGE_1]" in result.page_content - assert "[PAGE_2]" in result.page_content - assert "[PAGE_3]" in result.page_content diff --git a/tests/unit/services/workers/parsers/legacy_loaders/test_doc_loader.py b/tests/unit/services/workers/parsers/legacy_loaders/test_doc_loader.py deleted file mode 100644 index aa87a7f0f..000000000 --- a/tests/unit/services/workers/parsers/legacy_loaders/test_doc_loader.py +++ /dev/null @@ -1,146 +0,0 @@ -""" -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. -""" - -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest -from core.config.endpoints import VLMConfig -from core.config.indexation import LoaderConfig -from core.models.document import ProcessedDocument, TextBlock -from langchain_core.documents.base import Document as LCDocument - - -@pytest.fixture -def mock_config(): - """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) - return config - - -@pytest.fixture -def metadata(): - return {"file_id": "test-file-id", "partition": "test-partition"} - - -_PATCHES = [ - patch("services.workers.parsers.legacy_loaders.doc.DocParser"), - patch("services.workers.parsers.legacy_loaders.base.ChatOpenAI"), - patch("services.workers.parsers.legacy_loaders.base.load_config"), -] - - -def _start_patches(mock_config): - mocks = [p.start() for p in _PATCHES] - _mock_doc_parser_cls, _mock_chat, mock_load_config = mocks - mock_load_config.return_value = mock_config - - -def _stop_patches(): - for p in _PATCHES: - try: - p.stop() - except RuntimeError: - pass - - -@pytest.fixture(autouse=True) -def _patch_cleanup(): - yield - _stop_patches() - - -def _make_loader(mock_config): - from services.workers.parsers.legacy_loaders.doc import DocLoader - - return DocLoader(config=mock_config) - - -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, - ) - - -class TestDocLoaderShim: - """Shim-level integration: ``DocParser`` ↔ langchain ``Document`` ↔ ``BaseLoader``.""" - - @pytest.mark.asyncio - 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")) - - file_path = tmp_path / "x.doc" - file_path.write_bytes(b"\xd0\xcf\x11\xe0fake-doc") - - result = await loader.aload_document(str(file_path), metadata) - - assert isinstance(result, LCDocument) - assert result.page_content == "converted markdown" - assert result.metadata == metadata - - 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_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_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 = tmp_path / "x.doc" - file_path.write_bytes(b"x") - - with patch.object(loader, "save_content") as mock_save: - 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_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/tests/unit/services/workers/parsers/legacy_loaders/test_docx_loader.py b/tests/unit/services/workers/parsers/legacy_loaders/test_docx_loader.py deleted file mode 100644 index 362522a7d..000000000 --- a/tests/unit/services/workers/parsers/legacy_loaders/test_docx_loader.py +++ /dev/null @@ -1,213 +0,0 @@ -""" -Unit tests for DocxLoader.get_images_from_zip image extraction. - -Tests validate that unsupported media formats (EMF, WMF, OLE objects, etc.) -are gracefully skipped instead of crashing the entire DOCX ingestion pipeline. -""" - -import tempfile -import zipfile -from io import BytesIO -from pathlib import Path - -from PIL import Image -from services.workers.parsers.legacy_loaders.docx import DocxLoader, convert_to_png_image - - -def _create_png_bytes(width=10, height=10, color="red"): - """Create minimal PNG image bytes.""" - img = Image.new("RGBA", (width, height), color) - buf = BytesIO() - img.save(buf, format="PNG") - return buf.getvalue() - - -def _create_fake_docx(media_files: dict[str, bytes]) -> Path: - """Create a minimal .docx (zip) with given word/media/ entries. - - Args: - media_files: mapping of filename (e.g. "image1.png") to raw bytes. - - Returns: - Path to the temporary .docx file. - """ - 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 TestGetImagesFromZip: - """Test DocxLoader.get_images_from_zip with various media contents.""" - - def _make_loader(self): - """Create a DocxLoader without Hydra config (only get_images_from_zip is used).""" - loader = object.__new__(DocxLoader) - return loader - - def test_valid_images_returned_in_order(self): - """Valid PNG images are extracted and reordered by their number.""" - red = _create_png_bytes(color="red") - blue = _create_png_bytes(color="blue") - # Insert out of order: image2 before image1 - docx_path = _create_fake_docx({"image2.png": blue, "image1.png": red}) - - loader = self._make_loader() - images = loader.get_images_from_zip(docx_path) - - assert len(images) == 2 - # image1 (red) should come first - assert images[0] is not None - assert images[1] is not None - - def test_unsupported_format_skipped(self): - """Unsupported formats (EMF, WMF, etc.) are skipped, valid images preserved.""" - valid_png = _create_png_bytes() - fake_emf = b"\x01\x00\x00\x00EMF_GARBAGE_DATA" - - docx_path = _create_fake_docx( - { - "image1.png": valid_png, - "image2.emf": fake_emf, - "image3.png": valid_png, - } - ) - - loader = self._make_loader() - images = loader.get_images_from_zip(docx_path) - - # image2.emf is skipped; images list sized to max order (3) - # with None at position 2 (index 1) for the skipped EMF - assert len(images) == 3 - assert images[0] is not None # image1.png - assert images[1] is None # image2.emf was skipped - assert images[2] is not None # image3.png - - def test_non_image_media_files_skipped(self): - """Files like oleObject1.bin are skipped via try/except (not valid images).""" - valid_png = _create_png_bytes() - docx_path = _create_fake_docx( - { - "image1.png": valid_png, - "oleObject1.bin": b"OLE_DATA", - "hdphoto1.wdp": b"WDP_DATA", - } - ) - - loader = self._make_loader() - images = loader.get_images_from_zip(docx_path) - - # oleObject1.bin and hdphoto1.wdp fail Image.open() or order parsing - # Only image1.png succeeds - assert sum(1 for img in images if img is not None) == 1 - - def test_all_unsupported_returns_empty(self): - """When all image files are unsupported, returns empty list.""" - docx_path = _create_fake_docx( - { - "image1.emf": b"EMF_DATA", - "image2.wmf": b"WMF_DATA", - } - ) - - loader = self._make_loader() - images = loader.get_images_from_zip(docx_path) - - assert images == [] - - def test_no_media_returns_empty(self): - """DOCX with no word/media/ files returns empty list.""" - tmp = tempfile.NamedTemporaryFile(suffix=".docx", delete=False) - with zipfile.ZipFile(tmp, "w") as zf: - zf.writestr("word/document.xml", "") - - loader = self._make_loader() - images = loader.get_images_from_zip(Path(tmp.name)) - - assert images == [] - - -class TestCaptionBackslashInjection: - """Regression tests for #389. - - The old code passed VLM captions directly to re.sub as the replacement - string. Captions containing \\1 or \\g raised re.error. The fix - switched to str.replace(), which treats the replacement as a literal. - A leftover caption.replace("\\", "/") that mangled backslashes in captions - was also removed — this class verifies captions are now injected verbatim. - """ - - def _inject(self, ref: str, caption: str, content: str) -> str: - """Replicate the exact substitution used in DocxLoader and PPTXLoader.""" - return content.replace(ref, caption) - - def test_backreference_sequence_injected_verbatim(self): - r"""Caption containing \1 must appear as-is in the output.""" - ref = "![](pptx-image-0)" - caption = r"Figure with backreference \1 inside" - content = f"Slide text {ref} more text" - result = self._inject(ref, caption, content) - assert r"\1" in result - assert caption in result - - def test_named_group_reference_injected_verbatim(self): - r"""Caption containing \g must appear as-is in the output.""" - ref = "![](pptx-image-1)" - caption = r"Caption with \g group ref" - content = f"Before {ref} after" - result = self._inject(ref, caption, content) - assert r"\g" in result - assert caption in result - - def test_windows_path_in_caption_preserved(self): - r"""Caption containing C:\Users\name must not be mangled to C:/Users/name.""" - ref = "![](pptx-image-2)" - caption = r"Screenshot from C:\Users\alice\Desktop" - content = f"Intro {ref} end" - result = self._inject(ref, caption, content) - assert r"C:\Users\alice\Desktop" in result - assert r"C:/Users/alice/Desktop" not in result - - def test_no_exception_for_any_backslash_sequence(self): - """None of these captions should raise when injected into content.""" - ref = "![](img)" - problematic_captions = [ - r"\1", - r"\g", - r"\0", - r"\99", - r"C:\Windows\System32", - r"path\\to\\file", - "normal caption without backslashes", - ] - for caption in problematic_captions: - result = self._inject(ref, caption, f"text {ref} end") - assert caption in result, f"Caption not injected verbatim: {caption!r}" - - -class TestConvertToPngImage: - def test_rgb_image(self): - img = Image.new("RGB", (50, 50), "red") - result = convert_to_png_image(img) - assert result is not None - assert result.mode == "RGBA" - - def test_cmyk_image_converted(self): - """CMYK images should be converted to RGB then to PNG.""" - cmyk_img = Image.new("CMYK", (50, 50), (0, 0, 0, 0)) - result = convert_to_png_image(cmyk_img) - assert result is not None - assert result.mode == "RGBA" - - def test_rgba_image(self): - img = Image.new("RGBA", (50, 50), (255, 0, 0, 128)) - result = convert_to_png_image(img) - assert result is not None - assert result.mode == "RGBA" - - def test_palette_image(self): - img = Image.new("P", (50, 50)) - result = convert_to_png_image(img) - assert result is not None - assert result.mode == "RGBA" diff --git a/tests/unit/services/workers/parsers/legacy_loaders/test_eml_recursion.py b/tests/unit/services/workers/parsers/legacy_loaders/test_eml_recursion.py deleted file mode 100644 index d69b9c641..000000000 --- a/tests/unit/services/workers/parsers/legacy_loaders/test_eml_recursion.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Regression test for #364 — EmlLoader nested .eml recursion is capped. - -The loader recursed into .eml attachments via ``aload_document`` with no -depth counter, so .eml files nested inside one another recursed without -bound. The fix threads a ``_eml_recursion_depth`` keyword and stops -descending once the operator-tunable ``loader.eml_max_recursion_depth`` -is reached. -""" - -import base64 - -import pytest - - -def _make_eml_with_attached_eml(inner_bytes: bytes, subject: str = "outer") -> bytes: - """Build a multipart/mixed .eml whose attachment (Content-Disposition: - attachment; filename=nested.eml) is the given inner .eml bytes. - """ - encoded = base64.b64encode(inner_bytes).decode("ascii") - boundary = "============TEST_BOUNDARY============" - return ( - f"Subject: {subject}\r\n" - "From: a@example.com\r\n" - "To: b@example.com\r\n" - "MIME-Version: 1.0\r\n" - f'Content-Type: multipart/mixed; boundary="{boundary}"\r\n' - "\r\n" - f"--{boundary}\r\n" - "Content-Type: text/plain; charset=utf-8\r\n" - "\r\n" - "outer body\r\n" - f"--{boundary}\r\n" - 'Content-Type: application/octet-stream; name="nested.eml"\r\n' - 'Content-Disposition: attachment; filename="nested.eml"\r\n' - "Content-Transfer-Encoding: base64\r\n" - "\r\n" - f"{encoded}\r\n" - f"--{boundary}--\r\n" - ).encode("ascii") - - -def _make_leaf_eml() -> bytes: - return b"Subject: leaf\r\nFrom: a@example.com\r\nTo: b@example.com\r\n\r\nleaf body\r\n" - - -@pytest.mark.asyncio -async def test_eml_recursion_caps_at_max_depth(tmp_path): - from services.workers.parsers.legacy_loaders.eml_loader import EmlLoader - - # A single outer .eml whose attachment is another .eml. We seed the - # call at depth = cap - 1, so processing the outer's attachment (which - # would push us past the cap) trips the guard. This proves the depth - # cap stops the descent before the nested .eml is loaded again. - eml_path = tmp_path / "nested.eml" - eml_path.write_bytes(_make_eml_with_attached_eml(_make_leaf_eml())) - - loader = object.__new__(EmlLoader) - loader.loader_classes = {".eml": EmlLoader} - loader.kwargs = {} - loader.max_eml_recursion_depth = 5 - - seeded_depth = loader.max_eml_recursion_depth - 1 - doc = await loader.aload_document(str(eml_path), _eml_recursion_depth=seeded_depth) - assert "recursion depth limit" in doc.page_content - # The body of the outer .eml must still be retained - assert "outer body" in doc.page_content - - -@pytest.mark.asyncio -async def test_eml_below_cap_does_not_skip(tmp_path): - """At depth 0 the guard does not fire — attachments are still attempted.""" - from services.workers.parsers.legacy_loaders.eml_loader import EmlLoader - - eml_path = tmp_path / "nested.eml" - eml_path.write_bytes(_make_eml_with_attached_eml(_make_leaf_eml())) - - loader = object.__new__(EmlLoader) - loader.loader_classes = {".eml": EmlLoader} - loader.kwargs = {} - loader.max_eml_recursion_depth = 5 - - doc = await loader.aload_document(str(eml_path), _eml_recursion_depth=0) - # The guard message should NOT appear at low depth - assert "recursion depth limit" not in doc.page_content - - -def test_recursion_cap_lives_in_loader_config(): - """The cap must come from the loader config (operator-tunable), not be - hardcoded in the loader. Assert the contract: a positive integer.""" - from core.config import load_config - - cap = load_config().loader.eml_max_recursion_depth - assert isinstance(cap, int) - assert cap > 0 diff --git a/tests/unit/services/workers/parsers/test_doc_serializer_bridge.py b/tests/unit/services/workers/parsers/test_doc_serializer_bridge.py deleted file mode 100644 index 9fbb2ebed..000000000 --- a/tests/unit/services/workers/parsers/test_doc_serializer_bridge.py +++ /dev/null @@ -1,58 +0,0 @@ -from __future__ import annotations - -from typing import Any - -import pytest -from core.config.root import Settings -from core.models.document import Document, DocumentType -from langchain_core.documents.base import Document as LangChainDocument -from services.workers.parsers.doc_serializer_bridge import ( - INDEXATION_CONFIG_METADATA_KEY, - DocSerializerBridgeParser, -) - - -class _FakeLoader: - seen_config: Any = None - seen_metadata: dict[str, Any] | None = None - - def __init__(self, *, config: Any) -> None: - type(self).seen_config = config - - async def aload_document(self, *, file_path: str, metadata: dict | None = None, save_markdown: bool = False): - type(self).seen_metadata = dict(metadata or {}) - return LangChainDocument(page_content="hello", metadata=metadata or {}) - - -@pytest.mark.asyncio -async def test_bridge_disables_legacy_captioning_from_indexation_config(monkeypatch): - """Per-file Phase 14 config disables legacy loader image captioning.""" - - def fake_loader_classes(config): - return {".txt": _FakeLoader} - - monkeypatch.setattr("services.workers.parsers.legacy_loaders.get_loader_classes", fake_loader_classes) - - parser = DocSerializerBridgeParser( - Settings( - loader={ - "image_captioning": True, - "image_captioning_url": True, - } - ) - ) - document = Document( - filename="note.txt", - content_type=DocumentType.TEXT, - raw_bytes=b"hello", - metadata={ - "source": "note.txt", - INDEXATION_CONFIG_METADATA_KEY: {"enable_image_captioning": False}, - }, - ) - - await parser.parse(document) - - assert _FakeLoader.seen_config.loader.image_captioning is False - assert _FakeLoader.seen_config.loader.image_captioning_url is False - assert _FakeLoader.seen_metadata == {"source": "note.txt"} diff --git a/tests/unit/services/workers/parsers/test_parser_dispatcher.py b/tests/unit/services/workers/parsers/test_parser_dispatcher.py new file mode 100644 index 000000000..1e0be51ac --- /dev/null +++ b/tests/unit/services/workers/parsers/test_parser_dispatcher.py @@ -0,0 +1,123 @@ +"""Unit tests for the content-type → parser dispatch.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock +from services.workers.parsers.parser_dispatcher import ( + ParserDispatcher, + build_caption_vlm, +) + + +def _config( + *, pdf="MarkerLoader", audio="LocalWhisperLoader", image_captioning=True, vlm_base_url="" +) -> SimpleNamespace: + file_loaders = SimpleNamespace( + pdf=pdf, + mp3=audio, + wav=audio, + flac=audio, + mp4=audio, + ) + loader = SimpleNamespace(file_loaders=file_loaders, image_captioning=image_captioning) + vlm = SimpleNamespace(base_url=vlm_base_url, model="m", api_key="k", timeout=60) + return SimpleNamespace(loader=loader, vlm=vlm) + + +class _FakeParser: + def __init__(self) -> None: + self.seen: Document | None = None + + def supported_types(self) -> list[str]: + return [] + + async def parse(self, document: Document) -> ProcessedDocument: + self.seen = document + return ProcessedDocument(document_id=document.id, text_blocks=[TextBlock(text="ok")]) + + +@pytest.mark.parametrize( + ("filename", "content_type", "expected_backend"), + [ + ("a.txt", DocumentType.TEXT, "text"), + ("a.md", DocumentType.MARKDOWN, "markdown"), + ("a.html", DocumentType.HTML, "html"), + ("a.docx", DocumentType.DOCX, "docx"), + ("a.doc", DocumentType.DOC, "doc"), + ("a.pptx", DocumentType.PPTX, "pptx"), + ("a.eml", DocumentType.EML, "eml"), + ("a.png", DocumentType.IMAGE, "image"), + ("a.svg", DocumentType.IMAGE, "image"), + ("a.gif", DocumentType.IMAGE, "image"), + ("a.webp", DocumentType.IMAGE, "image"), + ("a.bmp", DocumentType.IMAGE, "image"), + ("a.pdf", DocumentType.PDF, "marker"), + ("a.mp3", DocumentType.AUDIO, "local_whisper"), + ("a.mp4", DocumentType.VIDEO, "local_whisper"), + ], +) +def test_resolve_backend(filename: str, content_type: DocumentType, expected_backend: str) -> None: + disp = ParserDispatcher(_config()) + from services.workers.parsers.parser_dispatcher import _suffix + + assert disp._resolve_backend(content_type, _suffix(filename)) == expected_backend + + +def test_resolve_pdf_backend_variants() -> None: + assert ParserDispatcher(_config(pdf="DoclingLoader"))._resolve_pdf_backend() == "docling" + assert ParserDispatcher(_config(pdf="PyMuPDFLoader"))._resolve_pdf_backend() == "pymupdf" + assert ParserDispatcher(_config(pdf="DotsOCRLoader"))._resolve_pdf_backend() == "pdf_client" + + +def test_resolve_audio_backend_openai() -> None: + disp = ParserDispatcher(_config(audio="OpenAIAudioLoader")) + assert disp._resolve_audio_backend("mp3") == "audio_client" + + +def test_unsupported_pdf_config_raises() -> None: + with pytest.raises(ValueError, match="Unsupported PDF loader"): + ParserDispatcher(_config(pdf="NopeLoader"))._resolve_pdf_backend() + + +@pytest.mark.asyncio +async def test_parse_dispatches_to_cached_backend() -> None: + disp = ParserDispatcher(_config()) + fake = _FakeParser() + disp._by_name["marker"] = fake # pre-seed so no real backend is built + + document = Document(filename="report.pdf", content_type=DocumentType.PDF, raw_bytes=b"%PDF-1.4") + result = await disp.parse(document) + + assert fake.seen is document + assert result.text_blocks[0].text == "ok" + + +def test_build_caption_vlm_requires_endpoint() -> None: + # No VLM endpoint configured -> unavailable, regardless of the captioning flag. + assert build_caption_vlm(_config(image_captioning=True, vlm_base_url="")) is None + assert build_caption_vlm(_config(image_captioning=False, vlm_base_url="")) is None + + +def test_build_caption_vlm_available_when_endpoint_set_even_if_globally_off() -> None: + # Availability is decoupled from the captioning policy: an endpoint is enough + # to build the VLM. Standalone-image captioning relies on this (the policy + # gate lives in the pipeline, not here). + assert build_caption_vlm(_config(image_captioning=False, vlm_base_url="http://vlm:8000/v1")) is not None + + +def test_build_eml_wires_nested_email_parser_with_depth_limit(monkeypatch: pytest.MonkeyPatch) -> None: + disp = ParserDispatcher(_config()) + fallback_parser = _FakeParser() + monkeypatch.setattr(disp, "_get", lambda name: fallback_parser) + + parser = disp._build_eml() + + assert parser._attachment_parsers["txt"] is fallback_parser + assert "eml" in parser._attachment_parsers + nested_1 = parser._attachment_parsers["eml"] + nested_2 = nested_1._attachment_parsers["eml"] + nested_3 = nested_2._attachment_parsers["eml"] + assert "eml" not in nested_3._attachment_parsers diff --git a/tests/unit/services/workers/test_indexer_pool.py b/tests/unit/services/workers/test_indexer_pool.py index e4e11d475..a51aacd64 100644 --- a/tests/unit/services/workers/test_indexer_pool.py +++ b/tests/unit/services/workers/test_indexer_pool.py @@ -236,7 +236,7 @@ def test_indexer_pool_wires_contextualizer_factory(monkeypatch: pytest.MonkeyPat import services.storage.milvus_store as milvus_store import services.storage.postgres_store as postgres_store import services.workers.indexer_pool as module - import services.workers.parsers.doc_serializer_bridge as parser_bridge + import services.workers.parsers.parser_dispatcher as parser_dispatcher import services.workers.pipeline_builder as pipeline_builder captured = {} @@ -257,6 +257,7 @@ def model_copy(self, *, update): batch_size=32, embed_concurrency=2, ), + loader=SimpleNamespace(image_captioning=True), vectordb=SimpleNamespace(collection_name="vdb_test"), rdb=RDBConfig(), ) @@ -281,7 +282,8 @@ def fake_build_pipeline(**kwargs): monkeypatch.setattr(core.embeddings.embedder_registry, "create", lambda *args, **kwargs: object()) monkeypatch.setattr(milvus_store, "MilvusVectorStore", lambda _cfg: object()) monkeypatch.setattr(postgres_store, "PostgresStore", lambda *args, **kwargs: Store()) - monkeypatch.setattr(parser_bridge, "DocSerializerBridgeParser", lambda **kwargs: object()) + monkeypatch.setattr(parser_dispatcher, "build_parser_dispatcher", lambda _cfg: object()) + monkeypatch.setattr(parser_dispatcher, "build_caption_vlm", lambda _cfg: object()) monkeypatch.setattr(pipeline_builder, "build_indexing_pipeline", fake_build_pipeline) actor_calls = [] diff --git a/tests/unit/services/workers/test_indexer_worker.py b/tests/unit/services/workers/test_indexer_worker.py index 5811d7a98..3b7220ec8 100644 --- a/tests/unit/services/workers/test_indexer_worker.py +++ b/tests/unit/services/workers/test_indexer_worker.py @@ -9,7 +9,6 @@ from core.models.chunk import Chunk from core.models.document import Document, DocumentType, ProcessedDocument, TextBlock from services.workers.indexer_actor import IndexerWorker, _load_document -from services.workers.parsers.doc_serializer_bridge import INDEXATION_CONFIG_METADATA_KEY from services.workers.pipeline_builder import build_indexing_pipeline # --------------------------------------------------------------------------- @@ -112,33 +111,58 @@ async def bulk_insert(self, tags: list[dict]) -> int: # --------------------------------------------------------------------------- -def test_load_document_reads_bytes_and_detects_type(tmp_path: Path) -> None: - p = tmp_path / "report.pdf" +@pytest.mark.asyncio +async def test_load_document_reads_bytes_and_detects_type_from_original_filename(tmp_path: Path) -> None: + p = tmp_path / "upload-without-extension" p.write_bytes(b"%PDF-1.4") - doc = _load_document(str(p), {"file_id": "fid-1"}, "tenant-a") + doc = await _load_document( + str(p), + {"file_id": "fid-1", "filename": "safe-name", "original_filename": "report.pdf"}, + "tenant-a", + ) assert doc.raw_bytes == b"%PDF-1.4" assert doc.content_type == DocumentType.PDF assert doc.partition == "tenant-a" - assert doc.filename == "fid-1" + assert doc.filename == "report.pdf" + # Document.id must be the file_id (not a random uuid): the chunker derives + # Chunk.document_id / file_id from ProcessedDocument.document_id == document.id. + assert doc.id == "fid-1" -def test_load_document_falls_back_to_filename_when_no_file_id(tmp_path: Path) -> None: +@pytest.mark.asyncio +async def test_load_document_requires_file_id(tmp_path: Path) -> None: p = tmp_path / "note.txt" p.write_bytes(b"hi") - doc = _load_document(str(p), {}, "p") - assert doc.filename == "note.txt" + # file_id is force-set upstream by IndexingService._build_metadata; if it is + # ever missing we fail loudly rather than persist chunks under a bad id. + with pytest.raises(ValueError, match="file_id"): + await _load_document(str(p), {}, "p") -def test_load_document_attaches_internal_indexation_config(tmp_path: Path) -> None: +@pytest.mark.asyncio +async def test_load_document_does_not_leak_internal_keys_into_metadata(tmp_path: Path) -> None: p = tmp_path / "note.txt" p.write_bytes(b"hi") - config = {"enable_image_captioning": False} - doc = _load_document(str(p), {"source": "note.txt"}, "p", indexation_config=config) + doc = await _load_document(str(p), {"file_id": "fid", "source": "note.txt"}, "p") + + # indexation_config reaches the pipeline via row["indexation_config"], never + # the document metadata, so it cannot leak into chunk metadata. + assert doc.metadata == {"file_id": "fid", "source": "note.txt"} + assert all(not key.startswith("_openrag") for key in doc.metadata) + + +@pytest.mark.asyncio +async def test_load_document_falls_back_to_stored_path_name(tmp_path: Path) -> None: + p = tmp_path / "audio.flac" + p.write_bytes(b"flac") + + doc = await _load_document(str(p), {"file_id": "fid"}, "p") - assert doc.metadata[INDEXATION_CONFIG_METADATA_KEY] == config + assert doc.filename == "audio.flac" + assert doc.content_type == DocumentType.AUDIO # --------------------------------------------------------------------------- @@ -196,7 +220,7 @@ def supported_types(self) -> list[str]: await worker.process_file( task_id="t2", path=str(path), - metadata={}, + metadata={"file_id": "f1"}, partition="p", ) @@ -218,7 +242,7 @@ async def test_process_file_missing_path_raises_and_sets_failed() -> None: await worker.process_file( task_id="t3", path="/nonexistent/file.txt", - metadata={}, + metadata={"file_id": "f1"}, partition="p", ) @@ -231,14 +255,23 @@ async def test_process_file_passes_partition_and_filename_to_row(tmp_path: Path) path.write_bytes(b"hello") seen_partitions: list[str] = [] + seen_documents: list[Document] = [] class TrackingChunker: def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: seen_partitions.append(partition) return [Chunk(id="c1", text="hello", partition=partition)] + class TrackingParser: + async def parse(self, document: Document) -> ProcessedDocument: + seen_documents.append(document) + return ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="hello")]) + + def supported_types(self) -> list[str]: + return [DocumentType.TEXT.value] + pipeline = build_indexing_pipeline( - parser=FakeParser(ProcessedDocument(document_id="d1", text_blocks=[TextBlock(text="hello")])), + parser=TrackingParser(), chunker=TrackingChunker(), embedder=FakeEmbedder(), vector_store=FakeVectorStore(), @@ -248,11 +281,12 @@ def chunk(self, document: ProcessedDocument, partition: str = "default") -> list await worker.process_file( task_id="t4", path=str(path), - metadata={"file_id": "fid"}, + metadata={"file_id": "fid", "original_filename": "original-note.txt"}, partition="tenant-b", ) assert seen_partitions == ["tenant-b"] + assert seen_documents[0].filename == "original-note.txt" @pytest.mark.asyncio diff --git a/tests/unit/services/workers/test_pipeline_builder.py b/tests/unit/services/workers/test_pipeline_builder.py index 40a75f676..89747e450 100644 --- a/tests/unit/services/workers/test_pipeline_builder.py +++ b/tests/unit/services/workers/test_pipeline_builder.py @@ -247,6 +247,57 @@ async def test_pipeline_row_indexation_config_selects_components(): assert row["chunks"][0].embedding == [0.5] +@pytest.mark.asyncio +async def test_pipeline_always_captions_standalone_image_even_when_globally_off(): + # A standalone image file's caption is its only text content, so it is + # captioned even when global image captioning is off (legacy parity). + document = Document(filename="logo.png", content_type=DocumentType.IMAGE, raw_bytes=b"img") + processed = ProcessedDocument( + document_id=document.id, + text_blocks=[], + images=[ImageBlock(image_bytes=b"png")], + ) + vlm = FakeVLM() + pipeline = build_indexing_pipeline( + parser=FakeParser(processed), + chunker=FakeChunker([Chunk(id="c1", text="caption", partition="tenant-a")]), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + vlm=vlm, + image_captioning=False, + ) + row = {"document": document, "partition": "tenant-a", "filename": "logo.png"} + + await pipeline.run(row) + + assert vlm.calls == [b"png"] + + +@pytest.mark.asyncio +async def test_pipeline_skips_embedded_image_caption_when_globally_off(): + # Images embedded in a non-image document stay gated by the global flag. + document = Document(filename="note.txt", text="hello", partition="tenant-a") + processed = ProcessedDocument( + document_id=document.id, + text_blocks=[TextBlock(text="hello")], + images=[ImageBlock(image_bytes=b"png")], + ) + vlm = FakeVLM() + pipeline = build_indexing_pipeline( + parser=FakeParser(processed), + chunker=FakeChunker([Chunk(id="c1", text="hello", partition="tenant-a")]), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + vlm=vlm, + image_captioning=False, + ) + row = {"document": document, "partition": "tenant-a", "filename": "note.txt"} + + await pipeline.run(row) + + assert vlm.calls == [] + + @pytest.mark.asyncio async def test_pipeline_indexation_config_disables_topic_tagging(): document = Document(filename="note.txt", text="hello", partition="tenant-a")