nv-ingest-api 2025.4.16.dev20250416__py3-none-any.whl → 2025.4.17.dev20250417__py3-none-any.whl

nv_ingest_api/util/metadata/aggregators.py

@@ -0,0 +1,469 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import base64
+import io
+import uuid
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+import pandas as pd
+import pypdfium2 as pdfium
+from PIL import Image
+from pypdfium2 import PdfImage
+
+from nv_ingest_api.internal.enums.common import ContentDescriptionEnum, DocumentTypeEnum
+from nv_ingest_api.internal.enums.common import ContentTypeEnum
+from nv_ingest_api.internal.schemas.meta.metadata_schema import NearbyObjectsSchema
+from nv_ingest_api.internal.enums.common import TableFormatEnum
+from nv_ingest_api.internal.schemas.meta.metadata_schema import validate_metadata
+from nv_ingest_api.util.converters import datetools
+from nv_ingest_api.util.detectors.language import detect_language
+from nv_ingest_api.util.exception_handlers.pdf import pdfium_exception_handler
+
+
+@dataclass
+class CroppedImageWithContent:
+    content: str
+    image: str
+    bbox: Tuple[int, int, int, int]
+    max_width: int
+    max_height: int
+    type_string: str
+    content_format: str = ""
+
+
+@dataclass
+class LatexTable:
+    latex: pd.DataFrame
+    bbox: Tuple[int, int, int, int]
+    max_width: int
+    max_height: int
+
+
+@dataclass
+class Base64Image:
+    image: str
+    bbox: Tuple[int, int, int, int]
+    width: int
+    height: int
+    max_width: int
+    max_height: int
+
+
+@dataclass
+class PDFMetadata:
+    """
+    A data object to store metadata information extracted from a PDF document.
+    """
+
+    page_count: int
+    filename: str
+    last_modified: str
+    date_created: str
+    keywords: List[str]
+    source_type: str = "PDF"
+
+
+def extract_pdf_metadata(doc: pdfium.PdfDocument, source_id: str) -> PDFMetadata:
+    """
+    Extracts metadata and relevant information from a PDF document.
+
+    Parameters
+    ----------
+    pdf_stream : bytes
+        The PDF document data as a byte stream.
+    source_id : str
+        The identifier for the source document, typically the filename.
+
+    Returns
+    -------
+    PDFMetadata
+        An object containing extracted metadata and information including:
+        - `page_count`: The total number of pages in the PDF.
+        - `filename`: The source filename or identifier.
+        - `last_modified`: The last modified date of the PDF document.
+        - `date_created`: The creation date of the PDF document.
+        - `keywords`: Keywords associated with the PDF document.
+        - `source_type`: The type/format of the source, e.g., "PDF".
+
+    Raises
+    ------
+    PdfiumError
+        If there is an issue processing the PDF document.
+    """
+    page_count: int = len(doc)
+    filename: str = source_id
+
+    # Extract document metadata
+    doc_meta = doc.get_metadata_dict()
+
+    # Extract and process the last modified date
+    last_modified: str = doc_meta.get("ModDate")
+    if last_modified in (None, ""):
+        last_modified = datetools.remove_tz(datetime.now()).isoformat()
+    else:
+        last_modified = datetools.datetimefrompdfmeta(last_modified)
+
+    # Extract and process the creation date
+    date_created: str = doc_meta.get("CreationDate")
+    if date_created in (None, ""):
+        date_created = datetools.remove_tz(datetime.now()).isoformat()
+    else:
+        date_created = datetools.datetimefrompdfmeta(date_created)
+
+    # Extract keywords, defaulting to an empty list if not found
+    keywords: List[str] = doc_meta.get("Keywords", [])
+
+    # Create the PDFMetadata object
+    metadata = PDFMetadata(
+        page_count=page_count,
+        filename=filename,
+        last_modified=last_modified,
+        date_created=date_created,
+        keywords=keywords,
+    )
+
+    return metadata
+
+
+def construct_text_metadata(
+    accumulated_text,
+    keywords,
+    page_idx,
+    block_idx,
+    line_idx,
+    span_idx,
+    page_count,
+    text_depth,
+    source_metadata,
+    base_unified_metadata,
+    delimiter=" ",
+    bbox_max_dimensions: Tuple[int, int] = (-1, -1),
+    nearby_objects: Optional[Dict[str, Any]] = None,
+):
+    extracted_text = delimiter.join(accumulated_text)
+
+    content_metadata = {
+        "type": ContentTypeEnum.TEXT,
+        "description": ContentDescriptionEnum.PDF_TEXT,
+        "page_number": page_idx,
+        "hierarchy": {
+            "page_count": page_count,
+            "page": page_idx,
+            "block": -1,
+            "line": -1,
+            "span": -1,
+            "nearby_objects": nearby_objects or NearbyObjectsSchema(),
+        },
+    }
+
+    language = detect_language(extracted_text)
+
+    # TODO(Devin) - Implement bounding box logic for text
+    bbox = (-1, -1, -1, -1)
+
+    text_metadata = {
+        "text_type": text_depth,
+        "summary": "",
+        "keywords": keywords,
+        "language": language,
+        "text_location": bbox,
+        "text_location_max_dimensions": bbox_max_dimensions,
+    }
+
+    ext_unified_metadata = base_unified_metadata.copy()
+
+    ext_unified_metadata.update(
+        {
+            "content": extracted_text,
+            "source_metadata": source_metadata,
+            "content_metadata": content_metadata,
+            "text_metadata": text_metadata,
+        }
+    )
+
+    validated_unified_metadata = validate_metadata(ext_unified_metadata)
+
+    return [ContentTypeEnum.TEXT, validated_unified_metadata.model_dump(), str(uuid.uuid4())]
+
+
+def construct_image_metadata_from_base64(
+    base64_image: str,
+    page_idx: int,
+    page_count: int,
+    source_metadata: Dict[str, Any],
+    base_unified_metadata: Dict[str, Any],
+) -> List[Any]:
+    """
+    Extracts image data from a base64-encoded image string, decodes the image to get
+    its dimensions and bounding box, and constructs metadata for the image.
+
+    Parameters
+    ----------
+    base64_image : str
+        A base64-encoded string representing the image.
+    page_idx : int
+        The index of the current page being processed.
+    page_count : int
+        The total number of pages in the PDF document.
+    source_metadata : Dict[str, Any]
+        Metadata related to the source of the PDF document.
+    base_unified_metadata : Dict[str, Any]
+        The base unified metadata structure to be updated with the extracted image information.
+
+    Returns
+    -------
+    List[Any]
+        A list containing the content type, validated metadata dictionary, and a UUID string.
+
+    Raises
+    ------
+    ValueError
+        If the image cannot be decoded from the base64 string.
+    """
+    # Decode the base64 image
+    try:
+        image_data = base64.b64decode(base64_image)
+        image = Image.open(io.BytesIO(image_data))
+    except Exception as e:
+        raise ValueError(f"Failed to decode image from base64: {e}")
+
+    # Extract image dimensions and bounding box
+    width, height = image.size
+    bbox = (0, 0, width, height)  # Assuming the full image as the bounding box
+
+    # Construct content metadata
+    content_metadata: Dict[str, Any] = {
+        "type": ContentTypeEnum.IMAGE,
+        "description": ContentDescriptionEnum.PDF_IMAGE,
+        "page_number": page_idx,
+        "hierarchy": {
+            "page_count": page_count,
+            "page": page_idx,
+            "block": -1,
+            "line": -1,
+            "span": -1,
+        },
+    }
+
+    # Construct image metadata
+    image_metadata: Dict[str, Any] = {
+        "image_type": DocumentTypeEnum.PNG,
+        "structured_image_type": ContentTypeEnum.UNKNOWN,
+        "caption": "",
+        "text": "",
+        "image_location": bbox,
+        "image_location_max_dimensions": (width, height),
+        "height": height,
+    }
+
+    # Update the unified metadata with the extracted image information
+    unified_metadata: Dict[str, Any] = base_unified_metadata.copy()
+    unified_metadata.update(
+        {
+            "content": base64_image,
+            "source_metadata": source_metadata,
+            "content_metadata": content_metadata,
+            "image_metadata": image_metadata,
+        }
+    )
+
+    # Validate and return the unified metadata
+    validated_unified_metadata = validate_metadata(unified_metadata)
+    return [ContentTypeEnum.IMAGE, validated_unified_metadata.model_dump(), str(uuid.uuid4())]
+
+
+def construct_image_metadata_from_pdf_image(
+    pdf_image: PdfImage,
+    page_idx: int,
+    page_count: int,
+    source_metadata: Dict[str, Any],
+    base_unified_metadata: Dict[str, Any],
+) -> List[Any]:
+    """
+    Extracts image data from a PdfImage object, converts it to a base64-encoded string,
+    and constructs metadata for the image.
+
+    Parameters
+    ----------
+    image_obj : PdfImage
+        The PdfImage object from which the image will be extracted.
+    page_idx : int
+        The index of the current page being processed.
+    page_count : int
+        The total number of pages in the PDF document.
+    source_metadata : dict
+        Metadata related to the source of the PDF document.
+    base_unified_metadata : dict
+        The base unified metadata structure to be updated with the extracted image information.
+
+    Returns
+    -------
+    List[Any]
+        A list containing the content type, validated metadata dictionary, and a UUID string.
+
+    Raises
+    ------
+    PdfiumError
+        If the image cannot be extracted due to an issue with the PdfImage object.
+        :param pdf_image:
+    """
+
+    # Construct content metadata
+    content_metadata: Dict[str, Any] = {
+        "type": ContentTypeEnum.IMAGE,
+        "description": ContentDescriptionEnum.PDF_IMAGE,
+        "page_number": page_idx,
+        "hierarchy": {
+            "page_count": page_count,
+            "page": page_idx,
+            "block": -1,
+            "line": -1,
+            "span": -1,
+        },
+    }
+
+    # Construct image metadata
+    image_metadata: Dict[str, Any] = {
+        "image_type": DocumentTypeEnum.PNG,
+        "structured_image_type": ContentTypeEnum.UNKNOWN,
+        "caption": "",
+        "text": "",
+        "image_location": pdf_image.bbox,
+        "image_location_max_dimensions": (max(pdf_image.max_width, 0), max(pdf_image.max_height, 0)),
+        "height": pdf_image.height,
+        "width": pdf_image.width,
+    }
+
+    # Update the unified metadata with the extracted image information
+    unified_metadata: Dict[str, Any] = base_unified_metadata.copy()
+    unified_metadata.update(
+        {
+            "content": pdf_image.image,
+            "source_metadata": source_metadata,
+            "content_metadata": content_metadata,
+            "image_metadata": image_metadata,
+        }
+    )
+
+    # Validate and return the unified metadata
+    validated_unified_metadata = validate_metadata(unified_metadata)
+    return [ContentTypeEnum.IMAGE, validated_unified_metadata.model_dump(), str(uuid.uuid4())]
+
+
+# TODO(Devin): Disambiguate tables and charts, create two distinct processing methods
+@pdfium_exception_handler(descriptor="pdfium")
+def construct_page_element_metadata(
+    structured_image: CroppedImageWithContent,
+    page_idx: int,
+    page_count: int,
+    source_metadata: Dict,
+    base_unified_metadata: Dict,
+):
+    """
+    +--------------------------------+--------------------------+------------+---+
+    | Table/Chart Metadata           |                          | Extracted  | Y |
+    | (tables within documents)      |                          |            |   |
+    +--------------------------------+--------------------------+------------+---+
+    | Table format                   | Structured (dataframe /  | Extracted  |   |
+    |                                | lists of rows and        |            |   |
+    |                                | columns), or serialized  |            |   |
+    |                                | as markdown, html,       |            |   |
+    |                                | latex, simple (cells     |            |   |
+    |                                | separated just as spaces)|            |   |
+    +--------------------------------+--------------------------+------------+---+
+    | Table content                  | Extracted text content   |            |   |
+    |                                |                          |            |   |
+    |                                | Important: Tables should |            |   |
+    |                                | not be chunked           |            |   |
+    +--------------------------------+--------------------------+------------+---+
+    | Table location                 | Bounding box of the table|            |   |
+    +--------------------------------+--------------------------+------------+---+
+    | Caption                        | Detected captions for    |            |   |
+    |                                | the table/chart          |            |   |
+    +--------------------------------+--------------------------+------------+---+
+    | uploaded_image_uri             | Mirrors                  |            |   |
+    |                                | source_metadata.         |            |   |
+    |                                | source_location          |            |   |
+    +--------------------------------+--------------------------+------------+---+
+    """
+
+    if structured_image.type_string in ("table",):
+        content = structured_image.image
+        structured_content_text = structured_image.content
+        structured_content_format = structured_image.content_format
+        table_format = TableFormatEnum.IMAGE
+        subtype = ContentTypeEnum.TABLE
+        description = ContentDescriptionEnum.PDF_TABLE
+        meta_name = "table_metadata"
+
+    elif structured_image.type_string in ("chart",):
+        content = structured_image.image
+        structured_content_text = structured_image.content
+        structured_content_format = structured_image.content_format
+        table_format = TableFormatEnum.IMAGE
+        subtype = ContentTypeEnum.CHART
+        description = ContentDescriptionEnum.PDF_CHART
+        # TODO(Devin) swap this to chart_metadata after we confirm metadata schema changes.
+        meta_name = "table_metadata"
+
+    elif structured_image.type_string in ("infographic",):
+        content = structured_image.image
+        structured_content_text = structured_image.content
+        structured_content_format = structured_image.content_format
+        table_format = TableFormatEnum.IMAGE
+        subtype = ContentTypeEnum.INFOGRAPHIC
+        description = ContentDescriptionEnum.PDF_INFOGRAPHIC
+        meta_name = "table_metadata"
+
+    else:
+        raise ValueError(f"Unknown table/chart/infographic type: {structured_image.type_string}")
+
+    content_metadata = {
+        "type": ContentTypeEnum.STRUCTURED,
+        "description": description,
+        "page_number": page_idx,
+        "hierarchy": {
+            "page_count": page_count,
+            "page": page_idx,
+            "line": -1,
+            "span": -1,
+        },
+        "subtype": subtype,
+    }
+
+    structured_metadata = {
+        "caption": "",
+        "table_format": table_format,
+        "table_content": structured_content_text,
+        "table_content_format": structured_content_format,
+        "table_location": structured_image.bbox,
+        "table_location_max_dimensions": (structured_image.max_width, structured_image.max_height),
+    }
+
+    ext_unified_metadata = base_unified_metadata.copy()
+
+    ext_unified_metadata.update(
+        {
+            "content": content,
+            "source_metadata": source_metadata,
+            "content_metadata": content_metadata,
+            meta_name: structured_metadata,
+        }
+    )
+
+    validated_unified_metadata = validate_metadata(ext_unified_metadata)
+
+    return [ContentTypeEnum.STRUCTURED, validated_unified_metadata.model_dump(), str(uuid.uuid4())]
+
+
+# TODO: remove this alias
+construct_table_and_chart_metadata = construct_page_element_metadata

nv_ingest_api/util/multi_processing/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+from .mp_pool_singleton import ProcessWorkerPoolSingleton
+
+__all__ = ["ProcessWorkerPoolSingleton"]

nv_ingest_api/util/multi_processing/mp_pool_singleton.py

@@ -0,0 +1,194 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+import math
+import multiprocessing as mp
+import os
+from threading import Lock
+from typing import Any, Callable, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class SimpleFuture:
+    """
+    A simplified future object that uses a multiprocessing Pipe to receive its result.
+
+    When the result() method is called, it blocks until the worker sends a tuple
+    (result, error) over the pipe.
+    """
+
+    def __init__(self, parent_conn: mp.connection.Connection) -> None:
+        """
+        Parameters
+        ----------
+        parent_conn : mp.connection.Connection
+            The parent end of the multiprocessing Pipe used to receive the result.
+        """
+        self._parent_conn: mp.connection.Connection = parent_conn
+
+    def result(self) -> Any:
+        """
+        Retrieve the result from the future, blocking until it is available.
+
+        Returns
+        -------
+        Any
+            The result returned by the worker function.
+
+        Raises
+        ------
+        Exception
+            If the worker function raised an exception, it is re-raised here.
+        """
+        result, error = self._parent_conn.recv()
+        if error is not None:
+            raise error
+        return result
+
+
+class ProcessWorkerPoolSingleton:
+    """
+    A singleton process worker pool using a dual-queue implementation.
+
+    Instead of a global result queue, each submitted task gets its own Pipe.
+    The submit_task() method returns a SimpleFuture, whose result() call blocks
+    until the task completes.
+    """
+
+    _instance: Optional["ProcessWorkerPoolSingleton"] = None
+    _lock: Lock = Lock()
+    _total_workers: int = 0
+
+    def __new__(cls) -> "ProcessWorkerPoolSingleton":
+        """
+        Create or return the singleton instance of ProcessWorkerPoolSingleton.
+
+        Returns
+        -------
+        ProcessWorkerPoolSingleton
+            The singleton instance.
+        """
+        logger.debug("Creating ProcessWorkerPoolSingleton instance...")
+        with cls._lock:
+            if cls._instance is None:
+                max_worker_limit: int = int(os.environ.get("MAX_INGEST_PROCESS_WORKERS", -1))
+                instance = super().__new__(cls)
+                # Determine available CPU count using affinity if possible
+                available: Optional[int] = (
+                    len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count()
+                )
+                # Use 40% of available CPUs, ensuring at least one worker
+                max_workers: int = math.floor(max(1, available * 0.4))
+                if (max_worker_limit > 0) and (max_workers > max_worker_limit):
+                    max_workers = max_worker_limit
+                logger.debug("Creating ProcessWorkerPoolSingleton instance with max workers: %d", max_workers)
+                instance._initialize(max_workers)
+                logger.debug("ProcessWorkerPoolSingleton instance created: %s", instance)
+                cls._instance = instance
+            else:
+                logger.debug("ProcessWorkerPoolSingleton instance already exists: %s", cls._instance)
+        return cls._instance
+
+    def _initialize(self, total_max_workers: int) -> None:
+        """
+        Initialize the worker pool with the specified number of worker processes.
+
+        Parameters
+        ----------
+        total_max_workers : int
+            The total number of worker processes to start.
+        """
+        self._total_workers = total_max_workers
+        self._context: mp.context.ForkContext = mp.get_context("fork")
+        # Bounded task queue: maximum tasks queued = 2 * total_max_workers.
+        self._task_queue: mp.Queue = self._context.Queue(maxsize=2 * total_max_workers)
+        self._next_task_id: int = 0
+        self._processes: list[mp.Process] = []
+        logger.debug(
+            "Initializing ProcessWorkerPoolSingleton with %d workers and queue size %d.",
+            total_max_workers,
+            2 * total_max_workers,
+        )
+        for i in range(total_max_workers):
+            p: mp.Process = self._context.Process(target=self._worker, args=(self._task_queue,))
+            p.start()
+            self._processes.append(p)
+            logger.debug("Started worker process %d/%d: PID %d", i + 1, total_max_workers, p.pid)
+        logger.debug("Initialized with max workers: %d", total_max_workers)
+
+    @staticmethod
+    def _worker(task_queue: mp.Queue) -> None:
+        """
+        Worker process that continuously processes tasks from the task queue.
+
+        Parameters
+        ----------
+        task_queue : mp.Queue
+            The queue from which tasks are retrieved.
+        """
+        logger.debug("Worker process started: PID %d", os.getpid())
+        while True:
+            task = task_queue.get()
+            if task is None:
+                # Stop signal received; exit the loop.
+                logger.debug("Worker process %d received stop signal.", os.getpid())
+                break
+            # Unpack task: (task_id, process_fn, args, child_conn)
+            task_id, process_fn, args, child_conn = task
+            try:
+                result = process_fn(*args)
+                child_conn.send((result, None))
+            except Exception as e:
+                logger.error("Task %d error in worker %d: %s", task_id, os.getpid(), e)
+                child_conn.send((None, e))
+            finally:
+                child_conn.close()
+
+    def submit_task(self, process_fn: Callable, *args: Any) -> SimpleFuture:
+        """
+        Submits a task to the worker pool for asynchronous execution.
+
+        If a single tuple is passed as the only argument, it is unpacked.
+
+        Parameters
+        ----------
+        process_fn : Callable
+            The function to be executed asynchronously.
+        *args : Any
+            The arguments to pass to the process function. If a single argument is a tuple,
+            it will be unpacked as the function arguments.
+
+        Returns
+        -------
+        SimpleFuture
+            A future object that can be used to retrieve the result of the task.
+        """
+        # Unpack tuple if a single tuple argument is provided.
+        if len(args) == 1 and isinstance(args[0], tuple):
+            args = args[0]
+        parent_conn, child_conn = mp.Pipe(duplex=False)
+        task_id: int = self._next_task_id
+        self._next_task_id += 1
+        self._task_queue.put((task_id, process_fn, args, child_conn))
+        return SimpleFuture(parent_conn)
+
+    def close(self) -> None:
+        """
+        Closes the worker pool and terminates all worker processes.
+
+        Sends a stop signal to each worker and waits for them to terminate.
+        """
+        logger.debug("Closing ProcessWorkerPoolSingleton...")
+        # Send a stop signal (None) for each worker.
+        for _ in range(self._total_workers):
+            self._task_queue.put(None)
+            logger.debug("Sent stop signal to worker.")
+        # Wait for all processes to finish.
+        for i, p in enumerate(self._processes):
+            p.join()
+            logger.debug("Worker process %d/%d joined: PID %d", i + 1, self._total_workers, p.pid)
+        logger.debug("ProcessWorkerPoolSingleton closed.")