Add embed model provider prefix for LiteLLM proxy compatability

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.embedModelProviderPrefix` | `""` | Optional LiteLLM provider prefix prepended to the remote embed model name. |
 
 #### VectorDB and the embed endpoint { #vectordb-and-the-embed-endpoint }
 

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_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 }}
@@ -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_provider_prefix: {{ if .Values.serviceConfig.vectordb.embedModelProviderPrefix }}{{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }}{{ else }}null{{ end }}
   vectordb_url: "http://{{ .vectordbSvc }}:{{ .vectordbPort }}"
 {{- else }}
 vectordb:

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.embedModelProviderPrefix }}
+            - --embed-model-provider-prefix
+            - {{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }}
+            {{- end }}
             - --port
             - {{ $vdb.port | quote }}
             {{- if $embedURL }}

nemo_retriever/helm/values.yaml

@@ -660,6 +660,7 @@ serviceConfig:
     lancedbUri: "/data/vectordb"
     tableName: "nemo_retriever"
     embedModel: "nvidia/llama-nemotron-embed-vl-1b-v2"
+    embedModelProviderPrefix: ""
 
   # Optional bearer-token authentication. When apiToken is set, every
   # request must carry "Authorization: Bearer <token>".

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_provider_prefix: opts.EmbedModelProviderPrefixOption = None,
     local_ingest_embed_backend: opts.LocalIngestEmbedBackendOption = None,
     embed_modality: opts.EmbedModalityOption = None,
     embed_granularity: opts.EmbedGranularityOption = None,

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.",
     ),
 ]
+EmbedModelProviderPrefixOption = Annotated[
+    str | None,
+    typer.Option(
+        "--embed-model-provider-prefix",
+        help="Optional LiteLLM provider prefix prepended to the remote embedding model name.",
+    ),
+]
 LocalIngestEmbedBackendOption = Annotated[
     LocalIngestEmbedBackendValue | None,
     typer.Option(

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_provider_prefix: opts.EmbedModelProviderPrefixOption = 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_provider_prefix=embed_model_provider_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_provider_prefix=embed_model_provider_prefix,
                 ),
                 rerank=QueryRerankOptions(
                     enabled=rerank,

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.",
     ),
 ]
+EmbedModelProviderPrefixOption = Annotated[
+    str | None,
+    typer.Option(
+        "--embed-model-provider-prefix",
+        help="Optional LiteLLM provider prefix prepended to the remote embedding model name.",
+    ),
+]
 RerankerInvokeUrlOption = Annotated[
     str | None,
     typer.Option("--reranker-invoke-url", help="Reranker endpoint URL."),

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_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_provider_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.

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_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

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_provider_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_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_provider_prefix(normalized[key], str(model_provider_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_provider_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_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:

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_provider_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_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,

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_provider_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_provider_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_provider_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_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:
@@ -353,6 +359,7 @@ def _make_async_request(
     api_key: Optional[str],
     embedding_nim_endpoint: str,
     embedding_model: str,
+    embedding_model_provider_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_provider_prefix=embedding_model_provider_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_provider_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_provider_prefix=embedding_model_provider_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_provider_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_provider_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_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
 
     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_provider_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_provider_prefix,
                     str(transform_config.encoding_format),
                     str(transform_config.input_type),
                     str(transform_config.truncate),

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_provider_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_provider_prefix=embed_model_provider_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_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),
@@ -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_provider_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_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),
@@ -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_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),

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)

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_provider_prefix,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -138,6 +141,7 @@ def infer_microservice(
     input_type: str = "passage",
     truncate: str = "END",
     batch_size: int = 8191,
+    model_provider_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_provider_prefix(model_name, model_provider_prefix)
         client = NimClient(
             model_interface=EmbeddingModelInterface(),
             protocol="http",