feat(backends): shared embedder-identity sidecar + qdrant identity (RFC 001, #1730)
Completes the embedder-identity contract for the last backend (qdrant) and unifies identity persistence across all local-path backends. Identity is stored in a small per-palace sidecar (mempalace_embedder.json), NOT in a backend's mismatch marker. The marker's presence signals "palace initialized" (reads raise CollectionNotInitializedError when the marker exists but the store doesn't), so recording identity at first empty open must not create it — a sidecar is unguarded, so a brand-new palace records identity immediately. This fixes a latent gap (caught in review) where pgvector/qdrant palaces stayed permanently "unknown" because the marker isn't written until the first real write. - New mempalace/backends/_sidecar.py: shared read/write_embedder_sidecar with the isinstance robustness the review bots taught us; chroma, pgvector, and qdrant all use it (chroma's inline copy is removed, pgvector switches off the marker, qdrant adds it). - QdrantCollection.get/set_embedder_identity delegate to the sidecar, so palace.get_collection enforces a model swap on a qdrant palace exactly like the other backends — no live server needed for the check. Tests: sidecar roundtrip + brand-new-palace recording (creates sidecar without a marker) + enforcement model-swap raise, for qdrant and pgvector, all server-free; chroma identity tests unchanged. Full suite: 2495 passed, 82.52%. Closes #1730. Refs #743, #1724.
This commit is contained in:
Igor Lins e Silva
parent
4ceb880d19
commit
ed77125ae3
|
|
@ -0,0 +1,71 @@
|
|||
"""Shared embedder-identity sidecar (RFC 001).
|
||||
|
||||
A small JSON file in the palace directory, keyed by collection name, recording
|
||||
the embedder identity (``model_name`` / ``dimension``). It is deliberately
|
||||
*separate* from a backend's mismatch marker: a marker's presence signals
|
||||
"palace initialized" (reads raise ``CollectionNotInitializedError`` when the
|
||||
marker exists but the store doesn't), so recording identity at first empty open
|
||||
must not create one. The sidecar is unguarded, so a brand-new palace can record
|
||||
identity immediately — the same approach the chroma backend uses.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
EMBEDDER_SIDECAR_FILENAME = "mempalace_embedder.json"
|
||||
|
||||
|
||||
def read_embedder_sidecar(path: Optional[str], collection_name: Optional[str]):
|
||||
"""Return the recorded :class:`EmbedderIdentity` for ``collection_name``, or None.
|
||||
|
||||
Robust to a missing, unreadable, or malformed (non-dict) sidecar — any of
|
||||
those degrade to ``None`` (the ``unknown`` state) rather than raising.
|
||||
"""
|
||||
from .base import EmbedderIdentity
|
||||
|
||||
if not path or not collection_name or not os.path.isfile(path):
|
||||
return None
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
entry = data.get(collection_name)
|
||||
if not isinstance(entry, dict) or not entry.get("model_name"):
|
||||
return None
|
||||
return EmbedderIdentity(
|
||||
model_name=str(entry["model_name"]),
|
||||
dimension=int(entry.get("dimension") or 0),
|
||||
)
|
||||
|
||||
|
||||
def write_embedder_sidecar(path: Optional[str], collection_name: Optional[str], identity) -> None:
|
||||
"""Record ``identity`` for ``collection_name`` in the sidecar, creating it if needed.
|
||||
|
||||
No-ops for a missing path, missing collection name, or a nameless identity.
|
||||
Preserves other collections' entries; never raises on I/O failure.
|
||||
"""
|
||||
if not path or not collection_name or not identity or not getattr(identity, "model_name", ""):
|
||||
return
|
||||
data: dict = {}
|
||||
if os.path.isfile(path):
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
loaded = json.load(f)
|
||||
if isinstance(loaded, dict):
|
||||
data = loaded
|
||||
except (OSError, json.JSONDecodeError):
|
||||
data = {}
|
||||
data[collection_name] = {
|
||||
"model_name": str(identity.model_name),
|
||||
"dimension": int(identity.dimension or 0),
|
||||
}
|
||||
try:
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
os.chmod(path, 0o600)
|
||||
except (OSError, NotImplementedError):
|
||||
pass
|
||||
|
|
@ -17,6 +17,7 @@ from typing import Any, Optional
|
|||
import chromadb
|
||||
from chromadb.errors import NotFoundError as _ChromaNotFoundError
|
||||
|
||||
from ._sidecar import EMBEDDER_SIDECAR_FILENAME, read_embedder_sidecar, write_embedder_sidecar
|
||||
from .base import (
|
||||
BaseBackend,
|
||||
BaseCollection,
|
||||
|
|
@ -1756,54 +1757,13 @@ class ChromaCollection(BaseCollection):
|
|||
def _embedder_sidecar_path(self) -> Optional[str]:
|
||||
if not self._palace_path:
|
||||
return None
|
||||
return os.path.join(self._palace_path, "mempalace_embedder.json")
|
||||
return os.path.join(self._palace_path, EMBEDDER_SIDECAR_FILENAME)
|
||||
|
||||
def get_stored_embedder_identity(self):
|
||||
from .base import EmbedderIdentity
|
||||
|
||||
path = self._embedder_sidecar_path()
|
||||
name = self._collection_name()
|
||||
if not path or not name or not os.path.isfile(path):
|
||||
return None
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
entry = data.get(name)
|
||||
if not isinstance(entry, dict) or not entry.get("model_name"):
|
||||
return None
|
||||
return EmbedderIdentity(
|
||||
model_name=str(entry["model_name"]),
|
||||
dimension=int(entry.get("dimension") or 0),
|
||||
)
|
||||
return read_embedder_sidecar(self._embedder_sidecar_path(), self._collection_name())
|
||||
|
||||
def set_embedder_identity(self, identity) -> None:
|
||||
path = self._embedder_sidecar_path()
|
||||
name = self._collection_name()
|
||||
if not path or not name or not identity or not identity.model_name:
|
||||
return
|
||||
data: dict = {}
|
||||
if os.path.isfile(path):
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
loaded = json.load(f)
|
||||
if isinstance(loaded, dict):
|
||||
data = loaded
|
||||
except (OSError, json.JSONDecodeError):
|
||||
data = {}
|
||||
data[name] = {
|
||||
"model_name": str(identity.model_name),
|
||||
"dimension": int(identity.dimension or 0),
|
||||
}
|
||||
try:
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
os.chmod(path, 0o600)
|
||||
except (OSError, NotImplementedError):
|
||||
pass
|
||||
write_embedder_sidecar(self._embedder_sidecar_path(), self._collection_name(), identity)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from urllib import parse as urlparse
|
|||
|
||||
import numpy as np
|
||||
|
||||
from ._sidecar import EMBEDDER_SIDECAR_FILENAME, read_embedder_sidecar, write_embedder_sidecar
|
||||
from .base import (
|
||||
BackendClosedError,
|
||||
BackendError,
|
||||
|
|
@ -747,6 +748,8 @@ class PgVectorCollection(BaseCollection):
|
|||
return self._backend._get_embedder_identity(self._palace, self._collection_name)
|
||||
|
||||
def set_embedder_identity(self, identity) -> None:
|
||||
# Sidecar-backed (see PgVectorBackend), so this records even on a
|
||||
# brand-new palace whose mismatch marker doesn't exist yet.
|
||||
self._backend._set_embedder_identity(self._palace, self._collection_name, identity)
|
||||
|
||||
def _ensure_table(self, dimension: int) -> None:
|
||||
|
|
@ -1255,14 +1258,6 @@ class PgVectorBackend(BaseBackend):
|
|||
"palace_id": palace.id,
|
||||
"pgvector": self._marker_target(palace, config),
|
||||
}
|
||||
# Preserve recorded embedder identities across marker rewrites — a
|
||||
# rebuilt marker must not wipe per-collection model-name tracking.
|
||||
try:
|
||||
existing = self._read_marker(palace)
|
||||
except BackendMismatchError:
|
||||
existing = None
|
||||
if isinstance(existing, dict) and isinstance(existing.get("embedders"), dict):
|
||||
marker["embedders"] = existing["embedders"]
|
||||
marker_path = self._marker_path(palace.local_path)
|
||||
with open(marker_path, "w", encoding="utf-8") as f:
|
||||
json.dump(marker, f, indent=2, ensure_ascii=False)
|
||||
|
|
@ -1271,50 +1266,22 @@ class PgVectorBackend(BaseBackend):
|
|||
except (OSError, NotImplementedError):
|
||||
pass
|
||||
|
||||
# Embedder identity lives in a sidecar, NOT the backend marker: the marker's
|
||||
# presence signals "palace initialized" (reads raise CollectionNotInitialized
|
||||
# when the marker exists but the remote table doesn't), so recording identity
|
||||
# at first empty open must not create it. The sidecar is unguarded — like the
|
||||
# chroma sidecar — so a brand-new palace can record identity immediately.
|
||||
@staticmethod
|
||||
def _embedder_sidecar_path(palace: PalaceRef) -> Optional[str]:
|
||||
if not palace.local_path:
|
||||
return None
|
||||
return os.path.join(palace.local_path, EMBEDDER_SIDECAR_FILENAME)
|
||||
|
||||
def _get_embedder_identity(self, palace: PalaceRef, collection_name: str):
|
||||
from .base import EmbedderIdentity
|
||||
|
||||
try:
|
||||
marker = self._read_marker(palace)
|
||||
except BackendMismatchError:
|
||||
return None
|
||||
if not isinstance(marker, dict):
|
||||
return None
|
||||
embedders = marker.get("embedders")
|
||||
if not isinstance(embedders, dict):
|
||||
return None
|
||||
entry = embedders.get(collection_name)
|
||||
if not isinstance(entry, dict) or not entry.get("model_name"):
|
||||
return None
|
||||
return EmbedderIdentity(
|
||||
model_name=str(entry["model_name"]),
|
||||
dimension=int(entry.get("dimension") or 0),
|
||||
)
|
||||
return read_embedder_sidecar(self._embedder_sidecar_path(palace), collection_name)
|
||||
|
||||
def _set_embedder_identity(self, palace: PalaceRef, collection_name: str, identity) -> None:
|
||||
if not palace.local_path or not identity or not identity.model_name:
|
||||
return
|
||||
marker_path = self._marker_path(palace.local_path)
|
||||
if not os.path.isfile(marker_path):
|
||||
return
|
||||
try:
|
||||
marker = self._read_marker(palace) or {}
|
||||
except BackendMismatchError:
|
||||
return
|
||||
embedders = marker.get("embedders")
|
||||
if not isinstance(embedders, dict):
|
||||
embedders = {}
|
||||
embedders[collection_name] = {
|
||||
"model_name": str(identity.model_name),
|
||||
"dimension": int(identity.dimension or 0),
|
||||
}
|
||||
marker["embedders"] = embedders
|
||||
with open(marker_path, "w", encoding="utf-8") as f:
|
||||
json.dump(marker, f, indent=2, ensure_ascii=False)
|
||||
try:
|
||||
os.chmod(marker_path, 0o600)
|
||||
except (OSError, NotImplementedError):
|
||||
pass
|
||||
write_embedder_sidecar(self._embedder_sidecar_path(palace), collection_name, identity)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
def _client(self, config: _PgVectorConfig) -> _PgVectorClient:
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from urllib import request as urlrequest
|
|||
|
||||
import numpy as np
|
||||
|
||||
from ._sidecar import EMBEDDER_SIDECAR_FILENAME, read_embedder_sidecar, write_embedder_sidecar
|
||||
from .base import (
|
||||
BackendClosedError,
|
||||
BackendMismatchError,
|
||||
|
|
@ -661,6 +662,14 @@ class QdrantCollection(BaseCollection):
|
|||
def _marker_exists(self) -> bool:
|
||||
return self._backend._marker_exists(self._palace)
|
||||
|
||||
def get_stored_embedder_identity(self):
|
||||
return self._backend._get_embedder_identity(self._palace, self._collection_name)
|
||||
|
||||
def set_embedder_identity(self, identity) -> None:
|
||||
# Sidecar-backed (see QdrantBackend), so this records even on a
|
||||
# brand-new palace whose mismatch marker doesn't exist yet.
|
||||
self._backend._set_embedder_identity(self._palace, self._collection_name, identity)
|
||||
|
||||
def _remote_dimension(self) -> Optional[int]:
|
||||
try:
|
||||
info = self._client.get_collection_info(self._remote_collection)
|
||||
|
|
@ -1175,6 +1184,23 @@ class QdrantBackend(BaseBackend):
|
|||
except (OSError, NotImplementedError):
|
||||
pass
|
||||
|
||||
# Embedder identity lives in a sidecar, NOT the backend marker: the marker's
|
||||
# presence signals "palace initialized" (reads raise CollectionNotInitialized
|
||||
# when the marker exists but the remote collection doesn't), so recording
|
||||
# identity at first empty open must not create it. The sidecar is unguarded,
|
||||
# so a brand-new palace can record identity immediately.
|
||||
@staticmethod
|
||||
def _embedder_sidecar_path(palace: PalaceRef) -> Optional[str]:
|
||||
if not palace.local_path:
|
||||
return None
|
||||
return os.path.join(palace.local_path, EMBEDDER_SIDECAR_FILENAME)
|
||||
|
||||
def _get_embedder_identity(self, palace: PalaceRef, collection_name: str):
|
||||
return read_embedder_sidecar(self._embedder_sidecar_path(palace), collection_name)
|
||||
|
||||
def _set_embedder_identity(self, palace: PalaceRef, collection_name: str, identity) -> None:
|
||||
write_embedder_sidecar(self._embedder_sidecar_path(palace), collection_name, identity)
|
||||
|
||||
def _client(self, config: _QdrantConfig) -> _QdrantRESTClient:
|
||||
if self._closed:
|
||||
raise BackendClosedError("QdrantBackend has been closed")
|
||||
|
|
|
|||
|
|
@ -119,16 +119,17 @@ def test_chroma_identity_roundtrip_via_sidecar(tmp_path):
|
|||
assert os.path.isfile(os.path.join(str(tmp_path), "mempalace_embedder.json"))
|
||||
|
||||
|
||||
def test_pgvector_marker_preserves_identity_across_rewrite(tmp_path):
|
||||
# The marker is rebuilt on every collection create; identity must survive.
|
||||
def test_pgvector_identity_survives_marker_rewrite(tmp_path):
|
||||
# Identity lives in a sidecar, separate from the mismatch marker, so a
|
||||
# marker rebuild (which happens on every write) must not affect it.
|
||||
from mempalace.backends.pgvector import PgVectorBackend, _PgVectorConfig
|
||||
|
||||
backend = PgVectorBackend()
|
||||
cfg = _PgVectorConfig(dsn="postgresql://example", namespace=None)
|
||||
ref = PalaceRef(id=str(tmp_path), local_path=str(tmp_path))
|
||||
backend._write_marker(ref, cfg)
|
||||
# No marker needed to record identity — the sidecar is unguarded.
|
||||
backend._set_embedder_identity(ref, "mempalace_drawers", EmbedderIdentity("minilm", 384))
|
||||
backend._write_marker(ref, cfg) # rebuild must not wipe embedders
|
||||
backend._write_marker(ref, cfg)
|
||||
got = backend._get_embedder_identity(ref, "mempalace_drawers")
|
||||
assert got is not None and got.model_name == "minilm" and got.dimension == 384
|
||||
|
||||
|
|
@ -308,3 +309,100 @@ def test_chroma_corrupt_sidecar_returns_none(tmp_path):
|
|||
# And a subsequent set still works (overwrites the junk).
|
||||
col.set_embedder_identity(EmbedderIdentity("minilm", 384))
|
||||
assert col.get_stored_embedder_identity().model_name == "minilm"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# qdrant: identity persisted in the local marker (no live qdrant needed)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _qdrant_collection(tmp_path, *, write_marker=True):
|
||||
from mempalace.backends.qdrant import QdrantBackend, QdrantCollection, _QdrantConfig
|
||||
|
||||
backend = QdrantBackend()
|
||||
config = _QdrantConfig(url="http://localhost:6333", api_key=None, namespace=None)
|
||||
ref = PalaceRef(id=str(tmp_path), local_path=str(tmp_path))
|
||||
if write_marker:
|
||||
backend._write_marker(ref, config)
|
||||
# The identity methods read/write the local marker only; the client is
|
||||
# never touched, so a placeholder stands in for a live REST connection.
|
||||
return QdrantCollection(
|
||||
backend=backend,
|
||||
client=object(),
|
||||
config=config,
|
||||
palace=ref,
|
||||
collection_name="mempalace_drawers",
|
||||
remote_collection="mp_drawers_remote",
|
||||
)
|
||||
|
||||
|
||||
def test_qdrant_identity_survives_marker_rewrite(tmp_path):
|
||||
from mempalace.backends.qdrant import QdrantBackend, _QdrantConfig
|
||||
|
||||
backend = QdrantBackend()
|
||||
config = _QdrantConfig(url="http://localhost:6333", api_key=None, namespace=None)
|
||||
ref = PalaceRef(id=str(tmp_path), local_path=str(tmp_path))
|
||||
backend._write_marker(ref, config)
|
||||
backend._set_embedder_identity(ref, "mempalace_drawers", EmbedderIdentity("minilm", 384))
|
||||
backend._write_marker(ref, config) # rebuild must not wipe embedders
|
||||
got = backend._get_embedder_identity(ref, "mempalace_drawers")
|
||||
assert got is not None and got.model_name == "minilm" and got.dimension == 384
|
||||
|
||||
|
||||
def test_qdrant_collection_delegates_identity(tmp_path):
|
||||
col = _qdrant_collection(tmp_path)
|
||||
assert col.get_stored_embedder_identity() is None
|
||||
col.set_embedder_identity(EmbedderIdentity("minilm", 384))
|
||||
got = col.get_stored_embedder_identity()
|
||||
assert got is not None and got.model_name == "minilm" and got.dimension == 384
|
||||
|
||||
|
||||
def test_qdrant_set_identity_creates_sidecar_when_missing(tmp_path):
|
||||
# Brand-new palace whose first write hasn't created the marker yet:
|
||||
# recording identity must create it, not silently no-op into permanent
|
||||
# "unknown" (the marker-on-write vs record-on-open timing gap).
|
||||
col = _qdrant_collection(tmp_path, write_marker=False)
|
||||
assert not col._marker_exists()
|
||||
col.set_embedder_identity(EmbedderIdentity("minilm", 384))
|
||||
got = col.get_stored_embedder_identity()
|
||||
assert got is not None and got.model_name == "minilm"
|
||||
|
||||
|
||||
def _pgvector_collection(tmp_path, *, write_marker=True):
|
||||
from mempalace.backends.pgvector import PgVectorBackend, PgVectorCollection, _PgVectorConfig
|
||||
|
||||
backend = PgVectorBackend()
|
||||
config = _PgVectorConfig(dsn="postgresql://example", namespace=None)
|
||||
ref = PalaceRef(id=str(tmp_path), local_path=str(tmp_path))
|
||||
if write_marker:
|
||||
backend._write_marker(ref, config)
|
||||
return PgVectorCollection(
|
||||
backend=backend,
|
||||
client=object(),
|
||||
config=config,
|
||||
palace=ref,
|
||||
collection_name="mempalace_drawers",
|
||||
table="mp_drawers_t",
|
||||
)
|
||||
|
||||
|
||||
def test_pgvector_set_identity_creates_sidecar_when_missing(tmp_path):
|
||||
# Same brand-new-palace timing gap as qdrant: recording must create the
|
||||
# marker rather than no-op.
|
||||
col = _pgvector_collection(tmp_path, write_marker=False)
|
||||
assert not col._marker_exists()
|
||||
col.set_embedder_identity(EmbedderIdentity("minilm", 384))
|
||||
got = col.get_stored_embedder_identity()
|
||||
assert got is not None and got.model_name == "minilm"
|
||||
|
||||
|
||||
def test_qdrant_enforcement_model_swap_raises(tmp_path, monkeypatch, clear_identity_cache):
|
||||
# The enforcement check reads the marker (no server) and compares to the
|
||||
# configured model — a swap raises just like the local backends.
|
||||
from mempalace import palace as P
|
||||
|
||||
col = _qdrant_collection(tmp_path)
|
||||
col.set_embedder_identity(EmbedderIdentity("minilm", 384))
|
||||
monkeypatch.setenv("MEMPALACE_EMBEDDING_MODEL", "embeddinggemma")
|
||||
with pytest.raises(EmbedderIdentityMismatchError):
|
||||
P._enforce_embedder_identity(col, str(tmp_path), "mempalace_drawers", create=False)
|
||||
|
|
|
|||
Loading…
Reference in New Issue