mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-20 02:32:27 +00:00
fix(python)!: don't trust embedding settings from table metadata (#4550)
BREAKING CHANGE: `SentenceTransformerEmbeddings.trust_remote_code` now
defaults to `False` (was `True`). When an embedding function is restored
from table metadata, `trust_remote_code=True` and custom endpoints
(OpenAI `base_url`, Ollama `host`/proxy, watsonx `url`) are only honored
if the process approved them; otherwise computing embeddings raises
`UntrustedEmbeddingConfigError`. Functions created in the same process
are approved automatically. To reopen such a table in a new process,
call `get_registry().allow_remote_code(model, revision=...)` or
`get_registry().allow_endpoint(url)` first.
## Summary
Embedding functions are rebuilt from table metadata, which anyone who
can write the table controls. Since sentence-transformers defaulted
`trust_remote_code` to `True`, opening and querying a crafted table ran
arbitrary code from a model repository (GHSA-x8v2-vj73-gfwm). Metadata
could likewise point `base_url` at another server, sending it queries
and the `OPENAI_API_KEY` from the environment.
The table can no longer grant itself trust:
```python
registry = get_registry()
registry.allow_remote_code("nomic-ai/nomic-embed-text-v1.5", revision="<commit>")
registry.allow_endpoint("http://my-ollama-host:11434")
```
Trust is keyed by model name (optionally pinned to a revision), so
approving one model doesn't let a table swap in another. Restored
functions that aren't approved are replaced by a placeholder that raises
on use without constructing the real function (`huggingface` loads its
model in `__init__`), so non-embedding operations on the table keep
working. `sentence-transformers` and `huggingface` gain a `revision`
parameter. Custom embedding functions can opt in by overriding the new
`remote_code_model()` / `endpoints()` hooks.
Not covered: TypeScript (its built-ins have neither remote code nor
configurable endpoints), Bedrock `assumed_role`, and third-party
functions that don't implement the hooks.
## Evidence
`test_embeddings_trust.py` covers restore with and without approval,
revision pinning, swapped model names, endpoint gating, and that the
table stays usable. Opening a table with crafted metadata and searching
it now gives:
```
UntrustedEmbeddingConfigError: Embedding function 'sentence-transformers' was restored from table metadata with settings this process has not approved:
- runs code from model 'attacker/model' at revision None. If you trust it, call get_registry().allow_remote_code('attacker/model', revision=None)
```
## Merge Risk
**Door:** One-way: changes a default and makes previously working
reopen-and-search flows require explicit approval.
**Blast Radius:** public-api (Python embeddings); users of remote-code
models or custom endpoints who reopen tables in a new process.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5.5
parent
7d52c7deb7
commit
9020640538
@@ -134,6 +134,29 @@ For non-sensitive settings such as inference device selection, you can also use
|
||||
|
||||
Find the full list of arguments for each provider in the [integrations](/integrations/embedding) section.
|
||||
|
||||
### Trusting remote code and custom endpoints {#trusting-remote-code-and-custom-endpoints}
|
||||
|
||||
In Python, the embedding definition is stored in the table, so anyone who can write to the table can change it.
|
||||
To keep a table from running code or sending data where you didn't intend, two kinds of settings are only honored
|
||||
when this process has approved them:
|
||||
|
||||
- `trust_remote_code=True`, which runs code from the model's Hugging Face repository.
|
||||
- Custom endpoints such as OpenAI `base_url`, Ollama `host`, or watsonx `url`, which receive your data and API keys.
|
||||
|
||||
Embedding functions you create in the current process are approved automatically. When you open a table in a new
|
||||
process, approve these settings explicitly before computing embeddings:
|
||||
|
||||
```python
|
||||
from lancedb.embeddings import get_registry
|
||||
|
||||
registry = get_registry()
|
||||
# Pinning a commit hash is recommended, since the code at a branch can change.
|
||||
registry.allow_remote_code("nomic-ai/nomic-embed-text-v1.5", revision="<commit-hash>")
|
||||
registry.allow_endpoint("http://my-ollama-host:11434")
|
||||
```
|
||||
|
||||
Until then, computing embeddings raises `UntrustedEmbeddingConfigError`. Other operations on the table still work.
|
||||
|
||||
## Multiple embedding columns {#multiple-embedding-columns}
|
||||
|
||||
A single table can include more than one embedding definition when you want to store multiple semantic views
|
||||
|
||||
@@ -5,7 +5,7 @@ sidebarTitle: Hugging Face
|
||||
|
||||
import { PyEmbeddingHuggingfaceUsage } from '/snippets/integrations.mdx';
|
||||
|
||||
We offer support for all Hugging Face models (which can be loaded via [transformers](https://huggingface.co/docs/transformers/en/index) library). The default model is `colbert-ir/colbertv2.0` which also has its own special callout - `registry.get("colbert")`. Some Hugging Face models might require custom models defined on the HuggingFace Hub in their own modeling files. You may enable this by setting `trust_remote_code=True`. This option should only be set to True for repositories you trust and in which you have read the code, as it will execute code present on the Hub on your local machine.
|
||||
We offer support for all Hugging Face models (which can be loaded via [transformers](https://huggingface.co/docs/transformers/en/index) library). The default model is `colbert-ir/colbertv2.0` which also has its own special callout - `registry.get("colbert")`. Some Hugging Face models might require custom models defined on the HuggingFace Hub in their own modeling files. You may enable this by setting `trust_remote_code=True`. This option should only be set to True for repositories you trust and in which you have read the code, as it will execute code present on the Hub on your local machine. Pin `revision` to a commit hash, and note that tables using such a model need it approved with `get_registry().allow_remote_code(...)` when opened in a new process. See [Trusting remote code and custom endpoints](/embedding#trusting-remote-code-and-custom-endpoints).
|
||||
|
||||
Example usage:
|
||||
|
||||
|
||||
@@ -12,7 +12,7 @@ Sentence transformer embeddings are normalized by default. It is recommended to
|
||||
</Info>
|
||||
|
||||
<Warning>
|
||||
The `trust_remote_code` parameter defaults to `True`, which allows models to execute arbitrary code from their Hugging Face repository during loading. If you are loading untrusted models, set `trust_remote_code=False` to prevent remote code execution.
|
||||
Setting `trust_remote_code=True` allows the model to execute arbitrary code from its Hugging Face repository during loading. Only enable it for repositories you trust, and pin `revision` to a commit hash. When a table using such a model is opened in a new process, the model must be approved with `get_registry().allow_remote_code(...)`. See [Trusting remote code and custom endpoints](/embedding#trusting-remote-code-and-custom-endpoints).
|
||||
</Warning>
|
||||
|
||||
| Parameter | Type | Default Value | Description |
|
||||
@@ -20,7 +20,8 @@ The `trust_remote_code` parameter defaults to `True`, which allows models to exe
|
||||
| `name` | `str` | `all-MiniLM-L6-v2` | The name of the model |
|
||||
| `device` | `str` | `cpu` | The device to run the model on (can be `cpu` or `gpu`) |
|
||||
| `normalize` | `bool` | `True` | Whether to normalize the input text before feeding it to the model |
|
||||
| `trust_remote_code` | `bool` | `True` | Whether to trust and execute remote code from the model's Huggingface repository |
|
||||
| `trust_remote_code` | `bool` | `False` | Whether to trust and execute remote code from the model's Huggingface repository |
|
||||
| `revision` | `str` | `None` | The model revision to load, such as a commit hash |
|
||||
|
||||
|
||||
<Details title="Check out available sentence-transformer models here!">
|
||||
|
||||
@@ -3,7 +3,12 @@
|
||||
|
||||
|
||||
# ruff: noqa: F401
|
||||
from .base import EmbeddingFunction, EmbeddingFunctionConfig, TextEmbeddingFunction
|
||||
from .base import (
|
||||
EmbeddingFunction,
|
||||
EmbeddingFunctionConfig,
|
||||
TextEmbeddingFunction,
|
||||
UntrustedEmbeddingConfigError,
|
||||
)
|
||||
from .bedrock import BedRockText
|
||||
from .cohere import CohereEmbeddingFunction
|
||||
from .gemini_text import GeminiText
|
||||
@@ -29,6 +34,7 @@ __all__ = [
|
||||
"EmbeddingFunction",
|
||||
"EmbeddingFunctionConfig",
|
||||
"TextEmbeddingFunction",
|
||||
"UntrustedEmbeddingConfigError",
|
||||
"EmbeddingFunctionRegistry",
|
||||
"get_registry",
|
||||
"register",
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
import copy
|
||||
from typing import List, Union
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
from lancedb.util import add_note
|
||||
import numpy as np
|
||||
@@ -13,6 +13,18 @@ from pydantic import BaseModel, Field, PrivateAttr
|
||||
from .utils import TEXT, retry_with_exponential_backoff
|
||||
|
||||
|
||||
class UntrustedEmbeddingConfigError(PermissionError):
|
||||
"""
|
||||
Raised when using an embedding function restored from table metadata whose
|
||||
settings this process has not approved.
|
||||
|
||||
See
|
||||
[allow_remote_code][lancedb.embeddings.EmbeddingFunctionRegistry.allow_remote_code]
|
||||
and
|
||||
[allow_endpoint][lancedb.embeddings.EmbeddingFunctionRegistry.allow_endpoint].
|
||||
"""
|
||||
|
||||
|
||||
class EmbeddingFunction(BaseModel, ABC):
|
||||
"""
|
||||
An ABC for embedding functions.
|
||||
@@ -37,11 +49,29 @@ class EmbeddingFunction(BaseModel, ABC):
|
||||
"""
|
||||
Create an instance of the embedding function
|
||||
"""
|
||||
from .registry import EmbeddingFunctionRegistry
|
||||
|
||||
resolved_kwargs = cls.__resolveVariables(kwargs)
|
||||
instance = cls(**resolved_kwargs)
|
||||
instance._original_args = kwargs
|
||||
instance._original_args = dict(kwargs)
|
||||
if (
|
||||
instance.remote_code_model() is not None
|
||||
and "trust_remote_code" in cls.model_fields
|
||||
):
|
||||
# Record the flag even when left at its default, so a missing flag
|
||||
# in table metadata only ever means a table written before this.
|
||||
instance._original_args.setdefault(
|
||||
"trust_remote_code", instance.trust_remote_code
|
||||
)
|
||||
# The caller wrote this config, so this process trusts it. Without this,
|
||||
# a table created with it could not be reopened in the same process.
|
||||
EmbeddingFunctionRegistry.get_instance()._trust(instance)
|
||||
return instance
|
||||
|
||||
@classmethod
|
||||
def _resolve_variables(cls, args: dict) -> dict:
|
||||
return cls.__resolveVariables(args)
|
||||
|
||||
@classmethod
|
||||
def __resolveVariables(cls, args: dict) -> dict:
|
||||
"""
|
||||
@@ -92,6 +122,38 @@ class EmbeddingFunction(BaseModel, ABC):
|
||||
"""
|
||||
return []
|
||||
|
||||
def remote_code_model(self) -> Optional[Tuple[str, Optional[str]]]:
|
||||
"""
|
||||
Return the model whose repository code this function runs when its
|
||||
`trust_remote_code` field is enabled, as a `(name, revision)` pair.
|
||||
|
||||
Return None if the function never runs code from a model repository.
|
||||
When the function is restored from table metadata, `trust_remote_code`
|
||||
is only enabled for models allowed with
|
||||
[allow_remote_code][lancedb.embeddings.EmbeddingFunctionRegistry.allow_remote_code].
|
||||
"""
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _legacy_trust_remote_code_default(cls) -> bool:
|
||||
"""
|
||||
The `trust_remote_code` to use for table metadata that omits it. Only
|
||||
tables written before the flag was always recorded omit it, so this is
|
||||
the default from that time.
|
||||
"""
|
||||
field = cls.model_fields.get("trust_remote_code")
|
||||
return bool(field.default) if field is not None else False
|
||||
|
||||
def endpoints(self) -> List[str]:
|
||||
"""
|
||||
Return the non-default URLs this function sends data or credentials to.
|
||||
|
||||
When the function is restored from table metadata, each of these must be
|
||||
allowed with
|
||||
[allow_endpoint][lancedb.embeddings.EmbeddingFunctionRegistry.allow_endpoint].
|
||||
"""
|
||||
return []
|
||||
|
||||
@abstractmethod
|
||||
def compute_query_embeddings(self, *args, **kwargs) -> list[Union[np.array, None]]:
|
||||
"""
|
||||
@@ -233,3 +295,35 @@ class TextEmbeddingFunction(EmbeddingFunction):
|
||||
) -> list[Union[np.array, None]]:
|
||||
"""Generate the embeddings for the given texts"""
|
||||
pass
|
||||
|
||||
|
||||
class UntrustedEmbeddingFunction(EmbeddingFunction):
|
||||
"""
|
||||
Stands in for an embedding function restored from table metadata whose
|
||||
settings this process has not approved.
|
||||
|
||||
The table stays usable for anything that doesn't compute embeddings, while
|
||||
computing embeddings raises
|
||||
[UntrustedEmbeddingConfigError][lancedb.embeddings.UntrustedEmbeddingConfigError].
|
||||
The real function is never constructed, since some load their model in
|
||||
`__init__`.
|
||||
"""
|
||||
|
||||
function_name: str
|
||||
reasons: Tuple[str, ...]
|
||||
|
||||
def _error(self) -> UntrustedEmbeddingConfigError:
|
||||
details = "\n".join(f" - {reason}" for reason in self.reasons)
|
||||
return UntrustedEmbeddingConfigError(
|
||||
f"Embedding function '{self.function_name}' was restored from table "
|
||||
"metadata with settings this process has not approved:\n" + details
|
||||
)
|
||||
|
||||
def compute_query_embeddings(self, *args, **kwargs):
|
||||
raise self._error()
|
||||
|
||||
def compute_source_embeddings(self, *args, **kwargs):
|
||||
raise self._error()
|
||||
|
||||
def ndims(self) -> int:
|
||||
raise self._error()
|
||||
|
||||
@@ -14,6 +14,9 @@ if TYPE_CHECKING:
|
||||
import ollama
|
||||
|
||||
|
||||
DEFAULT_OLLAMA_HOST = "http://localhost:11434"
|
||||
|
||||
|
||||
@register("ollama")
|
||||
class OllamaEmbeddings(TextEmbeddingFunction):
|
||||
"""
|
||||
@@ -24,11 +27,25 @@ class OllamaEmbeddings(TextEmbeddingFunction):
|
||||
"""
|
||||
|
||||
name: str = "nomic-embed-text"
|
||||
host: str = "http://localhost:11434"
|
||||
host: str = DEFAULT_OLLAMA_HOST
|
||||
options: Optional[dict] = None # type = ollama.Options
|
||||
keep_alive: Optional[Union[float, str]] = None
|
||||
ollama_client_kwargs: Optional[dict] = {}
|
||||
|
||||
def endpoints(self) -> List[str]:
|
||||
urls = []
|
||||
if self.host.rstrip("/") != DEFAULT_OLLAMA_HOST:
|
||||
urls.append(self.host)
|
||||
# These kwargs go to the underlying httpx client, where a proxy would
|
||||
# see every request.
|
||||
for key in ("proxy", "proxies"):
|
||||
proxy = (self.ollama_client_kwargs or {}).get(key)
|
||||
if isinstance(proxy, str):
|
||||
urls.append(proxy)
|
||||
elif isinstance(proxy, dict):
|
||||
urls.extend(v for v in proxy.values() if isinstance(v, str))
|
||||
return urls
|
||||
|
||||
def ndims(self) -> int:
|
||||
return len(self.generate_embeddings(["foo"])[0])
|
||||
|
||||
|
||||
@@ -58,6 +58,9 @@ class OpenAIEmbeddings(TextEmbeddingFunction):
|
||||
def sensitive_keys():
|
||||
return ["api_key"]
|
||||
|
||||
def endpoints(self) -> List[str]:
|
||||
return [self.base_url] if self.base_url else []
|
||||
|
||||
@staticmethod
|
||||
def model_names():
|
||||
return [
|
||||
|
||||
@@ -2,9 +2,15 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import json
|
||||
from typing import Dict, Optional, Type
|
||||
from typing import Dict, Optional, Set, Tuple, Type
|
||||
|
||||
from .base import EmbeddingFunction, EmbeddingFunctionConfig
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from .base import (
|
||||
EmbeddingFunction,
|
||||
EmbeddingFunctionConfig,
|
||||
UntrustedEmbeddingFunction,
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingFunctionRegistry:
|
||||
@@ -42,6 +48,9 @@ class EmbeddingFunctionRegistry:
|
||||
def __init__(self):
|
||||
self._functions = {}
|
||||
self._variables = {}
|
||||
# Revision is a str, None for the default branch, or _ANY_REVISION.
|
||||
self._allowed_remote_code: Set[Tuple[str, object]] = set()
|
||||
self._allowed_endpoints: Set[str] = set()
|
||||
|
||||
def register(self, alias: Optional[str] = None):
|
||||
"""
|
||||
@@ -122,20 +131,129 @@ class EmbeddingFunctionRegistry:
|
||||
obj["vector_column"]: EmbeddingFunctionConfig(
|
||||
vector_column=obj["vector_column"],
|
||||
source_column=obj["source_column"],
|
||||
function=self.get(obj["name"]).create(**obj["model"]),
|
||||
function=self._restore(obj["name"], obj["model"]),
|
||||
)
|
||||
for obj in raw_list
|
||||
}
|
||||
|
||||
def _restore(self, name: str, args: dict) -> EmbeddingFunction:
|
||||
"""
|
||||
Create an embedding function from table metadata. Anyone who can write
|
||||
the table controls `args`, so settings that run model code or send data
|
||||
elsewhere must already be trusted by this process.
|
||||
"""
|
||||
cls = self.get(name)
|
||||
resolved = cls._resolve_variables(args)
|
||||
# Inspect the config without running __init__, which may load a model.
|
||||
probe = cls.model_construct(**resolved)
|
||||
|
||||
reasons = []
|
||||
model = probe.remote_code_model()
|
||||
if model is not None:
|
||||
allowed = self._remote_code_allowed(*model)
|
||||
explicit = "trust_remote_code" in resolved
|
||||
if explicit:
|
||||
requested = TypeAdapter(bool).validate_python(
|
||||
resolved["trust_remote_code"]
|
||||
)
|
||||
else:
|
||||
requested = cls._legacy_trust_remote_code_default()
|
||||
# Always set the flag, since a subclass may default it to True.
|
||||
resolved["trust_remote_code"] = requested and allowed
|
||||
# A legacy table that omits the flag is downgraded rather than
|
||||
# blocked: most sentence-transformers models don't need remote code,
|
||||
# and those tables relied on the old True default without asking.
|
||||
if requested and not allowed and explicit:
|
||||
model_name, revision = model
|
||||
reasons.append(
|
||||
f"runs code from model '{model_name}' at revision "
|
||||
f"{revision!r}. If you trust it, call "
|
||||
f"get_registry().allow_remote_code({model_name!r}, "
|
||||
f"revision={revision!r})"
|
||||
)
|
||||
reasons.extend(
|
||||
f"sends data to endpoint '{url}'. If you trust it, call "
|
||||
f"get_registry().allow_endpoint({url!r})"
|
||||
for url in probe.endpoints()
|
||||
if _normalize_url(url) not in self._allowed_endpoints
|
||||
)
|
||||
|
||||
if reasons:
|
||||
func = UntrustedEmbeddingFunction(
|
||||
function_name=name, reasons=tuple(reasons)
|
||||
)
|
||||
else:
|
||||
func = cls(**resolved)
|
||||
func._original_args = args
|
||||
return func
|
||||
|
||||
def allow_remote_code(self, model: str, revision: Optional[str] = None) -> None:
|
||||
"""
|
||||
Trust the code in a model repository when an embedding function that
|
||||
loads it is restored from table metadata.
|
||||
|
||||
Table metadata cannot enable `trust_remote_code` by itself, since anyone
|
||||
who can write to the table could then run code in this process. Embedding
|
||||
functions created in this process with `trust_remote_code=True` are
|
||||
allowed automatically.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : str
|
||||
The model name, such as `"nomic-ai/nomic-embed-text-v1.5"`.
|
||||
revision : Optional[str]
|
||||
The revision the table's embedding function must pin. Pinning a
|
||||
commit hash is recommended, since the code at a branch or tag can
|
||||
change after you review it. If None, any revision is trusted,
|
||||
including pull request refs that anyone can push to the repository.
|
||||
"""
|
||||
self._allowed_remote_code.add(
|
||||
(model, _ANY_REVISION if revision is None else revision)
|
||||
)
|
||||
|
||||
def allow_endpoint(self, url: str) -> None:
|
||||
"""
|
||||
Trust a custom endpoint, such as an OpenAI `base_url` or Ollama `host`,
|
||||
when an embedding function that uses it is restored from table metadata.
|
||||
|
||||
Table metadata cannot pick an endpoint by itself, since it could send
|
||||
your data and API keys to a server you don't control. Embedding functions
|
||||
created in this process are allowed to use their endpoints automatically.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
url : str
|
||||
The endpoint URL. Trailing slashes are ignored.
|
||||
"""
|
||||
self._allowed_endpoints.add(_normalize_url(url))
|
||||
|
||||
def _remote_code_allowed(self, model: str, revision: Optional[str]) -> bool:
|
||||
return (model, _ANY_REVISION) in self._allowed_remote_code or (
|
||||
model,
|
||||
revision,
|
||||
) in self._allowed_remote_code
|
||||
|
||||
def _trust(self, func: EmbeddingFunction) -> None:
|
||||
model = func.remote_code_model()
|
||||
if model is not None and getattr(func, "trust_remote_code", False):
|
||||
# Only the revision this code selected, where None means the default
|
||||
# branch. Other revisions, such as PR refs, may hold other code.
|
||||
self._allowed_remote_code.add(model)
|
||||
for url in func.endpoints():
|
||||
self.allow_endpoint(url)
|
||||
|
||||
def function_to_metadata(self, conf: "EmbeddingFunctionConfig"):
|
||||
"""
|
||||
Convert the given embedding function and source / vector column configs
|
||||
into a config dictionary that can be serialized into arrow metadata
|
||||
"""
|
||||
func = conf.function
|
||||
name = getattr(
|
||||
func, "__embedding_function_registry_alias__", func.__class__.__name__
|
||||
)
|
||||
if isinstance(func, UntrustedEmbeddingFunction):
|
||||
name = func.function_name
|
||||
else:
|
||||
name = getattr(
|
||||
func, "__embedding_function_registry_alias__", func.__class__.__name__
|
||||
)
|
||||
json_data = func.safe_model_dump()
|
||||
return {
|
||||
"name": name,
|
||||
@@ -180,6 +298,13 @@ class EmbeddingFunctionRegistry:
|
||||
return self._variables[name]
|
||||
|
||||
|
||||
_ANY_REVISION = object()
|
||||
|
||||
|
||||
def _normalize_url(url: str) -> str:
|
||||
return url.rstrip("/")
|
||||
|
||||
|
||||
# Global instance
|
||||
__REGISTRY__ = EmbeddingFunctionRegistry()
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
from typing import List, Union
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -27,19 +27,32 @@ class SentenceTransformerEmbeddings(TextEmbeddingFunction):
|
||||
The device to use for the model
|
||||
normalize : bool, default True
|
||||
Whether to normalize the embeddings
|
||||
trust_remote_code : bool, default True
|
||||
Whether to trust the remote code
|
||||
trust_remote_code : bool, default False
|
||||
Whether to run code from the model's repository. Only enable this for
|
||||
repositories you trust. When restored from table metadata, this is only
|
||||
enabled for models allowed with
|
||||
[allow_remote_code][lancedb.embeddings.EmbeddingFunctionRegistry.allow_remote_code].
|
||||
revision : Optional[str], default None
|
||||
The model revision to load, such as a commit hash.
|
||||
"""
|
||||
|
||||
name: str = "all-MiniLM-L6-v2"
|
||||
device: str = "cpu"
|
||||
normalize: bool = True
|
||||
trust_remote_code: bool = True
|
||||
trust_remote_code: bool = False
|
||||
revision: Optional[str] = None
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._ndims = None
|
||||
|
||||
def remote_code_model(self) -> Optional[Tuple[str, Optional[str]]]:
|
||||
return (self.name, self.revision)
|
||||
|
||||
@classmethod
|
||||
def _legacy_trust_remote_code_default(cls) -> bool:
|
||||
return True
|
||||
|
||||
@property
|
||||
def embedding_model(self):
|
||||
"""
|
||||
@@ -84,5 +97,8 @@ class SentenceTransformerEmbeddings(TextEmbeddingFunction):
|
||||
"sentence_transformers", "sentence-transformers"
|
||||
)
|
||||
return sentence_transformers.SentenceTransformer(
|
||||
self.name, device=self.device, trust_remote_code=self.trust_remote_code
|
||||
self.name,
|
||||
device=self.device,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
revision=self.revision,
|
||||
)
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
|
||||
|
||||
from functools import cached_property
|
||||
from typing import List, Any
|
||||
from typing import List, Any, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -34,7 +34,11 @@ class TransformersEmbeddingFunction(EmbeddingFunction):
|
||||
Whether or not to allow for custom models defined on the HuggingFace
|
||||
Hub in their own modeling files. This option should only be set to True
|
||||
for repositories you trust and in which you have read the code, as it
|
||||
will execute code present on the Hub on your local machine.
|
||||
will execute code present on the Hub on your local machine. When
|
||||
restored from table metadata, this is only enabled for models allowed with
|
||||
[allow_remote_code][lancedb.embeddings.EmbeddingFunctionRegistry.allow_remote_code].
|
||||
revision : Optional[str]
|
||||
The model revision to load, such as a commit hash. Default is None.
|
||||
|
||||
to download package, run :
|
||||
`pip install transformers`
|
||||
@@ -45,6 +49,7 @@ class TransformersEmbeddingFunction(EmbeddingFunction):
|
||||
name: str = "colbert-ir/colbertv2.0"
|
||||
device: str = "cpu"
|
||||
trust_remote_code: bool = False
|
||||
revision: Optional[str] = None
|
||||
_tokenizer: Any = PrivateAttr()
|
||||
_model: Any = PrivateAttr()
|
||||
|
||||
@@ -52,12 +57,17 @@ class TransformersEmbeddingFunction(EmbeddingFunction):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._ndims = None
|
||||
transformers = attempt_import_or_raise("transformers")
|
||||
self._tokenizer = transformers.AutoTokenizer.from_pretrained(self.name)
|
||||
self._tokenizer = transformers.AutoTokenizer.from_pretrained(
|
||||
self.name, revision=self.revision
|
||||
)
|
||||
self._model = transformers.AutoModel.from_pretrained(
|
||||
self.name, trust_remote_code=self.trust_remote_code
|
||||
self.name, trust_remote_code=self.trust_remote_code, revision=self.revision
|
||||
)
|
||||
self._model.to(self.device)
|
||||
|
||||
def remote_code_model(self) -> Optional[Tuple[str, Optional[str]]]:
|
||||
return (self.name, self.revision)
|
||||
|
||||
model_config = ConfigDict(ignored_types=(cached_property,))
|
||||
|
||||
def ndims(self):
|
||||
|
||||
@@ -304,6 +304,11 @@ def _is_non_retryable_error(error: Exception) -> bool:
|
||||
Returns:
|
||||
True if the error should not be retried, False otherwise
|
||||
"""
|
||||
from .base import UntrustedEmbeddingConfigError
|
||||
|
||||
if isinstance(error, UntrustedEmbeddingConfigError):
|
||||
return True
|
||||
|
||||
# Check for OpenAI authentication errors
|
||||
error_type = type(error).__name__
|
||||
if error_type == "AuthenticationError":
|
||||
|
||||
@@ -87,6 +87,9 @@ class WatsonxEmbeddings(TextEmbeddingFunction):
|
||||
url: Optional[str] = None
|
||||
params: Optional[Dict] = None
|
||||
|
||||
def endpoints(self) -> List[str]:
|
||||
return [self.url] if self.url else []
|
||||
|
||||
@staticmethod
|
||||
def sensitive_keys():
|
||||
return ["api_key"]
|
||||
|
||||
@@ -0,0 +1,291 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""
|
||||
Restoring an embedding function from table metadata must not let the table grant
|
||||
itself trust: remote model code and custom endpoints need approval from the
|
||||
process that opens the table.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
import lancedb
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from lancedb.embeddings import (
|
||||
EmbeddingFunctionConfig,
|
||||
OllamaEmbeddings,
|
||||
OpenAIEmbeddings,
|
||||
SentenceTransformerEmbeddings,
|
||||
UntrustedEmbeddingConfigError,
|
||||
get_registry,
|
||||
)
|
||||
from lancedb.embeddings.base import UntrustedEmbeddingFunction
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def fresh_trust(monkeypatch):
|
||||
registry = get_registry()
|
||||
monkeypatch.setattr(registry, "_allowed_remote_code", set())
|
||||
monkeypatch.setattr(registry, "_allowed_endpoints", set())
|
||||
|
||||
|
||||
def metadata_for(name: str, model: dict) -> dict:
|
||||
entry = {
|
||||
"name": name,
|
||||
"model": model,
|
||||
"source_column": "text",
|
||||
"vector_column": "vector",
|
||||
}
|
||||
return {b"embedding_functions": json.dumps([entry]).encode()}
|
||||
|
||||
|
||||
def restore(name: str, model: dict):
|
||||
return get_registry().parse_functions(metadata_for(name, model))["vector"].function
|
||||
|
||||
|
||||
def test_sentence_transformers_does_not_trust_remote_code_by_default():
|
||||
assert SentenceTransformerEmbeddings.create().trust_remote_code is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["sentence-transformers", "huggingface", "colbert"])
|
||||
def test_restored_remote_code_requires_allowlist(name):
|
||||
model = {"name": "attacker/model", "trust_remote_code": True}
|
||||
func = restore(name, model)
|
||||
|
||||
with pytest.raises(UntrustedEmbeddingConfigError, match="attacker/model"):
|
||||
func.compute_query_embeddings("hello")
|
||||
with pytest.raises(UntrustedEmbeddingConfigError):
|
||||
func.ndims()
|
||||
|
||||
|
||||
def test_blocked_function_does_not_construct_model():
|
||||
# huggingface loads its model in __init__, so the check has to happen first.
|
||||
with patch(
|
||||
"lancedb.embeddings.transformers.attempt_import_or_raise"
|
||||
) as attempt_import:
|
||||
restore("huggingface", {"name": "attacker/model", "trust_remote_code": True})
|
||||
attempt_import.assert_not_called()
|
||||
|
||||
|
||||
def test_allowlisted_model_restores_with_remote_code():
|
||||
get_registry().allow_remote_code("org/model")
|
||||
func = restore(
|
||||
"sentence-transformers", {"name": "org/model", "trust_remote_code": True}
|
||||
)
|
||||
assert isinstance(func, SentenceTransformerEmbeddings)
|
||||
assert func.trust_remote_code is True
|
||||
|
||||
|
||||
def test_allowlist_does_not_cover_other_models():
|
||||
get_registry().allow_remote_code("org/model")
|
||||
func = restore(
|
||||
"sentence-transformers", {"name": "org/other", "trust_remote_code": True}
|
||||
)
|
||||
with pytest.raises(UntrustedEmbeddingConfigError, match="org/other"):
|
||||
func.compute_query_embeddings("hello")
|
||||
|
||||
|
||||
def test_pinned_revision_must_match():
|
||||
get_registry().allow_remote_code("org/model", revision="abc123")
|
||||
|
||||
pinned = restore(
|
||||
"sentence-transformers",
|
||||
{"name": "org/model", "revision": "abc123", "trust_remote_code": True},
|
||||
)
|
||||
assert pinned.trust_remote_code is True
|
||||
|
||||
for model in [
|
||||
{"name": "org/model", "revision": "def456", "trust_remote_code": True},
|
||||
{"name": "org/model", "trust_remote_code": True},
|
||||
]:
|
||||
func = restore("sentence-transformers", model)
|
||||
with pytest.raises(UntrustedEmbeddingConfigError, match="org/model"):
|
||||
func.compute_query_embeddings("hello")
|
||||
|
||||
|
||||
def test_restored_trust_flag_follows_allowlist_when_absent():
|
||||
# Tables written before the default changed don't record trust_remote_code.
|
||||
assert restore(
|
||||
"sentence-transformers", {"name": "org/model"}
|
||||
).trust_remote_code is (False)
|
||||
get_registry().allow_remote_code("org/model")
|
||||
assert restore("sentence-transformers", {"name": "org/model"}).trust_remote_code
|
||||
|
||||
|
||||
def test_created_function_records_trust_flag():
|
||||
# Otherwise a new table using the default would look like a pre-fix table,
|
||||
# where a missing flag meant True.
|
||||
conf = EmbeddingFunctionConfig(
|
||||
source_column="text",
|
||||
vector_column="vector",
|
||||
function=SentenceTransformerEmbeddings.create(name="org/model"),
|
||||
)
|
||||
assert get_registry().function_to_metadata(conf)["model"] == {
|
||||
"name": "org/model",
|
||||
"trust_remote_code": False,
|
||||
}
|
||||
|
||||
get_registry().allow_remote_code("org/model")
|
||||
metadata = get_registry().get_table_metadata([conf])
|
||||
restored = get_registry().parse_functions(metadata)["vector"].function
|
||||
assert restored.trust_remote_code is False
|
||||
|
||||
|
||||
def test_huggingface_missing_flag_keeps_false_default():
|
||||
# huggingface always defaulted to False, so approval must not flip it.
|
||||
get_registry().allow_remote_code("org/model")
|
||||
with patch("lancedb.embeddings.transformers.attempt_import_or_raise"):
|
||||
func = restore("huggingface", {"name": "org/model"})
|
||||
assert func.trust_remote_code is False
|
||||
|
||||
|
||||
def test_restored_explicit_false_stays_false():
|
||||
get_registry().allow_remote_code("org/model")
|
||||
func = restore(
|
||||
"sentence-transformers", {"name": "org/model", "trust_remote_code": False}
|
||||
)
|
||||
assert func.trust_remote_code is False
|
||||
|
||||
|
||||
def test_create_trusts_its_own_config(tmp_path):
|
||||
func = SentenceTransformerEmbeddings.create(
|
||||
name="org/model", trust_remote_code=True
|
||||
)
|
||||
db = lancedb.connect(tmp_path)
|
||||
schema = pa.schema(
|
||||
[pa.field("text", pa.string()), pa.field("vector", pa.list_(pa.float32(), 4))]
|
||||
)
|
||||
db.create_table(
|
||||
"t",
|
||||
schema=schema,
|
||||
embedding_functions=[
|
||||
EmbeddingFunctionConfig(
|
||||
source_column="text", vector_column="vector", function=func
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
restored = db.open_table("t").embedding_functions["vector"].function
|
||||
assert isinstance(restored, SentenceTransformerEmbeddings)
|
||||
assert restored.trust_remote_code is True
|
||||
|
||||
|
||||
def test_create_trusts_only_the_revision_it_selected():
|
||||
# Anyone can push a PR ref to a Hugging Face repo, so trusting a model's
|
||||
# default branch must not extend to its other revisions.
|
||||
SentenceTransformerEmbeddings.create(name="org/model", trust_remote_code=True)
|
||||
func = restore(
|
||||
"sentence-transformers",
|
||||
{"name": "org/model", "revision": "refs/pr/123", "trust_remote_code": True},
|
||||
)
|
||||
assert isinstance(func, UntrustedEmbeddingFunction)
|
||||
|
||||
same = restore(
|
||||
"sentence-transformers", {"name": "org/model", "trust_remote_code": True}
|
||||
)
|
||||
assert same.trust_remote_code is True
|
||||
|
||||
|
||||
def test_explicit_unpinned_approval_covers_any_revision():
|
||||
get_registry().allow_remote_code("org/model")
|
||||
func = restore(
|
||||
"sentence-transformers",
|
||||
{"name": "org/model", "revision": "abc123", "trust_remote_code": True},
|
||||
)
|
||||
assert func.trust_remote_code is True
|
||||
|
||||
|
||||
def test_missing_flag_does_not_inherit_unapproved_trust_default():
|
||||
class LegacySentenceTransformer(SentenceTransformerEmbeddings):
|
||||
trust_remote_code: bool = True
|
||||
|
||||
registry = get_registry()
|
||||
registry.register("legacy-trust-default")(LegacySentenceTransformer)
|
||||
try:
|
||||
func = restore("legacy-trust-default", {"name": "attacker/model"})
|
||||
assert not func.trust_remote_code
|
||||
finally:
|
||||
del registry._functions["legacy-trust-default"]
|
||||
|
||||
|
||||
def test_trusting_one_model_does_not_trust_a_swapped_name():
|
||||
SentenceTransformerEmbeddings.create(name="org/model", trust_remote_code=True)
|
||||
func = restore(
|
||||
"sentence-transformers", {"name": "attacker/model", "trust_remote_code": True}
|
||||
)
|
||||
with pytest.raises(UntrustedEmbeddingConfigError):
|
||||
func.compute_query_embeddings("hello")
|
||||
|
||||
|
||||
ATTACKER_URL = "https://attacker.example/v1"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name,model",
|
||||
[
|
||||
("openai", {"base_url": ATTACKER_URL}),
|
||||
("ollama", {"host": ATTACKER_URL}),
|
||||
("ollama", {"ollama_client_kwargs": {"proxy": ATTACKER_URL}}),
|
||||
("watsonx", {"url": ATTACKER_URL}),
|
||||
],
|
||||
)
|
||||
def test_restored_custom_endpoint_requires_allowlist(name, model):
|
||||
func = restore(name, model)
|
||||
with pytest.raises(UntrustedEmbeddingConfigError, match="attacker.example"):
|
||||
func.compute_query_embeddings("hello")
|
||||
|
||||
get_registry().allow_endpoint(ATTACKER_URL)
|
||||
assert isinstance(restore(name, model), get_registry().get(name))
|
||||
|
||||
|
||||
def test_default_endpoints_need_no_allowlist():
|
||||
assert isinstance(restore("openai", {}), OpenAIEmbeddings)
|
||||
assert isinstance(
|
||||
restore("ollama", {"host": "http://localhost:11434"}), OllamaEmbeddings
|
||||
)
|
||||
|
||||
|
||||
def test_endpoint_allowlist_ignores_trailing_slash():
|
||||
get_registry().allow_endpoint("http://my-host:11434/")
|
||||
assert isinstance(
|
||||
restore("ollama", {"host": "http://my-host:11434"}), OllamaEmbeddings
|
||||
)
|
||||
|
||||
|
||||
def test_create_trusts_its_own_endpoint():
|
||||
OpenAIEmbeddings.create(base_url="http://my-proxy/v1")
|
||||
assert isinstance(
|
||||
restore("openai", {"base_url": "http://my-proxy/v1"}), OpenAIEmbeddings
|
||||
)
|
||||
|
||||
|
||||
def test_blocked_function_round_trips_metadata():
|
||||
model = {"name": "attacker/model", "trust_remote_code": True}
|
||||
func = restore("sentence-transformers", model)
|
||||
conf = EmbeddingFunctionConfig(
|
||||
source_column="text", vector_column="vector", function=func
|
||||
)
|
||||
entry = get_registry().function_to_metadata(conf)
|
||||
assert entry["name"] == "sentence-transformers"
|
||||
assert entry["model"] == model
|
||||
|
||||
|
||||
def test_table_with_blocked_function_still_opens(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
schema = pa.schema(
|
||||
[pa.field("text", pa.string()), pa.field("vector", pa.list_(pa.float32(), 2))],
|
||||
metadata=metadata_for(
|
||||
"sentence-transformers",
|
||||
{"name": "attacker/model", "trust_remote_code": True},
|
||||
),
|
||||
)
|
||||
db.create_table("t", schema=schema)
|
||||
tbl = db.open_table("t")
|
||||
|
||||
tbl.add([{"text": "hi", "vector": [1.0, 2.0]}])
|
||||
assert tbl.search([1.0, 2.0]).limit(1).to_list()[0]["text"] == "hi"
|
||||
|
||||
with pytest.raises(UntrustedEmbeddingConfigError, match="allow_remote_code"):
|
||||
tbl.search("hi").limit(1).to_list()
|
||||
Reference in New Issue
Block a user