From 8de23acad9833abb2734f43c8f865904292b9f60 Mon Sep 17 00:00:00 2001 From: Chris Jarrett Date: Thu, 2 Jul 2026 09:48:27 -0700 Subject: [PATCH 1/2] Add embed model name prefix for LiteLLM proxy compatability --- nemo_retriever/helm/README.md | 1 + nemo_retriever/helm/templates/configmap.yaml | 2 + .../helm/templates/deployment-vectordb.yaml | 4 ++ nemo_retriever/helm/values.yaml | 1 + .../cli/ingest/graph_commands.py | 1 + .../src/nemo_retriever/cli/ingest/options.py | 7 +++ .../src/nemo_retriever/cli/query/app.py | 3 + .../src/nemo_retriever/cli/query/options.py | 7 +++ .../api/util/string_processing/__init__.py | 11 ++++ .../nemo_retriever/common/params/models.py | 1 + .../src/nemo_retriever/common/params/utils.py | 11 ++++ .../src/nemo_retriever/ingest/plan.py | 2 + .../models/inference/main_text_embed.py | 21 ++++++- .../models/inference/runtime.py | 6 ++ .../nemo_retriever/models/inference/shared.py | 5 +- .../models/nim/util/__init__.py | 7 ++- .../src/nemo_retriever/query/agentic.py | 2 + .../src/nemo_retriever/query/options.py | 1 + .../src/nemo_retriever/query/workflow.py | 8 ++- .../src/nemo_retriever/service/config.py | 8 +++ .../service/retriever-service.yaml | 3 + .../service/service_ingestor.py | 1 + .../service/services/pipeline_executor.py | 3 + .../nemo_retriever/service/vectordb_app.py | 9 +++ nemo_retriever/tests/test_embed_params.py | 57 +++++++++++++++++++ .../tests/test_root_cli_workflow.py | 6 +- nemo_retriever/tests/test_root_query_cli.py | 6 +- .../tests/test_service_local_models.py | 2 + .../tests/test_service_vectordb_app.py | 26 +++++++++ 29 files changed, 212 insertions(+), 10 deletions(-) diff --git a/nemo_retriever/helm/README.md b/nemo_retriever/helm/README.md index 8059ffc3d..a8e3e2923 100644 --- a/nemo_retriever/helm/README.md +++ b/nemo_retriever/helm/README.md @@ -331,6 +331,7 @@ The retriever service picks up the in-cluster ASR endpoint when `nimOperator.aud | `serviceConfig.vectordb.enabled` | `true` | Deploy the LanceDB vectordb Pod. When `true` the chart **requires** a resolvable embed endpoint (refer to [VectorDB and the embed endpoint](#vectordb-and-the-embed-endpoint)); `helm install` / `helm upgrade` fails fast otherwise. | | `serviceConfig.vectordb.lancedbUri` | `/data/vectordb` | LanceDB on the vectordb Pod's PVC. | | `serviceConfig.vectordb.embedModel` | `nvidia/llama-nemotron-embed-vl-1b-v2` | Passed to vectordb + worker `embed_model_name`. | +| `serviceConfig.vectordb.embedModelNamePrefix` | `""` | Optional LiteLLM route/org prefix prepended to the remote embed model name. | #### VectorDB and the embed endpoint { #vectordb-and-the-embed-endpoint } diff --git a/nemo_retriever/helm/templates/configmap.yaml b/nemo_retriever/helm/templates/configmap.yaml index a5c58bd77..8135a0315 100644 --- a/nemo_retriever/helm/templates/configmap.yaml +++ b/nemo_retriever/helm/templates/configmap.yaml @@ -67,6 +67,7 @@ nim_endpoints: ocr_invoke_url: {{ .ocrURL | quote }} embed_invoke_url: {{ .embedURL | quote }} embed_model_name: {{ .Values.serviceConfig.vectordb.embedModel | quote }} + embed_model_name_prefix: {{ if .Values.serviceConfig.vectordb.embedModelNamePrefix }}{{ .Values.serviceConfig.vectordb.embedModelNamePrefix | quote }}{{ else }}null{{ end }} caption_invoke_url: {{ if .captionURL }}{{ .captionURL | quote }}{{ else }}null{{ end }} caption_model_name: {{ if .captionModelName }}{{ .captionModelName | quote }}{{ else }}null{{ end }} audio_grpc_endpoint: {{ if .audioGrpcEndpoint }}{{ .audioGrpcEndpoint | quote }}{{ else }}null{{ end }} @@ -140,6 +141,7 @@ vectordb: lancedb_uri: {{ .Values.serviceConfig.vectordb.lancedbUri | quote }} table_name: {{ .Values.serviceConfig.vectordb.tableName | quote }} embed_model: {{ .Values.serviceConfig.vectordb.embedModel | quote }} + embed_model_name_prefix: {{ if .Values.serviceConfig.vectordb.embedModelNamePrefix }}{{ .Values.serviceConfig.vectordb.embedModelNamePrefix | quote }}{{ else }}null{{ end }} vectordb_url: "http://{{ .vectordbSvc }}:{{ .vectordbPort }}" {{- else }} vectordb: diff --git a/nemo_retriever/helm/templates/deployment-vectordb.yaml b/nemo_retriever/helm/templates/deployment-vectordb.yaml index ebee3b272..30df1cedc 100644 --- a/nemo_retriever/helm/templates/deployment-vectordb.yaml +++ b/nemo_retriever/helm/templates/deployment-vectordb.yaml @@ -95,6 +95,10 @@ spec: {{- end }} - --embed-model - {{ .Values.serviceConfig.vectordb.embedModel | default $localModels.embed.modelName | quote }} + {{- if .Values.serviceConfig.vectordb.embedModelNamePrefix }} + - --embed-model-name-prefix + - {{ .Values.serviceConfig.vectordb.embedModelNamePrefix | quote }} + {{- end }} - --port - {{ $vdb.port | quote }} {{- if $embedURL }} diff --git a/nemo_retriever/helm/values.yaml b/nemo_retriever/helm/values.yaml index 7811bd20f..fe19da545 100644 --- a/nemo_retriever/helm/values.yaml +++ b/nemo_retriever/helm/values.yaml @@ -660,6 +660,7 @@ serviceConfig: lancedbUri: "/data/vectordb" tableName: "nemo_retriever" embedModel: "nvidia/llama-nemotron-embed-vl-1b-v2" + embedModelNamePrefix: "" # Optional bearer-token authentication. When apiToken is set, every # request must carry "Authorization: Bearer ". diff --git a/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py b/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py index b0e36bedc..226c3fa39 100644 --- a/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py +++ b/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py @@ -266,6 +266,7 @@ def _graph_ingest_command( table_output_format: opts.TableOutputFormatOption = None, embed_invoke_url: opts.EmbedInvokeUrlOption = None, embed_model_name: opts.EmbedModelNameOption = None, + embed_model_name_prefix: opts.EmbedModelNamePrefixOption = None, local_ingest_embed_backend: opts.LocalIngestEmbedBackendOption = None, embed_modality: opts.EmbedModalityOption = None, embed_granularity: opts.EmbedGranularityOption = None, diff --git a/nemo_retriever/src/nemo_retriever/cli/ingest/options.py b/nemo_retriever/src/nemo_retriever/cli/ingest/options.py index 4ccd0c5cf..259380554 100644 --- a/nemo_retriever/src/nemo_retriever/cli/ingest/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/ingest/options.py @@ -288,6 +288,13 @@ help=f"Optional embedding model name override. Defaults to {DEFAULT_EMBED_MODEL} when omitted.", ), ] +EmbedModelNamePrefixOption = Annotated[ + str | None, + typer.Option( + "--embed-model-name-prefix", + help="Optional LiteLLM route/org prefix prepended to the remote embedding model name.", + ), +] LocalIngestEmbedBackendOption = Annotated[ LocalIngestEmbedBackendValue | None, typer.Option( diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 404c8c14d..758a19a19 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -174,6 +174,7 @@ def _local_command( table_name: opts.TableNameOption = "nemo-retriever", embed_invoke_url: opts.EmbedInvokeUrlOption = None, embed_model_name: opts.EmbedModelNameOption = None, + embed_model_name_prefix: opts.EmbedModelNamePrefixOption = None, reranker_invoke_url: opts.RerankerInvokeUrlOption = None, reranker_api_key_env: opts.RerankerApiKeyEnvOption = None, reranker_model_name: opts.RerankerModelNameOption = None, @@ -220,6 +221,7 @@ def _local_command( embed=QueryEmbedOptions( embed_invoke_url=embed_invoke_url, embed_model_name=embed_model_name, + embed_model_name_prefix=embed_model_name_prefix, ), rerank=QueryRerankOptions( enabled=rerank, @@ -261,6 +263,7 @@ def _request() -> QueryRequest: embed=QueryEmbedOptions( embed_invoke_url=embed_invoke_url, embed_model_name=embed_model_name, + embed_model_name_prefix=embed_model_name_prefix, ), rerank=QueryRerankOptions( enabled=rerank, diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index 0c60b78d2..f07dc33a8 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -72,6 +72,13 @@ help=f"Optional embedding model name override. Defaults to {DEFAULT_EMBED_MODEL} when omitted.", ), ] +EmbedModelNamePrefixOption = Annotated[ + str | None, + typer.Option( + "--embed-model-name-prefix", + help="Optional LiteLLM route/org prefix prepended to the remote embedding model name.", + ), +] RerankerInvokeUrlOption = Annotated[ str | None, typer.Option("--reranker-invoke-url", help="Reranker endpoint URL."), diff --git a/nemo_retriever/src/nemo_retriever/common/api/util/string_processing/__init__.py b/nemo_retriever/src/nemo_retriever/common/api/util/string_processing/__init__.py index 6fb7fb5e4..eae7d8b23 100644 --- a/nemo_retriever/src/nemo_retriever/common/api/util/string_processing/__init__.py +++ b/nemo_retriever/src/nemo_retriever/common/api/util/string_processing/__init__.py @@ -51,6 +51,17 @@ def ensure_openai_embeddings_http_url(endpoint_url: str) -> str: return urlunsplit((parts.scheme, parts.netloc, new_path, parts.query, parts.fragment)) +def prepend_model_name_prefix(model_name: str | None, model_name_prefix: str | None) -> str | None: + """Prepend a LiteLLM route/org prefix to a model identifier when configured.""" + if model_name is None: + return None + model = str(model_name).strip() + prefix = str(model_name_prefix or "").strip().strip("/") + if not model or not prefix: + return model + return f"{prefix}/{model.lstrip('/')}" + + def generate_url(url) -> str: """Examines the user defined URL for http*://. If that pattern is detected the URL is used as provided by the user. diff --git a/nemo_retriever/src/nemo_retriever/common/params/models.py b/nemo_retriever/src/nemo_retriever/common/params/models.py index 71c046cee..84491e693 100644 --- a/nemo_retriever/src/nemo_retriever/common/params/models.py +++ b/nemo_retriever/src/nemo_retriever/common/params/models.py @@ -384,6 +384,7 @@ class EmbedParams(_ParamsModel): embedding_endpoint: Optional[str] = None embed_invoke_url: Optional[str] = None embed_model_name: Optional[str] = None + embed_model_name_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 diff --git a/nemo_retriever/src/nemo_retriever/common/params/utils.py b/nemo_retriever/src/nemo_retriever/common/params/utils.py index f7c82f7c8..4290dfc22 100644 --- a/nemo_retriever/src/nemo_retriever/common/params/utils.py +++ b/nemo_retriever/src/nemo_retriever/common/params/utils.py @@ -8,6 +8,8 @@ from typing import TYPE_CHECKING, Any, Dict +from nemo_retriever.common.api.util.string_processing import prepend_model_name_prefix + if TYPE_CHECKING: from nemo_retriever.common.params.models import BatchTuningParams @@ -50,6 +52,12 @@ def normalize_embed_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: if "embed_invoke_url" in normalized: normalized.setdefault("embedding_endpoint", normalized["embed_invoke_url"]) + endpoint = normalized.get("embedding_endpoint") or normalized.get("embed_invoke_url") + model_name_prefix = normalized.pop("embed_model_name_prefix", None) + if endpoint and model_name_prefix: + for key in ("model_name", "embed_model_name"): + if key in normalized: + normalized[key] = prepend_model_name_prefix(normalized[key], str(model_name_prefix)) return normalized @@ -58,6 +66,7 @@ def build_embed_option_kwargs( embed_model_name: str | None, local_ingest_embed_backend: str | None = None, embed_api_key: str | None = None, + embed_model_name_prefix: str | None = None, embed_modality: str | None = None, text_elements_modality: str | None = None, structured_elements_modality: str | None = None, @@ -79,6 +88,8 @@ def build_embed_option_kwargs( 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_name_prefix is not None: + embed_kwargs["embed_model_name_prefix"] = embed_model_name_prefix if embed_modality is not None: embed_kwargs["embed_modality"] = embed_modality if text_elements_modality is not None: diff --git a/nemo_retriever/src/nemo_retriever/ingest/plan.py b/nemo_retriever/src/nemo_retriever/ingest/plan.py index e857f4385..c8ad53b49 100644 --- a/nemo_retriever/src/nemo_retriever/ingest/plan.py +++ b/nemo_retriever/src/nemo_retriever/ingest/plan.py @@ -185,6 +185,7 @@ class IngestEmbedBatchOptions: class IngestEmbedOptions: embed_invoke_url: str | None = None embed_model_name: str | None = None + embed_model_name_prefix: str | None = None local_ingest_embed_backend: LocalIngestEmbedBackendValue | None = None embed_api_key: str | None = None embed_modality: str | None = None @@ -642,6 +643,7 @@ def resolve_ingest_plan(request: IngestPlanRequest) -> ResolvedIngestPlan: embed.embed_model_name, local_ingest_embed_backend=embed.local_ingest_embed_backend, embed_api_key=embed.embed_api_key, + embed_model_name_prefix=embed.embed_model_name_prefix, embed_modality=embed.embed_modality, text_elements_modality=embed.text_elements_modality, structured_elements_modality=embed.structured_elements_modality, diff --git a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py index dd7c29c37..47cd5b818 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py @@ -48,7 +48,10 @@ def local_embedder(texts): from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple import pandas as pd -from nemo_retriever.common.api.util.string_processing import ensure_openai_embeddings_http_url +from nemo_retriever.common.api.util.string_processing import ( + ensure_openai_embeddings_http_url, + prepend_model_name_prefix, +) from nemo_retriever.models import _DEFAULT_EMBED_MODEL from nemo_retriever.common.params.models import IMAGE_MODALITIES @@ -72,6 +75,7 @@ class TextEmbeddingConfig: api_key: Optional[str] = None embedding_nim_endpoint: Optional[str] = None # e.g. "http://host:8000/v1" embedding_model: str = _DEFAULT_EMBED_MODEL + embedding_model_prefix: Optional[str] = None encoding_format: str = "float" # OpenAI-compatible embeddings often accept "float" input_type: str = "passage" truncate: str = "END" @@ -293,6 +297,7 @@ def _http_embed_openai_compat( encoding_format: str, input_type: str, truncate: str, + model_name_prefix: Optional[str] = None, dimensions: Optional[int] = None, timeout_s: float = 600.0, ) -> List[Optional[List[float]]]: @@ -308,6 +313,7 @@ def _http_embed_openai_compat( raise RuntimeError("Remote embedding requested but `httpx` is not installed.") from e url = _normalize_embeddings_endpoint(_pick_embed_endpoint(endpoint_url)) + model_name = prepend_model_name_prefix(model_name, model_name_prefix) or model_name headers: Dict[str, str] = {"accept": "application/json", "content-type": "application/json"} token = (api_key or "").strip() if token: @@ -353,6 +359,7 @@ def _make_async_request( api_key: Optional[str], embedding_nim_endpoint: str, embedding_model: str, + embedding_model_prefix: Optional[str], encoding_format: str, input_type: str, truncate: str, @@ -377,6 +384,7 @@ def _make_async_request( api_key=api_key, endpoint_url=str(embedding_nim_endpoint), model_name=str(embedding_model), + model_name_prefix=embedding_model_prefix, encoding_format=str(encoding_format), input_type=str(input_type), truncate=str(truncate), @@ -399,6 +407,7 @@ def _async_request_handler( api_key: Optional[str], embedding_nim_endpoint: str, embedding_model: str, + embedding_model_prefix: Optional[str], encoding_format: str, input_type: str, truncate: str, @@ -420,6 +429,7 @@ def _async_request_handler( api_key=api_key or None, embedding_nim_endpoint=str(embedding_nim_endpoint), embedding_model=str(embedding_model), + embedding_model_prefix=embedding_model_prefix, encoding_format=str(encoding_format), input_type=str(input_type), truncate=str(truncate), @@ -440,6 +450,7 @@ def _async_runner( api_key: Optional[str], embedding_nim_endpoint: str, embedding_model: str, + embedding_model_prefix: Optional[str], encoding_format: str, input_type: str, truncate: str, @@ -454,6 +465,7 @@ def _async_runner( api_key, embedding_nim_endpoint, embedding_model, + embedding_model_prefix, encoding_format, input_type, truncate, @@ -580,6 +592,11 @@ def create_text_embeddings_for_df( task_config["endpoint_url"] if "endpoint_url" in task_config else transform_config.embedding_nim_endpoint ) model_name = task_config["model_name"] if "model_name" in task_config else transform_config.embedding_model + model_name_prefix = ( + task_config["model_name_prefix"] + if "model_name_prefix" in task_config + else task_config.get("embed_model_name_prefix", transform_config.embedding_model_prefix) + ) dimensions = task_config["dimensions"] if "dimensions" in task_config else transform_config.dimensions endpoint_url = endpoint_url.strip() if isinstance(endpoint_url, str) else endpoint_url @@ -671,6 +688,7 @@ def _text_image_content(r: pd.Series) -> Optional[str]: api_key, str(endpoint_url), str(model_name), + model_name_prefix, str(transform_config.encoding_format), str(transform_config.input_type), str(transform_config.truncate), @@ -693,6 +711,7 @@ def _text_image_content(r: pd.Series) -> Optional[str]: api_key, str(endpoint_url), str(model_name), + model_name_prefix, str(transform_config.encoding_format), str(transform_config.input_type), str(transform_config.truncate), diff --git a/nemo_retriever/src/nemo_retriever/models/inference/runtime.py b/nemo_retriever/src/nemo_retriever/models/inference/runtime.py index 1206cbcee..e73a476f6 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/runtime.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/runtime.py @@ -30,6 +30,7 @@ def _embed_group( inference_batch_size: int, output_column: str, resolved_model_name: str, + embed_model_name_prefix: Optional[str] = None, nim_http_max_concurrent: int = 32, input_type: str = "passage", request_timeout_s: float = 600.0, @@ -76,6 +77,7 @@ def embedder(texts: Sequence[str]) -> Sequence[Sequence[float]]: # noqa dimensions=None, embedding_nim_endpoint=endpoint or "http://localhost:8012/v1", embedding_model=resolved_model_name or VL_EMBED_MODEL, + embedding_model_prefix=embed_model_name_prefix, embed_modality=group_modality, nim_http_max_concurrent=max(1, int(nim_http_max_concurrent)), ) @@ -87,6 +89,7 @@ def embedder(texts: Sequence[str]) -> Sequence[Sequence[float]]: # noqa "embedder": embedder, "multimodal_embedder": multimodal_embedder, "endpoint_url": endpoint, + "embed_model_name_prefix": embed_model_name_prefix, "local_batch_size": int(effective_batch_size), "nim_http_max_concurrent": max(1, int(nim_http_max_concurrent)), "request_timeout_s": float(request_timeout_s), @@ -110,6 +113,7 @@ def embed_text_main_text_embed( embedding_dim_column: str = "text_embeddings_1b_v2_dim", has_embedding_column: str = "text_embeddings_1b_v2_has_embedding", embed_modality: str = "text", + embed_model_name_prefix: Optional[str] = None, nim_http_max_concurrent: int = 32, input_type: str = "passage", request_timeout_s: float | None = None, @@ -148,6 +152,7 @@ def embed_text_main_text_embed( inference_batch_size=inference_batch_size, output_column=output_column, resolved_model_name=resolved_model_name, + embed_model_name_prefix=embed_model_name_prefix, nim_http_max_concurrent=nim_http_max_concurrent, input_type=input_type, request_timeout_s=float(request_timeout_s), @@ -169,6 +174,7 @@ def embed_text_main_text_embed( inference_batch_size=inference_batch_size, output_column=output_column, resolved_model_name=resolved_model_name, + embed_model_name_prefix=embed_model_name_prefix, nim_http_max_concurrent=nim_http_max_concurrent, input_type=input_type, request_timeout_s=float(request_timeout_s), diff --git a/nemo_retriever/src/nemo_retriever/models/inference/shared.py b/nemo_retriever/src/nemo_retriever/models/inference/shared.py index 21d042dbf..22066a1c0 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/shared.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/shared.py @@ -7,6 +7,7 @@ from __future__ import annotations from nemo_retriever.common.params import EmbedParams +from nemo_retriever.common.params.utils import normalize_embed_kwargs def _to_bool(v: object, default: bool = False) -> bool: @@ -22,6 +23,4 @@ def build_embed_kwargs(params: EmbedParams) -> dict[str, object]: **params.model_dump(mode="python", exclude={"runtime", "batch_tuning"}, exclude_none=True), **params.runtime.model_dump(mode="python", exclude_none=True), } - if "embedding_endpoint" not in kwargs and kwargs.get("embed_invoke_url"): - kwargs["embedding_endpoint"] = kwargs.get("embed_invoke_url") - return kwargs + return normalize_embed_kwargs(kwargs) diff --git a/nemo_retriever/src/nemo_retriever/models/nim/util/__init__.py b/nemo_retriever/src/nemo_retriever/models/nim/util/__init__.py index 20394de67..fd848a237 100644 --- a/nemo_retriever/src/nemo_retriever/models/nim/util/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/nim/util/__init__.py @@ -13,7 +13,10 @@ from nemo_retriever.models.nim.primitives.model_interface.text_embedding import EmbeddingModelInterface from nemo_retriever.models.nim.primitives.nim_client import get_nim_client_manager from nemo_retriever.models.nim.primitives.nim_model_interface import ModelInterface -from nemo_retriever.common.api.util.string_processing import ensure_openai_embeddings_http_url +from nemo_retriever.common.api.util.string_processing import ( + ensure_openai_embeddings_http_url, + prepend_model_name_prefix, +) logger = logging.getLogger(__name__) @@ -138,6 +141,7 @@ def infer_microservice( input_type: str = "passage", truncate: str = "END", batch_size: int = 8191, + model_name_prefix: str | None = None, grpc: bool = False, input_names: list = ["text"], output_names: list = ["embeddings"], @@ -200,6 +204,7 @@ def infer_microservice( ) else: embedding_endpoint = ensure_openai_embeddings_http_url(str(embedding_endpoint)) + model_name = prepend_model_name_prefix(model_name, model_name_prefix) client = NimClient( model_interface=EmbeddingModelInterface(), protocol="http", diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index aa7503447..2ca6e1b28 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic.py @@ -113,6 +113,7 @@ class AgenticRetrievalConfig: vdb_op: str = "lancedb" vdb_kwargs: dict[str, Any] = field(default_factory=dict) query_embedder: str = VL_EMBED_MODEL + query_embedder_prefix: Optional[str] = None embedding_endpoint: Optional[str] = None embedding_api_key: str = "" local_hf_batch_size: int = 32 @@ -174,6 +175,7 @@ def __init__( embed_kwargs={ "model_name": str(cfg.query_embedder or VL_EMBED_MODEL), "embed_model_name": str(cfg.query_embedder or VL_EMBED_MODEL), + "embed_model_name_prefix": cfg.query_embedder_prefix, "embedding_endpoint": cfg.embedding_endpoint, "api_key": cfg.embedding_api_key, "input_type": "query", diff --git a/nemo_retriever/src/nemo_retriever/query/options.py b/nemo_retriever/src/nemo_retriever/query/options.py index 237001a6e..6922eae18 100644 --- a/nemo_retriever/src/nemo_retriever/query/options.py +++ b/nemo_retriever/src/nemo_retriever/query/options.py @@ -26,6 +26,7 @@ class QueryRetrievalOptions: class QueryEmbedOptions: embed_invoke_url: str | None = None embed_model_name: str | None = None + embed_model_name_prefix: str | None = None @dataclass(frozen=True) diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index 5c777d537..4ed9c2f33 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/nemo_retriever/src/nemo_retriever/query/workflow.py @@ -102,7 +102,11 @@ def query_kwargs(self) -> dict[str, Any]: def resolve_query_plan(request: QueryRequest) -> ResolvedQueryPlan: """Resolve root query options once so callers can reuse a Retriever.""" - embed_kwargs = build_embed_option_kwargs(request.embed.embed_invoke_url, request.embed.embed_model_name) + embed_kwargs = build_embed_option_kwargs( + request.embed.embed_invoke_url, + request.embed.embed_model_name, + embed_model_name_prefix=request.embed.embed_model_name_prefix, + ) rerank_kwargs = _build_rerank_kwargs(request.rerank) if request.rerank.enabled else {} content_types = request.retrieval.content_types if content_types is not None and not isinstance(content_types, str): @@ -179,6 +183,8 @@ def agentic_query_documents(request: QueryRequest) -> list[dict[str, Any]]: } if request.embed.embed_model_name: cfg_kwargs["query_embedder"] = request.embed.embed_model_name + if request.embed.embed_model_name_prefix: + cfg_kwargs["query_embedder_prefix"] = request.embed.embed_model_name_prefix if request.rerank.enabled: # `reranker` doubles as the on/off gate (rerank=bool(cfg.reranker)) and the # model name, so fall back to the default model when only --rerank is given. diff --git a/nemo_retriever/src/nemo_retriever/service/config.py b/nemo_retriever/src/nemo_retriever/service/config.py index 08f306608..7b62e8e46 100644 --- a/nemo_retriever/src/nemo_retriever/service/config.py +++ b/nemo_retriever/src/nemo_retriever/service/config.py @@ -141,6 +141,13 @@ class NimEndpointsConfig(RichModel): "Server-owned — clients cannot override the deployed embed NIM SKU." ), ) + embed_model_name_prefix: str | None = Field( + default=None, + description=( + "Optional LiteLLM route/org prefix prepended to embed_model_name for " + "remote embedding endpoints that require namespaced model IDs." + ), + ) rerank_invoke_url: str | None = None audio_grpc_endpoint: str | None = Field( default=None, @@ -293,6 +300,7 @@ class VectorDbConfig(RichModel): lancedb_uri: str = "/data/vectordb" table_name: str = "nemo_retriever" embed_model: str = "nvidia/llama-nemotron-embed-vl-1b-v2" + embed_model_name_prefix: str | None = None vectordb_url: str = Field( default="http://nemo-retriever-vectordb:7671", description="URL of the vectordb service (for workers to POST embeddings to)", diff --git a/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml b/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml index 3a797bf7d..4277391fd 100644 --- a/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml +++ b/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml @@ -37,6 +37,9 @@ nim_endpoints: embed_invoke_url: null # Model name for the remote embed NIM (server-owned; must match the SKU). embed_model_name: null + # Optional LiteLLM route/org prefix prepended to embed_model_name for + # proxies that require namespaced model IDs (e.g. "team/model/provider"). + embed_model_name_prefix: null # gRPC endpoint for the Parakeet ASR NIM (e.g. parakeet-nim:50051). # When set, audio/video pipelines use remote ASR instead of loading # the local Parakeet model (which requires torch + GPU). diff --git a/nemo_retriever/src/nemo_retriever/service/service_ingestor.py b/nemo_retriever/src/nemo_retriever/service/service_ingestor.py index e53112c64..1e858cc93 100644 --- a/nemo_retriever/src/nemo_retriever/service/service_ingestor.py +++ b/nemo_retriever/src/nemo_retriever/service/service_ingestor.py @@ -219,6 +219,7 @@ def _normalize_files(files: Union[str, List[str], List[Path]]) -> list[Path]: "nemotron_parse_invoke_url", "embed_invoke_url", "embedding_endpoint", + "embed_model_name_prefix", "endpoint_url", "api_base", "auth_token", diff --git a/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py b/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py index 2e01fae0f..6841718f1 100644 --- a/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py +++ b/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py @@ -318,6 +318,7 @@ def _post_rows_to_vectordb(rows: list[dict[str, Any]], vectordb_url: str, filena "embedding_endpoint", "api_key", "embed_model_name", + "embed_model_name_prefix", "model_name", ) # Trust-owned caption keys. ``endpoint_url`` / ``api_key`` / @@ -851,6 +852,8 @@ def build_embed_params(nim: "NimEndpointsConfig", local: "LocalModelsConfig | No if nim.embed_model_name: kwargs["model_name"] = nim.embed_model_name kwargs["embed_model_name"] = nim.embed_model_name + if nim.embed_model_name_prefix: + kwargs["embed_model_name_prefix"] = nim.embed_model_name_prefix if nim.api_key: kwargs["api_key"] = nim.api_key return EmbedParams(**kwargs) diff --git a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py index 8221d2c29..a083abcf5 100644 --- a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py +++ b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py @@ -90,12 +90,14 @@ def _embed_queries_remote( embed_model: str, embed_endpoint: str, embed_api_key: str, + embed_model_name_prefix: str | None = None, ) -> list[list[float]]: from nemo_retriever.models.nim.util import infer_microservice return infer_microservice( texts, model_name=embed_model, + model_name_prefix=embed_model_name_prefix, embedding_endpoint=embed_endpoint, nvidia_api_key=embed_api_key or None, input_type="query", @@ -117,6 +119,7 @@ def __init__( embed_model: str, embed_api_key: str, *, + embed_model_name_prefix: str | None = None, local_embed: bool = False, local_embed_backend: str = "hf", hf_cache_dir: str | None = None, @@ -127,6 +130,7 @@ def __init__( self.table_name = table_name self.embed_endpoint = embed_endpoint self.embed_model = embed_model + self.embed_model_name_prefix = embed_model_name_prefix self.embed_api_key = embed_api_key self.local_embed = local_embed self.local_embed_backend = local_embed_backend @@ -244,6 +248,7 @@ def embed_queries(self, texts: list[str]) -> list[list[float]]: return _embed_queries_remote( texts, embed_model=self.embed_model, + embed_model_name_prefix=self.embed_model_name_prefix, embed_endpoint=self.embed_endpoint, embed_api_key=self.embed_api_key, ) @@ -268,6 +273,7 @@ def create_vectordb_app( table_name: str = "nemo_retriever", embed_endpoint: str = "", embed_model: str = "nvidia/llama-nemotron-embed-vl-1b-v2", + embed_model_name_prefix: str | None = None, embed_api_key: str = "", *, local_embed: bool = False, @@ -286,6 +292,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: table_name=table_name, embed_endpoint=embed_endpoint, embed_model=embed_model, + embed_model_name_prefix=embed_model_name_prefix, embed_api_key=embed_api_key, local_embed=local_embed, local_embed_backend=local_embed_backend, @@ -384,6 +391,7 @@ def main() -> None: parser.add_argument("--table-name", default="nemo_retriever", help="LanceDB table name") parser.add_argument("--embed-endpoint", default="", help="Remote NIM/OpenAI-compatible embed URL") parser.add_argument("--embed-model", default="nvidia/llama-nemotron-embed-vl-1b-v2") + parser.add_argument("--embed-model-name-prefix", default="", help="Optional LiteLLM route/org prefix") parser.add_argument("--embed-api-key", default="") parser.add_argument( "--local-embed", @@ -422,6 +430,7 @@ def main() -> None: table_name=args.table_name, embed_endpoint=args.embed_endpoint, embed_model=args.embed_model, + embed_model_name_prefix=args.embed_model_name_prefix or None, embed_api_key=args.embed_api_key, local_embed=args.local_embed, local_embed_backend=args.local_embed_backend, diff --git a/nemo_retriever/tests/test_embed_params.py b/nemo_retriever/tests/test_embed_params.py index 596c1b06f..cb58c8a04 100644 --- a/nemo_retriever/tests/test_embed_params.py +++ b/nemo_retriever/tests/test_embed_params.py @@ -11,6 +11,7 @@ import pytest from nemo_retriever.common.params.models import EmbedParams, IMAGE_MODALITIES +from nemo_retriever.common.params.utils import build_embed_option_kwargs def test_image_text_alias_is_rejected(): @@ -54,6 +55,62 @@ def test_image_modalities_constant(): assert isinstance(IMAGE_MODALITIES, frozenset) +def test_build_embed_option_kwargs_applies_remote_model_name_prefix(): + kwargs = build_embed_option_kwargs( + "https://litellm.example.com/v1/embeddings", + "nvidia/llama-nemotron-embed-vl-1b-v2", + embed_model_name_prefix="nvidia", + ) + + assert kwargs["model_name"] == "nvidia/nvidia/llama-nemotron-embed-vl-1b-v2" + assert kwargs["embed_model_name"] == "nvidia/nvidia/llama-nemotron-embed-vl-1b-v2" + assert "embed_model_name_prefix" not in kwargs + + +def test_build_embed_option_kwargs_leaves_model_unchanged_without_prefix(): + kwargs = build_embed_option_kwargs( + "https://integrate.api.nvidia.com/v1/embeddings", + "nvidia/llama-nemotron-embed-vl-1b-v2", + ) + + assert kwargs["model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert kwargs["embed_model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" + + +def test_build_embed_option_kwargs_prefix_supports_other_vendor_namespaces(): + kwargs = build_embed_option_kwargs( + "https://litellm.example.com/v1/embeddings", + "mistral/embed-small", + embed_model_name_prefix="acme", + ) + + assert kwargs["model_name"] == "acme/mistral/embed-small" + assert kwargs["embed_model_name"] == "acme/mistral/embed-small" + + +def test_build_embed_option_kwargs_prefix_supports_bare_model_name(): + kwargs = build_embed_option_kwargs( + "https://litellm.example.com/v1/embeddings", + "nv-embedqa-e5-v5", + embed_model_name_prefix="nvidia", + ) + + assert kwargs["model_name"] == "nvidia/nv-embedqa-e5-v5" + assert kwargs["embed_model_name"] == "nvidia/nv-embedqa-e5-v5" + + +def test_build_embed_option_kwargs_prefix_is_remote_only(): + kwargs = build_embed_option_kwargs( + None, + "nvidia/llama-nemotron-embed-vl-1b-v2", + embed_model_name_prefix="nvidia", + ) + + assert kwargs["model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert kwargs["embed_model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert "embed_model_name_prefix" not in kwargs + + # =================================================================== # embed_granularity # =================================================================== diff --git a/nemo_retriever/tests/test_root_cli_workflow.py b/nemo_retriever/tests/test_root_cli_workflow.py index f0decca4a..1eeb9af2d 100644 --- a/nemo_retriever/tests/test_root_cli_workflow.py +++ b/nemo_retriever/tests/test_root_cli_workflow.py @@ -434,6 +434,8 @@ def fake_create_ingestor(**_kwargs: Any) -> Any: "http://embed:8000/v1/embeddings", "--embed-model-name", "nvidia/llama-nemotron-embed-1b-v2", + "--embed-model-name-prefix", + "nvidia", ], ) @@ -450,8 +452,8 @@ def fake_create_ingestor(**_kwargs: Any) -> Any: assert isinstance(embed_params, EmbedParams) assert embed_params.embed_invoke_url == "http://embed:8000/v1/embeddings" assert embed_params.embedding_endpoint == "http://embed:8000/v1/embeddings" - assert embed_params.model_name == "nvidia/llama-nemotron-embed-1b-v2" - assert embed_params.embed_model_name == "nvidia/llama-nemotron-embed-1b-v2" + assert embed_params.model_name == "nvidia/nvidia/llama-nemotron-embed-1b-v2" + assert embed_params.embed_model_name == "nvidia/nvidia/llama-nemotron-embed-1b-v2" def test_root_ingest_passes_migrated_extraction_and_embedding_flags(monkeypatch, tmp_path) -> None: diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index 28bfeec9b..70d1fee29 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -136,6 +136,8 @@ def query(self, query: str, **_kwargs: Any) -> list[dict[str, Any]]: "http://embed:8000/v1/embeddings", "--embed-model-name", "nvidia/llama-nemotron-embed-1b-v2", + "--embed-model-name-prefix", + "nvidia", ], ) @@ -148,8 +150,8 @@ def query(self, query: str, **_kwargs: Any) -> list[dict[str, Any]]: "embed_kwargs": { "embed_invoke_url": "http://embed:8000/v1/embeddings", "embedding_endpoint": "http://embed:8000/v1/embeddings", - "model_name": "nvidia/llama-nemotron-embed-1b-v2", - "embed_model_name": "nvidia/llama-nemotron-embed-1b-v2", + "model_name": "nvidia/nvidia/llama-nemotron-embed-1b-v2", + "embed_model_name": "nvidia/nvidia/llama-nemotron-embed-1b-v2", }, } ] diff --git a/nemo_retriever/tests/test_service_local_models.py b/nemo_retriever/tests/test_service_local_models.py index e7abd6e3b..43555a656 100644 --- a/nemo_retriever/tests/test_service_local_models.py +++ b/nemo_retriever/tests/test_service_local_models.py @@ -27,6 +27,7 @@ def test_build_embed_params_from_nim_config() -> None: nim = NimEndpointsConfig( embed_invoke_url="http://embed-nim/v1/embeddings", embed_model_name="nvidia/llama-nemotron-embed-vl-1b-v2", + embed_model_name_prefix="nvidia", api_key="k", ) ep = build_embed_params(nim, LocalModelsConfig(enabled=True)) @@ -34,6 +35,7 @@ def test_build_embed_params_from_nim_config() -> None: assert ep.embed_invoke_url == "http://embed-nim/v1/embeddings" assert ep.model_name == "nvidia/llama-nemotron-embed-vl-1b-v2" assert ep.embed_model_name == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert ep.embed_model_name_prefix == "nvidia" assert ep.api_key == "k" diff --git a/nemo_retriever/tests/test_service_vectordb_app.py b/nemo_retriever/tests/test_service_vectordb_app.py index 1ea0a12d5..a9fd8d89f 100644 --- a/nemo_retriever/tests/test_service_vectordb_app.py +++ b/nemo_retriever/tests/test_service_vectordb_app.py @@ -10,6 +10,7 @@ from nemo_retriever.service.vectordb_app import ( VectorDBState, + _embed_queries_remote, _tensor_to_embedding_rows, create_vectordb_app, ) @@ -84,6 +85,31 @@ def test_vector_db_state_local_embed_queries() -> None: mock_embedder.embed_queries.assert_called_once_with(["hello"]) +def test_remote_embed_queries_delegates_model_prefix(monkeypatch) -> None: + calls = {} + + def fake_infer_microservice(data, **kwargs): + calls["data"] = data + calls.update(kwargs) + return [[0.1, 0.2]] + + monkeypatch.setattr("nemo_retriever.models.nim.util.infer_microservice", fake_infer_microservice) + + vectors = _embed_queries_remote( + ["hello"], + embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", + embed_endpoint="https://litellm.example.com/v1/embeddings", + embed_api_key="k", + embed_model_name_prefix="nvidia", + ) + + assert vectors == [[0.1, 0.2]] + assert calls["data"] == ["hello"] + assert calls["model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" + assert calls["model_name_prefix"] == "nvidia" + assert calls["embedding_endpoint"] == "https://litellm.example.com/v1/embeddings" + + _CANNED_HITS = [ { "text": "Revenue grew 12% year over year.", From 01182146892e927eeb433a69a07fa6c1d4c9da2c Mon Sep 17 00:00:00 2001 From: Chris Jarrett Date: Mon, 6 Jul 2026 12:26:00 -0700 Subject: [PATCH 2/2] Change parameter name to embed_model_provider_prefix --- nemo_retriever/helm/README.md | 2 +- nemo_retriever/helm/templates/configmap.yaml | 4 +-- .../helm/templates/deployment-vectordb.yaml | 6 ++-- nemo_retriever/helm/values.yaml | 2 +- .../cli/ingest/graph_commands.py | 2 +- .../src/nemo_retriever/cli/ingest/options.py | 6 ++-- .../src/nemo_retriever/cli/query/app.py | 6 ++-- .../src/nemo_retriever/cli/query/options.py | 6 ++-- .../api/util/string_processing/__init__.py | 6 ++-- .../nemo_retriever/common/params/models.py | 2 +- .../src/nemo_retriever/common/params/utils.py | 14 ++++---- .../src/nemo_retriever/ingest/plan.py | 4 +-- .../models/inference/main_text_embed.py | 32 +++++++++---------- .../models/inference/runtime.py | 12 +++---- .../models/nim/util/__init__.py | 6 ++-- .../src/nemo_retriever/query/agentic.py | 4 +-- .../src/nemo_retriever/query/options.py | 2 +- .../src/nemo_retriever/query/workflow.py | 6 ++-- .../src/nemo_retriever/service/config.py | 6 ++-- .../service/retriever-service.yaml | 6 ++-- .../service/service_ingestor.py | 2 +- .../service/services/pipeline_executor.py | 6 ++-- .../nemo_retriever/service/vectordb_app.py | 18 +++++------ nemo_retriever/tests/test_embed_params.py | 14 ++++---- .../tests/test_root_cli_workflow.py | 2 +- nemo_retriever/tests/test_root_query_cli.py | 2 +- .../tests/test_service_local_models.py | 4 +-- .../tests/test_service_vectordb_app.py | 4 +-- 28 files changed, 93 insertions(+), 93 deletions(-) diff --git a/nemo_retriever/helm/README.md b/nemo_retriever/helm/README.md index a8e3e2923..5ba9b9291 100644 --- a/nemo_retriever/helm/README.md +++ b/nemo_retriever/helm/README.md @@ -331,7 +331,7 @@ The retriever service picks up the in-cluster ASR endpoint when `nimOperator.aud | `serviceConfig.vectordb.enabled` | `true` | Deploy the LanceDB vectordb Pod. When `true` the chart **requires** a resolvable embed endpoint (refer to [VectorDB and the embed endpoint](#vectordb-and-the-embed-endpoint)); `helm install` / `helm upgrade` fails fast otherwise. | | `serviceConfig.vectordb.lancedbUri` | `/data/vectordb` | LanceDB on the vectordb Pod's PVC. | | `serviceConfig.vectordb.embedModel` | `nvidia/llama-nemotron-embed-vl-1b-v2` | Passed to vectordb + worker `embed_model_name`. | -| `serviceConfig.vectordb.embedModelNamePrefix` | `""` | Optional LiteLLM route/org prefix prepended to the remote embed model name. | +| `serviceConfig.vectordb.embedModelProviderPrefix` | `""` | Optional LiteLLM provider prefix prepended to the remote embed model name. | #### VectorDB and the embed endpoint { #vectordb-and-the-embed-endpoint } diff --git a/nemo_retriever/helm/templates/configmap.yaml b/nemo_retriever/helm/templates/configmap.yaml index 8135a0315..ba0410174 100644 --- a/nemo_retriever/helm/templates/configmap.yaml +++ b/nemo_retriever/helm/templates/configmap.yaml @@ -67,7 +67,7 @@ nim_endpoints: ocr_invoke_url: {{ .ocrURL | quote }} embed_invoke_url: {{ .embedURL | quote }} embed_model_name: {{ .Values.serviceConfig.vectordb.embedModel | quote }} - embed_model_name_prefix: {{ if .Values.serviceConfig.vectordb.embedModelNamePrefix }}{{ .Values.serviceConfig.vectordb.embedModelNamePrefix | quote }}{{ else }}null{{ end }} + embed_model_provider_prefix: {{ if .Values.serviceConfig.vectordb.embedModelProviderPrefix }}{{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }}{{ else }}null{{ end }} caption_invoke_url: {{ if .captionURL }}{{ .captionURL | quote }}{{ else }}null{{ end }} caption_model_name: {{ if .captionModelName }}{{ .captionModelName | quote }}{{ else }}null{{ end }} audio_grpc_endpoint: {{ if .audioGrpcEndpoint }}{{ .audioGrpcEndpoint | quote }}{{ else }}null{{ end }} @@ -141,7 +141,7 @@ vectordb: lancedb_uri: {{ .Values.serviceConfig.vectordb.lancedbUri | quote }} table_name: {{ .Values.serviceConfig.vectordb.tableName | quote }} embed_model: {{ .Values.serviceConfig.vectordb.embedModel | quote }} - embed_model_name_prefix: {{ if .Values.serviceConfig.vectordb.embedModelNamePrefix }}{{ .Values.serviceConfig.vectordb.embedModelNamePrefix | quote }}{{ else }}null{{ end }} + embed_model_provider_prefix: {{ if .Values.serviceConfig.vectordb.embedModelProviderPrefix }}{{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }}{{ else }}null{{ end }} vectordb_url: "http://{{ .vectordbSvc }}:{{ .vectordbPort }}" {{- else }} vectordb: diff --git a/nemo_retriever/helm/templates/deployment-vectordb.yaml b/nemo_retriever/helm/templates/deployment-vectordb.yaml index 30df1cedc..19c4efc00 100644 --- a/nemo_retriever/helm/templates/deployment-vectordb.yaml +++ b/nemo_retriever/helm/templates/deployment-vectordb.yaml @@ -95,9 +95,9 @@ spec: {{- end }} - --embed-model - {{ .Values.serviceConfig.vectordb.embedModel | default $localModels.embed.modelName | quote }} - {{- if .Values.serviceConfig.vectordb.embedModelNamePrefix }} - - --embed-model-name-prefix - - {{ .Values.serviceConfig.vectordb.embedModelNamePrefix | quote }} + {{- if .Values.serviceConfig.vectordb.embedModelProviderPrefix }} + - --embed-model-provider-prefix + - {{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }} {{- end }} - --port - {{ $vdb.port | quote }} diff --git a/nemo_retriever/helm/values.yaml b/nemo_retriever/helm/values.yaml index fe19da545..685d7bfd2 100644 --- a/nemo_retriever/helm/values.yaml +++ b/nemo_retriever/helm/values.yaml @@ -660,7 +660,7 @@ serviceConfig: lancedbUri: "/data/vectordb" tableName: "nemo_retriever" embedModel: "nvidia/llama-nemotron-embed-vl-1b-v2" - embedModelNamePrefix: "" + embedModelProviderPrefix: "" # Optional bearer-token authentication. When apiToken is set, every # request must carry "Authorization: Bearer ". diff --git a/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py b/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py index 226c3fa39..2c5eac1b6 100644 --- a/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py +++ b/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py @@ -266,7 +266,7 @@ def _graph_ingest_command( table_output_format: opts.TableOutputFormatOption = None, embed_invoke_url: opts.EmbedInvokeUrlOption = None, embed_model_name: opts.EmbedModelNameOption = None, - embed_model_name_prefix: opts.EmbedModelNamePrefixOption = None, + embed_model_provider_prefix: opts.EmbedModelProviderPrefixOption = None, local_ingest_embed_backend: opts.LocalIngestEmbedBackendOption = None, embed_modality: opts.EmbedModalityOption = None, embed_granularity: opts.EmbedGranularityOption = None, diff --git a/nemo_retriever/src/nemo_retriever/cli/ingest/options.py b/nemo_retriever/src/nemo_retriever/cli/ingest/options.py index 259380554..4115d3e55 100644 --- a/nemo_retriever/src/nemo_retriever/cli/ingest/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/ingest/options.py @@ -288,11 +288,11 @@ help=f"Optional embedding model name override. Defaults to {DEFAULT_EMBED_MODEL} when omitted.", ), ] -EmbedModelNamePrefixOption = Annotated[ +EmbedModelProviderPrefixOption = Annotated[ str | None, typer.Option( - "--embed-model-name-prefix", - help="Optional LiteLLM route/org prefix prepended to the remote embedding model name.", + "--embed-model-provider-prefix", + help="Optional LiteLLM provider prefix prepended to the remote embedding model name.", ), ] LocalIngestEmbedBackendOption = Annotated[ diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 758a19a19..9ae745553 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -174,7 +174,7 @@ def _local_command( table_name: opts.TableNameOption = "nemo-retriever", embed_invoke_url: opts.EmbedInvokeUrlOption = None, embed_model_name: opts.EmbedModelNameOption = None, - embed_model_name_prefix: opts.EmbedModelNamePrefixOption = None, + embed_model_provider_prefix: opts.EmbedModelProviderPrefixOption = None, reranker_invoke_url: opts.RerankerInvokeUrlOption = None, reranker_api_key_env: opts.RerankerApiKeyEnvOption = None, reranker_model_name: opts.RerankerModelNameOption = None, @@ -221,7 +221,7 @@ def _local_command( embed=QueryEmbedOptions( embed_invoke_url=embed_invoke_url, embed_model_name=embed_model_name, - embed_model_name_prefix=embed_model_name_prefix, + embed_model_provider_prefix=embed_model_provider_prefix, ), rerank=QueryRerankOptions( enabled=rerank, @@ -263,7 +263,7 @@ def _request() -> QueryRequest: embed=QueryEmbedOptions( embed_invoke_url=embed_invoke_url, embed_model_name=embed_model_name, - embed_model_name_prefix=embed_model_name_prefix, + embed_model_provider_prefix=embed_model_provider_prefix, ), rerank=QueryRerankOptions( enabled=rerank, diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index f07dc33a8..089e7e26d 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -72,11 +72,11 @@ help=f"Optional embedding model name override. Defaults to {DEFAULT_EMBED_MODEL} when omitted.", ), ] -EmbedModelNamePrefixOption = Annotated[ +EmbedModelProviderPrefixOption = Annotated[ str | None, typer.Option( - "--embed-model-name-prefix", - help="Optional LiteLLM route/org prefix prepended to the remote embedding model name.", + "--embed-model-provider-prefix", + help="Optional LiteLLM provider prefix prepended to the remote embedding model name.", ), ] RerankerInvokeUrlOption = Annotated[ diff --git a/nemo_retriever/src/nemo_retriever/common/api/util/string_processing/__init__.py b/nemo_retriever/src/nemo_retriever/common/api/util/string_processing/__init__.py index eae7d8b23..7182944db 100644 --- a/nemo_retriever/src/nemo_retriever/common/api/util/string_processing/__init__.py +++ b/nemo_retriever/src/nemo_retriever/common/api/util/string_processing/__init__.py @@ -51,12 +51,12 @@ def ensure_openai_embeddings_http_url(endpoint_url: str) -> str: return urlunsplit((parts.scheme, parts.netloc, new_path, parts.query, parts.fragment)) -def prepend_model_name_prefix(model_name: str | None, model_name_prefix: str | None) -> str | None: - """Prepend a LiteLLM route/org prefix to a model identifier when configured.""" +def prepend_model_provider_prefix(model_name: str | None, model_provider_prefix: str | None) -> str | None: + """Prepend a LiteLLM provider prefix to a model identifier when configured.""" if model_name is None: return None model = str(model_name).strip() - prefix = str(model_name_prefix or "").strip().strip("/") + prefix = str(model_provider_prefix or "").strip().strip("/") if not model or not prefix: return model return f"{prefix}/{model.lstrip('/')}" diff --git a/nemo_retriever/src/nemo_retriever/common/params/models.py b/nemo_retriever/src/nemo_retriever/common/params/models.py index 84491e693..3112249bf 100644 --- a/nemo_retriever/src/nemo_retriever/common/params/models.py +++ b/nemo_retriever/src/nemo_retriever/common/params/models.py @@ -384,7 +384,7 @@ class EmbedParams(_ParamsModel): embedding_endpoint: Optional[str] = None embed_invoke_url: Optional[str] = None embed_model_name: Optional[str] = None - embed_model_name_prefix: 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 diff --git a/nemo_retriever/src/nemo_retriever/common/params/utils.py b/nemo_retriever/src/nemo_retriever/common/params/utils.py index 4290dfc22..bf1b9b2d3 100644 --- a/nemo_retriever/src/nemo_retriever/common/params/utils.py +++ b/nemo_retriever/src/nemo_retriever/common/params/utils.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any, Dict -from nemo_retriever.common.api.util.string_processing import prepend_model_name_prefix +from nemo_retriever.common.api.util.string_processing import prepend_model_provider_prefix if TYPE_CHECKING: from nemo_retriever.common.params.models import BatchTuningParams @@ -53,11 +53,11 @@ def normalize_embed_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: if "embed_invoke_url" in normalized: normalized.setdefault("embedding_endpoint", normalized["embed_invoke_url"]) endpoint = normalized.get("embedding_endpoint") or normalized.get("embed_invoke_url") - model_name_prefix = normalized.pop("embed_model_name_prefix", None) - if endpoint and model_name_prefix: + model_provider_prefix = normalized.pop("embed_model_provider_prefix", None) + if endpoint and model_provider_prefix: for key in ("model_name", "embed_model_name"): if key in normalized: - normalized[key] = prepend_model_name_prefix(normalized[key], str(model_name_prefix)) + normalized[key] = prepend_model_provider_prefix(normalized[key], str(model_provider_prefix)) return normalized @@ -66,7 +66,7 @@ def build_embed_option_kwargs( embed_model_name: str | None, local_ingest_embed_backend: str | None = None, embed_api_key: str | None = None, - embed_model_name_prefix: 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, @@ -88,8 +88,8 @@ def build_embed_option_kwargs( 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_name_prefix is not None: - embed_kwargs["embed_model_name_prefix"] = embed_model_name_prefix + 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: diff --git a/nemo_retriever/src/nemo_retriever/ingest/plan.py b/nemo_retriever/src/nemo_retriever/ingest/plan.py index c8ad53b49..f6ff01f6e 100644 --- a/nemo_retriever/src/nemo_retriever/ingest/plan.py +++ b/nemo_retriever/src/nemo_retriever/ingest/plan.py @@ -185,7 +185,7 @@ class IngestEmbedBatchOptions: class IngestEmbedOptions: embed_invoke_url: str | None = None embed_model_name: str | None = None - embed_model_name_prefix: str | None = None + embed_model_provider_prefix: str | None = None local_ingest_embed_backend: LocalIngestEmbedBackendValue | None = None embed_api_key: str | None = None embed_modality: str | None = None @@ -643,7 +643,7 @@ def resolve_ingest_plan(request: IngestPlanRequest) -> ResolvedIngestPlan: embed.embed_model_name, local_ingest_embed_backend=embed.local_ingest_embed_backend, embed_api_key=embed.embed_api_key, - embed_model_name_prefix=embed.embed_model_name_prefix, + embed_model_provider_prefix=embed.embed_model_provider_prefix, embed_modality=embed.embed_modality, text_elements_modality=embed.text_elements_modality, structured_elements_modality=embed.structured_elements_modality, diff --git a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py index 47cd5b818..99fa2cde0 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py @@ -50,7 +50,7 @@ def local_embedder(texts): import pandas as pd from nemo_retriever.common.api.util.string_processing import ( ensure_openai_embeddings_http_url, - prepend_model_name_prefix, + prepend_model_provider_prefix, ) from nemo_retriever.models import _DEFAULT_EMBED_MODEL @@ -75,7 +75,7 @@ class TextEmbeddingConfig: api_key: Optional[str] = None embedding_nim_endpoint: Optional[str] = None # e.g. "http://host:8000/v1" embedding_model: str = _DEFAULT_EMBED_MODEL - embedding_model_prefix: Optional[str] = None + embedding_model_provider_prefix: Optional[str] = None encoding_format: str = "float" # OpenAI-compatible embeddings often accept "float" input_type: str = "passage" truncate: str = "END" @@ -297,7 +297,7 @@ def _http_embed_openai_compat( encoding_format: str, input_type: str, truncate: str, - model_name_prefix: Optional[str] = None, + model_provider_prefix: Optional[str] = None, dimensions: Optional[int] = None, timeout_s: float = 600.0, ) -> List[Optional[List[float]]]: @@ -313,7 +313,7 @@ def _http_embed_openai_compat( raise RuntimeError("Remote embedding requested but `httpx` is not installed.") from e url = _normalize_embeddings_endpoint(_pick_embed_endpoint(endpoint_url)) - model_name = prepend_model_name_prefix(model_name, model_name_prefix) or model_name + model_name = prepend_model_provider_prefix(model_name, model_provider_prefix) or model_name headers: Dict[str, str] = {"accept": "application/json", "content-type": "application/json"} token = (api_key or "").strip() if token: @@ -359,7 +359,7 @@ def _make_async_request( api_key: Optional[str], embedding_nim_endpoint: str, embedding_model: str, - embedding_model_prefix: Optional[str], + embedding_model_provider_prefix: Optional[str], encoding_format: str, input_type: str, truncate: str, @@ -384,7 +384,7 @@ def _make_async_request( api_key=api_key, endpoint_url=str(embedding_nim_endpoint), model_name=str(embedding_model), - model_name_prefix=embedding_model_prefix, + model_provider_prefix=embedding_model_provider_prefix, encoding_format=str(encoding_format), input_type=str(input_type), truncate=str(truncate), @@ -407,7 +407,7 @@ def _async_request_handler( api_key: Optional[str], embedding_nim_endpoint: str, embedding_model: str, - embedding_model_prefix: Optional[str], + embedding_model_provider_prefix: Optional[str], encoding_format: str, input_type: str, truncate: str, @@ -429,7 +429,7 @@ def _async_request_handler( api_key=api_key or None, embedding_nim_endpoint=str(embedding_nim_endpoint), embedding_model=str(embedding_model), - embedding_model_prefix=embedding_model_prefix, + embedding_model_provider_prefix=embedding_model_provider_prefix, encoding_format=str(encoding_format), input_type=str(input_type), truncate=str(truncate), @@ -450,7 +450,7 @@ def _async_runner( api_key: Optional[str], embedding_nim_endpoint: str, embedding_model: str, - embedding_model_prefix: Optional[str], + embedding_model_provider_prefix: Optional[str], encoding_format: str, input_type: str, truncate: str, @@ -465,7 +465,7 @@ def _async_runner( api_key, embedding_nim_endpoint, embedding_model, - embedding_model_prefix, + embedding_model_provider_prefix, encoding_format, input_type, truncate, @@ -592,10 +592,10 @@ def create_text_embeddings_for_df( task_config["endpoint_url"] if "endpoint_url" in task_config else transform_config.embedding_nim_endpoint ) model_name = task_config["model_name"] if "model_name" in task_config else transform_config.embedding_model - model_name_prefix = ( - task_config["model_name_prefix"] - if "model_name_prefix" in task_config - else task_config.get("embed_model_name_prefix", transform_config.embedding_model_prefix) + model_provider_prefix = ( + task_config["model_provider_prefix"] + if "model_provider_prefix" in task_config + else task_config.get("embed_model_provider_prefix", transform_config.embedding_model_provider_prefix) ) dimensions = task_config["dimensions"] if "dimensions" in task_config else transform_config.dimensions @@ -688,7 +688,7 @@ def _text_image_content(r: pd.Series) -> Optional[str]: api_key, str(endpoint_url), str(model_name), - model_name_prefix, + model_provider_prefix, str(transform_config.encoding_format), str(transform_config.input_type), str(transform_config.truncate), @@ -711,7 +711,7 @@ def _text_image_content(r: pd.Series) -> Optional[str]: api_key, str(endpoint_url), str(model_name), - model_name_prefix, + model_provider_prefix, str(transform_config.encoding_format), str(transform_config.input_type), str(transform_config.truncate), diff --git a/nemo_retriever/src/nemo_retriever/models/inference/runtime.py b/nemo_retriever/src/nemo_retriever/models/inference/runtime.py index e73a476f6..2edd8b9cd 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/runtime.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/runtime.py @@ -30,7 +30,7 @@ def _embed_group( inference_batch_size: int, output_column: str, resolved_model_name: str, - embed_model_name_prefix: Optional[str] = None, + embed_model_provider_prefix: Optional[str] = None, nim_http_max_concurrent: int = 32, input_type: str = "passage", request_timeout_s: float = 600.0, @@ -77,7 +77,7 @@ def embedder(texts: Sequence[str]) -> Sequence[Sequence[float]]: # noqa dimensions=None, embedding_nim_endpoint=endpoint or "http://localhost:8012/v1", embedding_model=resolved_model_name or VL_EMBED_MODEL, - embedding_model_prefix=embed_model_name_prefix, + embedding_model_provider_prefix=embed_model_provider_prefix, embed_modality=group_modality, nim_http_max_concurrent=max(1, int(nim_http_max_concurrent)), ) @@ -89,7 +89,7 @@ def embedder(texts: Sequence[str]) -> Sequence[Sequence[float]]: # noqa "embedder": embedder, "multimodal_embedder": multimodal_embedder, "endpoint_url": endpoint, - "embed_model_name_prefix": embed_model_name_prefix, + "embed_model_provider_prefix": embed_model_provider_prefix, "local_batch_size": int(effective_batch_size), "nim_http_max_concurrent": max(1, int(nim_http_max_concurrent)), "request_timeout_s": float(request_timeout_s), @@ -113,7 +113,7 @@ def embed_text_main_text_embed( embedding_dim_column: str = "text_embeddings_1b_v2_dim", has_embedding_column: str = "text_embeddings_1b_v2_has_embedding", embed_modality: str = "text", - embed_model_name_prefix: Optional[str] = None, + embed_model_provider_prefix: Optional[str] = None, nim_http_max_concurrent: int = 32, input_type: str = "passage", request_timeout_s: float | None = None, @@ -152,7 +152,7 @@ def embed_text_main_text_embed( inference_batch_size=inference_batch_size, output_column=output_column, resolved_model_name=resolved_model_name, - embed_model_name_prefix=embed_model_name_prefix, + embed_model_provider_prefix=embed_model_provider_prefix, nim_http_max_concurrent=nim_http_max_concurrent, input_type=input_type, request_timeout_s=float(request_timeout_s), @@ -174,7 +174,7 @@ def embed_text_main_text_embed( inference_batch_size=inference_batch_size, output_column=output_column, resolved_model_name=resolved_model_name, - embed_model_name_prefix=embed_model_name_prefix, + embed_model_provider_prefix=embed_model_provider_prefix, nim_http_max_concurrent=nim_http_max_concurrent, input_type=input_type, request_timeout_s=float(request_timeout_s), diff --git a/nemo_retriever/src/nemo_retriever/models/nim/util/__init__.py b/nemo_retriever/src/nemo_retriever/models/nim/util/__init__.py index fd848a237..e07ddcad3 100644 --- a/nemo_retriever/src/nemo_retriever/models/nim/util/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/nim/util/__init__.py @@ -15,7 +15,7 @@ from nemo_retriever.models.nim.primitives.nim_model_interface import ModelInterface from nemo_retriever.common.api.util.string_processing import ( ensure_openai_embeddings_http_url, - prepend_model_name_prefix, + prepend_model_provider_prefix, ) logger = logging.getLogger(__name__) @@ -141,7 +141,7 @@ def infer_microservice( input_type: str = "passage", truncate: str = "END", batch_size: int = 8191, - model_name_prefix: str | None = None, + model_provider_prefix: str | None = None, grpc: bool = False, input_names: list = ["text"], output_names: list = ["embeddings"], @@ -204,7 +204,7 @@ def infer_microservice( ) else: embedding_endpoint = ensure_openai_embeddings_http_url(str(embedding_endpoint)) - model_name = prepend_model_name_prefix(model_name, model_name_prefix) + model_name = prepend_model_provider_prefix(model_name, model_provider_prefix) client = NimClient( model_interface=EmbeddingModelInterface(), protocol="http", diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index 2ca6e1b28..b0d3fe5ad 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic.py @@ -113,7 +113,7 @@ class AgenticRetrievalConfig: vdb_op: str = "lancedb" vdb_kwargs: dict[str, Any] = field(default_factory=dict) query_embedder: str = VL_EMBED_MODEL - query_embedder_prefix: Optional[str] = None + query_embedder_provider_prefix: Optional[str] = None embedding_endpoint: Optional[str] = None embedding_api_key: str = "" local_hf_batch_size: int = 32 @@ -175,7 +175,7 @@ def __init__( embed_kwargs={ "model_name": str(cfg.query_embedder or VL_EMBED_MODEL), "embed_model_name": str(cfg.query_embedder or VL_EMBED_MODEL), - "embed_model_name_prefix": cfg.query_embedder_prefix, + "embed_model_provider_prefix": cfg.query_embedder_provider_prefix, "embedding_endpoint": cfg.embedding_endpoint, "api_key": cfg.embedding_api_key, "input_type": "query", diff --git a/nemo_retriever/src/nemo_retriever/query/options.py b/nemo_retriever/src/nemo_retriever/query/options.py index 6922eae18..abb978fdd 100644 --- a/nemo_retriever/src/nemo_retriever/query/options.py +++ b/nemo_retriever/src/nemo_retriever/query/options.py @@ -26,7 +26,7 @@ class QueryRetrievalOptions: class QueryEmbedOptions: embed_invoke_url: str | None = None embed_model_name: str | None = None - embed_model_name_prefix: str | None = None + embed_model_provider_prefix: str | None = None @dataclass(frozen=True) diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index 4ed9c2f33..c8af17281 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/nemo_retriever/src/nemo_retriever/query/workflow.py @@ -105,7 +105,7 @@ def resolve_query_plan(request: QueryRequest) -> ResolvedQueryPlan: embed_kwargs = build_embed_option_kwargs( request.embed.embed_invoke_url, request.embed.embed_model_name, - embed_model_name_prefix=request.embed.embed_model_name_prefix, + embed_model_provider_prefix=request.embed.embed_model_provider_prefix, ) rerank_kwargs = _build_rerank_kwargs(request.rerank) if request.rerank.enabled else {} content_types = request.retrieval.content_types @@ -183,8 +183,8 @@ def agentic_query_documents(request: QueryRequest) -> list[dict[str, Any]]: } if request.embed.embed_model_name: cfg_kwargs["query_embedder"] = request.embed.embed_model_name - if request.embed.embed_model_name_prefix: - cfg_kwargs["query_embedder_prefix"] = request.embed.embed_model_name_prefix + if request.embed.embed_model_provider_prefix: + cfg_kwargs["query_embedder_provider_prefix"] = request.embed.embed_model_provider_prefix if request.rerank.enabled: # `reranker` doubles as the on/off gate (rerank=bool(cfg.reranker)) and the # model name, so fall back to the default model when only --rerank is given. diff --git a/nemo_retriever/src/nemo_retriever/service/config.py b/nemo_retriever/src/nemo_retriever/service/config.py index 7b62e8e46..e7e301079 100644 --- a/nemo_retriever/src/nemo_retriever/service/config.py +++ b/nemo_retriever/src/nemo_retriever/service/config.py @@ -141,10 +141,10 @@ class NimEndpointsConfig(RichModel): "Server-owned — clients cannot override the deployed embed NIM SKU." ), ) - embed_model_name_prefix: str | None = Field( + embed_model_provider_prefix: str | None = Field( default=None, description=( - "Optional LiteLLM route/org prefix prepended to embed_model_name for " + "Optional LiteLLM provider prefix prepended to embed_model_name for " "remote embedding endpoints that require namespaced model IDs." ), ) @@ -300,7 +300,7 @@ class VectorDbConfig(RichModel): lancedb_uri: str = "/data/vectordb" table_name: str = "nemo_retriever" embed_model: str = "nvidia/llama-nemotron-embed-vl-1b-v2" - embed_model_name_prefix: str | None = None + embed_model_provider_prefix: str | None = None vectordb_url: str = Field( default="http://nemo-retriever-vectordb:7671", description="URL of the vectordb service (for workers to POST embeddings to)", diff --git a/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml b/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml index 4277391fd..449f53f1f 100644 --- a/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml +++ b/nemo_retriever/src/nemo_retriever/service/retriever-service.yaml @@ -37,9 +37,9 @@ nim_endpoints: embed_invoke_url: null # Model name for the remote embed NIM (server-owned; must match the SKU). embed_model_name: null - # Optional LiteLLM route/org prefix prepended to embed_model_name for - # proxies that require namespaced model IDs (e.g. "team/model/provider"). - embed_model_name_prefix: null + # Optional LiteLLM provider prefix prepended to embed_model_name for + # proxies that require provider/model IDs. + embed_model_provider_prefix: null # gRPC endpoint for the Parakeet ASR NIM (e.g. parakeet-nim:50051). # When set, audio/video pipelines use remote ASR instead of loading # the local Parakeet model (which requires torch + GPU). diff --git a/nemo_retriever/src/nemo_retriever/service/service_ingestor.py b/nemo_retriever/src/nemo_retriever/service/service_ingestor.py index 1e858cc93..af4d5cbb4 100644 --- a/nemo_retriever/src/nemo_retriever/service/service_ingestor.py +++ b/nemo_retriever/src/nemo_retriever/service/service_ingestor.py @@ -219,7 +219,7 @@ def _normalize_files(files: Union[str, List[str], List[Path]]) -> list[Path]: "nemotron_parse_invoke_url", "embed_invoke_url", "embedding_endpoint", - "embed_model_name_prefix", + "embed_model_provider_prefix", "endpoint_url", "api_base", "auth_token", diff --git a/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py b/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py index 6841718f1..e9f8c53b8 100644 --- a/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py +++ b/nemo_retriever/src/nemo_retriever/service/services/pipeline_executor.py @@ -318,7 +318,7 @@ def _post_rows_to_vectordb(rows: list[dict[str, Any]], vectordb_url: str, filena "embedding_endpoint", "api_key", "embed_model_name", - "embed_model_name_prefix", + "embed_model_provider_prefix", "model_name", ) # Trust-owned caption keys. ``endpoint_url`` / ``api_key`` / @@ -852,8 +852,8 @@ def build_embed_params(nim: "NimEndpointsConfig", local: "LocalModelsConfig | No if nim.embed_model_name: kwargs["model_name"] = nim.embed_model_name kwargs["embed_model_name"] = nim.embed_model_name - if nim.embed_model_name_prefix: - kwargs["embed_model_name_prefix"] = nim.embed_model_name_prefix + if nim.embed_model_provider_prefix: + kwargs["embed_model_provider_prefix"] = nim.embed_model_provider_prefix if nim.api_key: kwargs["api_key"] = nim.api_key return EmbedParams(**kwargs) diff --git a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py index a083abcf5..dd1ece41c 100644 --- a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py +++ b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py @@ -90,14 +90,14 @@ def _embed_queries_remote( embed_model: str, embed_endpoint: str, embed_api_key: str, - embed_model_name_prefix: str | None = None, + embed_model_provider_prefix: str | None = None, ) -> list[list[float]]: from nemo_retriever.models.nim.util import infer_microservice return infer_microservice( texts, model_name=embed_model, - model_name_prefix=embed_model_name_prefix, + model_provider_prefix=embed_model_provider_prefix, embedding_endpoint=embed_endpoint, nvidia_api_key=embed_api_key or None, input_type="query", @@ -119,7 +119,7 @@ def __init__( embed_model: str, embed_api_key: str, *, - embed_model_name_prefix: str | None = None, + embed_model_provider_prefix: str | None = None, local_embed: bool = False, local_embed_backend: str = "hf", hf_cache_dir: str | None = None, @@ -130,7 +130,7 @@ def __init__( self.table_name = table_name self.embed_endpoint = embed_endpoint self.embed_model = embed_model - self.embed_model_name_prefix = embed_model_name_prefix + self.embed_model_provider_prefix = embed_model_provider_prefix self.embed_api_key = embed_api_key self.local_embed = local_embed self.local_embed_backend = local_embed_backend @@ -248,7 +248,7 @@ def embed_queries(self, texts: list[str]) -> list[list[float]]: return _embed_queries_remote( texts, embed_model=self.embed_model, - embed_model_name_prefix=self.embed_model_name_prefix, + embed_model_provider_prefix=self.embed_model_provider_prefix, embed_endpoint=self.embed_endpoint, embed_api_key=self.embed_api_key, ) @@ -273,7 +273,7 @@ def create_vectordb_app( table_name: str = "nemo_retriever", embed_endpoint: str = "", embed_model: str = "nvidia/llama-nemotron-embed-vl-1b-v2", - embed_model_name_prefix: str | None = None, + embed_model_provider_prefix: str | None = None, embed_api_key: str = "", *, local_embed: bool = False, @@ -292,7 +292,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: table_name=table_name, embed_endpoint=embed_endpoint, embed_model=embed_model, - embed_model_name_prefix=embed_model_name_prefix, + embed_model_provider_prefix=embed_model_provider_prefix, embed_api_key=embed_api_key, local_embed=local_embed, local_embed_backend=local_embed_backend, @@ -391,7 +391,7 @@ def main() -> None: parser.add_argument("--table-name", default="nemo_retriever", help="LanceDB table name") parser.add_argument("--embed-endpoint", default="", help="Remote NIM/OpenAI-compatible embed URL") parser.add_argument("--embed-model", default="nvidia/llama-nemotron-embed-vl-1b-v2") - parser.add_argument("--embed-model-name-prefix", default="", help="Optional LiteLLM route/org prefix") + parser.add_argument("--embed-model-provider-prefix", default="", help="Optional LiteLLM provider prefix") parser.add_argument("--embed-api-key", default="") parser.add_argument( "--local-embed", @@ -430,7 +430,7 @@ def main() -> None: table_name=args.table_name, embed_endpoint=args.embed_endpoint, embed_model=args.embed_model, - embed_model_name_prefix=args.embed_model_name_prefix or None, + embed_model_provider_prefix=args.embed_model_provider_prefix or None, embed_api_key=args.embed_api_key, local_embed=args.local_embed, local_embed_backend=args.local_embed_backend, diff --git a/nemo_retriever/tests/test_embed_params.py b/nemo_retriever/tests/test_embed_params.py index cb58c8a04..db2b334c9 100644 --- a/nemo_retriever/tests/test_embed_params.py +++ b/nemo_retriever/tests/test_embed_params.py @@ -55,16 +55,16 @@ def test_image_modalities_constant(): assert isinstance(IMAGE_MODALITIES, frozenset) -def test_build_embed_option_kwargs_applies_remote_model_name_prefix(): +def test_build_embed_option_kwargs_applies_remote_model_provider_prefix(): kwargs = build_embed_option_kwargs( "https://litellm.example.com/v1/embeddings", "nvidia/llama-nemotron-embed-vl-1b-v2", - embed_model_name_prefix="nvidia", + embed_model_provider_prefix="nvidia", ) assert kwargs["model_name"] == "nvidia/nvidia/llama-nemotron-embed-vl-1b-v2" assert kwargs["embed_model_name"] == "nvidia/nvidia/llama-nemotron-embed-vl-1b-v2" - assert "embed_model_name_prefix" not in kwargs + assert "embed_model_provider_prefix" not in kwargs def test_build_embed_option_kwargs_leaves_model_unchanged_without_prefix(): @@ -81,7 +81,7 @@ def test_build_embed_option_kwargs_prefix_supports_other_vendor_namespaces(): kwargs = build_embed_option_kwargs( "https://litellm.example.com/v1/embeddings", "mistral/embed-small", - embed_model_name_prefix="acme", + embed_model_provider_prefix="acme", ) assert kwargs["model_name"] == "acme/mistral/embed-small" @@ -92,7 +92,7 @@ def test_build_embed_option_kwargs_prefix_supports_bare_model_name(): kwargs = build_embed_option_kwargs( "https://litellm.example.com/v1/embeddings", "nv-embedqa-e5-v5", - embed_model_name_prefix="nvidia", + embed_model_provider_prefix="nvidia", ) assert kwargs["model_name"] == "nvidia/nv-embedqa-e5-v5" @@ -103,12 +103,12 @@ def test_build_embed_option_kwargs_prefix_is_remote_only(): kwargs = build_embed_option_kwargs( None, "nvidia/llama-nemotron-embed-vl-1b-v2", - embed_model_name_prefix="nvidia", + embed_model_provider_prefix="nvidia", ) assert kwargs["model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" assert kwargs["embed_model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" - assert "embed_model_name_prefix" not in kwargs + assert "embed_model_provider_prefix" not in kwargs # =================================================================== diff --git a/nemo_retriever/tests/test_root_cli_workflow.py b/nemo_retriever/tests/test_root_cli_workflow.py index 1eeb9af2d..659904565 100644 --- a/nemo_retriever/tests/test_root_cli_workflow.py +++ b/nemo_retriever/tests/test_root_cli_workflow.py @@ -434,7 +434,7 @@ def fake_create_ingestor(**_kwargs: Any) -> Any: "http://embed:8000/v1/embeddings", "--embed-model-name", "nvidia/llama-nemotron-embed-1b-v2", - "--embed-model-name-prefix", + "--embed-model-provider-prefix", "nvidia", ], ) diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index 70d1fee29..f66a86f6a 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -136,7 +136,7 @@ def query(self, query: str, **_kwargs: Any) -> list[dict[str, Any]]: "http://embed:8000/v1/embeddings", "--embed-model-name", "nvidia/llama-nemotron-embed-1b-v2", - "--embed-model-name-prefix", + "--embed-model-provider-prefix", "nvidia", ], ) diff --git a/nemo_retriever/tests/test_service_local_models.py b/nemo_retriever/tests/test_service_local_models.py index 43555a656..fae41ed35 100644 --- a/nemo_retriever/tests/test_service_local_models.py +++ b/nemo_retriever/tests/test_service_local_models.py @@ -27,7 +27,7 @@ def test_build_embed_params_from_nim_config() -> None: nim = NimEndpointsConfig( embed_invoke_url="http://embed-nim/v1/embeddings", embed_model_name="nvidia/llama-nemotron-embed-vl-1b-v2", - embed_model_name_prefix="nvidia", + embed_model_provider_prefix="nvidia", api_key="k", ) ep = build_embed_params(nim, LocalModelsConfig(enabled=True)) @@ -35,7 +35,7 @@ def test_build_embed_params_from_nim_config() -> None: assert ep.embed_invoke_url == "http://embed-nim/v1/embeddings" assert ep.model_name == "nvidia/llama-nemotron-embed-vl-1b-v2" assert ep.embed_model_name == "nvidia/llama-nemotron-embed-vl-1b-v2" - assert ep.embed_model_name_prefix == "nvidia" + assert ep.embed_model_provider_prefix == "nvidia" assert ep.api_key == "k" diff --git a/nemo_retriever/tests/test_service_vectordb_app.py b/nemo_retriever/tests/test_service_vectordb_app.py index a9fd8d89f..04b0e54c9 100644 --- a/nemo_retriever/tests/test_service_vectordb_app.py +++ b/nemo_retriever/tests/test_service_vectordb_app.py @@ -100,13 +100,13 @@ def fake_infer_microservice(data, **kwargs): embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", embed_endpoint="https://litellm.example.com/v1/embeddings", embed_api_key="k", - embed_model_name_prefix="nvidia", + embed_model_provider_prefix="nvidia", ) assert vectors == [[0.1, 0.2]] assert calls["data"] == ["hello"] assert calls["model_name"] == "nvidia/llama-nemotron-embed-vl-1b-v2" - assert calls["model_name_prefix"] == "nvidia" + assert calls["model_provider_prefix"] == "nvidia" assert calls["embedding_endpoint"] == "https://litellm.example.com/v1/embeddings"