nv-ingest 2025.5.21.dev20250521__py3-none-any.whl

nv_ingest/framework/orchestration/ray/stages/extractors/image_extractor.py

@@ -0,0 +1,85 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import ray
+
+from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+from nv_ingest.framework.util.flow_control import filter_by_task
+from nv_ingest_api.internal.extract.image.image_extractor import extract_primitives_from_image_internal
+from nv_ingest_api.internal.primitives.ingest_control_message import IngestControlMessage, remove_task_by_type
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
+from nv_ingest_api.internal.schemas.extract.extract_image_schema import ImageExtractorSchema
+from nv_ingest_api.util.exception_handlers.decorators import (
+    nv_ingest_node_failure_try_except,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class ImageExtractorStage(RayActorStage):
+    """
+    A Ray actor stage that extracts primitives from image content.
+
+    It expects an IngestControlMessage containing a DataFrame with image data. It then:
+      1. Removes the "extract" task from the message.
+      2. Calls the image extraction logic (via extract_primitives_from_image_internal) using a validated configuration.
+      3. Updates the message payload with the extracted primitives DataFrame.
+    """
+
+    def __init__(self, config: ImageExtractorSchema) -> None:
+        super().__init__(config)
+        try:
+            self.validated_config = config
+            logger.info("ImageExtractorStage configuration validated successfully.")
+        except Exception as e:
+            logger.exception(f"Error validating Image Extractor config: {e}")
+            raise
+
+    @traceable("image_extraction")
+    @filter_by_task(required_tasks=[("extract", {"document_type": "regex:^(png|jpeg|jpg|tiff|bmp)$"})])
+    @nv_ingest_node_failure_try_except(annotation_id="image_extractor", raise_on_failure=False)
+    def on_data(self, control_message: IngestControlMessage) -> IngestControlMessage:
+        """
+        Process the control message by extracting primitives from images.
+
+        Parameters
+        ----------
+        control_message : IngestControlMessage
+            The message containing a DataFrame payload with image data.
+
+        Returns
+        -------
+        IngestControlMessage
+            The updated message with extracted image primitives.
+        """
+        logger.info("ImageExtractorStage.on_data: Starting image extraction process.")
+        try:
+            # Extract the DataFrame payload.
+            df_ledger = control_message.payload()
+            logger.debug("Extracted payload with %d rows.", len(df_ledger))
+
+            # Remove the "extract" task from the message to obtain task-specific configuration.
+            task_config = remove_task_by_type(control_message, "extract")
+            logger.debug("Extracted task config: %s", task_config)
+
+            # Perform image primitives extraction.
+            new_df, extraction_info = extract_primitives_from_image_internal(
+                df_extraction_ledger=df_ledger,
+                task_config=task_config,
+                extraction_config=self.validated_config,
+                execution_trace_log=None,
+            )
+            logger.info("Image extraction completed. Resulting DataFrame has %d rows.", len(new_df))
+
+            # Update the message payload with the extracted primitives DataFrame.
+            control_message.payload(new_df)
+            control_message.set_metadata("image_extraction_info", extraction_info)
+
+            return control_message
+        except Exception as e:
+            logger.exception(f"ImageExtractorStage failed processing control message: {e}")
+            raise

nv_ingest/framework/orchestration/ray/stages/extractors/infographic_extractor.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import ray
+
+from nv_ingest_api.internal.schemas.extract.extract_infographic_schema import InfographicExtractorSchema
+from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+from nv_ingest.framework.util.flow_control import filter_by_task
+from nv_ingest_api.internal.extract.image.infographic_extractor import extract_infographic_data_from_image_internal
+from nv_ingest_api.internal.primitives.ingest_control_message import IngestControlMessage, remove_task_by_type
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
+from nv_ingest_api.util.exception_handlers.decorators import nv_ingest_node_failure_try_except
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class InfographicExtractorStage(RayActorStage):
+    def __init__(self, config: InfographicExtractorSchema) -> None:
+        super().__init__(config)
+
+        try:
+            self.validated_config = config
+            logger.info("ImageExtractorStage configuration validated successfully.")
+        except Exception as e:
+            logger.exception(f"Error validating Image Extractor config: {e}")
+            raise
+
+    @traceable("infographic_extraction")
+    @filter_by_task(required_tasks=["infographic_data_extract"])
+    @nv_ingest_node_failure_try_except(annotation_id="infographic_extraction", raise_on_failure=False)
+    def on_data(self, control_message: IngestControlMessage) -> IngestControlMessage:
+        # Extract DataFrame payload
+        df_ledger = control_message.payload()
+
+        # Remove the "infographic_data_extract" task from the message
+        task_config = remove_task_by_type(control_message, "infographic_data_extract")
+
+        execution_trace_log = {}
+        new_df, extraction_info = extract_infographic_data_from_image_internal(
+            df_extraction_ledger=df_ledger,
+            task_config=task_config,
+            extraction_config=self.validated_config,
+            execution_trace_log=execution_trace_log,
+        )
+
+        control_message.payload(new_df)
+        control_message.set_metadata("infographic_extraction_info", extraction_info)
+
+        do_trace_tagging = control_message.get_metadata("config::add_trace_tagging") is True
+        if do_trace_tagging and execution_trace_log:
+            for key, ts in execution_trace_log.items():
+                control_message.set_timestamp(key, ts)
+
+        return control_message

nv_ingest/framework/orchestration/ray/stages/extractors/pdf_extractor.py

@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import pandas as pd
+from typing import Any, Dict, Tuple, Optional
+import ray
+
+# Assume these imports come from your project:
+from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+from nv_ingest.framework.util.flow_control import filter_by_task
+from nv_ingest_api.internal.extract.pdf.pdf_extractor import extract_primitives_from_pdf_internal
+from nv_ingest_api.internal.primitives.ingest_control_message import remove_task_by_type
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
+from nv_ingest_api.internal.schemas.extract.extract_pdf_schema import PDFExtractorSchema
+from nv_ingest_api.util.exception_handlers.decorators import (
+    nv_ingest_node_failure_try_except,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _inject_validated_config(
+    df_extraction_ledger: pd.DataFrame,
+    task_config: Dict,
+    execution_trace_log: Optional[Any] = None,
+    validated_config: Any = None,
+) -> Tuple[pd.DataFrame, Dict]:
+    """
+    Helper function that injects the validated_config into the configuration for PDF extraction
+    and calls extract_primitives_from_pdf_internal.
+    """
+    return extract_primitives_from_pdf_internal(
+        df_extraction_ledger=df_extraction_ledger,
+        task_config=task_config,
+        extractor_config=validated_config,
+        execution_trace_log=execution_trace_log,
+    )
+
+
+@ray.remote
+class PDFExtractorStage(RayActorStage):
+    """
+    A Ray actor stage that extracts PDF primitives from a DataFrame payload.
+
+    It expects an IngestControlMessage containing a DataFrame of PDF documents. It then:
+      1. Removes the "extract" task from the message.
+      2. Calls the PDF extraction logic (via _inject_validated_config) using a validated configuration.
+      3. Updates the message payload with the extracted DataFrame.
+      4. Optionally, stores additional extraction info in the message metadata.
+    """
+
+    def __init__(self, config: PDFExtractorSchema) -> None:
+        super().__init__(config)
+        try:
+            # Validate and store the PDF extractor configuration.
+            self.validated_config = config
+            logger.info("PDFExtractorStage configuration validated successfully.")
+        except Exception as e:
+            logger.exception(f"Error validating PDF extractor config: {e}")
+            raise
+
+    @traceable("pdf_extraction")
+    @filter_by_task(required_tasks=[("extract", {"document_type": "pdf"})])
+    @nv_ingest_node_failure_try_except(annotation_id="pdf_extractor", raise_on_failure=False)
+    def on_data(self, control_message: Any) -> Any:
+        """
+        Process the control message by extracting PDF content.
+
+        Parameters
+        ----------
+        control_message : IngestControlMessage
+            The message containing a DataFrame payload with PDF documents.
+
+        Returns
+        -------
+        IngestControlMessage
+            The updated message with the extracted DataFrame and extraction info in metadata.
+        """
+
+        logger.info("PDFExtractorStage.on_data: Starting PDF extraction process.")
+
+        # Extract the DataFrame payload.
+        df_extraction_ledger = control_message.payload()
+        logger.debug("Extracted payload with %d rows.", len(df_extraction_ledger))
+
+        # Remove the "extract" task from the message to obtain task-specific configuration.
+        task_config = remove_task_by_type(control_message, "extract")
+        logger.debug("Extracted task config: %s", task_config)
+
+        # Perform PDF extraction.
+        execution_trace_log = {}
+        new_df, extraction_info = _inject_validated_config(
+            df_extraction_ledger,
+            task_config,
+            execution_trace_log=execution_trace_log,
+            validated_config=self.validated_config,
+        )
+        logger.info("PDF extraction completed. Extracted %d rows.", len(new_df))
+
+        # Update the message payload with the extracted DataFrame.
+        control_message.payload(new_df)
+        # Optionally, annotate the message with extraction info.
+        control_message.set_metadata("pdf_extraction_info", extraction_info)
+        logger.info("PDF extraction metadata injected successfully.")
+
+        do_trace_tagging = control_message.get_metadata("config::add_trace_tagging") is True
+        if do_trace_tagging and execution_trace_log:
+            for key, ts in execution_trace_log.items():
+                control_message.set_timestamp(key, ts)
+
+        return control_message

nv_ingest/framework/orchestration/ray/stages/extractors/pptx_extractor.py

@@ -0,0 +1,85 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import ray
+
+from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+from nv_ingest.framework.util.flow_control import filter_by_task
+from nv_ingest_api.internal.extract.pptx.pptx_extractor import extract_primitives_from_pptx_internal
+from nv_ingest_api.internal.primitives.ingest_control_message import IngestControlMessage, remove_task_by_type
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
+from nv_ingest_api.internal.schemas.extract.extract_pptx_schema import PPTXExtractorSchema
+from nv_ingest_api.util.exception_handlers.decorators import nv_ingest_node_failure_try_except
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class PPTXExtractorStage(RayActorStage):
+    """
+    A Ray actor stage that extracts content from PPTX documents.
+
+    It expects an IngestControlMessage containing a DataFrame with PPTX document data. It then:
+      1. Removes the "pptx-extract" task from the message.
+      2. Calls the PPTX extraction logic (via extract_primitives_from_pptx_internal) using a validated configuration.
+      3. Updates the message payload with the extracted content DataFrame.
+    """
+
+    def __init__(self, config: PPTXExtractorSchema) -> None:
+        """
+        Initializes the PptxExtractorStage.
+
+        Parameters
+        ----------
+        config : PPTXExtractorSchema
+            The validated configuration object for PPTX extraction.
+        """
+        super().__init__(config)
+        try:
+            # The config passed in should already be validated, but storing it.
+            self.validated_config = config
+            logger.info("PptxExtractorStage configuration validated successfully.")
+        except Exception as e:
+            # If RayActorStage.__init__ or config access raises an issue.
+            logger.exception(f"Error initializing or validating PPTX Extractor config: {e}")
+            raise
+
+    @traceable("pptx_extractor")
+    @filter_by_task(required_tasks=[("extract", {"document_type": "pptx"})])
+    @nv_ingest_node_failure_try_except(annotation_id="pptx_extractor", raise_on_failure=False)
+    def on_data(self, control_message: IngestControlMessage) -> IngestControlMessage:
+        """
+        Process the control message by extracting content from PPTX documents.
+
+        Parameters
+        ----------
+        control_message : IngestControlMessage
+            The message containing a DataFrame payload with PPTX document data.
+
+        Returns
+        -------
+        IngestControlMessage
+            The updated message with extracted PPTX content.
+        """
+
+        # Extract the DataFrame payload.
+        df_ledger = control_message.payload()
+
+        # Remove the "pptx-extract" task from the message to obtain task-specific configuration.
+        task_config = remove_task_by_type(control_message, "extract")
+
+        new_df, extraction_info = extract_primitives_from_pptx_internal(
+            df_extraction_ledger=df_ledger,
+            task_config=task_config,
+            extraction_config=self.validated_config,
+            execution_trace_log=None,  # Assuming None is appropriate here as in DOCX example
+        )
+
+        # Update the message payload with the extracted PPTX content DataFrame.
+        control_message.payload(new_df)
+        control_message.set_metadata("pptx_extraction_info", extraction_info)  # <-- Changed metadata key
+
+        return control_message

nv_ingest/framework/orchestration/ray/stages/extractors/table_extractor.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Any
+import ray
+
+# These imports are assumed from your project.
+from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+from nv_ingest.framework.util.flow_control import filter_by_task
+from nv_ingest_api.internal.extract.image.table_extractor import extract_table_data_from_image_internal
+from nv_ingest_api.internal.primitives.ingest_control_message import remove_task_by_type
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
+from nv_ingest_api.internal.schemas.extract.extract_table_schema import TableExtractorSchema
+from nv_ingest_api.util.exception_handlers.decorators import (
+    nv_ingest_node_failure_try_except,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class TableExtractorStage(RayActorStage):
+    """
+    A Ray actor stage that extracts table data from PDF content.
+
+    It expects an IngestControlMessage containing a DataFrame payload with PDF documents.
+    The stage removes the "table_data_extract" task from the message, calls the internal
+    extraction function using a validated TableExtractorSchema, updates the message payload,
+    and annotates the message metadata with extraction info.
+    """
+
+    def __init__(self, config: TableExtractorSchema) -> None:
+        super().__init__(config)
+        try:
+            self.validated_config = config
+            logger.info("TableExtractorStage configuration validated successfully.")
+        except Exception as e:
+            logger.exception("Error validating table extractor config")
+            raise e
+
+    @traceable("table_extraction")
+    @filter_by_task(required_tasks=["table_data_extract"])
+    @nv_ingest_node_failure_try_except(annotation_id="table_extraction", raise_on_failure=False)
+    def on_data(self, control_message: Any) -> Any:
+        """
+        Process the control message by extracting table data from the PDF payload.
+
+        Parameters
+        ----------
+        control_message : IngestControlMessage
+            The incoming message containing the PDF payload.
+
+        Returns
+        -------
+        IngestControlMessage
+            The updated message with the extracted table data and extraction info in metadata.
+        """
+        logger.info("TableExtractorStage.on_data: Starting table extraction.")
+        # Extract the DataFrame payload.
+        df_payload = control_message.payload()
+        logger.debug("Extracted payload with %d rows.", len(df_payload))
+
+        # Remove the "table_data_extract" task to obtain task-specific configuration.
+        task_config = remove_task_by_type(control_message, "table_data_extract")
+        logger.debug("Extracted task configuration: %s", task_config)
+
+        # Perform table data extraction.
+        execution_trace_log = {}
+        new_df, extraction_info = extract_table_data_from_image_internal(
+            df_extraction_ledger=df_payload,
+            task_config=task_config,
+            extraction_config=self.validated_config,
+            execution_trace_log=execution_trace_log,
+        )
+        logger.info("Table extraction completed. Extracted %d rows.", len(new_df))
+
+        # Update the control message with the new DataFrame.
+        control_message.payload(new_df)
+        # Annotate the message with extraction info.
+        control_message.set_metadata("table_extraction_info", extraction_info)
+        logger.info("Table extraction metadata injected successfully.")
+
+        do_trace_tagging = control_message.get_metadata("config::add_trace_tagging") is True
+        if do_trace_tagging and execution_trace_log:
+            for key, ts in execution_trace_log.items():
+                control_message.set_timestamp(key, ts)
+
+        return control_message

nv_ingest/framework/orchestration/ray/stages/injectors/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

nv_ingest/framework/orchestration/ray/stages/injectors/metadata_injector.py

@@ -0,0 +1,97 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import pandas as pd
+from typing import Any
+from pydantic import BaseModel
+import ray
+
+from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+from nv_ingest_api.internal.enums.common import DocumentTypeEnum, ContentTypeEnum
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
+from nv_ingest_api.util.converters.type_mappings import doc_type_to_content_type
+from nv_ingest_api.util.exception_handlers.decorators import (
+    nv_ingest_node_failure_try_except,
+)
+
+# logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class MetadataInjectionStage(RayActorStage):
+    """
+    A Ray actor stage that performs metadata injection on IngestControlMessages.
+
+    This stage iterates over the rows of the DataFrame payload, checks if metadata
+    injection is required, and if so, injects the appropriate metadata.
+    """
+
+    def __init__(self, config: BaseModel) -> None:
+        # Call the base initializer to set attributes like self._running.
+        super().__init__(config)
+        # Additional initialization can be added here if necessary.
+        logger.info("MetadataInjectionStage initialized with config: %s", config)
+
+    @traceable("metadata_injector")
+    @nv_ingest_node_failure_try_except(annotation_id="metadata_injector", raise_on_failure=False)
+    def on_data(self, message: Any) -> Any:
+        """
+        Process an incoming IngestControlMessage by injecting metadata into its DataFrame payload.
+
+        Parameters
+        ----------
+        message : IngestControlMessage
+            The incoming message containing the payload DataFrame.
+
+        Returns
+        -------
+        IngestControlMessage
+            The message with updated metadata if injection was required.
+        """
+        df = message.payload()
+        update_required = False
+        rows = []
+        logger.info("Starting metadata injection on DataFrame with %d rows", len(df))
+
+        for _, row in df.iterrows():
+            try:
+                # Convert document type to content type using enums.
+                content_type = doc_type_to_content_type(DocumentTypeEnum(row["document_type"]))
+                # Check if metadata is missing or doesn't contain 'content'
+                if "metadata" not in row or not isinstance(row["metadata"], dict) or "content" not in row["metadata"]:
+                    update_required = True
+                    row["metadata"] = {
+                        "content": row.get("content"),
+                        "content_metadata": {
+                            "type": content_type.name.lower(),
+                        },
+                        "error_metadata": None,
+                        "audio_metadata": (
+                            None if content_type != ContentTypeEnum.AUDIO else {"audio_type": row["document_type"]}
+                        ),
+                        "image_metadata": (
+                            None if content_type != ContentTypeEnum.IMAGE else {"image_type": row["document_type"]}
+                        ),
+                        "source_metadata": {
+                            "source_id": row.get("source_id"),
+                            "source_name": row.get("source_name"),
+                            "source_type": row["document_type"],
+                        },
+                        "text_metadata": (None if content_type != ContentTypeEnum.TEXT else {"text_type": "document"}),
+                    }
+            except Exception as inner_e:
+                logger.exception("Failed to process row during metadata injection")
+                raise inner_e
+            rows.append(row)
+
+        if update_required:
+            docs = pd.DataFrame(rows)
+            message.payload(docs)
+            logger.info("Metadata injection updated payload with %d rows", len(docs))
+        else:
+            logger.info("No metadata update was necessary during metadata injection")
+
+        return message

nv_ingest/framework/orchestration/ray/stages/meta/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

nv_ingest/framework/orchestration/ray/stages/meta/ray_actor_edge_base.py

@@ -0,0 +1,70 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+
+# TODO(Devin): Early prototype. Not currently used anywhere
+
+
+class RayActorEdge(ABC):
+    """
+    Abstract base class for a Ray actor edge used in a RayPipeline.
+
+    Parameters
+    ----------
+    max_size : int
+        The maximum size of the edge's internal queue.
+    multi_reader : bool
+        Whether the edge supports multiple concurrent readers.
+    multi_writer : bool
+        Whether the edge supports multiple concurrent writers.
+    """
+
+    def __init__(self, max_size: int, multi_reader: bool = False, multi_writer: bool = False) -> None:
+        self.max_size = max_size
+        self.multi_reader = multi_reader
+        self.multi_writer = multi_writer
+
+    @abstractmethod
+    def write(self, item: Any) -> bool:
+        """
+        Write an item into the edge.
+
+        Parameters
+        ----------
+        item : Any
+            The item to enqueue.
+
+        Returns
+        -------
+        bool
+            True if the item was enqueued successfully.
+        """
+        pass
+
+    @abstractmethod
+    def read(self) -> Any:
+        """
+        Read an item from the edge.
+
+        Returns
+        -------
+        Any
+            The next item in the edge.
+        """
+        pass
+
+    @abstractmethod
+    def get_stats(self) -> Dict[str, int]:
+        """
+        Get current statistics for the edge.
+
+        Returns
+        -------
+        Dict[str, int]
+            A dictionary containing statistics (e.g. write_count, read_count, queue_full_count, current_size).
+        """
+        pass