nv-ingest-client 2025.8.13.dev20250813__py3-none-any.whl → 2025.8.15.dev20250815__py3-none-any.whl

nv_ingest_client/primitives/tasks/extract.py

@@ -12,12 +12,8 @@ from typing import Any
 from typing import Dict
 from typing import Literal
 from typing import Optional
-from typing import get_args
 
-from pydantic import BaseModel
-from pydantic import ConfigDict
-from pydantic import field_validator
-from pydantic import model_validator
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskExtractSchema
 
 from .task_base import Task
 
@@ -64,142 +60,9 @@ _Type_Extract_Method_PDF = Literal[
     "unstructured_io",
 ]
 
-_Type_Extract_Method_DOCX = Literal["python_docx", "haystack", "unstructured_local", "unstructured_service"]
+_Type_Extract_Images_Method = Literal["group", "yolox"]
 
-_Type_Extract_Method_PPTX = Literal["python_pptx", "haystack", "unstructured_local", "unstructured_service"]
-
-_Type_Extract_Method_Image = Literal["image"]
-
-_Type_Extract_Method_Audio = Literal["audio"]
-
-_Type_Extract_Method_Text = Literal["txt"]
-
-_Type_Extract_Method_Html = Literal["markitdown"]
-
-_Type_Extract_Method_Map = {
-    "bmp": get_args(_Type_Extract_Method_Image),
-    "docx": get_args(_Type_Extract_Method_DOCX),
-    "html": get_args(_Type_Extract_Method_Html),
-    "jpeg": get_args(_Type_Extract_Method_Image),
-    "jpg": get_args(_Type_Extract_Method_Image),
-    "pdf": get_args(_Type_Extract_Method_PDF),
-    "png": get_args(_Type_Extract_Method_Image),
-    "pptx": get_args(_Type_Extract_Method_PPTX),
-    "text": get_args(_Type_Extract_Method_Text),
-    "tiff": get_args(_Type_Extract_Method_Image),
-    "txt": get_args(_Type_Extract_Method_Text),
-    "mp3": get_args(_Type_Extract_Method_Audio),
-    "wav": get_args(_Type_Extract_Method_Audio),
-}
-
-_Type_Extract_Tables_Method_PDF = Literal["yolox", "pdfium", "nemoretriever_parse"]
-
-_Type_Extract_Tables_Method_DOCX = Literal["python_docx",]
-
-_Type_Extract_Tables_Method_PPTX = Literal["python_pptx",]
-
-_Type_Extract_Tables_Method_Map = {
-    "pdf": get_args(_Type_Extract_Tables_Method_PDF),
-    "docx": get_args(_Type_Extract_Tables_Method_DOCX),
-    "pptx": get_args(_Type_Extract_Tables_Method_PPTX),
-}
-
-_Type_Extract_Images_Method = Literal["simple", "group"]
-
-
-class ExtractTaskSchema(BaseModel):
-    document_type: str
-    extract_method: str = None  # Initially allow None to set a smart default
-    extract_text: bool = True
-    extract_images: bool = True
-    extract_images_method: str = "group"
-    extract_images_params: Optional[Dict[str, Any]] = None
-    extract_tables: bool = True
-    extract_tables_method: str = "yolox"
-    extract_charts: Optional[bool] = None  # Initially allow None to set a smart default
-    extract_infographics: bool = False
-    extract_page_as_image: bool = False
-    extract_audio_params: Optional[Dict[str, Any]] = None
-    text_depth: str = "document"
-    paddle_output_format: Optional[str] = None
-    table_output_format: str = "pseudo_markdown"
-
-    @model_validator(mode="after")
-    @classmethod
-    def set_default_extract_method(cls, values):
-        document_type = values.document_type.lower()  # Ensure case-insensitive comparison
-        extract_method = values.extract_method
-        paddle_output_format = values.paddle_output_format
-
-        if document_type not in _DEFAULT_EXTRACTOR_MAP:
-            raise ValueError(
-                f"Unsupported document type: {document_type}."
-                f" Supported types are: {list(_DEFAULT_EXTRACTOR_MAP.keys())}"
-            )
-
-        if extract_method is None:
-            values.extract_method = _DEFAULT_EXTRACTOR_MAP[document_type]
-
-        if paddle_output_format is not None:
-            logger.warning(
-                "`paddle_output_format` is deprecated and will be removed in a future release. "
-                "Please use `table_output_format` instead."
-            )
-            values.table_output_format = paddle_output_format
-
-        return values
-
-    @field_validator("extract_charts")
-    def set_default_extract_charts(cls, v, values):
-        # `extract_charts` is initially set to None for backward compatibility.
-        # {extract_tables: true, extract_charts: None} or {extract_tables: true, extract_charts: true} enables both
-        # table and chart extraction.
-        # {extract_tables: true, extract_charts: false} enables only the table extraction and disables chart extraction.
-        extract_charts = v
-        if extract_charts is None:
-            extract_charts = values.data.get("extract_tables")
-
-        return extract_charts
-
-    @field_validator("extract_method")
-    def extract_method_must_be_valid(cls, v, values, **kwargs):
-        document_type = values.data.get("document_type", "").lower()  # Ensure case-insensitive comparison
-
-        # Skip validation for text-like types, since they do not have 'extract' stages.
-        if document_type in ["txt", "text", "json", "md", "sh"]:
-            return
-
-        valid_methods = set(_Type_Extract_Method_Map[document_type])
-        if v not in valid_methods:
-            raise ValueError(f"extract_method must be one of {valid_methods}")
-
-        return v
-
-    @field_validator("document_type")
-    def document_type_must_be_supported(cls, v):
-        if v.lower() not in _DEFAULT_EXTRACTOR_MAP:
-            raise ValueError(
-                f"Unsupported document type '{v}'. Supported types are: {', '.join(_DEFAULT_EXTRACTOR_MAP.keys())}"
-            )
-        return v.lower()
-
-    @field_validator("extract_tables_method")
-    def extract_tables_method_must_be_valid(cls, v, values, **kwargs):
-        document_type = values.data.get("document_type", "").lower()  # Ensure case-insensitive comparison
-        valid_methods = set(_Type_Extract_Tables_Method_Map[document_type])
-        if v not in valid_methods:
-            raise ValueError(f"extract_method must be one of {valid_methods}")
-        return v
-
-    @field_validator("extract_images_method")
-    def extract_images_method_must_be_valid(cls, v):
-        if v.lower() not in get_args(_Type_Extract_Images_Method):
-            raise ValueError(
-                f"Unsupported document type '{v}'. Supported types are: {', '.join(_Type_Extract_Images_Method)}"
-            )
-        return v.lower()
-
-    model_config = ConfigDict(extra="forbid")
+_Type_Extract_Tables_Method_PDF = Literal["yolox", "paddle"]
 
 
 class ExtractTask(Task):
@@ -210,7 +73,7 @@ class ExtractTask(Task):
     def __init__(
         self,
         document_type,
-        extract_method: _Type_Extract_Method_PDF = "pdfium",
+        extract_method: _Type_Extract_Method_PDF = None,
         extract_text: bool = False,
         extract_images: bool = False,
         extract_tables: bool = False,
@@ -230,19 +93,62 @@ class ExtractTask(Task):
         """
         super().__init__()
 
-        self._document_type = document_type
+        # Set default extract_method if None
+        if extract_method is None:
+            # Handle both string and enum inputs
+            if hasattr(document_type, "value"):
+                document_type_str = document_type.value
+            else:
+                document_type_str = document_type
+            document_type_lower = document_type_str.lower()
+            if document_type_lower not in _DEFAULT_EXTRACTOR_MAP:
+                raise ValueError(
+                    f"Unsupported document type: {document_type}."
+                    f" Supported types are: {list(_DEFAULT_EXTRACTOR_MAP.keys())}"
+                )
+            extract_method = _DEFAULT_EXTRACTOR_MAP[document_type_lower]
+
+        # Set default extract_charts if None
+        if extract_charts is None:
+            extract_charts = extract_tables
+
+        # Build params dict for API schema validation
+        extract_params = {
+            "extract_text": extract_text,
+            "extract_images": extract_images,
+            "extract_images_method": extract_images_method,
+            "extract_tables": extract_tables,
+            "extract_tables_method": extract_tables_method,
+            "extract_charts": extract_charts,
+            "extract_infographics": extract_infographics,
+            "extract_page_as_image": extract_page_as_image,
+            "text_depth": text_depth,
+            "table_output_format": table_output_format,
+        }
+
+        # Add optional parameters if provided
+        if extract_images_params:
+            extract_params["extract_images_params"] = extract_images_params
+        if extract_audio_params:
+            extract_params["extract_audio_params"] = extract_audio_params
+
+        # Use the API schema for validation
+        validated_data = IngestTaskExtractSchema(
+            document_type=document_type,
+            method=extract_method,
+            params=extract_params,
+        )
+
+        # Store validated data
+        self._document_type = validated_data.document_type
+        self._extract_method = validated_data.method
         self._extract_audio_params = extract_audio_params
         self._extract_images = extract_images
-        self._extract_method = extract_method
         self._extract_tables = extract_tables
         self._extract_images_method = extract_images_method
         self._extract_images_params = extract_images_params
         self._extract_tables_method = extract_tables_method
-        # `extract_charts` is initially set to None for backward compatibility.
-        # {extract_tables: true, extract_charts: None} or {extract_tables: true, extract-charts: true} enables both
-        # table and chart extraction.
-        # {extract_tables: true, extract_charts: false} enables only the table extraction and disables chart extraction.
-        self._extract_charts = extract_charts if extract_charts is not None else extract_tables
+        self._extract_charts = extract_charts
         self._extract_infographics = extract_infographics
         self._extract_page_as_image = extract_page_as_image
         self._extract_text = extract_text
@@ -256,34 +162,27 @@ class ExtractTask(Task):
         """
         info = ""
         info += "Extract Task:\n"
-        info += f"  document type: {self._document_type}\n"
-        info += f"  extract method: {self._extract_method}\n"
-        info += f"  extract text: {self._extract_text}\n"
-        info += f"  extract images: {self._extract_images}\n"
-        info += f"  extract tables: {self._extract_tables}\n"
-        info += f"  extract charts: {self._extract_charts}\n"
-        info += f"  extract infographics: {self._extract_infographics}\n"
-        info += f"  extract page as image: {self._extract_page_as_image}\n"
-        info += f"  extract images method: {self._extract_images_method}\n"
-        info += f"  extract tables method: {self._extract_tables_method}\n"
-        info += f"  text depth: {self._text_depth}\n"
+        info += f"  document_type: {self._document_type.value}\n"
+        info += f"  extract_method: {self._extract_method}\n"
+        info += f"  extract_text: {self._extract_text}\n"
+        info += f"  extract_images: {self._extract_images}\n"
+        info += f"  extract_tables: {self._extract_tables}\n"
+        info += f"  extract_charts: {self._extract_charts}\n"
+        info += f"  extract_infographics: {self._extract_infographics}\n"
+        info += f"  extract_page_as_image: {self._extract_page_as_image}\n"
+        info += f"  text_depth: {self._text_depth}\n"
         info += f"  table_output_format: {self._table_output_format}\n"
-
-        if self._extract_images_params:
-            info += f"  extract images params: {self._extract_images_params}\n"
-        if self._extract_audio_params:
-            info += f"  extract audio params: {self._extract_audio_params}\n"
         return info
 
     def to_dict(self) -> Dict:
         """
-        Convert to a dict for submission to redis (fixme)
+        Convert to a dict for submission to redis
         """
         extract_params = {
             "extract_text": self._extract_text,
             "extract_images": self._extract_images,
-            "extract_tables": self._extract_tables,
             "extract_images_method": self._extract_images_method,
+            "extract_tables": self._extract_tables,
             "extract_tables_method": self._extract_tables_method,
             "extract_charts": self._extract_charts,
             "extract_infographics": self._extract_infographics,
@@ -306,7 +205,7 @@ class ExtractTask(Task):
 
         task_properties = {
             "method": self._extract_method,
-            "document_type": self._document_type,
+            "document_type": self._document_type.value,
             "params": extract_params,
         }
 
@@ -339,4 +238,4 @@ class ExtractTask(Task):
 
     @property
     def document_type(self):
-        return self._document_type
+        return self._document_type.value

nv_ingest_client/primitives/tasks/filter.py

@@ -11,31 +11,13 @@ from typing import Dict
 from typing import Literal
 from typing import Union
 
-from pydantic import BaseModel, field_validator
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskFilterSchema
 
 from .task_base import Task
 
 logger = logging.getLogger(__name__)
 
 
-class FilterTaskSchema(BaseModel):
-    content_type: str = "image"
-    min_size: int = 128
-    max_aspect_ratio: Union[float, int] = 5.0
-    min_aspect_ratio: Union[float, int] = 0.2
-    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 FilterTask(Task):
     """
     Object for document filter task
@@ -52,14 +34,26 @@ class FilterTask(Task):
         filter: bool = False,
     ) -> None:
         """
-        Setup Split Task Config
+        Setup Filter Task Config
         """
         super().__init__()
-        self._content_type = content_type
-        self._min_size = min_size
-        self._max_aspect_ratio = max_aspect_ratio
-        self._min_aspect_ratio = min_aspect_ratio
-        self._filter = filter
+
+        # Use the API schema for validation
+        validated_data = IngestTaskFilterSchema(
+            content_type=content_type,
+            params={
+                "min_size": min_size,
+                "max_aspect_ratio": max_aspect_ratio,
+                "min_aspect_ratio": min_aspect_ratio,
+                "filter": filter,
+            },
+        )
+
+        self._content_type = validated_data.content_type
+        self._min_size = validated_data.params.min_size
+        self._max_aspect_ratio = validated_data.params.max_aspect_ratio
+        self._min_aspect_ratio = validated_data.params.min_aspect_ratio
+        self._filter = validated_data.params.filter
 
     def __str__(self) -> str:
         """
@@ -67,7 +61,7 @@ class FilterTask(Task):
         """
         info = ""
         info += "Filter Task:\n"
-        info += f"  content_type: {self._content_type}\n"
+        info += f"  content_type: {self._content_type.value}\n"
         info += f"  min_size: {self._min_size}\n"
         info += f"  max_aspect_ratio: {self._max_aspect_ratio}\n"
         info += f"  min_aspect_ratio: {self._min_aspect_ratio}\n"
@@ -86,7 +80,7 @@ class FilterTask(Task):
         }
 
         task_properties = {
-            "content_type": self._content_type,
+            "content_type": self._content_type.value,
             "params": filter_params,
         }
 

nv_ingest_client/primitives/tasks/infographic_extraction.py

@@ -9,44 +9,47 @@
 import logging
 from typing import Dict
 
-from pydantic import BaseModel
-
-from .task_base import Task
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskInfographicExtraction
+from nv_ingest_client.primitives.tasks.task_base import Task
 
 logger = logging.getLogger(__name__)
 
 
-class InfographicExtractionSchema(BaseModel):
-    class Config:
-        extra = "forbid"
-
-
 class InfographicExtractionTask(Task):
     """
     Object for infographic extraction task
     """
 
-    def __init__(self) -> None:
+    def __init__(self, params: dict = None) -> None:
         """
-        Setup Dedup Task Config
+        Setup Infographic 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 = IngestTaskInfographicExtraction(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 += "infographic extraction task\n"
+        info += "Infographic Extraction Task:\n"
+        info += f"  params: {self._params}\n"
         return info
 
     def to_dict(self) -> Dict:
         """
         Convert to a dict for submission to redis
         """
-
         task_properties = {
-            "params": {},
+            "params": self._params,
         }
 
         return {"type": "infographic_data_extract", "task_properties": task_properties}

nv_ingest_client/primitives/tasks/split.py

@@ -8,25 +8,14 @@
 
 import logging
 from typing import Dict
-from typing import Optional
 
-from pydantic import BaseModel
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskSplitSchema
 
 from .task_base import Task
 
 logger = logging.getLogger(__name__)
 
 
-class SplitTaskSchema(BaseModel):
-    tokenizer: Optional[str] = None
-    chunk_size: int = 1024
-    chunk_overlap: int = 150
-    params: dict = {}
-
-    class Config:
-        extra = "forbid"
-
-
 class SplitTask(Task):
     """
     Object for document splitting task
@@ -37,16 +26,26 @@ class SplitTask(Task):
         tokenizer: str = None,
         chunk_size: int = 1024,
         chunk_overlap: int = 150,
-        params: dict = {},
-    ) -> None:
+        params: dict = None,
+    ):
         """
         Setup Split Task Config
         """
         super().__init__()
-        self._tokenizer = tokenizer
-        self._chunk_size = chunk_size
-        self._chunk_overlap = chunk_overlap
-        self._params = params
+
+        # Handle None params by converting to empty dict for backward compatibility
+        if params is None:
+            params = {}
+
+        # Use the API schema for validation
+        validated_data = IngestTaskSplitSchema(
+            tokenizer=tokenizer, chunk_size=chunk_size, chunk_overlap=chunk_overlap, params=params
+        )
+
+        self._tokenizer = validated_data.tokenizer
+        self._chunk_size = validated_data.chunk_size
+        self._chunk_overlap = validated_data.chunk_overlap
+        self._params = validated_data.params
 
     def __str__(self) -> str:
         """

nv_ingest_client/primitives/tasks/store.py

@@ -10,8 +10,8 @@ import logging
 from typing import Dict
 from typing import Literal
 
-from pydantic import BaseModel
-from pydantic import model_validator
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskStoreSchema
+from nv_ingest_api.internal.schemas.meta.ingest_job_schema import IngestTaskStoreEmbedSchema
 
 from .task_base import Task
 
@@ -20,28 +20,6 @@ logger = logging.getLogger(__name__)
 _DEFAULT_STORE_METHOD = "minio"
 
 
-class StoreEmbedTaskSchema(BaseModel):
-
-    class Config:
-        extra = "allow"
-
-
-class StoreTaskSchema(BaseModel):
-    store_method: str = None
-
-    @model_validator(mode="before")
-    @classmethod
-    def set_default_store_method(cls, values):
-        store_method = values.get("store_method")
-
-        if store_method is None:
-            values["store_method"] = _DEFAULT_STORE_METHOD
-        return values
-
-    class Config:
-        extra = "allow"
-
-
 class StoreTask(Task):
     """
     Object for image storage task.
@@ -64,10 +42,22 @@ class StoreTask(Task):
         """
         super().__init__()
 
-        self._structured = structured
-        self._images = images
-        self._store_method = store_method or "minio"
-        self._params = params
+        # Handle None params by converting to empty dict for backward compatibility
+        if params is None:
+            params = {}
+
+        # Merge extra_params into params for API schema compatibility
+        merged_params = {**params, **extra_params}
+
+        # Use the API schema for validation
+        validated_data = IngestTaskStoreSchema(
+            structured=structured, images=images, method=store_method or _DEFAULT_STORE_METHOD, params=merged_params
+        )
+
+        self._structured = validated_data.structured
+        self._images = validated_data.images
+        self._store_method = validated_data.method
+        self._params = validated_data.params
         self._extra_params = extra_params
 
     def __str__(self) -> str:
@@ -116,7 +106,17 @@ class StoreEmbedTask(Task):
         """
         super().__init__()
 
-        self._params = params or {}
+        # Handle None params by converting to empty dict for backward compatibility
+        if params is None:
+            params = {}
+
+        # Merge extra_params into params for API schema compatibility
+        merged_params = {**params, **extra_params}
+
+        # Use the API schema for validation
+        validated_data = IngestTaskStoreEmbedSchema(params=merged_params)
+
+        self._params = validated_data.params
         self._extra_params = extra_params
 
     def __str__(self) -> str:

nv_ingest_client/primitives/tasks/task_base.py

@@ -27,6 +27,7 @@ class TaskType(Enum):
     STORE_EMBEDDING = auto()
     TABLE_DATA_EXTRACT = auto()
     TRANSFORM = auto()
+    UDF = auto()
     VDB_UPLOAD = auto()
 
 
@@ -71,75 +72,3 @@ class Task:
         tasks that are then submitted to the redis client
         """
         return {}
-
-
-# class ExtractUnstructuredTask(ExtractTask):
-#    """
-#    Object for document unstructured extraction task
-#    extract_method = ["unstructured_local", "unstructured_service"]
-#    """
-#
-#    def __init__(
-#            self,
-#            extract_method: ExtractTask._Type_Extract_Method,
-#            document_type: ExtractTask._TypeDocumentType,
-#            api_key: str,
-#            uri: str,
-#    ) -> None:
-#        """
-#        Setup Extract Task Config
-#        """
-#        super().__init__(extract_method, document_type)
-#        self._api_key = api_key
-#        self._uri = uri
-#
-#    def __str__(self) -> str:
-#        """
-#        Returns a string with the object's config and run time state
-#        """
-#        info = ""
-#        info += super().__str__()
-#        info += f"unstructured uri: {self._uri}\n"
-#        return info
-#
-#    def to_dict(self) -> Dict:
-#        """
-#        Convert to a dict for submission to redis (fixme)
-#        """
-#        unstructured_properties = {
-#            "api_key": self._api_key,
-#            "unstructured_url": self._uri,
-#        }
-#        task_desc = super().to_dict()
-#        task_desc["task_properties"]["params"].update(unstructured_properties)
-#        return task_desc
-
-
-# class ExtractLlamaParseTask(ExtractTask):
-#    """
-#    Object for document llama extraction task
-#    extract_method = ["llama_parse"]
-#    """
-#
-#    def __init__(
-#            self,
-#            extract_method: ExtractTask._Type_Extract_Method,
-#            document_type: ExtractTask._TypeDocumentType,
-#            api_key: str,
-#    ) -> None:
-#        """
-#        Setup Extract Task Config
-#        """
-#        super().__init__(extract_method, document_type)
-#        self._api_key = api_key
-#
-#    def to_dict(self) -> Dict:
-#        """
-#        Convert to a dict for submission to redis (fixme)
-#        """
-#        llama_parse_properties = {
-#            "api_key": self._api_key,
-#        }
-#        task_desc = super().to_dict()
-#        task_desc["task_properties"]["params"].update(llama_parse_properties)
-#        return task_desc

nv_ingest_client/primitives/tasks/task_factory.py

@@ -19,6 +19,7 @@ from .store import StoreTask
 from .task_base import Task
 from .task_base import TaskType
 from .task_base import is_valid_task_type
+from .udf import UDFTask
 
 
 class TaskUnimplemented(Task):
@@ -42,6 +43,7 @@ _TASK_MAP: Dict[TaskType, Callable] = {
     TaskType.STORE_EMBEDDING: StoreEmbedTask,
     TaskType.STORE: StoreTask,
     TaskType.TRANSFORM: TaskUnimplemented,
+    TaskType.UDF: UDFTask,
 }