nv-ingest-api 2025.8.13.dev20250813__tar.gz → 2025.8.15.dev20250815__tar.gz

{nv_ingest_api-2025.8.13.dev20250813/src/nv_ingest_api.egg-info → nv_ingest_api-2025.8.15.dev20250815}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nv-ingest-api
-Version: 2025.8.13.dev20250813
+Version: 2025.8.15.dev20250815
 Summary: Python module with core document ingestion functions.
 Author-email: Jeremy Dyer <jdyer@nvidia.com>
 License:                                  Apache License

{nv_ingest_api-2025.8.13.dev20250813 → nv_ingest_api-2025.8.15.dev20250815}/src/nv_ingest_api/internal/enums/common.py

@@ -386,6 +386,40 @@ class StatusEnum(str, Enum):
     SUCCESS: str = "success"
 
 
+class PipelinePhase(int, Enum):
+    """
+    The logical phase of a pipeline stage.
+
+    Attributes
+    ----------
+    PRE_PROCESSING : int
+        Pre-processing phase.
+    EXTRACTION : int
+        Extraction phase.
+    POST_PROCESSING : int
+        Post-processing phase.
+    MUTATION : int
+        Mutation phase.
+    TRANSFORM : int
+        Transform phase.
+    RESPONSE : int
+        Response phase.
+    TELEMETRY : int
+        Telemetry phase.
+    DRAIN : int
+        Drain phase.
+    """
+
+    PRE_PROCESSING = 0
+    EXTRACTION = 1
+    POST_PROCESSING = 2
+    MUTATION = 3
+    TRANSFORM = 4
+    RESPONSE = 5
+    TELEMETRY = 6
+    DRAIN = 7
+
+
 class TableFormatEnum(str, Enum):
     """
     Enum for representing table formats.
@@ -446,6 +480,8 @@ class TaskTypeEnum(str, Enum):
         Represents a task for extracting chart data.
     INFOGRAPHIC_DATA_EXTRACT : str
         Represents a task for extracting infographic data.
+    UDF : str
+        Represents a user-defined function task.
     """
 
     AUDIO_DATA_EXTRACT: str = "audio_data_extract"
@@ -460,6 +496,7 @@ class TaskTypeEnum(str, Enum):
     STORE_EMBEDDING: str = "store_embedding"
     STORE: str = "store"
     TABLE_DATA_EXTRACT: str = "table_data_extract"
+    UDF: str = "udf"
     VDB_UPLOAD: str = "vdb_upload"
 
 

{nv_ingest_api-2025.8.13.dev20250813 → nv_ingest_api-2025.8.15.dev20250815}/src/nv_ingest_api/internal/extract/image/image_extractor.py

@@ -108,8 +108,12 @@ def _decode_and_extract_from_image(
             f"decode_and_extract: Extracting image content using image_extraction_config: "
             f"{validated_extraction_config}"
         )
+        # Ensure we pass the correct nested config type (ImageConfigSchema) to helpers.
+        # Some callers provide the full ImageExtractorSchema; extract its inner image_extraction_config.
         if validated_extraction_config is not None:
-            extract_params["image_extraction_config"] = validated_extraction_config
+            inner_cfg = getattr(validated_extraction_config, "image_extraction_config", validated_extraction_config)
+            if inner_cfg is not None:
+                extract_params["image_extraction_config"] = inner_cfg
 
         if execution_trace_log is not None:
             extract_params["trace_info"] = execution_trace_log

nv_ingest_api-2025.8.15.dev20250815/src/nv_ingest_api/internal/meta/udf.py

@@ -0,0 +1,232 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import hashlib
+import inspect
+import logging
+import time
+from typing import Any, Dict, List, Optional
+from dataclasses import dataclass
+
+from nv_ingest_api.internal.primitives.ingest_control_message import IngestControlMessage, remove_all_tasks_by_type
+from nv_ingest_api.internal.schemas.meta.udf import UDFStageSchema
+from nv_ingest_api.util.imports.callable_signatures import ingest_callable_signature
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CachedUDF:
+    """Cached UDF function with metadata"""
+
+    function: callable
+    function_name: str
+    signature_validated: bool
+    created_at: float
+    last_used: float
+    use_count: int
+
+
+class UDFCache:
+    """LRU cache for compiled and validated UDF functions"""
+
+    def __init__(self, max_size: int = 128, ttl_seconds: Optional[int] = 3600):
+        self.max_size = max_size
+        self.ttl_seconds = ttl_seconds
+        self.cache: Dict[str, CachedUDF] = {}
+        self.access_order: List[str] = []  # For LRU tracking
+
+    def _generate_cache_key(self, udf_function_str: str, udf_function_name: str) -> str:
+        """Generate cache key from UDF string and function name"""
+        content = f"{udf_function_str.strip()}:{udf_function_name}"
+        return hashlib.sha256(content.encode()).hexdigest()
+
+    def _evict_lru(self):
+        """Remove least recently used item"""
+        if self.access_order:
+            lru_key = self.access_order.pop(0)
+            self.cache.pop(lru_key, None)
+
+    def _cleanup_expired(self):
+        """Remove expired entries if TTL is configured"""
+        if not self.ttl_seconds:
+            return
+
+        current_time = time.time()
+        expired_keys = [
+            key for key, cached_udf in self.cache.items() if current_time - cached_udf.created_at > self.ttl_seconds
+        ]
+
+        for key in expired_keys:
+            self.cache.pop(key, None)
+            if key in self.access_order:
+                self.access_order.remove(key)
+
+    def get(self, udf_function_str: str, udf_function_name: str) -> Optional[CachedUDF]:
+        """Get cached UDF function if available"""
+        self._cleanup_expired()
+
+        cache_key = self._generate_cache_key(udf_function_str, udf_function_name)
+
+        if cache_key in self.cache:
+            # Update access tracking
+            if cache_key in self.access_order:
+                self.access_order.remove(cache_key)
+            self.access_order.append(cache_key)
+
+            # Update usage stats
+            cached_udf = self.cache[cache_key]
+            cached_udf.last_used = time.time()
+            cached_udf.use_count += 1
+
+            return cached_udf
+
+        return None
+
+    def put(
+        self, udf_function_str: str, udf_function_name: str, function: callable, signature_validated: bool = True
+    ) -> str:
+        """Cache a compiled and validated UDF function"""
+        cache_key = self._generate_cache_key(udf_function_str, udf_function_name)
+
+        # Evict LRU if at capacity
+        while len(self.cache) >= self.max_size:
+            self._evict_lru()
+
+        current_time = time.time()
+        cached_udf = CachedUDF(
+            function=function,
+            function_name=udf_function_name,
+            signature_validated=signature_validated,
+            created_at=current_time,
+            last_used=current_time,
+            use_count=1,
+        )
+
+        self.cache[cache_key] = cached_udf
+        self.access_order.append(cache_key)
+
+        return cache_key
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get cache statistics"""
+        total_uses = sum(udf.use_count for udf in self.cache.values())
+        most_used = max(self.cache.values(), key=lambda x: x.use_count, default=None)
+        return {
+            "size": len(self.cache),
+            "max_size": self.max_size,
+            "total_uses": total_uses,
+            "most_used_function": most_used.function_name if most_used else None,
+            "most_used_count": most_used.use_count if most_used else 0,
+        }
+
+
+# Global cache instance
+_udf_cache = UDFCache(max_size=128, ttl_seconds=3600)
+
+
+def compile_and_validate_udf(udf_function_str: str, udf_function_name: str, task_num: int) -> callable:
+    """Compile and validate UDF function (extracted for caching)"""
+    # Execute the UDF function string in a controlled namespace
+    namespace: Dict[str, Any] = {}
+    try:
+        exec(udf_function_str, namespace)
+    except Exception as e:
+        raise ValueError(f"UDF task {task_num} failed to execute: {str(e)}")
+
+    # Extract the specified function from the namespace
+    if udf_function_name in namespace and callable(namespace[udf_function_name]):
+        udf_function = namespace[udf_function_name]
+    else:
+        raise ValueError(f"UDF task {task_num}: Specified UDF function '{udf_function_name}' not found or not callable")
+
+    # Validate the UDF function signature
+    try:
+        ingest_callable_signature(inspect.signature(udf_function))
+    except Exception as e:
+        raise ValueError(f"UDF task {task_num} has invalid function signature: {str(e)}")
+
+    return udf_function
+
+
+def get_udf_cache_stats() -> Dict[str, Any]:
+    """Get UDF cache performance statistics"""
+    return _udf_cache.get_stats()
+
+
+def udf_stage_callable_fn(control_message: IngestControlMessage, stage_config: UDFStageSchema) -> IngestControlMessage:
+    """
+    UDF stage callable function that processes UDF tasks in a control message.
+
+    This function extracts all UDF tasks from the control message and executes them sequentially.
+
+    Parameters
+    ----------
+    control_message : IngestControlMessage
+        The control message containing UDF tasks to process
+    stage_config : UDFStageSchema
+        Configuration for the UDF stage
+
+    Returns
+    -------
+    IngestControlMessage
+        The control message after processing all UDF tasks
+    """
+    logger.debug("Starting UDF stage processing")
+
+    # Extract all UDF tasks from control message using free function
+    try:
+        all_task_configs = remove_all_tasks_by_type(control_message, "udf")
+    except ValueError:
+        # No UDF tasks found
+        if stage_config.ignore_empty_udf:
+            logger.debug("No UDF tasks found, ignoring as configured")
+            return control_message
+        else:
+            raise ValueError("No UDF tasks found in control message")
+
+    # Process each UDF task sequentially
+    for task_num, task_config in enumerate(all_task_configs, 1):
+        logger.debug(f"Processing UDF task {task_num} of {len(all_task_configs)}")
+
+        # Get UDF function string and function name from task properties
+        udf_function_str = task_config.get("udf_function", "").strip()
+        udf_function_name = task_config.get("udf_function_name", "").strip()
+
+        # Skip empty UDF functions if configured to ignore them
+        if not udf_function_str:
+            if stage_config.ignore_empty_udf:
+                logger.debug(f"UDF task {task_num} has empty function, skipping as configured")
+                continue
+            else:
+                raise ValueError(f"UDF task {task_num} has empty function string")
+
+        # Validate that function name is provided
+        if not udf_function_name:
+            raise ValueError(f"UDF task {task_num} missing required 'udf_function_name' property")
+
+        # Check if UDF function is cached
+        cached_udf = _udf_cache.get(udf_function_str, udf_function_name)
+        if cached_udf:
+            udf_function = cached_udf.function
+        else:
+            # Compile and validate UDF function
+            udf_function = compile_and_validate_udf(udf_function_str, udf_function_name, task_num)
+            # Cache the compiled UDF function
+            _udf_cache.put(udf_function_str, udf_function_name, udf_function)
+
+        # Execute the UDF function with the control message
+        try:
+            control_message = udf_function(control_message)
+        except Exception as e:
+            raise ValueError(f"UDF task {task_num} execution failed: {str(e)}")
+
+        # Validate that the UDF function returned an IngestControlMessage
+        if not isinstance(control_message, IngestControlMessage):
+            raise ValueError(f"UDF task {task_num} must return an IngestControlMessage, got {type(control_message)}")
+
+        logger.debug(f"UDF task {task_num} completed successfully")
+
+    logger.debug(f"UDF stage processing completed. Processed {len(all_task_configs)} UDF tasks")
+    return control_message

{nv_ingest_api-2025.8.13.dev20250813 → nv_ingest_api-2025.8.15.dev20250815}/src/nv_ingest_api/internal/primitives/ingest_control_message.py

@@ -5,10 +5,11 @@
 import copy
 import re
 from datetime import datetime
+from collections import defaultdict
+from typing import Any, Dict, Generator, List, Optional, Union
 
 import logging
 import pandas as pd
-from typing import Any, Dict, Generator, Union
 
 from nv_ingest_api.internal.primitives.control_message_task import ControlMessageTask
 
@@ -55,6 +56,52 @@ def remove_task_by_type(ctrl_msg, task: str):
     return removed_task.properties
 
 
+def remove_all_tasks_by_type(ctrl_msg, task: str):
+    """
+    Remove all tasks from the control message by matching their type.
+
+    This function iterates over the tasks in the control message, finds all tasks
+    whose type matches the provided task string, removes them, and returns their
+    properties as a list.
+
+    Parameters
+    ----------
+    ctrl_msg : IngestControlMessage
+        The control message from which to remove the tasks.
+    task : str
+        The task type to remove.
+
+    Returns
+    -------
+    list[dict]
+        A list of dictionaries of properties for all removed tasks.
+
+    Raises
+    ------
+    ValueError
+        If no tasks with the given type are found.
+    """
+    matching_tasks = []
+
+    # Find all tasks with matching type
+    for t in ctrl_msg.get_tasks():
+        if t.type == task:
+            matching_tasks.append(t)
+
+    if not matching_tasks:
+        err_msg = f"process_control_message: No tasks of type '{task}' found in control message."
+        logger.error(err_msg)
+        raise ValueError(err_msg)
+
+    # Remove all matching tasks and collect their properties
+    removed_task_properties = []
+    for task_obj in matching_tasks:
+        removed_task = ctrl_msg.remove_task(task_obj.id)
+        removed_task_properties.append(removed_task.properties)
+
+    return removed_task_properties
+
+
 class IngestControlMessage:
     """
     A control message class for ingesting tasks and managing associated metadata,
@@ -65,47 +112,41 @@ class IngestControlMessage:
         """
         Initialize a new IngestControlMessage instance.
         """
-        self._tasks: Dict[str, ControlMessageTask] = {}
+        self._tasks: Dict[str, List[ControlMessageTask]] = defaultdict(list)
         self._metadata: Dict[str, Any] = {}
         self._timestamps: Dict[str, datetime] = {}
-        self._payload: pd.DataFrame = pd.DataFrame()
+        self._payload: Optional[pd.DataFrame] = None
         self._config: Dict[str, Any] = {}
 
     def add_task(self, task: ControlMessageTask):
         """
-        Add a task to the control message, keyed by the task's unique 'id'.
-
-        Raises
-        ------
-        ValueError
-            If a task with the same 'id' already exists.
+        Add a task to the control message. Multiple tasks with the same ID are supported.
         """
-        if task.id in self._tasks:
-            raise ValueError(f"Task with id '{task.id}' already exists. Tasks must be unique.")
-        self._tasks[task.id] = task
+        self._tasks[task.id].append(task)
 
     def get_tasks(self) -> Generator[ControlMessageTask, None, None]:
         """
         Return all tasks as a generator.
         """
-        yield from self._tasks.values()
+        for task_list in self._tasks.values():
+            yield from task_list
 
     def has_task(self, task_id: str) -> bool:
         """
-        Check if a task with the given ID exists.
+        Check if any tasks with the given ID exist.
         """
-        return task_id in self._tasks
+        return task_id in self._tasks and len(self._tasks[task_id]) > 0
 
     def remove_task(self, task_id: str) -> ControlMessageTask:
         """
-        Remove a task from the control message. Logs a warning if the task does not exist.
+        Remove the first task with the given ID. Warns if no task exists.
         """
-        if task_id in self._tasks:
-            _task = self._tasks[task_id]
-
-            del self._tasks[task_id]
-
-            return _task
+        if task_id in self._tasks and self._tasks[task_id]:
+            task = self._tasks[task_id].pop(0)
+            # Clean up empty lists
+            if not self._tasks[task_id]:
+                del self._tasks[task_id]
+            return task
         else:
             raise RuntimeError(f"Attempted to remove non-existent task with id: {task_id}")
 

{nv_ingest_api-2025.8.13.dev20250813 → nv_ingest_api-2025.8.15.dev20250815}/src/nv_ingest_api/internal/primitives/tracing/tagging.py

@@ -5,24 +5,30 @@
 
 import functools
 import inspect
+import logging
 import string
 from datetime import datetime
+from typing import Optional
 
+logger = logging.getLogger(__name__)
 
-def traceable(trace_name=None):
+
+def traceable(trace_name: Optional[str] = None):
     """
     A decorator that adds entry and exit trace timestamps to a IngestControlMessage's metadata
     based on the presence of a 'config::add_trace_tagging' flag.
 
     This decorator checks if the 'config::add_trace_tagging' flag is set to True in the
     message's metadata. If so, it records the entry and exit timestamps of the function
-    execution, using either a provided custom trace name or the function's name by default.
+    execution, using either a provided custom trace name, auto-detected stage name from
+    self.stage_name, or the function's name as fallback.
 
     Parameters
     ----------
     trace_name : str, optional
-        A custom name for the trace entries in the message metadata. If not provided, the
-        function's name is used by default.
+        A custom name for the trace entries in the message metadata. If not provided,
+        attempts to use self.stage_name from the decorated method's instance,
+        falling back to the function's name if neither is available.
 
     Returns
     -------
@@ -41,26 +47,48 @@ def traceable(trace_name=None):
     - 'trace::entry::<trace_name>': The timestamp marking the function's entry.
     - 'trace::exit::<trace_name>': The timestamp marking the function's exit.
 
-    Example
-    -------
-    Applying the decorator without a custom trace name:
+    Examples
+    --------
+    Automatic stage name detection (recommended):
 
-    >>> @traceable()
-    ... def process_message(message):
+    >>> @traceable()  # Uses self.stage_name automatically
+    ... def process_message(self, message):
     ...     pass
 
-    Applying the decorator with a custom trace name on a class method:
+    Explicit trace name (override):
 
-    >>> class Processor:
-    ...     @traceable(trace_name="CustomTrace")
-    ...     def process(self, message):
-    ...         pass
+    >>> @traceable("custom_trace")
+    ... def process_message(self, message):
+    ...     pass
+
+    Function without instance (uses function name):
+
+    >>> @traceable()
+    ... def process_message(message):
+    ...     pass
     """
 
     def decorator_trace_tagging(func):
         @functools.wraps(func)
         def wrapper_trace_tagging(*args, **kwargs):
             ts_fetched = datetime.now()
+
+            # Determine the trace name to use
+            resolved_trace_name = trace_name
+
+            # If no explicit trace_name provided, try to get it from self.stage_name
+            if resolved_trace_name is None and len(args) >= 1:
+                stage_instance = args[0]  # 'self' in method calls
+                if hasattr(stage_instance, "stage_name") and stage_instance.stage_name:
+                    resolved_trace_name = stage_instance.stage_name
+                    logger.debug(f"Using auto-detected trace name: '{resolved_trace_name}'")
+                else:
+                    resolved_trace_name = func.__name__
+                    logger.debug(f"Using function name as trace name: '{resolved_trace_name}'")
+            elif resolved_trace_name is None:
+                resolved_trace_name = func.__name__
+                logger.debug(f"Using function name as trace name: '{resolved_trace_name}'")
+
             # Determine which argument is the message.
             if hasattr(args[0], "has_metadata"):
                 message = args[0]
@@ -73,7 +101,7 @@ def traceable(trace_name=None):
                 message.get_metadata("config::add_trace_tagging") is True
             )
 
-            trace_prefix = trace_name if trace_name else func.__name__
+            trace_prefix = resolved_trace_name
 
             if do_trace_tagging:
                 ts_send = message.get_timestamp("latency::ts_send")
@@ -199,3 +227,62 @@ def traceable_func(trace_name=None, dedupe=True):
         return wrapper_inject_trace_info
 
     return decorator_inject_trace_info
+
+
+def set_trace_timestamps_with_parent_context(control_message, execution_trace_log: dict, parent_name: str, logger=None):
+    """
+    Set trace timestamps on a control message with proper parent-child context.
+
+    This utility function processes trace timestamps from an execution_trace_log and
+    ensures that child traces are properly namespaced under their parent context.
+    This resolves OpenTelemetry span hierarchy issues where child spans cannot
+    find their expected parent contexts.
+
+    Parameters
+    ----------
+    control_message : IngestControlMessage
+        The control message to set timestamps on
+    execution_trace_log : dict
+        Dictionary of trace keys to timestamp values from internal operations
+    parent_name : str
+        The parent stage name to use as context for child traces
+    logger : logging.Logger, optional
+        Logger for debug output of key transformations
+
+    Examples
+    --------
+    Basic usage in a stage:
+
+    >>> execution_trace_log = {"trace::entry::yolox_inference": ts1, "trace::exit::yolox_inference": ts2}
+    >>> set_trace_timestamps_with_parent_context(
+    ...     control_message, execution_trace_log, "pdf_extractor", logger
+    ... )
+
+    This transforms:
+    - trace::entry::yolox_inference -> trace::entry::pdf_extractor::yolox_inference
+    - trace::exit::yolox_inference  -> trace::exit::pdf_extractor::yolox_inference
+    """
+    if not execution_trace_log:
+        return
+
+    for key, ts in execution_trace_log.items():
+        enhanced_key = key
+
+        # Check if this is a child trace that needs parent context
+        if key.startswith("trace::") and "::" in key:
+            # Parse the trace key to extract the base trace name
+            parts = key.split("::")
+            if len(parts) >= 3:  # e.g., ["trace", "entry", "yolox_inference"]
+                trace_type = parts[1]  # "entry" or "exit"
+                child_name = "::".join(parts[2:])  # everything after trace::entry:: or trace::exit::
+
+                # Only rewrite if it doesn't already include the parent context
+                if not child_name.startswith(f"{parent_name}::"):
+                    # Rewrite to include parent context: trace::entry::pdf_extractor::yolox_inference
+                    enhanced_key = f"trace::{trace_type}::{parent_name}::{child_name}"
+
+                    if logger:
+                        logger.debug(f"Enhanced trace key: {key} -> {enhanced_key}")
+
+        # Set the timestamp with the (possibly enhanced) key
+        control_message.set_timestamp(enhanced_key, ts)