nv-ingest-client 2025.8.14.dev20250814__py3-none-any.whl → 2025.8.16.dev20250816__py3-none-any.whl

nv_ingest_client/primitives/jobs/job_spec.py

@@ -150,23 +150,48 @@ class JobSpec:
 
     def add_task(self, task) -> None:
         """
-        Adds a task to the job specification.
+        Adds a task or list of tasks to the job specification.
 
         Parameters
         ----------
-        task
-            The task to add to the job specification. Assumes the task has a to_dict method.
+        task : Task or list of Task
+            The task(s) to add to the job specification. Can be a single task or a list of tasks.
+            Each task must derive from the Task class and have a to_dict method.
 
         Raises
         ------
         ValueError
-            If the task does not have a to_dict method.
+            If any task does not derive from the Task class.
+        """
+        # Handle both single tasks and lists of tasks
+        if isinstance(task, list):
+            # Process each task in the list
+            for single_task in task:
+                self._add_single_task(single_task)
+        else:
+            # Process single task
+            self._add_single_task(task)
+
+    def _add_single_task(self, task) -> None:
+        """
+        Adds a single task to the job specification with automatic task expansion.
+
+        Parameters
+        ----------
+        task : Task
+            The task to add to the job specification.
+
+        Raises
+        ------
+        ValueError
+            If the task does not derive from the Task class.
         """
         if not isinstance(task, Task):
             raise ValueError("Task must derive from nv_ingest_client.primitives.Task class")
 
         self._tasks.append(task)
 
+        # Automatic task expansion for ExtractTask
         if isinstance(task, ExtractTask) and (task._extract_tables is True):
             self._tasks.append(TableExtractionTask())
         if isinstance(task, ExtractTask) and (task._extract_charts is True):
@@ -322,11 +347,6 @@ class BatchJobSpec:
         document_type : str, optional
             The document type used to filter job specifications. If not provided, the
             `document_type` is inferred from the task, or the task is applied to all job specifications.
-
-        Raises
-        ------
-        ValueError
-            If the task does not derive from the `Task` class.
         """
         if not isinstance(task, Task):
             raise ValueError("Task must derive from nv_ingest_client.primitives.Task class")

nv_ingest_client/primitives/tasks/__init__.py

@@ -18,14 +18,18 @@ from .task_base import Task
 from .task_base import TaskType
 from .task_base import is_valid_task_type
 from .task_factory import task_factory
+from .udf import UDFTask
 
 __all__ = [
     "AudioExtractionTask",
     "CaptionTask",
     "ChartExtractionTask",
+    "DedupTask",
+    "EmbedTask",
     "ExtractTask",
-    "is_valid_task_type",
+    "FilterTask",
     "InfographicExtractionTask",
+    "is_valid_task_type",
     "SplitTask",
     "StoreEmbedTask",
     "StoreTask",
@@ -33,7 +37,5 @@ __all__ = [
     "Task",
     "task_factory",
     "TaskType",
-    "DedupTask",
-    "FilterTask",
-    "EmbedTask",
+    "UDFTask",
 ]

nv_ingest_client/primitives/tasks/audio_extraction.py

@@ -10,33 +10,19 @@ import logging
 from typing import Dict
 from typing import Optional
 
-from pydantic import BaseModel
-from pydantic import ConfigDict
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskAudioExtraction
 
 from .task_base import Task
 
 logger = logging.getLogger(__name__)
 
 
-class AudioExtractionSchema(BaseModel):
-    auth_token: Optional[str] = None
-    grpc_endpoint: Optional[str] = None
-    http_endpoint: Optional[str] = None
-    infer_protocol: Optional[str] = None
-    function_id: Optional[str] = None
-    use_ssl: Optional[bool] = None
-    ssl_cert: Optional[str] = None
-    segment_audio: Optional[bool] = None
-
-    model_config = ConfigDict(extra="forbid")
-    model_config["protected_namespaces"] = ()
-
-
 class AudioExtractionTask(Task):
     def __init__(
         self,
         auth_token: str = None,
         grpc_endpoint: str = None,
+        http_endpoint: str = None,
         infer_protocol: str = None,
         function_id: Optional[str] = None,
         use_ssl: bool = None,
@@ -45,13 +31,26 @@ class AudioExtractionTask(Task):
     ) -> None:
         super().__init__()
 
-        self._auth_token = auth_token
-        self._grpc_endpoint = grpc_endpoint
-        self._infer_protocol = infer_protocol
-        self._function_id = function_id
-        self._use_ssl = use_ssl
-        self._ssl_cert = ssl_cert
-        self._segment_audio = segment_audio
+        # Use the API schema for validation
+        validated_data = IngestTaskAudioExtraction(
+            auth_token=auth_token,
+            grpc_endpoint=grpc_endpoint,
+            http_endpoint=http_endpoint,
+            infer_protocol=infer_protocol,
+            function_id=function_id,
+            use_ssl=use_ssl,
+            ssl_cert=ssl_cert,
+            segment_audio=segment_audio,
+        )
+
+        self._auth_token = validated_data.auth_token
+        self._grpc_endpoint = validated_data.grpc_endpoint
+        self._http_endpoint = validated_data.http_endpoint
+        self._infer_protocol = validated_data.infer_protocol
+        self._function_id = validated_data.function_id
+        self._use_ssl = validated_data.use_ssl
+        self._ssl_cert = validated_data.ssl_cert
+        self._segment_audio = validated_data.segment_audio
 
     def __str__(self) -> str:
         """
@@ -64,6 +63,8 @@ class AudioExtractionTask(Task):
             info += "  auth_token: [redacted]\n"
         if self._grpc_endpoint:
             info += f"  grpc_endpoint: {self._grpc_endpoint}\n"
+        if self._http_endpoint:
+            info += f"  http_endpoint: {self._http_endpoint}\n"
         if self._infer_protocol:
             info += f"  infer_protocol: {self._infer_protocol}\n"
         if self._function_id:
@@ -89,6 +90,9 @@ class AudioExtractionTask(Task):
         if self._grpc_endpoint:
             task_properties["grpc_endpoint"] = self._grpc_endpoint
 
+        if self._http_endpoint:
+            task_properties["http_endpoint"] = self._http_endpoint
+
         if self._infer_protocol:
             task_properties["infer_protocol"] = self._infer_protocol
 

nv_ingest_client/primitives/tasks/caption.py

@@ -8,25 +8,14 @@
 
 import logging
 from typing import Dict
-from typing import Optional
 
-from pydantic import ConfigDict, BaseModel
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskCaptionSchema
 
 from .task_base import Task
 
 logger = logging.getLogger(__name__)
 
 
-class CaptionTaskSchema(BaseModel):
-    api_key: Optional[str] = None
-    endpoint_url: Optional[str] = None
-    prompt: Optional[str] = None
-    model_name: Optional[str] = None
-
-    model_config = ConfigDict(extra="forbid")
-    model_config["protected_namespaces"] = ()
-
-
 class CaptionTask(Task):
     def __init__(
         self,
@@ -37,10 +26,15 @@ class CaptionTask(Task):
     ) -> None:
         super().__init__()
 
-        self._api_key = api_key
-        self._endpoint_url = endpoint_url
-        self._prompt = prompt
-        self._model_name = model_name
+        # Use the API schema for validation
+        validated_data = IngestTaskCaptionSchema(
+            api_key=api_key, endpoint_url=endpoint_url, prompt=prompt, model_name=model_name
+        )
+
+        self._api_key = validated_data.api_key
+        self._endpoint_url = validated_data.endpoint_url
+        self._prompt = validated_data.prompt
+        self._model_name = validated_data.model_name
 
     def __str__(self) -> str:
         """

nv_ingest_client/primitives/tasks/chart_extraction.py

@@ -9,35 +9,41 @@
 import logging
 from typing import Dict
 
-from pydantic import BaseModel
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskChartExtraction
 
 from .task_base import Task
 
 logger = logging.getLogger(__name__)
 
 
-class ChartExtractionSchema(BaseModel):
-    class Config:
-        extra = "forbid"
-
-
 class ChartExtractionTask(Task):
     """
     Object for chart extraction task
     """
 
-    def __init__(self) -> None:
+    def __init__(self, params: dict = None) -> None:
         """
-        Setup Dedup Task Config
+        Setup Chart Extraction Task Config
         """
         super().__init__()
 
+        # Handle None params by converting to empty dict for backward compatibility
+        if params is None:
+            params = {}
+
+        # Use the API schema for validation
+        validated_data = IngestTaskChartExtraction(params=params)
+
+        self._params = validated_data.params
+
     def __str__(self) -> str:
         """
         Returns a string with the object's config and run time state
         """
         info = ""
-        info += "chart extraction task\n"
+        info += "Chart Extraction Task:\n"
+        if self._params:
+            info += f"  params: {self._params}\n"
         return info
 
     def to_dict(self) -> Dict:
@@ -46,7 +52,7 @@ class ChartExtractionTask(Task):
         """
 
         task_properties = {
-            "params": {},
+            "params": self._params,
         }
 
         return {"type": "chart_data_extract", "task_properties": task_properties}

nv_ingest_client/primitives/tasks/dedup.py

@@ -10,29 +10,13 @@ import logging
 from typing import Dict
 from typing import Literal
 
-from pydantic import BaseModel, field_validator
-
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskDedupSchema
 
 from .task_base import Task
 
 logger = logging.getLogger(__name__)
 
 
-class DedupTaskSchema(BaseModel):
-    content_type: str = "image"
-    filter: bool = False
-
-    @field_validator("content_type")
-    def content_type_must_be_valid(cls, v):
-        valid_criteria = ["image"]
-        if v not in valid_criteria:
-            raise ValueError(f"content_type must be one of {valid_criteria}")
-        return v
-
-    class Config:
-        extra = "forbid"
-
-
 class DedupTask(Task):
     """
     Object for document dedup task
@@ -49,8 +33,15 @@ class DedupTask(Task):
         Setup Dedup Task Config
         """
         super().__init__()
-        self._content_type = content_type
-        self._filter = filter
+
+        # Use the API schema for validation
+        validated_data = IngestTaskDedupSchema(
+            content_type=content_type,
+            params={"filter": filter},
+        )
+
+        self._content_type = validated_data.content_type
+        self._filter = validated_data.params.filter
 
     def __str__(self) -> str:
         """
@@ -58,7 +49,7 @@ class DedupTask(Task):
         """
         info = ""
         info += "Dedup Task:\n"
-        info += f"  content_type: {self._content_type}\n"
+        info += f"  content_type: {self._content_type.value}\n"
         info += f"  filter: {self._filter}\n"
         return info
 
@@ -69,7 +60,7 @@ class DedupTask(Task):
         dedup_params = {"filter": self._filter}
 
         task_properties = {
-            "content_type": self._content_type,
+            "content_type": self._content_type.value,
             "params": dedup_params,
         }
 

nv_ingest_client/primitives/tasks/embed.py

@@ -7,82 +7,15 @@
 import logging
 from typing import Any
 from typing import Dict
-from typing import Literal
 from typing import Optional
-from typing import Type
 
-from pydantic import BaseModel
-from pydantic import ConfigDict
-from pydantic import model_validator
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskEmbedSchema
 
 from .task_base import Task
 
 logger = logging.getLogger(__name__)
 
 
-class EmbedTaskSchema(BaseModel):
-    """
-    Schema for embed task configuration.
-
-    This schema contains configuration details for an embedding task,
-    including the endpoint URL, model name, API key, and error filtering flag.
-
-    Attributes
-    ----------
-    endpoint_url : Optional[str]
-        URL of the embedding endpoint. Default is None.
-    model_name : Optional[str]
-        Name of the embedding model. Default is None.
-    api_key : Optional[str]
-        API key for authentication with the embedding service. Default is None.
-    filter_errors : bool
-        Flag to indicate whether errors should be filtered. Default is False.
-    """
-
-    endpoint_url: Optional[str] = None
-    model_name: Optional[str] = None
-    api_key: Optional[str] = None
-    filter_errors: bool = False
-
-    text_elements_modality: Optional[Literal["text", "image", "text_image"]] = None
-    image_elements_modality: Optional[Literal["text", "image", "text_image"]] = None
-    structured_elements_modality: Optional[Literal["text", "image", "text_image"]] = None
-    audio_elements_modality: Optional[Literal["text"]] = None
-
-    @model_validator(mode="before")
-    def handle_deprecated_fields(cls: Type["EmbedTaskSchema"], values: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Handle deprecated fields before model validation.
-
-        This validator checks for the presence of deprecated keys ('text' and 'tables')
-        in the input dictionary and removes them. Warnings are issued if these keys are found.
-
-        Parameters
-        ----------
-        values : Dict[str, Any]
-            Input dictionary of model values.
-
-        Returns
-        -------
-        Dict[str, Any]
-            The updated dictionary with deprecated fields removed.
-        """
-        if "text" in values:
-            logger.warning(
-                "'text' parameter is deprecated and will be ignored. Future versions will remove this argument."
-            )
-            values.pop("text")
-        if "tables" in values:
-            logger.warning(
-                "'tables' parameter is deprecated and will be ignored. Future versions will remove this argument."
-            )
-            values.pop("tables")
-        return values
-
-    model_config = ConfigDict(extra="forbid")
-    model_config["protected_namespaces"] = ()
-
-
 class EmbedTask(Task):
     """
     Object for document embedding tasks.
@@ -133,14 +66,26 @@ class EmbedTask(Task):
                 "'tables' parameter is deprecated and will be ignored. Future versions will remove this argument."
             )
 
-        self._endpoint_url: Optional[str] = endpoint_url
-        self._model_name: Optional[str] = model_name
-        self._api_key: Optional[str] = api_key
-        self._filter_errors: bool = filter_errors
-        self._text_elements_modality: Optional[bool] = text_elements_modality
-        self._image_elements_modality: Optional[bool] = image_elements_modality
-        self._structured_elements_modality: Optional[bool] = structured_elements_modality
-        self._audio_elements_modality: Optional[bool] = audio_elements_modality
+        # Use the API schema for validation
+        validated_data = IngestTaskEmbedSchema(
+            endpoint_url=endpoint_url,
+            model_name=model_name,
+            api_key=api_key,
+            filter_errors=filter_errors,
+            text_elements_modality=text_elements_modality,
+            image_elements_modality=image_elements_modality,
+            structured_elements_modality=structured_elements_modality,
+            audio_elements_modality=audio_elements_modality,
+        )
+
+        self._endpoint_url = validated_data.endpoint_url
+        self._model_name = validated_data.model_name
+        self._api_key = validated_data.api_key
+        self._filter_errors = validated_data.filter_errors
+        self._text_elements_modality = validated_data.text_elements_modality
+        self._image_elements_modality = validated_data.image_elements_modality
+        self._structured_elements_modality = validated_data.structured_elements_modality
+        self._audio_elements_modality = validated_data.audio_elements_modality
 
     def __str__(self) -> str:
         """