diff --git a/mempalace/backends/chroma.py b/mempalace/backends/chroma.py index 9e2db52..02ff7e8 100644 --- a/mempalace/backends/chroma.py +++ b/mempalace/backends/chroma.py @@ -2153,6 +2153,7 @@ class ChromaBackend(BaseBackend): name = "chroma" capabilities = frozenset( { + "requires_explicit_embeddings", "supports_embeddings_in", "supports_embeddings_passthrough", "supports_embeddings_out", diff --git a/tests/conftest.py b/tests/conftest.py index 9f7c0e0..c7f1108 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,6 +11,9 @@ instead of the real user profile. """ import os +import hashlib +import math +import re import shutil import tempfile @@ -33,6 +36,60 @@ import pytest # noqa: E402 from mempalace.config import MempalaceConfig # noqa: E402 from mempalace.knowledge_graph import KnowledgeGraph # noqa: E402 +_TEST_EMBED_DIM = 384 +_TEST_TOKEN_RE = re.compile(r"\w+", re.UNICODE) +_REAL_EMBEDDING_TEST_MODULES = { + "test_embedding", + "test_embeddinggemma", +} + + +def _stable_test_embedding(text: str) -> list[float]: + """Small deterministic embedding for tests that do not test ONNX itself.""" + vec = [0.0] * _TEST_EMBED_DIM + tokens = _TEST_TOKEN_RE.findall((text or "").lower()) + if not tokens: + tokens = [""] + for token in tokens: + digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest() + vec[int.from_bytes(digest[:4], "little") % _TEST_EMBED_DIM] += 1.0 + norm = math.sqrt(sum(v * v for v in vec)) or 1.0 + return [v / norm for v in vec] + + +class _StableTestEmbeddingFunction: + @staticmethod + def name() -> str: + return "default" + + @staticmethod + def build_from_config(config): + _StableTestEmbeddingFunction.validate_config(config) + return _StableTestEmbeddingFunction() + + @staticmethod + def validate_config(config) -> None: + return + + def get_config(self) -> dict: + return {} + + def is_legacy(self) -> bool: + return False + + def default_space(self) -> str: + return "cosine" + + def supported_spaces(self) -> list[str]: + return ["cosine", "l2", "ip"] + + def embed_query(self, input): + return self(input=input) + + def __call__(self, input): + return [_stable_test_embedding(str(text)) for text in list(input or [])] + + # Redirect ChromaDB's ONNX model cache back to the real user's cache so tests # don't re-download the 79 MB model on every run. The HOME redirect above # would otherwise point ONNXMiniLM_L6_V2.DOWNLOAD_PATH at the empty temp dir. @@ -51,6 +108,35 @@ except ImportError: pass +@pytest.fixture(autouse=True) +def _stable_embedding_function_for_tests(request, monkeypatch): + """Keep ordinary tests off ChromaDB's native ONNX embedding path. + + Module-sized Windows runs were crashing inside onnxruntime after many raw + Chroma add/query calls. The embedding-specific tests opt out below; every + other test gets a deterministic in-process EF so it still exercises vector + writes/search without loading native ONNX sessions. + """ + module_name = getattr(getattr(request, "module", None), "__name__", "") + if module_name in _REAL_EMBEDDING_TEST_MODULES: + yield + return + + ef = _StableTestEmbeddingFunction() + + import mempalace.backends.chroma as chroma_mod + import mempalace.backends.embedding_wrapper as embedding_wrapper + import mempalace.embedding as embedding_mod + from chromadb.api.types import DefaultEmbeddingFunction + + monkeypatch.setattr(DefaultEmbeddingFunction, "__call__", lambda self, input: ef(input=input)) + monkeypatch.setattr(DefaultEmbeddingFunction, "embed_query", lambda self, input: ef(input=input)) + monkeypatch.setattr(embedding_mod, "get_embedding_function", lambda *_, **__: ef) + monkeypatch.setattr(chroma_mod.ChromaBackend, "_resolve_embedding_function", staticmethod(lambda: ef)) + monkeypatch.setattr(embedding_wrapper, "_embed_texts", lambda texts: ef(input=list(texts))) + yield + + @pytest.fixture(autouse=True) def _reset_mcp_cache(): """Reset cached MCP state between tests without importing mcp_server. diff --git a/tests/test_backends.py b/tests/test_backends.py index 56a8633..7823e10 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -509,6 +509,31 @@ def test_chroma_backend_create_true_creates_directory_and_collection(tmp_path): client.get_collection("mempalace_drawers") +def test_palace_wrapper_embeds_for_chroma(tmp_path, monkeypatch): + """Normal Chroma callers should not rely on Chroma's internal ONNX embedder.""" + import mempalace.backends.embedding_wrapper as embedding_wrapper + from mempalace.backends.embedding_wrapper import EmbeddingCollection + from mempalace.palace import get_collection + + calls = [] + + def fake_embed(texts): + texts = list(texts) + calls.append(texts) + return [[float(len(text)), 1.0, 0.0, 0.0] for text in texts] + + monkeypatch.setattr(embedding_wrapper, "_embed_texts", fake_embed) + + col = get_collection(str(tmp_path), create=True, backend="chroma") + assert isinstance(col, EmbeddingCollection) + + col.add(ids=["a"], documents=["alpha"], metadatas=[{"wing": "w"}]) + result = col.query(query_texts=["alpha"], n_results=1) + + assert result.ids == [["a"]] + assert calls == [["alpha"], ["alpha"]] + + def test_chroma_backend_creates_collection_with_cosine_distance(tmp_path): palace_path = tmp_path / "palace" diff --git a/tests/test_hybrid_search.py b/tests/test_hybrid_search.py index 98c11ae..e60f77c 100644 --- a/tests/test_hybrid_search.py +++ b/tests/test_hybrid_search.py @@ -50,6 +50,21 @@ def _seed_strong_closet_for(palace_path, drawer_id, source_file, topics): "generated_by": "test", }, ) + # Keep this fixture above Chroma's batch_size=2 persistence floor. A + # single-row closet collection can intermittently query as "Nothing found on + # disk" on Windows when the deterministic test embedder makes writes fast. + col.upsert( + ids=[f"closet_{drawer_id}_sentinel"], + documents=["test sentinel unrelated stabilization topic"], + metadatas=[ + { + "wing": "backend", + "room": "auth", + "source_file": f"{source_file}#sentinel", + "generated_by": "test", + } + ], + ) # ── core invariant: closets can only HELP, never HIDE ─────────────────────