Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 13 additions & 9 deletions mempalace/backends/chroma.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from typing import Any, Optional

import chromadb
from chromadb.errors import NotFoundError as _ChromaNotFoundError

from .base import (
BaseBackend,
Expand Down Expand Up @@ -1037,15 +1038,18 @@ def get_collection(
ef_kwargs = {"embedding_function": ef} if ef is not None else {}

if create:
collection = client.get_or_create_collection(
collection_name,
metadata={
"hnsw:space": hnsw_space,
"hnsw:num_threads": 1,
**_HNSW_BLOAT_GUARD,
},
**ef_kwargs,
)
try:
collection = client.get_collection(collection_name, **ef_kwargs)
except _ChromaNotFoundError:
collection = client.create_collection(
collection_name,
metadata={
"hnsw:space": hnsw_space,
"hnsw:num_threads": 1,
**_HNSW_BLOAT_GUARD,
},
**ef_kwargs,
)
else:
collection = client.get_collection(collection_name, **ef_kwargs)
_pin_hnsw_threads(collection)
Expand Down
26 changes: 26 additions & 0 deletions tests/test_backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -372,6 +372,32 @@ def test_chroma_backend_create_collection_sets_hnsw_bloat_guard(tmp_path):
assert col.metadata.get("hnsw:sync_threshold") == 50_000


def test_get_collection_create_true_is_idempotent(tmp_path):
"""Calling get_collection(create=True) twice on the same name must not crash.

ChromaDB 1.5.x's Rust bindings SIGSEGV when get_or_create_collection is
called with metadata that differs from the stored collection metadata. The
fix splits the call into get_collection -> fallback create_collection so the
metadata-comparison codepath in chromadb_rust_bindings is never reached for
existing collections. Regression guard for issue #1089.
"""
palace = str(tmp_path / "palace")
backend = ChromaBackend()
backend.get_collection(palace, collection_name="mempalace_drawers", create=True)
col2 = backend.get_collection(palace, collection_name="mempalace_drawers", create=True)
assert isinstance(col2, ChromaCollection)


def test_get_collection_create_true_preserves_existing_metadata(tmp_path):
"""Existing collection metadata is not overwritten when reopened with create=True."""
palace = str(tmp_path / "palace")
backend = ChromaBackend()
backend.get_collection(palace, collection_name="mempalace_drawers", create=True)
col = backend.get_collection(palace, collection_name="mempalace_drawers", create=True)
assert col._collection.metadata["hnsw:space"] == "cosine"
assert col._collection.metadata.get("hnsw:batch_size") == 50_000


def test_fix_blob_seq_ids_converts_blobs_to_integers(tmp_path):
"""Simulate a ChromaDB 0.6.x database with BLOB seq_ids and verify repair."""
db_path = tmp_path / "chroma.sqlite3"
Expand Down