Skip to content

Commit b060171

Browse files
authored
Merge pull request #852 from MemPalace/release/v4-prep
refactor: route all chromadb access through ChromaBackend (v4 prep)
2 parents 5a2f7db + 267a644 commit b060171

11 files changed

Lines changed: 215 additions & 189 deletions

File tree

mempalace/backends/base.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,11 @@ def upsert(
2727
) -> None:
2828
raise NotImplementedError
2929

30+
@abstractmethod
31+
def update(self, **kwargs: Any) -> None:
32+
"""Update existing records. Must raise if any ID is missing."""
33+
raise NotImplementedError
34+
3035
@abstractmethod
3136
def query(self, **kwargs: Any) -> Dict[str, Any]:
3237
raise NotImplementedError

mempalace/backends/chroma.py

Lines changed: 61 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,9 @@ def add(self, *, documents, ids, metadatas=None):
5555
def upsert(self, *, documents, ids, metadatas=None):
5656
self._collection.upsert(documents=documents, ids=ids, metadatas=metadatas)
5757

58+
def update(self, **kwargs):
59+
self._collection.update(**kwargs)
60+
5861
def query(self, **kwargs):
5962
return self._collection.query(**kwargs)
6063

@@ -71,6 +74,44 @@ def count(self):
7174
class ChromaBackend:
7275
"""Factory for MemPalace's default ChromaDB backend."""
7376

77+
def __init__(self):
78+
# Per-instance client cache: palace_path -> chromadb.PersistentClient
79+
self._clients: dict = {}
80+
81+
# ------------------------------------------------------------------
82+
# Internal helpers
83+
# ------------------------------------------------------------------
84+
85+
def _client(self, palace_path: str):
86+
"""Return a cached PersistentClient for *palace_path*, creating one if needed."""
87+
if palace_path not in self._clients:
88+
_fix_blob_seq_ids(palace_path)
89+
self._clients[palace_path] = chromadb.PersistentClient(path=palace_path)
90+
return self._clients[palace_path]
91+
92+
# ------------------------------------------------------------------
93+
# Public static helpers (for callers that manage their own caching)
94+
# ------------------------------------------------------------------
95+
96+
@staticmethod
97+
def make_client(palace_path: str):
98+
"""Create and return a fresh PersistentClient (fix BLOB seq_ids first).
99+
100+
Intended for long-lived callers (e.g. mcp_server) that keep their own
101+
inode/mtime-based client cache.
102+
"""
103+
_fix_blob_seq_ids(palace_path)
104+
return chromadb.PersistentClient(path=palace_path)
105+
106+
@staticmethod
107+
def backend_version() -> str:
108+
"""Return the installed chromadb package version string."""
109+
return chromadb.__version__
110+
111+
# ------------------------------------------------------------------
112+
# Collection lifecycle
113+
# ------------------------------------------------------------------
114+
74115
def get_collection(self, palace_path: str, collection_name: str, create: bool = False):
75116
if not create and not os.path.isdir(palace_path):
76117
raise FileNotFoundError(palace_path)
@@ -82,12 +123,30 @@ def get_collection(self, palace_path: str, collection_name: str, create: bool =
82123
except (OSError, NotImplementedError):
83124
pass
84125

85-
_fix_blob_seq_ids(palace_path)
86-
client = chromadb.PersistentClient(path=palace_path)
126+
client = self._client(palace_path)
87127
if create:
88128
collection = client.get_or_create_collection(
89129
collection_name, metadata={"hnsw:space": "cosine"}
90130
)
91131
else:
92132
collection = client.get_collection(collection_name)
93133
return ChromaCollection(collection)
134+
135+
def get_or_create_collection(
136+
self, palace_path: str, collection_name: str
137+
) -> "ChromaCollection":
138+
"""Shorthand for get_collection(..., create=True)."""
139+
return self.get_collection(palace_path, collection_name, create=True)
140+
141+
def delete_collection(self, palace_path: str, collection_name: str) -> None:
142+
"""Delete *collection_name* from the palace at *palace_path*."""
143+
self._client(palace_path).delete_collection(collection_name)
144+
145+
def create_collection(
146+
self, palace_path: str, collection_name: str, hnsw_space: str = "cosine"
147+
) -> "ChromaCollection":
148+
"""Create (not get-or-create) *collection_name* with cosine HNSW space."""
149+
collection = self._client(palace_path).create_collection(
150+
collection_name, metadata={"hnsw:space": hnsw_space}
151+
)
152+
return ChromaCollection(collection)

mempalace/cli.py

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -172,8 +172,8 @@ def cmd_status(args):
172172

173173
def cmd_repair(args):
174174
"""Rebuild palace vector index from SQLite metadata."""
175-
import chromadb
176175
import shutil
176+
from .backends.chroma import ChromaBackend
177177
from .migrate import confirm_destructive_action, contains_palace_database
178178

179179
palace_path = os.path.abspath(
@@ -193,10 +193,11 @@ def cmd_repair(args):
193193
print(f"{'=' * 55}\n")
194194
print(f" Palace: {palace_path}")
195195

196+
backend = ChromaBackend()
197+
196198
# Try to read existing drawers
197199
try:
198-
client = chromadb.PersistentClient(path=palace_path)
199-
col = client.get_collection("mempalace_drawers")
200+
col = backend.get_collection(palace_path, "mempalace_drawers")
200201
total = col.count()
201202
print(f" Drawers found: {total}")
202203
except Exception as e:
@@ -243,8 +244,8 @@ def cmd_repair(args):
243244
shutil.copytree(palace_path, backup_path)
244245

245246
print(" Rebuilding collection...")
246-
client.delete_collection("mempalace_drawers")
247-
new_col = client.create_collection("mempalace_drawers", metadata={"hnsw:space": "cosine"})
247+
backend.delete_collection(palace_path, "mempalace_drawers")
248+
new_col = backend.create_collection(palace_path, "mempalace_drawers")
248249

249250
filed = 0
250251
for i in range(0, len(all_ids), batch_size):
@@ -297,7 +298,7 @@ def cmd_mcp(args):
297298

298299
def cmd_compress(args):
299300
"""Compress drawers in a wing using AAAK Dialect."""
300-
import chromadb
301+
from .backends.chroma import ChromaBackend
301302
from .dialect import Dialect
302303

303304
palace_path = os.path.expanduser(args.palace) if args.palace else MempalaceConfig().palace_path
@@ -317,9 +318,9 @@ def cmd_compress(args):
317318
dialect = Dialect()
318319

319320
# Connect to palace
321+
backend = ChromaBackend()
320322
try:
321-
client = chromadb.PersistentClient(path=palace_path)
322-
col = client.get_collection("mempalace_drawers")
323+
col = backend.get_collection(palace_path, "mempalace_drawers")
323324
except Exception:
324325
print(f"\n No palace found at {palace_path}")
325326
print(" Run: mempalace init <dir> then mempalace mine <dir>")
@@ -394,9 +395,7 @@ def cmd_compress(args):
394395
# Store compressed versions (unless dry-run)
395396
if not args.dry_run:
396397
try:
397-
comp_col = client.get_or_create_collection(
398-
"mempalace_compressed", metadata={"hnsw:space": "cosine"}
399-
)
398+
comp_col = backend.get_or_create_collection(palace_path, "mempalace_compressed")
400399
for doc_id, compressed, meta, stats in compressed_entries:
401400
comp_meta = dict(meta)
402401
comp_meta["compression_ratio"] = round(stats["size_ratio"], 1)

mempalace/dedup.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
import time
2828
from collections import defaultdict
2929

30-
import chromadb
30+
from .backends.chroma import ChromaBackend
3131

3232

3333
COLLECTION_NAME = "mempalace_drawers"
@@ -130,8 +130,7 @@ def dedup_source_group(col, drawer_ids, threshold=DEFAULT_THRESHOLD, dry_run=Tru
130130
def show_stats(palace_path=None):
131131
"""Show duplication statistics without making changes."""
132132
palace_path = palace_path or _get_palace_path()
133-
client = chromadb.PersistentClient(path=palace_path)
134-
col = client.get_collection(COLLECTION_NAME)
133+
col = ChromaBackend().get_collection(palace_path, COLLECTION_NAME)
135134

136135
groups = get_source_groups(col)
137136

@@ -163,8 +162,7 @@ def dedup_palace(
163162
print(" MemPalace Deduplicator")
164163
print(f"{'=' * 55}")
165164

166-
client = chromadb.PersistentClient(path=palace_path)
167-
col = client.get_collection(COLLECTION_NAME)
165+
col = ChromaBackend().get_collection(palace_path, COLLECTION_NAME)
168166

169167
print(f" Palace: {palace_path}")
170168
print(f" Drawers: {col.count():,}")

mempalace/mcp_server.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232

3333
from .config import MempalaceConfig, sanitize_name, sanitize_content
3434
from .version import __version__
35-
import chromadb
35+
from .backends.chroma import ChromaBackend, ChromaCollection
3636
from .query_sanitizer import sanitize_query
3737
from .searcher import search_memories
3838
from .palace_graph import (
@@ -177,7 +177,7 @@ def _get_client():
177177
mtime_changed = current_mtime != 0.0 and abs(current_mtime - _palace_db_mtime) > 0.01
178178

179179
if _client_cache is None or inode_changed or mtime_changed:
180-
_client_cache = chromadb.PersistentClient(path=_config.palace_path)
180+
_client_cache = ChromaBackend.make_client(_config.palace_path)
181181
_collection_cache = None
182182
_metadata_cache = None
183183
_metadata_cache_time = 0
@@ -192,13 +192,15 @@ def _get_collection(create=False):
192192
try:
193193
client = _get_client()
194194
if create:
195-
_collection_cache = client.get_or_create_collection(
196-
_config.collection_name, metadata={"hnsw:space": "cosine"}
195+
_collection_cache = ChromaCollection(
196+
client.get_or_create_collection(
197+
_config.collection_name, metadata={"hnsw:space": "cosine"}
198+
)
197199
)
198200
_metadata_cache = None
199201
_metadata_cache_time = 0
200202
elif _collection_cache is None:
201-
_collection_cache = client.get_collection(_config.collection_name)
203+
_collection_cache = ChromaCollection(client.get_collection(_config.collection_name))
202204
_metadata_cache = None
203205
_metadata_cache_time = 0
204206
return _collection_cache

mempalace/migrate.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ def confirm_destructive_action(
134134

135135
def migrate(palace_path: str, dry_run: bool = False, confirm: bool = False):
136136
"""Migrate a palace to the currently installed ChromaDB version."""
137-
import chromadb
137+
from .backends.chroma import ChromaBackend
138138

139139
palace_path = os.path.abspath(os.path.expanduser(palace_path))
140140
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):
152152

153153
# Detect version
154154
source_version = detect_chromadb_version(db_path)
155+
target_version = ChromaBackend.backend_version()
155156
print(f" Source: ChromaDB {source_version}")
156-
print(f" Target: ChromaDB {chromadb.__version__}")
157+
print(f" Target: ChromaDB {target_version}")
157158

158159
# Try reading with current chromadb first
159160
try:
160-
client = chromadb.PersistentClient(path=palace_path)
161-
col = client.get_collection("mempalace_drawers")
161+
col = ChromaBackend().get_collection(palace_path, "mempalace_drawers")
162162
count = col.count()
163-
print(f"\n Palace is already readable by chromadb {chromadb.__version__}.")
163+
print(f"\n Palace is already readable by chromadb {target_version}.")
164164
print(f" {count} drawers found. No migration needed.")
165165
return True
166166
except Exception:
167-
print(f"\n Palace is NOT readable by chromadb {chromadb.__version__}.")
167+
print(f"\n Palace is NOT readable by chromadb {target_version}.")
168168
print(" Extracting from SQLite directly...")
169169

170170
# Extract all drawers via raw SQL
@@ -208,8 +208,8 @@ def migrate(palace_path: str, dry_run: bool = False, confirm: bool = False):
208208

209209
temp_palace = tempfile.mkdtemp(prefix="mempalace_migrate_")
210210
print(f" Creating fresh palace in {temp_palace}...")
211-
client = chromadb.PersistentClient(path=temp_palace)
212-
col = client.get_or_create_collection("mempalace_drawers", metadata={"hnsw:space": "cosine"})
211+
fresh_backend = ChromaBackend()
212+
col = fresh_backend.get_or_create_collection(temp_palace, "mempalace_drawers")
213213

214214
# Re-import in batches
215215
batch_size = 500
@@ -227,7 +227,7 @@ def migrate(palace_path: str, dry_run: bool = False, confirm: bool = False):
227227
# Verify before swapping
228228
final_count = col.count()
229229
del col
230-
del client
230+
del fresh_backend
231231

232232
# Swap: remove old palace, move new one into place
233233
print(" Swapping old palace for migrated version...")

mempalace/repair.py

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
import shutil
3333
import time
3434

35-
import chromadb
35+
from .backends.chroma import ChromaBackend
3636

3737

3838
COLLECTION_NAME = "mempalace_drawers"
@@ -90,8 +90,7 @@ def scan_palace(palace_path=None, only_wing=None):
9090
print(f"\n Palace: {palace_path}")
9191
print(" Loading...")
9292

93-
client = chromadb.PersistentClient(path=palace_path)
94-
col = client.get_collection(COLLECTION_NAME)
93+
col = ChromaBackend().get_collection(palace_path, COLLECTION_NAME)
9594

9695
where = {"wing": only_wing} if only_wing else None
9796
total = col.count()
@@ -174,8 +173,7 @@ def prune_corrupt(palace_path=None, confirm=False):
174173
print(" Re-run with --confirm to actually delete.")
175174
return
176175

177-
client = chromadb.PersistentClient(path=palace_path)
178-
col = client.get_collection(COLLECTION_NAME)
176+
col = ChromaBackend().get_collection(palace_path, COLLECTION_NAME)
179177
before = col.count()
180178
print(f" Collection size before: {before:,}")
181179

@@ -222,9 +220,9 @@ def rebuild_index(palace_path=None):
222220
print(f"{'=' * 55}\n")
223221
print(f" Palace: {palace_path}")
224222

225-
client = chromadb.PersistentClient(path=palace_path)
223+
backend = ChromaBackend()
226224
try:
227-
col = client.get_collection(COLLECTION_NAME)
225+
col = backend.get_collection(palace_path, COLLECTION_NAME)
228226
total = col.count()
229227
except Exception as e:
230228
print(f" Error reading palace: {e}")
@@ -264,8 +262,8 @@ def rebuild_index(palace_path=None):
264262

265263
# Rebuild with correct HNSW settings
266264
print(" Rebuilding collection with hnsw:space=cosine...")
267-
client.delete_collection(COLLECTION_NAME)
268-
new_col = client.create_collection(COLLECTION_NAME, metadata={"hnsw:space": "cosine"})
265+
backend.delete_collection(palace_path, COLLECTION_NAME)
266+
new_col = backend.create_collection(palace_path, COLLECTION_NAME)
269267

270268
filed = 0
271269
for i in range(0, len(all_ids), batch_size):

0 commit comments

Comments
 (0)