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

nv_ingest_api/internal/schemas/extract/extract_ocr_schema.py

@@ -0,0 +1,137 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Optional
+from typing import Tuple
+
+from pydantic import field_validator, model_validator, ConfigDict, BaseModel, Field
+
+from nv_ingest_api.internal.schemas.mixins import LowercaseProtocolMixin
+
+logger = logging.getLogger(__name__)
+
+
+class OCRExtractorConfigSchema(LowercaseProtocolMixin):
+    """
+    Configuration schema for text extraction service endpoints and options.
+
+    Parameters
+    ----------
+    auth_token : Optional[str], default=None
+        Authentication token required for secure services.
+
+    ocr_endpoints : Tuple[Optional[str], Optional[str]], default=(None, None)
+        A tuple containing the gRPC and HTTP services for the ocr endpoint.
+        Either the gRPC or HTTP service can be empty, but not both.
+
+    Methods
+    -------
+    validate_endpoints(values)
+        Validates that at least one of the gRPC or HTTP services is provided for each endpoint.
+
+    Raises
+    ------
+    ValueError
+        If both gRPC and HTTP services are empty for any endpoint.
+
+    Config
+    ------
+    extra : str
+        Pydantic config option to forbid extra fields.
+    """
+
+    auth_token: Optional[str] = Field(default=None, repr=False)
+
+    ocr_endpoints: Tuple[Optional[str], Optional[str]] = (None, None)
+    ocr_infer_protocol: str = ""
+
+    nim_batch_size: int = 2
+    workers_per_progress_engine: int = 5
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_endpoints(cls, values):
+        """
+        Validates the gRPC and HTTP services for all endpoints.
+
+        Ensures that at least one service (either gRPC or HTTP) is provided
+        for each endpoint in the configuration.
+
+        Parameters
+        ----------
+        values : dict
+            Dictionary containing the values of the attributes for the class.
+
+        Returns
+        -------
+        dict
+            The validated dictionary of values.
+
+        Raises
+        ------
+        ValueError
+            If both gRPC and HTTP services are empty for any endpoint.
+        """
+
+        def clean_service(service):
+            """Set service to None if it's an empty string or contains only spaces or quotes."""
+            if service is None or not service.strip() or service.strip(" \"'") == "":
+                return None
+            return service
+
+        for endpoint_name in ["ocr_endpoints"]:
+            grpc_service, http_service = values.get(endpoint_name, (None, None))
+            grpc_service = clean_service(grpc_service)
+            http_service = clean_service(http_service)
+
+            if not grpc_service and not http_service:
+                raise ValueError(f"Both gRPC and HTTP services cannot be empty for {endpoint_name}.")
+
+            values[endpoint_name] = (grpc_service, http_service)
+
+            # Auto-infer protocol from endpoints if not specified
+            protocol_name = endpoint_name.replace("_endpoints", "_infer_protocol")
+            protocol_value = values.get(protocol_name)
+            if not protocol_value:
+                protocol_value = "http" if http_service else "grpc" if grpc_service else ""
+            values[protocol_name] = protocol_value
+
+        return values
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class OCRExtractorSchema(BaseModel):
+    """
+    Configuration schema for text extraction processing settings.
+
+    Parameters
+    ----------
+    max_queue_size : int, default=1
+        The maximum number of items allowed in the processing queue.
+
+    n_workers : int, default=2
+        The number of worker threads to use for processing.
+
+    raise_on_failure : bool, default=False
+        A flag indicating whether to raise an exception if a failure occurs during text extraction.
+
+    stage_config : Optional[OCRExtractorConfigSchema], default=None
+        Configuration for the text extraction stage, including yolox and ocr service endpoints.
+    """
+
+    max_queue_size: int = 1
+    n_workers: int = 2
+    raise_on_failure: bool = False
+
+    endpoint_config: Optional[OCRExtractorConfigSchema] = None
+
+    @field_validator("max_queue_size", "n_workers")
+    def check_positive(cls, v, field):
+        if v <= 0:
+            raise ValueError(f"{field.field_name} must be greater than 0.")
+        return v
+
+    model_config = ConfigDict(extra="forbid")

nv_ingest_api/internal/schemas/extract/extract_pdf_schema.py

@@ -0,0 +1,220 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+from typing import Optional
+from typing import Tuple
+
+from pydantic import model_validator, ConfigDict, BaseModel, Field
+
+from nv_ingest_api.internal.schemas.mixins import LowercaseProtocolMixin
+
+logger = logging.getLogger(__name__)
+
+
+class PDFiumConfigSchema(LowercaseProtocolMixin):
+    """
+    Configuration schema for PDFium endpoints and options.
+
+    Parameters
+    ----------
+    auth_token : Optional[str], default=None
+        Authentication token required for secure services.
+
+    yolox_endpoints : Tuple[str, str]
+        A tuple containing the gRPC and HTTP services for the yolox endpoint.
+        Either the gRPC or HTTP service can be empty, but not both.
+
+    Methods
+    -------
+    validate_endpoints(values)
+        Validates that at least one of the gRPC or HTTP services is provided for each endpoint.
+
+    Raises
+    ------
+    ValueError
+        If both gRPC and HTTP services are empty for any endpoint.
+
+    Config
+    ------
+    extra : str
+        Pydantic config option to forbid extra fields.
+    """
+
+    auth_token: Optional[str] = Field(default=None, repr=False)
+
+    yolox_endpoints: Tuple[Optional[str], Optional[str]] = (None, None)
+    yolox_infer_protocol: str = ""
+
+    nim_batch_size: int = 4
+    workers_per_progress_engine: int = 5
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_endpoints(cls, values):
+        """
+        Validates the gRPC and HTTP services for all endpoints.
+
+        Parameters
+        ----------
+        values : dict
+            Dictionary containing the values of the attributes for the class.
+
+        Returns
+        -------
+        dict
+            The validated dictionary of values.
+
+        Raises
+        ------
+        ValueError
+            If both gRPC and HTTP services are empty for any endpoint.
+        """
+
+        for model_name in ["yolox"]:
+            endpoint_name = f"{model_name}_endpoints"
+            grpc_service, http_service = values.get(endpoint_name, ("", ""))
+            grpc_service = _clean_service(grpc_service)
+            http_service = _clean_service(http_service)
+
+            if not grpc_service and not http_service:
+                raise ValueError(f"Both gRPC and HTTP services cannot be empty for {endpoint_name}.")
+
+            values[endpoint_name] = (grpc_service, http_service)
+
+            # Auto-infer protocol from endpoints if not specified
+            protocol_name = f"{model_name}_infer_protocol"
+            protocol_value = values.get(protocol_name)
+            if not protocol_value:
+                protocol_value = "http" if http_service else "grpc" if grpc_service else ""
+            values[protocol_name] = protocol_value
+
+        return values
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class NemotronParseConfigSchema(LowercaseProtocolMixin):
+    """
+    Configuration schema for Nemotron Parse endpoints and options.
+
+    Parameters
+    ----------
+    auth_token : Optional[str], default=None
+        Authentication token required for secure services.
+
+    nemotron_parse_endpoints : Tuple[str, str]
+        A tuple containing the gRPC and HTTP services for the nemotron_parse endpoint.
+        Either the gRPC or HTTP service can be empty, but not both.
+
+    Methods
+    -------
+    validate_endpoints(values)
+        Validates that at least one of the gRPC or HTTP services is provided for each endpoint.
+
+    Raises
+    ------
+    ValueError
+        If both gRPC and HTTP services are empty for any endpoint.
+
+    Config
+    ------
+    extra : str
+        Pydantic config option to forbid extra fields.
+    """
+
+    auth_token: Optional[str] = Field(default=None, repr=False)
+
+    yolox_endpoints: Tuple[Optional[str], Optional[str]] = (None, None)
+    yolox_infer_protocol: str = ""
+
+    nemotron_parse_endpoints: Tuple[Optional[str], Optional[str]] = (None, None)
+    nemotron_parse_infer_protocol: str = ""
+
+    nemotron_parse_model_name: str = "nvidia/nemotron-parse"
+
+    timeout: float = 300.0
+
+    workers_per_progress_engine: int = 5
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_endpoints(cls, values):
+        """
+        Validates the gRPC and HTTP services for all endpoints.
+
+        Parameters
+        ----------
+        values : dict
+            Dictionary containing the values of the attributes for the class.
+
+        Returns
+        -------
+        dict
+            The validated dictionary of values.
+
+        Raises
+        ------
+        ValueError
+            If both gRPC and HTTP services are empty for any endpoint.
+        """
+
+        for model_name in ["nemotron_parse"]:
+            endpoint_name = f"{model_name}_endpoints"
+            grpc_service, http_service = values.get(endpoint_name, ("", ""))
+            grpc_service = _clean_service(grpc_service)
+            http_service = _clean_service(http_service)
+
+            if not grpc_service and not http_service:
+                raise ValueError(f"Both gRPC and HTTP services cannot be empty for {endpoint_name}.")
+
+            values[endpoint_name] = (grpc_service, http_service)
+
+            # Auto-infer protocol from endpoints if not specified
+            protocol_name = f"{model_name}_infer_protocol"
+            protocol_value = values.get(protocol_name)
+            if not protocol_value:
+                protocol_value = "http" if http_service else "grpc" if grpc_service else ""
+            values[protocol_name] = protocol_value
+
+        return values
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class PDFExtractorSchema(BaseModel):
+    """
+    Configuration schema for the PDF extractor settings.
+
+    Parameters
+    ----------
+    max_queue_size : int, default=1
+        The maximum number of items allowed in the processing queue.
+
+    n_workers : int, default=16
+        The number of worker threads to use for processing.
+
+    raise_on_failure : bool, default=False
+        A flag indicating whether to raise an exception on processing failure.
+
+    pdfium_config : Optional[PDFiumConfigSchema], default=None
+        Configuration for the PDFium service endpoints.
+    """
+
+    max_queue_size: int = 1
+    n_workers: int = 16
+    raise_on_failure: bool = False
+
+    pdfium_config: Optional[PDFiumConfigSchema] = None
+    nemotron_parse_config: Optional[NemotronParseConfigSchema] = None
+
+    model_config = ConfigDict(extra="forbid")
+
+
+def _clean_service(service):
+    """Set service to None if it's an empty string or contains only spaces or quotes."""
+    if service is None or not service.strip() or service.strip(" \"'") == "":
+        return None
+    return service

nv_ingest_api/internal/schemas/extract/extract_pptx_schema.py

@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+from typing import Optional
+from typing import Tuple
+
+from pydantic import model_validator, ConfigDict, BaseModel, Field
+
+from nv_ingest_api.internal.schemas.extract.extract_pdf_schema import PDFiumConfigSchema
+from nv_ingest_api.internal.schemas.mixins import LowercaseProtocolMixin
+
+logger = logging.getLogger(__name__)
+
+
+class PPTXConfigSchema(LowercaseProtocolMixin):
+    """
+    Configuration schema for docx extraction endpoints and options.
+
+    Parameters
+    ----------
+    auth_token : Optional[str], default=None
+        Authentication token required for secure services.
+
+    yolox_endpoints : Tuple[str, str]
+        A tuple containing the gRPC and HTTP services for the yolox endpoint.
+        Either the gRPC or HTTP service can be empty, but not both.
+
+    Methods
+    -------
+    validate_endpoints(values)
+        Validates that at least one of the gRPC or HTTP services is provided for each endpoint.
+
+    Raises
+    ------
+    ValueError
+        If both gRPC and HTTP services are empty for any endpoint.
+
+    Config
+    ------
+    extra : str
+        Pydantic config option to forbid extra fields.
+    """
+
+    auth_token: Optional[str] = Field(default=None, repr=False)
+
+    yolox_endpoints: Tuple[Optional[str], Optional[str]] = (None, None)
+    yolox_infer_protocol: str = ""
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_endpoints(cls, values):
+        """
+        Validates the gRPC and HTTP services for all endpoints.
+
+        Parameters
+        ----------
+        values : dict
+            Dictionary containing the values of the attributes for the class.
+
+        Returns
+        -------
+        dict
+            The validated dictionary of values.
+
+        Raises
+        ------
+        ValueError
+            If both gRPC and HTTP services are empty for any endpoint.
+        """
+
+        def clean_service(service):
+            """Set service to None if it's an empty string or contains only spaces or quotes."""
+            if service is None or not service.strip() or service.strip(" \"'") == "":
+                return None
+            return service
+
+        for model_name in ["yolox"]:
+            endpoint_name = f"{model_name}_endpoints"
+            grpc_service, http_service = values.get(endpoint_name)
+            grpc_service = clean_service(grpc_service)
+            http_service = clean_service(http_service)
+
+            if not grpc_service and not http_service:
+                raise ValueError(f"Both gRPC and HTTP services cannot be empty for {endpoint_name}.")
+
+            values[endpoint_name] = (grpc_service, http_service)
+
+            # Auto-infer protocol from endpoints if not specified
+            protocol_name = f"{model_name}_infer_protocol"
+            protocol_value = values.get(protocol_name)
+            if not protocol_value:
+                protocol_value = "http" if http_service else "grpc" if grpc_service else ""
+            values[protocol_name] = protocol_value
+
+        return values
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class PPTXExtractorSchema(BaseModel):
+    """
+    Configuration schema for the PDF extractor settings.
+
+    Parameters
+    ----------
+    max_queue_size : int, default=1
+        The maximum number of items allowed in the processing queue.
+
+    n_workers : int, default=16
+        The number of worker threads to use for processing.
+
+    raise_on_failure : bool, default=False
+        A flag indicating whether to raise an exception on processing failure.
+
+    image_extraction_config: Optional[ImageConfigSchema], default=None
+        Configuration schema for the image extraction stage.
+    """
+
+    max_queue_size: int = 1
+    n_workers: int = 16
+    raise_on_failure: bool = False
+
+    pptx_extraction_config: Optional[PPTXConfigSchema] = None
+    pdfium_config: Optional[PDFiumConfigSchema] = None
+    model_config = ConfigDict(extra="forbid")

nv_ingest_api/internal/schemas/extract/extract_table_schema.py

@@ -0,0 +1,137 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+from typing import Optional
+from typing import Tuple
+
+from pydantic import field_validator, model_validator, ConfigDict, BaseModel, Field
+
+from nv_ingest_api.internal.schemas.mixins import LowercaseProtocolMixin
+
+logger = logging.getLogger(__name__)
+
+
+class TableExtractorConfigSchema(LowercaseProtocolMixin):
+    """
+    Configuration schema for the table extraction stage settings.
+
+    Parameters
+    ----------
+    auth_token : Optional[str], default=None
+        Authentication token required for secure services.
+
+    ocr_endpoints : Tuple[Optional[str], Optional[str]], default=(None, None)
+        A tuple containing the gRPC and HTTP services for the ocr endpoint.
+        Either the gRPC or HTTP service can be empty, but not both.
+
+    Methods
+    -------
+    validate_endpoints(values)
+        Validates that at least one of the gRPC or HTTP services is provided for the yolox endpoint.
+
+    Raises
+    ------
+    ValueError
+        If both gRPC and HTTP services are empty for the yolox endpoint.
+
+    Config
+    ------
+    extra : str
+        Pydantic config option to forbid extra fields.
+    """
+
+    auth_token: Optional[str] = Field(default=None, repr=False)
+
+    yolox_endpoints: Tuple[Optional[str], Optional[str]] = (None, None)
+    yolox_infer_protocol: str = ""
+
+    ocr_endpoints: Tuple[Optional[str], Optional[str]] = (None, None)
+    ocr_infer_protocol: str = ""
+
+    nim_batch_size: int = 2
+    workers_per_progress_engine: int = 5
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_endpoints(cls, values):
+        """
+        Validates the gRPC and HTTP services for the yolox endpoint.
+
+        Parameters
+        ----------
+        values : dict
+            Dictionary containing the values of the attributes for the class.
+
+        Returns
+        -------
+        dict
+            The validated dictionary of values.
+
+        Raises
+        ------
+        ValueError
+            If both gRPC and HTTP services are empty for the yolox endpoint.
+        """
+
+        def clean_service(service):
+            """Set service to None if it's an empty string or contains only spaces or quotes."""
+            if service is None or not service.strip() or service.strip(" \"'") == "":
+                return None
+            return service
+
+        for endpoint_name in ["yolox_endpoints", "ocr_endpoints"]:
+            grpc_service, http_service = values.get(endpoint_name, (None, None))
+            grpc_service = clean_service(grpc_service)
+            http_service = clean_service(http_service)
+
+            if not grpc_service and not http_service:
+                raise ValueError(f"Both gRPC and HTTP services cannot be empty for {endpoint_name}.")
+
+            values[endpoint_name] = (grpc_service, http_service)
+
+            # Auto-infer protocol from endpoints if not specified
+            protocol_name = endpoint_name.replace("_endpoints", "_infer_protocol")
+            protocol_value = values.get(protocol_name)
+            if not protocol_value:
+                protocol_value = "http" if http_service else "grpc" if grpc_service else ""
+            values[protocol_name] = protocol_value
+
+        return values
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class TableExtractorSchema(BaseModel):
+    """
+    Configuration schema for the table extraction processing settings.
+
+    Parameters
+    ----------
+    max_queue_size : int, default=1
+        The maximum number of items allowed in the processing queue.
+
+    n_workers : int, default=2
+        The number of worker threads to use for processing.
+
+    raise_on_failure : bool, default=False
+        A flag indicating whether to raise an exception if a failure occurs during table extraction.
+
+    stage_config : Optional[TableExtractorConfigSchema], default=None
+        Configuration for the table extraction stage, including yolox service endpoints.
+    """
+
+    max_queue_size: int = 1
+    n_workers: int = 2
+    raise_on_failure: bool = False
+
+    @field_validator("max_queue_size", "n_workers")
+    def check_positive(cls, v, field):
+        if v <= 0:
+            raise ValueError(f"{field.field_name} must be greater than 0.")
+        return v
+
+    endpoint_config: Optional[TableExtractorConfigSchema] = None
+    model_config = ConfigDict(extra="forbid")

nv_ingest_api/internal/schemas/message_brokers/__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/schemas/message_brokers/message_broker_client_schema.py

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from pydantic import BaseModel, Field
+from typing import Optional, Literal, Annotated
+
+
+class MessageBrokerClientSchema(BaseModel):
+    """
+    Configuration schema for message broker client connections.
+    Supports Redis or simple in-memory clients.
+    """
+
+    host: str = Field(default="redis", description="Hostname of the broker service.")
+
+    port: Annotated[int, Field(gt=0, lt=65536)] = Field(
+        default=6379, description="Port to connect to. Must be between 1 and 65535."
+    )
+
+    client_type: Literal["redis", "simple"] = Field(
+        default="redis", description="Type of broker client. Supported values: 'redis', 'simple'."
+    )
+
+    broker_params: Optional[dict] = Field(
+        default_factory=dict, description="Optional parameters passed to the broker client."
+    )
+
+    connection_timeout: Annotated[int, Field(ge=0)] = Field(
+        default=300, description="Connection timeout in seconds. Must be >= 0."
+    )
+
+    max_backoff: Annotated[int, Field(ge=0)] = Field(
+        default=300, description="Maximum backoff time in seconds. Must be >= 0."
+    )
+
+    max_retries: Annotated[int, Field(ge=0)] = Field(default=0, description="Maximum number of retries. Must be >= 0.")

nv_ingest_api/internal/schemas/message_brokers/request_schema.py

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+from typing import Optional
+
+from pydantic import ConfigDict, BaseModel
+from pydantic import Field
+
+logger = logging.getLogger(__name__)
+
+
+# Define schemas for request validation
+class PushRequestSchema(BaseModel):
+    command: str
+    queue_name: str = Field(..., min_length=1)
+    message: str = Field(..., min_length=1)
+    timeout: Optional[float] = 100  # Optional timeout for blocking push
+    model_config = ConfigDict(extra="forbid")
+
+
+class PopRequestSchema(BaseModel):
+    command: str
+    queue_name: str = Field(..., min_length=1)
+    timeout: Optional[float] = 100  # Optional timeout for blocking pop
+    model_config = ConfigDict(extra="forbid")
+
+
+class SizeRequestSchema(BaseModel):
+    command: str
+    queue_name: str = Field(..., min_length=1)
+    model_config = ConfigDict(extra="forbid")