mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-30 15:26:52 +00:00
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.
This commit is contained in:
Taylor
@@ -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",
|
||||
|
||||
@@ -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).
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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])
|
||||
Reference in New Issue
Block a user