nv-ingest 2025.6.2.dev20250602__tar.gz → 2025.7.7.dev20250707__tar.gz

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nv-ingest
-Version: 2025.6.2.dev20250602
+Version: 2025.7.7.dev20250707
 Summary: Python module for multimodal document ingestion
 Author-email: Jeremy Dyer <jdyer@nvidia.com>
 License:                                  Apache License
@@ -240,13 +240,13 @@ Requires-Dist: python-docx>=1.1.2
 Requires-Dist: python-dotenv>=1.0.1
 Requires-Dist: python-pptx>=1.0.2
 Requires-Dist: prometheus-client
-Requires-Dist: torch==2.4.1
+Requires-Dist: torch>=2.4.1
 Requires-Dist: ray[all]>=2.37.0
 Requires-Dist: redis>=5.2.1
 Requires-Dist: requests>=2.28.2
 Requires-Dist: scikit-learn>=1.6.0
 Requires-Dist: scipy>=1.15.1
-Requires-Dist: setuptools>=58.2.0
+Requires-Dist: setuptools>=78.1.1
 Requires-Dist: tabulate>=0.9.0
 Requires-Dist: torchvision
 Requires-Dist: torchaudio
@@ -259,7 +259,7 @@ Requires-Dist: opencv-python
 Requires-Dist: pymilvus>=2.5.10
 Requires-Dist: pymilvus[bulk_writer,model]
 Requires-Dist: tritonclient
-Requires-Dist: nvidia-riva-client>=2.18.0
+Requires-Dist: nvidia-riva-client==2.20.0
 Requires-Dist: unstructured-client
 Requires-Dist: markitdown
 Dynamic: license-file

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/examples/pipeline_test_harness.py

@@ -63,7 +63,7 @@ def get_nim_service(env_var_prefix):
         "",
     )
     auth_token = os.environ.get(
-        "NVIDIA_BUILD_API_KEY",
+        "NVIDIA_API_KEY",
         "",
     ) or os.environ.get(
         "NGC_API_KEY",
@@ -151,11 +151,11 @@ if __name__ == "__main__":
     os.environ["PADDLE_INFER_PROTOCOL"] = "grpc"
     os.environ["NEMORETRIEVER_PARSE_HTTP_ENDPOINT"] = "https://integrate.api.nvidia.com/v1/chat/completions"
     os.environ["VLM_CAPTION_ENDPOINT"] = "https://integrate.api.nvidia.com/v1/chat/completions"
-    os.environ["VLM_CAPTION_MODEL_NAME"] = "meta/llama-3.2-11b-vision-instruct"
+    os.environ["VLM_CAPTION_MODEL_NAME"] = "nvidia/llama-3.1-nemotron-nano-vl-8b-v1"
     logger.info("Environment variables set.")
 
     image_caption_endpoint_url = "https://integrate.api.nvidia.com/v1/chat/completions"
-    image_caption_model_name = "meta/llama-3.2-11b-vision-instruct"
+    model_name = "nvidia/llama-3.1-nemotron-nano-vl-8b-v1"
     yolox_grpc, yolox_http, yolox_auth, yolox_protocol = get_nim_service("yolox")
     (
         yolox_table_structure_grpc,
@@ -228,7 +228,7 @@ if __name__ == "__main__":
     image_caption_config = {
         "api_key": yolox_auth,
         "endpoint_url": image_caption_endpoint_url,
-        "image_caption_model_name": image_caption_model_name,
+        "model_name": model_name,
         "prompt": "Caption the content of this image:",
     }
     logger.info("Service configuration retrieved from get_nim_service and environment variables.")

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/primitives/pipeline_topology.py

@@ -555,7 +555,7 @@ class PipelineTopology:
             return None
 
     def get_connections(self) -> Dict[str, List[Tuple[str, int]]]:
-        """Returns a shallow copy of the connections dictionary."""
+        """Returns a shallow copy of the connection dictionary."""
         with self._lock:
             # Shallow copy is usually sufficient here as tuples are immutable
             return self._connections.copy()
@@ -571,7 +571,7 @@ class PipelineTopology:
             return len(self._stage_actors.get(stage_name, []))
 
     def get_edge_queues(self) -> Dict[str, Tuple[Any, int]]:
-        """Returns a shallow copy of the edge queues dictionary."""
+        """Returns a shallow copy of the edge queues' dictionary."""
         with self._lock:
             return self._edge_queues.copy()
 

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/primitives/ray_pipeline.py

@@ -9,6 +9,7 @@ import threading
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from dataclasses import dataclass
+from types import FunctionType
 
 import psutil
 import uuid
@@ -24,6 +25,9 @@ import time
 from nv_ingest.framework.orchestration.ray.primitives.pipeline_topology import PipelineTopology, StageInfo
 from nv_ingest.framework.orchestration.ray.primitives.ray_stat_collector import RayStatsCollector
 from nv_ingest.framework.orchestration.ray.util.pipeline.pid_controller import PIDController, ResourceConstraintManager
+from nv_ingest.framework.orchestration.ray.util.pipeline.tools import wrap_callable_as_stage
+from nv_ingest_api.util.imports.callable_signatures import ingest_stage_callable_signature
+from nv_ingest_api.util.imports.dynamic_resolvers import resolve_callable_from_path
 
 logger = logging.getLogger(__name__)
 
@@ -43,7 +47,7 @@ class PipelineInterface(ABC):
         Parameters
         ----------
         monitor_poll_interval : float
-            Interval in seconds for monitoring poll (default: 5.0).
+            Interval in seconds for the monitoring poll (default: 5.0).
         scaling_poll_interval : float
             Interval in seconds for scaling decisions (default: 30.0).
         """
@@ -270,7 +274,7 @@ class RayPipeline(PipelineInterface):
 
         logger.info("RayStatsCollector initialized using StatsConfig.")
 
-    # --- Accessor Methods for Stats Collector (and internal use) ---
+    # --- Accessor Methods for Stat Collector (and internal use) ---
 
     def __del__(self):
         try:
@@ -428,15 +432,39 @@ class RayPipeline(PipelineInterface):
         return self
 
     def add_stage(
-        self, *, name: str, stage_actor: Any, config: BaseModel, min_replicas: int = 0, max_replicas: int = 1
+        self,
+        *,
+        name: str,
+        stage_actor: Any,
+        config: BaseModel,
+        min_replicas: int = 0,
+        max_replicas: int = 1,
     ) -> "RayPipeline":
         if min_replicas < 0:
             logger.warning(f"Stage '{name}': min_replicas cannot be negative. Overriding to 0.")
             min_replicas = 0
+
+        resolved_actor = stage_actor
+
+        # Support module path (e.g., "mypkg.mymodule:my_lambda")
+        if isinstance(stage_actor, str):
+            resolved_actor = resolve_callable_from_path(
+                callable_path=stage_actor, signature_schema=ingest_stage_callable_signature
+            )
+
+        # Wrap callables
+        if isinstance(resolved_actor, FunctionType):
+            schema_type = type(config)
+            resolved_actor = wrap_callable_as_stage(resolved_actor, schema_type)
+
         stage_info = StageInfo(
-            name=name, callable=stage_actor, config=config, min_replicas=min_replicas, max_replicas=max_replicas
+            name=name,
+            callable=resolved_actor,
+            config=config,
+            min_replicas=min_replicas,
+            max_replicas=max_replicas,
         )
-        self.topology.add_stage(stage_info)  # Delegate
+        self.topology.add_stage(stage_info)
 
         return self
 

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/primitives/ray_stat_collector.py

@@ -40,7 +40,7 @@ class RayStatsCollector:
             - `get_edge_queues() -> Dict[str, Tuple[Any, int]]`
             These methods should return snapshots suitable for iteration.
         interval : float, optional
-            The interval in seconds between stats collection attempts, by default 5.0.
+            The interval in seconds between stat collection attempts, by default 5.0.
         actor_timeout : float, optional
             Timeout in seconds for waiting for stats from a single actor, by default 5.0.
         queue_timeout : float, optional

nv_ingest-2025.7.7.dev20250707/nv_ingest/framework/orchestration/ray/stages/injectors/metadata_injector.py

@@ -0,0 +1,161 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from datetime import datetime
+import logging
+import pandas as pd
+from typing import Any
+from pydantic import BaseModel
+import ray
+
+from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
+from nv_ingest_api.internal.enums.common import (
+    DocumentTypeEnum,
+    ContentTypeEnum,
+    AccessLevelEnum,
+    TextTypeEnum,
+    LanguageEnum,
+)
+from nv_ingest_api.internal.primitives.tracing.tagging import traceable
+from nv_ingest_api.internal.schemas.meta.metadata_schema import ContentHierarchySchema
+from nv_ingest_api.util.converters.type_mappings import doc_type_to_content_type
+from nv_ingest_api.util.exception_handlers.decorators import (
+    nv_ingest_node_failure_try_except,
+)
+
+# logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class MetadataInjectionStage(RayActorStage):
+    """
+    A Ray actor stage that performs metadata injection on IngestControlMessages.
+
+    This stage iterates over the rows of the DataFrame payload, checks if metadata
+    injection is required, and if so, injects the appropriate metadata.
+    """
+
+    def __init__(self, config: BaseModel) -> None:
+        # Call the base initializer to set attributes like self._running.
+        super().__init__(config)
+        # Additional initialization can be added here if necessary.
+        logger.info("MetadataInjectionStage initialized with config: %s", config)
+
+    @traceable("metadata_injector")
+    @nv_ingest_node_failure_try_except(annotation_id="metadata_injector", raise_on_failure=False)
+    def on_data(self, message: Any) -> Any:
+        """
+        Process an incoming IngestControlMessage by injecting metadata into its DataFrame payload.
+
+        Parameters
+        ----------
+        message : IngestControlMessage
+            The incoming message containing the payload DataFrame.
+
+        Returns
+        -------
+        IngestControlMessage
+            The message with updated metadata if injection was required.
+        """
+        df = message.payload()
+        update_required = False
+        rows = []
+        logger.info("Starting metadata injection on DataFrame with %d rows", len(df))
+
+        for _, row in df.iterrows():
+            try:
+                # Convert document type to content type using enums.
+                content_type = doc_type_to_content_type(DocumentTypeEnum(row["document_type"]))
+                # Check if metadata is missing or doesn't contain 'content'
+                if (
+                    "metadata" not in row
+                    or not isinstance(row["metadata"], dict)
+                    or "content" not in row["metadata"].keys()
+                ):
+                    update_required = True
+
+                    # Initialize default structures based on MetaDataSchema
+                    default_source_metadata = {
+                        "source_id": row.get("source_id"),
+                        "source_name": row.get("source_name"),
+                        "source_type": row["document_type"],
+                        "source_location": "",
+                        "collection_id": "",
+                        "date_created": datetime.now().isoformat(),
+                        "last_modified": datetime.now().isoformat(),
+                        "summary": "",
+                        "partition_id": -1,
+                        "access_level": AccessLevelEnum.UNKNOWN.value,
+                    }
+
+                    default_content_metadata = {
+                        "type": content_type.name.lower(),
+                        "page_number": -1,
+                        "description": "",
+                        "hierarchy": ContentHierarchySchema().model_dump(),
+                        "subtype": "",
+                        "start_time": -1,
+                        "end_time": -1,
+                    }
+
+                    default_audio_metadata = None
+                    if content_type == ContentTypeEnum.AUDIO:
+                        default_audio_metadata = {
+                            "audio_type": row["document_type"],
+                            "audio_transcript": "",
+                        }
+
+                    default_image_metadata = None
+                    if content_type == ContentTypeEnum.IMAGE:
+                        default_image_metadata = {
+                            "image_type": row["document_type"],
+                            "structured_image_type": ContentTypeEnum.NONE.value,
+                            "caption": "",
+                            "text": "",
+                            "image_location": (0, 0, 0, 0),
+                            "image_location_max_dimensions": (0, 0),
+                            "uploaded_image_url": "",
+                            "width": 0,
+                            "height": 0,
+                        }
+
+                    default_text_metadata = None
+                    if content_type == ContentTypeEnum.TEXT:
+                        default_text_metadata = {
+                            "text_type": TextTypeEnum.DOCUMENT.value,
+                            "summary": "",
+                            "keywords": "",
+                            "language": LanguageEnum.UNKNOWN.value,
+                            "text_location": (0, 0, 0, 0),
+                            "text_location_max_dimensions": (0, 0, 0, 0),
+                        }
+
+                    row["metadata"] = {
+                        "content": row["content"],
+                        "content_metadata": default_content_metadata,
+                        "error_metadata": None,
+                        "audio_metadata": default_audio_metadata,
+                        "image_metadata": default_image_metadata,
+                        "source_metadata": default_source_metadata,
+                        "text_metadata": default_text_metadata,
+                    }
+                    logger.info(
+                        f"METADATA_INJECTOR_DEBUG: Rebuilt metadata for source_id='{row.get('source_id', 'N/A')}'. "
+                        f"Metadata keys: {list(row['metadata'].keys())}."
+                        f"'content' present: {'content' in row['metadata']}"
+                    )
+            except Exception as inner_e:
+                logger.exception("Failed to process row during metadata injection")
+                raise inner_e
+            rows.append(row)
+
+        if update_required:
+            docs = pd.DataFrame(rows)
+            message.payload(docs)
+            logger.info("Metadata injection updated payload with %d rows", len(docs))
+        else:
+            logger.info("No metadata update was necessary during metadata injection")
+
+        return message

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/stages/meta/ray_actor_source_stage_base.py

@@ -23,16 +23,13 @@ class RayActorSourceStage(RayActorStage, ABC):
         super().__init__(config, log_to_stdout=log_to_stdout)
         self.paused = False
 
+    def on_data(self, IngestControlMessage):
+        return NotImplemented("Source stages do not implement on_data().")
+
     @ray.method(num_returns=1)
     def set_input_queue(self, queue_handle: Any) -> bool:
         raise NotImplementedError("Source stages do not support an input queue.")
 
-    def get_input(self) -> Any:
-        """
-        Source stages must implement get_input() to fetch control messages from an external source.
-        """
-        pass
-
     @abstractmethod
     def _read_input(self) -> Any:
         """

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/stages/sources/message_broker_task_source.py

@@ -304,14 +304,6 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
 
         return control_message
 
-    def on_data(self, control_message: any) -> any:
-        """
-        Process the control message.
-        For this source stage, no additional processing is done, so simply return it.
-        """
-        self._logger.debug("on_data: Received control message for processing")
-        return control_message
-
     # In the processing loop, instead of checking a boolean, we wait on the event.
     def _processing_loop(self) -> None:
         """
@@ -336,7 +328,6 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
                 self._active_processing = True
 
                 self._logger.debug("Control message received; processing data")
-                updated_cm = self.on_data(control_message)
 
                 # Block until not paused using the pause event.
                 if self.output_queue is not None:
@@ -349,7 +340,7 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
 
                     while True:
                         try:
-                            self.output_queue.put(updated_cm)
+                            self.output_queue.put(control_message)
                             self.stats["successful_queue_writes"] += 1
                             break
                         except Exception:

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/stages/transforms/text_embed.py

@@ -32,7 +32,7 @@ class TextEmbeddingTransformStage(RayActorStage):
     """
 
     def __init__(self, config: TextEmbeddingSchema) -> None:
-        super().__init__(config)
+        super().__init__(config, log_to_stdout=False)
         try:
             self.validated_config = config
             logger.info("TextEmbeddingTransformStage configuration validated successfully.")

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/stages/transforms/text_splitter.py

@@ -1,11 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import logging
 from typing import Any
 import ray
 
-# Assume these imports come from your project:
 from nv_ingest.framework.orchestration.ray.stages.meta.ray_actor_stage_base import RayActorStage
 from nv_ingest.framework.util.flow_control import filter_by_task
-from nv_ingest_api.internal.primitives.ingest_control_message import remove_task_by_type
+from nv_ingest_api.internal.primitives.ingest_control_message import remove_task_by_type, IngestControlMessage
 from nv_ingest_api.internal.primitives.tracing.tagging import traceable
 from nv_ingest_api.internal.schemas.transform.transform_text_splitter_schema import TextSplitterSchema
 from nv_ingest_api.internal.transform.split_text import transform_text_split_and_tokenize_internal
@@ -72,3 +75,45 @@ class TextSplitterStage(RayActorStage):
         logger.info("TextSplitterStage.on_data: Finished processing, returning updated message.")
 
         return message
+
+
+def text_splitter_fn(control_message: IngestControlMessage, stage_config: TextSplitterSchema) -> IngestControlMessage:
+    """
+    Process an incoming IngestControlMessage by splitting and tokenizing its text.
+
+    Parameters
+    ----------
+    control_message : IngestControlMessage
+        The incoming message containing the payload DataFrame.
+
+    stage_config : BaseModel
+        The stage level configuration object
+
+    Returns
+    -------
+    IngestControlMessage
+        The updated message with its payload transformed.
+    """
+
+    # Extract the DataFrame payload.
+    df_payload = control_message.payload()
+    logger.debug("Extracted payload with %d rows.", len(df_payload))
+
+    # Remove the "split" task to obtain task-specific configuration.
+    task_config = remove_task_by_type(control_message, "split")
+    logger.debug("Extracted task config: %s", task_config)
+
+    # Transform the DataFrame (split text and tokenize).
+    df_updated = transform_text_split_and_tokenize_internal(
+        df_transform_ledger=df_payload,
+        task_config=task_config,
+        transform_config=stage_config,
+        execution_trace_log=None,
+    )
+    logger.info("TextSplitterStage.on_data: Transformation complete. Updated payload has %d rows.", len(df_updated))
+
+    # Update the message payload.
+    control_message.payload(df_updated)
+    logger.info("TextSplitterStage.on_data: Finished processing, returning updated message.")
+
+    return control_message

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/util/pipeline/pipeline_builders.py

@@ -174,9 +174,9 @@ def setup_ingestion_pipeline(pipeline: RayPipeline, ingest_config: Dict[str, Any
     pipeline.make_edge(image_dedup_stage_id, text_splitter_stage_id, queue_size=ingest_edge_buffer_size)
 
     ###### Primitive Transforms ########
-    pipeline.make_edge(text_splitter_stage_id, embed_extractions_stage_id, queue_size=ingest_edge_buffer_size)
-    pipeline.make_edge(embed_extractions_stage_id, image_caption_stage_id, queue_size=ingest_edge_buffer_size)
-    pipeline.make_edge(image_caption_stage_id, image_storage_stage_id, queue_size=ingest_edge_buffer_size)
+    pipeline.make_edge(text_splitter_stage_id, image_caption_stage_id, queue_size=ingest_edge_buffer_size)
+    pipeline.make_edge(image_caption_stage_id, embed_extractions_stage_id, queue_size=ingest_edge_buffer_size)
+    pipeline.make_edge(embed_extractions_stage_id, image_storage_stage_id, queue_size=ingest_edge_buffer_size)
 
     ###### Primitive Storage ########
     pipeline.make_edge(image_storage_stage_id, embedding_storage_stage_id, queue_size=ingest_edge_buffer_size)

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/util/pipeline/pipeline_runners.py

@@ -72,7 +72,7 @@ class PipelineCreationSchema(BaseModel):
 
     # API keys
     ngc_api_key: str = os.getenv("NGC_API_KEY", "")
-    nvidia_build_api_key: str = os.getenv("NVIDIA_BUILD_API_KEY", "")
+    nvidia_api_key: str = os.getenv("NVIDIA_API_KEY", "")
 
     # Observability settings
     otel_exporter_otlp_endpoint: str = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "localhost:4317")
@@ -86,9 +86,9 @@ class PipelineCreationSchema(BaseModel):
 
     # Vision language model settings
     vlm_caption_endpoint: str = os.getenv(
-        "VLM_CAPTION_ENDPOINT", "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-11b-vision-instruct/chat/completions"
+        "VLM_CAPTION_ENDPOINT", "https://ai.api.nvidia.com/v1/gr/nvidia/llama-3.1-nemotron-nano-vl-8b-v1/chat/completions"
     )
-    vlm_caption_model_name: str = os.getenv("VLM_CAPTION_MODEL_NAME", "meta/llama-3.2-11b-vision-instruct")
+    vlm_caption_model_name: str = os.getenv("VLM_CAPTION_MODEL_NAME", "nvidia/llama-3.1-nemotron-nano-vl-8b-v1")
 
     # YOLOX image processing settings
     yolox_graphic_elements_http_endpoint: str = os.getenv(
@@ -331,6 +331,10 @@ def run_pipeline(
     """
     if run_in_subprocess:
         logger.info("Launching pipeline in Python subprocess using multiprocessing.")
+        if (ingest_config.ngc_api_key is None or ingest_config.ngc_api_key == "") and (
+            ingest_config.nvidia_api_key is None or ingest_config.nvidia_api_key == ""
+        ):
+            logger.warning("NGC_API_KEY or NVIDIA_API_KEY are not set. NIM Related functions will not work.")
 
         ctx = multiprocessing.get_context("fork")
         process = ctx.Process(

{nv_ingest-2025.6.2.dev20250602 → nv_ingest-2025.7.7.dev20250707}/nv_ingest/framework/orchestration/ray/util/pipeline/stage_builders.py

@@ -2,8 +2,6 @@
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# TODO(Devin)
-# flake8: noqa
 import os
 
 import click
@@ -11,6 +9,7 @@ import logging
 
 from nv_ingest.framework.orchestration.ray.stages.sinks.default_drain import DefaultDrainSink
 from nv_ingest.framework.orchestration.ray.stages.telemetry.otel_tracer import OpenTelemetryTracerStage
+from nv_ingest.framework.orchestration.ray.stages.transforms.text_splitter import TextSplitterStage
 from nv_ingest.framework.schemas.framework_otel_tracer_schema import OpenTelemetryTracerSchema
 from nv_ingest_api.internal.schemas.extract.extract_infographic_schema import InfographicExtractorSchema
 
@@ -41,7 +40,6 @@ from nv_ingest.framework.orchestration.ray.stages.storage.image_storage import I
 from nv_ingest.framework.orchestration.ray.stages.storage.store_embeddings import EmbeddingStorageStage
 from nv_ingest.framework.orchestration.ray.stages.transforms.image_caption import ImageCaptionTransformStage
 from nv_ingest.framework.orchestration.ray.stages.transforms.text_embed import TextEmbeddingTransformStage
-from nv_ingest.framework.orchestration.ray.stages.transforms.text_splitter import TextSplitterStage
 from nv_ingest.framework.schemas.framework_metadata_injector_schema import MetadataInjectorSchema
 from nv_ingest_api.internal.schemas.extract.extract_audio_schema import AudioExtractorSchema
 from nv_ingest_api.internal.schemas.extract.extract_chart_schema import ChartExtractorSchema
@@ -107,7 +105,7 @@ def get_nim_service(env_var_prefix):
         "",
     )
     auth_token = os.environ.get(
-        "NVIDIA_BUILD_API_KEY",
+        "NVIDIA_API_KEY",
         "",
     ) or os.environ.get(
         "NGC_API_KEY",
@@ -137,7 +135,7 @@ def get_audio_retrieval_service(env_var_prefix):
         "",
     )
     auth_token = os.environ.get(
-        "NVIDIA_BUILD_API_KEY",
+        "NVIDIA_API_KEY",
         "",
     ) or os.environ.get(
         "NGC_API_KEY",
@@ -465,7 +463,7 @@ def add_text_splitter_stage(pipeline, default_cpu_count, stage_name="text_splitt
 
 def add_image_caption_stage(pipeline, default_cpu_count, stage_name="image_caption"):
     auth_token = os.environ.get(
-        "NVIDIA_BUILD_API_KEY",
+        "NVIDIA_API_KEY",
         "",
     ) or os.environ.get(
         "NGC_API_KEY",
@@ -473,13 +471,13 @@ def add_image_caption_stage(pipeline, default_cpu_count, stage_name="image_capti
     )
 
     endpoint_url = os.environ.get("VLM_CAPTION_ENDPOINT", "localhost:5000")
-    model_name = os.environ.get("VLM_CAPTION_MODEL_NAME", "meta/llama-3.2-11b-vision-instruct")
+    model_name = os.environ.get("VLM_CAPTION_MODEL_NAME", "nvidia/llama-3.1-nemotron-nano-vl-8b-v1")
 
     config = ImageCaptionExtractionSchema(
         **{
             "api_key": auth_token,
             "endpoint_url": endpoint_url,
-            "image_caption_model_name": model_name,
+            "model_name": model_name,
             "prompt": "Caption the content of this image:",
         }
     )
@@ -497,7 +495,7 @@ def add_image_caption_stage(pipeline, default_cpu_count, stage_name="image_capti
 
 def add_text_embedding_stage(pipeline, default_cpu_count, stage_name="text_embedding"):
     api_key = os.environ.get(
-        "NVIDIA_BUILD_API_KEY",
+        "NVIDIA_API_KEY",
         "",
     ) or os.environ.get(
         "NGC_API_KEY",