nv-ingest-api 25.4.2__py3-none-any.whl → 25.6.1__py3-none-any.whl

nv_ingest_api/internal/primitives/nim/model_interface/text_embedding.py

@@ -5,9 +5,9 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 from nv_ingest_api.internal.primitives.nim import ModelInterface
+import numpy as np
 
 
-# Assume ModelInterface is defined elsewhere in the project.
 class EmbeddingModelInterface(ModelInterface):
     """
     An interface for handling inference with an embedding model endpoint.
@@ -22,20 +22,13 @@ class EmbeddingModelInterface(ModelInterface):
 
     def prepare_data_for_inference(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
-        Prepare input data for embedding inference. Ensures that a 'prompts' key is provided
-        and that its value is a list.
-
-        Raises
-        ------
-        KeyError
-            If the 'prompts' key is missing.
+        Prepare input data for embedding inference. Returns a list of strings representing the text to be embedded.
         """
         if "prompts" not in data:
             raise KeyError("Input data must include 'prompts'.")
-        # Ensure the prompts are in list format.
         if not isinstance(data["prompts"], list):
             data["prompts"] = [data["prompts"]]
-        return data
+        return {"prompts": data["prompts"]}
 
     def format_input(
         self, data: Dict[str, Any], protocol: str, max_batch_size: int, **kwargs
@@ -63,29 +56,32 @@ class EmbeddingModelInterface(ModelInterface):
               - payloads is a list of JSON-serializable payload dictionaries.
               - batch_data_list is a list of dictionaries containing the key "prompts" corresponding to each batch.
         """
-        if protocol != "http":
-            raise ValueError("EmbeddingModelInterface only supports HTTP protocol.")
-
-        prompts = data.get("prompts", [])
 
         def chunk_list(lst, chunk_size):
+            lst = lst["prompts"]
             return [lst[i : i + chunk_size] for i in range(0, len(lst), chunk_size)]
 
-        batches = chunk_list(prompts, max_batch_size)
-        payloads = []
-        batch_data_list = []
-        for batch in batches:
-            payload = {
-                "model": kwargs.get("model_name"),
-                "input": batch,
-                "encoding_format": kwargs.get("encoding_format", "float"),
-                "extra_body": {
-                    "input_type": kwargs.get("input_type", "query"),
+        batches = chunk_list(data, max_batch_size)
+        if protocol == "http":
+            payloads = []
+            batch_data_list = []
+            for batch in batches:
+                payload = {
+                    "model": kwargs.get("model_name"),
+                    "input": batch,
+                    "encoding_format": kwargs.get("encoding_format", "float"),
+                    "input_type": kwargs.get("input_type", "passage"),
                     "truncate": kwargs.get("truncate", "NONE"),
-                },
-            }
-            payloads.append(payload)
-            batch_data_list.append({"prompts": batch})
+                }
+                payloads.append(payload)
+                batch_data_list.append({"prompts": batch})
+        elif protocol == "grpc":
+            payloads = []
+            batch_data_list = []
+            for batch in batches:
+                text_np = np.array([[text.encode("utf-8")] for text in batch], dtype=np.object_)
+                payloads.append(text_np)
+                batch_data_list.append({"prompts": batch})
         return payloads, batch_data_list
 
     def parse_output(self, response: Any, protocol: str, data: Optional[Dict[str, Any]] = None, **kwargs) -> Any:
@@ -108,16 +104,17 @@ class EmbeddingModelInterface(ModelInterface):
         list
             A list of generated embeddings extracted from the response.
         """
-        if protocol != "http":
-            raise ValueError("EmbeddingModelInterface only supports HTTP protocol.")
-        if isinstance(response, dict):
-            embeddings = response.get("data")
-            if not embeddings:
-                raise RuntimeError("Unexpected response format: 'data' key is missing or empty.")
-            # Each item in embeddings is expected to have an 'embedding' field.
-            return [item.get("embedding", None) for item in embeddings]
-        else:
-            return [str(response)]
+        if protocol == "http":
+            if isinstance(response, dict):
+                embeddings = response.get("data")
+                if not embeddings:
+                    raise RuntimeError("Unexpected response format: 'data' key is missing or empty.")
+                # Each item in embeddings is expected to have an 'embedding' field.
+                return [item.get("embedding", None) for item in embeddings]
+            else:
+                return [str(response)]
+        elif protocol == "grpc":
+            return [res.flatten() for res in response]
 
     def process_inference_results(self, output: Any, protocol: str, **kwargs) -> Any:
         """

nv_ingest_api/internal/primitives/nim/model_interface/yolox.py

@@ -709,7 +709,13 @@ def postprocess_results(
             raise ValueError(f"Error in postprocessing {result.shape} and {original_image_shape}: {e}")
 
         for box, score, label in zip(bboxes, scores, labels):
-            class_name = class_labels[int(label)]
+            # TODO(Devin): Sometimes we get back unexpected class labels?
+            if (label < 0) or (label >= len(class_labels)):
+                logger.warning(f"Invalid class label {label} found in postprocessing")
+                continue
+            else:
+                class_name = class_labels[int(label)]
+
             annotation_dict[class_name].append([round(float(x), 4) for x in np.concatenate((box, [score]))])
 
         out.append(annotation_dict)

nv_ingest_api/internal/primitives/nim/nim_client.py

@@ -129,7 +129,7 @@ class NimClient:
         """
         if self.protocol == "grpc":
             logger.debug("Performing gRPC inference for a batch...")
-            response = self._grpc_infer(batch_input, model_name)
+            response = self._grpc_infer(batch_input, model_name, **kwargs)
             logger.debug("gRPC inference received response for a batch")
         elif self.protocol == "http":
             logger.debug("Performing HTTP inference for a batch...")
@@ -221,7 +221,7 @@ class NimClient:
 
         return all_results
 
-    def _grpc_infer(self, formatted_input: np.ndarray, model_name: str) -> np.ndarray:
+    def _grpc_infer(self, formatted_input: np.ndarray, model_name: str, **kwargs) -> np.ndarray:
         """
         Perform inference using the gRPC protocol.
 
@@ -238,16 +238,24 @@ class NimClient:
             The output of the model as a numpy array.
         """
 
-        input_tensors = [grpcclient.InferInput("input", formatted_input.shape, datatype="FP32")]
-        input_tensors[0].set_data_from_numpy(formatted_input)
+        parameters = kwargs.get("parameters", {})
+        output_names = kwargs.get("outputs", ["output"])
+        dtype = kwargs.get("dtype", "FP32")
+        input_name = kwargs.get("input_name", "input")
 
-        outputs = [grpcclient.InferRequestedOutput("output")]
-        response = self.client.infer(model_name=model_name, inputs=input_tensors, outputs=outputs)
-        logger.debug(f"gRPC inference response: {response}")
+        input_tensors = grpcclient.InferInput(input_name, formatted_input.shape, datatype=dtype)
+        input_tensors.set_data_from_numpy(formatted_input)
 
-        # TODO(self.client.has_error(response)) => raise error
+        outputs = [grpcclient.InferRequestedOutput(output_name) for output_name in output_names]
+        response = self.client.infer(
+            model_name=model_name, parameters=parameters, inputs=[input_tensors], outputs=outputs
+        )
+        logger.debug(f"gRPC inference response: {response}")
 
-        return response.as_numpy("output")
+        if len(outputs) == 1:
+            return response.as_numpy(outputs[0].name())
+        else:
+            return [response.as_numpy(output.name()) for output in outputs]
 
     def _http_infer(self, formatted_input: dict) -> dict:
         """

nv_ingest_api/internal/primitives/tracing/tagging.py

@@ -31,13 +31,15 @@ def traceable(trace_name=None):
 
     Notes
     -----
-    The decorated function must accept a IngestControlMessage object as its first argument. The
-    IngestControlMessage object must implement `has_metadata`, `get_metadata`, and `set_metadata`
-    methods used by the decorator to check for the trace tagging flag and to add trace metadata.
+    The decorated function must accept a IngestControlMessage object as one of its arguments.
+    For a regular function, this is expected to be the first argument; for a class method,
+    this is expected to be the second argument (after 'self'). The IngestControlMessage object
+    must implement `has_metadata`, `get_metadata`, and `set_metadata` methods used by the decorator
+    to check for the trace tagging flag and to add trace metadata.
 
     The trace metadata added by the decorator includes two entries:
-    - 'trace::entry::<trace_name>': The monotonic timestamp marking the function's entry.
-    - 'trace::exit::<trace_name>': The monotonic timestamp marking the function's exit.
+    - 'trace::entry::<trace_name>': The timestamp marking the function's entry.
+    - 'trace::exit::<trace_name>': The timestamp marking the function's exit.
 
     Example
     -------
@@ -47,23 +49,25 @@ def traceable(trace_name=None):
     ... def process_message(message):
     ...     pass
 
-    Applying the decorator with a custom trace name:
-
-    >>> @traceable(custom_trace_name="CustomTraceName")
-    ... def process_message(message):
-    ...     pass
-
-    In both examples, `process_message` will have entry and exit timestamps added to the
-    IngestControlMessage's metadata if 'config::add_trace_tagging' is True.
+    Applying the decorator with a custom trace name on a class method:
 
+    >>> class Processor:
+    ...     @traceable(trace_name="CustomTrace")
+    ...     def process(self, message):
+    ...         pass
     """
 
     def decorator_trace_tagging(func):
         @functools.wraps(func)
         def wrapper_trace_tagging(*args, **kwargs):
-            # Assuming the first argument is always the message
             ts_fetched = datetime.now()
-            message = args[0]
+            # Determine which argument is the message.
+            if hasattr(args[0], "has_metadata"):
+                message = args[0]
+            elif len(args) > 1 and hasattr(args[1], "has_metadata"):
+                message = args[1]
+            else:
+                raise ValueError("traceable decorator could not find a message argument with 'has_metadata()'")
 
             do_trace_tagging = (message.has_metadata("config::add_trace_tagging") is True) and (
                 message.get_metadata("config::add_trace_tagging") is True
@@ -79,7 +83,7 @@ def traceable(trace_name=None):
                     message.set_timestamp(f"trace::entry::{trace_prefix}_channel_in", ts_send)
                     message.set_timestamp(f"trace::exit::{trace_prefix}_channel_in", ts_fetched)
 
-            # Call the decorated function
+            # Call the decorated function.
             result = func(*args, **kwargs)
 
             if do_trace_tagging:

nv_ingest_api/internal/schemas/extract/extract_chart_schema.py

@@ -129,7 +129,7 @@ class ChartExtractorSchema(BaseModel):
     @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 10.")
+            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_html_schema.py

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import logging
+
+from pydantic import ConfigDict, BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class HtmlExtractorSchema(BaseModel):
+    """
+    Configuration schema for the Html 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.
+
+    """
+
+    max_queue_size: int = 1
+    n_workers: int = 16
+    raise_on_failure: bool = False
+
+    model_config = ConfigDict(extra="forbid")

nv_ingest_api/internal/schemas/extract/extract_infographic_schema.py

@@ -122,7 +122,7 @@ class InfographicExtractorSchema(BaseModel):
     @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 10.")
+            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

@@ -131,7 +131,7 @@ class NemoRetrieverParseConfigSchema(BaseModel):
     nemoretriever_parse_endpoints: Tuple[Optional[str], Optional[str]] = (None, None)
     nemoretriever_parse_infer_protocol: str = ""
 
-    model_name: str = "nvidia/nemoretriever-parse"
+    nemoretriever_parse_model_name: str = "nvidia/nemoretriever-parse"
 
     timeout: float = 300.0
 

nv_ingest_api/internal/schemas/extract/extract_table_schema.py

@@ -122,7 +122,7 @@ class TableExtractorSchema(BaseModel):
     @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 10.")
+            raise ValueError(f"{field.field_name} must be greater than 0.")
         return v
 
     endpoint_config: Optional[TableExtractorConfigSchema] = None

nv_ingest_api/internal/schemas/message_brokers/message_broker_client_schema.py

@@ -2,22 +2,36 @@
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from pydantic import BaseModel, Field
+from typing import Optional, Literal, Annotated
 
-from typing import Optional, Literal
 
-from pydantic import Field, BaseModel
-from typing_extensions import 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.")
 
-class MessageBrokerClientSchema(BaseModel):
-    host: str = "redis"
-    port: Annotated[int, Field(gt=0, lt=65536)] = 6379
+    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."
+    )
 
-    # Update this for new broker types
-    client_type: Literal["redis", "simple"] = "redis"  # Restrict to 'redis' or 'simple'
+    connection_timeout: Annotated[int, Field(ge=0)] = Field(
+        default=300, description="Connection timeout in seconds. Must be >= 0."
+    )
 
-    broker_params: Optional[dict] = Field(default_factory=dict)
+    max_backoff: Annotated[int, Field(ge=0)] = Field(
+        default=300, description="Maximum backoff time in seconds. Must be >= 0."
+    )
 
-    connection_timeout: Optional[Annotated[int, Field(ge=0)]] = 300
-    max_backoff: Optional[Annotated[int, Field(ge=0)]] = 300
-    max_retries: Optional[Annotated[int, Field(ge=0)]] = 0
+    max_retries: Annotated[int, Field(ge=0)] = Field(default=0, description="Maximum number of retries. Must be >= 0.")

nv_ingest_api/internal/schemas/meta/ingest_job_schema.py

@@ -160,29 +160,40 @@ class IngestTaskSchema(BaseModelNoExt):
     @model_validator(mode="before")
     @classmethod
     def check_task_properties_type(cls, values):
-        task_type, task_properties = values.get("type"), values.get("task_properties", {})
-        if task_type and task_properties:
-            expected_type = {
-                TaskTypeEnum.CAPTION: IngestTaskCaptionSchema,
-                TaskTypeEnum.DEDUP: IngestTaskDedupSchema,
-                TaskTypeEnum.EMBED: IngestTaskEmbedSchema,
-                TaskTypeEnum.EXTRACT: IngestTaskExtractSchema,
-                TaskTypeEnum.FILTER: IngestTaskFilterSchema,  # Extend mapping as necessary
-                TaskTypeEnum.SPLIT: IngestTaskSplitSchema,
-                TaskTypeEnum.STORE_EMBEDDING: IngestTaskStoreEmbedSchema,
-                TaskTypeEnum.STORE: IngestTaskStoreSchema,
-                TaskTypeEnum.VDB_UPLOAD: IngestTaskVdbUploadSchema,
-                TaskTypeEnum.AUDIO_DATA_EXTRACT: IngestTaskAudioExtraction,
-                TaskTypeEnum.TABLE_DATA_EXTRACT: IngestTaskTableExtraction,
-                TaskTypeEnum.CHART_DATA_EXTRACT: IngestTaskChartExtraction,
-                TaskTypeEnum.INFOGRAPHIC_DATA_EXTRACT: IngestTaskInfographicExtraction,
-            }.get(
-                task_type
-            )  # Removed .upper()
-
-            # Validate task_properties against the expected schema.
-            validated_task_properties = expected_type(**task_properties)
-            values["task_properties"] = validated_task_properties
+        task_type = values.get("type")
+        task_properties = values.get("task_properties", {})
+
+        # Ensure task_type is lowercased and converted to enum early
+        if isinstance(task_type, str):
+            task_type = task_type.lower()
+            try:
+                task_type = TaskTypeEnum(task_type)
+            except ValueError:
+                raise ValueError(f"{task_type} is not a valid TaskTypeEnum value")
+
+        task_type_to_schema = {
+            TaskTypeEnum.CAPTION: IngestTaskCaptionSchema,
+            TaskTypeEnum.DEDUP: IngestTaskDedupSchema,
+            TaskTypeEnum.EMBED: IngestTaskEmbedSchema,
+            TaskTypeEnum.EXTRACT: IngestTaskExtractSchema,
+            TaskTypeEnum.FILTER: IngestTaskFilterSchema,
+            TaskTypeEnum.SPLIT: IngestTaskSplitSchema,
+            TaskTypeEnum.STORE_EMBEDDING: IngestTaskStoreEmbedSchema,
+            TaskTypeEnum.STORE: IngestTaskStoreSchema,
+            TaskTypeEnum.VDB_UPLOAD: IngestTaskVdbUploadSchema,
+            TaskTypeEnum.AUDIO_DATA_EXTRACT: IngestTaskAudioExtraction,
+            TaskTypeEnum.TABLE_DATA_EXTRACT: IngestTaskTableExtraction,
+            TaskTypeEnum.CHART_DATA_EXTRACT: IngestTaskChartExtraction,
+            TaskTypeEnum.INFOGRAPHIC_DATA_EXTRACT: IngestTaskInfographicExtraction,
+        }
+
+        expected_schema_cls = task_type_to_schema.get(task_type)
+        if expected_schema_cls is None:
+            raise ValueError(f"Unsupported or missing task_type '{task_type}'")
+
+        validated_task_properties = expected_schema_cls(**task_properties)
+        values["type"] = task_type  # ensure type is now always the enum
+        values["task_properties"] = validated_task_properties
         return values
 
     @field_validator("type", mode="before")

nv_ingest_api/internal/schemas/transform/transform_text_embedding_schema.py

@@ -5,7 +5,7 @@
 
 import logging
 
-from pydantic import ConfigDict, BaseModel
+from pydantic import ConfigDict, BaseModel, Field
 
 from nv_ingest_api.util.logging.configuration import LogLevel
 
@@ -13,13 +13,14 @@ logger = logging.getLogger(__name__)
 
 
 class TextEmbeddingSchema(BaseModel):
-    api_key: str = "api_key"
-    batch_size: int = 4
-    embedding_model: str = "nvidia/nv-embedqa-e5-v5"
-    embedding_nim_endpoint: str = "http://embedding:8000/v1"
-    encoding_format: str = "float"
-    httpx_log_level: LogLevel = LogLevel.WARNING
-    input_type: str = "passage"
-    raise_on_failure: bool = False
-    truncate: str = "END"
+    api_key: str = Field(default="api_key")
+    batch_size: int = Field(default=4)
+    embedding_model: str = Field(default="nvidia/llama-3.2-nv-embedqa-1b-v2")
+    embedding_nim_endpoint: str = Field(default="http://embedding:8000/v1")
+    encoding_format: str = Field(default="float")
+    httpx_log_level: LogLevel = Field(default=LogLevel.WARNING)
+    input_type: str = Field(default="passage")
+    raise_on_failure: bool = Field(default=False)
+    truncate: str = Field(default="END")
+
     model_config = ConfigDict(extra="forbid")

nv_ingest_api/internal/schemas/transform/transform_text_splitter_schema.py

@@ -2,21 +2,23 @@
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from pydantic import Field, BaseModel, field_validator
+from pydantic import Field, BaseModel, field_validator, ConfigDict
 
 from typing import Optional
 
-from typing_extensions import Annotated
-
 
 class TextSplitterSchema(BaseModel):
     tokenizer: Optional[str] = None
-    chunk_size: Annotated[int, Field(gt=0)] = 1024
-    chunk_overlap: Annotated[int, Field(ge=0)] = 150
+    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")
-    def check_chunk_overlap(cls, v, values, **kwargs):
-        if v is not None and "chunk_size" in values.data and v >= values.data["chunk_size"]:
+    @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/image_upload.py

@@ -116,6 +116,7 @@ def _upload_images_to_minio(df: pd.DataFrame, params: Dict[str, Any]) -> pd.Data
             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