From cef1c62fe7a24f4cb7700e5dc79b747b5bbfc6bc Mon Sep 17 00:00:00 2001 From: Igor Lins e Silva <4753812+igorls@users.noreply.github.com> Date: Thu, 14 May 2026 05:00:48 -0300 Subject: [PATCH] feat(embedding): EF-mismatch error helper, offline tests, migration docs Three follow-ups bundled for the embeddinggemma EF added in 51702e9: 1. Offline tests for EmbeddinggemmaONNX (10 tests, 0.08s, no network). Mocks huggingface_hub.hf_hub_download, tokenizers.Tokenizer.from_file, and onnxruntime.InferenceSession so CI never pulls the 300 MB model. Guarded with pytest.importorskip so the file is skipped when the multilingual extra isn't installed. Covers: stable name(), lazy-load runs exactly once, output shape (n, 384) after MRL truncation, L2 normalization, sim prefix applied, dispatch from get_embedding_function(model="embeddinggemma"), cache key separates models, helpful ImportError when deps missing, env override. 2. Friendlier ChromaDB EF-name-mismatch error. Switching MEMPALACE_EMBEDDING_MODEL on an existing palace previously surfaced ChromaDB's bare "Embedding function conflict: new: X vs persisted: Y" ValueError. Now ChromaBackend.get_collection() wraps that error and points users at the two recovery paths: revert the env var, or run `mempalace repair rebuild-index --palace `. New _explain_ef_mismatch helper + 3 tests (unit + end-to-end). 3. Docs: CHANGELOG [Unreleased] entry covers both the new EF and the error wrapper. README Requirements section mentions the multilingual extra and points at the embedding.py docstring for the migration note. --- CHANGELOG.md | 9 ++ README.md | 2 +- mempalace/backends/chroma.py | 45 ++++++- tests/test_backends.py | 59 +++++++++ tests/test_embeddinggemma.py | 229 +++++++++++++++++++++++++++++++++++ 5 files changed, 342 insertions(+), 2 deletions(-) create mode 100644 tests/test_embeddinggemma.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a4b298..8aec034 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,15 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), --- +## [Unreleased] + +### Features + +- **Opt-in multilingual embedding model: `embeddinggemma-300m` ONNX (q8, MRL→384-dim).** MemPalace's default embedder (`all-MiniLM-L6-v2`) is trained English-only — cross-lingual cosine similarity on parallel-translated text averages 0.35 across DE/FR/HI/IT/KO/RU (RU at 0.17, near-orthogonal). A Russian-speaking user effectively cannot find their own memories, which breaks the "100% recall" design promise from CLAUDE.md. New `EmbeddinggemmaONNX` class in [`mempalace/embedding.py`](mempalace/embedding.py) brings this to 0.88 average (validated lossless vs the Ollama gguf via direct ONNX-runtime test). Lazy-downloads `onnx-community/embeddinggemma-300m-ONNX` (~300 MB) on first use via `huggingface_hub`. Output is truncated to 384 dims via Matryoshka Representation Learning so the model is a drop-in for ChromaDB's 384-dim collections — no schema change. Sim prefix (`"task: sentence similarity | query: "`) is applied automatically. Opt-in via `MEMPALACE_EMBEDDING_MODEL=embeddinggemma` env var; default stays `minilm` for back-compat. Switching models on an existing palace requires re-embedding (different vector space) — run `mempalace repair rebuild-index` after changing the env var. Install with `pip install mempalace[multilingual]`. (#1483) +- **Friendlier ChromaDB EF-name-mismatch error.** Switching `MEMPALACE_EMBEDDING_MODEL` on an existing palace without running `rebuild-index` previously surfaced ChromaDB's bare `Embedding function conflict: new: X vs persisted: Y` `ValueError` — accurate but didn't tell users how to recover. `ChromaBackend.get_collection()` now wraps that error and points at both options: revert the env var, or run `mempalace repair rebuild-index --palace `. (#1483) + +--- + ## [3.3.5] — 2026-05-09 ### Bug Fixes diff --git a/README.md b/README.md index bc67637..6bea9ce 100644 --- a/README.md +++ b/README.md @@ -168,7 +168,7 @@ verbatim drawer per user/assistant message, idempotent and resume-safe. - Python 3.9+ - A vector-store backend (ChromaDB by default) -- ~300 MB disk for the default embedding model +- ~300 MB disk for the default English-only embedding model. For multilingual recall (100+ languages, including non-Latin scripts), install with `pip install mempalace[multilingual]` and set `MEMPALACE_EMBEDDING_MODEL=embeddinggemma` — see the docstring at [`mempalace/embedding.py`](mempalace/embedding.py) for details and migration notes. No API key is required for the core benchmark path. diff --git a/mempalace/backends/chroma.py b/mempalace/backends/chroma.py index fe36f34..0fb6758 100644 --- a/mempalace/backends/chroma.py +++ b/mempalace/backends/chroma.py @@ -1125,6 +1125,38 @@ class ChromaBackend(BaseBackend): logger.exception("Failed to build embedding function; using chromadb default") return None + @staticmethod + def _explain_ef_mismatch(error: Exception, palace_path: str) -> Optional[str]: + """If ``error`` looks like a ChromaDB EF-name mismatch, return a + user-friendly explanation. Otherwise return None so the caller can + re-raise unchanged. + + Triggered when ``MEMPALACE_EMBEDDING_MODEL`` is switched on an + existing palace — ChromaDB persists the EF name on the collection + and refuses reads with a different one. The bare ValueError + ChromaDB raises doesn't mention rebuild-index or the env var, so + users hit it and don't know how to recover. + """ + msg = str(error) + if "Embedding function conflict" not in msg and "embedding function" not in msg.lower(): + return None + try: + from ..config import MempalaceConfig + + current_model = MempalaceConfig().embedding_model + except Exception: + current_model = "unknown" + return ( + f"Embedding model mismatch reading palace at {palace_path!r}.\n" + f" Underlying ChromaDB error: {msg}\n" + f" Current MEMPALACE_EMBEDDING_MODEL={current_model!r}.\n" + f" The palace was built with a different embedding model. Either:\n" + f" (a) revert the model: unset MEMPALACE_EMBEDDING_MODEL (or set " + f"the previous value), or\n" + f" (b) re-embed in place: `mempalace repair rebuild-index " + f"--palace {palace_path}` (writes new vectors with the current model)." + ) + # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ @@ -1332,8 +1364,19 @@ class ChromaBackend(BaseBackend): }, **ef_kwargs, ) + except ValueError as e: + explanation = self._explain_ef_mismatch(e, palace_path) + if explanation: + raise ValueError(explanation) from e + raise else: - collection = client.get_collection(collection_name, **ef_kwargs) + try: + collection = client.get_collection(collection_name, **ef_kwargs) + except ValueError as e: + explanation = self._explain_ef_mismatch(e, palace_path) + if explanation: + raise ValueError(explanation) from e + raise _pin_hnsw_threads(collection) return ChromaCollection(collection, palace_path=palace_path) diff --git a/tests/test_backends.py b/tests/test_backends.py index 90cf128..68e8d86 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -1235,6 +1235,65 @@ def test_chroma_backend_requarantines_after_inode_replacement(tmp_path, monkeypa ] +def test_explain_ef_mismatch_recognizes_chromadb_conflict(): + """When ChromaDB rejects a collection read due to an EF-name mismatch + (user changed MEMPALACE_EMBEDDING_MODEL on an existing palace), the + backend wraps the bare ValueError with a message that tells the user + how to recover. Without this, users hit a stack trace and don't know + rebuild-index exists.""" + err = ValueError( + "An embedding function already exists in the collection configuration, " + "and a new one is provided. Embedding function conflict: new: " + "embeddinggemma_300m vs persisted: default" + ) + msg = ChromaBackend._explain_ef_mismatch(err, "/tmp/palace.db") + assert msg is not None + assert "/tmp/palace.db" in msg + assert "MEMPALACE_EMBEDDING_MODEL" in msg + assert "rebuild-index" in msg + + +def test_explain_ef_mismatch_returns_none_for_unrelated_errors(): + """Don't paper over unrelated ValueErrors with the EF-mismatch message — + the caller needs to re-raise unmodified so debugging stays sane.""" + err = ValueError("Some other ChromaDB problem") + assert ChromaBackend._explain_ef_mismatch(err, "/tmp/palace.db") is None + + +def test_get_collection_translates_ef_mismatch_to_helpful_error(tmp_path): + """End-to-end: create a palace with the default EF, then try to read it + with a different EF name and confirm we surface the rebuild-index hint.""" + backend = ChromaBackend() + palace_path = str(tmp_path / "palace") + os.makedirs(palace_path, exist_ok=True) + + # Create the collection using the default (minilm-based) EF. + coll = backend.get_collection(palace_path, "drawers", create=True) + coll.add(documents=["seed"], ids=["1"]) + + # Now swap in an incompatible EF name (simulates the user setting + # MEMPALACE_EMBEDDING_MODEL=embeddinggemma without rebuild-index). + class _ConflictingEF: + @staticmethod + def name() -> str: + return "embeddinggemma_300m" + + def __call__(self, input): + return [[0.0] * 384 for _ in input] + + original_resolver = backend._resolve_embedding_function + backend._resolve_embedding_function = lambda: _ConflictingEF() + # Drop the cached client so the next call goes through the open path. + backend.close_palace(palace_path) + + try: + with pytest.raises(ValueError, match=r"rebuild-index"): + backend.get_collection(palace_path, "drawers", create=False) + finally: + backend._resolve_embedding_function = original_resolver + backend.close_palace(palace_path) + + def test_palace_get_collection_uses_configured_collection_name(monkeypatch): from mempalace import palace diff --git a/tests/test_embeddinggemma.py b/tests/test_embeddinggemma.py new file mode 100644 index 0000000..95c56e8 --- /dev/null +++ b/tests/test_embeddinggemma.py @@ -0,0 +1,229 @@ +"""Offline tests for EmbeddinggemmaONNX. + +The real ONNX model is ~300 MB and pulled from HuggingFace on first use, so +these tests mock huggingface_hub.hf_hub_download, tokenizers.Tokenizer, and +onnxruntime.InferenceSession to keep CI fast and network-free. + +Skipped when the multilingual extra isn't installed (huggingface_hub/ +tokenizers/numpy) — CI runs only core deps by default. +""" +import sys + +import pytest + +np = pytest.importorskip("numpy") +pytest.importorskip("huggingface_hub") +pytest.importorskip("tokenizers") + +import mempalace.embedding as embedding # noqa: E402 (after importorskip) + + +@pytest.fixture(autouse=True) +def isolate_embedding_state(monkeypatch): + monkeypatch.setattr(embedding, "_EF_CACHE", {}) + monkeypatch.setattr(embedding, "_WARNED", set()) + + +def _make_fake_session(out_dim=768): + """Fake onnxruntime InferenceSession that returns a deterministic tensor. + + Shape: (batch, out_dim). The values aren't important — tests check shape, + truncation, and L2-normalization, not numerical correctness. + """ + + class _Output: + def __init__(self, name): + self.name = name + + class _Session: + def __init__(self, *args, **kwargs): + pass + + def get_outputs(self): + return [_Output("last_hidden_state"), _Output("sentence_embedding")] + + def run(self, _output_names, feed): + batch = feed["input_ids"].shape[0] + # Deterministic non-trivial values so L2-norm isn't degenerate. + sent = np.arange(batch * out_dim, dtype=np.float32).reshape(batch, out_dim) + 1.0 + last_hidden = np.zeros((batch, feed["input_ids"].shape[1], out_dim), dtype=np.float32) + return [last_hidden, sent] + + return _Session + + +class _FakeTokenizer: + """Stand-in for tokenizers.Tokenizer with the methods _lazy_load uses.""" + + def __init__(self): + self._padding_enabled = False + self._truncation_enabled = False + self._truncation_max = None + + def enable_padding(self): + self._padding_enabled = True + + def enable_truncation(self, max_length): + self._truncation_enabled = True + self._truncation_max = max_length + + def encode_batch(self, texts): + class _Enc: + def __init__(self, n): + self.ids = [0] * n + self.attention_mask = [1] * n + + # Same fixed length per batch — real tokenizers pad to the longest. + max_len = max(len(t.split()) for t in texts) + return [_Enc(max_len) for _ in texts] + + +@pytest.fixture +def patched_lazy_load(monkeypatch): + """Patch the third-party deps imported inside EmbeddinggemmaONNX._lazy_load. + + Returns a dict of recording counters so tests can assert how many times + each was called (e.g. confirm lazy-load caches after first call). + """ + calls = {"hf_hub_download": 0, "InferenceSession": 0, "Tokenizer.from_file": 0} + + def fake_download(repo, filename=None, subfolder=None, **kwargs): + calls["hf_hub_download"] += 1 + return f"/tmp/fake/{subfolder or ''}/{filename}" + + fake_session_cls = _make_fake_session() + + def fake_session_ctor(*args, **kwargs): + calls["InferenceSession"] += 1 + return fake_session_cls() + + def fake_tokenizer_from_file(_path): + calls["Tokenizer.from_file"] += 1 + return _FakeTokenizer() + + # huggingface_hub and tokenizers are real packages (installed via the + # multilingual extra), so we patch the functions in place rather than + # injecting stub modules. + import huggingface_hub + import onnxruntime + import tokenizers + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", fake_download) + monkeypatch.setattr(onnxruntime, "InferenceSession", fake_session_ctor) + monkeypatch.setattr(tokenizers.Tokenizer, "from_file", staticmethod(fake_tokenizer_from_file)) + + return calls + + +def test_name_is_stable(): + """ChromaDB persists this on the collection — changing it breaks reads.""" + assert embedding.EmbeddinggemmaONNX.name() == "embeddinggemma_300m" + + +def test_lazy_load_runs_once(patched_lazy_load): + ef = embedding.EmbeddinggemmaONNX() + ef(["one"]) + ef(["two"]) + ef(["three"]) + assert patched_lazy_load["hf_hub_download"] == 2 # model + tokenizer, once total + assert patched_lazy_load["InferenceSession"] == 1 + assert patched_lazy_load["Tokenizer.from_file"] == 1 + + +def test_output_shape_is_truncated_to_384(patched_lazy_load): + ef = embedding.EmbeddinggemmaONNX() + out = ef(["one", "two", "three"]) + arr = np.asarray(out) + assert arr.shape == (3, 384), f"expected (3, 384) after MRL truncation, got {arr.shape}" + + +def test_output_is_l2_normalized(patched_lazy_load): + ef = embedding.EmbeddinggemmaONNX() + out = ef(["hello world", "another sentence"]) + arr = np.asarray(out) + norms = np.linalg.norm(arr, axis=1) + assert np.allclose(norms, 1.0, atol=1e-5), f"vectors not unit-norm: {norms}" + + +def test_prefix_is_applied(patched_lazy_load, monkeypatch): + captured = [] + original_encode_batch = _FakeTokenizer.encode_batch + + def fake_encode_batch(self, texts): + captured.extend(texts) + return original_encode_batch(self, texts) + + monkeypatch.setattr(_FakeTokenizer, "encode_batch", fake_encode_batch) + ef = embedding.EmbeddinggemmaONNX() + ef(["raw text one", "raw text two"]) + assert all(t.startswith("task: sentence similarity | query: ") for t in captured) + # And the raw text is preserved after the prefix. + assert any("raw text one" in t for t in captured) + + +def test_get_embedding_function_dispatches_to_embeddinggemma(monkeypatch): + """model='embeddinggemma' must build EmbeddinggemmaONNX, not the MiniLM EF.""" + monkeypatch.setattr( + embedding, "_resolve_providers", lambda device: (["CPUExecutionProvider"], "cpu") + ) + ef = embedding.get_embedding_function(device="cpu", model="embeddinggemma") + assert isinstance(ef, embedding.EmbeddinggemmaONNX) + assert ef.name() == "embeddinggemma_300m" + + +def test_cache_key_separates_models(monkeypatch): + """Switching model must not return the cached EF for the other model. + + The cache key changed from `providers` to `(model, providers)` for exactly + this reason — without it, the second call would silently reuse the wrong EF. + """ + + class DummyMiniLM: + def __init__(self, preferred_providers=None): + self.kind = "minilm" + + monkeypatch.setattr(embedding, "_build_ef_class", lambda: DummyMiniLM) + monkeypatch.setattr( + embedding, "_resolve_providers", lambda device: (["CPUExecutionProvider"], "cpu") + ) + + ml = embedding.get_embedding_function(device="cpu", model="minilm") + eg = embedding.get_embedding_function(device="cpu", model="embeddinggemma") + ml_again = embedding.get_embedding_function(device="cpu", model="minilm") + + assert ml is ml_again, "minilm should cache-hit on second call" + assert isinstance(eg, embedding.EmbeddinggemmaONNX), "embeddinggemma should not collide with minilm cache" + assert ml is not eg + + +def test_missing_deps_raise_helpful_error(monkeypatch): + """If the user hasn't installed `mempalace[multilingual]`, the error must + name the extra rather than just spilling a bare ImportError.""" + + # Drop tokenizers from sys.modules and block re-import, simulating a user + # who didn't install the multilingual extra. huggingface_hub and onnxruntime + # are present (they ship with core), so the failure should land on tokenizers. + monkeypatch.setitem(sys.modules, "tokenizers", None) + + ef = embedding.EmbeddinggemmaONNX() + with pytest.raises(ImportError, match=r"mempalace\[multilingual\]"): + ef(["anything"]) + + +def test_config_embedding_model_env_override(monkeypatch): + """MEMPALACE_EMBEDDING_MODEL env var must override the config file default.""" + from mempalace.config import MempalaceConfig + + monkeypatch.setenv("MEMPALACE_EMBEDDING_MODEL", "embeddinggemma") + assert MempalaceConfig().embedding_model == "embeddinggemma" + + monkeypatch.setenv("MEMPALACE_EMBEDDING_MODEL", "MiniLM") # case-insensitive + assert MempalaceConfig().embedding_model == "minilm" + + +def test_config_embedding_model_default_is_minilm(monkeypatch): + """Back-compat: existing installs without explicit config get minilm.""" + from mempalace.config import MempalaceConfig + + monkeypatch.delenv("MEMPALACE_EMBEDDING_MODEL", raising=False) + assert MempalaceConfig().embedding_model == "minilm"