nv-ingest-api 26.1.0rc4__py3-none-any.whl

nv_ingest_api/internal/schemas/transform/transform_text_splitter_schema.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from pydantic import Field, BaseModel, field_validator, ConfigDict
+
+from typing import Optional
+
+
+class TextSplitterSchema(BaseModel):
+    tokenizer: Optional[str] = None
+    chunk_size: int = Field(default=1024, gt=0)
+    chunk_overlap: int = Field(default=150, ge=0)
+    raise_on_failure: bool = False
+
+    @field_validator("chunk_overlap")
+    @classmethod
+    def check_chunk_overlap(cls, v, values):
+        chunk_size = values.data.get("chunk_size")
+        if chunk_size is not None and v >= chunk_size:
+            raise ValueError("chunk_overlap must be less than chunk_size")
+        return v
+
+    model_config = ConfigDict(extra="forbid")

nv_ingest_api/internal/store/__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_api/internal/store/embed_text_upload.py

@@ -0,0 +1,236 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+from typing import Any, Union, Optional
+from typing import Dict
+
+import pandas as pd
+from minio import Minio
+from pymilvus import Collection
+from pymilvus import connections
+from pymilvus.bulk_writer.constants import BulkFileType
+from pymilvus.bulk_writer.remote_bulk_writer import RemoteBulkWriter
+from pydantic import BaseModel
+
+from nv_ingest_api.internal.enums.common import ContentTypeEnum
+from nv_ingest_api.internal.schemas.store.store_embedding_schema import EmbeddingStorageSchema
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_ENDPOINT = os.environ.get("MINIO_INTERNAL_ADDRESS", "minio:9000")
+_DEFAULT_BUCKET_NAME = os.environ.get("MINIO_BUCKET", "nv-ingest")
+
+
+def _upload_text_embeddings(df_store_ledger: pd.DataFrame, task_config: Dict[str, Any]) -> pd.DataFrame:
+    """
+    Uploads embeddings to MinIO for contents (e.g., images) contained in a DataFrame.
+    The image metadata in the "metadata" column is updated with the URL (or path) of the uploaded data.
+
+    This function performs the following steps:
+      1. Initializes a MinIO client using the provided task configuration parameters.
+      2. Connects to a Milvus instance and retrieves the collection schema.
+      3. Ensures that the target bucket exists (creating it if necessary).
+      4. Configures a RemoteBulkWriter to upload embedding data in PARQUET format.
+      5. Iterates over each row in the DataFrame, updates the metadata with the bucket path, and appends
+         rows to the writer if an embedding is present.
+      6. Commits the writer, finalizing the upload process.
+
+    Parameters
+    ----------
+    df_store_ledger : pd.DataFrame
+        DataFrame containing the data to upload. Each row is expected to have:
+          - A "metadata" column (a dictionary) that includes keys such as "content", "embedding",
+            "source_metadata", and "content_metadata".
+          - A "document_type" column indicating the type of document (e.g., IMAGE, STRUCTURED).
+    task_config : Dict[str, Any]
+        Dictionary of parameters for the upload. Expected keys include:
+          - "minio_access_key": Optional[str]
+                Access key for MinIO.
+          - "minio_secret_key": Optional[str]
+                Secret key for MinIO.
+          - "minio_endpoint": str, default _DEFAULT_ENDPOINT
+                MinIO endpoint URL.
+          - "minio_bucket_name": str, default _DEFAULT_BUCKET_NAME
+                Name of the bucket in MinIO.
+          - "minio_bucket_path": str, default "embeddings"
+                Path within the bucket where embeddings are stored.
+          - "minio_session_token": Optional[str]
+                (Optional) Session token for MinIO.
+          - "minio_secure": bool, default False
+                Whether to use a secure connection to MinIO.
+          - "minio_region": Optional[str]
+                (Optional) Region for the MinIO service.
+          - "milvus_address": str, default "milvus"
+                Address of the Milvus service.
+          - "milvus_uri": str, default "http://milvus:19530"
+                URI for Milvus.
+          - "milvus_host": str, default "milvus"
+                Host for Milvus.
+          - "milvus_port": int, default 19530
+                Port for Milvus.
+          - "collection_name": str, default "nv_ingest_collection"
+                Name of the Milvus collection from which to retrieve the schema.
+
+    Returns
+    -------
+    pd.DataFrame
+        The input DataFrame with updated "metadata" columns containing the uploaded embedding URL
+        (or bucket path).
+
+    Raises
+    ------
+    Exception
+        Propagates any exception encountered during the upload process, wrapping it with additional context.
+    """
+    try:
+        # Retrieve connection parameters for MinIO
+        minio_access_key: Optional[str] = task_config.get("minio_access_key")
+        minio_secret_key: Optional[str] = task_config.get("minio_secret_key")
+        minio_endpoint: str = task_config.get("minio_endpoint", _DEFAULT_ENDPOINT)
+        minio_bucket_name: str = task_config.get("minio_bucket_name", _DEFAULT_BUCKET_NAME)
+        minio_bucket_path: str = task_config.get("minio_bucket_path", "embeddings")
+
+        # Retrieve connection parameters for Milvus
+        milvus_address: str = task_config.get("milvus_address", "milvus")
+        milvus_uri: str = task_config.get("milvus_uri", "http://milvus:19530")
+        milvus_host: str = task_config.get("milvus_host", "milvus")
+        milvus_port: int = task_config.get("milvus_port", 19530)
+        milvus_collection_name: str = task_config.get("collection_name", "nv_ingest_collection")
+
+        # Initialize MinIO client
+        client = Minio(
+            minio_endpoint,
+            access_key=minio_access_key,
+            secret_key=minio_secret_key,
+            session_token=task_config.get("minio_session_token"),
+            secure=task_config.get("minio_secure", False),
+            region=task_config.get("minio_region"),
+        )
+
+        # Connect to Milvus and retrieve collection schema
+        connections.connect(
+            address=milvus_address,
+            uri=f"{milvus_uri}:{milvus_port}",
+            host=milvus_host,
+            port=milvus_port,
+        )
+        schema = Collection(milvus_collection_name).schema
+
+        # Ensure bucket exists
+        if not client.bucket_exists(minio_bucket_name):
+            client.make_bucket(minio_bucket_name)
+            logger.debug("Created bucket %s", minio_bucket_name)
+        else:
+            logger.debug("Bucket %s already exists", minio_bucket_name)
+
+        # Setup connection parameters for RemoteBulkWriter
+        conn = RemoteBulkWriter.ConnectParam(
+            endpoint=minio_endpoint,
+            access_key=minio_access_key,
+            secret_key=minio_secret_key,
+            bucket_name=minio_bucket_name,
+            secure=False,
+        )
+        writer = RemoteBulkWriter(
+            schema=schema,
+            remote_path=minio_bucket_path,
+            connect_param=conn,
+            file_type=BulkFileType.PARQUET,
+        )
+
+        # Process each row in the DataFrame
+        for idx, row in df_store_ledger.iterrows():
+            metadata: Dict[str, Any] = row["metadata"].copy()
+            # Update embedding metadata with the bucket path
+            metadata["embedding_metadata"] = {"uploaded_embedding_url": minio_bucket_path}
+
+            doc_type = row["document_type"]
+            content_replace: bool = doc_type in [ContentTypeEnum.IMAGE, ContentTypeEnum.STRUCTURED]
+            location: str = metadata["source_metadata"]["source_location"]
+            content = metadata["content"]
+
+            # If an embedding exists, update metadata and append the row for upload
+            if metadata.get("embedding") is not None:
+                logger.error(f"row type: {doc_type} -  {location} -  {len(content)}")
+                df_store_ledger.at[idx, "metadata"] = metadata
+
+                writer.append_row(
+                    {
+                        "text": location if content_replace else content,
+                        "source": metadata["source_metadata"],
+                        "content_metadata": metadata["content_metadata"],
+                        "vector": metadata["embedding"],
+                    }
+                )
+
+        writer.commit()
+        return df_store_ledger
+
+    except Exception as e:
+        err_msg = f"upload_embeddings: Error uploading embeddings. Original error: {e}"
+        logger.error(err_msg, exc_info=True)
+        raise type(e)(err_msg) from e
+
+
+def store_text_embeddings_internal(
+    df_store_ledger: pd.DataFrame,
+    task_config: Union[BaseModel, Dict[str, Any]],
+    store_config: EmbeddingStorageSchema,
+    execution_trace_log: Optional[Dict[str, Any]] = None,
+) -> pd.DataFrame:
+    """
+    Stores embeddings by uploading content from a DataFrame to MinIO.
+
+    This function prepares the necessary parameters for the upload based on the task configuration,
+    invokes the upload routine, and returns the updated DataFrame.
+
+    Parameters
+    ----------
+    df_store_ledger : pd.DataFrame
+        DataFrame containing the data whose embeddings need to be stored.
+    task_config : Union[BaseModel, Dict[str, Any]]
+        Task configuration. If it is a Pydantic model, it will be converted to a dictionary.
+    store_config : Dict[str, Any]
+        Configuration parameters for storage (not directly used in the current implementation).
+    execution_trace_log : Optional[Dict[str, Any]], default=None
+        Optional dictionary for trace logging information.
+
+    Returns
+    -------
+    pd.DataFrame
+        The updated DataFrame after embeddings have been uploaded and metadata updated.
+
+    Raises
+    ------
+    Exception
+        If any error occurs during the storage process, it is logged and re-raised with additional context.
+    """
+
+    _ = store_config  # Unused
+    _ = execution_trace_log  # Unused
+
+    try:
+        # Convert Pydantic model to dict if necessary
+        if isinstance(task_config, BaseModel):
+            task_config = task_config.model_dump()
+
+        # Set content types for embeddings and update params
+        content_types = {ContentTypeEnum.EMBEDDING: True}
+        params: Dict[str, Any] = task_config.get("params", {})
+        params["content_types"] = content_types
+
+        # Perform the upload of embeddings
+        df_store_ledger = _upload_text_embeddings(df_store_ledger, params)
+
+        result, execution_trace_log = df_store_ledger, {}
+        _ = execution_trace_log  # Unused
+
+        return result
+
+    except Exception as e:
+        err_msg = f"_store_embeddings: Failed to store embeddings: {e}"
+        logger.error(err_msg, exc_info=True)
+        raise type(e)(err_msg) from e

nv_ingest_api/internal/store/image_upload.py

@@ -0,0 +1,251 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import base64
+import logging
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+import pandas as pd
+from upath import UPath
+
+from nv_ingest_api.internal.enums.common import ContentTypeEnum
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_storage_root(storage_uri: str, storage_options: Dict[str, Any]) -> Tuple[UPath, str]:
+    """
+    Construct a UPath instance rooted at the configured URI and return both the root path and protocol.
+    """
+    storage_root = UPath(storage_uri, **storage_options)
+    protocol = storage_root._url.scheme or "file"
+    return storage_root, protocol
+
+
+def _extract_image_type(doc_type: Any, metadata: Dict[str, Any]) -> str:
+    """
+    Determine the image type to use when writing the decoded content based on the document type.
+    """
+
+    def _normalize(raw_value: Any, default: str = "png") -> str:
+        if raw_value is None:
+            return default
+        if hasattr(raw_value, "value"):
+            return str(raw_value.value).lower()
+        string_value = str(raw_value).strip()
+        return string_value.lower() if string_value else default
+
+    if doc_type == ContentTypeEnum.IMAGE:
+        image_metadata = metadata.get("image_metadata", {})
+        return _normalize(image_metadata.get("image_type"))
+
+    if doc_type == ContentTypeEnum.STRUCTURED:
+        table_metadata = metadata.get("table_metadata", {})
+        return _normalize(table_metadata.get("image_type"))
+
+    return "png"
+
+
+def _build_destination_path(storage_root: UPath, source_id: str, row_index: int, image_type: str) -> Tuple[UPath, str]:
+    """
+    Build the destination UPath for the decoded content and return both the destination and relative key.
+    """
+    safe_source_name = os.path.basename(source_id.rstrip("/")) or "source"
+    clean_source_name = safe_source_name.replace("/", "_")
+
+    destination: UPath = storage_root / clean_source_name / f"{row_index}.{image_type}"
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    relative_key = destination.relative_to(storage_root).as_posix()
+    return destination, relative_key
+
+
+def _upload_images_via_fsspec(df: pd.DataFrame, params: Dict[str, Any]) -> pd.DataFrame:
+    """
+    Identifies content within a DataFrame and persists it using an fsspec-compatible filesystem, updating
+    metadata with the resulting URIs.
+
+    This function iterates over rows of the provided DataFrame. For rows whose "document_type" is listed
+    in the provided 'content_types' configuration, it decodes the base64-encoded content, writes the object
+    via fsspec/UPath, and updates the metadata with the resolved URL. Errors during individual row processing
+    are logged and skipped so the process continues for remaining rows.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        The DataFrame containing rows with content and associated metadata.
+    params : Dict[str, Any]
+        A flat dictionary of configuration parameters for the upload. Expected keys include:
+            - "content_types": Dict mapping document types to booleans.
+            - "storage_uri": Base URI (file://, s3://, etc.) where images should be written.
+            - "storage_options": Optional dictionary forwarded to UPath/fsspec constructors.
+            - "public_base_url": Optional HTTP(s) base used to surface stored objects.
+
+    Returns
+    -------
+    pd.DataFrame
+        The updated DataFrame with metadata reflecting the uploaded URLs. Rows that encountered errors
+        during processing will remain unchanged.
+
+    Raises
+    ------
+    ValueError
+        If the required "content_types" key is missing or is not a dictionary.
+    Exception
+        Propagates any critical exceptions not handled at the row level.
+    """
+    # Validate required configuration
+    content_types = params.get("content_types")
+    if not isinstance(content_types, dict):
+        raise ValueError("Invalid configuration: 'content_types' must be provided as a dictionary in params")
+
+    storage_uri: Optional[str] = params.get("storage_uri")
+    if not storage_uri or not storage_uri.strip():
+        raise ValueError("`storage_uri` must be provided in task params.")
+
+    storage_options: Dict[str, Any] = params.get("storage_options") or {}
+    public_base_url: Optional[str] = params.get("public_base_url")
+
+    storage_root, protocol = _resolve_storage_root(storage_uri, storage_options)
+
+    # Process each row and attempt to upload images
+    for idx, row in df.iterrows():
+        try:
+            doc_type = row.get("document_type")
+            if doc_type not in content_types:
+                continue
+
+            metadata = row.get("metadata")
+            if not isinstance(metadata, dict):
+                logger.error("Row %s: 'metadata' is not a dictionary", idx)
+                continue
+
+            # Validate required metadata fields
+            if "content" not in metadata:
+                logger.error("Row %s: missing 'content' in metadata", idx)
+                continue
+
+            if "source_metadata" not in metadata or not isinstance(metadata["source_metadata"], dict):
+                logger.error("Row %s: missing or invalid 'source_metadata' in metadata", idx)
+                continue
+
+            source_metadata = metadata["source_metadata"]
+            if "source_id" not in source_metadata:
+                logger.error("Row %s: missing 'source_id' in source_metadata", idx)
+                continue
+
+            # Decode the content from base64
+            content = base64.b64decode(metadata["content"].encode())
+            source_id = source_metadata["source_id"]
+
+            image_type = _extract_image_type(doc_type, metadata)
+
+            # Construct destination file path
+            destination, relative_key = _build_destination_path(
+                storage_root=storage_root,
+                source_id=source_id,
+                row_index=idx,
+                image_type=image_type,
+            )
+            with destination.open("wb") as target_file:
+                target_file.write(content)
+
+            destination_uri = destination.as_uri()
+            public_url: Optional[str] = None
+            if public_base_url:
+                public_url = f"{public_base_url.rstrip('/')}/{relative_key}"
+
+            primary_uri = public_url or destination_uri
+            source_metadata["source_location"] = primary_uri
+
+            local_uri: Optional[str] = None
+            if protocol == "file":
+                local_uri = destination.path
+                source_metadata["local_source_location"] = local_uri
+
+            if doc_type == ContentTypeEnum.IMAGE:
+                logger.debug("Persisting image data for row %s", idx)
+                image_metadata = metadata.get("image_metadata", {})
+                if public_url is not None:
+                    image_metadata["uploaded_image_url"] = public_url
+                if local_uri is not None:
+                    image_metadata["uploaded_image_local_path"] = local_uri
+                metadata["image_metadata"] = image_metadata
+            elif doc_type == ContentTypeEnum.STRUCTURED:
+                logger.debug("Persisting structured image data for row %s", idx)
+                table_metadata = metadata.get("table_metadata", {})
+                if public_url is not None:
+                    table_metadata["uploaded_image_url"] = public_url
+                if local_uri is not None:
+                    table_metadata["uploaded_image_local_path"] = local_uri
+                metadata["table_metadata"] = table_metadata
+
+            df.at[idx, "metadata"] = metadata
+
+        except Exception as e:
+            logger.exception("Failed to process row %s: %s", idx, e)
+            # Continue processing the remaining rows
+
+    return df
+
+
+def store_images_to_minio_internal(
+    df_storage_ledger: pd.DataFrame,
+    task_config: Dict[str, Any],
+    storage_config: Dict[str, Any],
+    execution_trace_log: Optional[List[Any]] = None,
+) -> pd.DataFrame:
+    """
+    Processes a storage ledger DataFrame to persist images (and structured content) via an fsspec-compatible
+    filesystem.
+
+    This function validates the input DataFrame and task configuration, then creates a mask to select rows
+    where the "document_type" is among the desired types specified in the configuration. If matching rows are
+    found, it calls the internal upload function to process and update the DataFrame; otherwise, it returns the
+    original DataFrame unmodified.
+
+    Parameters
+    ----------
+    df_storage_ledger : pd.DataFrame
+        The DataFrame containing storage ledger information, which must include at least the columns
+        "document_type" and "metadata".
+    task_config : Dict[str, Any]
+        A flat dictionary containing configuration parameters for image storage. Expected to include the key
+        "content_types" (a dict mapping document types to booleans) along with `storage_uri`,
+        `storage_options`, and optional presentation hints such as `public_base_url`.
+    storage_config : Dict[str, Any]
+        A dictionary reserved for additional storage configuration (currently unused).
+    execution_trace_log : Optional[List[Any]], optional
+        An optional list for capturing execution trace details (currently unused), by default None.
+
+    Returns
+    -------
+    pd.DataFrame
+        The updated DataFrame after attempting to upload images for rows with matching document types. Rows
+        that do not match remain unchanged.
+
+    Raises
+    ------
+    ValueError
+        If the input DataFrame is missing required columns or if the task configuration is invalid.
+    """
+    # Validate that required keys and columns exist
+    if "content_types" not in task_config or not isinstance(task_config["content_types"], dict):
+        raise ValueError("Task configuration must include a valid 'content_types' dictionary.")
+
+    if "document_type" not in df_storage_ledger.columns:
+        raise ValueError("Input DataFrame must contain a 'document_type' column.")
+
+    content_types = task_config["content_types"]
+
+    # Create a mask for rows where "document_type" is one of the desired types
+    storage_obj_mask = df_storage_ledger["document_type"].isin(list(content_types.keys()))
+    if (~storage_obj_mask).all():
+        logger.debug("No storage objects matching %s found in the DataFrame.", content_types)
+        return df_storage_ledger
+
+    result, execution_trace_log = _upload_images_via_fsspec(df_storage_ledger, task_config), {}
+    _ = execution_trace_log
+
+    return result

nv_ingest_api/internal/transform/__init__.py

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