NeMo Retriever API Reference

This page is the public Python SDK contract for NeMo Retriever Library. The generated signatures document create_ingestor() and the concrete objects it returns. They also document Retriever query and answer helpers, generation operators, and parameter models.

Import the factory and GraphIngestionError from nemo_retriever. Import generation operators from nemo_retriever.operators.generation. Import LLM client, task, and result types from nemo_retriever.models.llm.

Error and failure contract

The Python API does not define a separate set of numeric NeMo Retriever extraction error codes. Depending on the run mode and failing stage, callers observe one or more of the following:

• Python configuration or dependency exceptions, such as ValueError, ImportError, or RuntimeError.
• GraphIngestionError for row-level failures from explicitly configured remote NIM stages in run_mode="inprocess" or "batch" when error_policy="raise" (the default).
• HTTP status codes or gRPC errors returned by a remote NIM or by the Retriever service. These are transport or upstream-service statuses, not NeMo Retriever-specific error codes.
• Per-document failures in ServiceIngestResult.failures when run_mode="service".

The generated API signatures and parameter models below are the API contract. Exception text and upstream response bodies can change between releases; do not parse them as stable codes. The stable text-generation codes documented in One-shot text generation apply to generation operator output columns, not to document extraction.

Configure at least one input source

Before you call .ingest(), .ingest_stream(), or .aingest_stream(), configure at least one input source by calling .files(), .texts(), or .buffers() with a nonempty value. Omitting input configuration or passing an empty collection raises ValueError before pipeline execution.

A configured source can legitimately produce blank text or an empty result. For example, OCR can find no text on an image-only page. This outcome does not raise the missing-input error.

A nonempty optional glob passed to .files() also counts as a configured source. If it matches no files, .ingest() can return an empty result, and the streaming methods can yield no results.

Select a supported extraction method

ExtractParams validates method when you construct the model. For PDF extraction, use pdfium, pdfium_hybrid, ocr, or nemotron_parse. The audio value remains available for the legacy params-driven audio path. For new audio pipelines, use GraphIngestor.extract_audio() instead.

Any other value raises a Pydantic ValidationError before pipeline setup. The error lists the supported values, so spelling and configuration errors do not silently select another extraction path.

For local nemotron_parse extraction in NeMo Retriever Library 26.08, omit nemotron_parse_model to use the default model, or set it to nvidia/NVIDIA-Nemotron-Parse-v1.2. Other local model values, including nvidia/NVIDIA-Nemotron-Parse-2.0, raise a Pydantic ValidationError before pipeline execution. To use a remote Nemotron Parse endpoint, configure nemotron_parse_invoke_url or invoke_url and select the model that matches the endpoint contract.

Choose raise or collect behavior

For graph run modes, error_policy="raise" raises GraphIngestionError when an explicitly configured remote NIM stage reports a row-level error. The exception retains the underlying records in exc.records. When available, its message identifies the stage, invoke URL, and HTTP status in a form similar to [stage=OCR NIM url=https://... http=503], followed by a troubleshooting hint.

Use error_policy="collect" when partial results are useful and your application inspects the error fields in every returned row. Alternatively, pass return_failures=True to .ingest() to receive a (result, failures) tuple. When no remote invoke URL is configured, return_failures=True scans all output columns for row-level error fields so local failures are still visible. In service mode, failures are also available from ServiceIngestResult.failures.

What the raise error policy covers

The strict policy applies only to stages where you explicitly configure a remote NIM invoke URL. It does not raise for local-only PDFium parsing, caption, audio or video, or ASR failures, even when those stages populate row-level error fields.

Configured invoke URL	DataFrame column scanned	Stage label in messages
page_elements_invoke_url	output_column (default page_elements_v3)	Page Elements NIM
ocr_invoke_url	ocr	OCR NIM
table_structure_invoke_url	table_structure_ocr_v1	Table Structure NIM
nemotron_parse_invoke_url or invoke_url	nemotron_parse_v1_2	Nemotron Parse NIM
embed_invoke_url or embedding_endpoint	output_column (default text_embeddings_1b_v2)	Embedding NIM

Caption and ASR use remote endpoints but are outside this raise path today. Remote caption failures can abort the whole ingest instead of returning a partial DataFrame. ASR failures can omit affected rows while logging a warning, which can look like an empty transcript unless you inspect logs.

Row-level error payloads

Most extraction stages write errors into the result row instead of raising immediately. The common nested shape is:

{
  "error": {
    "stage": "ocr_page_elements",
    "type": "HTTPError",
    "message": "HTTP 503 from https://example/v1/infer: ...",
    "traceback": "..."
  }
}

The stage string is a semi-stable operator identifier (for example remote_inference, nemotron_parse_pages, or split_pdf). It is not a product-wide error-code enum. HTTP status codes usually appear inside message text rather than as a separate status_code field; when a structured status is present, GraphIngestionError can include it in the rendered exception.

import os

from nemo_retriever import GraphIngestionError, create_ingestor
from nemo_retriever.common.params import ExtractParams

pipeline = (
    create_ingestor(run_mode="inprocess", error_policy="raise")
    .files(["document.pdf"])
    .extract(
        ExtractParams(
            method="ocr",
            ocr_invoke_url=os.environ["OCR_INVOKE_URL"],
        )
    )
)

try:
    result = pipeline.ingest()
except GraphIngestionError as exc:
    # Records can contain source paths, endpoint details, and upstream
    # response text. Extract only known-safe diagnostic fields before
    # logging or sending them to your support workflow.
    for record in exc.records:
        payload = record.get("error") if isinstance(record, dict) else record
        if isinstance(payload, dict):
            print(
                {
                    "column": record.get("column"),
                    "stage": payload.get("stage"),
                    "type": payload.get("type"),
                    "message": payload.get("message"),
                }
            )
        else:
            print(
                {
                    "column": record.get("column") if isinstance(record, dict) else None,
                    "message": str(payload),
                }
            )

For a support-oriented mapping of extraction paths, error signals, corrective actions, and escalation criteria, refer to Python API error triage.

Version-specific behavior

This reference describes the current NeMo Retriever Library. Older NV-Ingest releases, including 25.4.2, can use different exception text and result shapes and might not include enriched GraphIngestionError diagnostics. When troubleshooting an older deployment, use the package and container versions from that deployment and include them in the support case.

PDF pre-splitting for parallel ingest

Large PDFs are split into page batches before Ray processing so extraction can run in parallel. This happens on the default ingest path; you do not need extra configuration for typical workloads.

To tune splitter throughput from the CLI, use --pdf-split-batch-size (Ray actor batch size for the splitter stage). Refer to Local and batch ingest in the CLI reference.

Python client (pdf_split_config): Only ServiceIngestor.pdf_split_config() records page-chunking settings in the request pipeline spec for the remote gateway. Obtain that object with create_ingestor(run_mode="service"). Local graph ingest (run_mode="inprocess" or "batch") does not implement this method. PDFs are split automatically on the default graph ingest path without client-side configuration.

One-shot text generation

TextGenerationOperator is the reusable base for synchronous, one-request-per-row text generation. It is a provisional text-only API: it does not support tool calls, agent loops, streaming, multiple choices, or structured domain results.

Concrete operators construct an immutable TextGenerationTask and provide reconstructible constructor state. Runtime task and client objects must not be included in graph constructor arguments. A custom completion client must be safe for concurrent calls or report that it does not support concurrent calls so the operator serializes access.

Embedding and captioning remain separate operator families because they use modality grouping, native batching, and specialized CPU/GPU lifecycles.

Generic generation and summarization

Both operators consume a pandas DataFrame and add text, latency, model, and error columns without changing the input rows:

import pandas as pd

from nemo_retriever.common.params import TextGenerationParams
from nemo_retriever.operators.generation import GenericGenerationOperator, SummarizationOperator

summary_params = TextGenerationParams.from_kwargs(
    model="openai/gpt-4o-mini",
    api_key="os.environ/OPENAI_API_KEY",
    temperature=0.0,
    max_tokens=512,
)
summaries = SummarizationOperator(summary_params).run(
    pd.DataFrame({"text": ["A long document to summarize."]})
)

prompt_params = TextGenerationParams.from_kwargs(
    model="openai/gpt-4o-mini",
    api_key="os.environ/OPENAI_API_KEY",
    prompt="Write a {tone} title for: {text}",
)
titles = GenericGenerationOperator(
    prompt_params,
    input_columns={"tone": "style", "text": "document"},
    output_column="title",
).run(pd.DataFrame({"style": ["concise"], "document": ["Quarterly results"]}))

SummarizationOperator defaults to text, summary, summary_latency_s, summary_model, and summary_error. GenericGenerationOperator maps each named prompt placeholder to a physical DataFrame column and derives the metadata column names from output_column. Prompt contracts are validated when the operator is constructed, before any provider request runs. SummarizeTask inherits from TextGenerationTask and supplies the built-in summarization prompt unless you override prompt on TextGenerationParams.

TextGenerationParams configuration

Construct TextGenerationParams with TextGenerationParams.from_kwargs(...). The following fields are supported.

Field	Purpose
model	Required provider model identifier.
api_base	Optional OpenAI-compatible API base URL.
api_key	Optional credential or os.environ/<NAME> reference. Literal keys are not written to persisted graph JSON.
temperature	Optional sampling override. Valid values are 0.0 through 2.0.
top_p	Optional sampling override. Valid values are 0.0 through 1.0.
max_tokens	Optional positive token-limit override. Omit the field to inherit the task default.
extra_params	Optional provider-specific request keys. Dedicated fields such as model, messages, sampling, and credentials must not be duplicated here.
num_retries	Transport retry count. The default is 3. The value must be 0 or greater.
timeout	Transport timeout in seconds. The default is 120.0. The value must be greater than 0.
prompt	Optional user prompt or prompt template. GenericGenerationOperator requires this field.
system_prompt	Optional system prompt.
rag_system_prompt	Optional retrieval-augmented generation (RAG) system prompt.
rag_system_prompt_prefix	Optional prefix applied to the RAG system prompt.
reasoning_enabled	Optional reasoning toggle. When unset, transport reasoning defaults to true.
max_workers	Concurrent row workers. The default is 8. The value must be 1 or greater.

Sampling fields that you omit inherit the task defaults. temperature, top_p, and max_tokens are applied only when you pass them explicitly.

To define another one-request/one-text-result task, subclass TextGenerationTask, declare required_inputs, and implement build_request(). Then construct it from a TextGenerationOperator subclass with explicit logical-input-to-DataFrame-column mappings. This abstraction is intentionally text-only; use a separate operator family for embeddings, captioning, tools, streaming, or structured domain results.

Generation failures are collected per row using stable error codes: empty_input, request_error, transport_error, unsupported_response, parse_error, empty_output, and the RAG-specific thinking_truncated. Raw provider exceptions and credentials are not written to DataFrame outputs.

Generated signatures for TextGenerationOperator, GenericGenerationOperator, and SummarizationOperator appear in Generation operators. Generated signatures for TextGenerationTask, LiteLLMClient, LLMClient, AnswerJudge, AnswerResult, and related types appear in LLM clients, tasks, and results. Retriever.answer() requires an LLMClient and returns AnswerResult.

Persisted graphs are trusted configuration

Graph loading imports operator classes and invokes their constructors. Load graph JSON only from trusted sources; do not expose graph payloads, callable references, or class names as model- or user-controlled agent tools.

Version 2 graph files preserve shared-node DAG identity and reject cycles. Constructor state must consist of supported JSON-native values, typed Pydantic models, paths, sets and tuples, or importable type/callable references. Runtime data such as DataFrames and opaque client objects is not persistable.

API keys are never written into graph JSON. Use an explicit environment reference in persisted configuration:

QAGenerationOperator(
    model="openai/gpt-4o-mini",
    api_key="os.environ/OPENAI_API_KEY",
)

For LiteLLM nvidia_nim/... models, including the default LiteLLMClient and LLMJudge models, use os.environ/NVIDIA_API_KEY. Refer to Authentication and API keys.

Serializing a graph containing a literal API key fails with a contextual error instead of guessing which provider credential should be used on a worker.

Generated Python API

The signatures below are the supported public ingest, retrieve, generation, and parameter surfaces. Private helpers, module loggers, and duplicate re-exports are omitted.

Use the following public import paths:

• Import create_ingestor and GraphIngestionError from nemo_retriever.
• Import GraphIngestor from nemo_retriever.ingestor.graph_ingestor.
• Import ServiceIngestor from nemo_retriever.service.service_ingestor.
• Import generation operators from nemo_retriever.operators.generation.
• Import LLM client, task, and result types from nemo_retriever.models.llm.
• Import parameter models from nemo_retriever.common.params.
• Import RetrieveVdbOperator from nemo_retriever.operators.vdb.
• Import NemotronRerankActor from nemo_retriever.operators.rerank.

Public ingestion factory

create_ingestor() returns a concrete ingestion client from run_mode. The supported values are inprocess, batch, and service.

run_mode	Runtime type	Execution
inprocess	GraphIngestor	Local in-process graph. This is the default.
batch	GraphIngestor	Ray Data graph.
service	ServiceIngestor	Remote Retriever service.

The function is annotated as returning the shared Ingestor interface. At runtime it returns GraphIngestor or ServiceIngestor. Use the generated class entries below for run-mode-specific methods.

Factory keyword arguments merge into IngestorCreateParams. Common fields include documents, base_url, api_key, error_policy, ray_address, and max_concurrency. Refer to IngestorCreateParams in Parameter models.

from nemo_retriever import GraphIngestionError, create_ingestor

graph = create_ingestor(run_mode="inprocess")
batch = create_ingestor(run_mode="batch")
service = create_ingestor(run_mode="service", base_url="http://localhost:7670")

nemo_retriever.ingestor.core.create_ingestor(*, run_mode='inprocess', params=None, **kwargs)

Source code in nemo_retriever/ingestor/core.py

def create_ingestor(
    *,
    run_mode: IngestorRunMode = "inprocess",
    params: IngestorCreateParams | None = None,
    **kwargs: Any,
) -> "Ingestor":
    """
    Graph-only ingestion factory.
    """
    merged = _merge_params(params, kwargs)
    if isinstance(merged, IngestorCreateParams):
        parsed = merged
    else:
        parsed = IngestorCreateParams(**merged)

    if run_mode == "service":
        from nemo_retriever.service.service_ingestor import ServiceIngestor

        service_kwargs: dict[str, Any] = {
            "base_url": parsed.base_url,
            "documents": parsed.documents,
            "api_token": parsed.api_key,
        }
        if parsed.max_concurrency is not None:
            service_kwargs["max_concurrency"] = parsed.max_concurrency
        return ServiceIngestor(**service_kwargs)

    if run_mode not in {"batch", "inprocess"}:
        raise ValueError(f"create_ingestor supports run modes 'inprocess', 'batch', and 'service'; got {run_mode!r}.")

    from nemo_retriever.ingestor.graph_ingestor import GraphIngestor

    return GraphIngestor(
        run_mode=run_mode,
        documents=parsed.documents,
        ray_address=parsed.ray_address,
        ray_log_to_driver=parsed.ray_log_to_driver,
        debug=parsed.debug,
        allow_no_gpu=parsed.allow_no_gpu,
        node_overrides=parsed.node_overrides,
        error_policy=parsed.error_policy,
    )

Graph ingest

create_ingestor(run_mode="inprocess") and create_ingestor(run_mode="batch") return GraphIngestor. Import the class from nemo_retriever.ingestor.graph_ingestor when you need the type explicitly. Import GraphIngestionError from nemo_retriever.

GraphIngestor methods include extract_html(), extract_audio(), extract_video(), get_error_rows(), and get_dataset(). get_error_rows() filters rows that contain stage error payloads from a pandas DataFrame or Ray Dataset. If you omit dataset, it uses the dataset retained from the last ingest() call. get_dataset() returns that retained dataset.

nemo_retriever.ingestor.graph_ingestor.GraphIngestor

Bases: ingestor

Parameters:

Name	Type	Description	Default
run_mode	str	"inprocess" (single-process pandas, default) or "batch" (Ray Data).	'inprocess'
ray_address	Optional[str]	Ray cluster address. None starts a local cluster.	None
batch_size	int	Default map_batches batch size for RayDataExecutor.	1
num_cpus	float	Default CPU resources per operator node (batch mode).	1
num_gpus	float	Default GPU resources per operator node (batch mode).	0
node_overrides	Optional[Dict[str, Dict[str, Any]]]	Per-node resource/batching overrides forwarded to :class:~nemo_retriever.graph.RayDataExecutor. Keys are node names (e.g. "OCRActor"); values are dicts accepted by RayDataExecutor.__init__ (num_gpus, batch_size, etc.).	None
show_progress	bool	Show a tqdm progress bar when running in inprocess mode.	True
error_policy	str	"raise" raises when explicitly configured remote NIM stages report row-level errors. "collect" returns partial results with the stage error payloads preserved.	'raise'

Source code in nemo_retriever/ingestor/graph_ingestor.py

class GraphIngestor(ingestor):
    """Ingestor that constructs and executes operator graphs directly.

    The fluent builder methods record pipeline stages. When :meth:`ingest` is
    called it builds a :class:`~nemo_retriever.graph.Graph` and feeds it to
    the appropriate executor.

    Parameters
    ----------
    run_mode
        ``"inprocess"`` (single-process pandas, default) or ``"batch"`` (Ray
        Data).
    ray_address
        Ray cluster address. ``None`` starts a local cluster.
    batch_size
        Default ``map_batches`` batch size for ``RayDataExecutor``.
    num_cpus
        Default CPU resources per operator node (batch mode).
    num_gpus
        Default GPU resources per operator node (batch mode).
    node_overrides
        Per-node resource/batching overrides forwarded to
        :class:`~nemo_retriever.graph.RayDataExecutor`.  Keys are node names
        (e.g. ``"OCRActor"``); values are dicts accepted by
        ``RayDataExecutor.__init__`` (``num_gpus``, ``batch_size``, etc.).
    show_progress
        Show a tqdm progress bar when running in inprocess mode.
    error_policy
        ``"raise"`` raises when explicitly configured remote NIM stages report
        row-level errors. ``"collect"`` returns partial results with the stage
        error payloads preserved.
    """

    RUN_MODE = "graph"

    def __init__(
        self,
        *,
        run_mode: str = "inprocess",
        documents: Optional[List[str]] = None,
        ray_address: Optional[str] = None,
        ray_log_to_driver: bool = True,
        debug: bool = False,
        allow_no_gpu: bool = False,
        batch_size: int = 1,
        num_cpus: float = 1,
        num_gpus: float = 0,
        node_overrides: Optional[Dict[str, Dict[str, Any]]] = None,
        show_progress: bool = True,
        error_policy: str = "raise",
    ) -> None:
        super().__init__(documents=documents)
        if run_mode not in {"batch", "inprocess"}:
            raise ValueError(f"run_mode must be 'batch' or 'inprocess', got {run_mode!r}")
        if error_policy not in {"raise", "collect"}:
            raise ValueError(f"error_policy must be 'raise' or 'collect', got {error_policy!r}")
        self._run_mode = run_mode
        self._ray_address = ray_address
        self._ray_log_to_driver = ray_log_to_driver
        self._debug = debug
        self._allow_no_gpu = allow_no_gpu
        self._batch_size = batch_size
        self._num_cpus = num_cpus
        self._num_gpus = num_gpus
        self._node_overrides: Dict[str, Dict[str, Any]] = node_overrides or {}
        self._show_progress = show_progress
        self._error_policy = error_policy
        self._rd_dataset: Any = None
        self._buffers: list[tuple[str, BytesIO]] = []
        self._inline_texts: list[str] | None = None

        # Pipeline configuration accumulated by fluent methods
        self._extraction_mode: str | None = None
        self._extract_params: Any = None
        self._text_params: Any = None
        self._html_params: Any = None
        self._audio_chunk_params: Any = None
        self._asr_params: Any = None
        self._video_frame_params: Any = None
        self._video_text_dedup_params: Any = None
        self._av_fuse_params: Any = None
        self._embed_params: Any = None
        self._split_config: dict[str, Any] = dict.fromkeys(SPLIT_CONFIG_VALID_KEYS, None)
        self._caption_params: Any = None
        self._dedup_params: Any = None
        self._store_params: Any = None
        self._vdb_upload_params: Any = None
        self._webhook_params: Any = None
        # Ordered list of stage names; "extract" is tracked but excluded from
        # the post-extraction stage_order passed to graph builders.
        self._stage_order: List[str] = []

    # ------------------------------------------------------------------
    # Input configuration
    # ------------------------------------------------------------------

    def files(self, documents: Union[str, List[str]]) -> "GraphIngestor":
        """Set the input file paths or glob patterns."""
        self._documents = [documents] if isinstance(documents, str) else list(documents)
        return self

    def texts(self, texts: Union[str, Sequence[str]]) -> Self:
        """Set raw inline text documents as the graph input.

        Each string is one logical source document. It receives a deterministic
        ``inline://`` identifier and flows through the normal text splitter,
        embedding, and sink stages without being written to a temporary file.
        Inline text may be combined with file or buffer inputs; the manifest
        planner routes each source through its matching extraction branch.
        """
        self._inline_texts = normalize_inline_texts(texts)
        return self

    def buffers(
        self,
        buffers: Union[Tuple[str, BytesIO], List[Tuple[str, BytesIO]]],
    ) -> "GraphIngestor":
        """Set in-memory buffers for processing.

        Each buffer is a ``(name, BytesIO)`` pair where *name* carries the
        original filename (including extension) so downstream operators can
        detect file type.  Accepts a single tuple or a list of tuples.

        Only supported for ``run_mode='inprocess'``.
        """
        if isinstance(buffers, tuple) and len(buffers) == 2 and isinstance(buffers[0], str):
            self._buffers = [buffers]
        else:
            self._buffers = list(buffers)
        return self

    # ------------------------------------------------------------------
    # Extraction stage (sets extraction_mode and primary params)
    # ------------------------------------------------------------------

    def extract(
        self,
        params: Optional[ExtractParams] = None,
        *,
        split_config: dict[str, Any] | None = None,
        extraction_mode: str | None = None,
        text_params: Optional[TextChunkParams] = None,
        html_params: Optional[HtmlChunkParams] = None,
        audio_chunk_params: Optional[AudioChunkParams] = None,
        asr_params: Optional[ASRParams] = None,
        video_frame_params: Optional[VideoFrameParams] = None,
        video_text_dedup_params: Optional[VideoFrameTextDedupParams] = None,
        av_fuse_params: Optional[AudioVisualFuseParams] = None,
        **kwargs: Any,
    ) -> "GraphIngestor":
        """Configure extraction.

        By default, the effective extraction mode is inferred from the input
        file extensions immediately before graph construction. Pass
        ``extraction_mode='pdf'`` to force the dedicated PDF/document graph, or
        ``extraction_mode='auto'`` to dispatch a mixed folder through
        :class:`MultiTypeExtractOperator`.
        Chunking is opt-in: pass ``split_config={"<key>": {...}}`` to enable
        post-extract token chunking for that source type.

        Unknown ``**kwargs`` raise :class:`TypeError`. Only fields declared
        on :class:`ExtractParams` are accepted as extra kwargs; ASR / audio
        configuration belongs on :class:`ASRParams` (pass ``asr_params=``
        or use :meth:`extract_audio`).
        """
        unknown = set(kwargs) - set(ExtractParams.model_fields)
        if unknown:
            raise TypeError(
                f"extract() got unexpected keyword argument(s) {sorted(unknown)!r}. "
                f"Allowed extra kwargs must be fields of ExtractParams. "
                f"For ASR / audio configuration, pass asr_params=ASRParams(...) "
                f"or use .extract_audio(asr_params=ASRParams(...)) "
                f"(see docs/extraction/audio-video.md)."
            )
        self._extraction_mode = extraction_mode
        self._extract_params = _resolve_api_key(_coerce(params, kwargs, default_factory=ExtractParams))
        if text_params is not None:
            self._text_params = text_params
        if html_params is not None:
            self._html_params = html_params
        if audio_chunk_params is not None:
            self._audio_chunk_params = audio_chunk_params
        if asr_params is not None:
            self._asr_params = asr_params
        if video_frame_params is not None:
            self._video_frame_params = video_frame_params
        if video_text_dedup_params is not None:
            self._video_text_dedup_params = video_text_dedup_params
        if av_fuse_params is not None:
            self._av_fuse_params = av_fuse_params
        self._apply_split_config(split_config)
        self._record_stage("extract")
        return self

    def extract_image_files(
        self,
        params: Optional[ExtractParams] = None,
        *,
        split_config: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> "GraphIngestor":
        """Configure image extraction (extraction_mode='image')."""
        self._extraction_mode = "image"
        self._extract_params = _resolve_api_key(_coerce(params, kwargs, default_factory=ExtractParams))
        self._apply_split_config(split_config)
        self._record_stage("extract")
        return self

    def extract_html(self, params: Optional[HtmlChunkParams] = None, **kwargs: Any) -> "GraphIngestor":
        """Configure HTML extraction (extraction_mode='html')."""
        self._extraction_mode = "html"
        self._html_params = _coerce(params, kwargs, default_factory=HtmlChunkParams)
        self._record_stage("extract")
        return self

    def extract_audio(
        self,
        params: Optional[AudioChunkParams] = None,
        *,
        asr_params: Optional[ASRParams] = None,
        split_config: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> "GraphIngestor":
        """Configure audio extraction (extraction_mode='audio')."""
        self._extraction_mode = "audio"
        self._audio_chunk_params = _coerce(params, kwargs, default_factory=AudioChunkParams)
        self._asr_params = asr_params or ASRParams()
        self._apply_split_config(split_config)
        self._record_stage("extract")
        return self

    def extract_video(
        self,
        params: Optional[AudioChunkParams] = None,
        *,
        asr_params: Optional[ASRParams] = None,
        video_frame_params: Optional[VideoFrameParams] = None,
        video_text_dedup_params: Optional[VideoFrameTextDedupParams] = None,
        av_fuse_params: Optional[AudioVisualFuseParams] = None,
        extract_params: Optional[ExtractParams] = None,
        split_config: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> "GraphIngestor":
        """Configure video extraction.

        Sets ``extraction_mode='auto'`` so :class:`MultiTypeExtractOperator`
        dispatches by file extension; ``.mp4``/``.mov``/``.mkv``
        files are routed to a combined audio-from-video ASR + frame OCR +
        scene fusion pipeline.

        Frame OCR config (``ocr_invoke_url``, ``ocr_api_key``,
        ``inference_batch_size``, ``ocr_request_timeout_s``) is read from
        :class:`ExtractParams` — the same object the PDF/image pipelines
        use — so the user only configures OCR once.

        The ``split_config`` keyword honors the ``"video"`` key (chunking the
        fused audio+visual transcript). The ``"audio"`` key is ignored on the
        video pipeline — for audio-only chunking, use :meth:`extract_audio`
        directly with that file.
        """
        self._extraction_mode = "auto"
        self._audio_chunk_params = _coerce(params, kwargs, default_factory=AudioChunkParams)
        self._asr_params = asr_params or ASRParams()
        self._video_frame_params = video_frame_params or VideoFrameParams()
        self._video_text_dedup_params = video_text_dedup_params or VideoFrameTextDedupParams()
        self._av_fuse_params = av_fuse_params or AudioVisualFuseParams()
        if extract_params is not None:
            self._extract_params = _resolve_api_key(extract_params)
        elif self._extract_params is None:
            self._extract_params = ExtractParams()
        self._apply_split_config(split_config)
        self._record_stage("extract")
        return self

    # ------------------------------------------------------------------
    # Post-extraction transform stages
    # ------------------------------------------------------------------

    def dedup(self, params: Optional[DedupParams] = None, **kwargs: Any) -> "GraphIngestor":
        """Record a dedup stage."""
        self._dedup_params = _coerce(params, kwargs, default_factory=DedupParams)
        self._record_stage("dedup")
        return self

    def caption(self, params: Optional[CaptionParams] = None, **kwargs: Any) -> "GraphIngestor":
        """Record a caption stage."""
        self._caption_params = _resolve_api_key(_coerce(params, kwargs, default_factory=CaptionParams))
        self._record_stage("caption")
        return self

    def store(self, params: Optional[StoreParams] = None, **kwargs: Any) -> "GraphIngestor":
        """Record a store stage for persisting extracted image assets to storage."""
        self._store_params = _coerce(params, kwargs, default_factory=StoreParams)
        self._record_stage("store")
        return self

    def embed(self, params: Optional[EmbedParams] = None, **kwargs: Any) -> "GraphIngestor":
        """Record an embedding stage."""
        self._embed_params = _resolve_api_key(_coerce(params, kwargs, default_factory=EmbedParams))
        self._record_stage("embed")
        return self

    def vdb_upload(self, params: Optional[VdbUploadParams] = None, **kwargs: Any) -> "GraphIngestor":
        """Record a vector DB upload **sink** (in-graph after embed/store, before webhook).

        Does not call :meth:`_record_stage`: ``stage_order`` only lists
        ``dedup`` / ``caption`` / ``store`` / ``embed`` for reordering; VDB is
        always appended from ``_vdb_upload_params`` in
        :func:`~nemo_retriever.graph.ingestor_runtime._append_ordered_transform_stages`.
        Plan builders that round-trip sinks use :meth:`~nemo_retriever.ingest_plans.BaseIngestPlan.record_sink`.
        """
        self._vdb_upload_params = _coerce(params, kwargs, default_factory=VdbUploadParams)
        return self

    def webhook(self, params: Optional[WebhookParams] = None, **kwargs: Any) -> "GraphIngestor":
        """Record a webhook notification stage (always runs last).

        When ``endpoint_url`` is set, processed results are HTTP-POSTed to
        that URL.  If ``endpoint_url`` is ``None`` the stage is a no-op.
        """
        self._webhook_params = _coerce(params, kwargs, default_factory=WebhookParams)
        self._record_stage("webhook")
        return self

    # ------------------------------------------------------------------
    # Execution
    # ------------------------------------------------------------------

    def ingest(self, params: Any = None, **kwargs: Any) -> Any:
        """Build the operator graph and run it through the configured executor.

        Parameters
        ----------
        params
            Optional :class:`IngestExecuteParams` (or plain ``dict``) carrying
            execute-time flags. Graph run modes honor ``return_failures``.
        **kwargs
            Execute-time flags passed directly. ``return_failures`` may be
            passed here and takes precedence over the value in ``params``.
        return_failures
            When ``True`` (default ``False``), return ``(result, failures)``
            instead of raising collected row-level stage errors. If no explicit
            remote-stage diagnostics are configured, all output columns are
            scanned for populated error fields so local collected failures can
            still be returned; the default raise path remains scoped to
            explicitly configured remote stages.

        Returns
        -------
        ``run_mode='batch'`` or ``run_mode='inprocess'``
            A ``pandas.DataFrame``.
        ``return_failures=True``
            ``(result, failures)`` where ``failures`` is a list of
            service-style ``(source, error)`` tuples.
        """
        return_failures = self._resolve_return_failures(params, kwargs)
        self._validate_input_sources(self._inline_texts)
        if not self._documents and not self._buffers and is_blank_inline_corpus(self._inline_texts):
            result = empty_text_chunks_df()
            if self._run_mode == "batch":
                self._rd_dataset = result
            else:
                self._rd_dataset = None
            return self._finalize_ingest_result(result, return_failures=return_failures)

        default_branches = self._plan_default_extraction_branches()
        execute_branches = default_branches is not None and (
            len(default_branches) > 1 or self._has_mixed_inline_sources()
        )
        if default_branches is None:
            single_effective = self._resolve_effective_extraction_inputs()
        elif not execute_branches:
            single_effective = self._resolve_branch_extraction_inputs(default_branches[0])
        else:
            single_effective = None

        # Auto-enable dedup before captioning so that images overlapping
        # with table/chart/infographic detections are removed first.
        # Skip for image-only extraction — the image IS the content.
        image_only = single_effective is not None and single_effective.extraction_mode == "image"
        if self._caption_params is not None and self._dedup_params is None and not image_only:
            self._dedup_params = DedupParams()
            if "dedup" not in self._stage_order:
                try:
                    idx = self._stage_order.index("caption")
                except ValueError:
                    idx = len(self._stage_order)
                self._stage_order.insert(idx, "dedup")

        post_extract_order = tuple(s for s in self._stage_order if s != "extract")

        if execute_branches:
            result = self._execute_extraction_branches(default_branches, post_extract_order=post_extract_order)
        else:
            if single_effective is None:
                raise RuntimeError("Internal error: extraction inputs were not resolved.")
            result = self._execute_single_graph(single_effective, post_extract_order=post_extract_order)

        return self._finalize_ingest_result(result, return_failures=return_failures)

    def _execute_single_graph(
        self,
        effective_extraction: ResolvedExtractionInputs,
        *,
        post_extract_order: tuple[str, ...],
    ) -> Any:
        if self._run_mode == "batch":
            return self._execute_single_graph_batch(effective_extraction, post_extract_order=post_extract_order)
        return self._execute_single_graph_inprocess(effective_extraction, post_extract_order=post_extract_order)

    def _execute_single_graph_batch(
        self,
        effective_extraction: ResolvedExtractionInputs,
        *,
        post_extract_order: tuple[str, ...],
    ) -> Any:
        ray, cluster_resources = self._ensure_batch_runtime()
        graph = build_graph(
            extraction_mode=effective_extraction.extraction_mode,
            extract_params=effective_extraction.extract_params,
            text_params=effective_extraction.text_params,
            html_params=effective_extraction.html_params,
            audio_chunk_params=effective_extraction.audio_chunk_params,
            asr_params=effective_extraction.asr_params,
            video_frame_params=effective_extraction.video_frame_params,
            video_text_dedup_params=effective_extraction.video_text_dedup_params,
            av_fuse_params=effective_extraction.av_fuse_params,
            embed_params=self._embed_params,
            split_config=self._split_config,
            caption_params=self._caption_params,
            dedup_params=self._dedup_params,
            store_params=self._store_params,
            vdb_upload_params=self._vdb_upload_params,
            webhook_params=self._webhook_params,
            stage_order=post_extract_order,
        )
        effective_allow_no_gpu = self._allow_no_gpu or cluster_resources.total_gpu_count() == 0
        derived_overrides = batch_tuning_to_node_overrides(
            effective_extraction.extract_params,
            self._embed_params,
            store_params=self._store_params,
            cluster_resources=cluster_resources,
            allow_no_gpu=effective_allow_no_gpu,
            caption_params=self._caption_params,
            video_frame_params=effective_extraction.video_frame_params,
        )
        executor = RayDataExecutor(
            graph,
            ray_address=self._ray_address,
            batch_size=self._batch_size,
            num_cpus=self._num_cpus,
            num_gpus=self._num_gpus,
            node_overrides=merge_node_overrides(derived_overrides, self._node_overrides),
            auto_concurrency_nodes=(
                default_concurrency_node_names(
                    effective_extraction.extract_params,
                    self._embed_params,
                    self._store_params,
                    self._caption_params,
                )
                - set(self._node_overrides)
            ),
        )
        executor_input = self._inline_text_dataset(ray.data) if self._inline_texts else self._documents
        result = executor.ingest(executor_input)
        self._rd_dataset = result
        return result

    def _execute_single_graph_inprocess(
        self,
        effective_extraction: ResolvedExtractionInputs,
        *,
        post_extract_order: tuple[str, ...],
    ) -> Any:
        graph = build_graph(
            extraction_mode=effective_extraction.extraction_mode,
            extract_params=effective_extraction.extract_params,
            text_params=effective_extraction.text_params,
            html_params=effective_extraction.html_params,
            audio_chunk_params=effective_extraction.audio_chunk_params,
            asr_params=effective_extraction.asr_params,
            video_frame_params=effective_extraction.video_frame_params,
            video_text_dedup_params=effective_extraction.video_text_dedup_params,
            av_fuse_params=effective_extraction.av_fuse_params,
            embed_params=self._embed_params,
            split_config=self._split_config,
            caption_params=self._caption_params,
            dedup_params=self._dedup_params,
            store_params=self._store_params,
            vdb_upload_params=self._vdb_upload_params,
            webhook_params=self._webhook_params,
            stage_order=post_extract_order,
        )
        executor = InprocessExecutor(graph, show_progress=self._show_progress)
        self._rd_dataset = None
        if self._inline_texts:
            return executor.ingest(self._inline_text_dataframe())
        if self._buffers:
            import pandas as pd

            df = pd.DataFrame([{"bytes": buf.getvalue(), "path": name} for name, buf in self._buffers])
            return executor.ingest(df)
        return executor.ingest(self._documents)

    def _execute_extraction_branches(
        self,
        branches: tuple[ExtractionBranchPlan, ...],
        *,
        post_extract_order: tuple[str, ...],
    ) -> Any:
        result = ExtractionBranchExecutor(
            run_mode=self._run_mode,
            branches=branches,
            documents=self._documents,
            buffers=self._buffers,
            inline_rows=self._inline_text_rows(),
            split_config=self._split_config,
            extract_params=self._extract_params,
            text_params=self._text_params,
            html_params=self._html_params,
            audio_chunk_params=self._audio_chunk_params,
            asr_params=self._asr_params,
            video_frame_params=self._video_frame_params,
            video_text_dedup_params=self._video_text_dedup_params,
            av_fuse_params=self._av_fuse_params,
            embed_params=self._embed_params,
            caption_params=self._caption_params,
            dedup_params=self._dedup_params,
            store_params=self._store_params,
            vdb_upload_params=self._vdb_upload_params,
            webhook_params=self._webhook_params,
            post_extract_order=post_extract_order,
            ray_address=self._ray_address,
            batch_size=self._batch_size,
            num_cpus=self._num_cpus,
            num_gpus=self._num_gpus,
            node_overrides=self._node_overrides,
            show_progress=self._show_progress,
            allow_no_gpu=self._allow_no_gpu,
            ensure_batch_runtime=self._ensure_batch_runtime,
        ).execute()
        self._rd_dataset = result if self._run_mode == "batch" else None
        return result

    def _ensure_batch_runtime(self) -> tuple[Any, Any]:
        ray = ensure_local_ray_runtime(self._ray_address, log_to_driver=self._ray_log_to_driver)
        return ray, gather_cluster_resources(ray)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _has_mixed_inline_sources(self) -> bool:
        return bool(self._inline_texts) and bool(self._documents or self._buffers)

    def _inline_text_rows(self) -> list[dict[str, str]]:
        return [
            {"text": text, "path": inline_text_source_id(index)} for index, text in enumerate(self._inline_texts or [])
        ]

    def _inline_text_dataframe(self) -> Any:
        import pandas as pd

        return pd.DataFrame(self._inline_text_rows(), columns=["text", "path"])

    def _inline_text_dataset(self, ray_data: Any) -> Any:
        return ray_data.from_items(self._inline_text_rows())

    def _configured_input_paths(self) -> list[str]:
        paths: list[str] = []
        for document in self._documents:
            try:
                paths.extend(expand_input_file_patterns([document]))
            except FileNotFoundError:
                paths.append(os.fspath(document))
        paths.extend(name for name, _ in self._buffers)
        paths.extend(inline_text_source_id(index) for index, _ in enumerate(self._inline_texts or []))
        return paths

    def _classified_input_paths(self) -> list[tuple[str, str | None]]:
        # Service workers receive inline text as a named byte buffer. Keep the
        # logical URI as its source path while classifying it as decoded text.
        return [
            (path, "txt" if is_inline_text_source(path) else input_type_for_path(path))
            for path in self._configured_input_paths()
        ]

    @staticmethod
    def _input_type_examples(paths: Iterable[str], *, limit: int = 3) -> str:
        examples = list(paths)[:limit]
        return ", ".join(examples)

    def _validate_explicit_extraction_mode_inputs(
        self,
        extraction_mode: str,
        classified: list[tuple[str, str | None]],
    ) -> None:
        allowed_types = _EXPLICIT_MODE_INPUT_TYPES.get(extraction_mode)
        if allowed_types is None:
            return

        mismatched = [
            path
            for path, input_type in classified
            if not _is_explicit_glob_path(path) and (input_type is None or input_type not in allowed_types)
        ]
        if mismatched:
            examples = self._input_type_examples(mismatched)
            raise ValueError(f"Input file type(s) do not match extraction_mode={extraction_mode!r}: {examples}")

    def _plan_default_extraction_branches(self) -> tuple[ExtractionBranchPlan, ...] | None:
        if self._extraction_mode is not None and not self._has_mixed_inline_sources():
            return None
        manifest = build_input_manifest(self._configured_input_paths())
        branches = plan_extraction_branches(manifest)
        if self._debug:
            logger.info(
                "Retriever ingest manifest planned %d extraction branches: %s",
                len(branches),
                format_branch_summary(branches),
            )
        return branches

    def _resolve_branch_extraction_inputs(self, branch: ExtractionBranchPlan) -> ResolvedExtractionInputs:
        return resolve_branch_extraction_inputs(
            branch,
            extract_params=self._extract_params,
            text_params=self._text_params,
            html_params=self._html_params,
            audio_chunk_params=self._audio_chunk_params,
            asr_params=self._asr_params,
            video_frame_params=self._video_frame_params,
            video_text_dedup_params=self._video_text_dedup_params,
            av_fuse_params=self._av_fuse_params,
        )

    def _resolve_effective_extraction_inputs(self) -> ResolvedExtractionInputs:
        extraction_mode = self._extraction_mode
        classified = self._classified_input_paths()
        if extraction_mode is not None:
            self._validate_explicit_extraction_mode_inputs(extraction_mode, classified)
            text_params = self._text_params
            html_params = self._html_params
            if extraction_mode == "auto":
                observed_input_types = {input_type for _, input_type in classified if input_type is not None}
                if "txt" in observed_input_types:
                    text_params = text_params or TextChunkParams()
                if "html" in observed_input_types:
                    html_params = html_params or HtmlChunkParams()
            return ResolvedExtractionInputs(
                extraction_mode=extraction_mode,
                extract_params=self._extract_params,
                text_params=text_params,
                html_params=html_params,
                audio_chunk_params=self._audio_chunk_params,
                asr_params=self._asr_params,
                video_frame_params=self._video_frame_params,
                video_text_dedup_params=self._video_text_dedup_params,
                av_fuse_params=self._av_fuse_params,
            )

        branches = self._plan_default_extraction_branches()
        if branches is None:
            raise RuntimeError("Internal error: default extraction planning did not return branches.")
        if len(branches) == 1:
            return self._resolve_branch_extraction_inputs(branches[0])

        # Compatibility fallback for private callers that still ask for a
        # scalar effective mode directly. The public ingest path executes the
        # branches instead of using this MultiType fallback.
        return ResolvedExtractionInputs(
            extraction_mode="auto",
            extract_params=self._extract_params or ExtractParams(),
            text_params=self._text_params or TextChunkParams(),
            html_params=self._html_params or HtmlChunkParams(),
            audio_chunk_params=self._audio_chunk_params,
            asr_params=self._asr_params,
            video_frame_params=self._video_frame_params,
            video_text_dedup_params=self._video_text_dedup_params,
            av_fuse_params=self._av_fuse_params,
        )

    @staticmethod
    def _is_populated_error_field(key: str, value: Any) -> bool:
        return is_populated_error_field(key, value)

    @classmethod
    def _iter_stage_errors_from_value(cls, value: Any, *, path: str = "") -> Iterator[dict[str, Any]]:
        yield from iter_stage_errors_from_value(value, path=path)

    @staticmethod
    def _row_value(row: Any, key: str) -> Any:
        if isinstance(row, dict):
            return row.get(key)
        getter = getattr(row, "get", None)
        if callable(getter):
            try:
                return getter(key)
            except Exception as exc:  # noqa: BLE001 - row metadata lookup is best-effort diagnostic context.
                logger.debug(
                    "Failed to read source identifier field %r from row type %s: %s",
                    key,
                    type(row).__name__,
                    exc,
                    exc_info=True,
                )
                return None
        return None

    @staticmethod
    def _nested_mapping_value(value: Any, path: tuple[str, ...]) -> Any:
        current = value
        for key in path:
            if not isinstance(current, dict):
                return None
            current = current.get(key)
        return current

    @classmethod
    def _source_identifier_from_row(cls, row: Any, row_index: Any) -> str:
        for field in ("document_id", "path", "source_path"):
            identifier = _coerce_source_identifier(cls._row_value(row, field))
            if identifier is not None:
                return identifier

        metadata = cls._row_value(row, "metadata")
        for nested_path in (
            ("source_path",),
            ("source_metadata", "source_id"),
            ("source_metadata", "source_name"),
        ):
            identifier = _coerce_source_identifier(cls._nested_mapping_value(metadata, nested_path))
            if identifier is not None:
                return identifier

        for field in ("source_id", "source_name"):
            identifier = _coerce_source_identifier(cls._row_value(row, field))
            if identifier is not None:
                return identifier

        return f"row {row_index}" if row_index is not None else "row ?"

    @staticmethod
    def _public_failure_tuple(record: dict[str, Any]) -> tuple[str, str]:
        identifier = _coerce_source_identifier(record.get("source_identifier"))
        if identifier is None:
            row_index = record.get("row_index")
            identifier = f"row {row_index}" if row_index is not None else "row ?"
        return identifier, _format_public_failure_message(record)

    @classmethod
    def _stage_error_records(cls, batch: Any, *, columns: Iterable[str] | None = None) -> list[dict[str, Any]]:
        iter_batches = getattr(batch, "iter_batches", None)
        if getattr(batch, "columns", None) is None and not callable(iter_batches):
            return []
        requested_columns = list(columns) if columns is not None else None

        batches = iter_batches(batch_format="pyarrow") if callable(iter_batches) else (batch,)

        records: list[dict[str, Any]] = []
        for raw_batch in batches:
            available_columns = (
                getattr(raw_batch, "column_names", None)
                if callable(iter_batches)
                else getattr(raw_batch, "columns", None)
            )
            if available_columns is None:
                continue
            target_columns = (
                list(available_columns)
                if requested_columns is None
                else [c for c in requested_columns if c in available_columns]
            )
            if not target_columns:
                continue
            # Finalization only needs diagnostic payloads and source context.
            # Reading the full frame can iterate unrelated Arrow-backed image
            # columns that pandas cannot safely materialize row by row.
            scan_columns = list(
                dict.fromkeys(
                    [*target_columns, *(column for column in _SOURCE_IDENTIFIER_COLUMNS if column in available_columns)]
                )
            )
            batch_df = (
                arrow_table_to_pandas(raw_batch.select(scan_columns))
                if callable(iter_batches)
                else raw_batch.loc[:, scan_columns]
            )
            column_values = {column: batch_df[column].array for column in scan_columns}
            for row_position, row_index in enumerate(batch_df.index):
                row = {column: values[row_position] for column, values in column_values.items()}
                source_identifier = cls._source_identifier_from_row(row, row_index)
                for column in target_columns:
                    for record in cls._iter_stage_errors_from_value(row[column]):
                        records.append(
                            {
                                "row_index": row_index,
                                "source_identifier": source_identifier,
                                "column": column,
                                **record,
                            }
                        )
        return records

    @staticmethod
    def _has_error(v: Any) -> bool:
        return any(GraphIngestor._iter_stage_errors_from_value(v))

    @staticmethod
    def _param_value(params: Any, field: str) -> Any:
        if params is None:
            return None
        if isinstance(params, dict):
            return params.get(field)
        return getattr(params, field, None)

    @classmethod
    def _is_configured(cls, value: Any) -> bool:
        if value is None:
            return False
        if isinstance(value, str):
            return bool(value.strip())
        if isinstance(value, (list, tuple, set)):
            return any(cls._is_configured(v) for v in value)
        return bool(value)

    @classmethod
    def _params_has_configured_field(cls, params: Any, fields: tuple[str, ...]) -> bool:
        return any(cls._is_configured(cls._param_value(params, field)) for field in fields)

    def _remote_stage_error_columns(self) -> set[str]:
        """Backwards-compatible thin shim over :meth:`_remote_stage_diagnostics`.

        Older callers (and existing tests) consume the set of columns
        the strict-error-policy will gate on. The richer
        :meth:`_remote_stage_diagnostics` mapping carries the same set
        of keys plus per-stage NIM URL / display-name diagnostics that
        :class:`GraphIngestionError` uses to format actionable messages.
        """
        return set(self._remote_stage_diagnostics().keys())

    def _remote_stage_diagnostics(self) -> dict[str, _StageDiagnostic]:
        """Build a column → :class:`_StageDiagnostic` map for remote-NIM stages.

        Only stages that have an explicitly configured invoke URL appear
        here — the ``"raise"`` error policy is scoped to remote endpoints
        the operator opted into. The map's keys are the dataframe column
        names emitted by each stage; the values carry the resolved
        display name and URL so :class:`GraphIngestionError` can render
        ``stage=… url=…`` per row and a ``Troubleshooting:`` footer.
        """
        diagnostics: dict[str, _StageDiagnostic] = {}

        extract = self._extract_params
        if self._params_has_configured_field(extract, ("page_elements_invoke_url",)):
            column = self._param_value(extract, "output_column") or _DEFAULT_PAGE_ELEMENTS_COLUMN
            diagnostics[column] = _StageDiagnostic(
                column=column,
                display_name="Page Elements NIM",
                invoke_url=self._param_value(extract, "page_elements_invoke_url"),
                role="page_elements",
            )
        if self._params_has_configured_field(extract, ("ocr_invoke_url",)):
            diagnostics["ocr"] = _StageDiagnostic(
                column="ocr",
                display_name="OCR NIM",
                invoke_url=self._param_value(extract, "ocr_invoke_url"),
                role="ocr",
            )
        if self._params_has_configured_field(extract, ("table_structure_invoke_url",)):
            diagnostics["table_structure_ocr_v1"] = _StageDiagnostic(
                column="table_structure_ocr_v1",
                display_name="Table Structure NIM",
                invoke_url=self._param_value(extract, "table_structure_invoke_url"),
                role="table_structure",
            )
        if self._params_has_configured_field(extract, ("invoke_url", "nemotron_parse_invoke_url")):
            url = self._param_value(extract, "nemotron_parse_invoke_url") or self._param_value(extract, "invoke_url")
            diagnostics["nemotron_parse_v1_2"] = _StageDiagnostic(
                column="nemotron_parse_v1_2",
                display_name="Nemotron Parse NIM",
                invoke_url=url,
                role="nemotron_parse",
            )
        if self._params_has_configured_field(self._embed_params, _REMOTE_EMBED_ENDPOINT_FIELDS):
            column = self._param_value(self._embed_params, "output_column") or _DEFAULT_EMBED_COLUMN
            url = self._param_value(self._embed_params, "embed_invoke_url") or self._param_value(
                self._embed_params, "embedding_endpoint"
            )
            diagnostics[column] = _StageDiagnostic(
                column=column,
                display_name="Embedding NIM",
                invoke_url=url,
                model_name=self._param_value(self._embed_params, "model_name"),
                role="embed",
            )
        return diagnostics

    def _raise_for_stage_errors(self, result: Any) -> None:
        if self._error_policy == "collect":
            return
        diagnostics = self._remote_stage_diagnostics()
        if not diagnostics:
            return
        records = self._stage_error_records(result, columns=set(diagnostics.keys()))
        if records:
            raise GraphIngestionError(records, stage_diagnostics=diagnostics)

    @staticmethod
    def _resolve_return_failures(params: Any, kwargs: dict[str, Any]) -> bool:
        if "return_failures" in kwargs:
            return bool(kwargs["return_failures"])
        if isinstance(params, IngestExecuteParams):
            return bool(params.return_failures)
        if isinstance(params, dict) and "return_failures" in params:
            return bool(params["return_failures"])
        return False

    def _collect_failure_records(self, result: Any) -> list[dict[str, Any]]:
        diagnostics = self._remote_stage_diagnostics()
        # With explicit remote stages, report only their diagnostic columns.
        # Without them, scan all columns so ``return_failures=True`` can expose
        # local collected failures instead of silently returning an empty list.
        columns = set(diagnostics.keys()) if diagnostics else None
        return self._stage_error_records(result, columns=columns)

    def _collect_failure_tuples(self, result: Any) -> list[tuple[str, str]]:
        return [self._public_failure_tuple(record) for record in self._collect_failure_records(result)]

    def _finalize_ingest_result(self, result: Any, *, return_failures: bool) -> Any:
        if return_failures:
            return result, self._collect_failure_tuples(result)
        self._raise_for_stage_errors(result)
        return result

    @staticmethod
    def extract_error_rows(batch: Any) -> Any:
        if batch is None:
            return batch
        columns = getattr(batch, "columns", None)
        if columns is None:
            return batch
        if len(columns) == 0:
            return batch.iloc[0:0]

        mask = batch[columns[0]].apply(GraphIngestor._has_error).astype(bool)
        for c in columns[1:]:
            mask = mask | batch[c].apply(GraphIngestor._has_error).astype(bool)
        return batch[mask]

    def get_error_rows(self, dataset: Any = None) -> Any:
        import pandas as pd

        target = dataset if dataset is not None else self._rd_dataset
        if target is None:
            raise RuntimeError("No Ray Dataset available to inspect for errors.")
        if isinstance(target, pd.DataFrame):
            return self.extract_error_rows(target)
        return target.map_batches(
            call_pandas_function_on_arrow,
            batch_format="pyarrow",
            fn_kwargs={"fn": self.extract_error_rows},
        )

    def get_dataset(self) -> Any:
        return self._rd_dataset

    def _record_stage(self, name: str) -> None:
        """Append *name* to the stage order list (deduplicated in place)."""
        if name not in self._stage_order:
            self._stage_order.append(name)

    def _apply_split_config(self, split_config: dict[str, Any] | None) -> None:
        """Resolve an explicitly supplied split configuration."""
        if split_config is not None:
            self._split_config = resolve_split_params(split_config)

RUN_MODE = 'graph' class-attribute instance-attribute

files(documents)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def files(self, documents: Union[str, List[str]]) -> "GraphIngestor":
    """Set the input file paths or glob patterns."""
    self._documents = [documents] if isinstance(documents, str) else list(documents)
    return self

texts(texts)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def texts(self, texts: Union[str, Sequence[str]]) -> Self:
    """Set raw inline text documents as the graph input.

    Each string is one logical source document. It receives a deterministic
    ``inline://`` identifier and flows through the normal text splitter,
    embedding, and sink stages without being written to a temporary file.
    Inline text may be combined with file or buffer inputs; the manifest
    planner routes each source through its matching extraction branch.
    """
    self._inline_texts = normalize_inline_texts(texts)
    return self

buffers(buffers)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def buffers(
    self,
    buffers: Union[Tuple[str, BytesIO], List[Tuple[str, BytesIO]]],
) -> "GraphIngestor":
    """Set in-memory buffers for processing.

    Each buffer is a ``(name, BytesIO)`` pair where *name* carries the
    original filename (including extension) so downstream operators can
    detect file type.  Accepts a single tuple or a list of tuples.

    Only supported for ``run_mode='inprocess'``.
    """
    if isinstance(buffers, tuple) and len(buffers) == 2 and isinstance(buffers[0], str):
        self._buffers = [buffers]
    else:
        self._buffers = list(buffers)
    return self

extract(params=None, *, split_config=None, extraction_mode=None, text_params=None, html_params=None, audio_chunk_params=None, asr_params=None, video_frame_params=None, video_text_dedup_params=None, av_fuse_params=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def extract(
    self,
    params: Optional[ExtractParams] = None,
    *,
    split_config: dict[str, Any] | None = None,
    extraction_mode: str | None = None,
    text_params: Optional[TextChunkParams] = None,
    html_params: Optional[HtmlChunkParams] = None,
    audio_chunk_params: Optional[AudioChunkParams] = None,
    asr_params: Optional[ASRParams] = None,
    video_frame_params: Optional[VideoFrameParams] = None,
    video_text_dedup_params: Optional[VideoFrameTextDedupParams] = None,
    av_fuse_params: Optional[AudioVisualFuseParams] = None,
    **kwargs: Any,
) -> "GraphIngestor":
    """Configure extraction.

    By default, the effective extraction mode is inferred from the input
    file extensions immediately before graph construction. Pass
    ``extraction_mode='pdf'`` to force the dedicated PDF/document graph, or
    ``extraction_mode='auto'`` to dispatch a mixed folder through
    :class:`MultiTypeExtractOperator`.
    Chunking is opt-in: pass ``split_config={"<key>": {...}}`` to enable
    post-extract token chunking for that source type.

    Unknown ``**kwargs`` raise :class:`TypeError`. Only fields declared
    on :class:`ExtractParams` are accepted as extra kwargs; ASR / audio
    configuration belongs on :class:`ASRParams` (pass ``asr_params=``
    or use :meth:`extract_audio`).
    """
    unknown = set(kwargs) - set(ExtractParams.model_fields)
    if unknown:
        raise TypeError(
            f"extract() got unexpected keyword argument(s) {sorted(unknown)!r}. "
            f"Allowed extra kwargs must be fields of ExtractParams. "
            f"For ASR / audio configuration, pass asr_params=ASRParams(...) "
            f"or use .extract_audio(asr_params=ASRParams(...)) "
            f"(see docs/extraction/audio-video.md)."
        )
    self._extraction_mode = extraction_mode
    self._extract_params = _resolve_api_key(_coerce(params, kwargs, default_factory=ExtractParams))
    if text_params is not None:
        self._text_params = text_params
    if html_params is not None:
        self._html_params = html_params
    if audio_chunk_params is not None:
        self._audio_chunk_params = audio_chunk_params
    if asr_params is not None:
        self._asr_params = asr_params
    if video_frame_params is not None:
        self._video_frame_params = video_frame_params
    if video_text_dedup_params is not None:
        self._video_text_dedup_params = video_text_dedup_params
    if av_fuse_params is not None:
        self._av_fuse_params = av_fuse_params
    self._apply_split_config(split_config)
    self._record_stage("extract")
    return self

extract_image_files(params=None, *, split_config=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def extract_image_files(
    self,
    params: Optional[ExtractParams] = None,
    *,
    split_config: dict[str, Any] | None = None,
    **kwargs: Any,
) -> "GraphIngestor":
    """Configure image extraction (extraction_mode='image')."""
    self._extraction_mode = "image"
    self._extract_params = _resolve_api_key(_coerce(params, kwargs, default_factory=ExtractParams))
    self._apply_split_config(split_config)
    self._record_stage("extract")
    return self

extract_html(params=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def extract_html(self, params: Optional[HtmlChunkParams] = None, **kwargs: Any) -> "GraphIngestor":
    """Configure HTML extraction (extraction_mode='html')."""
    self._extraction_mode = "html"
    self._html_params = _coerce(params, kwargs, default_factory=HtmlChunkParams)
    self._record_stage("extract")
    return self

extract_audio(params=None, *, asr_params=None, split_config=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def extract_audio(
    self,
    params: Optional[AudioChunkParams] = None,
    *,
    asr_params: Optional[ASRParams] = None,
    split_config: dict[str, Any] | None = None,
    **kwargs: Any,
) -> "GraphIngestor":
    """Configure audio extraction (extraction_mode='audio')."""
    self._extraction_mode = "audio"
    self._audio_chunk_params = _coerce(params, kwargs, default_factory=AudioChunkParams)
    self._asr_params = asr_params or ASRParams()
    self._apply_split_config(split_config)
    self._record_stage("extract")
    return self

extract_video(params=None, *, asr_params=None, video_frame_params=None, video_text_dedup_params=None, av_fuse_params=None, extract_params=None, split_config=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def extract_video(
    self,
    params: Optional[AudioChunkParams] = None,
    *,
    asr_params: Optional[ASRParams] = None,
    video_frame_params: Optional[VideoFrameParams] = None,
    video_text_dedup_params: Optional[VideoFrameTextDedupParams] = None,
    av_fuse_params: Optional[AudioVisualFuseParams] = None,
    extract_params: Optional[ExtractParams] = None,
    split_config: dict[str, Any] | None = None,
    **kwargs: Any,
) -> "GraphIngestor":
    """Configure video extraction.

    Sets ``extraction_mode='auto'`` so :class:`MultiTypeExtractOperator`
    dispatches by file extension; ``.mp4``/``.mov``/``.mkv``
    files are routed to a combined audio-from-video ASR + frame OCR +
    scene fusion pipeline.

    Frame OCR config (``ocr_invoke_url``, ``ocr_api_key``,
    ``inference_batch_size``, ``ocr_request_timeout_s``) is read from
    :class:`ExtractParams` — the same object the PDF/image pipelines
    use — so the user only configures OCR once.

    The ``split_config`` keyword honors the ``"video"`` key (chunking the
    fused audio+visual transcript). The ``"audio"`` key is ignored on the
    video pipeline — for audio-only chunking, use :meth:`extract_audio`
    directly with that file.
    """
    self._extraction_mode = "auto"
    self._audio_chunk_params = _coerce(params, kwargs, default_factory=AudioChunkParams)
    self._asr_params = asr_params or ASRParams()
    self._video_frame_params = video_frame_params or VideoFrameParams()
    self._video_text_dedup_params = video_text_dedup_params or VideoFrameTextDedupParams()
    self._av_fuse_params = av_fuse_params or AudioVisualFuseParams()
    if extract_params is not None:
        self._extract_params = _resolve_api_key(extract_params)
    elif self._extract_params is None:
        self._extract_params = ExtractParams()
    self._apply_split_config(split_config)
    self._record_stage("extract")
    return self

dedup(params=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def dedup(self, params: Optional[DedupParams] = None, **kwargs: Any) -> "GraphIngestor":
    """Record a dedup stage."""
    self._dedup_params = _coerce(params, kwargs, default_factory=DedupParams)
    self._record_stage("dedup")
    return self

caption(params=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def caption(self, params: Optional[CaptionParams] = None, **kwargs: Any) -> "GraphIngestor":
    """Record a caption stage."""
    self._caption_params = _resolve_api_key(_coerce(params, kwargs, default_factory=CaptionParams))
    self._record_stage("caption")
    return self

store(params=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def store(self, params: Optional[StoreParams] = None, **kwargs: Any) -> "GraphIngestor":
    """Record a store stage for persisting extracted image assets to storage."""
    self._store_params = _coerce(params, kwargs, default_factory=StoreParams)
    self._record_stage("store")
    return self

embed(params=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def embed(self, params: Optional[EmbedParams] = None, **kwargs: Any) -> "GraphIngestor":
    """Record an embedding stage."""
    self._embed_params = _resolve_api_key(_coerce(params, kwargs, default_factory=EmbedParams))
    self._record_stage("embed")
    return self

vdb_upload(params=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def vdb_upload(self, params: Optional[VdbUploadParams] = None, **kwargs: Any) -> "GraphIngestor":
    """Record a vector DB upload **sink** (in-graph after embed/store, before webhook).

    Does not call :meth:`_record_stage`: ``stage_order`` only lists
    ``dedup`` / ``caption`` / ``store`` / ``embed`` for reordering; VDB is
    always appended from ``_vdb_upload_params`` in
    :func:`~nemo_retriever.graph.ingestor_runtime._append_ordered_transform_stages`.
    Plan builders that round-trip sinks use :meth:`~nemo_retriever.ingest_plans.BaseIngestPlan.record_sink`.
    """
    self._vdb_upload_params = _coerce(params, kwargs, default_factory=VdbUploadParams)
    return self

webhook(params=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def webhook(self, params: Optional[WebhookParams] = None, **kwargs: Any) -> "GraphIngestor":
    """Record a webhook notification stage (always runs last).

    When ``endpoint_url`` is set, processed results are HTTP-POSTed to
    that URL.  If ``endpoint_url`` is ``None`` the stage is a no-op.
    """
    self._webhook_params = _coerce(params, kwargs, default_factory=WebhookParams)
    self._record_stage("webhook")
    return self

ingest(params=None, **kwargs)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def ingest(self, params: Any = None, **kwargs: Any) -> Any:
    """Build the operator graph and run it through the configured executor.

    Parameters
    ----------
    params
        Optional :class:`IngestExecuteParams` (or plain ``dict``) carrying
        execute-time flags. Graph run modes honor ``return_failures``.
    **kwargs
        Execute-time flags passed directly. ``return_failures`` may be
        passed here and takes precedence over the value in ``params``.
    return_failures
        When ``True`` (default ``False``), return ``(result, failures)``
        instead of raising collected row-level stage errors. If no explicit
        remote-stage diagnostics are configured, all output columns are
        scanned for populated error fields so local collected failures can
        still be returned; the default raise path remains scoped to
        explicitly configured remote stages.

    Returns
    -------
    ``run_mode='batch'`` or ``run_mode='inprocess'``
        A ``pandas.DataFrame``.
    ``return_failures=True``
        ``(result, failures)`` where ``failures`` is a list of
        service-style ``(source, error)`` tuples.
    """
    return_failures = self._resolve_return_failures(params, kwargs)
    self._validate_input_sources(self._inline_texts)
    if not self._documents and not self._buffers and is_blank_inline_corpus(self._inline_texts):
        result = empty_text_chunks_df()
        if self._run_mode == "batch":
            self._rd_dataset = result
        else:
            self._rd_dataset = None
        return self._finalize_ingest_result(result, return_failures=return_failures)

    default_branches = self._plan_default_extraction_branches()
    execute_branches = default_branches is not None and (
        len(default_branches) > 1 or self._has_mixed_inline_sources()
    )
    if default_branches is None:
        single_effective = self._resolve_effective_extraction_inputs()
    elif not execute_branches:
        single_effective = self._resolve_branch_extraction_inputs(default_branches[0])
    else:
        single_effective = None

    # Auto-enable dedup before captioning so that images overlapping
    # with table/chart/infographic detections are removed first.
    # Skip for image-only extraction — the image IS the content.
    image_only = single_effective is not None and single_effective.extraction_mode == "image"
    if self._caption_params is not None and self._dedup_params is None and not image_only:
        self._dedup_params = DedupParams()
        if "dedup" not in self._stage_order:
            try:
                idx = self._stage_order.index("caption")
            except ValueError:
                idx = len(self._stage_order)
            self._stage_order.insert(idx, "dedup")

    post_extract_order = tuple(s for s in self._stage_order if s != "extract")

    if execute_branches:
        result = self._execute_extraction_branches(default_branches, post_extract_order=post_extract_order)
    else:
        if single_effective is None:
            raise RuntimeError("Internal error: extraction inputs were not resolved.")
        result = self._execute_single_graph(single_effective, post_extract_order=post_extract_order)

    return self._finalize_ingest_result(result, return_failures=return_failures)

extract_error_rows(batch) staticmethod

Source code in nemo_retriever/ingestor/graph_ingestor.py

@staticmethod
def extract_error_rows(batch: Any) -> Any:
    if batch is None:
        return batch
    columns = getattr(batch, "columns", None)
    if columns is None:
        return batch
    if len(columns) == 0:
        return batch.iloc[0:0]

    mask = batch[columns[0]].apply(GraphIngestor._has_error).astype(bool)
    for c in columns[1:]:
        mask = mask | batch[c].apply(GraphIngestor._has_error).astype(bool)
    return batch[mask]

get_error_rows(dataset=None)

Source code in nemo_retriever/ingestor/graph_ingestor.py

def get_error_rows(self, dataset: Any = None) -> Any:
    import pandas as pd

    target = dataset if dataset is not None else self._rd_dataset
    if target is None:
        raise RuntimeError("No Ray Dataset available to inspect for errors.")
    if isinstance(target, pd.DataFrame):
        return self.extract_error_rows(target)
    return target.map_batches(
        call_pandas_function_on_arrow,
        batch_format="pyarrow",
        fn_kwargs={"fn": self.extract_error_rows},
    )

get_dataset()

Source code in nemo_retriever/ingestor/graph_ingestor.py

def get_dataset(self) -> Any:
    return self._rd_dataset

nemo_retriever.ingestor.graph_ingestor.GraphIngestionError

Bases: RuntimeError

Source code in nemo_retriever/ingestor/graph_ingestor.py

class GraphIngestionError(RuntimeError):
    """Raised when graph ingestion stages report structured row-level errors.

    The exception message is built to be self-diagnosing: when the
    caller provides ``stage_diagnostics`` (a mapping from the dataframe
    column the error landed in to a :class:`_StageDiagnostic` describing
    the originating NIM), each row in the rendered message names the
    stage and the configured invoke URL, and the message gains a
    ``Troubleshooting:`` footer with concrete next steps for the
    observed (stage, HTTP status) tuples.

    Backwards compatible signature: ``GraphIngestionError(records)``
    still works and produces the legacy message shape.
    """

    def __init__(
        self,
        records: list[Any],
        stage_diagnostics: dict[str, _StageDiagnostic] | None = None,
    ) -> None:
        self.records = records
        self.stage_diagnostics = dict(stage_diagnostics) if stage_diagnostics else {}
        super().__init__(_format_stage_error_message(records, self.stage_diagnostics))

records = records instance-attribute

stage_diagnostics = dict(stage_diagnostics) if stage_diagnostics else {} instance-attribute

Service ingest

create_ingestor(run_mode="service") returns ServiceIngestor. Import the class from nemo_retriever.service.service_ingestor. ingest() returns ServiceIngestResult.

Service-only methods include split(), pdf_split_config(), save_to_disk(), ingest_stream(), aingest_stream(), and cancel(). cancel() is part of the public class. It currently raises NotImplementedError because the service does not expose a cancel endpoint.

nemo_retriever.service.service_ingestor.ServiceIngestor

Bases: ingestor

Parameters:

Name	Type	Description	Default
base_url	str	Base URL of the retriever service (default http://localhost:7670).	'http://localhost:7670'
documents	Optional[List[str]]	Initial list of file paths to ingest; may also be set/extended via :meth:files and :meth:buffers.	None
max_concurrency	int	Maximum concurrent document uploads (default 8).	8
request_timeout_s	float	Per-request HTTP timeout (default 600s for large documents).	600.0
api_token	str | None	Optional bearer token for service authentication.	None

Source code in nemo_retriever/service/service_ingestor.py

class ServiceIngestor(ingestor):
    """Ingestor that submits work to a running ``retriever service``.

    Parameters
    ----------
    base_url
        Base URL of the retriever service (default ``http://localhost:7670``).
    documents
        Initial list of file paths to ingest; may also be set/extended via
        :meth:`files` and :meth:`buffers`.
    max_concurrency
        Maximum concurrent document uploads (default 8).
    request_timeout_s
        Per-request HTTP timeout (default 600s for large documents).
    api_token
        Optional bearer token for service authentication.
    """

    RUN_MODE = "service"

    def __init__(
        self,
        *,
        base_url: str = "http://localhost:7670",
        documents: Optional[List[str]] = None,
        max_concurrency: int = 8,
        request_timeout_s: float = 600.0,
        api_token: str | None = None,
    ) -> None:
        super().__init__(documents=documents)
        self._base_url = base_url.rstrip("/")
        self._max_concurrency = max_concurrency
        self._request_timeout_s = request_timeout_s
        self._api_token = (api_token or "").strip() or None
        self._inline_texts: list[str] | None = None
        self._document_ids: list[str] = []
        self._last_run_elapsed_s: float = 0.0
        self._last_job_id: str | None = None
        self._pipeline_spec: dict[str, Any] = {
            "extraction_mode": "auto",
            "stage_order": [],
        }
        # save_to_disk state (populated by .save_to_disk(...); None when disabled)
        self._save_to_disk_dir: Path | None = None
        self._save_to_disk_compression: str | None = None
        self._save_to_disk_cleanup: bool = True

    # ------------------------------------------------------------------
    # Pipeline-spec helpers
    # ------------------------------------------------------------------

    def _record_stage(self, name: str) -> None:
        order = self._pipeline_spec["stage_order"]
        if name not in order:
            order.append(name)

    def _new_result_fetch_client(self) -> httpx.Client:
        """Create a client for retained-result status requests."""
        return httpx.Client(timeout=self._request_timeout_s, headers=self._auth_headers)

    def _fetch_document_result_data(
        self,
        document_id: str,
        *,
        client: httpx.Client | None = None,
    ) -> list[dict[str, Any]]:
        """Fetch ``result_data`` for *document_id* from the status endpoint.

        The status endpoint retains ``result_data`` through the job retention
        window, so retrying this read is safe. When the caller supplies a
        client it is reused for the first attempt. A transient failure receives
        one retry through a fresh client so a stale pooled connection cannot be
        selected again.
        """
        if not document_id:
            raise ValueError("_fetch_document_result_data(): empty document_id")

        if client is None:
            with self._new_result_fetch_client() as scoped_client:
                return self._fetch_document_result_data(document_id, client=scoped_client)

        url = f"{self._base_url}/v1/ingest/status/{document_id}"
        try:
            resp = client.get(url)
        except _RESULT_FETCH_TRANSIENT_ERRORS as exc:
            logger.debug(
                "Transient %s fetching retained result for %s; retrying on a fresh connection",
                type(exc).__name__,
                document_id,
            )
            with self._new_result_fetch_client() as retry_client:
                resp = retry_client.get(url)
        resp.raise_for_status()
        body = resp.json()
        return list(body.get("result_data") or [])

    def _write_result_data_to_disk(self, document_id: str, result_data: list[dict[str, Any]]) -> Path:
        """Write *result_data* for *document_id* to the configured output directory."""
        import gzip
        import json as _json

        if self._save_to_disk_dir is None:
            raise RuntimeError("_write_result_data_to_disk(): save_to_disk was never enabled")

        suffix = ".json.gz" if self._save_to_disk_compression == "gzip" else ".json"
        out_path = self._save_to_disk_dir / f"{document_id}{suffix}"
        payload = _json.dumps(
            {"document_id": document_id, "rows": result_data},
            ensure_ascii=False,
        ).encode("utf-8")
        if self._save_to_disk_compression == "gzip":
            with gzip.open(out_path, "wb") as fh:
                fh.write(payload)
        else:
            out_path.write_bytes(payload)
        return out_path

    def _save_document_to_disk(
        self,
        document_id: str,
        *,
        client: httpx.Client | None = None,
    ) -> Path:
        """Fetch ``result_data`` for *document_id* and write a JSON artifact.

        Returns the path that was written. Raises if the document_id is
        missing or the fetch fails.
        """
        if self._save_to_disk_dir is None:
            raise RuntimeError("_save_document_to_disk(): save_to_disk was never enabled")
        result_data = self._fetch_document_result_data(document_id, client=client)
        return self._write_result_data_to_disk(document_id, result_data)

    def _materialize_completed_document(
        self,
        document_id: str,
        *,
        return_results: bool,
        client: httpx.Client | None = None,
    ) -> list[dict[str, Any]] | None:
        """Fetch (once) and optionally persist rows for a completed document."""
        if not return_results and self._save_to_disk_dir is None:
            return None
        result_data = self._fetch_document_result_data(document_id, client=client)
        if self._save_to_disk_dir is not None:
            self._write_result_data_to_disk(document_id, result_data)
        return result_data if return_results else None

    def _pipeline_payload(
        self,
        *,
        result_schema: ResultSchema = "legacy",
        return_embeddings: bool = False,
        return_images: bool = False,
    ) -> dict[str, Any] | None:
        """Return the spec dict to send on the wire, or ``None`` when empty.

        The "empty" check mirrors :meth:`PipelineSpec.is_empty` server-side
        so the worker can short-circuit identically.
        """
        spec = dict(self._pipeline_spec)
        if self._has_mixed_inline_sources():
            spec["extraction_mode"] = "auto"
        spec["result_schema"] = result_schema
        spec["return_embeddings"] = bool(return_embeddings or spec.get("return_embeddings", False))
        spec["return_images"] = bool(return_images or spec.get("return_images", False))
        is_empty = (
            spec.get("extraction_mode", "auto") in ("pdf", "auto")
            and not spec.get("stage_order")
            and not any(
                spec.get(k)
                for k in (
                    "extract_params",
                    "embed_params",
                    "dedup_params",
                    "caption_params",
                    "store_params",
                    "vdb_upload_params",
                    "webhook_params",
                    "split_config",
                    "pdf_split",
                )
            )
            and spec.get("result_schema", "legacy") == "legacy"
            and not spec.get("return_embeddings", False)
            and not spec.get("return_images", False)
        )
        return None if is_empty else spec

    @property
    def _auth_headers(self) -> dict[str, str]:
        return {"Authorization": f"Bearer {self._api_token}"} if self._api_token else {}

    # ------------------------------------------------------------------
    # Input configuration (these ARE meaningful client-side)
    # ------------------------------------------------------------------

    def files(self, documents: Union[str, List[str]]) -> "ServiceIngestor":
        """Add document paths/URIs for processing."""
        if isinstance(documents, str):
            self._documents.append(documents)
        else:
            self._documents.extend(documents)
        return self

    def texts(self, texts: Union[str, Sequence[str]]) -> Self:
        """Set raw inline text documents, optionally alongside file or buffer uploads."""
        self._inline_texts = normalize_inline_texts(texts)
        return self

    def buffers(
        self,
        buffers: Union[Tuple[str, BytesIO], List[Tuple[str, BytesIO]]],
    ) -> "ServiceIngestor":
        """Add in-memory buffers for processing.

        Each buffer must be ``(filename, BytesIO)`` so the server can record
        a meaningful source filename.
        """
        if isinstance(buffers, tuple):
            buffers = [buffers]
        for name, buf in buffers:
            self._buffers.append((name, buf))
        return self

    def load(self) -> "ServiceIngestor":
        """No-op for service mode."""
        return self

    # ------------------------------------------------------------------
    # Phase 1: pipeline-shape stages — sent via PipelineSpec
    # ------------------------------------------------------------------

    def all_tasks(self) -> "ServiceIngestor":
        """Record the default chain: extract → dedup → embed.

        Concrete params come from server config; ``all_tasks()`` only
        controls *stage order* and is the closest in-process equivalent
        of "run everything the server is configured to do".
        """
        self._record_stage("extract")
        self._record_stage("dedup")
        self._record_stage("embed")
        return self

    def dedup(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
        """Record a dedup stage with optional :class:`DedupParams` overrides."""
        if params is not None or kwargs:
            from nemo_retriever.common.policy import _DEFAULT_ALLOWED_DEDUP_KEYS

            merged = _merge_params(params, kwargs)
            _wire_client_stage_params(
                self._pipeline_spec,
                "dedup_params",
                merged,
                method="dedup",
                allowed=_DEFAULT_ALLOWED_DEDUP_KEYS,
            )
        self._record_stage("dedup")
        return self

    def embed(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
        """Record an embed stage with optional :class:`EmbedParams` overrides.

        Embedding endpoint URL and API key are server-owned and will be
        rejected if set here.
        """
        if params is not None or kwargs:
            from nemo_retriever.common.policy import _DEFAULT_ALLOWED_EMBED_KEYS

            merged = _merge_params(params, kwargs)
            _wire_client_stage_params(
                self._pipeline_spec,
                "embed_params",
                merged,
                method="embed",
                allowed=_DEFAULT_ALLOWED_EMBED_KEYS,
            )
        self._record_stage("embed")
        return self

    def extract(
        self,
        params: Any = None,
        *,
        split_config: Optional[dict[str, Any]] = None,
        extraction_mode: str = "auto",
        **kwargs: Any,
    ) -> "ServiceIngestor":
        """Record a generic extraction stage.

        ``extraction_mode`` selects the worker's extraction path
        (``'auto'`` default — dispatches by file extension; ``'pdf'``
        forces the PDF path for all inputs, etc.).

        When no ``ExtractParams`` overrides are supplied, ``extract_params``
        is omitted from the wire payload so the worker applies the
        service's server-owned defaults (and the allow-list is not tripped
        by client-side model defaults).
        """
        if params is not None or kwargs:
            from nemo_retriever.common.policy import _DEFAULT_ALLOWED_EXTRACT_KEYS

            merged = _merge_params(params, kwargs)
            _wire_client_stage_params(
                self._pipeline_spec,
                "extract_params",
                merged,
                method="extract",
                allowed=_DEFAULT_ALLOWED_EXTRACT_KEYS,
            )
        self._pipeline_spec["extraction_mode"] = extraction_mode
        if split_config is not None:
            self._pipeline_spec["split_config"] = split_config
        self._record_stage("extract")
        return self

    def extract_image_files(
        self, params: Any = None, *, split_config: Optional[dict[str, Any]] = None, **kwargs: Any
    ) -> "ServiceIngestor":
        """Record image-file extraction (``extraction_mode='image'``)."""
        if params is not None or kwargs:
            from nemo_retriever.common.policy import _DEFAULT_ALLOWED_EXTRACT_KEYS

            merged = _merge_params(params, kwargs)
            _wire_client_stage_params(
                self._pipeline_spec,
                "extract_params",
                merged,
                method="extract_image_files",
                allowed=_DEFAULT_ALLOWED_EXTRACT_KEYS,
            )
        self._pipeline_spec["extraction_mode"] = "image"
        if split_config is not None:
            self._pipeline_spec["split_config"] = split_config
        self._record_stage("extract")
        return self

    def filter(self) -> "ServiceIngestor":
        """Record a filter stage."""
        self._record_stage("filter")
        return self

    def split(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
        """Record post-extract split / chunking configuration.

        Accepts the same dict shape as :meth:`GraphIngestor.extract`'s
        ``split_config`` keyword (``{"<source_type>": {"max_tokens": …}}``).
        """
        merged: dict[str, Any]
        if isinstance(params, dict):
            merged = dict(params)
        elif params is None:
            merged = {}
        else:
            merged = _params_to_dict(params)
        merged.update(kwargs)
        self._pipeline_spec["split_config"] = merged
        return self

    def pdf_split_config(self, pages_per_chunk: int = 32) -> "ServiceIngestor":
        """Record PDF page-chunking config (per-request).

        The gateway uses this to decide realtime-vs-batch routing
        (chunked docs always go to batch).
        """
        PdfSplitParams.model_validate({})  # cheap sanity touch
        self._pipeline_spec["pdf_split"] = {"pages_per_chunk": int(pages_per_chunk)}
        return self

    # ------------------------------------------------------------------
    # Phase 2: remote sinks — sent via PipelineSpec, gated by SinkUrlAllowlist
    # ------------------------------------------------------------------

    def store(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
        """Record an image-asset store stage targeting a remote URI.

        ``storage_uri`` must be a non-local URI (``s3://``, ``gs://``,
        ``azure://``, …) — the worker pod has no view into the caller's
        filesystem. The server's ``sinks.storage_uri_schemes`` allowlist
        gates which schemes are admissible.
        """
        merged = _merge_params(params, kwargs) if (params or kwargs) else StoreParams()
        params_dict = _params_to_dict(merged)
        uri = params_dict.get("storage_uri")
        if uri is not None:
            _require_remote_uri(uri, "store", "storage_uri")
        # ``storage_uri`` is the legitimate sink destination, so we let
        # it through the local denylist check.
        for k in list(params_dict):
            if k != "storage_uri" and k in _SERVER_OWNED_KEYS:
                raise ValueError(f"ServiceIngestor.store(): key {k!r} is server-owned in " "run_mode='service'.")
        from nemo_retriever.common.policy import _DEFAULT_ALLOWED_STORE_KEYS

        params_dict = _filter_policy_allowed(params_dict, _DEFAULT_ALLOWED_STORE_KEYS)
        _set_stage_params(self._pipeline_spec, "store_params", params_dict)
        self._record_stage("store")
        return self

    def store_embed(self) -> "ServiceIngestor":
        _raise_unsupported(
            "store_embed",
            phase_hint=(
                "By design — service run_mode persists embeddings via "
                ".store(...) / .vdb_upload(...) sinks, not the in-process "
                "store_embed helper. Wire a remote storage_uri instead."
            ),
        )

    def udf(
        self,
        udf_function: str,
        udf_function_name: Optional[str] = None,
        phase: Optional[Union[int, str]] = None,
        target_stage: Optional[str] = None,
        run_before: bool = False,
        run_after: bool = False,
    ) -> "ServiceIngestor":
        _raise_unsupported(
            "udf",
            phase_hint=(
                "Phase 5 deferred to a follow-up. Service run_mode requires "
                "the operator to register Python callables in "
                "retriever-service.yaml under 'udfs:' (clients reference "
                "them by name; arbitrary code never crosses the trust "
                "boundary). Until that ships, run UDFs locally via "
                "run_mode='inprocess'."
            ),
        )

    def vdb_upload(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
        """Record a vector-DB upload sink targeting a remote LanceDB URI.

        ``vdb_kwargs.lancedb_uri`` must be a non-local URI matching the
        server's ``sinks.vdb_uri_schemes`` allowlist.

        Sidecar metadata (``meta_dataframe`` + ``meta_source_field`` +
        ``meta_fields``) is uploaded eagerly via ``POST /v1/ingest/sidecar``
        and the returned id is shipped on the spec as ``meta_dataframe_id``.
        The original ``meta_dataframe`` (path or in-memory DataFrame) is
        never sent on the wire — the worker pod cannot read it.
        """
        merged = _merge_params(params, kwargs) if (params or kwargs) else VdbUploadParams()
        params_dict = _params_to_dict(merged)

        # Resolve sidecar metadata: if the caller supplied a path or an
        # in-memory DataFrame, upload it now and substitute the returned id.
        meta_df = params_dict.pop("meta_dataframe", None)
        meta_source = params_dict.pop("meta_source_field", None)
        meta_fields = params_dict.pop("meta_fields", None)
        meta_join = params_dict.pop("meta_join_key", "auto")
        if meta_df is not None or meta_source is not None or meta_fields is not None:
            if meta_df is None or meta_source is None or not meta_fields:
                raise ValueError(
                    "ServiceIngestor.vdb_upload(): sidecar metadata requires all "
                    "three of meta_dataframe / meta_source_field / meta_fields."
                )
            sidecar_id = self._upload_sidecar(meta_df)
            params_dict["meta_dataframe_id"] = sidecar_id
            params_dict["meta_source_field"] = str(meta_source)
            params_dict["meta_fields"] = [str(x) for x in meta_fields]
            params_dict["meta_join_key"] = meta_join

        vdb_kwargs = params_dict.get("vdb_kwargs") or {}
        if vdb_kwargs:
            uri = vdb_kwargs.get("lancedb_uri") or vdb_kwargs.get("uri")
            if uri is not None:
                _require_remote_uri(uri, "vdb_upload", "vdb_kwargs.lancedb_uri")
        self._pipeline_spec["vdb_upload_params"] = params_dict
        return self

    def _upload_sidecar(self, meta_df: Any) -> str:
        """POST sidecar metadata to ``/v1/ingest/sidecar`` and return the id.

        Accepts a path (string / PathLike) or an in-memory ``pandas.DataFrame``.
        DataFrames are serialised as parquet to keep the payload compact;
        local paths are streamed as their on-disk bytes with content-type
        inferred from the suffix.
        """
        import io
        import json as _json
        import urllib.request

        from pathlib import Path as _Path

        filename: str
        content_type: str
        payload: bytes

        # In-memory DataFrame (or duck-typed pandas-like) → parquet bytes.
        if hasattr(meta_df, "to_parquet"):
            buf = io.BytesIO()
            try:
                meta_df.to_parquet(buf, index=False)
            except Exception as exc:
                raise ValueError(
                    f"ServiceIngestor.vdb_upload(): failed to serialise sidecar " f"DataFrame to parquet: {exc}"
                ) from exc
            payload = buf.getvalue()
            filename = "sidecar.parquet"
            content_type = "application/x-parquet"
        else:
            # Treat as filesystem path.
            path = _Path(str(meta_df))
            if not path.is_file():
                raise FileNotFoundError(f"ServiceIngestor.vdb_upload(): sidecar metadata file not found: {path}")
            payload = path.read_bytes()
            filename = path.name
            suf = path.suffix.lower()
            if suf == ".parquet" or suf == ".pq":
                content_type = "application/x-parquet"
            elif suf in (".json", ".jsonl"):
                content_type = "application/json"
            else:
                content_type = "text/csv"

        # Build a minimal multipart/form-data request — avoids dragging in
        # an httpx dependency where urllib already works.
        import re
        import secrets

        safe_filename = re.sub(r"[^\w.\-]", "_", filename) or "upload"
        boundary = f"----nrlib-sidecar-{secrets.token_hex(16)}"
        body = io.BytesIO()
        body.write(f"--{boundary}\r\n".encode())
        body.write(f'Content-Disposition: form-data; name="file"; filename="{safe_filename}"\r\n'.encode())
        body.write(f"Content-Type: {content_type}\r\n\r\n".encode())
        body.write(payload)
        body.write(f"\r\n--{boundary}--\r\n".encode())

        url = f"{self._base_url}/v1/ingest/sidecar"
        req = urllib.request.Request(url, data=body.getvalue(), method="POST")
        req.add_header("Content-Type", f"multipart/form-data; boundary={boundary}")
        if self._api_token:
            req.add_header("Authorization", f"Bearer {self._api_token}")
        try:
            with urllib.request.urlopen(req, timeout=self._request_timeout_s) as resp:
                body_json = _json.loads(resp.read().decode("utf-8"))
        except Exception as exc:
            raise RuntimeError(f"ServiceIngestor.vdb_upload(): failed to upload sidecar to {url}: {exc}") from exc

        sidecar_id = body_json.get("sidecar_id")
        if not sidecar_id:
            raise RuntimeError(
                f"ServiceIngestor.vdb_upload(): sidecar upload response missing sidecar_id: {body_json!r}"
            )
        logger.debug("Uploaded sidecar %s (%d bytes)", sidecar_id, len(payload))
        return sidecar_id

    def save_intermediate_results(self, output_dir: str) -> "ServiceIngestor":
        _raise_unsupported(
            "save_intermediate_results",
            phase_hint=(
                "By design — service workers don't expose stage-by-stage "
                "outputs; they run the whole pipeline to completion before "
                "returning a result. For per-stage debugging use "
                "run_mode='inprocess'. To capture final outputs use "
                ".save_to_disk(output_directory=...) instead."
            ),
        )

    def save_to_disk(
        self,
        output_directory: Optional[str] = None,
        cleanup: bool = True,
        compression: Optional[str] = "gzip",
    ) -> "ServiceIngestor":
        """Stream per-document results to ``output_directory`` as they finish.

        Each completed document produces one JSON file (or ``.json.gz`` when
        ``compression='gzip'``) named ``<document_id>.json[.gz]`` whose
        contents are the worker's transport-serialized pipeline rows
        (see :mod:`nemo_retriever.ingest_results`) — the same column
        layout as ``GraphIngestor.ingest()`` in local run modes.

        Important differences from graph mode:

        * Large binary columns (``bytes``, ``page_image``, ``images``,
          ``charts``, ``tables``) are stripped server-side before the
          rows leave the worker. Use :meth:`store` to persist image
          assets to a remote URI; the local-disk artifact only carries
          the structured metadata.
        * The client does the writing — the server has no view into the
          caller's filesystem. ``cleanup`` is accepted for API parity
          with graph mode but has no server-side effect today.
        """
        if output_directory is None:
            raise ValueError("ServiceIngestor.save_to_disk(): output_directory is required.")
        if compression not in (None, "gzip"):
            raise ValueError(
                f"save_to_disk(compression={compression!r}): only None or 'gzip' " "are supported in service run_mode."
            )
        target = Path(output_directory)
        target.mkdir(parents=True, exist_ok=True)
        self._save_to_disk_dir = target
        self._save_to_disk_compression = compression
        self._save_to_disk_cleanup = cleanup
        return self

    def caption(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
        """Record a caption stage backed by the server's remote VLM endpoint.

        Behavioural knobs — ``prompt``, ``system_prompt``, ``batch_size``,
        ``context_text_max_chars``, ``caption_infographics``, and generic
        sampling params (``temperature``, ``max_tokens``, ``top_p``,
        ``top_k``) — are honored. Trust-sensitive fields
        (``endpoint_url``, ``api_key``, ``model_name``) and
        local-execution fields (``device``, ``hf_cache_dir``,
        ``tensor_parallel_size``, ``gpu_memory_utilization``) are
        rejected on the client; the operator-configured remote endpoint
        is the only path to a caption NIM.

        We use Pydantic's ``model_fields_set`` to distinguish fields
        the caller *explicitly* set from fields carrying their
        ``CaptionParams`` default — only the former are rejected.
        """
        trust_sensitive = {"endpoint_url", "api_key", "model_name"}
        local_only = {
            "device",
            "hf_cache_dir",
            "tensor_parallel_size",
            "gpu_memory_utilization",
        }

        # Identify which keys the caller actually meant to pass. The
        # signal for kwargs is unambiguous (any key in **kwargs is by
        # definition caller-provided); for a passed-in CaptionParams
        # instance we compare against class defaults, with one wrinkle:
        # ``api_key`` is auto-populated by the model validator from the
        # NVIDIA_API_KEY env var, so we cannot distinguish "caller set
        # this" from "validator set this" — we conservatively strip the
        # value either way and only raise when the caller used kwargs.
        explicit_keys: set[str] = set(kwargs.keys())
        if isinstance(params, CaptionParams):
            class_defaults = {name: field.default for name, field in CaptionParams.model_fields.items()}
            for k in trust_sensitive | local_only:
                if k == "api_key":
                    continue  # see comment above; the env-var auto-fill is ambiguous.
                val = getattr(params, k, None)
                if val is not None and val != class_defaults.get(k):
                    explicit_keys.add(k)

        bad_trust = sorted(explicit_keys & trust_sensitive)
        if bad_trust:
            raise ValueError(
                f"ServiceIngestor.caption(): keys {bad_trust!r} are server-owned in "
                "run_mode='service'. The operator configures the caption "
                "endpoint via retriever-service.yaml (nim_endpoints.caption_invoke_url)."
            )
        bad_local = sorted(explicit_keys & local_only)
        if bad_local:
            raise ValueError(
                f"ServiceIngestor.caption(): keys {bad_local!r} configure local "
                "in-process GPU execution and have no effect against a remote "
                "caption endpoint. Remove them and rely on the server-owned "
                "model / endpoint."
            )

        merged = _merge_params(params, kwargs) if (params or kwargs) else CaptionParams()
        params_dict = _params_to_dict(merged)
        # Drop both classes of keys before the spec leaves the client —
        # the server's allowlist rejects them anyway, but failing fast
        # at the boundary keeps the network message small and the policy
        # error rare in practice.
        scrubbed = {k: v for k, v in params_dict.items() if k not in trust_sensitive | local_only}
        self._pipeline_spec["caption_params"] = scrubbed
        self._record_stage("caption")
        return self

    def webhook(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
        """Record a webhook-notification stage targeting a remote URL.

        ``endpoint_url`` must match one of the server's
        ``sinks.webhook_url_prefixes``. Without that allowlist the
        service rejects webhook requests entirely so worker egress
        cannot be steered by clients.
        """
        merged = _merge_params(params, kwargs) if (params or kwargs) else WebhookParams()
        params_dict = _params_to_dict(merged)
        endpoint = params_dict.get("endpoint_url")
        if endpoint is None:
            raise ValueError(
                "ServiceIngestor.webhook(): endpoint_url is required "
                "(unlike inprocess run_mode, an empty endpoint_url is treated "
                "as misconfiguration in service mode)."
            )
        if not endpoint.startswith(("http://", "https://")):
            raise ValueError(
                "ServiceIngestor.webhook(): endpoint_url must be a fully-qualified "
                f"http(s):// URL; got {endpoint!r}."
            )
        self._pipeline_spec["webhook_params"] = params_dict
        self._record_stage("webhook")
        return self

    def _ingest_events_with_result_client(
        self,
        *,
        retain_results: bool,
        result_schema: ResultSchema,
        return_embeddings: bool,
        return_images: bool,
    ) -> Iterator[tuple[dict[str, Any], httpx.Client | None]]:
        """Yield ingest events while owning the optional shared result client."""
        client_context = self._new_result_fetch_client() if retain_results else nullcontext(None)
        with client_context as result_client:
            for evt in self.ingest_stream(
                retain_results=retain_results,
                result_schema=result_schema,
                return_embeddings=return_embeddings,
                return_images=return_images,
            ):
                yield evt, result_client

    # ------------------------------------------------------------------
    # Execution — sync materialized
    # ------------------------------------------------------------------

    def ingest(self, params: Any = None, **kwargs: Any) -> Any:
        """Block until every document has finished processing on the server.

        Internally opens exactly one server-side job aggregate for the
        full input set (sized to ``len(documents)``). The aggregate
        ``job_id`` is captured from the first ``job_created`` event and
        exposed on :class:`ServiceIngestResult` so the caller can call
        ``GET /v1/ingest/job/{job_id}`` for follow-up status.

        Parameters
        ----------
        params
            Optional :class:`IngestExecuteParams` (or plain ``dict``)
            carrying execute-time flags.  In service run_mode only
            ``return_failures`` / ``return_traces`` / ``return_results`` /
            ``result_schema`` / ``return_embeddings`` / ``return_images``
            are honored — every other field is recorded on the server-side
            pipeline spec.
        **kwargs
            Same execute-time flags may be passed individually.  Anything
            not recognised is silently ignored (server-side execution
            in service mode is driven by the pipeline spec, not by
            execute-time knobs).
        return_results
            When ``True`` (default), fetch each completed document's
            ``result_data`` from ``GET /v1/ingest/status/{id}`` and
            expose the combined rows on ``result.dataframe`` as a
            ``pandas.DataFrame``. Set to ``False`` to skip those HTTP
            round-trips when only job metadata is needed.
        result_schema
            ``"legacy"`` (default) preserves the existing service
            DataFrame column layout with bulky values stripped and emits
            a deprecation warning when result rows are retained.
            ``"compact"`` opts into the future compact row schema.
        return_embeddings, return_images
            When using legacy result rows, include embedding vectors and
            raw image payloads instead of stripping them from transport
            cells. Defaults remain ``False`` to avoid large responses.

        Returns
        -------
        ServiceIngestResult
            When neither ``return_failures`` nor ``return_traces`` is
            set — a list subclass of per-document completion events with
            extra ``job_id`` / ``failures`` / ``document_ids`` /
            ``elapsed_s`` / ``job_status`` / ``dataframe`` attributes.
        tuple
            With ``return_failures=True`` only — ``(result, failures)``.
            With ``return_traces=True`` only — ``(result, traces)``.
            With both — ``(result, failures, traces)``.  ``failures``
            mirrors ``result.failures``; ``traces`` is the ordered list
            of raw SSE event dicts observed during the run, useful for
            debugging pipeline behaviour without re-running the job.
        """
        return_failures, return_traces, return_results, result_schema, return_embeddings, return_images = (
            self._resolve_execute_flags(params, kwargs)
        )
        del params, kwargs
        self._validate_input_sources(self._inline_texts)
        if not self._documents and not self._buffers and is_blank_inline_corpus(self._inline_texts):
            self._document_ids.clear()
            self._last_run_elapsed_s = 0.0
            self._last_job_id = None
            result = ServiceIngestResult()
            if return_results:
                result.dataframe = _empty_service_result_dataframe(result_schema)
            if return_failures and return_traces:
                return result, [], []
            if return_failures or return_traces:
                return result, []
            return result

        retain_results = return_results or self._save_to_disk_dir is not None
        if retain_results and result_schema == "legacy":
            warnings.warn(_LEGACY_RESULT_SCHEMA_DEPRECATION, DeprecationWarning, stacklevel=2)
        result = ServiceIngestResult()
        traces: list[dict[str, Any]] = []
        rows_by_document: dict[str, list[dict[str, Any]]] = {}
        t0 = time.monotonic()

        documents_completed = 0
        documents_failed = 0
        total_uploaded = 0

        for evt, result_client in self._ingest_events_with_result_client(
            retain_results=retain_results,
            result_schema=result_schema,
            return_embeddings=return_embeddings,
            return_images=return_images,
        ):
            if return_traces:
                traces.append(evt)
            event_type = evt.get("event")

            if event_type == "job_created":
                result.job_id = evt.get("job_id") or result.job_id
                result.trace_id = evt.get("trace_id") or result.trace_id
                continue

            if event_type in ("job_finalized", "job_partial", "job_failed"):
                if event_type == "job_finalized":
                    result.job_status = "completed"
                elif event_type == "job_partial":
                    result.job_status = "partial_success"
                else:
                    result.job_status = "failed"
                continue

            if event_type == "job_progress" or event_type == "job_started":
                continue

            if event_type == "upload_complete":
                total_uploaded += 1
                document_id = evt.get("document_id")
                filename = evt.get("filename")
                if document_id and filename:
                    result.document_filenames[str(document_id)] = str(filename)
                if result.job_id is None:
                    # Race: SSE delivered an upload_complete before the
                    # generator yielded job_created. Fall back to the
                    # job_id stamped on the per-doc event by the client.
                    result.job_id = evt.get("job_id") or result.job_id
                print(
                    f"\r  Job {result.job_id or '?'}  |  "
                    f"Uploaded: {total_uploaded}  |  "
                    f"Completed: {documents_completed}  |  "
                    f"Failed: {documents_failed}",
                    end="",
                    flush=True,
                )

            elif event_type == "document_complete":
                status = evt.get("status", "completed")
                if status not in ("completed", "failed"):
                    continue
                if status == "failed":
                    documents_failed += 1
                    error = evt.get("error", "unknown error")
                    doc_id = evt.get("document_id", "?")
                    result.failures.append((doc_id, error))
                else:
                    documents_completed += 1
                    doc_id = evt.get("document_id", "")
                    if return_results or self._save_to_disk_dir is not None:
                        try:
                            rows = self._materialize_completed_document(
                                doc_id,
                                return_results=return_results,
                                client=result_client,
                            )
                            if rows is not None and return_results:
                                rows_by_document[doc_id] = rows
                        except Exception as exc:
                            label = "return_results" if return_results else "save_to_disk"
                            logger.warning("%s: failed to fetch/persist %s: %s", label, doc_id, exc)
                            result.failures.append((doc_id, f"{label}: {exc}"))
                result.append(evt)
                print(
                    f"\r  Job {result.job_id or '?'}  |  "
                    f"Uploaded: {total_uploaded}  |  "
                    f"Completed: {documents_completed}  |  "
                    f"Failed: {documents_failed}",
                    end="",
                    flush=True,
                )

            elif event_type == "upload_failed":
                fname = evt.get("filename", "?")
                error = evt.get("error", "unknown")
                result.failures.append((fname, f"upload failed: {error}"))

        if total_uploaded > 0:
            print()

        result.document_ids = list(self._document_ids)
        result.elapsed_s = time.monotonic() - t0
        if return_results:
            doc_order = [d for d in self._document_ids if d in rows_by_document] or list(rows_by_document)
            result.dataframe = concat_ingest_results(rows_by_document, doc_order)
        self._last_run_elapsed_s = result.elapsed_s
        # Cache the job_id on the ingestor for the get_status() /
        # remaining_jobs() accessors so they can target the job
        # aggregate endpoints once J6 wiring is opted in (kept
        # backwards compatible — get_status() still uses document_ids).
        self._last_job_id = result.job_id

        if return_failures and return_traces:
            return result, list(result.failures), traces
        if return_failures:
            return result, list(result.failures)
        if return_traces:
            return result, traces
        return result

    @staticmethod
    def _normalize_result_schema(value: Any) -> ResultSchema:
        schema = str(value or "legacy").strip().lower()
        if schema not in ("legacy", "compact"):
            raise ValueError("result_schema must be 'legacy' or 'compact'")
        return schema  # type: ignore[return-value]

    @classmethod
    def _resolve_execute_flags(
        cls, params: Any, kwargs: dict[str, Any]
    ) -> tuple[bool, bool, bool, ResultSchema, bool, bool]:
        """Read execute-time flags from ``params`` and/or ``kwargs``.

        kwargs take precedence over fields on ``params`` when both supply
        the same flag, mirroring the precedence used by
        :func:`nemo_retriever.ingestor._merge_params`.
        """

        def _from_params(name: str, *, default: bool) -> bool:
            if isinstance(params, IngestExecuteParams):
                return bool(getattr(params, name, default))
            if isinstance(params, dict):
                if name in params:
                    return bool(params[name])
                return default
            return default

        def _from_params_value(name: str, *, default: Any) -> Any:
            if isinstance(params, IngestExecuteParams):
                return getattr(params, name, default)
            if isinstance(params, dict):
                return params.get(name, default)
            return default

        return_failures = (
            bool(kwargs["return_failures"])
            if "return_failures" in kwargs
            else _from_params("return_failures", default=False)
        )
        return_traces = (
            bool(kwargs["return_traces"]) if "return_traces" in kwargs else _from_params("return_traces", default=False)
        )
        return_results = (
            bool(kwargs["return_results"])
            if "return_results" in kwargs
            else _from_params("return_results", default=True)
        )
        result_schema = cls._normalize_result_schema(
            kwargs["result_schema"]
            if "result_schema" in kwargs
            else _from_params_value("result_schema", default="legacy")
        )
        return_embeddings = (
            bool(kwargs["return_embeddings"])
            if "return_embeddings" in kwargs
            else _from_params("return_embeddings", default=False)
        )
        return_images = (
            bool(kwargs["return_images"]) if "return_images" in kwargs else _from_params("return_images", default=False)
        )
        return return_failures, return_traces, return_results, result_schema, return_embeddings, return_images

    # ------------------------------------------------------------------
    # Execution — sync streaming
    # ------------------------------------------------------------------

    def ingest_stream(
        self,
        *,
        retain_results: bool = False,
        result_schema: ResultSchema = "legacy",
        return_embeddings: bool = False,
        return_images: bool = False,
    ) -> Iterator[dict[str, Any]]:
        """Sync generator yielding events as documents are processed.

        Yields dicts with:

        * ``{"event": "job_created", "job_id": ..., "expected_documents": ...}``
        * ``{"event": "upload_complete", "filename": ..., "document_id": ..., "job_id": ...}``
        * ``{"event": "document_complete", "document_id": ..., "status": ..., "job_id": ..., ...}``
        * ``{"event": "upload_failed", "filename": ..., "error": ..., "job_id": ...}``
        * ``{"event": "job_progress", "job_id": ..., "completed": ..., "failed": ..., ...}``
        * ``{"event": "job_finalized"|"job_partial"|"job_failed", "job_id": ..., ...}``
        """
        result_schema = self._normalize_result_schema(result_schema)
        return self._ingest_stream_with_retain(
            retain_results,
            result_schema=result_schema,
            return_embeddings=return_embeddings,
            return_images=return_images,
        )

    # ------------------------------------------------------------------
    # Execution — async streaming
    # ------------------------------------------------------------------

    async def aingest_stream(
        self,
        *,
        retain_results: bool = False,
        result_schema: ResultSchema = "legacy",
        return_embeddings: bool = False,
        return_images: bool = False,
    ) -> AsyncIterator[dict[str, Any]]:
        """Async generator yielding events as documents are processed."""
        result_schema = self._normalize_result_schema(result_schema)
        files = self._collect_inputs()
        if not files:
            return

        self._document_ids.clear()
        async for evt in self._aingest_stream_impl(
            files,
            retain_results=retain_results,
            result_schema=result_schema,
            return_embeddings=return_embeddings,
            return_images=return_images,
        ):
            if evt.get("event") == "upload_complete":
                did = evt.get("document_id")
                if did:
                    self._document_ids.append(did)
            yield evt

    # ------------------------------------------------------------------
    # Async helper used by both sync and async streaming entry points
    # ------------------------------------------------------------------

    def _ingest_stream_with_retain(
        self,
        retain_results: bool,
        *,
        result_schema: ResultSchema = "legacy",
        return_embeddings: bool = False,
        return_images: bool = False,
    ) -> Iterator[dict[str, Any]]:
        """Like :meth:`ingest_stream` but passes server-side retention to the HTTP client."""
        files = self._collect_inputs()
        if not files:
            return iter(())

        self._document_ids.clear()

        def _record_doc_id(evt: dict[str, Any]) -> None:
            if evt.get("event") == "upload_complete":
                did = evt.get("document_id")
                if did:
                    self._document_ids.append(did)

        def _factory():
            return self._wrap_for_capture(
                self._aingest_stream_impl(
                    files,
                    retain_results=retain_results,
                    result_schema=result_schema,
                    return_embeddings=return_embeddings,
                    return_images=return_images,
                ),
                _record_doc_id,
            )

        bridge = _AsyncToSyncBridge(_factory)
        return iter(bridge)

    async def _aingest_stream_impl(
        self,
        files: list[UploadInput],
        *,
        retain_results: bool = False,
        result_schema: ResultSchema = "legacy",
        return_embeddings: bool = False,
        return_images: bool = False,
    ) -> AsyncIterator[dict[str, Any]]:
        client = RetrieverServiceClient(
            base_url=self._base_url,
            max_concurrency=self._max_concurrency,
            api_token=self._api_token,
        )
        pipeline_payload = self._pipeline_payload(
            result_schema=result_schema,
            return_embeddings=return_embeddings,
            return_images=return_images,
        )
        async for evt in client.aingest_documents_stream(
            files=files,
            pipeline_spec=pipeline_payload,
            retain_results=retain_results,
        ):
            yield evt

    @staticmethod
    async def _wrap_for_capture(
        agen: AsyncIterator[dict[str, Any]],
        on_event,
    ) -> AsyncIterator[dict[str, Any]]:
        """Pass-through wrapper that lets the sync bridge capture document_ids."""
        async for evt in agen:
            on_event(evt)
            yield evt

    # ------------------------------------------------------------------
    # Async-future API
    # ------------------------------------------------------------------

    def ingest_async(
        self,
        *,
        return_failures: bool = False,
        return_traces: bool = False,
        return_results: bool = True,
        result_schema: ResultSchema = "legacy",
        return_embeddings: bool = False,
        return_images: bool = False,
    ) -> Any:
        """Run :meth:`ingest` on a background thread; return a ``Future``.

        The flags are forwarded to :meth:`ingest`, so calling
        ``future.result()`` produces the same tuple/list shape that a
        direct synchronous call with the same flags would return.
        """
        from concurrent.futures import ThreadPoolExecutor

        executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="ServiceIngestorAsync")
        return executor.submit(
            self.ingest,
            return_failures=return_failures,
            return_traces=return_traces,
            return_results=return_results,
            result_schema=result_schema,
            return_embeddings=return_embeddings,
            return_images=return_images,
        )

    # ------------------------------------------------------------------
    # Status & document-counter accessors
    # ------------------------------------------------------------------

    def get_status(self) -> dict[str, str]:
        """Return ``{document_id: status}`` for every document submitted so far."""
        if not self._document_ids:
            return {}
        url = f"{self._base_url}/v1/ingest/status/batch"
        with httpx.Client(timeout=30.0, headers=self._auth_headers) as client:
            try:
                resp = client.post(url, json={"ids": self._document_ids})
                resp.raise_for_status()
                items = resp.json().get("items", {})
                return {did: info.get("status", "unknown") for did, info in items.items()}
            except Exception as exc:
                logger.warning("Could not fetch bulk status: %s", exc)
                return {did: "unknown" for did in self._document_ids}

    def completed_jobs(self) -> int:
        return sum(1 for s in self.get_status().values() if s == "completed")

    def failed_jobs(self) -> int:
        return sum(1 for s in self.get_status().values() if s == "failed")

    def cancelled_jobs(self) -> int:
        return 0

    def remaining_jobs(self) -> int:
        return sum(1 for s in self.get_status().values() if s in ("processing", "unknown"))

    # ------------------------------------------------------------------
    # Cancel — not supported (no server endpoint)
    # ------------------------------------------------------------------

    def cancel(self, job_id: str | None = None) -> dict[str, Any]:
        """Not supported — the server does not expose a cancel endpoint."""
        raise NotImplementedError(
            "Cancel is not supported in service mode. " "The server does not currently expose a cancel endpoint."
        )

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _has_mixed_inline_sources(self) -> bool:
        return bool(self._inline_texts) and bool(self._documents or self._buffers)

    def _collect_inputs(self) -> list[UploadInput]:
        """Gather filesystem and in-memory inputs for the service client."""
        self._validate_input_sources(self._inline_texts)
        if not self._documents and not self._buffers and is_blank_inline_corpus(self._inline_texts):
            return []

        files: list[UploadInput] = [Path(p) for p in self._documents]

        if self._buffers:
            import tempfile

            tmp_dir = Path(tempfile.mkdtemp(prefix="service_ingestor_"))
            for name, buf in self._buffers:
                target = tmp_dir / name
                target.write_bytes(buf.getvalue())
                files.append(target)

        for index, text in enumerate(self._inline_texts or []):
            source_id = inline_text_source_id(index)
            files.append(
                InMemoryUpload(
                    filename=source_id,
                    content=text.encode("utf-8"),
                    content_type="text/plain; charset=utf-8",
                    classification_filename=f"inline-{index:08d}.txt",
                )
            )

        return files

RUN_MODE = 'service' class-attribute instance-attribute

files(documents)

Source code in nemo_retriever/service/service_ingestor.py

def files(self, documents: Union[str, List[str]]) -> "ServiceIngestor":
    """Add document paths/URIs for processing."""
    if isinstance(documents, str):
        self._documents.append(documents)
    else:
        self._documents.extend(documents)
    return self

texts(texts)

Source code in nemo_retriever/service/service_ingestor.py

def texts(self, texts: Union[str, Sequence[str]]) -> Self:
    """Set raw inline text documents, optionally alongside file or buffer uploads."""
    self._inline_texts = normalize_inline_texts(texts)
    return self

buffers(buffers)

Source code in nemo_retriever/service/service_ingestor.py

def buffers(
    self,
    buffers: Union[Tuple[str, BytesIO], List[Tuple[str, BytesIO]]],
) -> "ServiceIngestor":
    """Add in-memory buffers for processing.

    Each buffer must be ``(filename, BytesIO)`` so the server can record
    a meaningful source filename.
    """
    if isinstance(buffers, tuple):
        buffers = [buffers]
    for name, buf in buffers:
        self._buffers.append((name, buf))
    return self

load()

Source code in nemo_retriever/service/service_ingestor.py

def load(self) -> "ServiceIngestor":
    """No-op for service mode."""
    return self

all_tasks()

Source code in nemo_retriever/service/service_ingestor.py

def all_tasks(self) -> "ServiceIngestor":
    """Record the default chain: extract → dedup → embed.

    Concrete params come from server config; ``all_tasks()`` only
    controls *stage order* and is the closest in-process equivalent
    of "run everything the server is configured to do".
    """
    self._record_stage("extract")
    self._record_stage("dedup")
    self._record_stage("embed")
    return self

dedup(params=None, **kwargs)

Source code in nemo_retriever/service/service_ingestor.py

def dedup(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
    """Record a dedup stage with optional :class:`DedupParams` overrides."""
    if params is not None or kwargs:
        from nemo_retriever.common.policy import _DEFAULT_ALLOWED_DEDUP_KEYS

        merged = _merge_params(params, kwargs)
        _wire_client_stage_params(
            self._pipeline_spec,
            "dedup_params",
            merged,
            method="dedup",
            allowed=_DEFAULT_ALLOWED_DEDUP_KEYS,
        )
    self._record_stage("dedup")
    return self

embed(params=None, **kwargs)

Source code in nemo_retriever/service/service_ingestor.py

def embed(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
    """Record an embed stage with optional :class:`EmbedParams` overrides.

    Embedding endpoint URL and API key are server-owned and will be
    rejected if set here.
    """
    if params is not None or kwargs:
        from nemo_retriever.common.policy import _DEFAULT_ALLOWED_EMBED_KEYS

        merged = _merge_params(params, kwargs)
        _wire_client_stage_params(
            self._pipeline_spec,
            "embed_params",
            merged,
            method="embed",
            allowed=_DEFAULT_ALLOWED_EMBED_KEYS,
        )
    self._record_stage("embed")
    return self

extract(params=None, *, split_config=None, extraction_mode='auto', **kwargs)

Source code in nemo_retriever/service/service_ingestor.py

def extract(
    self,
    params: Any = None,
    *,
    split_config: Optional[dict[str, Any]] = None,
    extraction_mode: str = "auto",
    **kwargs: Any,
) -> "ServiceIngestor":
    """Record a generic extraction stage.

    ``extraction_mode`` selects the worker's extraction path
    (``'auto'`` default — dispatches by file extension; ``'pdf'``
    forces the PDF path for all inputs, etc.).

    When no ``ExtractParams`` overrides are supplied, ``extract_params``
    is omitted from the wire payload so the worker applies the
    service's server-owned defaults (and the allow-list is not tripped
    by client-side model defaults).
    """
    if params is not None or kwargs:
        from nemo_retriever.common.policy import _DEFAULT_ALLOWED_EXTRACT_KEYS

        merged = _merge_params(params, kwargs)
        _wire_client_stage_params(
            self._pipeline_spec,
            "extract_params",
            merged,
            method="extract",
            allowed=_DEFAULT_ALLOWED_EXTRACT_KEYS,
        )
    self._pipeline_spec["extraction_mode"] = extraction_mode
    if split_config is not None:
        self._pipeline_spec["split_config"] = split_config
    self._record_stage("extract")
    return self

extract_image_files(params=None, *, split_config=None, **kwargs)

Source code in nemo_retriever/service/service_ingestor.py

def extract_image_files(
    self, params: Any = None, *, split_config: Optional[dict[str, Any]] = None, **kwargs: Any
) -> "ServiceIngestor":
    """Record image-file extraction (``extraction_mode='image'``)."""
    if params is not None or kwargs:
        from nemo_retriever.common.policy import _DEFAULT_ALLOWED_EXTRACT_KEYS

        merged = _merge_params(params, kwargs)
        _wire_client_stage_params(
            self._pipeline_spec,
            "extract_params",
            merged,
            method="extract_image_files",
            allowed=_DEFAULT_ALLOWED_EXTRACT_KEYS,
        )
    self._pipeline_spec["extraction_mode"] = "image"
    if split_config is not None:
        self._pipeline_spec["split_config"] = split_config
    self._record_stage("extract")
    return self

filter()

Source code in nemo_retriever/service/service_ingestor.py

def filter(self) -> "ServiceIngestor":
    """Record a filter stage."""
    self._record_stage("filter")
    return self

split(params=None, **kwargs)

Source code in nemo_retriever/service/service_ingestor.py

def split(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
    """Record post-extract split / chunking configuration.

    Accepts the same dict shape as :meth:`GraphIngestor.extract`'s
    ``split_config`` keyword (``{"<source_type>": {"max_tokens": …}}``).
    """
    merged: dict[str, Any]
    if isinstance(params, dict):
        merged = dict(params)
    elif params is None:
        merged = {}
    else:
        merged = _params_to_dict(params)
    merged.update(kwargs)
    self._pipeline_spec["split_config"] = merged
    return self

pdf_split_config(pages_per_chunk=32)

Source code in nemo_retriever/service/service_ingestor.py

def pdf_split_config(self, pages_per_chunk: int = 32) -> "ServiceIngestor":
    """Record PDF page-chunking config (per-request).

    The gateway uses this to decide realtime-vs-batch routing
    (chunked docs always go to batch).
    """
    PdfSplitParams.model_validate({})  # cheap sanity touch
    self._pipeline_spec["pdf_split"] = {"pages_per_chunk": int(pages_per_chunk)}
    return self

store(params=None, **kwargs)

Source code in nemo_retriever/service/service_ingestor.py

def store(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
    """Record an image-asset store stage targeting a remote URI.

    ``storage_uri`` must be a non-local URI (``s3://``, ``gs://``,
    ``azure://``, …) — the worker pod has no view into the caller's
    filesystem. The server's ``sinks.storage_uri_schemes`` allowlist
    gates which schemes are admissible.
    """
    merged = _merge_params(params, kwargs) if (params or kwargs) else StoreParams()
    params_dict = _params_to_dict(merged)
    uri = params_dict.get("storage_uri")
    if uri is not None:
        _require_remote_uri(uri, "store", "storage_uri")
    # ``storage_uri`` is the legitimate sink destination, so we let
    # it through the local denylist check.
    for k in list(params_dict):
        if k != "storage_uri" and k in _SERVER_OWNED_KEYS:
            raise ValueError(f"ServiceIngestor.store(): key {k!r} is server-owned in " "run_mode='service'.")
    from nemo_retriever.common.policy import _DEFAULT_ALLOWED_STORE_KEYS

    params_dict = _filter_policy_allowed(params_dict, _DEFAULT_ALLOWED_STORE_KEYS)
    _set_stage_params(self._pipeline_spec, "store_params", params_dict)
    self._record_stage("store")
    return self

store_embed()

Source code in nemo_retriever/service/service_ingestor.py

def store_embed(self) -> "ServiceIngestor":
    _raise_unsupported(
        "store_embed",
        phase_hint=(
            "By design — service run_mode persists embeddings via "
            ".store(...) / .vdb_upload(...) sinks, not the in-process "
            "store_embed helper. Wire a remote storage_uri instead."
        ),
    )

udf(udf_function, udf_function_name=None, phase=None, target_stage=None, run_before=False, run_after=False)

Source code in nemo_retriever/service/service_ingestor.py

def udf(
    self,
    udf_function: str,
    udf_function_name: Optional[str] = None,
    phase: Optional[Union[int, str]] = None,
    target_stage: Optional[str] = None,
    run_before: bool = False,
    run_after: bool = False,
) -> "ServiceIngestor":
    _raise_unsupported(
        "udf",
        phase_hint=(
            "Phase 5 deferred to a follow-up. Service run_mode requires "
            "the operator to register Python callables in "
            "retriever-service.yaml under 'udfs:' (clients reference "
            "them by name; arbitrary code never crosses the trust "
            "boundary). Until that ships, run UDFs locally via "
            "run_mode='inprocess'."
        ),
    )

vdb_upload(params=None, **kwargs)

Source code in nemo_retriever/service/service_ingestor.py

def vdb_upload(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
    """Record a vector-DB upload sink targeting a remote LanceDB URI.

    ``vdb_kwargs.lancedb_uri`` must be a non-local URI matching the
    server's ``sinks.vdb_uri_schemes`` allowlist.

    Sidecar metadata (``meta_dataframe`` + ``meta_source_field`` +
    ``meta_fields``) is uploaded eagerly via ``POST /v1/ingest/sidecar``
    and the returned id is shipped on the spec as ``meta_dataframe_id``.
    The original ``meta_dataframe`` (path or in-memory DataFrame) is
    never sent on the wire — the worker pod cannot read it.
    """
    merged = _merge_params(params, kwargs) if (params or kwargs) else VdbUploadParams()
    params_dict = _params_to_dict(merged)

    # Resolve sidecar metadata: if the caller supplied a path or an
    # in-memory DataFrame, upload it now and substitute the returned id.
    meta_df = params_dict.pop("meta_dataframe", None)
    meta_source = params_dict.pop("meta_source_field", None)
    meta_fields = params_dict.pop("meta_fields", None)
    meta_join = params_dict.pop("meta_join_key", "auto")
    if meta_df is not None or meta_source is not None or meta_fields is not None:
        if meta_df is None or meta_source is None or not meta_fields:
            raise ValueError(
                "ServiceIngestor.vdb_upload(): sidecar metadata requires all "
                "three of meta_dataframe / meta_source_field / meta_fields."
            )
        sidecar_id = self._upload_sidecar(meta_df)
        params_dict["meta_dataframe_id"] = sidecar_id
        params_dict["meta_source_field"] = str(meta_source)
        params_dict["meta_fields"] = [str(x) for x in meta_fields]
        params_dict["meta_join_key"] = meta_join

    vdb_kwargs = params_dict.get("vdb_kwargs") or {}
    if vdb_kwargs:
        uri = vdb_kwargs.get("lancedb_uri") or vdb_kwargs.get("uri")
        if uri is not None:
            _require_remote_uri(uri, "vdb_upload", "vdb_kwargs.lancedb_uri")
    self._pipeline_spec["vdb_upload_params"] = params_dict
    return self

save_intermediate_results(output_dir)

Source code in nemo_retriever/service/service_ingestor.py

def save_intermediate_results(self, output_dir: str) -> "ServiceIngestor":
    _raise_unsupported(
        "save_intermediate_results",
        phase_hint=(
            "By design — service workers don't expose stage-by-stage "
            "outputs; they run the whole pipeline to completion before "
            "returning a result. For per-stage debugging use "
            "run_mode='inprocess'. To capture final outputs use "
            ".save_to_disk(output_directory=...) instead."
        ),
    )

save_to_disk(output_directory=None, cleanup=True, compression='gzip')

Source code in nemo_retriever/service/service_ingestor.py

def save_to_disk(
    self,
    output_directory: Optional[str] = None,
    cleanup: bool = True,
    compression: Optional[str] = "gzip",
) -> "ServiceIngestor":
    """Stream per-document results to ``output_directory`` as they finish.

    Each completed document produces one JSON file (or ``.json.gz`` when
    ``compression='gzip'``) named ``<document_id>.json[.gz]`` whose
    contents are the worker's transport-serialized pipeline rows
    (see :mod:`nemo_retriever.ingest_results`) — the same column
    layout as ``GraphIngestor.ingest()`` in local run modes.

    Important differences from graph mode:

    * Large binary columns (``bytes``, ``page_image``, ``images``,
      ``charts``, ``tables``) are stripped server-side before the
      rows leave the worker. Use :meth:`store` to persist image
      assets to a remote URI; the local-disk artifact only carries
      the structured metadata.
    * The client does the writing — the server has no view into the
      caller's filesystem. ``cleanup`` is accepted for API parity
      with graph mode but has no server-side effect today.
    """
    if output_directory is None:
        raise ValueError("ServiceIngestor.save_to_disk(): output_directory is required.")
    if compression not in (None, "gzip"):
        raise ValueError(
            f"save_to_disk(compression={compression!r}): only None or 'gzip' " "are supported in service run_mode."
        )
    target = Path(output_directory)
    target.mkdir(parents=True, exist_ok=True)
    self._save_to_disk_dir = target
    self._save_to_disk_compression = compression
    self._save_to_disk_cleanup = cleanup
    return self

caption(params=None, **kwargs)

Source code in nemo_retriever/service/service_ingestor.py

def caption(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
    """Record a caption stage backed by the server's remote VLM endpoint.

    Behavioural knobs — ``prompt``, ``system_prompt``, ``batch_size``,
    ``context_text_max_chars``, ``caption_infographics``, and generic
    sampling params (``temperature``, ``max_tokens``, ``top_p``,
    ``top_k``) — are honored. Trust-sensitive fields
    (``endpoint_url``, ``api_key``, ``model_name``) and
    local-execution fields (``device``, ``hf_cache_dir``,
    ``tensor_parallel_size``, ``gpu_memory_utilization``) are
    rejected on the client; the operator-configured remote endpoint
    is the only path to a caption NIM.

    We use Pydantic's ``model_fields_set`` to distinguish fields
    the caller *explicitly* set from fields carrying their
    ``CaptionParams`` default — only the former are rejected.
    """
    trust_sensitive = {"endpoint_url", "api_key", "model_name"}
    local_only = {
        "device",
        "hf_cache_dir",
        "tensor_parallel_size",
        "gpu_memory_utilization",
    }

    # Identify which keys the caller actually meant to pass. The
    # signal for kwargs is unambiguous (any key in **kwargs is by
    # definition caller-provided); for a passed-in CaptionParams
    # instance we compare against class defaults, with one wrinkle:
    # ``api_key`` is auto-populated by the model validator from the
    # NVIDIA_API_KEY env var, so we cannot distinguish "caller set
    # this" from "validator set this" — we conservatively strip the
    # value either way and only raise when the caller used kwargs.
    explicit_keys: set[str] = set(kwargs.keys())
    if isinstance(params, CaptionParams):
        class_defaults = {name: field.default for name, field in CaptionParams.model_fields.items()}
        for k in trust_sensitive | local_only:
            if k == "api_key":
                continue  # see comment above; the env-var auto-fill is ambiguous.
            val = getattr(params, k, None)
            if val is not None and val != class_defaults.get(k):
                explicit_keys.add(k)

    bad_trust = sorted(explicit_keys & trust_sensitive)
    if bad_trust:
        raise ValueError(
            f"ServiceIngestor.caption(): keys {bad_trust!r} are server-owned in "
            "run_mode='service'. The operator configures the caption "
            "endpoint via retriever-service.yaml (nim_endpoints.caption_invoke_url)."
        )
    bad_local = sorted(explicit_keys & local_only)
    if bad_local:
        raise ValueError(
            f"ServiceIngestor.caption(): keys {bad_local!r} configure local "
            "in-process GPU execution and have no effect against a remote "
            "caption endpoint. Remove them and rely on the server-owned "
            "model / endpoint."
        )

    merged = _merge_params(params, kwargs) if (params or kwargs) else CaptionParams()
    params_dict = _params_to_dict(merged)
    # Drop both classes of keys before the spec leaves the client —
    # the server's allowlist rejects them anyway, but failing fast
    # at the boundary keeps the network message small and the policy
    # error rare in practice.
    scrubbed = {k: v for k, v in params_dict.items() if k not in trust_sensitive | local_only}
    self._pipeline_spec["caption_params"] = scrubbed
    self._record_stage("caption")
    return self

webhook(params=None, **kwargs)

Source code in nemo_retriever/service/service_ingestor.py

def webhook(self, params: Any = None, **kwargs: Any) -> "ServiceIngestor":
    """Record a webhook-notification stage targeting a remote URL.

    ``endpoint_url`` must match one of the server's
    ``sinks.webhook_url_prefixes``. Without that allowlist the
    service rejects webhook requests entirely so worker egress
    cannot be steered by clients.
    """
    merged = _merge_params(params, kwargs) if (params or kwargs) else WebhookParams()
    params_dict = _params_to_dict(merged)
    endpoint = params_dict.get("endpoint_url")
    if endpoint is None:
        raise ValueError(
            "ServiceIngestor.webhook(): endpoint_url is required "
            "(unlike inprocess run_mode, an empty endpoint_url is treated "
            "as misconfiguration in service mode)."
        )
    if not endpoint.startswith(("http://", "https://")):
        raise ValueError(
            "ServiceIngestor.webhook(): endpoint_url must be a fully-qualified "
            f"http(s):// URL; got {endpoint!r}."
        )
    self._pipeline_spec["webhook_params"] = params_dict
    self._record_stage("webhook")
    return self

ingest(params=None, **kwargs)

Source code in nemo_retriever/service/service_ingestor.py

def ingest(self, params: Any = None, **kwargs: Any) -> Any:
    """Block until every document has finished processing on the server.

    Internally opens exactly one server-side job aggregate for the
    full input set (sized to ``len(documents)``). The aggregate
    ``job_id`` is captured from the first ``job_created`` event and
    exposed on :class:`ServiceIngestResult` so the caller can call
    ``GET /v1/ingest/job/{job_id}`` for follow-up status.

    Parameters
    ----------
    params
        Optional :class:`IngestExecuteParams` (or plain ``dict``)
        carrying execute-time flags.  In service run_mode only
        ``return_failures`` / ``return_traces`` / ``return_results`` /
        ``result_schema`` / ``return_embeddings`` / ``return_images``
        are honored — every other field is recorded on the server-side
        pipeline spec.
    **kwargs
        Same execute-time flags may be passed individually.  Anything
        not recognised is silently ignored (server-side execution
        in service mode is driven by the pipeline spec, not by
        execute-time knobs).
    return_results
        When ``True`` (default), fetch each completed document's
        ``result_data`` from ``GET /v1/ingest/status/{id}`` and
        expose the combined rows on ``result.dataframe`` as a
        ``pandas.DataFrame``. Set to ``False`` to skip those HTTP
        round-trips when only job metadata is needed.
    result_schema
        ``"legacy"`` (default) preserves the existing service
        DataFrame column layout with bulky values stripped and emits
        a deprecation warning when result rows are retained.
        ``"compact"`` opts into the future compact row schema.
    return_embeddings, return_images
        When using legacy result rows, include embedding vectors and
        raw image payloads instead of stripping them from transport
        cells. Defaults remain ``False`` to avoid large responses.

    Returns
    -------
    ServiceIngestResult
        When neither ``return_failures`` nor ``return_traces`` is
        set — a list subclass of per-document completion events with
        extra ``job_id`` / ``failures`` / ``document_ids`` /
        ``elapsed_s`` / ``job_status`` / ``dataframe`` attributes.
    tuple
        With ``return_failures=True`` only — ``(result, failures)``.
        With ``return_traces=True`` only — ``(result, traces)``.
        With both — ``(result, failures, traces)``.  ``failures``
        mirrors ``result.failures``; ``traces`` is the ordered list
        of raw SSE event dicts observed during the run, useful for
        debugging pipeline behaviour without re-running the job.
    """
    return_failures, return_traces, return_results, result_schema, return_embeddings, return_images = (
        self._resolve_execute_flags(params, kwargs)
    )
    del params, kwargs
    self._validate_input_sources(self._inline_texts)
    if not self._documents and not self._buffers and is_blank_inline_corpus(self._inline_texts):
        self._document_ids.clear()
        self._last_run_elapsed_s = 0.0
        self._last_job_id = None
        result = ServiceIngestResult()
        if return_results:
            result.dataframe = _empty_service_result_dataframe(result_schema)
        if return_failures and return_traces:
            return result, [], []
        if return_failures or return_traces:
            return result, []
        return result

    retain_results = return_results or self._save_to_disk_dir is not None
    if retain_results and result_schema == "legacy":
        warnings.warn(_LEGACY_RESULT_SCHEMA_DEPRECATION, DeprecationWarning, stacklevel=2)
    result = ServiceIngestResult()
    traces: list[dict[str, Any]] = []
    rows_by_document: dict[str, list[dict[str, Any]]] = {}
    t0 = time.monotonic()

    documents_completed = 0
    documents_failed = 0
    total_uploaded = 0

    for evt, result_client in self._ingest_events_with_result_client(
        retain_results=retain_results,
        result_schema=result_schema,
        return_embeddings=return_embeddings,
        return_images=return_images,
    ):
        if return_traces:
            traces.append(evt)
        event_type = evt.get("event")

        if event_type == "job_created":
            result.job_id = evt.get("job_id") or result.job_id
            result.trace_id = evt.get("trace_id") or result.trace_id
            continue

        if event_type in ("job_finalized", "job_partial", "job_failed"):
            if event_type == "job_finalized":
                result.job_status = "completed"
            elif event_type == "job_partial":
                result.job_status = "partial_success"
            else:
                result.job_status = "failed"
            continue

        if event_type == "job_progress" or event_type == "job_started":
            continue

        if event_type == "upload_complete":
            total_uploaded += 1
            document_id = evt.get("document_id")
            filename = evt.get("filename")
            if document_id and filename:
                result.document_filenames[str(document_id)] = str(filename)
            if result.job_id is None:
                # Race: SSE delivered an upload_complete before the
                # generator yielded job_created. Fall back to the
                # job_id stamped on the per-doc event by the client.
                result.job_id = evt.get("job_id") or result.job_id
            print(
                f"\r  Job {result.job_id or '?'}  |  "
                f"Uploaded: {total_uploaded}  |  "
                f"Completed: {documents_completed}  |  "
                f"Failed: {documents_failed}",
                end="",
                flush=True,
            )

        elif event_type == "document_complete":
            status = evt.get("status", "completed")
            if status not in ("completed", "failed"):
                continue
            if status == "failed":
                documents_failed += 1
                error = evt.get("error", "unknown error")
                doc_id = evt.get("document_id", "?")
                result.failures.append((doc_id, error))
            else:
                documents_completed += 1
                doc_id = evt.get("document_id", "")
                if return_results or self._save_to_disk_dir is not None:
                    try:
                        rows = self._materialize_completed_document(
                            doc_id,
                            return_results=return_results,
                            client=result_client,
                        )
                        if rows is not None and return_results:
                            rows_by_document[doc_id] = rows
                    except Exception as exc:
                        label = "return_results" if return_results else "save_to_disk"
                        logger.warning("%s: failed to fetch/persist %s: %s", label, doc_id, exc)
                        result.failures.append((doc_id, f"{label}: {exc}"))
            result.append(evt)
            print(
                f"\r  Job {result.job_id or '?'}  |  "
                f"Uploaded: {total_uploaded}  |  "
                f"Completed: {documents_completed}  |  "
                f"Failed: {documents_failed}",
                end="",
                flush=True,
            )

        elif event_type == "upload_failed":
            fname = evt.get("filename", "?")
            error = evt.get("error", "unknown")
            result.failures.append((fname, f"upload failed: {error}"))

    if total_uploaded > 0:
        print()

    result.document_ids = list(self._document_ids)
    result.elapsed_s = time.monotonic() - t0
    if return_results:
        doc_order = [d for d in self._document_ids if d in rows_by_document] or list(rows_by_document)
        result.dataframe = concat_ingest_results(rows_by_document, doc_order)
    self._last_run_elapsed_s = result.elapsed_s
    # Cache the job_id on the ingestor for the get_status() /
    # remaining_jobs() accessors so they can target the job
    # aggregate endpoints once J6 wiring is opted in (kept
    # backwards compatible — get_status() still uses document_ids).
    self._last_job_id = result.job_id

    if return_failures and return_traces:
        return result, list(result.failures), traces
    if return_failures:
        return result, list(result.failures)
    if return_traces:
        return result, traces
    return result

ingest_stream(*, retain_results=False, result_schema='legacy', return_embeddings=False, return_images=False)

Source code in nemo_retriever/service/service_ingestor.py

def ingest_stream(
    self,
    *,
    retain_results: bool = False,
    result_schema: ResultSchema = "legacy",
    return_embeddings: bool = False,
    return_images: bool = False,
) -> Iterator[dict[str, Any]]:
    """Sync generator yielding events as documents are processed.

    Yields dicts with:

    * ``{"event": "job_created", "job_id": ..., "expected_documents": ...}``
    * ``{"event": "upload_complete", "filename": ..., "document_id": ..., "job_id": ...}``
    * ``{"event": "document_complete", "document_id": ..., "status": ..., "job_id": ..., ...}``
    * ``{"event": "upload_failed", "filename": ..., "error": ..., "job_id": ...}``
    * ``{"event": "job_progress", "job_id": ..., "completed": ..., "failed": ..., ...}``
    * ``{"event": "job_finalized"|"job_partial"|"job_failed", "job_id": ..., ...}``
    """
    result_schema = self._normalize_result_schema(result_schema)
    return self._ingest_stream_with_retain(
        retain_results,
        result_schema=result_schema,
        return_embeddings=return_embeddings,
        return_images=return_images,
    )

aingest_stream(*, retain_results=False, result_schema='legacy', return_embeddings=False, return_images=False) async

Source code in nemo_retriever/service/service_ingestor.py

async def aingest_stream(
    self,
    *,
    retain_results: bool = False,
    result_schema: ResultSchema = "legacy",
    return_embeddings: bool = False,
    return_images: bool = False,
) -> AsyncIterator[dict[str, Any]]:
    """Async generator yielding events as documents are processed."""
    result_schema = self._normalize_result_schema(result_schema)
    files = self._collect_inputs()
    if not files:
        return

    self._document_ids.clear()
    async for evt in self._aingest_stream_impl(
        files,
        retain_results=retain_results,
        result_schema=result_schema,
        return_embeddings=return_embeddings,
        return_images=return_images,
    ):
        if evt.get("event") == "upload_complete":
            did = evt.get("document_id")
            if did:
                self._document_ids.append(did)
        yield evt

ingest_async(*, return_failures=False, return_traces=False, return_results=True, result_schema='legacy', return_embeddings=False, return_images=False)

Source code in nemo_retriever/service/service_ingestor.py

def ingest_async(
    self,
    *,
    return_failures: bool = False,
    return_traces: bool = False,
    return_results: bool = True,
    result_schema: ResultSchema = "legacy",
    return_embeddings: bool = False,
    return_images: bool = False,
) -> Any:
    """Run :meth:`ingest` on a background thread; return a ``Future``.

    The flags are forwarded to :meth:`ingest`, so calling
    ``future.result()`` produces the same tuple/list shape that a
    direct synchronous call with the same flags would return.
    """
    from concurrent.futures import ThreadPoolExecutor

    executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="ServiceIngestorAsync")
    return executor.submit(
        self.ingest,
        return_failures=return_failures,
        return_traces=return_traces,
        return_results=return_results,
        result_schema=result_schema,
        return_embeddings=return_embeddings,
        return_images=return_images,
    )

get_status()

Source code in nemo_retriever/service/service_ingestor.py

def get_status(self) -> dict[str, str]:
    """Return ``{document_id: status}`` for every document submitted so far."""
    if not self._document_ids:
        return {}
    url = f"{self._base_url}/v1/ingest/status/batch"
    with httpx.Client(timeout=30.0, headers=self._auth_headers) as client:
        try:
            resp = client.post(url, json={"ids": self._document_ids})
            resp.raise_for_status()
            items = resp.json().get("items", {})
            return {did: info.get("status", "unknown") for did, info in items.items()}
        except Exception as exc:
            logger.warning("Could not fetch bulk status: %s", exc)
            return {did: "unknown" for did in self._document_ids}

completed_jobs()

Source code in nemo_retriever/service/service_ingestor.py

def completed_jobs(self) -> int:
    return sum(1 for s in self.get_status().values() if s == "completed")

failed_jobs()

Source code in nemo_retriever/service/service_ingestor.py

def failed_jobs(self) -> int:
    return sum(1 for s in self.get_status().values() if s == "failed")

cancelled_jobs()

Source code in nemo_retriever/service/service_ingestor.py

def cancelled_jobs(self) -> int:
    return 0

remaining_jobs()

Source code in nemo_retriever/service/service_ingestor.py

def remaining_jobs(self) -> int:
    return sum(1 for s in self.get_status().values() if s in ("processing", "unknown"))

cancel(job_id=None)

Source code in nemo_retriever/service/service_ingestor.py

def cancel(self, job_id: str | None = None) -> dict[str, Any]:
    """Not supported — the server does not expose a cancel endpoint."""
    raise NotImplementedError(
        "Cancel is not supported in service mode. " "The server does not currently expose a cancel endpoint."
    )

nemo_retriever.service.service_ingestor.ServiceIngestResult

Bases: list

Attributes:

Name	Type	Description
job_id	str | None	The server-assigned job aggregate id for this ingest() call. Every :meth:ServiceIngestor.ingest invocation opens exactly one job, sized to len(documents); this is the handle to drive GET /v1/ingest/job/{job_id} follow-ups.
failures	list[tuple[str, str]]	(document_id_or_filename, error_message) pairs for documents that failed during upload or pipeline processing.
document_ids	list[str]	Document identifiers returned by the server, in upload order.
document_filenames	dict[str, str]	Mapping from server document id to the source filename submitted for that document.
elapsed_s	float	Wall-clock seconds from first upload to last result.
job_status	str | None	Final aggregate status reported by the server (completed / failed / partial_success) when a job lifecycle event was observed during the run. None if the stream closed without a terminal job event (e.g. SSE fallback only delivered per-document completions).
trace_id	str | None	Trace id returned by the server on the job_created event. None when tracing is disabled or the service does not include a trace id in the job creation event.
dataframe	Any	When :meth:ServiceIngestor.ingest is called with return_results=True (the default), a pandas.DataFrame of all successfully ingested rows fetched from the service via GET /v1/ingest/status/{document_id}, concatenated in upload order. The current default result_schema="legacy" preserves the same column layout as GraphIngestor.ingest() in inprocess / batch run modes, with bulky raw image and embedding values stripped from cells before transport. Pass result_schema="compact" to opt into the future compact schema. None when return_results=False.

Source code in nemo_retriever/service/service_ingestor.py

class ServiceIngestResult(list):
    """Materialized result returned by :meth:`ServiceIngestor.ingest`.

    Subclasses ``list`` so it satisfies the existing
    ``ingestor.ingest()`` return-type annotation (``List[Any]``); callers
    can iterate it just like a normal list.  Each entry is a per-document
    completion event dict.

    Attributes
    ----------
    job_id
        The server-assigned job aggregate id for this ``ingest()`` call.
        Every :meth:`ServiceIngestor.ingest` invocation opens exactly one
        job, sized to ``len(documents)``; this is the handle to drive
        ``GET /v1/ingest/job/{job_id}`` follow-ups.
    failures
        ``(document_id_or_filename, error_message)`` pairs for documents
        that failed during upload or pipeline processing.
    document_ids
        Document identifiers returned by the server, in upload order.
    document_filenames
        Mapping from server document id to the source filename submitted for
        that document.
    elapsed_s
        Wall-clock seconds from first upload to last result.
    job_status
        Final aggregate status reported by the server
        (``completed`` / ``failed`` / ``partial_success``) when a job
        lifecycle event was observed during the run. ``None`` if the
        stream closed without a terminal job event (e.g. SSE fallback
        only delivered per-document completions).
    trace_id
        Trace id returned by the server on the ``job_created`` event.
        ``None`` when tracing is disabled or the service does not include
        a trace id in the job creation event.
    dataframe
        When :meth:`ServiceIngestor.ingest` is called with
        ``return_results=True`` (the default), a ``pandas.DataFrame``
        of all successfully ingested rows fetched from the service via
        ``GET /v1/ingest/status/{document_id}``, concatenated in upload
        order. The current default ``result_schema="legacy"`` preserves
        the same column layout as ``GraphIngestor.ingest()`` in
        ``inprocess`` / ``batch`` run modes, with bulky raw image and
        embedding values stripped from cells before transport. Pass
        ``result_schema="compact"`` to opt into the future compact schema.
        ``None`` when ``return_results=False``.
    """

    def __init__(self, items: list[dict[str, Any]] | None = None) -> None:
        super().__init__(items or [])
        self.job_id: str | None = None
        self.failures: list[tuple[str, str]] = []
        self.document_ids: list[str] = []
        self.document_filenames: dict[str, str] = {}
        self.elapsed_s: float = 0.0
        self.job_status: str | None = None
        self.trace_id: str | None = None
        self.dataframe: Any = None

    def __repr__(self) -> str:
        return (
            f"ServiceIngestResult(job_id={self.job_id!r}, "
            f"documents={len(self)}, "
            f"failures={len(self.failures)}, "
            f"elapsed_s={self.elapsed_s:.2f})"
        )

job_id = None instance-attribute

failures = [] instance-attribute

document_ids = [] instance-attribute

document_filenames = {} instance-attribute

elapsed_s = 0.0 instance-attribute

job_status = None instance-attribute

trace_id = None instance-attribute

dataframe = None instance-attribute

Retrieve

Retriever embeds queries, retrieves from a vector database, and can rerank results. Pass embed_kwargs that match EmbedParams, vdb_kwargs for RetrieveVdbOperator, and rerank_kwargs for NemotronRerankActor.

Retriever.answer() requires an LLMClient and returns AnswerResult. Pass an AnswerJudge only when you also pass reference. Those types are documented in LLM clients, tasks, and results.

Retriever.pipeline() returns RetrieverPipelineBuilder. generate() accepts a LiteLLMClient instance or model= keyword arguments. judge() accepts an LLMJudge instance or model= keyword arguments.

The package exports nemo_retriever.retriever as a default Retriever() instance. Construct Retriever directly when you need a configured object.

nemo_retriever.graph.retriever.Retriever dataclass

Source code in nemo_retriever/graph/retriever.py

@dataclass
class Retriever:
    """Graph-based query helper: batch embed → VDB retrieve [→ Nemotron rerank].

    Configuration is passed through ``embed_kwargs`` (:class:`~nemo_retriever.params.EmbedParams`),
    ``vdb_kwargs`` (constructor kwargs for :class:`~nemo_retriever.vdb.operators.RetrieveVdbOperator`),
    and optional ``rerank_kwargs`` for :class:`~nemo_retriever.rerank.rerank.NemotronRerankActor`.

    See ``retriever.md`` for examples.
    """

    run_mode: Literal["local", "service"] = "local"
    """``local`` uses archetype batch embed resolution; ``service`` forces CPU HTTP embed."""

    top_k: int = 10
    rerank: bool = False
    """When ``True``, append :class:`~nemo_retriever.rerank.rerank.NemotronRerankActor` after retrieval."""

    graph: Any = None
    """Custom :class:`~nemo_retriever.graph.pipeline_graph.Graph`. When set, ``embed_kwargs`` /
    ``vdb_kwargs`` default-graph fields are ignored for construction (you still pass execute kwargs)."""

    embed_kwargs: dict[str, Any] = field(default_factory=dict)
    vdb_kwargs: dict[str, Any] = field(default_factory=dict)
    rerank_kwargs: dict[str, Any] = field(default_factory=dict)

    _cached_graph: Any = field(default=None, init=False, repr=False, compare=False)
    _cache_key: Any = field(default=None, init=False, repr=False, compare=False)
    _lancedb_capabilities_cache: dict[tuple[str, str], LanceTableCapabilities] = field(
        default_factory=dict, init=False, repr=False, compare=False
    )

    def __post_init__(self) -> None:
        if self.run_mode not in ("local", "service"):
            raise ValueError("run_mode must be 'local' or 'service'")

    def _merge_embed_params(self, extra: Optional[dict[str, Any]] = None) -> Any:
        from nemo_retriever.models import _LOCAL_INGEST_EMBED_BACKENDS, normalize_backend
        from nemo_retriever.common.params import EmbedParams

        base: dict[str, Any] = {
            "model_name": VL_EMBED_MODEL,
            "embed_model_name": VL_EMBED_MODEL,
            "input_type": "query",
            "text_column": "text",
            "inference_batch_size": 32,
            "embed_inference_batch_size": 32,
            "local_ingest_embed_backend": "hf",
        }
        overrides = {**dict(self.embed_kwargs or {}), **dict(extra or {})}
        merged = {**base, **overrides}
        endpoint = str(merged.get("embedding_endpoint") or merged.get("embed_invoke_url") or "").strip()
        if self.run_mode == "local" and not endpoint and overrides.get("local_ingest_embed_backend") is None:
            model_id = str(merged.get("embed_model_name") or merged.get("model_name") or "").strip()
            spec = resolve_embed_model_spec(model_id, revision=merged.get("embed_model_revision"))
            merged["local_ingest_embed_backend"] = "vllm" if spec.requires_vllm else "hf"
            merged["embed_model_revision"] = spec.revision
        if "local_ingest_embed_backend" in merged and merged["local_ingest_embed_backend"] is not None:
            merged["local_ingest_embed_backend"] = normalize_backend(
                str(merged["local_ingest_embed_backend"]),
                _LOCAL_INGEST_EMBED_BACKENDS,
                field_name="local_ingest_embed_backend",
                default="vllm",
            )
        params = EmbedParams.model_validate(merged)
        if self.run_mode == "service":
            url = (params.embedding_endpoint or params.embed_invoke_url or "").strip()
            if not url:
                raise ValueError(
                    "run_mode='service' requires a non-empty HTTP embedding URL. "
                    "Set ``embedding_endpoint`` or ``embed_invoke_url`` inside ``embed_kwargs``."
                )
        return params

    def _merge_rerank_actor_kwargs(self) -> dict[str, Any]:
        return {**_default_rerank_actor_kwargs(), **dict(self.rerank_kwargs or {})}

    def _refine_factor(self) -> int:
        if not self.rerank:
            return 1
        return int(self._merge_rerank_actor_kwargs().get("refine_factor", 4))

    def _build_default_graph(self, *, embed_extra: Optional[dict[str, Any]] = None) -> Any:
        from nemo_retriever.operators.rerank import NemotronRerankActor
        from nemo_retriever.operators.embed.cpu_operator import _BatchEmbedCPUActor
        from nemo_retriever.operators.embed.operators import _BatchEmbedActor

        embed_params = self._merge_embed_params(embed_extra)
        if self.run_mode == "service":
            embed_op = _BatchEmbedCPUActor(params=embed_params)
        else:
            embed_op = _BatchEmbedActor(params=embed_params)

        vdb_init = _coerce_vdb_init(self.vdb_kwargs)
        retrieve = RetrieveVdbOperator(
            explode_for_rerank=self.rerank,
            **vdb_init,
        )

        chain = embed_op >> retrieve
        if self.rerank:
            rk = self._merge_rerank_actor_kwargs()
            rk.pop("refine_factor", None)
            chain = chain >> NemotronRerankActor(**rk)

        return chain

    def _get_graph(self, *, embed_extra: Optional[dict[str, Any]] = None) -> Any:
        if self.graph is not None:
            return self.graph

        key = (
            self.run_mode,
            self.rerank,
            json.dumps(self.vdb_kwargs, sort_keys=True, default=str),
            json.dumps(self.embed_kwargs, sort_keys=True, default=str),
            json.dumps(self.rerank_kwargs, sort_keys=True, default=str),
            json.dumps(embed_extra or {}, sort_keys=True, default=str),
        )
        if self._cached_graph is not None and self._cache_key == key:
            return self._cached_graph
        g = self._build_default_graph(embed_extra=embed_extra)
        self._cached_graph = g
        self._cache_key = key
        return g

    def _execute_queries_graph(
        self,
        query_texts: list[str],
        *,
        effective_top_k: int,
        retrieval_top_k: int,
        vdb_call_kwargs: Optional[dict[str, Any]],
        embed_extra: Optional[dict[str, Any]],
    ) -> list[list[dict[str, Any]]]:
        if self.graph is None:
            embed_params = self._merge_embed_params(embed_extra)
            text_col = str(embed_params.text_column)
        else:
            # A caller-owned graph controls its own operators and does not
            # require default embedding configuration. In particular, avoid
            # resolving a local embedding model merely to choose its input
            # column. Agentic result graphs use ``query_text`` and delegate
            # actual retrieval to their configured inner retriever.
            text_col = str({**dict(self.embed_kwargs or {}), **dict(embed_extra or {})}.get("text_column") or "text")
        df = pd.DataFrame({text_col: query_texts})

        # Hybrid retrieval relies on these ordered query strings staying aligned
        # with the embedded rows produced from ``df``. If this query graph grows
        # distributed/shuffled stages, carry row-local query text or IDs instead.
        graph = self._get_graph(embed_extra=embed_extra)

        exec_kwargs: dict[str, Any] = {
            **filter_retrieval_kwargs(dict(vdb_call_kwargs or {})),
            "top_k": int(retrieval_top_k),
            "query_texts": query_texts,
        }
        if self.graph is None:
            leaves = graph.execute_in_place(df, **exec_kwargs)
        else:
            # Preserve resolve-per-query behavior for caller-owned graphs, which
            # may be mutated between calls.
            resolve = getattr(graph, "resolve_for_local_execution", None)
            if not callable(resolve):
                raise TypeError("graph must provide resolve_for_local_execution() (e.g. pipeline_graph.Graph)")
            resolved = resolve()
            leaves = resolved.execute(df, **exec_kwargs)
        if len(leaves) != 1:
            raise RuntimeError(
                f"Retriever query graph must yield exactly one leaf output; got {len(leaves)}. "
                "Use a linear graph or adjust your custom ``graph``."
            )
        out = leaves[0]

        if isinstance(out, pd.DataFrame):
            if not self.rerank:
                raise TypeError(
                    "Graph returned a DataFrame but ``rerank`` is False; expected list[list[dict]] from retrieval."
                )
            rk = self._merge_rerank_actor_kwargs()
            score_col = str(rk.get("score_column", "rerank_score"))
            return rerank_long_dataframe_to_hits(
                out, query_texts=query_texts, top_k=int(effective_top_k), score_column=score_col
            )
        if not isinstance(out, list):
            raise TypeError(f"Unexpected query graph output type: {type(out).__name__}")
        return out

    def _inspect_lancedb_capabilities(self, uri: str, table_name: str) -> LanceTableCapabilities:
        key = (uri, table_name)
        caps = self._lancedb_capabilities_cache.get(key)
        if caps is None:
            caps = inspect_lancedb_table(uri, table_name)
            self._lancedb_capabilities_cache[key] = caps
        return caps

    def _resolve_lancedb_query_mode(
        self,
        runtime_vdb_kwargs: Optional[dict[str, Any]],
    ) -> tuple[str, LanceTableCapabilities, str, str, bool] | None:
        if self.graph is not None:
            return None

        lancedb_kwargs = dict(self.vdb_kwargs or {})
        if "vdb" in lancedb_kwargs:
            return None
        if "vdb_op" in lancedb_kwargs:
            if str(lancedb_kwargs.get("vdb_op") or "").strip().lower() != "lancedb":
                return None
            lancedb_kwargs = dict(lancedb_kwargs.get("vdb_kwargs") or {})
        lancedb_kwargs.update(dict(runtime_vdb_kwargs or {}))

        uri = str(
            lancedb_kwargs.get("table_path")
            or lancedb_kwargs.get("uri")
            or lancedb_kwargs.get("lancedb_uri")
            or "lancedb"
        )
        table_name = str(lancedb_kwargs.get("table_name") or lancedb_kwargs.get("lancedb_table") or "nv-ingest")
        caps = self._inspect_lancedb_capabilities(uri, table_name)

        mode_override = str(lancedb_kwargs.get("retrieval_mode") or "auto").strip().lower()
        if mode_override not in {"auto", "dense", "hybrid", "sparse"}:
            raise ValueError(
                f"Unsupported LanceDB retrieval mode {mode_override!r}; " "use 'auto', 'dense', 'hybrid', or 'sparse'."
            )
        if "hybrid" in lancedb_kwargs:
            mode_override = "hybrid" if bool(lancedb_kwargs["hybrid"]) else "dense"
        mode = caps.retrieval_mode if mode_override == "auto" else cast(LanceRetrievalMode, mode_override)

        if mode == "unknown":
            raise ValueError(
                f"LanceDB table {table_name!r} at {uri!r} is not queryable: "
                "no vector column or FTS index was detected."
            )
        if mode == "dense" and not caps.has_vector:
            raise ValueError(
                f"LanceDB table {table_name!r} at {uri!r} cannot run dense retrieval: " "no vector column was detected."
            )
        if mode == "hybrid" and (not caps.has_vector or not caps.has_fts):
            raise ValueError(
                f"LanceDB table {table_name!r} at {uri!r} cannot run hybrid retrieval: "
                "both a vector column and FTS index are required."
            )
        if mode == "sparse" and not caps.has_fts:
            raise ValueError(
                f"LanceDB table {table_name!r} at {uri!r} cannot run sparse retrieval: " "no FTS index was detected."
            )

        return mode, caps, uri, table_name, mode_override != "auto"

    @staticmethod
    def _embedding_model_from_kwargs(kwargs: Optional[dict[str, Any]]) -> str | None:
        values = dict(kwargs or {})
        for key in ("model_name", "embed_model_name"):
            value = str(values.get(key) or "").strip()
            if value:
                return value
        return None

    def _resolve_embed_kwargs(
        self,
        index_model: str | None,
        runtime_embed_kwargs: Optional[dict[str, Any]],
        index_revision: str | None = None,
    ) -> dict[str, Any]:
        """Choose the query model snapshot: explicit override, index metadata, or default."""
        resolved = dict(runtime_embed_kwargs or {})
        runtime_model = self._embedding_model_from_kwargs(runtime_embed_kwargs)
        configured_model = self._embedding_model_from_kwargs(self.embed_kwargs)
        explicit_model = runtime_model or configured_model
        model_name = explicit_model or index_model
        model_name = resolve_embed_model(model_name)
        resolved["model_name"] = model_name
        resolved["embed_model_name"] = model_name
        if runtime_model and "embed_model_revision" not in resolved:
            resolved_configured = resolve_embed_model(configured_model) if configured_model else None
            if resolved_configured != model_name:
                resolved["embed_model_revision"] = None
        if explicit_model is None and index_revision:
            resolved.setdefault("embed_model_revision", index_revision)
        return resolved

    def _execute_sparse_lancedb_queries(
        self,
        query_texts: list[str],
        *,
        retrieval_top_k: int,
        vdb_call_kwargs: Optional[dict[str, Any]],
        caps: LanceTableCapabilities,
        uri: str,
        table_name: str,
    ) -> list[list[dict[str, Any]]]:
        from nemo_retriever.common.vdb.lancedb import LanceDB

        text_column = caps.text_column or "text"
        retrieval_kwargs = {
            **filter_retrieval_kwargs(dict(vdb_call_kwargs or {})),
            "top_k": int(retrieval_top_k),
            "table_path": uri,
            "table_name": table_name,
            "text_column_name": text_column,
        }
        vdb = LanceDB(uri=uri, table_name=table_name, overwrite=False, sparse=True)
        return normalize_retrieval_results(vdb.sparse_retrieval(query_texts, **retrieval_kwargs))

    def query(
        self,
        query: str,
        *,
        top_k: Optional[int] = None,
        candidate_k: Optional[int] = None,
        page_dedup: bool = False,
        content_types: str | Sequence[str] | None = None,
        vdb_kwargs: Optional[dict[str, Any]] = None,
        embed_kwargs: Optional[dict[str, Any]] = None,
    ) -> list[RetrievalHit]:
        """Run one retrieval query and return shaped hits.

        ``top_k`` is the final number of hits to return. ``candidate_k`` is the
        wider pre-filter/pre-dedup candidate pool and must be greater than or
        equal to ``top_k``. Increase it when page deduplication or content-type
        filtering would otherwise reduce the final hit count. ``page_dedup``
        keeps the first hit per document page. ``content_types`` accepts a
        comma-separated string or sequence of content types to keep, such as
        ``"text,table"``, and normalizes values to the canonical content types
        stored in hit metadata. Hits with missing or unknown content types are
        excluded while this filter is active. Page deduplication and
        content-type filtering are applied after vector retrieval, preserving
        retriever ranking order.
        """
        return self.queries(
            [query],
            top_k=top_k,
            candidate_k=candidate_k,
            page_dedup=page_dedup,
            content_types=content_types,
            vdb_kwargs=vdb_kwargs,
            embed_kwargs=embed_kwargs,
        )[0]

    def queries(
        self,
        queries: Sequence[str],
        *,
        top_k: Optional[int] = None,
        candidate_k: Optional[int] = None,
        page_dedup: bool = False,
        content_types: str | Sequence[str] | None = None,
        vdb_kwargs: Optional[dict[str, Any]] = None,
        embed_kwargs: Optional[dict[str, Any]] = None,
    ) -> list[list[RetrievalHit]]:
        """Run retrieval for multiple query strings and return shaped hits.

        ``top_k`` is the final number of hits to return. ``candidate_k`` is the
        wider pre-filter/pre-dedup candidate pool and must be greater than or
        equal to ``top_k``. Increase it when page deduplication or content-type
        filtering would otherwise reduce the final hit count. ``page_dedup``
        keeps the first hit per document page. ``content_types`` accepts a
        comma-separated string or sequence of content types to keep, such as
        ``"text,table"``, and normalizes values to the canonical content types
        stored in hit metadata. Hits with missing or unknown content types are
        excluded while this filter is active. Page deduplication and
        content-type filtering are applied after vector retrieval, preserving
        retriever ranking order.
        """
        query_texts = [str(q) for q in queries]
        if not query_texts:
            return []

        effective_top_k = int(top_k) if top_k is not None else int(self.top_k)
        candidate_top_k = int(candidate_k) if candidate_k is not None else effective_top_k
        if candidate_top_k < effective_top_k:
            raise ValueError(
                f"candidate_k ({candidate_top_k}) must be greater than or equal to top_k ({effective_top_k})."
            )
        refine = self._refine_factor()
        retrieval_top_k = candidate_top_k * refine if self.rerank else candidate_top_k

        vdb_call_kwargs = dict(vdb_kwargs or {})
        index_model: str | None = None
        index_revision: str | None = None
        explicit_model = self._embedding_model_from_kwargs(embed_kwargs) or self._embedding_model_from_kwargs(
            self.embed_kwargs
        )
        if self.graph is None:
            metadata_reader = RetrieveVdbOperator(**_coerce_vdb_init(self.vdb_kwargs))
            index_model = metadata_reader.get_index_metadata("embedding_model_name", **vdb_call_kwargs)
            index_revision = metadata_reader.get_index_metadata("embedding_model_revision", **vdb_call_kwargs)
            if explicit_model and index_model:
                resolved_explicit_model = resolve_embed_model(explicit_model)
                resolved_index_model = resolve_embed_model(index_model)
                if resolved_explicit_model != resolved_index_model:
                    logger.warning(
                        "The explicitly configured query embedding model %r differs from the model %r "
                        "recorded on the index. Results may be unreliable because different embedding "
                        "models can use incompatible vector spaces. Use the index model or rebuild the "
                        "index with the query model. Continuing with the explicitly configured model.",
                        resolved_explicit_model,
                        resolved_index_model,
                    )

        lancedb_mode = self._resolve_lancedb_query_mode(vdb_call_kwargs)
        for key in _QUERY_ROUTING_VDB_KWARGS:
            vdb_call_kwargs.pop(key, None)
        if lancedb_mode is not None:
            mode, caps, uri, table_name, has_mode_override = lancedb_mode
            if mode == "sparse":
                raw_hits = self._execute_sparse_lancedb_queries(
                    query_texts,
                    retrieval_top_k=retrieval_top_k,
                    vdb_call_kwargs=vdb_call_kwargs,
                    caps=caps,
                    uri=uri,
                    table_name=table_name,
                )
                return [
                    shape_query_hits(
                        hits,
                        top_k=effective_top_k,
                        page_dedup=page_dedup,
                        content_types=content_types,
                    )
                    for hits in raw_hits
                ]
            if mode == "hybrid":
                vdb_call_kwargs["hybrid"] = True
                vdb_call_kwargs.setdefault("hybrid_fusion", DEFAULT_HYBRID_FUSION_POLICY)
            elif mode == "dense" and has_mode_override:
                vdb_call_kwargs["hybrid"] = False
            if caps.vector_column and caps.vector_column != "vector":
                vdb_call_kwargs.setdefault("vector_column_name", caps.vector_column)
        if self.graph is None:
            embed_kwargs = self._resolve_embed_kwargs(index_model, embed_kwargs, index_revision)

        raw_hits = self._execute_queries_graph(
            query_texts,
            effective_top_k=candidate_top_k,
            retrieval_top_k=retrieval_top_k,
            vdb_call_kwargs=vdb_call_kwargs,
            embed_extra=embed_kwargs,
        )
        return [
            shape_query_hits(
                hits,
                top_k=effective_top_k,
                page_dedup=page_dedup,
                content_types=content_types,
            )
            for hits in raw_hits
        ]

    def retrieve(
        self,
        query: str,
        top_k: Optional[int] = None,
        *,
        vdb_kwargs: Optional[dict[str, Any]] = None,
        embed_kwargs: Optional[dict[str, Any]] = None,
    ) -> "RetrievalResult":
        from nemo_retriever.models.llm.types import RetrievalResult

        hits = self.query(query, top_k=top_k, vdb_kwargs=vdb_kwargs, embed_kwargs=embed_kwargs)

        chunks: list[str] = []
        metadata: list[dict[str, Any]] = []
        for hit in hits:
            chunks.append(str(hit.get("text", "")))
            metadata.append({k: v for k, v in hit.items() if k != "text"})
        return RetrievalResult(chunks=chunks, metadata=metadata)

    def retrieve_batch(
        self,
        queries: Sequence[str],
        *,
        top_k: Optional[int] = None,
        vdb_kwargs: Optional[dict[str, Any]] = None,
        embed_kwargs: Optional[dict[str, Any]] = None,
    ) -> list["RetrievalResult"]:
        from nemo_retriever.models.llm.types import RetrievalResult

        query_texts = [str(q) for q in queries]
        if not query_texts:
            return []

        hits_per_query = self.queries(query_texts, top_k=top_k, vdb_kwargs=vdb_kwargs, embed_kwargs=embed_kwargs)

        results: list[RetrievalResult] = []
        for hits in hits_per_query:
            chunks = [str(hit.get("text", "")) for hit in hits]
            metadata = [{k: v for k, v in hit.items() if k != "text"} for hit in hits]
            results.append(RetrievalResult(chunks=chunks, metadata=metadata))
        return results

    def answer(
        self,
        query: str,
        *,
        llm: "LLMClient",
        judge: Optional["AnswerJudge"] = None,
        reference: Optional[str] = None,
        top_k: Optional[int] = None,
        reasoning_enabled: Optional[bool] = None,
        vdb_kwargs: Optional[dict[str, Any]] = None,
        embed_kwargs: Optional[dict[str, Any]] = None,
    ) -> "AnswerResult":
        from nemo_retriever.models.llm.types import (
            AnswerRequest,
            build_answer_result,
        )

        if judge is not None and reference is None:
            raise ValueError("judge requires reference")

        answer_req = AnswerRequest(
            query=query,
            top_k=int(top_k) if top_k is not None else int(self.top_k),
            reasoning_enabled=reasoning_enabled,
            reference=reference,
            judge_enabled=judge is not None,
        )
        retrieved = self.retrieve(
            answer_req.query,
            top_k=answer_req.top_k,
            vdb_kwargs=vdb_kwargs,
            embed_kwargs=embed_kwargs,
        )

        generate_kwargs: dict[str, Any] = {}
        if answer_req.reasoning_enabled is not None:
            generate_kwargs["reasoning_enabled"] = answer_req.reasoning_enabled
        gen = llm.generate(answer_req.query, retrieved.chunks, **generate_kwargs)

        return build_answer_result(
            query=answer_req.query,
            retrieval=retrieved,
            generation=gen,
            reference=answer_req.reference,
            judge=judge if answer_req.judge_enabled else None,
        )

    def pipeline(self, *, top_k: Optional[int] = None) -> "RetrieverPipelineBuilder":
        effective_top_k = int(top_k) if top_k is not None else int(self.top_k)
        return RetrieverPipelineBuilder(self, top_k=effective_top_k)

    def generate_sql(self, query: str) -> str:
        from nemo_retriever.tabular_data.retrieval import generate_sql

        return generate_sql(query)

run_mode = 'local' class-attribute instance-attribute

top_k = 10 class-attribute instance-attribute

rerank = False class-attribute instance-attribute

graph = None class-attribute instance-attribute

embed_kwargs = field(default_factory=dict) class-attribute instance-attribute

vdb_kwargs = field(default_factory=dict) class-attribute instance-attribute

rerank_kwargs = field(default_factory=dict) class-attribute instance-attribute

query(query, *, top_k=None, candidate_k=None, page_dedup=False, content_types=None, vdb_kwargs=None, embed_kwargs=None)

Source code in nemo_retriever/graph/retriever.py

def query(
    self,
    query: str,
    *,
    top_k: Optional[int] = None,
    candidate_k: Optional[int] = None,
    page_dedup: bool = False,
    content_types: str | Sequence[str] | None = None,
    vdb_kwargs: Optional[dict[str, Any]] = None,
    embed_kwargs: Optional[dict[str, Any]] = None,
) -> list[RetrievalHit]:
    """Run one retrieval query and return shaped hits.

    ``top_k`` is the final number of hits to return. ``candidate_k`` is the
    wider pre-filter/pre-dedup candidate pool and must be greater than or
    equal to ``top_k``. Increase it when page deduplication or content-type
    filtering would otherwise reduce the final hit count. ``page_dedup``
    keeps the first hit per document page. ``content_types`` accepts a
    comma-separated string or sequence of content types to keep, such as
    ``"text,table"``, and normalizes values to the canonical content types
    stored in hit metadata. Hits with missing or unknown content types are
    excluded while this filter is active. Page deduplication and
    content-type filtering are applied after vector retrieval, preserving
    retriever ranking order.
    """
    return self.queries(
        [query],
        top_k=top_k,
        candidate_k=candidate_k,
        page_dedup=page_dedup,
        content_types=content_types,
        vdb_kwargs=vdb_kwargs,
        embed_kwargs=embed_kwargs,
    )[0]

queries(queries, *, top_k=None, candidate_k=None, page_dedup=False, content_types=None, vdb_kwargs=None, embed_kwargs=None)

Source code in nemo_retriever/graph/retriever.py

def queries(
    self,
    queries: Sequence[str],
    *,
    top_k: Optional[int] = None,
    candidate_k: Optional[int] = None,
    page_dedup: bool = False,
    content_types: str | Sequence[str] | None = None,
    vdb_kwargs: Optional[dict[str, Any]] = None,
    embed_kwargs: Optional[dict[str, Any]] = None,
) -> list[list[RetrievalHit]]:
    """Run retrieval for multiple query strings and return shaped hits.

    ``top_k`` is the final number of hits to return. ``candidate_k`` is the
    wider pre-filter/pre-dedup candidate pool and must be greater than or
    equal to ``top_k``. Increase it when page deduplication or content-type
    filtering would otherwise reduce the final hit count. ``page_dedup``
    keeps the first hit per document page. ``content_types`` accepts a
    comma-separated string or sequence of content types to keep, such as
    ``"text,table"``, and normalizes values to the canonical content types
    stored in hit metadata. Hits with missing or unknown content types are
    excluded while this filter is active. Page deduplication and
    content-type filtering are applied after vector retrieval, preserving
    retriever ranking order.
    """
    query_texts = [str(q) for q in queries]
    if not query_texts:
        return []

    effective_top_k = int(top_k) if top_k is not None else int(self.top_k)
    candidate_top_k = int(candidate_k) if candidate_k is not None else effective_top_k
    if candidate_top_k < effective_top_k:
        raise ValueError(
            f"candidate_k ({candidate_top_k}) must be greater than or equal to top_k ({effective_top_k})."
        )
    refine = self._refine_factor()
    retrieval_top_k = candidate_top_k * refine if self.rerank else candidate_top_k

    vdb_call_kwargs = dict(vdb_kwargs or {})
    index_model: str | None = None
    index_revision: str | None = None
    explicit_model = self._embedding_model_from_kwargs(embed_kwargs) or self._embedding_model_from_kwargs(
        self.embed_kwargs
    )
    if self.graph is None:
        metadata_reader = RetrieveVdbOperator(**_coerce_vdb_init(self.vdb_kwargs))
        index_model = metadata_reader.get_index_metadata("embedding_model_name", **vdb_call_kwargs)
        index_revision = metadata_reader.get_index_metadata("embedding_model_revision", **vdb_call_kwargs)
        if explicit_model and index_model:
            resolved_explicit_model = resolve_embed_model(explicit_model)
            resolved_index_model = resolve_embed_model(index_model)
            if resolved_explicit_model != resolved_index_model:
                logger.warning(
                    "The explicitly configured query embedding model %r differs from the model %r "
                    "recorded on the index. Results may be unreliable because different embedding "
                    "models can use incompatible vector spaces. Use the index model or rebuild the "
                    "index with the query model. Continuing with the explicitly configured model.",
                    resolved_explicit_model,
                    resolved_index_model,
                )

    lancedb_mode = self._resolve_lancedb_query_mode(vdb_call_kwargs)
    for key in _QUERY_ROUTING_VDB_KWARGS:
        vdb_call_kwargs.pop(key, None)
    if lancedb_mode is not None:
        mode, caps, uri, table_name, has_mode_override = lancedb_mode
        if mode == "sparse":
            raw_hits = self._execute_sparse_lancedb_queries(
                query_texts,
                retrieval_top_k=retrieval_top_k,
                vdb_call_kwargs=vdb_call_kwargs,
                caps=caps,
                uri=uri,
                table_name=table_name,
            )
            return [
                shape_query_hits(
                    hits,
                    top_k=effective_top_k,
                    page_dedup=page_dedup,
                    content_types=content_types,
                )
                for hits in raw_hits
            ]
        if mode == "hybrid":
            vdb_call_kwargs["hybrid"] = True
            vdb_call_kwargs.setdefault("hybrid_fusion", DEFAULT_HYBRID_FUSION_POLICY)
        elif mode == "dense" and has_mode_override:
            vdb_call_kwargs["hybrid"] = False
        if caps.vector_column and caps.vector_column != "vector":
            vdb_call_kwargs.setdefault("vector_column_name", caps.vector_column)
    if self.graph is None:
        embed_kwargs = self._resolve_embed_kwargs(index_model, embed_kwargs, index_revision)

    raw_hits = self._execute_queries_graph(
        query_texts,
        effective_top_k=candidate_top_k,
        retrieval_top_k=retrieval_top_k,
        vdb_call_kwargs=vdb_call_kwargs,
        embed_extra=embed_kwargs,
    )
    return [
        shape_query_hits(
            hits,
            top_k=effective_top_k,
            page_dedup=page_dedup,
            content_types=content_types,
        )
        for hits in raw_hits
    ]

retrieve(query, top_k=None, *, vdb_kwargs=None, embed_kwargs=None)

Source code in nemo_retriever/graph/retriever.py

def retrieve(
    self,
    query: str,
    top_k: Optional[int] = None,
    *,
    vdb_kwargs: Optional[dict[str, Any]] = None,
    embed_kwargs: Optional[dict[str, Any]] = None,
) -> "RetrievalResult":
    from nemo_retriever.models.llm.types import RetrievalResult

    hits = self.query(query, top_k=top_k, vdb_kwargs=vdb_kwargs, embed_kwargs=embed_kwargs)

    chunks: list[str] = []
    metadata: list[dict[str, Any]] = []
    for hit in hits:
        chunks.append(str(hit.get("text", "")))
        metadata.append({k: v for k, v in hit.items() if k != "text"})
    return RetrievalResult(chunks=chunks, metadata=metadata)

retrieve_batch(queries, *, top_k=None, vdb_kwargs=None, embed_kwargs=None)

Source code in nemo_retriever/graph/retriever.py

def retrieve_batch(
    self,
    queries: Sequence[str],
    *,
    top_k: Optional[int] = None,
    vdb_kwargs: Optional[dict[str, Any]] = None,
    embed_kwargs: Optional[dict[str, Any]] = None,
) -> list["RetrievalResult"]:
    from nemo_retriever.models.llm.types import RetrievalResult

    query_texts = [str(q) for q in queries]
    if not query_texts:
        return []

    hits_per_query = self.queries(query_texts, top_k=top_k, vdb_kwargs=vdb_kwargs, embed_kwargs=embed_kwargs)

    results: list[RetrievalResult] = []
    for hits in hits_per_query:
        chunks = [str(hit.get("text", "")) for hit in hits]
        metadata = [{k: v for k, v in hit.items() if k != "text"} for hit in hits]
        results.append(RetrievalResult(chunks=chunks, metadata=metadata))
    return results

answer(query, *, llm, judge=None, reference=None, top_k=None, reasoning_enabled=None, vdb_kwargs=None, embed_kwargs=None)

Source code in nemo_retriever/graph/retriever.py

def answer(
    self,
    query: str,
    *,
    llm: "LLMClient",
    judge: Optional["AnswerJudge"] = None,
    reference: Optional[str] = None,
    top_k: Optional[int] = None,
    reasoning_enabled: Optional[bool] = None,
    vdb_kwargs: Optional[dict[str, Any]] = None,
    embed_kwargs: Optional[dict[str, Any]] = None,
) -> "AnswerResult":
    from nemo_retriever.models.llm.types import (
        AnswerRequest,
        build_answer_result,
    )

    if judge is not None and reference is None:
        raise ValueError("judge requires reference")

    answer_req = AnswerRequest(
        query=query,
        top_k=int(top_k) if top_k is not None else int(self.top_k),
        reasoning_enabled=reasoning_enabled,
        reference=reference,
        judge_enabled=judge is not None,
    )
    retrieved = self.retrieve(
        answer_req.query,
        top_k=answer_req.top_k,
        vdb_kwargs=vdb_kwargs,
        embed_kwargs=embed_kwargs,
    )

    generate_kwargs: dict[str, Any] = {}
    if answer_req.reasoning_enabled is not None:
        generate_kwargs["reasoning_enabled"] = answer_req.reasoning_enabled
    gen = llm.generate(answer_req.query, retrieved.chunks, **generate_kwargs)

    return build_answer_result(
        query=answer_req.query,
        retrieval=retrieved,
        generation=gen,
        reference=answer_req.reference,
        judge=judge if answer_req.judge_enabled else None,
    )

pipeline(*, top_k=None)

Source code in nemo_retriever/graph/retriever.py

def pipeline(self, *, top_k: Optional[int] = None) -> "RetrieverPipelineBuilder":
    effective_top_k = int(top_k) if top_k is not None else int(self.top_k)
    return RetrieverPipelineBuilder(self, top_k=effective_top_k)

generate_sql(query)

Source code in nemo_retriever/graph/retriever.py

def generate_sql(self, query: str) -> str:
    from nemo_retriever.tabular_data.retrieval import generate_sql

    return generate_sql(query)

nemo_retriever.graph.retriever.RetrieverPipelineBuilder

Source code in nemo_retriever/graph/retriever.py

class RetrieverPipelineBuilder:
    """Fluent builder for live-RAG batch operator graphs.

    Returned from :meth:`Retriever.pipeline`.  Each builder method appends
    an :class:`~nemo_retriever.evaluation.eval_operator.EvalOperator` to an
    internal list; :meth:`run` composes them into a graph via the existing
    ``>>`` chaining and executes it on a DataFrame built from the provided
    queries.

    Example:
        >>> builder = retriever.pipeline()  # doctest: +SKIP
        >>> df = builder.generate(llm).score().judge(judge).run(  # doctest: +SKIP
        ...     queries=["q1", "q2"],
        ...     reference=["r1", "r2"],
        ... )
    """

    def __init__(self, retriever: "Retriever", *, top_k: int = 5) -> None:
        self._retriever = retriever
        self._top_k = int(top_k)
        self._steps: list[Any] = []

    def with_retrieval(self, *, top_k: int) -> "RetrieverPipelineBuilder":
        """Override the ``top_k`` used for the live retrieval source."""
        self._top_k = int(top_k)
        return self

    def generate(
        self,
        llm: Optional[Any] = None,
        /,
        *,
        model: Optional[str] = None,
        **kwargs: Any,
    ) -> "RetrieverPipelineBuilder":
        """Append a :class:`QAGenerationOperator` step.

        Accepts either a pre-built
        :class:`~nemo_retriever.llm.clients.LiteLLMClient` (whose transport
        and sampling params are unpacked onto the operator) or the flat
        ``model=..., api_base=..., ...`` kwargs forwarded to the operator
        constructor directly.

        Raises:
            ValueError: If neither ``llm`` nor ``model`` is provided.
        """
        from nemo_retriever.tools.evaluation.generation import QAGenerationOperator

        if llm is None and model is None:
            raise ValueError("generate() requires either llm= or model=")

        if llm is not None:
            transport = llm.transport
            sampling = llm.sampling
            operator = QAGenerationOperator(
                model=transport.model,
                api_base=transport.api_base,
                api_key=transport.api_key,
                temperature=sampling.temperature,
                top_p=sampling.top_p,
                max_tokens=sampling.max_tokens,
                extra_params=dict(transport.extra_params) if transport.extra_params else None,
                num_retries=transport.num_retries,
                timeout=transport.timeout,
                rag_system_prompt=transport.rag_system_prompt,
                rag_system_prompt_prefix=transport.rag_system_prompt_prefix,
                reasoning_enabled=getattr(transport, "reasoning_enabled", True),
            )
        else:
            operator = QAGenerationOperator(model=model, **kwargs)

        self._steps.append(operator)
        return self

    def score(self) -> "RetrieverPipelineBuilder":
        """Append a :class:`ScoringOperator` step (Tier 1 + Tier 2)."""
        from nemo_retriever.operators.graph_ops.scoring_operator import ScoringOperator

        self._steps.append(ScoringOperator())
        return self

    def judge(
        self,
        judge: Optional[Any] = None,
        /,
        *,
        model: Optional[str] = None,
        **kwargs: Any,
    ) -> "RetrieverPipelineBuilder":
        """Append a :class:`JudgingOperator` step (Tier 3).

        Accepts either a pre-built
        :class:`~nemo_retriever.llm.clients.judge.LLMJudge` (whose transport params
        are unpacked onto the operator) or the flat ``model=...`` kwargs
        forwarded to the operator constructor.

        Raises:
            ValueError: If neither ``judge`` nor ``model`` is provided.
        """
        from nemo_retriever.tools.evaluation.judging import JudgingOperator

        if judge is None and model is None:
            raise ValueError("judge() requires either judge= or model=")

        if judge is not None:
            transport = judge.transport
            operator = JudgingOperator(
                model=transport.model,
                api_base=transport.api_base,
                api_key=transport.api_key,
                extra_params=dict(transport.extra_params) if transport.extra_params else None,
                num_retries=transport.num_retries,
                timeout=transport.timeout,
            )
        else:
            operator = JudgingOperator(model=model, **kwargs)

        self._steps.append(operator)
        return self

    def run(
        self,
        queries: Any,
        *,
        reference: Any = None,
    ) -> "pd.DataFrame":
        """Execute the composed graph on ``queries``.

        Args:
            queries: A single query string, a list of query strings, or a
                pre-built ``pandas.DataFrame`` (which must contain a
                ``query`` column and, when judging/scoring, a
                ``reference_answer`` column).
            reference: Optional ground-truth answer(s).  Accepts a single
                string (applied to all queries), a list aligned with
                ``queries``, or ``None``.  Ignored when ``queries`` is
                already a DataFrame.

        Returns:
            A ``pandas.DataFrame`` with the columns contributed by each
            appended step (always ``query``, ``context``, and
            ``context_metadata``; plus ``answer``/``latency_s``/... when
            ``.generate()`` ran, and so on).

        Raises:
            ValueError: If ``reference`` is a list whose length does not
                match ``queries``.
        """
        import pandas as pd

        from nemo_retriever.tools.evaluation.live_retrieval import LiveRetrievalOperator

        if isinstance(queries, str):
            query_list = [queries]
            df = pd.DataFrame({"query": query_list})
            if reference is not None:
                refs = reference if isinstance(reference, list) else [reference]
                if len(refs) != len(query_list):
                    raise ValueError("reference length must match queries length")
                df["reference_answer"] = refs
        elif isinstance(queries, list):
            df = pd.DataFrame({"query": list(queries)})
            if reference is not None:
                refs = reference if isinstance(reference, list) else [reference] * len(queries)
                if len(refs) != len(queries):
                    raise ValueError("reference length must match queries length")
                df["reference_answer"] = refs
        elif isinstance(queries, pd.DataFrame):
            df = queries.copy()
        else:
            raise TypeError("queries must be a str, list[str], or pandas.DataFrame; " f"got {type(queries).__name__}")

        retrieval_op = LiveRetrievalOperator(self._retriever, top_k=self._top_k)
        if not self._steps:
            out = retrieval_op.run(df)
        else:
            graph = retrieval_op
            for step in self._steps:
                graph = graph >> step
            # Linear live-RAG pipelines have exactly one leaf.
            leaves = graph.execute(df)
            if len(leaves) != 1:
                raise RuntimeError(f"Unexpected pipeline fan-out: got {len(leaves)} leaf outputs")
            out = leaves[0]

        # Expose the generation failure rate on ``df.attrs`` for downstream aggregators.
        if "gen_error" in out.columns and len(out) > 0:
            out.attrs["generation_failure_rate"] = float(out["gen_error"].notna().mean())

        return out

with_retrieval(*, top_k)

Source code in nemo_retriever/graph/retriever.py

def with_retrieval(self, *, top_k: int) -> "RetrieverPipelineBuilder":
    """Override the ``top_k`` used for the live retrieval source."""
    self._top_k = int(top_k)
    return self

generate(llm=None, /, *, model=None, **kwargs)

Raises:

Type	Description
ValueError	If neither llm nor model is provided.

Source code in nemo_retriever/graph/retriever.py

def generate(
    self,
    llm: Optional[Any] = None,
    /,
    *,
    model: Optional[str] = None,
    **kwargs: Any,
) -> "RetrieverPipelineBuilder":
    """Append a :class:`QAGenerationOperator` step.

    Accepts either a pre-built
    :class:`~nemo_retriever.llm.clients.LiteLLMClient` (whose transport
    and sampling params are unpacked onto the operator) or the flat
    ``model=..., api_base=..., ...`` kwargs forwarded to the operator
    constructor directly.

    Raises:
        ValueError: If neither ``llm`` nor ``model`` is provided.
    """
    from nemo_retriever.tools.evaluation.generation import QAGenerationOperator

    if llm is None and model is None:
        raise ValueError("generate() requires either llm= or model=")

    if llm is not None:
        transport = llm.transport
        sampling = llm.sampling
        operator = QAGenerationOperator(
            model=transport.model,
            api_base=transport.api_base,
            api_key=transport.api_key,
            temperature=sampling.temperature,
            top_p=sampling.top_p,
            max_tokens=sampling.max_tokens,
            extra_params=dict(transport.extra_params) if transport.extra_params else None,
            num_retries=transport.num_retries,
            timeout=transport.timeout,
            rag_system_prompt=transport.rag_system_prompt,
            rag_system_prompt_prefix=transport.rag_system_prompt_prefix,
            reasoning_enabled=getattr(transport, "reasoning_enabled", True),
        )
    else:
        operator = QAGenerationOperator(model=model, **kwargs)

    self._steps.append(operator)
    return self

score()

Source code in nemo_retriever/graph/retriever.py

def score(self) -> "RetrieverPipelineBuilder":
    """Append a :class:`ScoringOperator` step (Tier 1 + Tier 2)."""
    from nemo_retriever.operators.graph_ops.scoring_operator import ScoringOperator

    self._steps.append(ScoringOperator())
    return self

judge(judge=None, /, *, model=None, **kwargs)

Raises:

Type	Description
ValueError	If neither judge nor model is provided.

Source code in nemo_retriever/graph/retriever.py

def judge(
    self,
    judge: Optional[Any] = None,
    /,
    *,
    model: Optional[str] = None,
    **kwargs: Any,
) -> "RetrieverPipelineBuilder":
    """Append a :class:`JudgingOperator` step (Tier 3).

    Accepts either a pre-built
    :class:`~nemo_retriever.llm.clients.judge.LLMJudge` (whose transport params
    are unpacked onto the operator) or the flat ``model=...`` kwargs
    forwarded to the operator constructor.

    Raises:
        ValueError: If neither ``judge`` nor ``model`` is provided.
    """
    from nemo_retriever.tools.evaluation.judging import JudgingOperator

    if judge is None and model is None:
        raise ValueError("judge() requires either judge= or model=")

    if judge is not None:
        transport = judge.transport
        operator = JudgingOperator(
            model=transport.model,
            api_base=transport.api_base,
            api_key=transport.api_key,
            extra_params=dict(transport.extra_params) if transport.extra_params else None,
            num_retries=transport.num_retries,
            timeout=transport.timeout,
        )
    else:
        operator = JudgingOperator(model=model, **kwargs)

    self._steps.append(operator)
    return self

run(queries, *, reference=None)

Parameters:

Name	Type	Description	Default
queries	Any	A single query string, a list of query strings, or a pre-built pandas.DataFrame (which must contain a query column and, when judging/scoring, a reference_answer column).	required
reference	Any	Optional ground-truth answer(s). Accepts a single string (applied to all queries), a list aligned with queries, or None. Ignored when queries is already a DataFrame.	None

Returns:

Type	Description
'pd.DataFrame'	A pandas.DataFrame with the columns contributed by each
'pd.DataFrame'	appended step (always query, context, and
'pd.DataFrame'	context_metadata; plus answer/latency_s/... when
'pd.DataFrame'	.generate() ran, and so on).

Raises:

Type	Description
ValueError	If reference is a list whose length does not match queries.

Source code in nemo_retriever/graph/retriever.py

def run(
    self,
    queries: Any,
    *,
    reference: Any = None,
) -> "pd.DataFrame":
    """Execute the composed graph on ``queries``.

    Args:
        queries: A single query string, a list of query strings, or a
            pre-built ``pandas.DataFrame`` (which must contain a
            ``query`` column and, when judging/scoring, a
            ``reference_answer`` column).
        reference: Optional ground-truth answer(s).  Accepts a single
            string (applied to all queries), a list aligned with
            ``queries``, or ``None``.  Ignored when ``queries`` is
            already a DataFrame.

    Returns:
        A ``pandas.DataFrame`` with the columns contributed by each
        appended step (always ``query``, ``context``, and
        ``context_metadata``; plus ``answer``/``latency_s``/... when
        ``.generate()`` ran, and so on).

    Raises:
        ValueError: If ``reference`` is a list whose length does not
            match ``queries``.
    """
    import pandas as pd

    from nemo_retriever.tools.evaluation.live_retrieval import LiveRetrievalOperator

    if isinstance(queries, str):
        query_list = [queries]
        df = pd.DataFrame({"query": query_list})
        if reference is not None:
            refs = reference if isinstance(reference, list) else [reference]
            if len(refs) != len(query_list):
                raise ValueError("reference length must match queries length")
            df["reference_answer"] = refs
    elif isinstance(queries, list):
        df = pd.DataFrame({"query": list(queries)})
        if reference is not None:
            refs = reference if isinstance(reference, list) else [reference] * len(queries)
            if len(refs) != len(queries):
                raise ValueError("reference length must match queries length")
            df["reference_answer"] = refs
    elif isinstance(queries, pd.DataFrame):
        df = queries.copy()
    else:
        raise TypeError("queries must be a str, list[str], or pandas.DataFrame; " f"got {type(queries).__name__}")

    retrieval_op = LiveRetrievalOperator(self._retriever, top_k=self._top_k)
    if not self._steps:
        out = retrieval_op.run(df)
    else:
        graph = retrieval_op
        for step in self._steps:
            graph = graph >> step
        # Linear live-RAG pipelines have exactly one leaf.
        leaves = graph.execute(df)
        if len(leaves) != 1:
            raise RuntimeError(f"Unexpected pipeline fan-out: got {len(leaves)} leaf outputs")
        out = leaves[0]

    # Expose the generation failure rate on ``df.attrs`` for downstream aggregators.
    if "gen_error" in out.columns and len(out) > 0:
        out.attrs["generation_failure_rate"] = float(out["gen_error"].notna().mean())

    return out

Generation operators

Import these classes from nemo_retriever.operators.generation. Usage notes and parameter-field tables appear in One-shot text generation.

nemo_retriever.operators.generation.TextGenerationOperator

Bases: AbstractOperator, CPUOperator

Source code in nemo_retriever/operators/generation/base.py

class TextGenerationOperator(AbstractOperator, CPUOperator):
    """Base operator for one text-generation request per DataFrame row.

    Concrete operators construct an immutable :class:`TextGenerationTask` before
    calling this base. The task and client are runtime-only state; graph
    reconstruction uses only defensive constructor state. The base owns
    validation, safe bounded execution, positional ordering, and stable
    output metadata.

    ``input_columns`` maps each task-level input name to a physical DataFrame
    column. Results are tracked by row position rather than index label so
    duplicate DataFrame indices remain safe.
    """

    required_columns: ClassVar[tuple[str, ...]] = ()
    output_columns: ClassVar[tuple[str, ...]] = ()

    def __init__(
        self,
        params: TextGenerationParams,
        *,
        task: TextGenerationTask,
        input_columns: Mapping[str, str],
        output_column: str,
        latency_column: str | None = None,
        model_column: str | None = None,
        error_column: str | None = None,
        overwrite: bool = False,
        client: TextCompletionClient | None = None,
    ) -> None:
        copied_params = params.model_copy(deep=True)
        logical_columns = dict(input_columns)
        self._validate_input_mapping(logical_columns)

        resolved_latency_column = latency_column if latency_column is not None else f"{output_column}_latency_s"
        resolved_model_column = model_column if model_column is not None else f"{output_column}_model"
        resolved_error_column = error_column if error_column is not None else f"{output_column}_error"
        output_columns = (
            output_column,
            resolved_latency_column,
            resolved_model_column,
            resolved_error_column,
        )
        self._validate_output_columns(output_columns)

        super().__init__()

        self._params = copied_params
        self._input_columns = logical_columns
        self._output_column = output_column
        self._latency_column = resolved_latency_column
        self._model_column = resolved_model_column
        self._error_column = resolved_error_column
        self._latency_column_arg = latency_column
        self._model_column_arg = model_column
        self._error_column_arg = error_column
        self._overwrite = overwrite
        self._max_workers = copied_params.max_workers
        self._configured_model = copied_params.transport.model

        self.required_columns = tuple(dict.fromkeys(logical_columns.values()))
        self.output_columns = output_columns

        self._task = task
        missing_inputs = [name for name in self._task.required_inputs if name not in logical_columns]
        if missing_inputs:
            raise ValueError(f"{type(self).__name__} is missing task input mappings: {missing_inputs}")

        if client is None:
            sampling = copied_params.resolve_sampling(self._task.default_sampling)
            self._client: TextCompletionClient = self._create_client(copied_params, sampling)
        else:
            self._client = client

    def _create_client(
        self,
        params: TextGenerationParams,
        sampling: LLMInferenceParams,
    ) -> TextCompletionClient:
        """Create the default client without introducing a global registry."""
        return LiteLLMClient(transport=params.transport, sampling=sampling)

    @abstractmethod
    def _get_generation_constructor_kwargs(self) -> dict[str, Any]:
        """Return reconstructible public constructor state for this operator."""
        ...

    @classmethod
    def _contains_runtime_object(
        cls,
        value: Any,
        targets: tuple[object, ...],
        seen: set[int],
    ) -> bool:
        if any(value is target for target in targets):
            return True
        if isinstance(value, (str, bytes, bytearray, memoryview)):
            return False
        value_id = id(value)
        if value_id in seen:
            return False
        seen.add(value_id)
        if isinstance(value, BaseModel):
            return any(
                cls._contains_runtime_object(getattr(value, name), targets, seen) for name in type(value).model_fields
            )
        if isinstance(value, Mapping):
            return any(cls._contains_runtime_object(item, targets, seen) for pair in value.items() for item in pair)
        if isinstance(value, (list, tuple, set, frozenset)):
            return any(cls._contains_runtime_object(item, targets, seen) for item in value)
        return False

    def get_constructor_kwargs(self) -> dict[str, Any]:
        """Return validated graph state without capturing a live task or client."""
        kwargs = dict(self._get_generation_constructor_kwargs())
        forbidden_keys = {"client", "task", "_client", "_task"}.intersection(kwargs)
        if forbidden_keys:
            raise TypeError(
                f"{type(self).__name__} graph constructor hook returned runtime-only keys: " f"{sorted(forbidden_keys)}"
            )
        if self._contains_runtime_object(
            kwargs,
            (self._client, self._task),
            set(),
        ):
            raise TypeError(f"{type(self).__name__} graph constructor hook captured a live client or task")

        signature = inspect.signature(type(self).__init__)
        try:
            signature.bind(None, **kwargs)
        except TypeError as exc:
            raise TypeError(f"{type(self).__name__} returned invalid graph constructor kwargs: {exc}") from exc
        try:
            return deepcopy(kwargs)
        except Exception as exc:
            raise TypeError(f"{type(self).__name__} graph constructor kwargs could not be copied safely") from exc

    @staticmethod
    def _validate_input_mapping(input_columns: Mapping[str, str]) -> None:
        if not input_columns:
            raise ValueError("input_columns must contain at least one task input mapping")
        for logical_name, column_name in input_columns.items():
            if not isinstance(logical_name, str) or not logical_name:
                raise ValueError("input_columns task input names must be non-empty strings")
            if not isinstance(column_name, str) or not column_name:
                raise ValueError("input_columns DataFrame column names must be non-empty strings")

    @staticmethod
    def _validate_output_columns(output_columns: tuple[str, ...]) -> None:
        if any(not isinstance(column, str) or not column for column in output_columns):
            raise ValueError("output column names must be non-empty strings")
        if len(set(output_columns)) != len(output_columns):
            raise ValueError(f"output column names must be distinct: {list(output_columns)}")

    @staticmethod
    def _label_positions(data: pd.DataFrame, label: str) -> list[int]:
        return [int(position) for position in data.columns.get_indexer_for([label]) if position >= 0]

    def _validate_and_resolve_dataframe(
        self,
        data: Any,
    ) -> tuple[pd.DataFrame, dict[str, int]]:
        if not isinstance(data, pd.DataFrame):
            raise TypeError(f"{type(self).__name__} requires a pandas DataFrame")

        input_positions: dict[str, int] = {}
        missing: list[str] = []
        ambiguous_inputs: list[str] = []
        for logical_name, column_name in self._input_columns.items():
            positions = self._label_positions(data, column_name)
            if not positions:
                missing.append(column_name)
            elif len(positions) > 1:
                ambiguous_inputs.append(column_name)
            else:
                input_positions[logical_name] = positions[0]
        if missing:
            missing = list(dict.fromkeys(missing))
            raise ValueError(f"{type(self).__name__} requires missing columns: {missing}")
        if ambiguous_inputs:
            ambiguous_inputs = list(dict.fromkeys(ambiguous_inputs))
            raise ValueError(
                f"{type(self).__name__} mapped input columns are ambiguous because their labels "
                f"are duplicated: {ambiguous_inputs}"
            )

        if not self._overwrite:
            collisions = [column for column in self.output_columns if self._label_positions(data, column)]
            if collisions:
                raise ValueError(
                    f"{type(self).__name__} output columns already exist: {collisions}; "
                    "set overwrite=True to replace them"
                )
        else:
            ambiguous_outputs = [
                column for column in self.output_columns if len(self._label_positions(data, column)) > 1
            ]
            if ambiguous_outputs:
                raise ValueError(
                    f"{type(self).__name__} cannot overwrite ambiguous duplicate output " f"labels: {ambiguous_outputs}"
                )
        return data, input_positions

    def preprocess(self, data: Any, **kwargs: Any) -> pd.DataFrame:
        df, _ = self._validate_and_resolve_dataframe(data)
        return df

    def _execute_task(self, inputs: dict[str, Any]) -> GeneratedTextResult:
        """Execute the configured task; subclasses may adapt legacy clients."""
        return self._task.invoke(self._client, **inputs)

    def _execute_row(self, position: int, inputs: dict[str, Any]) -> tuple[int, GeneratedTextResult]:
        started_at = time.monotonic()
        try:
            result = self._execute_task(inputs)
        except GenerationTaskError as exc:
            # Keep the strict task's measured lifecycle while covering custom
            # adapters that report a shorter or zero failure duration.
            elapsed = max(exc.latency_s, time.monotonic() - started_at)
            result = self._failure_result(exc.code, elapsed)
            logger.warning("Row %d generation failed (%s)", position, exc.code)
        except Exception:
            # Unexpected adapter/client failures remain isolated by row. Raw
            # provider exception text is never persisted or logged.
            result = self._failure_result("request_error", time.monotonic() - started_at)
            logger.warning("Row %d generation failed (request_error)", position)
        return position, result

    def _failure_model(self) -> str:
        try:
            model = self._client.model
        except Exception as exc:
            # Failure reporting must remain best-effort, but a broken client
            # property should still be diagnosable. Log only the exception
            # type: provider messages may contain request data or credentials.
            exc_type = f"{type(exc).__module__}.{type(exc).__qualname__}"
            logger.debug(
                "Unable to read generation client model metadata (%s); using configured model",
                exc_type,
            )
            return self._configured_model
        return model if isinstance(model, str) and model else self._configured_model

    def _failure_result(self, error: str, latency_s: float) -> GeneratedTextResult:
        return GeneratedTextResult(
            text="",
            latency_s=latency_s,
            model=self._failure_model(),
            error=error,
        )

    def _effective_max_workers(self, row_count: int) -> int:
        """Return safe concurrency for the current runtime client."""
        try:
            supports_concurrency = getattr(self._client, "supports_concurrent_calls", False) is True
        except Exception:
            supports_concurrency = False
        configured_workers = self._max_workers if supports_concurrency else 1
        return min(configured_workers, row_count)

    def process(self, data: Any, **kwargs: Any) -> pd.DataFrame:
        df, input_positions = self._validate_and_resolve_dataframe(data)
        results: list[GeneratedTextResult | None] = [None] * len(df)

        if len(df):
            futures: dict[Future[tuple[int, GeneratedTextResult]], int] = {}
            with ThreadPoolExecutor(max_workers=self._effective_max_workers(len(df))) as pool:
                for position in range(len(df)):
                    inputs = {
                        name: df.iat[position, column_position] for name, column_position in input_positions.items()
                    }
                    future = pool.submit(self._execute_row, position, inputs)
                    futures[future] = position

                for future in as_completed(futures):
                    position = futures[future]
                    # _execute_row owns per-row failure collection. Exceptions
                    # or position mismatches here are executor/programming
                    # failures and must not be silently converted into row data.
                    result_position, result = future.result()
                    if result_position != position:
                        raise RuntimeError(
                            f"generation result position {result_position} does not match "
                            f"submitted position {position}"
                        )
                    results[position] = result

        # Every non-empty row is assigned either a task result or a failure
        # result above. The cast-free local assertion catches future changes to
        # that invariant before partially writing output columns.
        if any(result is None for result in results):
            raise RuntimeError("generation completed without a result for every row")
        completed_results = [result for result in results if result is not None]

        out = df.copy()
        # Explicit dtypes keep empty and non-empty Ray/Arrow batches
        # schema-compatible while retaining positional duplicate-index writes.
        out[self._output_column] = pd.array([result.text for result in completed_results], dtype="object")
        out[self._latency_column] = pd.array([result.latency_s for result in completed_results], dtype="float64")
        out[self._model_column] = pd.array([result.model for result in completed_results], dtype="object")
        out[self._error_column] = pd.array([result.error for result in completed_results], dtype="object")
        return out

    def postprocess(self, data: Any, **kwargs: Any) -> Any:
        return data

required_columns = tuple(dict.fromkeys(logical_columns.values())) class-attribute instance-attribute

output_columns = output_columns class-attribute instance-attribute

get_constructor_kwargs()

Source code in nemo_retriever/operators/generation/base.py

def get_constructor_kwargs(self) -> dict[str, Any]:
    """Return validated graph state without capturing a live task or client."""
    kwargs = dict(self._get_generation_constructor_kwargs())
    forbidden_keys = {"client", "task", "_client", "_task"}.intersection(kwargs)
    if forbidden_keys:
        raise TypeError(
            f"{type(self).__name__} graph constructor hook returned runtime-only keys: " f"{sorted(forbidden_keys)}"
        )
    if self._contains_runtime_object(
        kwargs,
        (self._client, self._task),
        set(),
    ):
        raise TypeError(f"{type(self).__name__} graph constructor hook captured a live client or task")

    signature = inspect.signature(type(self).__init__)
    try:
        signature.bind(None, **kwargs)
    except TypeError as exc:
        raise TypeError(f"{type(self).__name__} returned invalid graph constructor kwargs: {exc}") from exc
    try:
        return deepcopy(kwargs)
    except Exception as exc:
        raise TypeError(f"{type(self).__name__} graph constructor kwargs could not be copied safely") from exc

preprocess(data, **kwargs)

Source code in nemo_retriever/operators/generation/base.py

def preprocess(self, data: Any, **kwargs: Any) -> pd.DataFrame:
    df, _ = self._validate_and_resolve_dataframe(data)
    return df

process(data, **kwargs)

Source code in nemo_retriever/operators/generation/base.py

def process(self, data: Any, **kwargs: Any) -> pd.DataFrame:
    df, input_positions = self._validate_and_resolve_dataframe(data)
    results: list[GeneratedTextResult | None] = [None] * len(df)

    if len(df):
        futures: dict[Future[tuple[int, GeneratedTextResult]], int] = {}
        with ThreadPoolExecutor(max_workers=self._effective_max_workers(len(df))) as pool:
            for position in range(len(df)):
                inputs = {
                    name: df.iat[position, column_position] for name, column_position in input_positions.items()
                }
                future = pool.submit(self._execute_row, position, inputs)
                futures[future] = position

            for future in as_completed(futures):
                position = futures[future]
                # _execute_row owns per-row failure collection. Exceptions
                # or position mismatches here are executor/programming
                # failures and must not be silently converted into row data.
                result_position, result = future.result()
                if result_position != position:
                    raise RuntimeError(
                        f"generation result position {result_position} does not match "
                        f"submitted position {position}"
                    )
                results[position] = result

    # Every non-empty row is assigned either a task result or a failure
    # result above. The cast-free local assertion catches future changes to
    # that invariant before partially writing output columns.
    if any(result is None for result in results):
        raise RuntimeError("generation completed without a result for every row")
    completed_results = [result for result in results if result is not None]

    out = df.copy()
    # Explicit dtypes keep empty and non-empty Ray/Arrow batches
    # schema-compatible while retaining positional duplicate-index writes.
    out[self._output_column] = pd.array([result.text for result in completed_results], dtype="object")
    out[self._latency_column] = pd.array([result.latency_s for result in completed_results], dtype="float64")
    out[self._model_column] = pd.array([result.model for result in completed_results], dtype="object")
    out[self._error_column] = pd.array([result.error for result in completed_results], dtype="object")
    return out

postprocess(data, **kwargs)

Source code in nemo_retriever/operators/generation/base.py

def postprocess(self, data: Any, **kwargs: Any) -> Any:
    return data

nemo_retriever.operators.generation.GenericGenerationOperator

Bases: TextGenerationOperator

Source code in nemo_retriever/operators/generation/generic.py

class GenericGenerationOperator(TextGenerationOperator):
    """Generate text from a validated prompt template and mapped row inputs."""

    def __init__(
        self,
        params: TextGenerationParams,
        input_columns: Mapping[str, str],
        output_column: str = "generated_text",
        *,
        latency_column: str | None = None,
        model_column: str | None = None,
        error_column: str | None = None,
        overwrite: bool = False,
        client: TextCompletionClient | None = None,
    ) -> None:
        normalized_input_columns = dict(input_columns)
        if params.prompt is None:
            raise ValueError("GenericGenerationOperator requires params.prompt")
        reasoning_enabled = (
            params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled
        )
        task = GenericPromptTask(
            prompt=params.prompt,
            input_names=tuple(normalized_input_columns),
            system_prompt=params.system_prompt,
            reasoning_enabled=reasoning_enabled,
        )
        super().__init__(
            params,
            task=task,
            input_columns=normalized_input_columns,
            output_column=output_column,
            latency_column=latency_column,
            model_column=model_column,
            error_column=error_column,
            overwrite=overwrite,
            client=client,
        )

    def _get_generation_constructor_kwargs(self) -> dict[str, object]:
        return {
            "params": self._params.model_copy(deep=True),
            "input_columns": self._input_columns.copy(),
            "output_column": self._output_column,
            "latency_column": self._latency_column_arg,
            "model_column": self._model_column_arg,
            "error_column": self._error_column_arg,
            "overwrite": self._overwrite,
        }

nemo_retriever.operators.generation.SummarizationOperator

Bases: TextGenerationOperator

Source code in nemo_retriever/operators/generation/summarization.py

class SummarizationOperator(TextGenerationOperator):
    """Summarize the text in one DataFrame column per row."""

    def __init__(
        self,
        params: TextGenerationParams,
        input_column: str = "text",
        output_column: str = "summary",
        *,
        latency_column: str | None = None,
        model_column: str | None = None,
        error_column: str | None = None,
        overwrite: bool = False,
        client: TextCompletionClient | None = None,
    ) -> None:
        reasoning_enabled = (
            params.reasoning_enabled if params.reasoning_enabled is not None else params.transport.reasoning_enabled
        )
        task = SummarizeTask(
            prompt=params.prompt,
            system_prompt=params.system_prompt,
            reasoning_enabled=reasoning_enabled,
        )
        super().__init__(
            params,
            task=task,
            input_columns={"text": input_column},
            output_column=output_column,
            latency_column=latency_column,
            model_column=model_column,
            error_column=error_column,
            overwrite=overwrite,
            client=client,
        )

    def _get_generation_constructor_kwargs(self) -> dict[str, object]:
        return {
            "params": self._params.model_copy(deep=True),
            "input_column": self._input_columns["text"],
            "output_column": self._output_column,
            "latency_column": self._latency_column_arg,
            "model_column": self._model_column_arg,
            "error_column": self._error_column_arg,
            "overwrite": self._overwrite,
        }

LLM clients, tasks, and results

Import these names from nemo_retriever.models.llm. That package path is the supported public import. LiteLLMClient and LLMJudge load from implementation submodules. Import those classes from nemo_retriever.models.llm. Use from_kwargs() for a flat constructor.

nemo_retriever.models.llm

AnswerJudge

Bases: Protocol

Source code in nemo_retriever/models/llm/types.py

@runtime_checkable
class AnswerJudge(Protocol):
    """Pluggable answer scoring interface."""

    def judge(self, query: str, reference: str, candidate: str) -> "JudgeResult": ...

judge(query, reference, candidate)

Source code in nemo_retriever/models/llm/types.py

def judge(self, query: str, reference: str, candidate: str) -> "JudgeResult": ...

AnswerRequest

Bases: BaseModel

Source code in nemo_retriever/models/llm/types.py

class AnswerRequest(BaseModel):
    """Shared internal request model for answer generation.

    Public callers may continue using ergonomic keyword arguments on
    ``Retriever.answer``. Service and local code normalize into this model so
    query, retrieval, and per-call generation controls stay aligned.
    """

    model_config = ConfigDict(extra="forbid")

    query: str
    top_k: int = Field(default=5, ge=1)
    reasoning_enabled: Optional[bool] = None
    reference: Optional[str] = None
    judge_enabled: bool = False

model_config = ConfigDict(extra='forbid') class-attribute instance-attribute

query instance-attribute

top_k = Field(default=5, ge=1) class-attribute instance-attribute

reasoning_enabled = None class-attribute instance-attribute

reference = None class-attribute instance-attribute

judge_enabled = False class-attribute instance-attribute

AnswerResult

Bases: BaseModel

Attributes:

Name	Type	Description
query	str	The question that was answered.
answer	str	The generated answer text.
chunks	Optional[list[str]]	Retrieved chunk texts used as context, in rank order.
metadata	Optional[list[dict[str, Any]]]	Per-chunk metadata (source, page_number, etc.), aligned with chunks.
model	str	Model identifier that produced answer.
latency_s	float	Wall-clock latency of the generation call in seconds.
chunk_count	int	Number of retrieved chunks used for generation.
error	Optional[str]	Non-None when generation failed. Scoring and judge are skipped when error is set.
judge_score	Optional[float]	ragas AnswerAccuracy Tier-3 score (0.0-1.0) when a judge was run.
judge_reasoning	Optional[str]	Empty -- AnswerAccuracy emits only a numeric rating.
judge_error	Optional[str]	Non-None when the judge call failed.
token_f1	Optional[float]	Tier-2 token-level F1 between answer and the reference answer (0.0-1.0).
exact_match	Optional[bool]	Tier-2 normalised exact-match flag.
answer_in_context	Optional[bool]	Tier-1 flag -- True if at least half of the reference answer's content words appear in the retrieved chunks.
failure_mode	Optional[str]	Classification produced by :func:~nemo_retriever.evaluation.scoring.classify_failure.

Source code in nemo_retriever/models/llm/types.py

class AnswerResult(BaseModel):
    """Result from a single live-RAG call to ``Retriever.answer``.

    Holds the generated answer alongside the retrieved context that was used
    to produce it and -- when a ``reference`` answer and/or ``judge`` are
    supplied -- the Tier-1 / Tier-2 / Tier-3 scoring artefacts produced by
    :mod:`nemo_retriever.evaluation.scoring` and
    :class:`~nemo_retriever.models.llm.clients.judge.LLMJudge`.

    Attributes:
        query: The question that was answered.
        answer: The generated answer text.
        chunks: Retrieved chunk texts used as context, in rank order.
        metadata: Per-chunk metadata (source, page_number, etc.), aligned
            with ``chunks``.
        model: Model identifier that produced ``answer``.
        latency_s: Wall-clock latency of the generation call in seconds.
        chunk_count: Number of retrieved chunks used for generation.
        error: Non-None when generation failed. Scoring and judge are
            skipped when ``error`` is set.
        judge_score: ragas AnswerAccuracy Tier-3 score (0.0-1.0) when a
            judge was run.
        judge_reasoning: Empty -- AnswerAccuracy emits only a numeric rating.
        judge_error: Non-None when the judge call failed.
        token_f1: Tier-2 token-level F1 between ``answer`` and the
            reference answer (0.0-1.0).
        exact_match: Tier-2 normalised exact-match flag.
        answer_in_context: Tier-1 flag -- True if at least half of the
            reference answer's content words appear in the retrieved chunks.
        failure_mode: Classification produced by
            :func:`~nemo_retriever.evaluation.scoring.classify_failure`.
    """

    model_config = ConfigDict(extra="forbid")

    query: str
    answer: str
    model: str
    latency_s: float
    chunk_count: int
    chunks: Optional[list[str]] = None
    metadata: Optional[list[dict[str, Any]]] = None
    error: Optional[str] = None
    judge_score: Optional[float] = None
    judge_reasoning: Optional[str] = None
    judge_error: Optional[str] = None
    token_f1: Optional[float] = None
    exact_match: Optional[bool] = None
    answer_in_context: Optional[bool] = None
    failure_mode: Optional[str] = None

model_config = ConfigDict(extra='forbid') class-attribute instance-attribute

query instance-attribute

answer instance-attribute

model instance-attribute

latency_s instance-attribute

chunk_count instance-attribute

chunks = None class-attribute instance-attribute

metadata = None class-attribute instance-attribute

error = None class-attribute instance-attribute

judge_score = None class-attribute instance-attribute

judge_reasoning = None class-attribute instance-attribute

judge_error = None class-attribute instance-attribute

token_f1 = None class-attribute instance-attribute

exact_match = None class-attribute instance-attribute

answer_in_context = None class-attribute instance-attribute

failure_mode = None class-attribute instance-attribute

GeneratedTextResult dataclass

Source code in nemo_retriever/models/llm/types.py

@dataclass(frozen=True)
class GeneratedTextResult:
    """Task-neutral result from a single text-generation request."""

    text: str
    latency_s: float
    model: str
    error: Optional[str] = None

text instance-attribute

latency_s instance-attribute

model instance-attribute

error = None class-attribute instance-attribute

GenerationRequest dataclass

Source code in nemo_retriever/models/llm/types.py

@dataclass(frozen=True)
class GenerationRequest:
    """One text-only request produced by a generation task.

    Tools, streaming, multiple choices, and structured domain results are not
    supported by this provisional contract.
    """

    messages: list[dict[str, Any]]
    max_tokens: Optional[int] = None
    extra_params: Optional[dict[str, Any]] = None

    def __post_init__(self) -> None:
        """Snapshot mutable inputs and reject non-text or protected state."""
        # Keep this types module lightweight and avoid a package import cycle.
        from nemo_retriever.common.params.models import validate_llm_extra_params

        messages = deepcopy(self.messages)
        extra_params = deepcopy(self.extra_params)
        if not isinstance(messages, list) or not all(isinstance(message, dict) for message in messages):
            raise TypeError("GenerationRequest.messages must be a list of message dictionaries")
        for message in messages:
            if not isinstance(message.get("role"), str) or not isinstance(message.get("content"), str):
                raise TypeError("GenerationRequest messages require string role and content fields")
            if {"tool_calls", "function_call", "tool_call_id"}.intersection(message):
                raise ValueError("GenerationRequest does not support tool messages or tool calls")
        validate_llm_extra_params(extra_params or {}, source="GenerationRequest.extra_params")
        object.__setattr__(self, "messages", messages)
        object.__setattr__(self, "extra_params", extra_params)

messages instance-attribute

max_tokens = None class-attribute instance-attribute

extra_params = None class-attribute instance-attribute

GenerationResult dataclass

Source code in nemo_retriever/models/llm/types.py

@dataclass
class GenerationResult:
    """Result from a single LLM generation call."""

    answer: str
    latency_s: float
    model: str
    error: Optional[str] = None

answer instance-attribute

latency_s instance-attribute

model instance-attribute

error = None class-attribute instance-attribute

GenerationTaskError

Bases: RuntimeError

Source code in nemo_retriever/models/llm/tasks/base.py

class GenerationTaskError(RuntimeError):
    """Sanitized failure raised by the strict generation-task lifecycle."""

    def __init__(
        self,
        *,
        code: str,
        phase: Literal["request", "transport", "response", "parse"],
        retryable: bool,
        public_message: str,
        latency_s: float,
    ) -> None:
        super().__init__(public_message)
        self.code = code
        self.phase = phase
        self.retryable = retryable
        self.public_message = public_message
        self.latency_s = latency_s

code = code instance-attribute

phase = phase instance-attribute

retryable = retryable instance-attribute

public_message = public_message instance-attribute

latency_s = latency_s instance-attribute

GenericPromptTask dataclass

Bases: TextGenerationTask

Source code in nemo_retriever/models/llm/tasks/generic.py

@dataclass(frozen=True, init=False)
class GenericPromptTask(TextGenerationTask):
    """Render declared row inputs into a validated prompt template."""

    prompt: str
    required_inputs: tuple[str, ...]
    system_prompt: Optional[str]
    reasoning_enabled: Optional[bool]

    _default_sampling: ClassVar[dict[str, Any]] = {
        "temperature": 1.0,
        "top_p": None,
        "max_tokens": 1024,
    }

    def __init__(
        self,
        *,
        prompt: str,
        input_names: Sequence[str],
        system_prompt: Optional[str] = None,
        reasoning_enabled: Optional[bool] = None,
    ) -> None:
        if isinstance(input_names, str):
            raise TypeError("input_names must be a sequence of names, not a string")
        names = tuple(input_names)
        _validate_prompt_template(prompt, names)
        object.__setattr__(self, "prompt", prompt)
        object.__setattr__(self, "required_inputs", names)
        object.__setattr__(self, "system_prompt", system_prompt)
        object.__setattr__(self, "reasoning_enabled", reasoning_enabled)

    def build_request(self, **inputs: object) -> GenerationRequest:
        """Render declared inputs and build one completion request."""
        missing = [name for name in self.required_inputs if name not in inputs]
        if missing:
            raise KeyError(f"missing required inputs: {missing}")
        values = {name: inputs[name] for name in self.required_inputs}
        user_content = self.prompt.format(**values)
        messages: list[dict[str, Any]] = []
        if self.system_prompt is not None:
            messages.append({"role": "system", "content": self.system_prompt})
        messages.append({"role": "user", "content": user_content})
        messages, extra_params = _apply_reasoning_control(messages, self.reasoning_enabled)
        return GenerationRequest(messages=messages, extra_params=extra_params)

    def parse(self, raw_text: str) -> str:
        """Remove visible model reasoning from generated text."""
        return strip_think_tags(raw_text)

prompt instance-attribute

required_inputs instance-attribute

system_prompt instance-attribute

reasoning_enabled instance-attribute

build_request(**inputs)

Source code in nemo_retriever/models/llm/tasks/generic.py

def build_request(self, **inputs: object) -> GenerationRequest:
    """Render declared inputs and build one completion request."""
    missing = [name for name in self.required_inputs if name not in inputs]
    if missing:
        raise KeyError(f"missing required inputs: {missing}")
    values = {name: inputs[name] for name in self.required_inputs}
    user_content = self.prompt.format(**values)
    messages: list[dict[str, Any]] = []
    if self.system_prompt is not None:
        messages.append({"role": "system", "content": self.system_prompt})
    messages.append({"role": "user", "content": user_content})
    messages, extra_params = _apply_reasoning_control(messages, self.reasoning_enabled)
    return GenerationRequest(messages=messages, extra_params=extra_params)

parse(raw_text)

Source code in nemo_retriever/models/llm/tasks/generic.py

def parse(self, raw_text: str) -> str:
    """Remove visible model reasoning from generated text."""
    return strip_think_tags(raw_text)

JudgeResult dataclass

Source code in nemo_retriever/models/llm/types.py

@dataclass
class JudgeResult:
    """Result from a single judge evaluation.

    ``score`` is ``None`` when the judge could not produce a score
    (API error, empty candidate, or no valid rating). Valid scores are
    ragas ``AnswerAccuracy`` values on a ``0.0-1.0`` scale (higher is
    better). ``reasoning`` is empty -- ``AnswerAccuracy`` emits only a
    numeric rating.
    """

    score: Optional[float] = None
    reasoning: str = ""
    error: Optional[str] = None

score = None class-attribute instance-attribute

reasoning = '' class-attribute instance-attribute

error = None class-attribute instance-attribute

LLMClient

Bases: Protocol

Source code in nemo_retriever/models/llm/types.py

@runtime_checkable
class LLMClient(Protocol):
    """Pluggable LLM answer generation interface."""

    def generate(
        self,
        query: str,
        chunks: list[str],
        *,
        reasoning_enabled: Optional[bool] = None,
    ) -> "GenerationResult": ...

generate(query, chunks, *, reasoning_enabled=None)

Source code in nemo_retriever/models/llm/types.py

def generate(
    self,
    query: str,
    chunks: list[str],
    *,
    reasoning_enabled: Optional[bool] = None,
) -> "GenerationResult": ...

RagAnswerTask dataclass

Bases: TextGenerationTask

Source code in nemo_retriever/models/llm/tasks/rag_answer.py

@dataclass(frozen=True)
class RagAnswerTask(TextGenerationTask):
    """Generate a grounded answer from a query and retrieved text chunks."""

    prompt: Optional[str] = None
    system_prompt: Optional[str] = None
    system_prompt_prefix: Optional[str] = None
    reasoning_enabled: Optional[bool] = None

    required_inputs: ClassVar[tuple[str, ...]] = ("query", "chunks")
    _default_sampling: ClassVar[dict[str, Any]] = {
        "temperature": 0.0,
        "top_p": None,
        "max_tokens": 4096,
    }
    empty_output_error: ClassVar[str] = "thinking_truncated"

    def __post_init__(self) -> None:
        if self.prompt is not None:
            _validate_rag_prompt(self.prompt)

    def build_request(self, **inputs: object) -> GenerationRequest:
        """Build a grounded answer request, including optional reasoning controls."""
        query = inputs["query"]
        chunks = inputs["chunks"]
        if not isinstance(query, str):
            raise TypeError("query must be a string")
        if not isinstance(chunks, list) or not all(isinstance(chunk, str) for chunk in chunks):
            raise TypeError("chunks must be a list of strings")

        formatted_system_prompt = _format_rag_system_prompt(
            rag_system_prompt=self.system_prompt,
            rag_system_prompt_prefix=self.system_prompt_prefix,
        )
        messages = _build_rag_prompt(
            query,
            chunks,
            formatted_rag_system_prompt=formatted_system_prompt,
        )
        if self.prompt is not None:
            context = "\n\n---\n\n".join(chunks) if chunks else "(no context retrieved)"
            messages[-1]["content"] = self.prompt.format(context=context, query=query)

        per_request_reasoning = inputs.get("reasoning_enabled")
        effective_reasoning = self.reasoning_enabled if per_request_reasoning is None else bool(per_request_reasoning)
        messages, extra_params = _apply_reasoning_control(messages, effective_reasoning)
        return GenerationRequest(messages=messages, extra_params=extra_params)

    def parse(self, raw_text: str) -> str:
        """Remove visible model reasoning from the answer."""
        return strip_think_tags(raw_text)

prompt = None class-attribute instance-attribute

system_prompt = None class-attribute instance-attribute

system_prompt_prefix = None class-attribute instance-attribute

reasoning_enabled = None class-attribute instance-attribute

required_inputs = ('query', 'chunks') class-attribute

empty_output_error = 'thinking_truncated' class-attribute

build_request(**inputs)

Source code in nemo_retriever/models/llm/tasks/rag_answer.py

def build_request(self, **inputs: object) -> GenerationRequest:
    """Build a grounded answer request, including optional reasoning controls."""
    query = inputs["query"]
    chunks = inputs["chunks"]
    if not isinstance(query, str):
        raise TypeError("query must be a string")
    if not isinstance(chunks, list) or not all(isinstance(chunk, str) for chunk in chunks):
        raise TypeError("chunks must be a list of strings")

    formatted_system_prompt = _format_rag_system_prompt(
        rag_system_prompt=self.system_prompt,
        rag_system_prompt_prefix=self.system_prompt_prefix,
    )
    messages = _build_rag_prompt(
        query,
        chunks,
        formatted_rag_system_prompt=formatted_system_prompt,
    )
    if self.prompt is not None:
        context = "\n\n---\n\n".join(chunks) if chunks else "(no context retrieved)"
        messages[-1]["content"] = self.prompt.format(context=context, query=query)

    per_request_reasoning = inputs.get("reasoning_enabled")
    effective_reasoning = self.reasoning_enabled if per_request_reasoning is None else bool(per_request_reasoning)
    messages, extra_params = _apply_reasoning_control(messages, effective_reasoning)
    return GenerationRequest(messages=messages, extra_params=extra_params)

parse(raw_text)

Source code in nemo_retriever/models/llm/tasks/rag_answer.py

def parse(self, raw_text: str) -> str:
    """Remove visible model reasoning from the answer."""
    return strip_think_tags(raw_text)

RetrievalResult dataclass

Source code in nemo_retriever/models/llm/types.py

@dataclass
class RetrievalResult:
    """Result from a retrieval operation."""

    chunks: list[str]
    metadata: list[dict[str, Any]] = field(default_factory=list)

chunks instance-attribute

metadata = field(default_factory=list) class-attribute instance-attribute

RetrieverStrategy

Bases: Protocol

Source code in nemo_retriever/models/llm/types.py

@runtime_checkable
class RetrieverStrategy(Protocol):
    """Pluggable retrieval strategy interface."""

    def retrieve(self, query: str, top_k: int) -> "RetrievalResult": ...

retrieve(query, top_k)

Source code in nemo_retriever/models/llm/types.py

def retrieve(self, query: str, top_k: int) -> "RetrievalResult": ...

SummarizeTask dataclass

Bases: TextGenerationTask

Source code in nemo_retriever/models/llm/tasks/summarize.py

@dataclass(frozen=True)
class SummarizeTask(TextGenerationTask):
    """Summarize one text value without truncation or hidden map-reduce."""

    prompt: Optional[str] = None
    system_prompt: Optional[str] = None
    reasoning_enabled: Optional[bool] = None

    required_inputs: ClassVar[tuple[str, ...]] = ("text",)
    _default_sampling: ClassVar[dict[str, Any]] = {
        "temperature": 0.0,
        "top_p": None,
        "max_tokens": 1024,
    }

    def __post_init__(self) -> None:
        if self.prompt is not None:
            _summary_prompt_fields(self.prompt)

    def _preflight_error(self, **inputs: object) -> Optional[str]:
        text = inputs.get("text")
        if isinstance(text, str) and not text.strip():
            return "empty_input"
        return None

    def build_request(self, **inputs: object) -> GenerationRequest:
        """Build one faithful-summary request for the supplied text."""
        text = inputs["text"]
        if not isinstance(text, str):
            raise TypeError("text must be a string")

        prompt = self.prompt if self.prompt is not None else _SUMMARIZE_USER_TEMPLATE
        fields = _summary_prompt_fields(prompt)
        user_content = prompt.format(text=text) if fields else f"{prompt}\n\n{text}"
        system_content = self.system_prompt if self.system_prompt is not None else _SUMMARIZE_SYSTEM_PROMPT
        messages = [
            {"role": "system", "content": system_content},
            {"role": "user", "content": user_content},
        ]
        messages, extra_params = _apply_reasoning_control(messages, self.reasoning_enabled)
        return GenerationRequest(messages=messages, extra_params=extra_params)

    def parse(self, raw_text: str) -> str:
        """Remove visible model reasoning from the summary."""
        return strip_think_tags(raw_text)

prompt = None class-attribute instance-attribute

system_prompt = None class-attribute instance-attribute

reasoning_enabled = None class-attribute instance-attribute

required_inputs = ('text',) class-attribute

build_request(**inputs)

Source code in nemo_retriever/models/llm/tasks/summarize.py

def build_request(self, **inputs: object) -> GenerationRequest:
    """Build one faithful-summary request for the supplied text."""
    text = inputs["text"]
    if not isinstance(text, str):
        raise TypeError("text must be a string")

    prompt = self.prompt if self.prompt is not None else _SUMMARIZE_USER_TEMPLATE
    fields = _summary_prompt_fields(prompt)
    user_content = prompt.format(text=text) if fields else f"{prompt}\n\n{text}"
    system_content = self.system_prompt if self.system_prompt is not None else _SUMMARIZE_SYSTEM_PROMPT
    messages = [
        {"role": "system", "content": system_content},
        {"role": "user", "content": user_content},
    ]
    messages, extra_params = _apply_reasoning_control(messages, self.reasoning_enabled)
    return GenerationRequest(messages=messages, extra_params=extra_params)

parse(raw_text)

Source code in nemo_retriever/models/llm/tasks/summarize.py

def parse(self, raw_text: str) -> str:
    """Remove visible model reasoning from the summary."""
    return strip_think_tags(raw_text)

TextCompletionClient

Bases: Protocol

Source code in nemo_retriever/models/llm/types.py

@runtime_checkable
class TextCompletionClient(Protocol):
    """Provisional synchronous, thread-safe, single-turn text client contract.

    Implementations return exactly one text completion. Tools, streaming,
    multiple choices, and structured domain responses are intentionally
    outside this contract.
    """

    @property
    def model(self) -> str:
        """Return the model identifier used for generated results."""
        ...

    def complete(
        self,
        messages: list[dict[str, Any]],
        max_tokens: Optional[int] = None,
        extra_params: Optional[dict[str, Any]] = None,
    ) -> tuple[str, float]:
        """Return generated text and wall-clock latency in seconds."""
        ...

model property

complete(messages, max_tokens=None, extra_params=None)

Source code in nemo_retriever/models/llm/types.py

def complete(
    self,
    messages: list[dict[str, Any]],
    max_tokens: Optional[int] = None,
    extra_params: Optional[dict[str, Any]] = None,
) -> tuple[str, float]:
    """Return generated text and wall-clock latency in seconds."""
    ...

TextGenerationTask

Bases: ABC

Source code in nemo_retriever/models/llm/tasks/base.py

class TextGenerationTask(ABC):
    """Stateless strategy that turns logical inputs into one completion call."""

    required_inputs: tuple[str, ...] = ()
    _default_sampling: ClassVar[dict[str, Any]] = {
        "temperature": 1.0,
        "top_p": None,
        "max_tokens": 1024,
    }
    empty_output_error: ClassVar[str] = "empty_output"

    @property
    def default_sampling(self) -> LLMInferenceParams:
        """Return a fresh copy of this task's sampling defaults."""
        return LLMInferenceParams(**self._default_sampling)

    @abstractmethod
    def build_request(self, **inputs: object) -> GenerationRequest:
        """Build one provider-neutral request from logical task inputs."""

    def parse(self, raw_text: str) -> str:
        """Parse completion text into the task's text result."""
        return raw_text.strip()

    def _preflight_error(self, **inputs: object) -> Optional[str]:
        """Return an error code when no provider request should be made."""
        return None

    @staticmethod
    def _client_model(client: object) -> str:
        """Read a client model identifier without making error handling fail."""
        try:
            model = getattr(client, "model", "")
        except Exception:
            return ""
        return model if isinstance(model, str) else ""

    @staticmethod
    def _elapsed(started_at: float) -> float:
        return time.monotonic() - started_at

    def invoke(self, client: TextCompletionClient, **inputs: object) -> GeneratedTextResult:
        """Strictly build, execute, and parse one text request.

        Failures are raised as :class:`GenerationTaskError` with stable codes
        and sanitized messages. Batch operators collect those errors at the row
        boundary; callers that need the historical collecting behavior can use
        :meth:`execute`.
        """
        started_at = time.monotonic()

        preflight_failure: Optional[GenerationTaskError] = None
        try:
            preflight_error = self._preflight_error(**inputs)
        except Exception:
            preflight_failure = GenerationTaskError(
                code="request_error",
                phase="request",
                retryable=False,
                public_message="generation request validation failed",
                latency_s=self._elapsed(started_at),
            )
        if preflight_failure is not None:
            raise preflight_failure
        if preflight_error is not None:
            raise GenerationTaskError(
                code=preflight_error,
                phase="request",
                retryable=False,
                public_message="generation request was skipped",
                latency_s=self._elapsed(started_at),
            )

        request_failure: Optional[GenerationTaskError] = None
        try:
            request = self.build_request(**inputs)
            if not isinstance(request, GenerationRequest):
                raise TypeError("build_request must return GenerationRequest")
            request = GenerationRequest(
                messages=request.messages,
                max_tokens=request.max_tokens,
                extra_params=request.extra_params,
            )
        except Exception:
            request_failure = GenerationTaskError(
                code="request_error",
                phase="request",
                retryable=False,
                public_message="generation request construction failed",
                latency_s=self._elapsed(started_at),
            )
        if request_failure is not None:
            raise request_failure

        transport_failure: Optional[GenerationTaskError] = None
        try:
            raw_text, latency_s = client.complete(
                request.messages,
                max_tokens=request.max_tokens,
                extra_params=request.extra_params,
            )
        except UnsupportedTextResponseError:
            transport_failure = GenerationTaskError(
                code="unsupported_response",
                phase="response",
                retryable=False,
                public_message="provider response is not representable as a text completion",
                latency_s=self._elapsed(started_at),
            )
        except Exception as exc:
            try:
                retryable = bool(getattr(exc, "retryable", False))
            except Exception:
                retryable = False
            transport_failure = GenerationTaskError(
                code="transport_error",
                phase="transport",
                retryable=retryable,
                public_message="text completion request failed",
                latency_s=self._elapsed(started_at),
            )
        if transport_failure is not None:
            raise transport_failure

        parse_failure: Optional[GenerationTaskError] = None
        try:
            text = self.parse(raw_text)
            if not isinstance(text, str):
                raise TypeError("parse must return text")
        except Exception:
            parse_failure = GenerationTaskError(
                code="parse_error",
                phase="parse",
                retryable=False,
                public_message="generation response parsing failed",
                latency_s=self._elapsed(started_at),
            )
        if parse_failure is not None:
            raise parse_failure
        if not text:
            raise GenerationTaskError(
                code=self.empty_output_error,
                phase="response",
                retryable=False,
                public_message="generation produced no usable text",
                latency_s=latency_s,
            )
        return GeneratedTextResult(
            text=text,
            latency_s=latency_s,
            model=self._client_model(client),
            error=None,
        )

    def execute(self, client: TextCompletionClient, **inputs: object) -> GeneratedTextResult:
        """Compatibility facade that collects strict failures into a result."""
        try:
            return self.invoke(client, **inputs)
        except GenerationTaskError as exc:
            return GeneratedTextResult(
                text="",
                latency_s=exc.latency_s,
                model=self._client_model(client),
                error=exc.code,
            )

required_inputs = () class-attribute instance-attribute

empty_output_error = 'empty_output' class-attribute

default_sampling property

build_request(**inputs) abstractmethod

Source code in nemo_retriever/models/llm/tasks/base.py

@abstractmethod
def build_request(self, **inputs: object) -> GenerationRequest:
    """Build one provider-neutral request from logical task inputs."""

parse(raw_text)

Source code in nemo_retriever/models/llm/tasks/base.py

def parse(self, raw_text: str) -> str:
    """Parse completion text into the task's text result."""
    return raw_text.strip()

invoke(client, **inputs)

Source code in nemo_retriever/models/llm/tasks/base.py

def invoke(self, client: TextCompletionClient, **inputs: object) -> GeneratedTextResult:
    """Strictly build, execute, and parse one text request.

    Failures are raised as :class:`GenerationTaskError` with stable codes
    and sanitized messages. Batch operators collect those errors at the row
    boundary; callers that need the historical collecting behavior can use
    :meth:`execute`.
    """
    started_at = time.monotonic()

    preflight_failure: Optional[GenerationTaskError] = None
    try:
        preflight_error = self._preflight_error(**inputs)
    except Exception:
        preflight_failure = GenerationTaskError(
            code="request_error",
            phase="request",
            retryable=False,
            public_message="generation request validation failed",
            latency_s=self._elapsed(started_at),
        )
    if preflight_failure is not None:
        raise preflight_failure
    if preflight_error is not None:
        raise GenerationTaskError(
            code=preflight_error,
            phase="request",
            retryable=False,
            public_message="generation request was skipped",
            latency_s=self._elapsed(started_at),
        )

    request_failure: Optional[GenerationTaskError] = None
    try:
        request = self.build_request(**inputs)
        if not isinstance(request, GenerationRequest):
            raise TypeError("build_request must return GenerationRequest")
        request = GenerationRequest(
            messages=request.messages,
            max_tokens=request.max_tokens,
            extra_params=request.extra_params,
        )
    except Exception:
        request_failure = GenerationTaskError(
            code="request_error",
            phase="request",
            retryable=False,
            public_message="generation request construction failed",
            latency_s=self._elapsed(started_at),
        )
    if request_failure is not None:
        raise request_failure

    transport_failure: Optional[GenerationTaskError] = None
    try:
        raw_text, latency_s = client.complete(
            request.messages,
            max_tokens=request.max_tokens,
            extra_params=request.extra_params,
        )
    except UnsupportedTextResponseError:
        transport_failure = GenerationTaskError(
            code="unsupported_response",
            phase="response",
            retryable=False,
            public_message="provider response is not representable as a text completion",
            latency_s=self._elapsed(started_at),
        )
    except Exception as exc:
        try:
            retryable = bool(getattr(exc, "retryable", False))
        except Exception:
            retryable = False
        transport_failure = GenerationTaskError(
            code="transport_error",
            phase="transport",
            retryable=retryable,
            public_message="text completion request failed",
            latency_s=self._elapsed(started_at),
        )
    if transport_failure is not None:
        raise transport_failure

    parse_failure: Optional[GenerationTaskError] = None
    try:
        text = self.parse(raw_text)
        if not isinstance(text, str):
            raise TypeError("parse must return text")
    except Exception:
        parse_failure = GenerationTaskError(
            code="parse_error",
            phase="parse",
            retryable=False,
            public_message="generation response parsing failed",
            latency_s=self._elapsed(started_at),
        )
    if parse_failure is not None:
        raise parse_failure
    if not text:
        raise GenerationTaskError(
            code=self.empty_output_error,
            phase="response",
            retryable=False,
            public_message="generation produced no usable text",
            latency_s=latency_s,
        )
    return GeneratedTextResult(
        text=text,
        latency_s=latency_s,
        model=self._client_model(client),
        error=None,
    )

execute(client, **inputs)

Source code in nemo_retriever/models/llm/tasks/base.py

def execute(self, client: TextCompletionClient, **inputs: object) -> GeneratedTextResult:
    """Compatibility facade that collects strict failures into a result."""
    try:
        return self.invoke(client, **inputs)
    except GenerationTaskError as exc:
        return GeneratedTextResult(
            text="",
            latency_s=exc.latency_s,
            model=self._client_model(client),
            error=exc.code,
        )

nemo_retriever.models.llm.clients.litellm.LiteLLMClient

Source code in nemo_retriever/models/llm/clients/litellm.py

class LiteLLMClient:
    """Unified LLM client backed by litellm.

    A single model string change routes to any supported provider:
    - NVIDIA NIM:  nvidia_nim/<org>/<model>
    - OpenAI:      openai/<model>
    - Any OpenAI-compatible server (vLLM, Ollama): openai/<model> + api_base
    - HuggingFace: huggingface/<org>/<model>

    Provider API keys are read from environment variables automatically
    (NVIDIA_API_KEY, OPENAI_API_KEY, HUGGINGFACE_API_KEY, etc.).

    Configuration is split into two orthogonal Pydantic objects:

    * ``transport``: :class:`~nemo_retriever.common.params.LLMRemoteClientParams`
      owns provider endpoint, authentication, retry, and timeout.
    * ``sampling``: :class:`~nemo_retriever.common.params.LLMInferenceParams`
      owns ``temperature``, ``top_p``, and ``max_tokens``.

    Use :meth:`from_kwargs` for a flat, backwards-compatible constructor.
    """

    _DEFAULT_MODEL: str = "nvidia_nim/nvidia/llama-3.3-nemotron-super-49b-v1.5"
    supports_concurrent_calls: bool = True

    def __init__(
        self,
        transport: LLMRemoteClientParams,
        sampling: Optional[LLMInferenceParams] = None,
    ):
        self.transport = transport
        # Default to ``temperature=0.0, max_tokens=4096`` so the structured
        # constructor matches ``from_kwargs`` and keeps RAG-eval runs
        # deterministic.  ``LLMInferenceParams`` itself defaults to
        # ``max_tokens=1024`` for captioning/summarization workloads; RAG
        # answers routinely exceed that, so the client overrides it.
        self.sampling = sampling if sampling is not None else LLMInferenceParams(temperature=0.0, max_tokens=4096)

    @property
    def model(self) -> str:
        """Return the model identifier from the transport params."""
        return self.transport.model

    @classmethod
    def from_kwargs(
        cls,
        *,
        model: str = _DEFAULT_MODEL,
        api_base: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: Optional[float] = 0.0,
        top_p: Optional[float] = None,
        max_tokens: int = 4096,
        extra_params: Optional[dict[str, Any]] = None,
        num_retries: int = 3,
        timeout: float = 120.0,
        rag_system_prompt: Optional[str] = None,
        rag_system_prompt_prefix: Optional[str] = None,
        reasoning_enabled: bool = True,
    ) -> "LiteLLMClient":
        """Flat-kwarg constructor for zero-churn migration from the old signature.

        Splits the flat kwargs into the two structured params objects. All
        validation (temperature range, ``num_retries >= 0``, ``timeout > 0``)
        is delegated to the Pydantic models.
        """
        transport = LLMRemoteClientParams(
            model=model,
            api_base=api_base,
            api_key=api_key,
            num_retries=num_retries,
            timeout=timeout,
            extra_params=extra_params or {},
            rag_system_prompt=rag_system_prompt,
            rag_system_prompt_prefix=rag_system_prompt_prefix,
            reasoning_enabled=reasoning_enabled,
        )
        sampling = LLMInferenceParams(
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
        )
        return cls(transport=transport, sampling=sampling)

    def complete(
        self,
        messages: list[dict],
        max_tokens: Optional[int] = None,
        extra_params: Optional[dict[str, Any]] = None,
    ) -> tuple[str, float]:
        """Raw litellm completion call. Returns (content_text, latency_s)."""
        validate_llm_extra_params(self.transport.extra_params, source="LLMRemoteClientParams.extra_params")
        validate_llm_extra_params(extra_params, source="GenerationRequest.extra_params")
        import litellm

        sampling_kwargs = self.sampling.to_sampling_kwargs()
        if max_tokens is not None:
            sampling_kwargs["max_tokens"] = max_tokens

        call_kwargs: dict[str, Any] = {
            "model": self.transport.model,
            "messages": messages,
            "num_retries": self.transport.num_retries,
            "timeout": self.transport.timeout,
            **sampling_kwargs,
        }
        if self.transport.api_base:
            call_kwargs["api_base"] = self.transport.api_base
        if self.transport._uses_no_api_key("api_key"):
            call_kwargs["api_key"] = _NO_AUTH_API_KEY
        elif self.transport.api_key is not None:
            resolved_api_key = resolve_environment_reference(self.transport.api_key)
            if resolved_api_key:
                call_kwargs["api_key"] = resolved_api_key
        call_kwargs.update(_deep_merge_dicts(self.transport.extra_params, extra_params or {}))

        t0 = time.monotonic()
        try:
            response = litellm.completion(**call_kwargs)
        except Exception as exc:
            err = str(exc)
            if "temperature" in err and "top_p" in err:
                logger.error(
                    "Model %s rejected the request because both `temperature` "
                    "and `top_p` were specified. Some providers (e.g. Bedrock) "
                    "only accept one. Either remove `top_p` from the model "
                    "config or set `temperature` to null. Sent: "
                    "temperature=%s, top_p=%s",
                    self.transport.model,
                    call_kwargs.get("temperature"),
                    call_kwargs.get("top_p"),
                )
            raise
        latency = time.monotonic() - t0

        choices = _field(response, "choices")
        if not isinstance(choices, (list, tuple)) or len(choices) != 1:
            raise UnsupportedTextResponseError("provider response must contain exactly one choice")
        choice = choices[0]
        message = _field(choice, "message")
        if message is None:
            raise UnsupportedTextResponseError("provider response choice has no message")
        if _field(choice, "finish_reason") in {"tool_calls", "function_call"}:
            raise UnsupportedTextResponseError("tool-call responses are unsupported by the text completion contract")
        if _field(message, "tool_calls") or _field(message, "function_call"):
            raise UnsupportedTextResponseError("tool-call responses are unsupported by the text completion contract")
        if _field(message, "refusal"):
            raise UnsupportedTextResponseError("refusal responses are unsupported by the text completion contract")
        if any(_field(message, field_name) is not None for field_name in ("audio", "images", "videos")):
            raise UnsupportedTextResponseError("non-text response modalities are unsupported")
        content = _field(message, "content")
        if not isinstance(content, str):
            raise UnsupportedTextResponseError("provider response content must be plain text")
        content = content.strip()
        return content, latency

    def generate(
        self,
        query: str,
        chunks: list[str],
        *,
        reasoning_enabled: Optional[bool] = None,
    ) -> GenerationResult:
        """Generate an answer for the given query using retrieved chunks as context."""
        effective_reasoning_enabled = (
            self.transport.reasoning_enabled if reasoning_enabled is None else reasoning_enabled
        )
        task = RagAnswerTask(
            system_prompt=self.transport.rag_system_prompt,
            system_prompt_prefix=self.transport.rag_system_prompt_prefix,
            reasoning_enabled=effective_reasoning_enabled,
        )
        result = task.execute(self, query=query, chunks=chunks)
        return GenerationResult(
            answer=result.text,
            latency_s=result.latency_s,
            model=result.model,
            error=result.error,
        )

supports_concurrent_calls = True class-attribute instance-attribute

transport = transport instance-attribute

sampling = sampling if sampling is not None else LLMInferenceParams(temperature=0.0, max_tokens=4096) instance-attribute

model property

from_kwargs(*, model=_DEFAULT_MODEL, api_base=None, api_key=None, temperature=0.0, top_p=None, max_tokens=4096, extra_params=None, num_retries=3, timeout=120.0, rag_system_prompt=None, rag_system_prompt_prefix=None, reasoning_enabled=True) classmethod

Source code in nemo_retriever/models/llm/clients/litellm.py

@classmethod
def from_kwargs(
    cls,
    *,
    model: str = _DEFAULT_MODEL,
    api_base: Optional[str] = None,
    api_key: Optional[str] = None,
    temperature: Optional[float] = 0.0,
    top_p: Optional[float] = None,
    max_tokens: int = 4096,
    extra_params: Optional[dict[str, Any]] = None,
    num_retries: int = 3,
    timeout: float = 120.0,
    rag_system_prompt: Optional[str] = None,
    rag_system_prompt_prefix: Optional[str] = None,
    reasoning_enabled: bool = True,
) -> "LiteLLMClient":
    """Flat-kwarg constructor for zero-churn migration from the old signature.

    Splits the flat kwargs into the two structured params objects. All
    validation (temperature range, ``num_retries >= 0``, ``timeout > 0``)
    is delegated to the Pydantic models.
    """
    transport = LLMRemoteClientParams(
        model=model,
        api_base=api_base,
        api_key=api_key,
        num_retries=num_retries,
        timeout=timeout,
        extra_params=extra_params or {},
        rag_system_prompt=rag_system_prompt,
        rag_system_prompt_prefix=rag_system_prompt_prefix,
        reasoning_enabled=reasoning_enabled,
    )
    sampling = LLMInferenceParams(
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens,
    )
    return cls(transport=transport, sampling=sampling)

complete(messages, max_tokens=None, extra_params=None)

Source code in nemo_retriever/models/llm/clients/litellm.py

def complete(
    self,
    messages: list[dict],
    max_tokens: Optional[int] = None,
    extra_params: Optional[dict[str, Any]] = None,
) -> tuple[str, float]:
    """Raw litellm completion call. Returns (content_text, latency_s)."""
    validate_llm_extra_params(self.transport.extra_params, source="LLMRemoteClientParams.extra_params")
    validate_llm_extra_params(extra_params, source="GenerationRequest.extra_params")
    import litellm

    sampling_kwargs = self.sampling.to_sampling_kwargs()
    if max_tokens is not None:
        sampling_kwargs["max_tokens"] = max_tokens

    call_kwargs: dict[str, Any] = {
        "model": self.transport.model,
        "messages": messages,
        "num_retries": self.transport.num_retries,
        "timeout": self.transport.timeout,
        **sampling_kwargs,
    }
    if self.transport.api_base:
        call_kwargs["api_base"] = self.transport.api_base
    if self.transport._uses_no_api_key("api_key"):
        call_kwargs["api_key"] = _NO_AUTH_API_KEY
    elif self.transport.api_key is not None:
        resolved_api_key = resolve_environment_reference(self.transport.api_key)
        if resolved_api_key:
            call_kwargs["api_key"] = resolved_api_key
    call_kwargs.update(_deep_merge_dicts(self.transport.extra_params, extra_params or {}))

    t0 = time.monotonic()
    try:
        response = litellm.completion(**call_kwargs)
    except Exception as exc:
        err = str(exc)
        if "temperature" in err and "top_p" in err:
            logger.error(
                "Model %s rejected the request because both `temperature` "
                "and `top_p` were specified. Some providers (e.g. Bedrock) "
                "only accept one. Either remove `top_p` from the model "
                "config or set `temperature` to null. Sent: "
                "temperature=%s, top_p=%s",
                self.transport.model,
                call_kwargs.get("temperature"),
                call_kwargs.get("top_p"),
            )
        raise
    latency = time.monotonic() - t0

    choices = _field(response, "choices")
    if not isinstance(choices, (list, tuple)) or len(choices) != 1:
        raise UnsupportedTextResponseError("provider response must contain exactly one choice")
    choice = choices[0]
    message = _field(choice, "message")
    if message is None:
        raise UnsupportedTextResponseError("provider response choice has no message")
    if _field(choice, "finish_reason") in {"tool_calls", "function_call"}:
        raise UnsupportedTextResponseError("tool-call responses are unsupported by the text completion contract")
    if _field(message, "tool_calls") or _field(message, "function_call"):
        raise UnsupportedTextResponseError("tool-call responses are unsupported by the text completion contract")
    if _field(message, "refusal"):
        raise UnsupportedTextResponseError("refusal responses are unsupported by the text completion contract")
    if any(_field(message, field_name) is not None for field_name in ("audio", "images", "videos")):
        raise UnsupportedTextResponseError("non-text response modalities are unsupported")
    content = _field(message, "content")
    if not isinstance(content, str):
        raise UnsupportedTextResponseError("provider response content must be plain text")
    content = content.strip()
    return content, latency

generate(query, chunks, *, reasoning_enabled=None)

Source code in nemo_retriever/models/llm/clients/litellm.py

def generate(
    self,
    query: str,
    chunks: list[str],
    *,
    reasoning_enabled: Optional[bool] = None,
) -> GenerationResult:
    """Generate an answer for the given query using retrieved chunks as context."""
    effective_reasoning_enabled = (
        self.transport.reasoning_enabled if reasoning_enabled is None else reasoning_enabled
    )
    task = RagAnswerTask(
        system_prompt=self.transport.rag_system_prompt,
        system_prompt_prefix=self.transport.rag_system_prompt_prefix,
        reasoning_enabled=effective_reasoning_enabled,
    )
    result = task.execute(self, query=query, chunks=chunks)
    return GenerationResult(
        answer=result.text,
        latency_s=result.latency_s,
        model=result.model,
        error=result.error,
    )

nemo_retriever.models.llm.clients.judge.LLMJudge

Source code in nemo_retriever/models/llm/clients/judge.py

class LLMJudge:
    """LLM-as-judge that scores candidate answers on a ``0.0-1.0`` scale.

    Ports ragas' dual-judge ``AnswerAccuracy`` onto ``LiteLLMClient``. Two
    paraphrased judges each rate the answer ``0/2/4`` against the reference --
    the second with reference/candidate roles swapped -- and the normalised
    ratings are averaged.

    Configuration is split into two Pydantic objects:

    * ``transport``: :class:`~nemo_retriever.common.params.LLMRemoteClientParams` owns
      the endpoint, api_key, retries, and timeout. ``num_retries`` is reused as
      the per-judge attempt budget for obtaining a valid ``0/2/4`` rating.
    * ``sampling``: :class:`~nemo_retriever.common.params.LLMInferenceParams` owns
      ``temperature`` / ``top_p`` / ``max_tokens``. Defaults to
      ``temperature=0.1, max_tokens=4096`` for judge consistency.

    Use :meth:`from_kwargs` for a flat, backwards-compatible constructor.
    """

    _DEFAULT_MODEL: str = "nvidia_nim/nvidia/llama-3.3-nemotron-super-49b-v1.5"
    # max_tokens must accommodate the Nemotron reasoning block + the final
    # {"rating": X}; NVIDIA's llm-judge recipe uses 32768. 4096 truncated mid-think.
    _DEFAULT_SAMPLING: LLMInferenceParams = LLMInferenceParams(temperature=0.1, max_tokens=32768)

    def __init__(
        self,
        transport: LLMRemoteClientParams,
        sampling: Optional[LLMInferenceParams] = None,
    ):
        self.transport = transport
        self.sampling = sampling if sampling is not None else self._DEFAULT_SAMPLING
        self._client = LiteLLMClient(transport=transport, sampling=self.sampling)

    @property
    def model(self) -> str:
        """Return the judge model identifier from the transport params."""
        return self.transport.model

    @classmethod
    def from_kwargs(
        cls,
        *,
        model: str = _DEFAULT_MODEL,
        api_base: Optional[str] = None,
        api_key: Optional[str] = None,
        extra_params: Optional[dict[str, Any]] = None,
        num_retries: int = 3,
        timeout: float = 120.0,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
    ) -> "LLMJudge":
        """Flat-kwarg constructor for zero-churn migration from the old signature.

        Sampling is left at the class default unless ``temperature`` or
        ``max_tokens`` is supplied. Use the two-arg constructor to override
        the full sampling object.
        """
        transport = LLMRemoteClientParams(
            model=model,
            api_base=api_base,
            api_key=api_key,
            num_retries=num_retries,
            timeout=timeout,
            extra_params=extra_params or {},
        )
        sampling = None
        if temperature is not None or max_tokens is not None:
            sampling = LLMInferenceParams(
                temperature=cls._DEFAULT_SAMPLING.temperature if temperature is None else temperature,
                top_p=cls._DEFAULT_SAMPLING.top_p,
                max_tokens=cls._DEFAULT_SAMPLING.max_tokens if max_tokens is None else max_tokens,
            )
        return cls(transport=transport, sampling=sampling)

    def _rate(self, prefix: str, query: str, user_answer: str, reference_answer: str) -> tuple[float, Optional[str]]:
        """Run one judge prompt.

        Returns ``(score, error)`` where ``score`` is the normalised
        0.0/0.5/1.0 rating (or NaN) and ``error`` is the last transport error
        string when every attempt failed, else ``None``. Mirrors ragas
        ``_get_judge_rating``: retry on an invalid rating or a transport error
        up to ``num_retries`` attempts, then give up with NaN.
        """
        # The prompt already forbids explanation and demands `{"rating": X}`, and
        # _parse_rating strips any <think> block — i.e. the judge is designed for the
        # Nemotron reasoning model. The only requirement is a large enough max_tokens
        # for the reasoning block to finish and still emit the rating (NVIDIA's
        # llm-judge recipe uses 32768); 4096 truncated mid-think -> null content.
        messages = [{"role": "user", "content": _render_prompt(prefix, query, user_answer, reference_answer)}]
        attempts = max(1, self.transport.num_retries)
        last_exc: Optional[Exception] = None
        for attempt in range(attempts):
            try:
                raw, _ = self._client.complete(messages)
            except Exception as exc:  # noqa: BLE001 - retried; surfaced below if all attempts fail
                last_exc = exc
                logger.warning("Judge transport error on attempt %d/%d: %s", attempt + 1, attempts, exc)
                continue
            rating = _parse_rating(raw)
            if rating in _VALID_RATINGS:
                return rating / 4.0, None
        if last_exc is not None:
            logger.warning("All %d judge attempts failed; last error: %s", attempts, last_exc)
            return float("nan"), str(last_exc)
        # Attempts succeeded at the transport layer but never produced a valid rating.
        return float("nan"), None

    def judge(self, query: str, reference: str, candidate: str) -> JudgeResult:
        """Score a candidate answer against the reference answer (0.0-1.0)."""
        if not candidate or not candidate.strip():
            return JudgeResult(score=None, reasoning="Candidate answer was empty.", error="empty_candidate")

        try:
            # Judge 1: candidate as the user answer, reference as ground truth.
            rating1, err1 = self._rate(_JUDGE1_PREFIX, query, candidate, reference)
            # Judge 2: roles swapped (bidirectional check) under the paraphrased rubric.
            rating2, err2 = self._rate(_JUDGE2_PREFIX, query, reference, candidate)
        except Exception as exc:
            return JudgeResult(score=None, reasoning="", error=f"judge_api_error: {exc}")

        score = _average_scores(rating1, rating2)
        if math.isnan(score):
            transport_err = err1 or err2
            error = "judge_no_score: neither judge produced a valid rating"
            if transport_err is not None:
                error += f" (last transport error: {transport_err})"
            return JudgeResult(score=None, reasoning="", error=error)
        # AnswerAccuracy emits only a numeric rating, so there is no rationale.
        return JudgeResult(score=float(score), reasoning="")

transport = transport instance-attribute

sampling = sampling if sampling is not None else self._DEFAULT_SAMPLING instance-attribute

model property

from_kwargs(*, model=_DEFAULT_MODEL, api_base=None, api_key=None, extra_params=None, num_retries=3, timeout=120.0, temperature=None, max_tokens=None) classmethod

Source code in nemo_retriever/models/llm/clients/judge.py

@classmethod
def from_kwargs(
    cls,
    *,
    model: str = _DEFAULT_MODEL,
    api_base: Optional[str] = None,
    api_key: Optional[str] = None,
    extra_params: Optional[dict[str, Any]] = None,
    num_retries: int = 3,
    timeout: float = 120.0,
    temperature: Optional[float] = None,
    max_tokens: Optional[int] = None,
) -> "LLMJudge":
    """Flat-kwarg constructor for zero-churn migration from the old signature.

    Sampling is left at the class default unless ``temperature`` or
    ``max_tokens`` is supplied. Use the two-arg constructor to override
    the full sampling object.
    """
    transport = LLMRemoteClientParams(
        model=model,
        api_base=api_base,
        api_key=api_key,
        num_retries=num_retries,
        timeout=timeout,
        extra_params=extra_params or {},
    )
    sampling = None
    if temperature is not None or max_tokens is not None:
        sampling = LLMInferenceParams(
            temperature=cls._DEFAULT_SAMPLING.temperature if temperature is None else temperature,
            top_p=cls._DEFAULT_SAMPLING.top_p,
            max_tokens=cls._DEFAULT_SAMPLING.max_tokens if max_tokens is None else max_tokens,
        )
    return cls(transport=transport, sampling=sampling)

judge(query, reference, candidate)

Source code in nemo_retriever/models/llm/clients/judge.py

def judge(self, query: str, reference: str, candidate: str) -> JudgeResult:
    """Score a candidate answer against the reference answer (0.0-1.0)."""
    if not candidate or not candidate.strip():
        return JudgeResult(score=None, reasoning="Candidate answer was empty.", error="empty_candidate")

    try:
        # Judge 1: candidate as the user answer, reference as ground truth.
        rating1, err1 = self._rate(_JUDGE1_PREFIX, query, candidate, reference)
        # Judge 2: roles swapped (bidirectional check) under the paraphrased rubric.
        rating2, err2 = self._rate(_JUDGE2_PREFIX, query, reference, candidate)
    except Exception as exc:
        return JudgeResult(score=None, reasoning="", error=f"judge_api_error: {exc}")

    score = _average_scores(rating1, rating2)
    if math.isnan(score):
        transport_err = err1 or err2
        error = "judge_no_score: neither judge produced a valid rating"
        if transport_err is not None:
            error += f" (last transport error: {transport_err})"
        return JudgeResult(score=None, reasoning="", error=error)
    # AnswerAccuracy emits only a numeric rating, so there is no rationale.
    return JudgeResult(score=float(score), reasoning="")

Parameter models

nemo_retriever.common.params

IngestorRunMode = Literal['inprocess', 'batch', 'service'] module-attribute

NO_API_KEY = '' module-attribute

MetaJoinKey = Literal['auto', 'source_id', 'source_name'] module-attribute

SPLIT_CONFIG_VALID_KEYS = frozenset({'text', 'html', 'pdf', 'audio', 'image', 'video'}) module-attribute

ASRParams

Bases: _ParamsModel

Params for ASR (Parakeet/Riva gRPC or local transformers backend).

Choice of remote-NIM vs local-model is made by the :class:ASRActor archetype (CPU variant = remote, GPU variant = local), not by a flag here. Pass audio_endpoints to force the remote variant on any host; leave them empty to let the archetype pick GPU (local) when a GPU is present and fall back to remote (NVCF default) when not.

Source code in nemo_retriever/common/params/models.py

class ASRParams(_ParamsModel):
    """Params for ASR (Parakeet/Riva gRPC or local transformers backend).

    Choice of remote-NIM vs local-model is made by the :class:`ASRActor`
    archetype (CPU variant = remote, GPU variant = local), not by a flag here.
    Pass ``audio_endpoints`` to force the remote variant on any host; leave
    them empty to let the archetype pick GPU (local) when a GPU is present
    and fall back to remote (NVCF default) when not.
    """

    audio_endpoints: Tuple[Optional[str], Optional[str]] = (None, None)
    audio_infer_protocol: str = "grpc"
    # ``auto``: streaming (online) for NVCF; offline recognize for other gRPC
    # endpoints (e.g. Helm Parakeet NIM with ``mode=ofl``).
    audio_infer_mode: Literal["auto", "online", "offline"] = "auto"
    function_id: Optional[str] = None
    auth_token: Optional[str] = None
    segment_audio: bool = False

audio_endpoints = (None, None) class-attribute instance-attribute

audio_infer_protocol = 'grpc' class-attribute instance-attribute

audio_infer_mode = 'auto' class-attribute instance-attribute

function_id = None class-attribute instance-attribute

auth_token = None class-attribute instance-attribute

segment_audio = False class-attribute instance-attribute

AudioChunkParams

Bases: _ParamsModel

Params for media chunking (audio/video split). Aligned with nemo_retriever.api dataloader.

Set enabled=False (when wired through VideoSplitActor) to skip audio chunking and ASR on a video pipeline — useful for visual-only recall benchmarks. MediaChunkActor ignores this flag for the audio-only pipeline since chunking is the whole point there.

audio_only=True on a video input extracts only the audio track, runs ASR over it, and skips the visual branch entirely — no frame extraction, no OCR, no audio/visual fusion.

video_audio_separate is accepted for compatibility but ignored by MediaChunkActor on video inputs: this ASR chunking path always demuxes videos to ASR-safe audio chunks and does not emit video-container chunks. Use VideoSplitActor or the video pipeline when you need audio+visual video processing.

Source code in nemo_retriever/common/params/models.py

class AudioChunkParams(_ParamsModel):
    """Params for media chunking (audio/video split). Aligned with `nemo_retriever.api` dataloader.

    Set ``enabled=False`` (when wired through ``VideoSplitActor``) to skip
    audio chunking and ASR on a video pipeline — useful for visual-only
    recall benchmarks. ``MediaChunkActor`` ignores this flag for the
    audio-only pipeline since chunking is the whole point there.

    ``audio_only=True`` on a video input extracts only the audio track,
    runs ASR over it, and skips the visual branch entirely — no frame
    extraction, no OCR, no audio/visual fusion.

    ``video_audio_separate`` is accepted for compatibility but ignored by
    ``MediaChunkActor`` on video inputs: this ASR chunking path always demuxes
    videos to ASR-safe audio chunks and does not emit video-container chunks.
    Use ``VideoSplitActor`` or the video pipeline when you need audio+visual
    video processing.
    """

    enabled: bool = True
    split_type: Literal["size", "time", "frame"] = "size"
    split_interval: int = 450
    audio_only: bool = False
    video_audio_separate: bool = False

enabled = True class-attribute instance-attribute

split_type = 'size' class-attribute instance-attribute

split_interval = 450 class-attribute instance-attribute

audio_only = False class-attribute instance-attribute

video_audio_separate = False class-attribute instance-attribute

AudioVisualFuseParams

Bases: _ParamsModel

Toggle for :class:~nemo_retriever.video.AudioVisualFuser.

Source code in nemo_retriever/common/params/models.py

class AudioVisualFuseParams(_ParamsModel):
    """Toggle for :class:`~nemo_retriever.video.AudioVisualFuser`."""

    enabled: bool = True

enabled = True class-attribute instance-attribute

BatchTuningParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class BatchTuningParams(_ParamsModel):
    debug_run_id: str = "unknown"
    pdf_split_batch_size: int = 1
    pdf_extract_batch_size: int = 4
    pdf_extract_num_cpus: float = 2
    pdf_extract_workers: Optional[int] = None
    page_elements_batch_size: int = 24
    detect_batch_size: int = 24
    ocr_inference_batch_size: Optional[int] = None
    page_elements_workers: Optional[int] = None
    ocr_workers: Optional[int] = None
    detect_workers: Optional[int] = None
    page_elements_cpus_per_actor: float = 1
    ocr_cpus_per_actor: float = 1
    table_structure_workers: Optional[int] = None
    table_structure_batch_size: Optional[int] = None
    table_structure_cpus_per_actor: float = 1
    embed_workers: Optional[int] = None
    embed_batch_size: int = 32
    embed_cpus_per_actor: float = 1
    gpu_page_elements: Optional[float] = None
    gpu_ocr: Optional[float] = None
    gpu_table_structure: Optional[float] = None
    gpu_embed: Optional[float] = None
    nemotron_parse_workers: Optional[int] = None
    gpu_nemotron_parse: Optional[float] = None
    nemotron_parse_batch_size: Optional[int] = None
    store_workers: Optional[int] = None
    inference_batch_size: int = 8

debug_run_id = 'unknown' class-attribute instance-attribute

pdf_split_batch_size = 1 class-attribute instance-attribute

pdf_extract_batch_size = 4 class-attribute instance-attribute

pdf_extract_num_cpus = 2 class-attribute instance-attribute

pdf_extract_workers = None class-attribute instance-attribute

page_elements_batch_size = 24 class-attribute instance-attribute

detect_batch_size = 24 class-attribute instance-attribute

ocr_inference_batch_size = None class-attribute instance-attribute

page_elements_workers = None class-attribute instance-attribute

ocr_workers = None class-attribute instance-attribute

detect_workers = None class-attribute instance-attribute

page_elements_cpus_per_actor = 1 class-attribute instance-attribute

ocr_cpus_per_actor = 1 class-attribute instance-attribute

table_structure_workers = None class-attribute instance-attribute

table_structure_batch_size = None class-attribute instance-attribute

table_structure_cpus_per_actor = 1 class-attribute instance-attribute

embed_workers = None class-attribute instance-attribute

embed_batch_size = 32 class-attribute instance-attribute

embed_cpus_per_actor = 1 class-attribute instance-attribute

gpu_page_elements = None class-attribute instance-attribute

gpu_ocr = None class-attribute instance-attribute

gpu_table_structure = None class-attribute instance-attribute

gpu_embed = None class-attribute instance-attribute

nemotron_parse_workers = None class-attribute instance-attribute

gpu_nemotron_parse = None class-attribute instance-attribute

nemotron_parse_batch_size = None class-attribute instance-attribute

store_workers = None class-attribute instance-attribute

inference_batch_size = 8 class-attribute instance-attribute

CaptionParams

Bases: LLMInferenceParams

Source code in nemo_retriever/common/params/models.py

class CaptionParams(LLMInferenceParams):
    endpoint_url: Optional[str] = None
    model_name: str = Field(
        default=DEFAULT_LOCAL_CAPTION_MODEL_ID,
        description=(
            "Caption model identifier. The default local BF16 checkpoint has approximately 62 GiB of weights; "
            "set this explicitly to select a smaller local model or an API model for a remote endpoint."
        ),
    )
    api_key: Optional[str] = None
    prompt: str = "Caption the content of this image:"
    system_prompt: Optional[str] = "/no_think"
    batch_size: int = 8
    device: Optional[str] = None
    hf_cache_dir: Optional[str] = None
    context_text_max_chars: int = 0
    tensor_parallel_size: int = 1
    gpu_memory_utilization: Optional[float] = Field(
        default=None,
        gt=0,
        le=1,
        description="Fraction of GPU memory reserved for local vLLM captioning; defaults to the model profile.",
    )
    caption_infographics: bool = False
    extra_body: dict[str, Any] = Field(default_factory=dict)

    @field_validator("temperature")
    @classmethod
    def _require_temperature(cls, value: Optional[float]) -> float:
        if value is None:
            raise ValueError("temperature cannot be None for captioning")
        return value

endpoint_url = None class-attribute instance-attribute

model_name = Field(default=DEFAULT_LOCAL_CAPTION_MODEL_ID, description='Caption model identifier. The default local BF16 checkpoint has approximately 62 GiB of weights; set this explicitly to select a smaller local model or an API model for a remote endpoint.') class-attribute instance-attribute

api_key = None class-attribute instance-attribute

prompt = 'Caption the content of this image:' class-attribute instance-attribute

system_prompt = '/no_think' class-attribute instance-attribute

batch_size = 8 class-attribute instance-attribute

device = None class-attribute instance-attribute

hf_cache_dir = None class-attribute instance-attribute

context_text_max_chars = 0 class-attribute instance-attribute

tensor_parallel_size = 1 class-attribute instance-attribute

gpu_memory_utilization = Field(default=None, gt=0, le=1, description='Fraction of GPU memory reserved for local vLLM captioning; defaults to the model profile.') class-attribute instance-attribute

caption_infographics = False class-attribute instance-attribute

extra_body = Field(default_factory=dict) class-attribute instance-attribute

ChartParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class ChartParams(_ParamsModel):
    remote: RemoteInvokeParams = Field(default_factory=RemoteInvokeParams)
    remote_retry: RemoteRetryParams = Field(default_factory=RemoteRetryParams)
    inference_batch_size: int = 8

remote = Field(default_factory=RemoteInvokeParams) class-attribute instance-attribute

remote_retry = Field(default_factory=RemoteRetryParams) class-attribute instance-attribute

inference_batch_size = 8 class-attribute instance-attribute

DedupParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class DedupParams(_ParamsModel):
    content_hash: bool = True
    bbox_iou: bool = True
    iou_threshold: float = Field(default=0.45, ge=0.0, le=1.0)

content_hash = True class-attribute instance-attribute

bbox_iou = True class-attribute instance-attribute

iou_threshold = Field(default=0.45, ge=0.0, le=1.0) class-attribute instance-attribute

EmbedParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class EmbedParams(_ParamsModel):
    model_name: Optional[str] = None
    embedding_endpoint: Optional[str] = None
    embed_invoke_url: Optional[str] = None
    embed_model_name: Optional[str] = None
    embed_model_revision: Optional[str] = None
    embed_model_provider_prefix: Optional[str] = None
    api_key: Optional[str] = None
    input_type: str = "passage"
    embed_modality: str = "text"  # "text", "image", or "text_image" — default for all element types
    embed_granularity: Literal["element", "page"] = "element"  # "element" = per-element rows, "page" = one row per page
    text_elements_modality: Optional[str] = None  # per-type override for page-text rows
    structured_elements_modality: Optional[str] = None  # per-type override for table/chart/infographic rows
    text_column: str = "text"
    inference_batch_size: int = 32
    output_column: str = "text_embeddings_1b_v2"
    embedding_dim_column: str = "text_embeddings_1b_v2_dim"
    has_embedding_column: str = "text_embeddings_1b_v2_has_embedding"
    embed_output_column: str = "text_embeddings_1b_v2"
    embed_inference_batch_size: int = 16

    local_ingest_embed_backend: str = (
        "vllm"  # "vllm" or "hf" — selects ingest-time embedder backend for both text and VL models
    )
    query_max_length: int = 128
    dimensions: Optional[int] = None

    # Concurrent HTTP embedding requests per Ray batch (OpenAI-compatible NIM).
    nim_http_max_concurrent: int = 32
    request_timeout_s: float = 600.0

    runtime: ModelRuntimeParams = Field(default_factory=ModelRuntimeParams)
    batch_tuning: BatchTuningParams = Field(default_factory=BatchTuningParams)

    @field_validator("local_ingest_embed_backend", mode="before")
    @classmethod
    def _validate_local_ingest_embed_backend(cls, v: str) -> str:
        from nemo_retriever.models import (
            _LOCAL_INGEST_EMBED_BACKENDS,
            normalize_backend,
        )

        return normalize_backend(
            str(v) if v is not None else None,
            _LOCAL_INGEST_EMBED_BACKENDS,
            field_name="local_ingest_embed_backend",
            default="vllm",
        )

    @field_validator(
        "embed_modality",
        "text_elements_modality",
        "structured_elements_modality",
        mode="before",
    )
    @classmethod
    def _validate_modality(cls, v: str | None) -> str | None:
        if v is None:
            return None
        modality = str(v).strip()
        if modality == "image_text":
            raise ValueError("Use 'text_image' instead of 'image_text'.")
        if modality not in VALID_EMBED_MODALITIES:
            raise ValueError(f"Modality must be one of {sorted(VALID_EMBED_MODALITIES)}")
        return modality

    @model_validator(mode="after")
    def _warn_page_granularity_overrides(self) -> "EmbedParams":
        if self.embed_granularity == "page" and (
            self.text_elements_modality is not None or self.structured_elements_modality is not None
        ):
            warnings.warn(
                "text_elements_modality and structured_elements_modality are ignored when "
                "embed_granularity='page' (only embed_modality is used).",
                UserWarning,
                stacklevel=2,
            )
        return self

model_name = None class-attribute instance-attribute

embedding_endpoint = None class-attribute instance-attribute

embed_invoke_url = None class-attribute instance-attribute

embed_model_name = None class-attribute instance-attribute

embed_model_revision = None class-attribute instance-attribute

embed_model_provider_prefix = None class-attribute instance-attribute

api_key = None class-attribute instance-attribute

input_type = 'passage' class-attribute instance-attribute

embed_modality = 'text' class-attribute instance-attribute

embed_granularity = 'element' class-attribute instance-attribute

text_elements_modality = None class-attribute instance-attribute

structured_elements_modality = None class-attribute instance-attribute

text_column = 'text' class-attribute instance-attribute

inference_batch_size = 32 class-attribute instance-attribute

output_column = 'text_embeddings_1b_v2' class-attribute instance-attribute

embedding_dim_column = 'text_embeddings_1b_v2_dim' class-attribute instance-attribute

has_embedding_column = 'text_embeddings_1b_v2_has_embedding' class-attribute instance-attribute

embed_output_column = 'text_embeddings_1b_v2' class-attribute instance-attribute

embed_inference_batch_size = 16 class-attribute instance-attribute

local_ingest_embed_backend = 'vllm' class-attribute instance-attribute

query_max_length = 128 class-attribute instance-attribute

dimensions = None class-attribute instance-attribute

nim_http_max_concurrent = 32 class-attribute instance-attribute

request_timeout_s = 600.0 class-attribute instance-attribute

runtime = Field(default_factory=ModelRuntimeParams) class-attribute instance-attribute

batch_tuning = Field(default_factory=BatchTuningParams) class-attribute instance-attribute

ExtractParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class ExtractParams(_ParamsModel):
    # Extraction flags
    extract_text: bool = True
    extract_images: bool = True
    extract_tables: bool = True
    extract_charts: bool = True
    extract_infographics: bool = False
    extract_page_as_image: Optional[bool] = True

    # Extraction options
    method: Literal["pdfium", "pdfium_hybrid", "ocr", "nemotron_parse", "audio"] = Field(
        default="pdfium",
        description=(
            "Extraction method. PDF extraction supports 'pdfium', 'pdfium_hybrid', 'ocr', and "
            "'nemotron_parse'; 'audio' is retained for the legacy params-driven audio path."
        ),
    )
    # Run PageElementDetection (layout/yolox). Required by TableStructure and
    # OCR. Safe to disable for text-only ingests.
    use_page_elements: bool = True
    use_table_structure: bool = False
    table_output_format: Optional[Literal["pseudo_markdown", "markdown"]] = None
    dpi: int = 200
    image_format: str = "jpeg"
    jpeg_quality: int = 100
    render_mode: Literal["full_dpi", "fit_to_model"] = "fit_to_model"
    inference_batch_size: int = 8
    ocr_model_dir: Optional[str] = None
    ocr_version: Literal["v1", "v2"] = "v2"
    ocr_lang: Optional[Literal["multi", "english"]] = None

    # Service endpoints
    invoke_url: Optional[str] = None
    api_key: Optional[str] = None
    request_timeout_s: float = 60.0
    page_elements_invoke_url: Optional[str] = None
    page_elements_api_key: Optional[str] = None
    page_elements_request_timeout_s: Optional[float] = None
    ocr_invoke_url: Optional[str] = None
    ocr_api_key: Optional[str] = None
    ocr_request_timeout_s: Optional[float] = None
    table_structure_invoke_url: Optional[str] = None
    nemotron_parse_invoke_url: Optional[str] = None
    nemotron_parse_model: Optional[str] = None

    # Output columns
    output_column: str = "page_elements_v3"
    num_detections_column: str = "page_elements_v3_num_detections"
    counts_by_label_column: str = "page_elements_v3_counts_by_label"

    remote_retry: RemoteRetryParams = Field(default_factory=RemoteRetryParams)
    batch_tuning: BatchTuningParams = Field(default_factory=BatchTuningParams)

    @model_validator(mode="after")
    def _auto_enable_features(self) -> "ExtractParams":
        """Auto-configure feature flags from remote endpoints.

        * Enable ``use_table_structure`` when ``table_structure_invoke_url``
          is provided.
        * Default ``table_output_format`` to ``"markdown"`` when the stage is
          enabled and the caller did not explicitly choose a format.
        """
        if self.table_structure_invoke_url and not self.use_table_structure:
            self.use_table_structure = True
        if self.table_output_format is None:
            self.table_output_format = "markdown" if self.use_table_structure else "pseudo_markdown"
        if self.ocr_version == "v1" and self.ocr_lang is not None:
            raise ValueError("ocr_lang is only supported when ocr_version='v2'.")
        if self.method != "nemotron_parse" and (
            self.nemotron_parse_invoke_url is not None or self.nemotron_parse_model is not None
        ):
            raise ValueError(
                "`nemotron_parse_invoke_url` and `nemotron_parse_model` require "
                "`method='nemotron_parse'`; Parse-specific configuration is otherwise ignored."
            )
        if self.method == "nemotron_parse":
            parse_endpoints = validate_nemotron_parse_endpoint_list(self.nemotron_parse_invoke_url or self.invoke_url)
            if (
                not parse_endpoints
                and self.nemotron_parse_model is not None
                and self.nemotron_parse_model != NEMOTRON_PARSE_LOCAL_DEFAULT_MODEL
            ):
                raise ValueError(
                    f"Local Nemotron Parse supports only `{NEMOTRON_PARSE_LOCAL_DEFAULT_MODEL}` in this release; "
                    f"received `{self.nemotron_parse_model}`. Configure `nemotron_parse_invoke_url` or `invoke_url` "
                    "to use a compatible remote model."
                )
        if not self.use_page_elements:
            consumers = [("use_table_structure", self.use_table_structure and self.extract_tables)]
            enabled = [name for name, on in consumers if on]
            if enabled:
                raise ValueError(f"use_page_elements=False is incompatible with: {', '.join(enabled)}")
        return self

extract_text = True class-attribute instance-attribute

extract_images = True class-attribute instance-attribute

extract_tables = True class-attribute instance-attribute

extract_charts = True class-attribute instance-attribute

extract_infographics = False class-attribute instance-attribute

extract_page_as_image = True class-attribute instance-attribute

method = Field(default='pdfium', description="Extraction method. PDF extraction supports 'pdfium', 'pdfium_hybrid', 'ocr', and 'nemotron_parse'; 'audio' is retained for the legacy params-driven audio path.") class-attribute instance-attribute

use_page_elements = True class-attribute instance-attribute

use_table_structure = False class-attribute instance-attribute

table_output_format = None class-attribute instance-attribute

dpi = 200 class-attribute instance-attribute

image_format = 'jpeg' class-attribute instance-attribute

jpeg_quality = 100 class-attribute instance-attribute

render_mode = 'fit_to_model' class-attribute instance-attribute

inference_batch_size = 8 class-attribute instance-attribute

ocr_model_dir = None class-attribute instance-attribute

ocr_version = 'v2' class-attribute instance-attribute

ocr_lang = None class-attribute instance-attribute

invoke_url = None class-attribute instance-attribute

api_key = None class-attribute instance-attribute

request_timeout_s = 60.0 class-attribute instance-attribute

page_elements_invoke_url = None class-attribute instance-attribute

page_elements_api_key = None class-attribute instance-attribute

page_elements_request_timeout_s = None class-attribute instance-attribute

ocr_invoke_url = None class-attribute instance-attribute

ocr_api_key = None class-attribute instance-attribute

ocr_request_timeout_s = None class-attribute instance-attribute

table_structure_invoke_url = None class-attribute instance-attribute

nemotron_parse_invoke_url = None class-attribute instance-attribute

nemotron_parse_model = None class-attribute instance-attribute

output_column = 'page_elements_v3' class-attribute instance-attribute

num_detections_column = 'page_elements_v3_num_detections' class-attribute instance-attribute

counts_by_label_column = 'page_elements_v3_counts_by_label' class-attribute instance-attribute

remote_retry = Field(default_factory=RemoteRetryParams) class-attribute instance-attribute

batch_tuning = Field(default_factory=BatchTuningParams) class-attribute instance-attribute

GpuAllocationParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class GpuAllocationParams(_ParamsModel):
    gpu_devices: list[str] = Field(default_factory=list)
    startup_timeout: float = 600.0

gpu_devices = Field(default_factory=list) class-attribute instance-attribute

startup_timeout = 600.0 class-attribute instance-attribute

HtmlChunkParams

Bases: TextChunkParams

Source code in nemo_retriever/common/params/models.py

class HtmlChunkParams(TextChunkParams):
    pass

IngestExecuteParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class IngestExecuteParams(_ParamsModel):
    show_progress: bool = False
    return_failures: bool = False
    return_traces: bool = False
    return_results: bool = True
    result_schema: Literal["legacy", "compact"] = "legacy"
    return_embeddings: bool = False
    return_images: bool = False
    parallel: bool = False
    max_workers: Optional[int] = None
    gpu_devices: list[str] = Field(default_factory=list)
    page_chunk_size: int = 32
    runtime_metrics_dir: Optional[str] = None
    runtime_metrics_prefix: Optional[str] = None

show_progress = False class-attribute instance-attribute

return_failures = False class-attribute instance-attribute

return_traces = False class-attribute instance-attribute

return_results = True class-attribute instance-attribute

result_schema = 'legacy' class-attribute instance-attribute

return_embeddings = False class-attribute instance-attribute

return_images = False class-attribute instance-attribute

parallel = False class-attribute instance-attribute

max_workers = None class-attribute instance-attribute

gpu_devices = Field(default_factory=list) class-attribute instance-attribute

page_chunk_size = 32 class-attribute instance-attribute

runtime_metrics_dir = None class-attribute instance-attribute

runtime_metrics_prefix = None class-attribute instance-attribute

IngestorCreateParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class IngestorCreateParams(_ParamsModel):
    documents: list[str] = Field(default_factory=list)
    ray_address: Optional[str] = None
    ray_log_to_driver: bool = True
    debug: bool = False
    base_url: str = "http://localhost:7670"
    allow_no_gpu: bool = False
    node_overrides: Optional[dict[str, dict[str, Any]]] = None
    api_key: Optional[str] = None
    error_policy: Literal["raise", "collect"] = "raise"
    # service run mode: maximum number of concurrent page uploads.  Lower
    # values (e.g. 2-4) reduce burst pressure on Kubernetes NodePort /
    # kube-proxy paths that otherwise reset connections under heavy load.
    max_concurrency: Optional[int] = None

documents = Field(default_factory=list) class-attribute instance-attribute

ray_address = None class-attribute instance-attribute

ray_log_to_driver = True class-attribute instance-attribute

debug = False class-attribute instance-attribute

base_url = 'http://localhost:7670' class-attribute instance-attribute

allow_no_gpu = False class-attribute instance-attribute

node_overrides = None class-attribute instance-attribute

api_key = None class-attribute instance-attribute

error_policy = 'raise' class-attribute instance-attribute

max_concurrency = None class-attribute instance-attribute

LanceDbParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class LanceDbParams(_ParamsModel):
    lancedb_uri: str = "lancedb"
    table_name: str = "nv-ingest"
    overwrite: bool = True
    create_index: bool = True
    index_type: str = "IVF_HNSW_SQ"
    metric: str = "l2"
    num_partitions: int = 16
    num_sub_vectors: int = 256
    embedding_column: str = "text_embeddings_1b_v2"
    embedding_key: str = "embedding"
    include_text: bool = True
    text_column: str = "text"
    hybrid: bool = False
    fts_language: str = "English"

lancedb_uri = 'lancedb' class-attribute instance-attribute

table_name = 'nv-ingest' class-attribute instance-attribute

overwrite = True class-attribute instance-attribute

create_index = True class-attribute instance-attribute

index_type = 'IVF_HNSW_SQ' class-attribute instance-attribute

metric = 'l2' class-attribute instance-attribute

num_partitions = 16 class-attribute instance-attribute

num_sub_vectors = 256 class-attribute instance-attribute

embedding_column = 'text_embeddings_1b_v2' class-attribute instance-attribute

embedding_key = 'embedding' class-attribute instance-attribute

include_text = True class-attribute instance-attribute

text_column = 'text' class-attribute instance-attribute

hybrid = False class-attribute instance-attribute

fts_language = 'English' class-attribute instance-attribute

LLMInferenceParams

Bases: _ParamsModel

Reusable LLM sampling / generation parameters.

Inherit from this model to add temperature, top_p, and max_tokens to any task that invokes an LLM (captioning, summarization, etc.).

Source code in nemo_retriever/common/params/models.py

class LLMInferenceParams(_ParamsModel):
    """Reusable LLM sampling / generation parameters.

    Inherit from this model to add temperature, top_p, and max_tokens
    to any task that invokes an LLM (captioning, summarization, etc.).
    """

    temperature: Optional[float] = 1.0
    top_p: Optional[float] = None
    max_tokens: int = 1024

    @field_validator("temperature")
    @classmethod
    def _check_temperature(cls, v: Optional[float]) -> Optional[float]:
        if v is not None and not (0.0 <= v <= 2.0):
            raise ValueError("temperature must be between 0.0 and 2.0")
        return v

    @field_validator("top_p")
    @classmethod
    def _check_top_p(cls, v: Optional[float]) -> Optional[float]:
        if v is not None and not (0.0 <= v <= 1.0):
            raise ValueError("top_p must be between 0.0 and 1.0")
        return v

    @field_validator("max_tokens")
    @classmethod
    def _check_max_tokens(cls, v: int) -> int:
        if v <= 0:
            raise ValueError("max_tokens must be > 0")
        return v

    def to_sampling_kwargs(self) -> dict[str, Any]:
        """Build a dict of sampling parameters suitable for LLM inference calls.

        ``top_p`` is only included when explicitly set (not ``None``), because
        many backends (vLLM, OpenAI, NIM) change behaviour when the key is
        present vs. absent.
        """
        kw: dict[str, Any] = {"max_tokens": self.max_tokens}
        if self.temperature is not None:
            kw["temperature"] = self.temperature
        if self.top_p is not None:
            kw["top_p"] = self.top_p
        return kw

temperature = 1.0 class-attribute instance-attribute

top_p = None class-attribute instance-attribute

max_tokens = 1024 class-attribute instance-attribute

to_sampling_kwargs()

Build a dict of sampling parameters suitable for LLM inference calls.

top_p is only included when explicitly set (not None), because many backends (vLLM, OpenAI, NIM) change behaviour when the key is present vs. absent.

Source code in nemo_retriever/common/params/models.py

def to_sampling_kwargs(self) -> dict[str, Any]:
    """Build a dict of sampling parameters suitable for LLM inference calls.

    ``top_p`` is only included when explicitly set (not ``None``), because
    many backends (vLLM, OpenAI, NIM) change behaviour when the key is
    present vs. absent.
    """
    kw: dict[str, Any] = {"max_tokens": self.max_tokens}
    if self.temperature is not None:
        kw["temperature"] = self.temperature
    if self.top_p is not None:
        kw["top_p"] = self.top_p
    return kw

LLMRemoteClientParams

Bases: _ParamsModel

Transport / connection parameters for any remote LLM client.

Pairs with :class:LLMInferenceParams (sampling) to fully specify a call. api_key=None is left unset so LiteLLM can perform provider-native environment lookup on the worker.

Source code in nemo_retriever/common/params/models.py

class LLMRemoteClientParams(_ParamsModel):
    """Transport / connection parameters for any remote LLM client.

    Pairs with :class:`LLMInferenceParams` (sampling) to fully specify a
    call. ``api_key=None`` is left unset so LiteLLM can perform provider-native
    environment lookup on the worker.
    """

    _auto_resolve_unset_api_keys: ClassVar[bool] = False

    model: str
    api_base: Optional[str] = None
    api_key: Optional[str] = None
    num_retries: int = 3
    timeout: float = 120.0
    extra_params: dict[str, Any] = Field(default_factory=dict)
    rag_system_prompt: Optional[str] = None
    rag_system_prompt_prefix: Optional[str] = None
    reasoning_enabled: bool = True

    @field_validator("extra_params")
    @classmethod
    def _check_extra_params(cls, value: dict[str, Any]) -> dict[str, Any]:
        validate_llm_extra_params(value, source="LLMRemoteClientParams.extra_params")
        return value

    @field_validator("num_retries")
    @classmethod
    def _check_retries(cls, v: int) -> int:
        if v < 0:
            raise ValueError("num_retries must be >= 0")
        return v

    @field_validator("timeout")
    @classmethod
    def _check_timeout(cls, v: float) -> float:
        if v <= 0:
            raise ValueError("timeout must be > 0")
        return v

model instance-attribute

api_base = None class-attribute instance-attribute

api_key = None class-attribute instance-attribute

num_retries = 3 class-attribute instance-attribute

timeout = 120.0 class-attribute instance-attribute

extra_params = Field(default_factory=dict) class-attribute instance-attribute

rag_system_prompt = None class-attribute instance-attribute

rag_system_prompt_prefix = None class-attribute instance-attribute

reasoning_enabled = True class-attribute instance-attribute

LLMSamplingOverrides

Bases: _ParamsModel

Partial sampling overrides resolved on top of task-specific defaults.

Source code in nemo_retriever/common/params/models.py

class LLMSamplingOverrides(_ParamsModel):
    """Partial sampling overrides resolved on top of task-specific defaults."""

    temperature: Optional[float] = None
    top_p: Optional[float] = None
    max_tokens: Optional[int] = None

    @field_validator("temperature")
    @classmethod
    def _check_temperature(cls, v: Optional[float]) -> Optional[float]:
        if v is not None and not (0.0 <= v <= 2.0):
            raise ValueError("temperature must be between 0.0 and 2.0")
        return v

    @field_validator("top_p")
    @classmethod
    def _check_top_p(cls, v: Optional[float]) -> Optional[float]:
        if v is not None and not (0.0 <= v <= 1.0):
            raise ValueError("top_p must be between 0.0 and 1.0")
        return v

    @field_validator("max_tokens")
    @classmethod
    def _check_max_tokens(cls, v: Optional[int]) -> Optional[int]:
        if v is not None and v <= 0:
            raise ValueError("max_tokens must be > 0")
        return v

    @model_validator(mode="after")
    def _reject_explicit_null_max_tokens(self) -> "LLMSamplingOverrides":
        if "max_tokens" in self.model_fields_set and self.max_tokens is None:
            raise ValueError("max_tokens cannot be None; omit it to inherit the task default")
        return self

    @model_serializer(mode="plain")
    def _serialize_only_explicit_overrides(self) -> dict[str, Any]:
        """Preserve omitted-vs-null state across model and JSON round trips."""
        return {
            name: getattr(self, name)
            for name in ("temperature", "top_p", "max_tokens")
            if name in self.model_fields_set
        }

    def __eq__(self, other: object) -> bool:
        if isinstance(other, LLMSamplingOverrides):
            return self.model_fields_set == other.model_fields_set and super().__eq__(other)
        return super().__eq__(other)

    def resolve(self, defaults: LLMInferenceParams) -> LLMInferenceParams:
        """Apply explicitly supplied fields to defaults."""
        values = defaults.model_dump()
        for name in self.model_fields_set:
            value = getattr(self, name)
            values[name] = value
        return LLMInferenceParams(**values)

temperature = None class-attribute instance-attribute

top_p = None class-attribute instance-attribute

max_tokens = None class-attribute instance-attribute

resolve(defaults)

Apply explicitly supplied fields to defaults.

Source code in nemo_retriever/common/params/models.py

def resolve(self, defaults: LLMInferenceParams) -> LLMInferenceParams:
    """Apply explicitly supplied fields to defaults."""
    values = defaults.model_dump()
    for name in self.model_fields_set:
        value = getattr(self, name)
        values[name] = value
    return LLMInferenceParams(**values)

ModelRuntimeParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class ModelRuntimeParams(_ParamsModel):
    device: Optional[str] = None
    hf_cache_dir: Optional[str] = None
    normalize: bool = True
    max_length: int = 8192
    model_name: Optional[str] = None
    gpu_memory_utilization: float = 0.45
    enforce_eager: bool = False

device = None class-attribute instance-attribute

hf_cache_dir = None class-attribute instance-attribute

normalize = True class-attribute instance-attribute

max_length = 8192 class-attribute instance-attribute

model_name = None class-attribute instance-attribute

gpu_memory_utilization = 0.45 class-attribute instance-attribute

enforce_eager = False class-attribute instance-attribute

OcrParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class OcrParams(_ParamsModel):
    remote: RemoteInvokeParams = Field(default_factory=RemoteInvokeParams)
    remote_retry: RemoteRetryParams = Field(default_factory=RemoteRetryParams)
    inference_batch_size: int = 8
    extract_tables: bool = False
    extract_charts: bool = False
    extract_infographics: bool = False

remote = Field(default_factory=RemoteInvokeParams) class-attribute instance-attribute

remote_retry = Field(default_factory=RemoteRetryParams) class-attribute instance-attribute

inference_batch_size = 8 class-attribute instance-attribute

extract_tables = False class-attribute instance-attribute

extract_charts = False class-attribute instance-attribute

extract_infographics = False class-attribute instance-attribute

PageElementsParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class PageElementsParams(_ParamsModel):
    remote: RemoteInvokeParams = Field(default_factory=RemoteInvokeParams)
    remote_retry: RemoteRetryParams = Field(default_factory=RemoteRetryParams)
    inference_batch_size: int = 8
    output_column: str = "page_elements_v3"
    num_detections_column: str = "page_elements_v3_num_detections"
    counts_by_label_column: str = "page_elements_v3_counts_by_label"

remote = Field(default_factory=RemoteInvokeParams) class-attribute instance-attribute

remote_retry = Field(default_factory=RemoteRetryParams) class-attribute instance-attribute

inference_batch_size = 8 class-attribute instance-attribute

output_column = 'page_elements_v3' class-attribute instance-attribute

num_detections_column = 'page_elements_v3_num_detections' class-attribute instance-attribute

counts_by_label_column = 'page_elements_v3_counts_by_label' class-attribute instance-attribute

PdfSplitParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class PdfSplitParams(_ParamsModel):
    start_page: Optional[int] = None
    end_page: Optional[int] = None

start_page = None class-attribute instance-attribute

end_page = None class-attribute instance-attribute

RemoteInvokeParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class RemoteInvokeParams(_ParamsModel):
    invoke_url: Optional[str] = None
    api_key: Optional[str] = None
    request_timeout_s: float = 60.0

invoke_url = None class-attribute instance-attribute

api_key = None class-attribute instance-attribute

request_timeout_s = 60.0 class-attribute instance-attribute

RemoteRetryParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class RemoteRetryParams(_ParamsModel):
    remote_max_pool_workers: int = 32
    remote_max_retries: int = 5
    remote_max_429_retries: int = 3

remote_max_pool_workers = 32 class-attribute instance-attribute

remote_max_retries = 5 class-attribute instance-attribute

remote_max_429_retries = 3 class-attribute instance-attribute

StoreParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class StoreParams(_ParamsModel):
    storage_uri: str = "stored_images"
    storage_options: dict[str, Any] = Field(default_factory=dict)
    image_format: str = "png"
    strip_base64: bool = True
    batch_tuning: BatchTuningParams = Field(default_factory=BatchTuningParams)

    @model_validator(mode="after")
    def _resolve_local_storage_uri(self) -> "StoreParams":
        """Resolve relative local paths to absolute so they survive Ray serialization."""
        if not urlparse(self.storage_uri).scheme:
            self.storage_uri = str(UPath(self.storage_uri).resolve())
        return self

storage_uri = 'stored_images' class-attribute instance-attribute

storage_options = Field(default_factory=dict) class-attribute instance-attribute

image_format = 'png' class-attribute instance-attribute

strip_base64 = True class-attribute instance-attribute

batch_tuning = Field(default_factory=BatchTuningParams) class-attribute instance-attribute

TabularExtractParams

Bases: _ParamsModel

Params for step 1: extract schema metadata and write to Neo4j.

Covers SQLAlchemy reflection of a live database and/or parsing of pre-existing SQL DDL/query files. Produces Database, Schema, Table, Column, View and Query nodes together with their relationships. The Neo4j connection is provided by get_neo4j_conn() (see tabular_data.neo4j) and is not configured here.

Source code in nemo_retriever/common/params/models.py

class TabularExtractParams(_ParamsModel):
    """Params for step 1: extract schema metadata and write to Neo4j.

    Covers SQLAlchemy reflection of a live database and/or parsing of
    pre-existing SQL DDL/query files.  Produces Database, Schema, Table,
    Column, View and Query nodes together with their relationships.
    The Neo4j connection is provided by get_neo4j_conn() (see
    tabular_data.neo4j) and is not configured here.
    """

    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)

    connector: Optional[SQLDatabase] = None

model_config = ConfigDict(extra='forbid', arbitrary_types_allowed=True) class-attribute instance-attribute

connector = None class-attribute instance-attribute

TableParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class TableParams(_ParamsModel):
    remote: RemoteInvokeParams = Field(default_factory=RemoteInvokeParams)
    remote_retry: RemoteRetryParams = Field(default_factory=RemoteRetryParams)
    inference_batch_size: int = 8
    output_column: str = "table_structure_v1"
    num_detections_column: str = "table_structure_v1_num_detections"
    counts_by_label_column: str = "table_structure_v1_counts_by_label"

remote = Field(default_factory=RemoteInvokeParams) class-attribute instance-attribute

remote_retry = Field(default_factory=RemoteRetryParams) class-attribute instance-attribute

inference_batch_size = 8 class-attribute instance-attribute

output_column = 'table_structure_v1' class-attribute instance-attribute

num_detections_column = 'table_structure_v1_num_detections' class-attribute instance-attribute

counts_by_label_column = 'table_structure_v1_counts_by_label' class-attribute instance-attribute

TextChunkParams

Bases: _ParamsModel

Source code in nemo_retriever/common/params/models.py

class TextChunkParams(_ParamsModel):
    max_tokens: int = 1024
    overlap_tokens: int = 0
    tokenizer_model_id: Optional[str] = None
    encoding: str = "utf-8"
    tokenizer_cache_dir: Optional[str] = None

max_tokens = 1024 class-attribute instance-attribute

overlap_tokens = 0 class-attribute instance-attribute

tokenizer_model_id = None class-attribute instance-attribute

encoding = 'utf-8' class-attribute instance-attribute

tokenizer_cache_dir = None class-attribute instance-attribute

TextGenerationParams

Bases: _ParamsModel

Transport, task controls, and partial sampling for text generation.

Source code in nemo_retriever/common/params/models.py

class TextGenerationParams(_ParamsModel):
    """Transport, task controls, and partial sampling for text generation."""

    transport: LLMRemoteClientParams
    sampling: LLMSamplingOverrides = Field(default_factory=LLMSamplingOverrides)
    prompt: Optional[str] = None
    system_prompt: Optional[str] = None
    reasoning_enabled: Optional[bool] = None
    max_workers: int = Field(default=8, ge=1)

    def resolve_sampling(self, defaults: LLMInferenceParams) -> LLMInferenceParams:
        """Resolve explicit sampling fields over a task's defaults."""
        return self.sampling.resolve(defaults)

    @classmethod
    def from_kwargs(
        cls,
        *,
        model: str,
        api_base: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: Any = _SAMPLING_UNSET,
        top_p: Any = _SAMPLING_UNSET,
        max_tokens: Any = _SAMPLING_UNSET,
        extra_params: Optional[dict[str, Any]] = None,
        num_retries: int = 3,
        timeout: float = 120.0,
        rag_system_prompt: Optional[str] = None,
        rag_system_prompt_prefix: Optional[str] = None,
        reasoning_enabled: Optional[bool] = None,
        prompt: Optional[str] = None,
        system_prompt: Optional[str] = None,
        max_workers: int = 8,
    ) -> "TextGenerationParams":
        """Construct structured text-generation params from flat kwargs."""
        sampling_values: dict[str, Any] = {}
        for name, value in (
            ("temperature", temperature),
            ("top_p", top_p),
            ("max_tokens", max_tokens),
        ):
            if value is not _SAMPLING_UNSET:
                sampling_values[name] = value

        transport_reasoning = True if reasoning_enabled is None else reasoning_enabled
        return cls(
            transport=LLMRemoteClientParams(
                model=model,
                api_base=api_base,
                api_key=api_key,
                num_retries=num_retries,
                timeout=timeout,
                extra_params=extra_params or {},
                rag_system_prompt=rag_system_prompt,
                rag_system_prompt_prefix=rag_system_prompt_prefix,
                reasoning_enabled=transport_reasoning,
            ),
            sampling=LLMSamplingOverrides(**sampling_values),
            prompt=prompt,
            system_prompt=system_prompt,
            reasoning_enabled=reasoning_enabled,
            max_workers=max_workers,
        )

transport instance-attribute

sampling = Field(default_factory=LLMSamplingOverrides) class-attribute instance-attribute

prompt = None class-attribute instance-attribute

system_prompt = None class-attribute instance-attribute

reasoning_enabled = None class-attribute instance-attribute

max_workers = Field(default=8, ge=1) class-attribute instance-attribute

resolve_sampling(defaults)

Resolve explicit sampling fields over a task's defaults.

Source code in nemo_retriever/common/params/models.py

def resolve_sampling(self, defaults: LLMInferenceParams) -> LLMInferenceParams:
    """Resolve explicit sampling fields over a task's defaults."""
    return self.sampling.resolve(defaults)

from_kwargs(*, model, api_base=None, api_key=None, temperature=_SAMPLING_UNSET, top_p=_SAMPLING_UNSET, max_tokens=_SAMPLING_UNSET, extra_params=None, num_retries=3, timeout=120.0, rag_system_prompt=None, rag_system_prompt_prefix=None, reasoning_enabled=None, prompt=None, system_prompt=None, max_workers=8) classmethod

Construct structured text-generation params from flat kwargs.

Source code in nemo_retriever/common/params/models.py

@classmethod
def from_kwargs(
    cls,
    *,
    model: str,
    api_base: Optional[str] = None,
    api_key: Optional[str] = None,
    temperature: Any = _SAMPLING_UNSET,
    top_p: Any = _SAMPLING_UNSET,
    max_tokens: Any = _SAMPLING_UNSET,
    extra_params: Optional[dict[str, Any]] = None,
    num_retries: int = 3,
    timeout: float = 120.0,
    rag_system_prompt: Optional[str] = None,
    rag_system_prompt_prefix: Optional[str] = None,
    reasoning_enabled: Optional[bool] = None,
    prompt: Optional[str] = None,
    system_prompt: Optional[str] = None,
    max_workers: int = 8,
) -> "TextGenerationParams":
    """Construct structured text-generation params from flat kwargs."""
    sampling_values: dict[str, Any] = {}
    for name, value in (
        ("temperature", temperature),
        ("top_p", top_p),
        ("max_tokens", max_tokens),
    ):
        if value is not _SAMPLING_UNSET:
            sampling_values[name] = value

    transport_reasoning = True if reasoning_enabled is None else reasoning_enabled
    return cls(
        transport=LLMRemoteClientParams(
            model=model,
            api_base=api_base,
            api_key=api_key,
            num_retries=num_retries,
            timeout=timeout,
            extra_params=extra_params or {},
            rag_system_prompt=rag_system_prompt,
            rag_system_prompt_prefix=rag_system_prompt_prefix,
            reasoning_enabled=transport_reasoning,
        ),
        sampling=LLMSamplingOverrides(**sampling_values),
        prompt=prompt,
        system_prompt=system_prompt,
        reasoning_enabled=reasoning_enabled,
        max_workers=max_workers,
    )

VdbUploadParams

Bases: _ParamsModel

Post-graph vector DB upload configuration.

Sidecar metadata (meta_*) matches nv_ingest_client / metadata_and_filtered_search.ipynb: all three fields must be set together to merge columns into each chunk's content_metadata.

Source code in nemo_retriever/common/params/models.py

class VdbUploadParams(_ParamsModel):
    """Post-graph vector DB upload configuration.

    Sidecar metadata (``meta_*``) matches ``nv_ingest_client`` / ``metadata_and_filtered_search.ipynb``:
    all three fields must be set together to merge columns into each chunk's ``content_metadata``.
    """

    vdb_op: str = "lancedb"
    vdb_kwargs: dict[str, Any] = Field(default_factory=dict)
    meta_dataframe: Optional[Any] = None
    """Path to csv/json/parquet or an in-memory :class:`pandas.DataFrame`."""
    meta_source_field: Optional[str] = None
    meta_fields: Optional[list[str]] = None
    meta_join_key: MetaJoinKey = "auto"
    """How to match rows to documents: ``source_id`` (full path), ``source_name`` (basename), or ``auto`` (try both)."""

    @model_validator(mode="after")
    def _validate_sidecar_triplet(self) -> "VdbUploadParams":
        trio = (self.meta_dataframe, self.meta_source_field, self.meta_fields)
        if all(x is None for x in trio):
            return self
        if any(x is None for x in trio):
            raise ValueError(
                "meta_dataframe, meta_source_field, and meta_fields must all be set together "
                "when attaching sidecar metadata."
            )
        if not self.meta_fields:
            raise ValueError("meta_fields must be a non-empty list when sidecar metadata is enabled.")
        return self

    def to_ingest_operator_kwargs(self) -> dict[str, Any]:
        """Flatten into kwargs for :class:`~nemo_retriever.vdb.IngestVdbOperator`."""
        out = dict(self.vdb_kwargs or {})
        if self.meta_dataframe is not None:
            out["meta_dataframe"] = self.meta_dataframe
            out["meta_source_field"] = self.meta_source_field
            out["meta_fields"] = list(self.meta_fields or [])
            out["meta_join_key"] = self.meta_join_key
        return out

vdb_op = 'lancedb' class-attribute instance-attribute

vdb_kwargs = Field(default_factory=dict) class-attribute instance-attribute

meta_dataframe = None class-attribute instance-attribute

Path to csv/json/parquet or an in-memory :class:pandas.DataFrame.

meta_source_field = None class-attribute instance-attribute

meta_fields = None class-attribute instance-attribute

meta_join_key = 'auto' class-attribute instance-attribute

How to match rows to documents: source_id (full path), source_name (basename), or auto (try both).

to_ingest_operator_kwargs()

Flatten into kwargs for :class:~nemo_retriever.vdb.IngestVdbOperator.

Source code in nemo_retriever/common/params/models.py

def to_ingest_operator_kwargs(self) -> dict[str, Any]:
    """Flatten into kwargs for :class:`~nemo_retriever.vdb.IngestVdbOperator`."""
    out = dict(self.vdb_kwargs or {})
    if self.meta_dataframe is not None:
        out["meta_dataframe"] = self.meta_dataframe
        out["meta_source_field"] = self.meta_source_field
        out["meta_fields"] = list(self.meta_fields or [])
        out["meta_join_key"] = self.meta_join_key
    return out

VideoFrameParams

Bases: _ParamsModel

Params for video frame extraction (ffmpeg fps + perceptual-hash dedup).

Set enabled=False to skip frame extraction entirely; the video pipeline then produces only audio (ASR) rows — no frame OCR, no audio+visual fusion. Useful for ablating the visual modality or for audio-only recall benchmarks against video corpora.

dedup activates perceptual-hash (dhash) dedup before OCR. dhash catches visually-identical adjacent frames that byte-level hashing misses (encoder noise, brightness drift, etc.). On a 60s slide-heavy sample we measured ~91% duplicates collapsed at distance 5 vs ~11% for MD5 — a near-10x cut in OCR cost on slide content. Tune dedup_max_hamming_distance upward for more aggressive merging or down to 0 to require exact perceptual-hash matches.

Source code in nemo_retriever/common/params/models.py

class VideoFrameParams(_ParamsModel):
    """Params for video frame extraction (ffmpeg fps + perceptual-hash dedup).

    Set ``enabled=False`` to skip frame extraction entirely; the video
    pipeline then produces only audio (ASR) rows — no frame OCR, no
    audio+visual fusion. Useful for ablating the visual modality or for
    audio-only recall benchmarks against video corpora.

    ``dedup`` activates perceptual-hash (dhash) dedup before OCR. dhash
    catches visually-identical adjacent frames that byte-level hashing
    misses (encoder noise, brightness drift, etc.). On a 60s slide-heavy
    sample we measured ~91% duplicates collapsed at distance 5 vs ~11%
    for MD5 — a near-10x cut in OCR cost on slide content. Tune
    ``dedup_max_hamming_distance`` upward for more aggressive merging or
    down to 0 to require exact perceptual-hash matches.
    """

    enabled: bool = True
    fps: float = Field(default=1.0, gt=0.0)
    max_frames: Optional[int] = None
    dedup: bool = True
    dedup_max_hamming_distance: int = 5
    dedup_max_dropped_frames: int = 2

enabled = True class-attribute instance-attribute

fps = Field(default=1.0, gt=0.0) class-attribute instance-attribute

max_frames = None class-attribute instance-attribute

dedup = True class-attribute instance-attribute

dedup_max_hamming_distance = 5 class-attribute instance-attribute

dedup_max_dropped_frames = 2 class-attribute instance-attribute

VideoFrameTextDedupParams

Bases: _ParamsModel

Params for merging consecutive video_frame rows with identical OCR text.

After full-frame OCR, slides that are visible for many seconds produce a flood of frames with the same text (image-hash dedup misses them when encoder noise differs frame-to-frame). This stage groups by (source_path, text) and merges adjacent runs into a single row whose segment_start_seconds / segment_end_seconds cover the union of the run.

Tolerance is expressed in dropped frames, not seconds, so it scales with video_frame_fps: at runtime the dedup reads each group's metadata.fps and converts to max_gap_seconds = max_dropped_frames / fps. Default 2 means we bridge gaps of up to 2 missing frames in a run — a typical safety margin for image-hash dedup leaving small holes.

Source code in nemo_retriever/common/params/models.py

class VideoFrameTextDedupParams(_ParamsModel):
    """Params for merging consecutive video_frame rows with identical OCR text.

    After full-frame OCR, slides that are visible for many seconds produce a
    flood of frames with the same text (image-hash dedup misses them when
    encoder noise differs frame-to-frame). This stage groups by
    ``(source_path, text)`` and merges adjacent runs into a single row whose
    ``segment_start_seconds`` / ``segment_end_seconds`` cover the union of
    the run.

    Tolerance is expressed in **dropped frames**, not seconds, so it scales
    with ``video_frame_fps``: at runtime the dedup reads each group's
    ``metadata.fps`` and converts to ``max_gap_seconds = max_dropped_frames / fps``.
    Default 2 means we bridge gaps of up to 2 missing frames in a run —
    a typical safety margin for image-hash dedup leaving small holes.
    """

    enabled: bool = True
    max_dropped_frames: int = 2

enabled = True class-attribute instance-attribute

max_dropped_frames = 2 class-attribute instance-attribute

WebhookParams

Bases: _ParamsModel

Configuration for the webhook notification stage.

When endpoint_url is set, selected columns from the processed batch are serialised to JSON and HTTP-POSTed to that URL. If endpoint_url is None the stage is a no-op.

Source code in nemo_retriever/common/params/models.py

class WebhookParams(_ParamsModel):
    """Configuration for the webhook notification stage.

    When ``endpoint_url`` is set, selected columns from the processed batch
    are serialised to JSON and HTTP-POSTed to that URL.  If ``endpoint_url``
    is ``None`` the stage is a no-op.
    """

    endpoint_url: Optional[str] = None
    columns: list[str] = Field(default_factory=list)
    headers: dict[str, str] = Field(default_factory=dict)
    timeout_s: float = 30.0
    max_retries: int = 3

endpoint_url = None class-attribute instance-attribute

columns = Field(default_factory=list) class-attribute instance-attribute

headers = Field(default_factory=dict) class-attribute instance-attribute

timeout_s = 30.0 class-attribute instance-attribute

max_retries = 3 class-attribute instance-attribute

build_embed_option_kwargs(embed_invoke_url, embed_model_name, local_ingest_embed_backend=None, embed_api_key=None, embed_model_provider_prefix=None, embed_modality=None, text_elements_modality=None, structured_elements_modality=None, embed_granularity=None, embed_workers=None, embed_batch_size=None, embed_cpus_per_actor=None, embed_gpus_per_actor=None, embed_model_revision=None)

Build EmbedParams kwargs from CLI/request option values.

Source code in nemo_retriever/common/params/utils.py

def build_embed_option_kwargs(
    embed_invoke_url: str | None,
    embed_model_name: str | None,
    local_ingest_embed_backend: str | None = None,
    embed_api_key: str | None = None,
    embed_model_provider_prefix: str | None = None,
    embed_modality: str | None = None,
    text_elements_modality: str | None = None,
    structured_elements_modality: str | None = None,
    embed_granularity: str | None = None,
    embed_workers: int | None = None,
    embed_batch_size: int | None = None,
    embed_cpus_per_actor: float | None = None,
    embed_gpus_per_actor: float | None = None,
    embed_model_revision: str | None = None,
) -> Dict[str, Any]:
    """Build ``EmbedParams`` kwargs from CLI/request option values."""
    embed_kwargs: Dict[str, Any] = {}
    if embed_invoke_url is not None:
        embed_kwargs["embed_invoke_url"] = embed_invoke_url
    if embed_model_name is not None:
        # Remote HTTP embedding reads model_name; local/GPU paths read embed_model_name.
        embed_kwargs["model_name"] = embed_model_name
        embed_kwargs["embed_model_name"] = embed_model_name
    if embed_model_revision is not None:
        embed_kwargs["embed_model_revision"] = embed_model_revision
    if local_ingest_embed_backend is not None:
        embed_kwargs["local_ingest_embed_backend"] = local_ingest_embed_backend
    if embed_api_key is not None:
        embed_kwargs["api_key"] = embed_api_key
    if embed_model_provider_prefix is not None:
        embed_kwargs["embed_model_provider_prefix"] = embed_model_provider_prefix
    if embed_modality is not None:
        embed_kwargs["embed_modality"] = embed_modality
    if text_elements_modality is not None:
        embed_kwargs["text_elements_modality"] = text_elements_modality
    if structured_elements_modality is not None:
        embed_kwargs["structured_elements_modality"] = structured_elements_modality
    if embed_granularity is not None:
        embed_kwargs["embed_granularity"] = embed_granularity
    embed_tuning = _build_embed_batch_tuning(
        embed_workers=embed_workers,
        embed_batch_size=embed_batch_size,
        embed_cpus_per_actor=embed_cpus_per_actor,
        embed_gpus_per_actor=embed_gpus_per_actor,
    )
    if embed_tuning is not None:
        embed_kwargs["batch_tuning"] = embed_tuning
    return normalize_embed_kwargs(embed_kwargs)

resolve_split_params(split_config)

Resolve a user-supplied split_config dict into per-key effective params.

Returns a dict keyed by every entry in SPLIT_CONFIG_VALID_KEYS. Each value is one of: a TextChunkParams / HtmlChunkParams instance (chunking enabled for that key), None (key absent — chunking off via the default), or False (explicit opt-out sentinel).

Per-key values supplied by the caller may be a plain dict of chunk-params fields, a pre-built TextChunkParams / HtmlChunkParams instance (passed through verbatim), None, or False.

Source code in nemo_retriever/common/params/utils.py

def resolve_split_params(
    split_config: dict[str, Any] | None,
) -> dict[str, Any]:
    """Resolve a user-supplied split_config dict into per-key effective params.

    Returns a dict keyed by every entry in ``SPLIT_CONFIG_VALID_KEYS``. Each
    value is one of: a ``TextChunkParams`` / ``HtmlChunkParams`` instance
    (chunking enabled for that key), ``None`` (key absent — chunking off
    via the default), or ``False`` (explicit opt-out sentinel).

    Per-key values supplied by the caller may be a plain dict of
    chunk-params fields, a pre-built ``TextChunkParams`` /
    ``HtmlChunkParams`` instance (passed through verbatim), ``None``, or
    ``False``.
    """
    from nemo_retriever.common.params.models import HtmlChunkParams, TextChunkParams

    cfg = split_config or {}
    unknown = set(cfg) - SPLIT_CONFIG_VALID_KEYS
    if unknown:
        raise ValueError(
            f"Unknown split_config key(s): {sorted(unknown)}; " f"expected one of {sorted(SPLIT_CONFIG_VALID_KEYS)}"
        )

    out: dict[str, Any] = {}
    for key in SPLIT_CONFIG_VALID_KEYS:
        v = cfg.get(key)
        if v is None:
            out[key] = None
            continue
        if v is False:
            out[key] = False  # explicit opt-out (distinct from None / absent)
            continue
        if isinstance(v, TextChunkParams):  # HtmlChunkParams is a TextChunkParams subclass
            out[key] = v
            continue
        if isinstance(v, dict):
            cls = HtmlChunkParams if key == "html" else TextChunkParams
            out[key] = cls(**v)
            continue
        raise TypeError(
            f"split_config['{key}'] must be a TextChunkParams, dict, None, or False; got {type(v).__name__}"
        )
    return out