nv-ingest-client 2025.7.24.dev20250724__py3-none-any.whl → 2025.11.2.dev20251102__py3-none-any.whl

nv_ingest_client/util/document_analysis.py

@@ -0,0 +1,314 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Utility functions for analyzing document-level chunk composition from nv-ingest results.
+
+This module provides analysis capabilities for understanding the distribution and types
+of extracted content elements across individual documents. It enables customers to
+gain visibility into their document composition for performance optimization and
+capacity planning decisions.
+"""
+
+import logging
+import os
+from collections import defaultdict
+from typing import Any, Dict, List, Union
+
+logger = logging.getLogger(__name__)
+
+
+def analyze_document_chunks(
+    results: Union[List[List[Dict[str, Any]]], List[Dict[str, Any]]],
+) -> Dict[str, Dict[str, Dict[str, int]]]:
+    """
+    Analyze ingestor results to count elements by type and page for each document.
+
+    This function processes results from nv-ingest ingestion and provides a per-document,
+    per-page breakdown of extracted content types, enabling customers to understand document
+    composition and page-level distribution for optimization and planning purposes.
+
+    Parameters
+    ----------
+    results : Union[List[List[Dict[str, Any]]], List[Dict[str, Any]]]
+        Ingestor results from ingestor.ingest() in standard List[List[Dict]] format,
+        or flattened List[Dict] format. Handles both regular lists and
+        LazyLoadedList objects automatically.
+
+    Returns
+    -------
+    Dict[str, Dict[str, Dict[str, int]]]
+        Dictionary mapping document names to page-level element type counts with structure:
+        {
+            "document1.pdf": {
+                "total": {
+                    "text": 7, "charts": 1, "tables": 1,
+                    "unstructured_images": 0, "infographics": 0, "page_images": 0
+                },
+                "1": {
+                    "text": 3, "charts": 1, "tables": 0,
+                    "unstructured_images": 0, "infographics": 0, "page_images": 0
+                },
+                "2": {
+                    "text": 4, "charts": 0, "tables": 1,
+                    "unstructured_images": 0, "infographics": 0, "page_images": 0
+                }
+            },
+            "document2.pdf": {...}
+        }
+
+    Notes
+    -----
+    - Requires purge_results_after_upload=False in vdb_upload() configuration
+    - Automatically handles LazyLoadedList objects from nv-ingest client
+    - Returns zero counts for missing element types
+    - Assumes valid nv-ingest output format with guaranteed metadata structure
+
+    Examples
+    --------
+    >>> from nv_ingest_client.util.document_analysis import analyze_document_chunks
+    >>>
+    >>> # After running ingestion
+    >>> results, failures = ingestor.ingest(show_progress=True, return_failures=True)
+    >>>
+    >>> # Analyze document composition by page
+    >>> breakdown = analyze_document_chunks(results)
+    >>>
+    >>> for doc_name, pages in breakdown.items():
+    ...     total_counts = pages["total"]
+    ...     total_elements = sum(total_counts.values())
+    ...     page_count = len(pages) - 1  # Subtract 1 for "total" key
+    ...     print(f"{doc_name}: {total_elements} elements across {page_count} pages")
+    ...     print(f"  total: {total_elements} elements ({total_counts['text']} text, {total_counts['charts']} charts)")
+    ...     for page_name, counts in pages.items():
+    ...         if page_name != "total":  # Skip total when listing pages
+    ...             page_total = sum(counts.values())
+    ...             print(
+                f"  page {page_name}: {page_total} elements "
+                f"({counts['text']} text, {counts['charts']} charts)"
+            )
+    """
+
+    if not results:
+        logger.warning("No results provided for analysis")
+        return {}
+
+    # Normalize input format to handle both List[List[Dict]] and List[Dict] structures
+    normalized_results = _normalize_results_format(results)
+
+    # Group elements by document name and page number
+    document_page_elements = defaultdict(lambda: defaultdict(list))
+
+    for doc_results in normalized_results:
+        # Handle LazyLoadedList and other iterable types
+        elements = _extract_elements_from_doc(doc_results)
+
+        for element in elements:
+            doc_name = _extract_document_name(element)
+            page_key = _extract_page_key(element)
+            document_page_elements[doc_name][page_key].append(element)
+
+    # Count element types per page within each document and calculate totals
+    document_page_counts = {}
+
+    for doc_name, pages in document_page_elements.items():
+        document_page_counts[doc_name] = {}
+        total_counts = _initialize_element_counts()
+
+        for page_key, elements in pages.items():
+            counts = _initialize_element_counts()
+
+            for element in elements:
+                element_type = _categorize_element(element)
+                counts[element_type] += 1
+                total_counts[element_type] += 1  # Add to document total
+
+            document_page_counts[doc_name][page_key] = counts
+
+        # Add the total counts for this document
+        document_page_counts[doc_name]["total"] = total_counts
+
+    if document_page_counts:
+        total_docs = len(document_page_counts)
+        total_pages = sum(len(pages) - 1 for pages in document_page_counts.values())  # Subtract 1 for "total" key
+        total_elements = sum(sum(page_counts["total"].values()) for page_counts in document_page_counts.values())
+        logger.info(f"Analyzed {total_elements} elements across {total_pages} pages in {total_docs} documents")
+    else:
+        logger.warning("No valid documents found for analysis")
+
+    return document_page_counts
+
+
+def _normalize_results_format(results: Union[List[List[Dict]], List[Dict]]) -> List[List[Dict]]:
+    """
+    Normalize various input formats to consistent List[List[Dict]] structure.
+
+    Parameters
+    ----------
+    results : Union[List[List[Dict]], List[Dict]]
+        Input results in various formats
+
+    Returns
+    -------
+    List[List[Dict]]
+        Normalized results in standard format
+    """
+
+    if not results:
+        return []
+
+    # Handle List[List[Dict]] or List[LazyLoadedList] formats
+    if isinstance(results, list) and len(results) > 0:
+        first_elem = results[0]
+        # Check for list, LazyLoadedList, or any sequence-like object
+        if isinstance(first_elem, list) or (
+            hasattr(first_elem, "__iter__") and hasattr(first_elem, "__len__") and not isinstance(first_elem, dict)
+        ):
+            return results
+
+    # Handle flattened List[Dict] format by grouping elements by document
+    if isinstance(results, list) and len(results) > 0 and isinstance(results[0], dict):
+        doc_groups = defaultdict(list)
+        for element in results:
+            doc_name = _extract_document_name(element)
+            doc_groups[doc_name].append(element)
+
+        return list(doc_groups.values())
+
+    # Fallback for unexpected formats
+    return [[item] for item in results if item]
+
+
+def _extract_elements_from_doc(doc_results) -> List[Dict]:
+    """
+    Extract elements from document results, handling various data types.
+
+    Parameters
+    ----------
+    doc_results : Any
+        Document results which may be a list, LazyLoadedList, or other iterable
+
+    Returns
+    -------
+    List[Dict]
+        List of element dictionaries
+    """
+
+    if isinstance(doc_results, list):
+        return doc_results
+    elif hasattr(doc_results, "__iter__") and hasattr(doc_results, "__len__"):
+        # Handle LazyLoadedList and other sequence-like objects
+        return list(doc_results)
+    else:
+        # Single element case
+        return [doc_results] if doc_results else []
+
+
+def _extract_document_name(element: Dict[str, Any]) -> str:
+    """
+    Extract clean document name from element metadata.
+
+    Parameters
+    ----------
+    element : Dict[str, Any]
+        Element dictionary containing metadata
+
+    Returns
+    -------
+    str
+        Clean document filename (basename of source_id)
+    """
+
+    # nv-ingest guarantees this structure exists
+    source_id = element["metadata"]["source_metadata"]["source_id"]
+    return os.path.basename(source_id)
+
+
+def _extract_page_key(element: Dict[str, Any]) -> str:
+    """
+    Extract page key from element metadata for consistent page naming.
+
+    Parameters
+    ----------
+    element : Dict[str, Any]
+        Element dictionary containing metadata
+
+    Returns
+    -------
+    str
+        Page number as string (e.g., "1", "2", or "unknown")
+    """
+
+    try:
+        page_number = element["metadata"]["content_metadata"]["page_number"]
+        if page_number is not None and page_number >= 0:
+            return str(page_number)
+        else:
+            return "unknown"
+    except (KeyError, TypeError):
+        logger.warning("Missing or invalid page_number in element metadata")
+        return "unknown"
+
+
+def _categorize_element(element: Dict[str, Any]) -> str:
+    """
+    Categorize element by type using document_type and content metadata.
+
+    Parameters
+    ----------
+    element : Dict[str, Any]
+        Element dictionary with document_type and metadata fields
+
+    Returns
+    -------
+    str
+        Element category: "text", "charts", "tables", "unstructured_images",
+        "infographics", or "page_images"
+    """
+
+    doc_type = element["document_type"]
+
+    # Text elements
+    if doc_type == "text":
+        return "text"
+
+    # Structured elements with subtypes
+    elif doc_type == "structured":
+        subtype = element["metadata"]["content_metadata"]["subtype"]
+        if subtype == "chart":
+            return "charts"
+        elif subtype == "table":
+            return "tables"
+        elif subtype == "infographic":
+            return "infographics"
+        elif subtype == "page_image":
+            return "page_images"
+
+    # Image elements (unstructured)
+    elif doc_type == "image":
+        return "unstructured_images"
+
+    # Should not reach here with valid nv-ingest output
+    logger.warning(f"Unexpected element type: {doc_type}")
+    return "text"  # Default to text for safety
+
+
+def _initialize_element_counts() -> Dict[str, int]:
+    """
+    Initialize element counts dictionary with all supported types.
+
+    Returns
+    -------
+    Dict[str, int]
+        Dictionary with zero counts for all element types
+    """
+
+    return {
+        "text": 0,
+        "charts": 0,
+        "tables": 0,
+        "unstructured_images": 0,
+        "infographics": 0,
+        "page_images": 0,
+    }

nv_ingest_client/util/image_disk_utils.py

@@ -0,0 +1,300 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Utility functions for saving images from ingestion results to disk as actual image files.
+
+This module provides comprehensive utilities for extracting and saving base64-encoded
+images from nv-ingest results to local filesystem. Features include:
+- Configurable filtering by image type (charts, tables, infographics, etc.)
+- Descriptive filename generation with source and page information
+- Organized directory structure by image type
+- Detailed image counting and statistics
+
+Typical use cases:
+- Debugging and visual inspection of extracted content
+- Quality assessment of image extraction pipeline
+"""
+
+import logging
+import os
+from typing import Any, Dict, List
+
+from nv_ingest_client.client.util.processing import get_valid_filename
+from nv_ingest_api.util.image_processing.transforms import save_image_to_disk, _detect_base64_image_format
+
+logger = logging.getLogger(__name__)
+
+
+def _detect_extension_from_content(image_content: str) -> str:
+    """
+    Get file extension by detecting original image format.
+    Falls back to .jpeg if detection fails or format is unknown.
+    """
+    DEFAULT_EXT = "jpg"  # must be either "jpg" or "png"
+    try:
+        fmt = _detect_base64_image_format(image_content).upper()
+    except Exception:
+        logger.warning("Image format detection failed; falling back to default '%s'.", DEFAULT_EXT)
+        return DEFAULT_EXT
+    ext_map = {
+        "JPEG": "jpg",
+        "JPG": "jpg",
+        "PNG": "png",
+    }
+    ext = ext_map.get(fmt, None)
+    if ext:
+        return ext
+    logger.warning("Unsupported image format '%s'; falling back to default '%s'.", fmt, DEFAULT_EXT)
+    return DEFAULT_EXT
+
+
+def save_images_to_disk(
+    response_data: List[Dict[str, Any]],
+    output_directory: str,
+    save_charts: bool = True,
+    save_tables: bool = True,
+    save_infographics: bool = True,
+    save_page_images: bool = False,
+    save_raw_images: bool = False,
+    organize_by_type: bool = True,
+    output_format: str = "auto",
+) -> Dict[str, int]:
+    """
+    Save base64-encoded images from ingestion results to disk as actual image files.
+
+    This utility extracts images from ingestion response data and saves them to disk
+    with descriptive filenames that include the image subtype and page information.
+    It provides granular control over which types of images to save.
+
+    Parameters
+    ----------
+    response_data : List[Dict[str, Any]]
+        List of document results from ingestion, each containing metadata with base64 images.
+    output_directory : str
+        Base directory where images will be saved.
+    save_charts : bool, optional
+        Whether to save chart images. Default is True.
+    save_tables : bool, optional
+        Whether to save table images. Default is True.
+    save_infographics : bool, optional
+        Whether to save infographic images. Default is True.
+    save_page_images : bool, optional
+        Whether to save page-as-image files. Default is False.
+    save_raw_images : bool, optional
+        Whether to save raw/natural images. Default is False.
+    organize_by_type : bool, optional
+        Whether to organize images into subdirectories by type. Default is True.
+    output_format : str, optional
+        Output image format for saved files. Default is "auto".
+        - "auto": Preserve original format (fastest, no conversion)
+        - "jpeg": Convert to JPEG (smaller files, good compression)
+        - "png": Convert to PNG (lossless quality)
+        Use "auto" for maximum speed by avoiding format conversion.
+
+    Returns
+    -------
+    Dict[str, int]
+        Dictionary with counts of images saved by type.
+
+    Raises
+    ------
+    ValueError
+        If output_format is not supported.
+
+    Examples
+    --------
+    >>> from nv_ingest_client.util.image_disk_utils import save_images_to_disk
+    >>>
+    >>> # Save only charts and tables
+    >>> counts = save_images_to_disk(
+    ...     response_data,
+    ...     "./output/images",
+    ...     save_charts=True,
+    ...     save_tables=True,
+    ...     save_page_images=False
+    ... )
+    >>> print(f"Saved {counts['chart']} charts and {counts['table']} tables")
+    """
+
+    if not response_data:
+        logger.warning("No response data provided")
+        return {}
+
+    # Validate format upfront to fail fast
+    normalized_format = output_format.lower()
+    if normalized_format not in ["auto", "png", "jpeg", "jpg"]:
+        raise ValueError(
+            f"Unsupported output format: '{output_format}'. Supported formats: 'auto', 'png', 'jpeg', 'jpg'"
+        )
+
+    # Initialize counters
+    image_counts = {"chart": 0, "table": 0, "infographic": 0, "page_image": 0, "image": 0, "total": 0}
+
+    # Create output directory
+    os.makedirs(output_directory, exist_ok=True)
+
+    for doc_idx, document in enumerate(response_data):
+        try:
+            metadata = document.get("metadata", {})
+            doc_type = document.get("document_type", "unknown")
+
+            # Skip documents without image content
+            image_content = metadata.get("content")
+            if not image_content:
+                continue
+
+            # Get document info for naming
+            source_metadata = metadata.get("source_metadata", {})
+            source_id = source_metadata.get("source_id", f"document_{doc_idx}")
+            clean_source_name = get_valid_filename(os.path.basename(source_id))
+
+            content_metadata = metadata.get("content_metadata", {})
+            subtype = content_metadata.get("subtype", "image")
+            page_number = content_metadata.get("page_number", 0)
+
+            # Apply filtering based on image subtype and user preferences
+            should_save = False
+            if subtype == "chart" and save_charts:
+                should_save = True
+            elif subtype == "table" and save_tables:
+                should_save = True
+            elif subtype == "infographic" and save_infographics:
+                should_save = True
+            elif subtype == "page_image" and save_page_images:
+                should_save = True
+            elif (
+                doc_type == "image"
+                and subtype not in ["chart", "table", "infographic", "page_image"]
+                and save_raw_images
+            ):
+                should_save = True
+                subtype = "image"  # Normalize subtype for consistent counting
+
+            if not should_save:
+                continue
+
+            # Determine file extension and target format (format already validated upfront)
+            if normalized_format in ["jpeg", "jpg"]:
+                file_ext, target_format = "jpeg", "jpeg"
+            elif normalized_format == "png":
+                file_ext, target_format = "png", "png"
+            else:  # normalized_format == "auto" - detect once and use result
+                detected_ext = _detect_extension_from_content(image_content)
+                if detected_ext == "png":
+                    file_ext, target_format = "png", "png"
+                else:  # detected_ext == "jpeg"
+                    file_ext, target_format = "jpeg", "jpeg"
+
+            if organize_by_type:
+                # Organize into subdirectories by image type
+                type_dir = os.path.join(output_directory, subtype)
+                os.makedirs(type_dir, exist_ok=True)
+                image_filename = f"{clean_source_name}_p{page_number}_{doc_idx}.{file_ext}"
+                image_path = os.path.join(type_dir, image_filename)
+            else:
+                # Flat directory structure with type in filename
+                image_filename = f"{clean_source_name}_{subtype}_p{page_number}_{doc_idx}.{file_ext}"
+                image_path = os.path.join(output_directory, image_filename)
+
+            # Save image using centralized API function
+            try:
+                success = save_image_to_disk(image_content, image_path, target_format)
+
+                if success:
+                    # Update image type counters
+                    image_counts[subtype] += 1
+                    image_counts["total"] += 1
+                    logger.debug(f"Saved {subtype} image: {image_path}")
+                else:
+                    logger.error(f"Failed to save {subtype} image for {clean_source_name}")
+
+            except Exception as e:
+                logger.error(f"Failed to save {subtype} image for {clean_source_name}: {e}")
+
+        except Exception as e:
+            logger.error(f"Failed to process document {doc_idx}: {e}")
+            continue
+
+    # Log summary statistics
+    if image_counts["total"] > 0:
+        logger.info(f"Successfully saved {image_counts['total']} images to {output_directory}")
+        for img_type, count in image_counts.items():
+            if img_type != "total" and count > 0:
+                logger.info(f"  - {img_type}: {count}")
+    else:
+        logger.info("No images were saved (none met filter criteria)")
+
+    return image_counts
+
+
+def save_images_from_response(response: Dict[str, Any], output_directory: str, **kwargs) -> Dict[str, int]:
+    """
+    Convenience function to save images from a full API response.
+
+    Parameters
+    ----------
+    response : Dict[str, Any]
+        Full API response containing a "data" field with document results.
+    output_directory : str
+        Directory where images will be saved.
+    **kwargs
+        Additional arguments passed to save_images_to_disk().
+        Includes output_format ("auto", "png", or "jpeg") and other filtering options.
+
+    Returns
+    -------
+    Dict[str, int]
+        Dictionary with counts of images saved by type.
+    """
+
+    if "data" not in response or not response["data"]:
+        logger.warning("No data found in response")
+        return {}
+
+    return save_images_to_disk(response["data"], output_directory, **kwargs)
+
+
+def save_images_from_ingestor_results(
+    results: List[List[Dict[str, Any]]], output_directory: str, **kwargs
+) -> Dict[str, int]:
+    """
+    Save images from Ingestor.ingest() results.
+
+    Parameters
+    ----------
+    results : List[List[Dict[str, Any]]]
+        Results from Ingestor.ingest(), where each inner list contains
+        document results for one source file. Can also handle LazyLoadedList
+        objects when save_to_disk=True is used.
+    output_directory : str
+        Directory where images will be saved.
+    **kwargs
+        Additional arguments passed to save_images_to_disk().
+        Includes output_format ("auto", "png", or "jpeg") and other filtering options.
+
+    Returns
+    -------
+    Dict[str, int]
+        Dictionary with counts of images saved by type.
+    """
+
+    # Flatten results from multiple documents into single list
+    all_documents = []
+    for doc_results in results:
+        if isinstance(doc_results, list):
+            # Standard list of document results
+            all_documents.extend(doc_results)
+        elif hasattr(doc_results, "__iter__") and hasattr(doc_results, "__len__"):
+            # Handle LazyLoadedList or other sequence-like objects
+            try:
+                all_documents.extend(list(doc_results))
+            except Exception as e:
+                logger.warning(f"Failed to process document results: {e}")
+                continue
+        else:
+            # Handle single document case
+            all_documents.append(doc_results)
+
+    return save_images_to_disk(all_documents, output_directory, **kwargs)

nv_ingest_client/util/transport.py

@@ -15,35 +15,41 @@ def infer_microservice(
     truncate: str = "END",
     batch_size: int = 8191,
     grpc: bool = False,
+    input_names: list = ["text"],
+    output_names: list = ["embeddings"],
+    dtypes: list = ["BYTES"],
 ):
     """
     This function takes the input data and creates a list of embeddings
     using the NVIDIA embedding microservice.
     """
-    data = {"prompts": [res["metadata"]["content"] for res in data]}
+    if isinstance(data[0], str):
+        data = {"prompts": data}
+    else:
+        data = {"prompts": [res["metadata"]["content"] for res in data]}
     if grpc:
         model_name = re.sub(r"[^a-zA-Z0-9]", "_", model_name)
         client = NimClient(
             model_interface=EmbeddingModelInterface(),
             protocol="grpc",
-            endpoints=(embedding_endpoint, embedding_endpoint),
+            endpoints=(embedding_endpoint, None),
             auth_token=nvidia_api_key,
         )
         return client.infer(
             data,
             model_name,
             parameters={"input_type": input_type, "truncate": truncate},
-            outputs=["embeddings"],
-            dtype=["BYTES"],
-            input_name=["text"],
+            dtypes=dtypes,
+            input_names=input_names,
             batch_size=batch_size,
+            output_names=output_names,
         )
     else:
         embedding_endpoint = f"{embedding_endpoint}/embeddings"
         client = NimClient(
             model_interface=EmbeddingModelInterface(),
             protocol="http",
-            endpoints=(embedding_endpoint, embedding_endpoint),
+            endpoints=(None, embedding_endpoint),
             auth_token=nvidia_api_key,
         )
         return client.infer(data, model_name, input_type=input_type, truncate=truncate, batch_size=batch_size)