Embedding backend for memory + zettelkasten semantic search.
We use model2vec <https://github.com/MinishLab/model2vec>_ static embeddings
instead of a transformer (fastembed/onnxruntime). model2vec is torch-free at
inference and its runtime deps (numpy, tokenizers, safetensors) ship abi3 /
pure-Python wheels with near-universal platform coverage (Intel + Apple-Silicon
macOS, manylinux2014, musllinux/Alpine, Windows), which onnxruntime does not.
The single adapter here implements kglite's duck-typed EmbeddingModel
protocol (a dimension attribute plus embed(texts) -> list[list[float]],
with optional load/unload), so it can be handed straight to
KnowledgeGraph.set_embedder.
Model source. The ~120MB model is fetched from a mirror we host on the angelo
GitHub release (the same pattern as the vendored kglite wheels) and cached in a
shared, cross-repo location on disk. The mirror avoids HuggingFace's anonymous
rate limits (which can bite when many machines sit behind one IP). If the mirror
is unreachable we transparently fall back to pulling the model from HuggingFace
by id. For air-gapped use, pre-fetch once (angelo doctor / angelo init
both prime the cache) or set HF_HUB_OFFLINE=1 after it's cached.
Model2VecAdapter
kglite EmbeddingModel adapter backed by a model2vec StaticModel.
The model is resolved via :func:resolve_model_source (mirror-first, with a
HuggingFace fallback) on first :meth:load, then cached on disk. Loading is
lazy so merely constructing the adapter (e.g. at import time) never touches
the network; the model is pulled on the first :meth:embed call.
Source code in memory/embedders.py
| class Model2VecAdapter:
"""kglite ``EmbeddingModel`` adapter backed by a model2vec ``StaticModel``.
The model is resolved via :func:`resolve_model_source` (mirror-first, with a
HuggingFace fallback) on first :meth:`load`, then cached on disk. Loading is
lazy so merely constructing the adapter (e.g. at import time) never touches
the network; the model is pulled on the first :meth:`embed` call.
"""
def __init__(self, model_name: str = EMBED_MODEL):
self._name = model_name
self._model = None
# Advertised up front so callers can compare against a stored dimension
# before the (potentially slow) model load; corrected to the model's
# real dim after load in case the model id ever changes.
self.dimension = EMBED_DIM
def load(self) -> None:
"""Instantiate the underlying StaticModel (idempotent)."""
if self._model is not None:
return
from model2vec import StaticModel
# Only resolve via the mirror for the default model; a caller that asked
# for a specific other model id gets it straight from HuggingFace.
source = resolve_model_source() if self._name == EMBED_MODEL else self._name
self._model = StaticModel.from_pretrained(source)
try:
self.dimension = int(self._model.dim)
except (AttributeError, TypeError, ValueError):
# Fall back to the declared dimension if the model doesn't expose
# ``dim`` for some reason; potion-retrieval-32M is 512-d.
self.dimension = EMBED_DIM
def unload(self) -> None:
"""Drop the model reference so its memory can be reclaimed."""
self._model = None
def embed(self, texts: Iterable[str]) -> list[list[float]]:
"""Return one embedding vector (list of floats) per input text."""
self.load()
vectors = self._model.encode(list(texts))
return vectors.tolist()
|
load
Instantiate the underlying StaticModel (idempotent).
Source code in memory/embedders.py
| def load(self) -> None:
"""Instantiate the underlying StaticModel (idempotent)."""
if self._model is not None:
return
from model2vec import StaticModel
# Only resolve via the mirror for the default model; a caller that asked
# for a specific other model id gets it straight from HuggingFace.
source = resolve_model_source() if self._name == EMBED_MODEL else self._name
self._model = StaticModel.from_pretrained(source)
try:
self.dimension = int(self._model.dim)
except (AttributeError, TypeError, ValueError):
# Fall back to the declared dimension if the model doesn't expose
# ``dim`` for some reason; potion-retrieval-32M is 512-d.
self.dimension = EMBED_DIM
|
unload
Drop the model reference so its memory can be reclaimed.
Source code in memory/embedders.py
| def unload(self) -> None:
"""Drop the model reference so its memory can be reclaimed."""
self._model = None
|
embed
embed(texts: Iterable[str]) -> list[list[float]]
Return one embedding vector (list of floats) per input text.
Source code in memory/embedders.py
| def embed(self, texts: Iterable[str]) -> list[list[float]]:
"""Return one embedding vector (list of floats) per input text."""
self.load()
vectors = self._model.encode(list(texts))
return vectors.tolist()
|
resolve_model_source
resolve_model_source(allow_download: bool = True) -> str
Return what to hand StaticModel.from_pretrained.
Resolution order
- The cached mirror dir, if already present (no network).
- Download from our mirror into the cache (unless
allow_download is
False), then return that dir.
- The HuggingFace model id, so model2vec fetches/caches it itself.
Step 3 is the graceful fallback when the mirror is unreachable; it also
covers HF_HUB_OFFLINE setups where model2vec finds the model in the HF
cache. Setting ANGELO_EMBED_SOURCE=hf skips the mirror entirely.
Source code in memory/embedders.py
| def resolve_model_source(allow_download: bool = True) -> str:
"""Return what to hand ``StaticModel.from_pretrained``.
Resolution order:
1. The cached mirror dir, if already present (no network).
2. Download from our mirror into the cache (unless ``allow_download`` is
False), then return that dir.
3. The HuggingFace model id, so model2vec fetches/caches it itself.
Step 3 is the graceful fallback when the mirror is unreachable; it also
covers ``HF_HUB_OFFLINE`` setups where model2vec finds the model in the HF
cache. Setting ``ANGELO_EMBED_SOURCE=hf`` skips the mirror entirely.
"""
if os.environ.get("ANGELO_EMBED_SOURCE", "").lower() == "hf":
return EMBED_MODEL
model_dir = _model_cache_dir() / _MODEL_DIRNAME
if _model_present(model_dir):
return str(model_dir)
if allow_download and _download_from_mirror(model_dir):
return str(model_dir)
return EMBED_MODEL
|
prefetch_model
prefetch_model() -> tuple[bool, str]
Prime the model cache ahead of first use (called by init/doctor).
Returns (ok, message). ok is True if the model is available locally
afterwards (mirror or HF cache). Never raises — a failure just means search
degrades to keyword matching until the model can be fetched.
Source code in memory/embedders.py
| def prefetch_model() -> tuple[bool, str]:
"""Prime the model cache ahead of first use (called by init/doctor).
Returns ``(ok, message)``. ``ok`` is True if the model is available locally
afterwards (mirror or HF cache). Never raises — a failure just means search
degrades to keyword matching until the model can be fetched.
"""
if os.environ.get("MEMORY_SKIP_EMBEDDINGS", "") == "1":
return False, "skipped (MEMORY_SKIP_EMBEDDINGS=1)"
model_dir = _model_cache_dir() / _MODEL_DIRNAME
if _model_present(model_dir):
return True, f"already cached at {model_dir}"
try:
source = resolve_model_source(allow_download=True)
except Exception as exc: # pragma: no cover - defensive
return False, f"failed: {exc}"
if source != EMBED_MODEL and _model_present(Path(source)):
return True, f"downloaded to {source}"
# Mirror unavailable; try to materialize via HuggingFace so a later run is
# offline-capable. Loading the model triggers the HF download + cache.
try:
from model2vec import StaticModel
StaticModel.from_pretrained(EMBED_MODEL)
return True, "downloaded from HuggingFace (mirror unavailable)"
except Exception as exc:
return False, f"unavailable (mirror + HuggingFace both failed: {exc})"
|
register_embedder
register_embedder(name: str, factory: Callable[[], Any]) -> None
Register (or replace) an embedder factory under name for the seam.
factory is a zero-arg callable returning an unloaded object conforming
to kglite's EmbeddingModel protocol (a dimension attribute plus
embed(texts) -> list[list[float]], with optional load/unload).
Selection is then via ANGELO_EMBEDDER=<name>. Registering a heavy backend
(e.g. a transformer) here keeps the swap to ONE place; nothing loads a model
until :func:resolve_embedder is called and the caller invokes load.
Source code in memory/embedders.py
| def register_embedder(name: str, factory: Callable[[], Any]) -> None:
"""Register (or replace) an embedder factory under ``name`` for the seam.
``factory`` is a zero-arg callable returning an *unloaded* object conforming
to kglite's ``EmbeddingModel`` protocol (a ``dimension`` attribute plus
``embed(texts) -> list[list[float]]``, with optional ``load``/``unload``).
Selection is then via ``ANGELO_EMBEDDER=<name>``. Registering a heavy backend
(e.g. a transformer) here keeps the swap to ONE place; nothing loads a model
until :func:`resolve_embedder` is called and the caller invokes ``load``.
"""
_EMBEDDER_REGISTRY[str(name).strip().lower()] = factory
|
resolve_embedder
resolve_embedder() -> Any
Resolve the process embedder from ANGELO_EMBEDDER (default model2vec).
Returns an UNLOADED adapter (the caller invokes load) so this never
touches the network. An unknown or empty selector degrades to the default
model2vec backend rather than raising — a bad env value must not take down
the whole search path.
Source code in memory/embedders.py
| def resolve_embedder() -> Any:
"""Resolve the process embedder from ``ANGELO_EMBEDDER`` (default model2vec).
Returns an UNLOADED adapter (the caller invokes ``load``) so this never
touches the network. An unknown or empty selector degrades to the default
model2vec backend rather than raising — a bad env value must not take down
the whole search path.
"""
name = os.environ.get("ANGELO_EMBEDDER", _DEFAULT_EMBEDDER).strip().lower()
factory = _EMBEDDER_REGISTRY.get(name) or _EMBEDDER_REGISTRY[_DEFAULT_EMBEDDER]
return factory()
|