diff --git a/nemo_retriever/helm/README.md b/nemo_retriever/helm/README.md index 8059ffc3da..5ba9b9291c 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.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 a5c58bd773..ba04101745 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_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: diff --git a/nemo_retriever/helm/templates/deployment-vectordb.yaml b/nemo_retriever/helm/templates/deployment-vectordb.yaml index ebee3b2729..19c4efc009 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.embedModelProviderPrefix }} + - --embed-model-provider-prefix + - {{ .Values.serviceConfig.vectordb.embedModelProviderPrefix | quote }} + {{- end }} - --port - {{ $vdb.port | quote }} {{- if $embedURL }} diff --git a/nemo_retriever/helm/values.yaml b/nemo_retriever/helm/values.yaml index 7811bd20fe..685d7bfd2d 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" + 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 b0e36bedcb..2c5eac1b69 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_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 4ccd0c5cfc..4115d3e555 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.", ), ] +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( diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 404c8c14d8..9ae7455537 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_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, diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index 0c60b78d2e..089e7e26dc 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.", ), ] +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."), 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 6fb7fb5e45..7182944dba 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_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. diff --git a/nemo_retriever/src/nemo_retriever/common/params/models.py b/nemo_retriever/src/nemo_retriever/common/params/models.py index 71c046cee8..3112249bfb 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_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 f7c82f7c88..bf1b9b2d36 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_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: diff --git a/nemo_retriever/src/nemo_retriever/ingest/plan.py b/nemo_retriever/src/nemo_retriever/ingest/plan.py index e857f4385f..f6ff01f6e7 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_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, 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 dd7c29c373..99fa2cde0a 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_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), diff --git a/nemo_retriever/src/nemo_retriever/models/inference/runtime.py b/nemo_retriever/src/nemo_retriever/models/inference/runtime.py index 1206cbcee2..2edd8b9cd4 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_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), diff --git a/nemo_retriever/src/nemo_retriever/models/inference/shared.py b/nemo_retriever/src/nemo_retriever/models/inference/shared.py index 21d042dbfd..22066a1c01 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 20394de67f..e07ddcad36 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_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", diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index aa75034470..b0d3fe5adc 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_provider_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_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 237001a6e9..abb978fddd 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_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 5c777d5378..c8af172819 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_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 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_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 08f306608b..e7e3010792 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_provider_prefix: str | None = Field( + default=None, + description=( + "Optional LiteLLM provider 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_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 3a797bf7db..449f53f1fa 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 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 e53112c641..af4d5cbb47 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_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 2e01fae0f4..e9f8c53b81 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_provider_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_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 8221d2c298..dd1ece41c5 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_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_provider_prefix=embed_model_provider_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_provider_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_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 @@ -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_provider_prefix=self.embed_model_provider_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_provider_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_provider_prefix=embed_model_provider_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-provider-prefix", default="", help="Optional LiteLLM provider 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_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 596c1b06fc..db2b334c99 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_provider_prefix(): + kwargs = build_embed_option_kwargs( + "https://litellm.example.com/v1/embeddings", + "nvidia/llama-nemotron-embed-vl-1b-v2", + 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_provider_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_provider_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_provider_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_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_provider_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 f0decca4aa..659904565b 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-provider-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 28bfeec9bf..f66a86f6aa 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-provider-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 e7abd6e3ba..fae41ed351 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_provider_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_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 1ea0a12d55..04b0e54c92 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_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_provider_prefix"] == "nvidia" + assert calls["embedding_endpoint"] == "https://litellm.example.com/v1/embeddings" + + _CANNED_HITS = [ { "text": "Revenue grew 12% year over year.",