refactor: route all chromadb access through ChromaBackend (v4 prep)

mempalace/backends/base.py

@@ -27,6 +27,11 @@ def upsert(
     ) -> None:
         raise NotImplementedError
 
+    @abstractmethod
+    def update(self, **kwargs: Any) -> None:
+        """Update existing records. Must raise if any ID is missing."""
+        raise NotImplementedError
+
     @abstractmethod
     def query(self, **kwargs: Any) -> Dict[str, Any]:
         raise NotImplementedError

mempalace/backends/chroma.py

@@ -55,6 +55,9 @@ def add(self, *, documents, ids, metadatas=None):
     def upsert(self, *, documents, ids, metadatas=None):
         self._collection.upsert(documents=documents, ids=ids, metadatas=metadatas)
 
+    def update(self, **kwargs):
+        self._collection.update(**kwargs)
+
     def query(self, **kwargs):
         return self._collection.query(**kwargs)
 
@@ -71,6 +74,44 @@ def count(self):
 class ChromaBackend:
     """Factory for MemPalace's default ChromaDB backend."""
 
+    def __init__(self):
+        # Per-instance client cache: palace_path -> chromadb.PersistentClient
+        self._clients: dict = {}
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _client(self, palace_path: str):
+        """Return a cached PersistentClient for *palace_path*, creating one if needed."""
+        if palace_path not in self._clients:
+            _fix_blob_seq_ids(palace_path)
+            self._clients[palace_path] = chromadb.PersistentClient(path=palace_path)
+        return self._clients[palace_path]
+
+    # ------------------------------------------------------------------
+    # Public static helpers (for callers that manage their own caching)
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def make_client(palace_path: str):
+        """Create and return a fresh PersistentClient (fix BLOB seq_ids first).
+
+        Intended for long-lived callers (e.g. mcp_server) that keep their own
+        inode/mtime-based client cache.
+        """
+        _fix_blob_seq_ids(palace_path)
+        return chromadb.PersistentClient(path=palace_path)
+
+    @staticmethod
+    def backend_version() -> str:
+        """Return the installed chromadb package version string."""
+        return chromadb.__version__
+
+    # ------------------------------------------------------------------
+    # Collection lifecycle
+    # ------------------------------------------------------------------
+
     def get_collection(self, palace_path: str, collection_name: str, create: bool = False):
         if not create and not os.path.isdir(palace_path):
             raise FileNotFoundError(palace_path)
@@ -82,12 +123,30 @@ def get_collection(self, palace_path: str, collection_name: str, create: bool =
             except (OSError, NotImplementedError):
                 pass
 
-        _fix_blob_seq_ids(palace_path)
-        client = chromadb.PersistentClient(path=palace_path)
+        client = self._client(palace_path)
         if create:
             collection = client.get_or_create_collection(
                 collection_name, metadata={"hnsw:space": "cosine"}
             )
         else:
             collection = client.get_collection(collection_name)
         return ChromaCollection(collection)
+
+    def get_or_create_collection(
+        self, palace_path: str, collection_name: str
+    ) -> "ChromaCollection":
+        """Shorthand for get_collection(..., create=True)."""
+        return self.get_collection(palace_path, collection_name, create=True)
+
+    def delete_collection(self, palace_path: str, collection_name: str) -> None:
+        """Delete *collection_name* from the palace at *palace_path*."""
+        self._client(palace_path).delete_collection(collection_name)
+
+    def create_collection(
+        self, palace_path: str, collection_name: str, hnsw_space: str = "cosine"
+    ) -> "ChromaCollection":
+        """Create (not get-or-create) *collection_name* with cosine HNSW space."""
+        collection = self._client(palace_path).create_collection(
+            collection_name, metadata={"hnsw:space": hnsw_space}
+        )
+        return ChromaCollection(collection)

mempalace/cli.py

@@ -172,8 +172,8 @@ def cmd_status(args):
 
 def cmd_repair(args):
     """Rebuild palace vector index from SQLite metadata."""
-    import chromadb
     import shutil
+    from .backends.chroma import ChromaBackend
     from .migrate import confirm_destructive_action, contains_palace_database
 
     palace_path = os.path.abspath(
@@ -193,10 +193,11 @@ def cmd_repair(args):
     print(f"{'=' * 55}\n")
     print(f"  Palace: {palace_path}")
 
+    backend = ChromaBackend()
+
     # Try to read existing drawers
     try:
-        client = chromadb.PersistentClient(path=palace_path)
-        col = client.get_collection("mempalace_drawers")
+        col = backend.get_collection(palace_path, "mempalace_drawers")
         total = col.count()
         print(f"  Drawers found: {total}")
     except Exception as e:
@@ -243,8 +244,8 @@ def cmd_repair(args):
     shutil.copytree(palace_path, backup_path)
 
     print("  Rebuilding collection...")
-    client.delete_collection("mempalace_drawers")
-    new_col = client.create_collection("mempalace_drawers", metadata={"hnsw:space": "cosine"})
+    backend.delete_collection(palace_path, "mempalace_drawers")
+    new_col = backend.create_collection(palace_path, "mempalace_drawers")
 
     filed = 0
     for i in range(0, len(all_ids), batch_size):
@@ -297,7 +298,7 @@ def cmd_mcp(args):
 
 def cmd_compress(args):
     """Compress drawers in a wing using AAAK Dialect."""
-    import chromadb
+    from .backends.chroma import ChromaBackend
     from .dialect import Dialect
 
     palace_path = os.path.expanduser(args.palace) if args.palace else MempalaceConfig().palace_path
@@ -317,9 +318,9 @@ def cmd_compress(args):
         dialect = Dialect()
 
     # Connect to palace
+    backend = ChromaBackend()
     try:
-        client = chromadb.PersistentClient(path=palace_path)
-        col = client.get_collection("mempalace_drawers")
+        col = backend.get_collection(palace_path, "mempalace_drawers")
     except Exception:
         print(f"\n  No palace found at {palace_path}")
         print("  Run: mempalace init <dir> then mempalace mine <dir>")
@@ -394,9 +395,7 @@ def cmd_compress(args):
     # Store compressed versions (unless dry-run)
     if not args.dry_run:
         try:
-            comp_col = client.get_or_create_collection(
-                "mempalace_compressed", metadata={"hnsw:space": "cosine"}
-            )
+            comp_col = backend.get_or_create_collection(palace_path, "mempalace_compressed")
             for doc_id, compressed, meta, stats in compressed_entries:
                 comp_meta = dict(meta)
                 comp_meta["compression_ratio"] = round(stats["size_ratio"], 1)

mempalace/dedup.py

@@ -27,7 +27,7 @@
 import time
 from collections import defaultdict
 
-import chromadb
+from .backends.chroma import ChromaBackend
 
 
 COLLECTION_NAME = "mempalace_drawers"
@@ -130,8 +130,7 @@ def dedup_source_group(col, drawer_ids, threshold=DEFAULT_THRESHOLD, dry_run=Tru
 def show_stats(palace_path=None):
     """Show duplication statistics without making changes."""
     palace_path = palace_path or _get_palace_path()
-    client = chromadb.PersistentClient(path=palace_path)
-    col = client.get_collection(COLLECTION_NAME)
+    col = ChromaBackend().get_collection(palace_path, COLLECTION_NAME)
 
     groups = get_source_groups(col)
 
@@ -163,8 +162,7 @@ def dedup_palace(
     print("  MemPalace Deduplicator")
     print(f"{'=' * 55}")
 
-    client = chromadb.PersistentClient(path=palace_path)
-    col = client.get_collection(COLLECTION_NAME)
+    col = ChromaBackend().get_collection(palace_path, COLLECTION_NAME)
 
     print(f"  Palace: {palace_path}")
     print(f"  Drawers: {col.count():,}")

mempalace/mcp_server.py

@@ -32,7 +32,7 @@
 
 from .config import MempalaceConfig, sanitize_name, sanitize_content
 from .version import __version__
-import chromadb
+from .backends.chroma import ChromaBackend, ChromaCollection
 from .query_sanitizer import sanitize_query
 from .searcher import search_memories
 from .palace_graph import (
@@ -177,7 +177,7 @@ def _get_client():
     mtime_changed = current_mtime != 0.0 and abs(current_mtime - _palace_db_mtime) > 0.01
 
     if _client_cache is None or inode_changed or mtime_changed:
-        _client_cache = chromadb.PersistentClient(path=_config.palace_path)
+        _client_cache = ChromaBackend.make_client(_config.palace_path)
         _collection_cache = None
         _metadata_cache = None
         _metadata_cache_time = 0
@@ -192,13 +192,15 @@ def _get_collection(create=False):
     try:
         client = _get_client()
         if create:
-            _collection_cache = client.get_or_create_collection(
-                _config.collection_name, metadata={"hnsw:space": "cosine"}
+            _collection_cache = ChromaCollection(
+                client.get_or_create_collection(
+                    _config.collection_name, metadata={"hnsw:space": "cosine"}
+                )
             )
             _metadata_cache = None
             _metadata_cache_time = 0
         elif _collection_cache is None:
-            _collection_cache = client.get_collection(_config.collection_name)
+            _collection_cache = ChromaCollection(client.get_collection(_config.collection_name))
             _metadata_cache = None
             _metadata_cache_time = 0
         return _collection_cache

mempalace/migrate.py

@@ -134,7 +134,7 @@ def confirm_destructive_action(
 
 def migrate(palace_path: str, dry_run: bool = False, confirm: bool = False):
     """Migrate a palace to the currently installed ChromaDB version."""
-    import chromadb
+    from .backends.chroma import ChromaBackend
 
     palace_path = os.path.abspath(os.path.expanduser(palace_path))
     db_path = os.path.join(palace_path, "chroma.sqlite3")
@@ -152,19 +152,19 @@ def migrate(palace_path: str, dry_run: bool = False, confirm: bool = False):
 
     # Detect version
     source_version = detect_chromadb_version(db_path)
+    target_version = ChromaBackend.backend_version()
     print(f"  Source:    ChromaDB {source_version}")
-    print(f"  Target:    ChromaDB {chromadb.__version__}")
+    print(f"  Target:    ChromaDB {target_version}")
 
     # Try reading with current chromadb first
     try:
-        client = chromadb.PersistentClient(path=palace_path)
-        col = client.get_collection("mempalace_drawers")
+        col = ChromaBackend().get_collection(palace_path, "mempalace_drawers")
         count = col.count()
-        print(f"\n  Palace is already readable by chromadb {chromadb.__version__}.")
+        print(f"\n  Palace is already readable by chromadb {target_version}.")
         print(f"  {count} drawers found. No migration needed.")
         return True
     except Exception:
-        print(f"\n  Palace is NOT readable by chromadb {chromadb.__version__}.")
+        print(f"\n  Palace is NOT readable by chromadb {target_version}.")
         print("  Extracting from SQLite directly...")
 
     # Extract all drawers via raw SQL
@@ -208,8 +208,8 @@ def migrate(palace_path: str, dry_run: bool = False, confirm: bool = False):
 
     temp_palace = tempfile.mkdtemp(prefix="mempalace_migrate_")
     print(f"  Creating fresh palace in {temp_palace}...")
-    client = chromadb.PersistentClient(path=temp_palace)
-    col = client.get_or_create_collection("mempalace_drawers", metadata={"hnsw:space": "cosine"})
+    fresh_backend = ChromaBackend()
+    col = fresh_backend.get_or_create_collection(temp_palace, "mempalace_drawers")
 
     # Re-import in batches
     batch_size = 500
@@ -227,7 +227,7 @@ def migrate(palace_path: str, dry_run: bool = False, confirm: bool = False):
     # Verify before swapping
     final_count = col.count()
     del col
-    del client
+    del fresh_backend
 
     # Swap: remove old palace, move new one into place
     print("  Swapping old palace for migrated version...")

mempalace/repair.py

@@ -32,7 +32,7 @@
 import shutil
 import time
 
-import chromadb
+from .backends.chroma import ChromaBackend
 
 
 COLLECTION_NAME = "mempalace_drawers"
@@ -90,8 +90,7 @@ def scan_palace(palace_path=None, only_wing=None):
     print(f"\n  Palace: {palace_path}")
     print("  Loading...")
 
-    client = chromadb.PersistentClient(path=palace_path)
-    col = client.get_collection(COLLECTION_NAME)
+    col = ChromaBackend().get_collection(palace_path, COLLECTION_NAME)
 
     where = {"wing": only_wing} if only_wing else None
     total = col.count()
@@ -174,8 +173,7 @@ def prune_corrupt(palace_path=None, confirm=False):
         print("  Re-run with --confirm to actually delete.")
         return
 
-    client = chromadb.PersistentClient(path=palace_path)
-    col = client.get_collection(COLLECTION_NAME)
+    col = ChromaBackend().get_collection(palace_path, COLLECTION_NAME)
     before = col.count()
     print(f"  Collection size before: {before:,}")
 
@@ -222,9 +220,9 @@ def rebuild_index(palace_path=None):
     print(f"{'=' * 55}\n")
     print(f"  Palace: {palace_path}")
 
-    client = chromadb.PersistentClient(path=palace_path)
+    backend = ChromaBackend()
     try:
-        col = client.get_collection(COLLECTION_NAME)
+        col = backend.get_collection(palace_path, COLLECTION_NAME)
         total = col.count()
     except Exception as e:
         print(f"  Error reading palace: {e}")
@@ -264,8 +262,8 @@ def rebuild_index(palace_path=None):
 
     # Rebuild with correct HNSW settings
     print("  Rebuilding collection with hnsw:space=cosine...")
-    client.delete_collection(COLLECTION_NAME)
-    new_col = client.create_collection(COLLECTION_NAME, metadata={"hnsw:space": "cosine"})
+    backend.delete_collection(palace_path, COLLECTION_NAME)
+    new_col = backend.create_collection(palace_path, COLLECTION_NAME)
 
     filed = 0
     for i in range(0, len(all_ids), batch_size):