nv-ingest-api 2025.4.21.dev20250421__py3-none-any.whl → 2025.4.23.dev20250423__py3-none-any.whl

nv_ingest_api/internal/extract/docx/docx_extractor.py

@@ -0,0 +1,205 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import base64
+import functools
+import io
+import logging
+from typing import Optional, Dict, Any, Union
+
+import pandas as pd
+from pydantic import BaseModel
+
+from nv_ingest_api.internal.extract.docx.engines.docxreader_helpers.docx_helper import python_docx
+from nv_ingest_api.internal.schemas.extract.extract_docx_schema import DocxExtractorSchema
+from nv_ingest_api.util.exception_handlers.decorators import unified_exception_handler
+
+logger = logging.getLogger(__name__)
+
+
+def _prepare_task_props(
+    task_config: Union[Dict[str, Any], BaseModel], base64_row: pd.Series
+) -> (Dict[str, Any], Optional[str]):
+    """
+    Prepares the task properties by converting a Pydantic model to a dictionary (if needed)
+    and injecting row-specific data.
+
+    Parameters
+    ----------
+    task_config : Union[Dict[str, Any], BaseModel]
+        A dictionary or Pydantic model containing instructions and parameters for extraction.
+    base64_row : pd.Series
+        A Series representing a row from the DataFrame that contains at least the "content"
+        key and optionally "source_id".
+
+    Returns
+    -------
+    Tuple[Dict[str, Any], Optional[str]]
+        A tuple where the first element is the prepared task properties dictionary with the key
+        "row_data" added under its "params" key, and the second element is the source_id (if present),
+        otherwise None.
+    """
+
+    if isinstance(task_config, BaseModel):
+        task_config = task_config.model_dump()
+    else:
+        task_config = dict(task_config)
+
+    # Extract all row data except the "content" field.
+    row_data = base64_row.drop(labels=["content"], errors="ignore")
+    if "params" not in task_config:
+        task_config["params"] = {}
+
+    task_config["params"]["row_data"] = row_data
+
+    source_id = base64_row.get("source_id", None)
+
+    return task_config, source_id
+
+
+@unified_exception_handler
+def _decode_and_extract_from_docx(
+    base64_row: pd.Series,
+    task_config: Union[Dict[str, Any], BaseModel],
+    extraction_config: Any,
+    execution_trace_log: Optional[Dict[str, Any]] = None,
+) -> Any:
+    """
+    Decodes base64 content from a DataFrame row and extracts data using the specified extraction method.
+
+    The function decodes the base64-encoded content from the "content" key in the row, prepares
+    extraction parameters (including additional row data and configuration), and invokes the extraction
+    function from the docx module. If an error occurs, an exception tag is returned.
+
+    Parameters
+    ----------
+    base64_row : pd.Series
+        A Series containing the base64-encoded content under the key "content" and optionally a "source_id".
+    task_config : Union[Dict[str, Any], BaseModel]
+        A dictionary or Pydantic model containing extraction instructions and parameters.
+        Expected to have a "params" key for additional parameters and optionally a "method" key specifying
+        the extraction method.
+    extraction_config : Any
+        A configuration object that contains extraction-specific settings, such as `docx_extraction_config`.
+    execution_trace_log : Optional[Dict[str, Any]], default=None
+        A dictionary containing trace information for debugging or logging.
+    default : str, optional
+        The default extraction method to use if the specified method is not available (default is "python_docx").
+
+    Returns
+    -------
+    Any
+        The extracted data, or an exception tag if extraction fails.
+
+    Raises
+    ------
+    Exception
+        If an unhandled exception occurs during extraction, it is logged and a tagged error is returned.
+    """
+    # Prepare task properties and extract source_id
+    task_config, source_id = _prepare_task_props(task_config, base64_row)
+
+    # Retrieve base64 content and decode it into a byte stream.
+    base64_content: str = base64_row["content"]
+    doc_bytes: bytes = base64.b64decode(base64_content)
+    doc_stream: io.BytesIO = io.BytesIO(doc_bytes)
+
+    extract_params: Dict[str, Any] = task_config.get("params", {})
+
+    # Extract required boolean flags from params.
+    try:
+        extract_text = extract_params.pop("extract_text", False)
+        extract_images = extract_params.pop("extract_images", False)
+        extract_tables = extract_params.pop("extract_tables", False)
+        extract_charts = extract_params.pop("extract_charts", False)
+        extract_infographics = extract_params.pop("extract_infographics", False)
+    except KeyError as e:
+        raise ValueError(f"Missing required extraction flag: {e}")
+
+    # Inject configuration and trace info into extraction parameters.
+    if getattr(extraction_config, "docx_extraction_config", None) is not None:
+        extract_params["docx_extraction_config"] = extraction_config.docx_extraction_config
+
+    if execution_trace_log is not None:
+        extract_params["trace_info"] = execution_trace_log
+
+    # extraction_func: Callable = _get_extraction_function(extract_method, default)
+    extracted_data: Any = python_docx(
+        docx_stream=doc_stream,
+        extract_text=extract_text,
+        extract_images=extract_images,
+        extract_infographics=extract_infographics,
+        extract_tables=extract_tables,
+        extract_charts=extract_charts,
+        extraction_config=extract_params,
+        execution_trace_log=None,
+    )
+
+    return extracted_data
+
+
+@unified_exception_handler
+def extract_primitives_from_docx_internal(
+    df_extraction_ledger: pd.DataFrame,
+    task_config: Union[Dict[str, Any], BaseModel],
+    extraction_config: DocxExtractorSchema,
+    execution_trace_log: Optional[Dict[str, Any]] = None,
+) -> pd.DataFrame:
+    """
+    Processes a pandas DataFrame containing DOCX files encoded in base64, extracting text from
+    each document and replacing the original content with the extracted text.
+
+    This function applies a decoding and extraction routine to each row of the input DataFrame.
+    The routine is provided via the `decode_and_extract` function, which is partially applied with
+    task configuration, extraction configuration, and optional trace information. The results are
+    exploded and any missing values are dropped, then compiled into a new DataFrame with columns
+    for document type, metadata, and a UUID identifier.
+
+    Parameters
+    ----------
+    df_extraction_ledger : pd.DataFrame
+        The input DataFrame containing DOCX files in base64 encoding. Expected columns include
+        'source_id' and 'content'.
+    task_config : Union[Dict[str, Any], BaseModel]
+        Configuration instructions for the document processing task. This can be provided as a
+        dictionary or a Pydantic model.
+    extraction_config : Any
+        A configuration object for document extraction that guides the extraction process.
+    execution_trace_log : Optional[Dict[str, Any]], default=None
+        An optional dictionary containing trace information for debugging or logging.
+
+    Returns
+    -------
+    pd.DataFrame
+        A DataFrame with the original DOCX content replaced by the extracted text. The resulting
+        DataFrame contains the columns "document_type", "metadata", and "uuid".
+
+    Raises
+    ------
+    Exception
+        If an error occurs during the document extraction process, the exception is logged and
+        re-raised.
+    """
+    # Create a partial function to decode and extract using the provided configurations.
+    _decode_and_extract = functools.partial(
+        _decode_and_extract_from_docx,
+        task_config=task_config,
+        extraction_config=extraction_config,
+        execution_trace_log=execution_trace_log,
+    )
+
+    # Apply the decode_and_extract function to each row in the DataFrame.
+    sr_extraction = df_extraction_ledger.apply(_decode_and_extract, axis=1)
+
+    # Explode any list results and drop missing values.
+    sr_extraction = sr_extraction.explode().dropna()
+
+    # Convert the extraction results to a DataFrame if available.
+    if not sr_extraction.empty:
+        extracted_df = pd.DataFrame(sr_extraction.to_list(), columns=["document_type", "metadata", "uuid"])
+    else:
+        extracted_df = pd.DataFrame({"document_type": [], "metadata": [], "uuid": []})
+
+    return extracted_df

nv_ingest_api/internal/extract/docx/engines/docxreader_helpers/__init__.py

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

nv_ingest_api/internal/extract/docx/engines/docxreader_helpers/docx_helper.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2024, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=too-many-locals
+
+
+import logging
+from typing import IO, Optional, List
+
+from nv_ingest_api.internal.enums.common import AccessLevelEnum, DocumentTypeEnum
+from nv_ingest_api.internal.enums.common import TextTypeEnum
+from nv_ingest_api.internal.extract.docx.engines.docxreader_helpers.docxreader import DocxReader
+
+logger = logging.getLogger(__name__)
+
+
+def python_docx(
+    *,
+    docx_stream: IO,
+    extract_text: bool,
+    extract_images: bool,
+    extract_infographics: bool,
+    extract_tables: bool,
+    extract_charts: bool,
+    extraction_config: dict,
+    execution_trace_log: Optional[List] = None,
+):
+    """
+    Helper function that use python-docx to extract text from a bytestream document
+
+    A document has three levels - document, paragraphs and runs. To align with the
+    pdf extraction paragraphs are aliased as block. python-docx leaves the page number
+    and line number to the renderer so we assume that the entire document is a single
+    page.
+
+    Run level parsing has been skipped but can be added as needed.
+
+    Parameters
+    ----------
+    docx_stream:
+        Bytestream
+    extract_text : bool
+        Specifies whether to extract text.
+    extract_images : bool
+        Specifies whether to extract images.
+    extract_infographics : bool
+        Specifies whether to extract infographics.
+    extract_tables : bool
+        Specifies whether to extract tables.
+    extract_charts : bool
+        Specifies whether to extract charts.
+    extraction_config : dict
+        A dictionary of configuration parameters for the extraction process.
+    execution_trace_log : list, optional
+        A list for accumulating trace information during extraction. Defaults to None.
+
+    Returns
+    -------
+    str
+        A string of extracted text.
+    """
+
+    _ = execution_trace_log
+    _ = extract_infographics
+
+    row_data = extraction_config.get("row_data")
+    # get source_id
+    source_id = row_data["source_id"]
+    # get text_depth
+    text_depth = extraction_config.get("text_depth", "document")
+    text_depth = TextTypeEnum(text_depth)
+    # get base metadata
+    metadata_col = "metadata"
+
+    docx_extractor_config = extraction_config.get("docx_extraction_config", {})
+
+    base_unified_metadata = row_data[metadata_col] if metadata_col in row_data.index else {}
+
+    # get base source_metadata
+    base_source_metadata = base_unified_metadata.get("source_metadata", {})
+    # get source_location
+    source_location = base_source_metadata.get("source_location", "")
+    # get collection_id (assuming coming in from source_metadata...)
+    collection_id = base_source_metadata.get("collection_id", "")
+    # get partition_id (assuming coming in from source_metadata...)
+    partition_id = base_source_metadata.get("partition_id", -1)
+    # get access_level (assuming coming in from source_metadata...)
+    access_level = base_source_metadata.get("access_level", AccessLevelEnum.UNKNOWN)
+
+    # python-docx doesn't maintain filename; re-use source_id
+    source_metadata = {
+        "source_name": source_id,
+        "source_id": source_id,
+        "source_location": source_location,
+        "source_type": DocumentTypeEnum.DOCX,
+        "collection_id": collection_id,
+        "partition_id": partition_id,
+        "access_level": access_level,
+        "summary": "",
+    }
+
+    # Extract data from the document using python-docx
+    doc = DocxReader(docx_stream, source_metadata, extraction_config=docx_extractor_config)
+    extracted_data = doc.extract_data(
+        base_unified_metadata, text_depth, extract_text, extract_charts, extract_tables, extract_images
+    )
+
+    return extracted_data