From a065b4b9c658206f175e00b28462f946088ebfdc Mon Sep 17 00:00:00 2001 From: Taylor <28880387+tsmithv11@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:10:50 -0700 Subject: [PATCH] feat(python): add NVIDIA Nemotron embedding provider Adds a Python `nvidia` embedding provider for Nemotron 3.5 Embed 1B and 8B. It works with NVIDIA's hosted API and self-hosted NIM `/v1` endpoints, and uses the existing `openai` client. The OpenAI provider can't send NVIDIA's passage/query input types or pick Nemotron's dimensions. - Documents are embedded as `passage`, searches as `query`. - Default dimensions are 2048 (1B) and 4096 (8B), and `dim` is checked against each model's supported range. Self-hosted model IDs need an explicit `dim`. - Response count, order, dimensions and values are checked. Null and empty text get a null vector without a request. - Each request is retried on its own with LanceDB's retry policy. Invalid input or responses fail immediately. - `NVIDIA_API_KEY` is only sent to the hosted API. Custom endpoints need an explicit `api_key` (a `$var:` reference), and endpoints restored from table metadata must be allowed with `allow_endpoint`. - Request dimensions are sent through `extra_body`, preserving the NVIDIA wire contract while supporting the declared minimum `openai==1.6.1`. Python and text only. Docs are in `integrations/embedding/nvidia`. Originally authored by Taylor (@tsmithv11). ### Testing - `test_nvidia_embeddings.py`: 45 passed, 2 deselected with both the locked `openai==2.44.0` and minimum `openai==1.6.1`. These use the real OpenAI client with mock HTTP transport and real local tables. - Embedding, registry, field-metadata and trust suites: 163 passed, 4 skipped, 10 deselected. - Historical live experiments: hosted 8B passed a five-request smoke test; hosted 1B returned HTTP 500. No live request was made in this stage, so current hosted 1B availability was not verified. Enterprise coverage uses the repository mock HTTP server, not a live deployment. --- docs/web/docs.json | 1 + docs/web/integrations/embedding/nvidia.mdx | 107 ++++ python/python/lancedb/embeddings/__init__.py | 2 + python/python/lancedb/embeddings/nvidia.py | 194 +++++++ python/python/tests/test_nvidia_embeddings.py | 522 ++++++++++++++++++ 5 files changed, 826 insertions(+) create mode 100644 docs/web/integrations/embedding/nvidia.mdx create mode 100644 python/python/lancedb/embeddings/nvidia.py create mode 100644 python/python/tests/test_nvidia_embeddings.py diff --git a/docs/web/docs.json b/docs/web/docs.json index 66ce3acd3..917eb6406 100644 --- a/docs/web/docs.json +++ b/docs/web/docs.json @@ -446,6 +446,7 @@ "integrations/embedding/imagebind", "integrations/embedding/instructor", "integrations/embedding/jina", + "integrations/embedding/nvidia", "integrations/embedding/ollama", "integrations/embedding/openai", "integrations/embedding/openclip", diff --git a/docs/web/integrations/embedding/nvidia.mdx b/docs/web/integrations/embedding/nvidia.mdx new file mode 100644 index 000000000..c0664684a --- /dev/null +++ b/docs/web/integrations/embedding/nvidia.mdx @@ -0,0 +1,107 @@ +--- +title: NVIDIA Nemotron +sidebarTitle: NVIDIA +description: Text embeddings with Nemotron 3.5 Embed 1B and 8B in the Python SDK. +--- + +The Python SDK registers `nvidia` for client-side text embedding with NVIDIA-hosted +and self-hosted NIM endpoints. Documents use `input_type="passage"`; searches use +`input_type="query"`. This provider supports text only, with floating-point output. +Null and empty text values get a null vector without a request; pass +`on_bad_vectors="null"` or `"drop"` to `add` to keep or skip those rows. +Image and combined text/image inputs are not supported by this integration. + +Install the optional client and set your NVIDIA API key: + +```shell +pip install openai +export NVIDIA_API_KEY="your-api-key" +``` + +```python +import lancedb +from lancedb.embeddings import get_registry +from lancedb.pydantic import LanceModel, Vector + +func = get_registry().get("nvidia").create( + name="nvidia/nemotron-3.5-embed-1b", +) + +class Document(LanceModel): + text: str = func.SourceField() + vector: Vector(func.ndims()) = func.VectorField() + +db = lancedb.connect("./nemotron-db") +table = db.create_table("documents", schema=Document) +table.add([{"text": "Paris is the capital of France."}]) +results = table.search("What is the capital of France?").limit(3).to_list() +``` + +For a LanceDB Enterprise table, connect as usual, for example +`lancedb.connect("db://my-database", api_key=..., region=...)`. The same schema, +`add`, and `search` calls compute embeddings in the Python client. The NVIDIA key +is separate from the LanceDB database key. Use a fresh table name when creating a +table; reopening an existing table restores its embedding configuration. + +## Models and configuration + +The hosted API documents these exact IDs and dimensions: + +| Model ID | Default dimensions | Supported `dim` | +|---|---:|---:| +| `nvidia/nemotron-3.5-embed-1b` | 2048 | 128–2048 | +| `nvidia/nemotron-3.5-embed-8b` | 4096 | 1–4096 | + +| Parameter | Default | Description | +|---|---|---| +| `name` | `nvidia/nemotron-3.5-embed-1b` | Model ID sent to the endpoint. | +| `dim` | Model default above | Output dimensions; required for other served model IDs. | +| `base_url` | `https://integrate.api.nvidia.com/v1` | API base including `/v1`. | +| `api_key` | `NVIDIA_API_KEY` for the hosted API | Alternatively, a registry reference such as `$var:nvidia-key`. | +| `batch_size` | `32` | Maximum inputs per request. Reduce for your deployment's limits. | +| `truncate` | `END` | Truncate overlong text at the end, or use `NONE` to reject it. | +| `max_retries` | `7` | LanceDB's retry policy, applied to each request during automatic document/query embedding. | + +Requests specify `modality="text"` and `encoding_format="float"`. Returned vectors +are checked for count, input indices, dimensions, and finite numeric values. +The hosted input limit is 10,240 tokens per item; self-hosted limits depend on the +served profile. This provider batches by item count, not token count. + +## Self-hosted NIM + +Set `name` to the exact ID returned by your deployment's `GET /v1/models` and +`dim` to an output dimension that deployment supports: + +```python +func = get_registry().get("nvidia").create( + name="nvidia/Nemotron-3.5-Embed-1B-BF16", + base_url="http://localhost:8000/v1", + dim=2048, +) +``` + +Custom endpoints may run without a key. `NVIDIA_API_KEY` is only sent to the +hosted API, so for authenticated endpoints use +`get_registry().set_var("nvidia-key", key)` with `api_key="$var:nvidia-key"`. Registry references, rather than secret values, are +stored in the schema. Restore registry variables in each process. Before using +a custom endpoint restored from table metadata in a new process, call +`get_registry().allow_endpoint("http://localhost:8000/v1")`. + +The NIM 2.4 support matrix lists `nvidia/Nemotron-3.5-Embed-8B-BF16` at 2048 +dimensions, whereas the hosted 8B reference specifies 4096. For self-hosted model +IDs, this provider requires an explicit `dim` and checks the response instead of +assuming the hosted defaults. + +SDK coverage: Python. TypeScript, Rust, and Java do not gain a NVIDIA provider +from this integration. Automated tests use mocked NVIDIA and Enterprise HTTP +responses plus real local tables. The opt-in hosted smoke test runs with +`NVIDIA_API_KEY` and `pytest python/tests/test_nvidia_embeddings.py -m slow -k nvidia_live` +from the repository's Python environment. It sends five small requests per model +(ten total), with retries disabled, to check local ingestion, reopening, search, +and reduced-dimensional embeddings. + +NVIDIA references (checked October 23, 2026): +[hosted 1B API](https://docs.api.nvidia.com/nim/reference/nvidia-nemotron-3_5-embed-1b-infer), +[hosted 8B API](https://docs.api.nvidia.com/nim/reference/nvidia-nemotron-3_5-embed-8b-infer), +[NIM request/response format](https://docs.nvidia.com/nim/nemo-retriever/embedding/2.4/use-the-api-openai.html), +[NIM support matrix](https://docs.nvidia.com/nim/nemo-retriever/embedding/2.4/support-matrix.html). diff --git a/python/python/lancedb/embeddings/__init__.py b/python/python/lancedb/embeddings/__init__.py index 636502c9e..20e8c331b 100644 --- a/python/python/lancedb/embeddings/__init__.py +++ b/python/python/lancedb/embeddings/__init__.py @@ -16,6 +16,7 @@ from .instructor import InstructorEmbeddingFunction from .ollama import OllamaEmbeddings from .open_clip import OpenClipEmbeddings from .openai import OpenAIEmbeddings +from .nvidia import NvidiaEmbeddings from .registry import EmbeddingFunctionRegistry, get_registry, register from .sentence_transformers import SentenceTransformerEmbeddings from .gte import GteEmbeddings @@ -40,6 +41,7 @@ __all__ = [ "register", "SentenceTransformerEmbeddings", "OpenAIEmbeddings", + "NvidiaEmbeddings", "OpenClipEmbeddings", "BedRockText", "CohereEmbeddingFunction", diff --git a/python/python/lancedb/embeddings/nvidia.py b/python/python/lancedb/embeddings/nvidia.py new file mode 100644 index 000000000..11fae0d34 --- /dev/null +++ b/python/python/lancedb/embeddings/nvidia.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import math +import os +from functools import cached_property +from typing import List, Literal, Optional + +from pydantic import Field, model_validator + +from ..util import attempt_import_or_raise +from .base import TextEmbeddingFunction +from .registry import register +from .utils import TEXT, MissingApiKeyError, retry_with_exponential_backoff + + +NVIDIA_BASE_URL = "https://integrate.api.nvidia.com/v1" +# Hosted model ID -> (default dimensions, minimum dimensions) +MODEL_DIMENSIONS = { + "nvidia/nemotron-3.5-embed-1b": (2048, 128), + "nvidia/nemotron-3.5-embed-8b": (4096, 1), +} + + +@register("nvidia") +class NvidiaEmbeddings(TextEmbeddingFunction): + """Text embeddings using NVIDIA Nemotron 3.5 Embed 1B or 8B. + + Uses the optional ``openai`` package for NVIDIA-hosted and self-hosted NIM + ``/v1/embeddings`` endpoints. Documents use ``input_type="passage"`` and + queries use ``input_type="query"``. Null and empty strings get a null + embedding. Image inputs are not supported. + + Parameters + ---------- + name : str + Hosted model ID, or the served model ID from a NIM ``/v1/models`` response. + dim : int, optional + Requested output dimensions. Defaults to 2048 for the hosted 1B model + and 4096 for the hosted 8B model. Required for other served model IDs; + use a dimension supported by your deployment. + base_url : str + API base URL including ``/v1``. Defaults to NVIDIA's hosted API. + api_key : str, optional + Registry variable reference, such as ``$var:nvidia-key``. Otherwise the + hosted API uses ``NVIDIA_API_KEY``. Custom endpoints only receive an + explicit ``api_key`` and can run without authentication. + batch_size : int + Maximum inputs per request (default 32). Reduce for deployment limits. + truncate : str + ``END`` (default) truncates overlong text; ``NONE`` rejects it. + max_retries : int + Inherited LanceDB retry limit, applied to each request. The OpenAI + client's own retries are disabled. + + Examples + -------- + Set ``NVIDIA_API_KEY`` before using the hosted API: + + >>> from lancedb.embeddings import get_registry + >>> from lancedb.pydantic import LanceModel, Vector + >>> func = get_registry().get("nvidia").create() + >>> class Document(LanceModel): + ... text: str = func.SourceField() + ... vector: Vector(func.ndims()) = func.VectorField() + + Embedding runs client-side for both local and remote tables. For a custom + endpoint restored from table metadata in a new process, first call + [allow_endpoint][lancedb.embeddings.EmbeddingFunctionRegistry.allow_endpoint]. + """ + + name: str = "nvidia/nemotron-3.5-embed-1b" + dim: Optional[int] = Field(default=None, gt=0) + base_url: str = NVIDIA_BASE_URL + api_key: Optional[str] = Field(default=None, repr=False) + batch_size: int = Field(default=32, gt=0) + truncate: Literal["END", "NONE"] = "END" + + @model_validator(mode="after") + def _validate_dimensions(self): + dimensions = MODEL_DIMENSIONS.get(self.name) + if dimensions is None and self.dim is None: + raise ValueError("Set dim explicitly for self-hosted NIM model IDs") + if dimensions is not None and self.dim is not None: + native_dim, minimum = dimensions + if not minimum <= self.dim <= native_dim: + raise ValueError( + f"{self.name} supports dimensions from {minimum} to {native_dim}" + ) + return self + + @staticmethod + def sensitive_keys(): + return ["api_key"] + + def endpoints(self) -> List[str]: + return [] if self.base_url.rstrip("/") == NVIDIA_BASE_URL else [self.base_url] + + def ndims(self) -> int: + return self.dim if self.dim is not None else MODEL_DIMENSIONS[self.name][0] + + def compute_query_embeddings( + self, query: str, *args, **kwargs + ) -> List[Optional[List[float]]]: + return self._embed(query, input_type="query") + + def generate_embeddings(self, texts: TEXT) -> List[Optional[List[float]]]: + return self._embed(texts, input_type="passage") + + # Retry each request rather than the whole call, so a transient failure does + # not resend batches that already succeeded, and invalid inputs or responses + # fail at once instead of being retried. + def compute_query_embeddings_with_retry( + self, query: str, *args, **kwargs + ) -> List[Optional[List[float]]]: + return self._embed(query, input_type="query", max_retries=self.max_retries) + + def compute_source_embeddings_with_retry( + self, texts: TEXT, *args, **kwargs + ) -> List[Optional[List[float]]]: + return self._embed(texts, input_type="passage", max_retries=self.max_retries) + + def _embed( + self, + texts: TEXT, + *, + input_type: Literal["passage", "query"], + max_retries: int = 0, + ) -> List[Optional[List[float]]]: + texts = list(self.sanitize_input(texts)) + if any(text is not None and not isinstance(text, str) for text in texts): + raise ValueError("NVIDIA embeddings support text strings only") + request = self._request + if max_retries > 0: + request = retry_with_exponential_backoff(request, max_retries=max_retries) + embeddings = [None] * len(texts) + valid = [index for index, text in enumerate(texts) if text] + for start in range(0, len(valid), self.batch_size): + rows = valid[start : start + self.batch_size] + data = request([texts[row] for row in rows], input_type) + if not isinstance(data, list) or len(data) != len(rows): + raise ValueError("NVIDIA returned an unexpected number of embeddings") + indices = [getattr(item, "index", None) for item in data] + if any(type(index) is not int for index in indices) or sorted( + indices + ) != list(range(len(rows))): + raise ValueError("NVIDIA returned invalid embedding indices") + for row, item in zip(rows, sorted(data, key=lambda item: item.index)): + vector = getattr(item, "embedding", None) + if ( + not isinstance(vector, list) + or len(vector) != self.ndims() + or any( + type(value) not in (int, float) or not math.isfinite(value) + for value in vector + ) + ): + raise ValueError( + "NVIDIA returned an invalid embedding: expected " + f"{self.ndims()} finite numeric values" + ) + embeddings[row] = vector + return embeddings + + def _request(self, batch: List[str], input_type: str): + response = self._client.embeddings.create( + model=self.name, + input=batch, + encoding_format="float", + # The typed dimensions parameter is unavailable in openai 1.6.1. + extra_body={ + "dimensions": self.ndims(), + "input_type": input_type, + "modality": "text", + "truncate": self.truncate, + }, + ) + return getattr(response, "data", None) + + @cached_property + def _client(self): + openai = attempt_import_or_raise("openai") + hosted = not self.endpoints() + # Never send the hosted API key from the environment to a custom endpoint + api_key = self.api_key or (os.environ.get("NVIDIA_API_KEY") if hosted else None) + if not api_key and hosted: + raise MissingApiKeyError( + "Set NVIDIA_API_KEY or provide api_key via a registry variable" + ) + return openai.OpenAI( + base_url=self.base_url, + api_key=api_key or "not-required", + max_retries=0, + ) diff --git a/python/python/tests/test_nvidia_embeddings.py b/python/python/tests/test_nvidia_embeddings.py new file mode 100644 index 000000000..3ae07c330 --- /dev/null +++ b/python/python/tests/test_nvidia_embeddings.py @@ -0,0 +1,522 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import json +import os +from typing import Optional + +import lancedb +import numpy as np +import pyarrow as pa +import pytest + +from lancedb.embeddings import ( + EmbeddingFunctionConfig, + NvidiaEmbeddings, + UntrustedEmbeddingConfigError, + get_registry, +) +from lancedb.embeddings.nvidia import MODEL_DIMENSIONS, NVIDIA_BASE_URL +from lancedb.pydantic import LanceModel, Vector + + +@pytest.fixture +def nvidia_api(monkeypatch): + """Exercise the real OpenAI client against a deterministic NVIDIA endpoint.""" + httpx = pytest.importorskip("httpx") + openai = pytest.importorskip("openai") + requests, responses, clients = [], [], [] + original_client = openai.OpenAI + monkeypatch.setenv("NVIDIA_API_KEY", "test-nvidia-key") + registry = get_registry() + monkeypatch.setattr(registry, "_variables", {}) + monkeypatch.setattr(registry, "_allowed_endpoints", set()) + + def respond(request): + requests.append(request) + if responses: + return responses.pop(0) + body = json.loads(request.content) + data = [] + for index, text in enumerate(body["input"]): + vector = [0.0] * body["dimensions"] + vector[0 if text == "alpha" else 1] = 1.0 + data.append({"index": index, "embedding": vector, "object": "embedding"}) + return httpx.Response( + 200, + json={ + "object": "list", + "model": body["model"], + # The API's index, not response order, determines the input row. + "data": data[::-1], + "usage": {"prompt_tokens": len(data), "total_tokens": len(data)}, + }, + ) + + def client(**kwargs): + assert kwargs["max_retries"] == 0 + instance = original_client( + **kwargs, + http_client=httpx.Client( + transport=httpx.MockTransport(respond), trust_env=False + ), + ) + clients.append(instance) + return instance + + monkeypatch.setattr(openai, "OpenAI", client) + yield requests, responses + for client in clients: + client.close() + + +@pytest.mark.parametrize( + "name,dimensions", [(name, dims[0]) for name, dims in MODEL_DIMENSIONS.items()] +) +def test_hosted_models(nvidia_api, name, dimensions): + requests, _ = nvidia_api + func = get_registry().get("nvidia").create(name=name) + assert isinstance(func, NvidiaEmbeddings) + assert func.ndims() == dimensions + assert len(func.compute_source_embeddings(["alpha"])[0]) == dimensions + assert len(func.compute_query_embeddings("alpha")[0]) == dimensions + for request, input_type in zip(requests, ["passage", "query"]): + assert str(request.url) == f"{NVIDIA_BASE_URL}/embeddings" + assert request.headers["authorization"] == "Bearer test-nvidia-key" + assert json.loads(request.content) == { + "model": name, + "input": ["alpha"], + "dimensions": dimensions, + "encoding_format": "float", + "input_type": input_type, + "modality": "text", + "truncate": "END", + } + + +@pytest.mark.parametrize( + "texts", + [ + ["alpha", "beta", "gamma"], + np.array(["alpha", "beta", "gamma"]), + pa.array(["alpha", "beta", "gamma"]), + pa.chunked_array([["alpha"], ["beta", "gamma"]]), + ], +) +def test_batching_and_response_order(nvidia_api, texts): + requests, _ = nvidia_api + func = NvidiaEmbeddings.create(dim=128, batch_size=2, truncate="NONE") + result = func.compute_source_embeddings(texts) + assert [vector[:2] for vector in result] == [[1, 0], [0, 1], [0, 1]] + assert [json.loads(request.content)["input"] for request in requests] == [ + ["alpha", "beta"], + ["gamma"], + ] + assert all(json.loads(r.content)["truncate"] == "NONE" for r in requests) + assert func.compute_source_embeddings([]) == [] + assert len(requests) == 2 + + +@pytest.mark.parametrize( + "texts", + [ + ["alpha", None, "", "beta"], + pa.array(["alpha", None, "", "beta"]), + ], +) +def test_null_and_empty_inputs(nvidia_api, texts): + requests, _ = nvidia_api + func = NvidiaEmbeddings.create(dim=128) + result = func.compute_source_embeddings(texts) + assert [vector and vector[:2] for vector in result] == [[1, 0], None, None, [0, 1]] + assert [json.loads(r.content)["input"] for r in requests] == [["alpha", "beta"]] + assert func.compute_source_embeddings([None, ""]) == [None, None] + assert len(requests) == 1 + + +def test_table_with_null_text(nvidia_api, tmp_path): + func = NvidiaEmbeddings.create(dim=128) + + class Document(LanceModel): + text: Optional[str] = func.SourceField() + vector: Optional[Vector(func.ndims())] = func.VectorField() + + db = lancedb.connect(tmp_path) + table = db.create_table("docs", schema=Document) + table.add([{"text": "alpha"}, {"text": None}], on_bad_vectors="null") + rows = table.to_arrow().to_pylist() + assert rows[0]["vector"][:2] == [1, 0] + assert rows[1]["vector"] is None + + +@pytest.mark.parametrize("texts", [[b"image"], [["nested"]], [123]]) +def test_rejects_non_text(nvidia_api, monkeypatch, texts): + requests, _ = nvidia_api + monkeypatch.setattr("lancedb.embeddings.utils.time.sleep", pytest.fail) + func = NvidiaEmbeddings.create() + with pytest.raises(ValueError, match="text strings only"): + func.compute_source_embeddings(texts) + with pytest.raises(ValueError, match="text strings only"): + func.compute_source_embeddings_with_retry(texts) + assert requests == [] + + +@pytest.mark.parametrize( + "options", + [ + {"batch_size": 0}, + {"dim": 0}, + {"dim": 127}, + {"dim": 2049}, + {"name": "nvidia/nemotron-3.5-embed-8b", "dim": 4097}, + {"name": "nvidia/Nemotron-3.5-Embed-8B-BF16"}, + {"truncate": "START"}, + ], +) +def test_invalid_configuration(options): + with pytest.raises(ValueError): + NvidiaEmbeddings.create(**options) + + +@pytest.mark.parametrize( + "data,error", + [ + ([], "number of embeddings"), + (None, "number of embeddings"), + ([{"index": 1, "embedding": [0.0] * 128}], "indices"), + ([{"embedding": [0.0] * 128}], "indices"), + ([{"index": 0, "embedding": [0.0]}], "finite numeric"), + ([{"index": 0, "embedding": "base64"}], "finite numeric"), + ([{"index": 0, "embedding": ["bad"] * 128}], "finite numeric"), + ([{"index": 0}], "finite numeric"), + ], +) +def test_invalid_response(nvidia_api, data, error): + import httpx + + _, responses = nvidia_api + responses.append(httpx.Response(200, json={"data": data})) + with pytest.raises(ValueError, match=error): + NvidiaEmbeddings.create(dim=128).compute_source_embeddings(["alpha"]) + + +def test_invalid_response_is_not_retried(nvidia_api, monkeypatch): + import httpx + + requests, responses = nvidia_api + monkeypatch.setattr("lancedb.embeddings.utils.time.sleep", pytest.fail) + responses.append(httpx.Response(200, json={"data": []})) + func = NvidiaEmbeddings.create(dim=128) + with pytest.raises(ValueError, match="number of embeddings"): + func.compute_source_embeddings_with_retry(["alpha"]) + assert len(requests) == 1 + + +def test_missing_response_data(nvidia_api): + import httpx + + _, responses = nvidia_api + responses.append(httpx.Response(200, json={})) + with pytest.raises(ValueError, match="number of embeddings"): + NvidiaEmbeddings.create().compute_query_embeddings("alpha") + + +@pytest.mark.parametrize("bad_value", ["NaN", "Infinity", "-Infinity"]) +def test_nonfinite_response(nvidia_api, bad_value): + import httpx + + _, responses = nvidia_api + # Raw JSON is needed because httpx refuses to encode non-finite numbers. + vector = ",".join([bad_value] + ["0"] * 127) + responses.append( + httpx.Response(200, text='{"data":[{"index":0,"embedding":[' + vector + "]}]}") + ) + with pytest.raises(ValueError, match="finite numeric"): + NvidiaEmbeddings.create(dim=128).compute_source_embeddings(["alpha"]) + + +def test_duplicate_indices(nvidia_api): + import httpx + + _, responses = nvidia_api + responses.append( + httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.0] * 128}] * 2}) + ) + with pytest.raises(ValueError, match="indices"): + NvidiaEmbeddings.create(dim=128).compute_source_embeddings(["alpha", "beta"]) + + +@pytest.mark.parametrize( + "status,expected_calls", [(429, 2), (503, 2), (401, 1), (400, 1)] +) +def test_retry_policy(nvidia_api, monkeypatch, status, expected_calls): + import httpx + import openai + + requests, responses = nvidia_api + monkeypatch.setattr("lancedb.embeddings.utils.time.sleep", lambda _: None) + responses.append(httpx.Response(status, json={"error": {"message": "test error"}})) + func = NvidiaEmbeddings.create(dim=128, max_retries=1) + if expected_calls == 2: + assert len(func.compute_query_embeddings_with_retry("alpha")[0]) == 128 + else: + with pytest.raises(openai.APIStatusError) as error: + func.compute_query_embeddings_with_retry("alpha") + assert error.value.status_code == status + assert len(requests) == expected_calls + + +def test_retry_resends_only_failed_batch(nvidia_api, monkeypatch): + import httpx + + requests, responses = nvidia_api + monkeypatch.setattr("lancedb.embeddings.utils.time.sleep", lambda _: None) + vector = [1.0] + [0.0] * 127 + responses.append( + httpx.Response(200, json={"data": [{"index": 0, "embedding": vector}]}) + ) + responses.append(httpx.Response(429, json={"error": {"message": "slow down"}})) + func = NvidiaEmbeddings.create(dim=128, batch_size=1, max_retries=1) + result = func.compute_source_embeddings_with_retry(["alpha", "beta"]) + assert [v[:2] for v in result] == [[1, 0], [0, 1]] + assert [json.loads(r.content)["input"] for r in requests] == [ + ["alpha"], + ["beta"], + ["beta"], + ] + + +def test_credentials_and_metadata(nvidia_api, monkeypatch): + requests, _ = nvidia_api + registry = get_registry() + registry.set_var("nvidia-key", "secret-from-registry") + options = { + "name": "nvidia/Nemotron-3.5-Embed-8B-BF16", + "base_url": "http://localhost:8000/v1", + "dim": 2048, + "api_key": "$var:nvidia-key", + "batch_size": 2, + "truncate": "NONE", + } + func = NvidiaEmbeddings.create(**options) + config = EmbeddingFunctionConfig( + source_column="text", vector_column="vector", function=func + ) + schema = pa.schema([], metadata=registry.get_table_metadata([config])) + # Round-trip the configuration through Arrow, as table schemas do. + restored_schema = pa.ipc.read_schema(pa.BufferReader(schema.serialize())) + metadata = restored_schema.metadata + assert b"secret-from-registry" not in metadata[b"embedding_functions"] + restored = registry.parse_functions(metadata)["vector"].function + assert restored.safe_model_dump() == options + restored.compute_source_embeddings(["alpha"]) + restored.compute_query_embeddings("alpha") + assert str(requests[0].url) == "http://localhost:8000/v1/embeddings" + assert requests[0].headers["authorization"] == "Bearer secret-from-registry" + assert json.loads(requests[0].content)["model"] == options["name"] + assert [json.loads(r.content)["input_type"] for r in requests] == [ + "passage", + "query", + ] + + monkeypatch.setattr(registry, "_allowed_endpoints", set()) + blocked = registry.parse_functions(metadata)["vector"].function + with pytest.raises(UntrustedEmbeddingConfigError): + blocked.compute_query_embeddings("alpha") + registry.allow_endpoint(options["base_url"]) + assert isinstance( + registry.parse_functions(metadata)["vector"].function, NvidiaEmbeddings + ) + with pytest.raises(ValueError, match="Sensitive key"): + NvidiaEmbeddings.create(api_key="hardcoded-secret") + + +def test_missing_key_and_unauthenticated_nim(nvidia_api, monkeypatch): + from lancedb.embeddings.utils import MissingApiKeyError + + requests, _ = nvidia_api + # The hosted key in NVIDIA_API_KEY is not sent to custom endpoints. + func = NvidiaEmbeddings.create( + name="nvidia/Nemotron-3.5-Embed-1B-BF16", + base_url="http://localhost:8000/v1", + dim=128, + ) + func.compute_query_embeddings("alpha") + assert requests[0].headers["authorization"] == "Bearer not-required" + monkeypatch.delenv("NVIDIA_API_KEY") + with pytest.raises(MissingApiKeyError, match="NVIDIA_API_KEY"): + NvidiaEmbeddings.create().compute_query_embeddings_with_retry("alpha") + assert len(requests) == 1 + + +@pytest.mark.parametrize("name", MODEL_DIMENSIONS) +def test_local_table_document_query_and_reopen(nvidia_api, tmp_path, name): + requests, _ = nvidia_api + func = NvidiaEmbeddings.create(name=name, dim=128) + + class Document(LanceModel): + text: str = func.SourceField() + vector: Vector(func.ndims()) = func.VectorField() + + db = lancedb.connect(tmp_path) + table = db.create_table("docs", schema=Document, data=[{"text": "alpha"}]) + table.add([{"text": "beta"}]) + reopened = lancedb.connect(tmp_path).open_table("docs") + assert reopened.schema.metadata == table.schema.metadata + assert isinstance(reopened.embedding_functions["vector"].function, NvidiaEmbeddings) + assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha" + assert [json.loads(request.content)["input_type"] for request in requests] == [ + "passage", + "passage", + "query", + ] + + +@pytest.mark.parametrize("name", MODEL_DIMENSIONS) +def test_remote_table_document_query_and_reopen(nvidia_api, name): + from test_remote_db import mock_lancedb_connection, read_json_body, send_json + + requests, _ = nvidia_api + func = NvidiaEmbeddings.create(name=name, dim=128) + config = EmbeddingFunctionConfig( + source_column="text", vector_column="vector", function=func + ) + writes, queries = [], [] + + def handler(request): + if "/create/" in request.path or "/insert/" in request.path: + if request.headers.get("Transfer-Encoding") == "chunked": + body = bytearray() + while True: + size = int(request.rfile.readline(), 16) + if size == 0: + request.rfile.readline() + break + body.extend(request.rfile.read(size)) + request.rfile.read(2) + else: + body = request.rfile.read(int(request.headers["Content-Length"])) + writes.append(pa.ipc.open_stream(body).read_all()) + send_json(request, {}) + elif "/describe/" in request.path: + send_json( + request, + { + "version": 1, + "schema": { + "fields": [ + { + "name": "text", + "type": {"type": "string"}, + "nullable": True, + }, + { + "name": "vector", + "nullable": True, + "type": { + "type": "fixed_size_list", + "length": 128, + "fields": [ + { + "name": "item", + "type": {"type": "float"}, + "nullable": True, + } + ], + }, + }, + ], + "metadata": { + k.decode(): v.decode() + for k, v in writes[0].schema.metadata.items() + }, + }, + }, + ) + elif "/query/" in request.path: + queries.append(read_json_body(request)) + result = pa.table({"text": ["alpha"], "_distance": [0.0]}) + request.send_response(200) + request.send_header("Content-Type", "application/vnd.apache.arrow.file") + request.end_headers() + with pa.ipc.new_file(request.wfile, result.schema) as writer: + writer.write_table(result) + else: + request.send_response(404) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + table = db.create_table( + "docs", data=[{"text": "alpha"}], embedding_functions=[config] + ) + table.add([{"text": "beta"}]) + reopened = db.open_table("docs") + assert reopened.schema.metadata == writes[0].schema.metadata + assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha" + assert len(writes) == 2 + assert writes[0]["vector"].to_pylist()[0] == [1.0] + [0.0] * 127 + assert writes[1]["vector"].to_pylist()[0] == [0.0, 1.0] + [0.0] * 126 + assert queries[0]["vector"] == [1.0] + [0.0] * 127 + assert [json.loads(r.content)["input_type"] for r in requests] == [ + "passage", + "passage", + "query", + ] + + +@pytest.mark.asyncio +async def test_async_local_table(nvidia_api, tmp_path): + requests, _ = nvidia_api + func = NvidiaEmbeddings.create(dim=128) + + class Document(LanceModel): + text: str = func.SourceField() + vector: Vector(func.ndims()) = func.VectorField() + + with await lancedb.connect_async(tmp_path) as db: + table = await db.create_table("docs", schema=Document) + await table.add([{"text": "alpha"}, {"text": "beta"}]) + reopened = await db.open_table("docs") + results = await (await reopened.search("alpha")).limit(1).to_list() + assert results[0]["text"] == "alpha" + assert [json.loads(r.content)["input_type"] for r in requests] == [ + "passage", + "query", + ] + + +@pytest.mark.slow +@pytest.mark.skipif( + not os.environ.get("NVIDIA_API_KEY"), reason="NVIDIA_API_KEY not set" +) +@pytest.mark.parametrize("name", MODEL_DIMENSIONS) +def test_nvidia_live(name, tmp_path): + """Five small requests per model, with retries disabled for free-tier tests.""" + pytest.importorskip("openai") + func = NvidiaEmbeddings.create(name=name, max_retries=0) + passages = ["Paris is the capital of France.", "Bananas are yellow fruit."] + question = "What is the capital of France?" + + class Document(LanceModel): + text: str = func.SourceField() + vector: Vector(func.ndims()) = func.VectorField() + + db = lancedb.connect(tmp_path) + table = db.create_table( + "docs", schema=Document, data=[{"text": text} for text in passages] + ) + table.add([{"text": "Saturn is a planet with rings."}]) + reopened = lancedb.connect(tmp_path).open_table("docs") + assert reopened.schema.metadata == table.schema.metadata + assert reopened.schema.field("vector").type.list_size == func.ndims() + result = reopened.search(question).limit(1).to_list()[0] + assert result["text"] == passages[0] + assert len(result["vector"]) == func.ndims() + assert np.isfinite(result["vector"]).all() + + reduced = NvidiaEmbeddings.create(name=name, dim=128, max_retries=0) + vectors = reduced.compute_source_embeddings_with_retry(passages) + query = reduced.compute_query_embeddings_with_retry(question)[0] + assert np.shape(vectors) == (2, 128) + assert np.shape(query) == (128,) + assert np.dot(query, vectors[0]) > np.dot(query, vectors[1])