项目文件夹

文件
wehub-resource-sync c3749daf48
Tests / test-linux (3.13) (push) Failing after 0s
Tests / test-linux (3.11) (push) Failing after 1s
Tests / lint (push) Failing after 0s
Tests / test-linux (3.9) (push) Failing after 1s
Docker / build (push) Failing after 1s
Docker / build-gpu (push) Failing after 2s
Tests / test-windows (push) Has been cancelled
Tests / test-macos (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:03:03 +08:00

2349 行
93 KiB
Python

此文件含有模棱两可的 Unicode 字符
此文件含有可能会与其他字符混淆的 Unicode 字符。 如果您是想特意这样的,可以安全地忽略该警告。 使用 Escape 按钮显示他们。
"""ChromaDB-backed MemPalace storage backend (RFC 001 reference implementation)."""
import contextlib
import datetime as _dt
import json
import logging
import math
import os
import pickle
import re
import shlex
import sqlite3
import time
from collections import defaultdict
from numbers import Integral
from pathlib import Path
from typing import Any, Optional
import chromadb
from chromadb.errors import NotFoundError as _ChromaNotFoundError
from ..config import sqlite_read_uri
from ._sidecar import EMBEDDER_SIDECAR_FILENAME, read_embedder_sidecar, write_embedder_sidecar
from .base import (
BaseBackend,
BaseCollection,
CollectionNotInitializedError,
GetResult,
HealthStatus,
LexicalHit,
LexicalResult,
PalaceNotFoundError,
PalaceRef,
QueryResult,
UnsupportedFilterError,
_IncludeSpec,
)
logger = logging.getLogger(__name__)
_REQUIRED_OPERATORS = frozenset({"$eq", "$ne", "$in", "$nin", "$and", "$or", "$contains"})
_OPTIONAL_OPERATORS = frozenset({"$gt", "$gte", "$lt", "$lte"})
_SUPPORTED_OPERATORS = _REQUIRED_OPERATORS | _OPTIONAL_OPERATORS
_TOKEN_RE = re.compile(r"\w{2,}", re.UNICODE)
# A healthy HNSW payload should keep link_lists.bin proportional to
# data_level0.bin. When link_lists.bin grows orders of magnitude larger than
# data_level0.bin, Chroma/HNSW can segfault while opening the segment even if
# index_metadata.pickle is structurally valid.
#
# The report in #1218 showed ratios above 300x, while healthy snapshots were far below 1x.
# Treat only >10x as corruption so normal flush lag or small segments do not get
# quarantined.
_HNSW_LINK_TO_DATA_MAX_RATIO = 10.0
def _hnsw_link_to_data_ratio(seg_dir: str) -> Optional[float]:
"""Return link_lists.bin / data_level0.bin size ratio for a segment.
``None`` means the ratio is not meaningful, usually because one file is
missing or data_level0.bin is empty. ``float("inf")`` means the files were
present but could not be statted safely, which should be treated as
suspicious by callers.
"""
link_path = os.path.join(seg_dir, "link_lists.bin")
data_path = os.path.join(seg_dir, "data_level0.bin")
if not (os.path.isfile(link_path) and os.path.isfile(data_path)):
return None
try:
data_size = os.path.getsize(data_path)
link_size = os.path.getsize(link_path)
except OSError:
return float("inf")
if data_size <= 0:
return None
return link_size / data_size
def _hnsw_metadata_marker_intact(seg_dir: str) -> bool:
"""Return True when ``index_metadata.pickle`` bears a complete envelope.
ChromaDB writes ``index_metadata.pickle`` last during a persist, so an
intact pickle envelope — protocol marker ``0x80`` at the head, ``STOP``
byte ``0x2e`` at the tail — proves the flush finished. Used as a
persist-completion marker: when present, the segment's on-disk shape
(including an all-layer-0 index with an empty ``link_lists.bin``) is the
one chromadb intentionally serialized, not a half-written one.
Deliberately byte-sniffs only; never deserializes. Deserialization can
execute arbitrary code, and the byte-sniff is enough to tell a complete
write from truncation or zero-fill. Assumes pickle protocol >= 2.
"""
meta_path = os.path.join(seg_dir, "index_metadata.pickle")
try:
if not os.path.isfile(meta_path) or os.path.getsize(meta_path) < 16:
return False
with open(meta_path, "rb") as f:
head = f.read(2)
f.seek(-1, 2) # last byte
tail = f.read(1)
except OSError:
return False
return len(head) == 2 and head[0] == 0x80 and tail == b"\x2e"
def _hnsw_link_lists_is_usable_for_payload(seg_dir: str) -> bool:
"""Return False when a non-trivial HNSW payload looks like a partial flush.
A zero-byte ``link_lists.bin`` is *not* corruption on its own. hnswlib
stores the entire layer-0 graph inside ``data_level0.bin`` and only writes
``link_lists.bin`` for elements promoted to level > 0. A small or
low-fanout index where every element stays on layer 0 therefore
serializes an empty ``link_lists.bin`` and loads/searches fine (#1716).
Treating that shape as corruption produced a self-perpetuating quarantine
loop: repair rebuilt the byte-identical all-layer-0 segment, which the next
cold start quarantined again, accumulating drift dirs without bound.
An empty ``link_lists.bin`` only signals trouble when the persist was
interrupted before it could finish. ChromaDB writes
``index_metadata.pickle`` last, so an intact metadata envelope proves the
flush completed and the empty ``link_lists.bin`` is the legitimate
all-layer-0 shape. Only when there is real payload, an empty
``link_lists.bin``, *and* no completion marker do we treat the segment as a
partial flush.
"""
data_path = os.path.join(seg_dir, "data_level0.bin")
link_path = os.path.join(seg_dir, "link_lists.bin")
try:
if not os.path.isfile(data_path):
return True
data_size = os.path.getsize(data_path)
if data_size <= _HNSW_MISSING_METADATA_DATA_FLOOR:
return True
if os.path.isfile(link_path) and os.path.getsize(link_path) > 0:
return True
except OSError:
return False
# Real payload with an empty/absent link_lists.bin: legitimate only when
# the persist completed (all-layer-0 index), proven by an intact metadata
# marker. Otherwise it is a half-written segment chromadb could segfault on.
return _hnsw_metadata_marker_intact(seg_dir)
def _hnsw_payload_appears_sane(seg_dir: str) -> bool:
"""Return False when HNSW payload files are structurally implausible."""
if not _hnsw_link_lists_is_usable_for_payload(seg_dir):
return False
ratio = _hnsw_link_to_data_ratio(seg_dir)
return ratio is None or ratio <= _HNSW_LINK_TO_DATA_MAX_RATIO
# HNSW batch/sync thresholds applied at collection creation.
#
# chromadb's Rust HNSW segment writes index_metadata.pickle and
# link_lists.bin only when internal counters cross both thresholds
# (batch_size gates _apply_batch; sync_threshold gates _persist).
# Records below both thresholds stay in memory and are lost on exit.
#
# Previously 50k/50k to work around link_lists.bin sparse-file bloat
# in pre-1.5.x Python chromadb (#344). chromadb >=1.5.4 Rust bindings
# (the minimum mempalace supports) do not exhibit that bloat; verified
# at batch_size=2 with 20k records: link_lists.bin = 171 KB, no
# sparse-file inflation.
#
# The 50k guard caused #1579: mines under 50k drawers never triggered
# _persist(), leaving index_metadata.pickle absent and link_lists.bin
# empty. quarantine_stale_hnsw then renamed the segment on every cold
# open after a 300s mtime gap, accumulating .drift-* directories.
#
# Lowered to 2 (empirical Rust-side minimum for chromadb >=1.5.4; the
# Rust bindings reject 1 with InvalidArgumentError) so any mine of 2+
# drawers triggers a natural persist. Existing palaces created under
# the old 50k guard keep those thresholds in their collection metadata
# until the user runs repair --mode from-sqlite --archive-existing.
_HNSW_BLOAT_GUARD = {
"hnsw:batch_size": 2,
"hnsw:sync_threshold": 2,
}
# Below this size, data_level0.bin is too small for a meaningful HNSW graph.
# Used by _hnsw_link_lists_is_usable_for_payload (empty link_lists is fine
# when data is trivially small) and _missing_dimensionality_appears_recoverable
# (don't attempt recovery on segments with negligible data).
_HNSW_MISSING_METADATA_DATA_FLOOR = 1024
def _validate_where(where: Optional[dict]) -> None:
"""Scan a where-clause for unknown operators and raise ``UnsupportedFilterError``.
Spec (RFC 001 §1.4): silent dropping of unknown operators is forbidden.
"""
if not where:
return
stack = [where]
while stack:
node = stack.pop()
if not isinstance(node, dict):
continue
for k, v in node.items():
if k.startswith("$") and k not in _SUPPORTED_OPERATORS:
raise UnsupportedFilterError(f"operator {k!r} not supported by chroma backend")
if isinstance(v, dict):
stack.append(v)
elif isinstance(v, list):
stack.extend(x for x in v if isinstance(x, dict))
def _tokenize(text: str) -> list[str]:
if not text:
return []
return _TOKEN_RE.findall(text.lower())
def _bm25_scores(
query: str,
documents: list[str],
k1: float = 1.5,
b: float = 0.75,
) -> list[float]:
query_terms = set(_tokenize(query))
n_docs = len(documents)
if not query_terms or n_docs == 0:
return [0.0] * n_docs
tokenized = [_tokenize(doc) for doc in documents]
doc_lens = [len(toks) for toks in tokenized]
if not any(doc_lens):
return [0.0] * n_docs
avgdl = sum(doc_lens) / n_docs or 1.0
df = {term: 0 for term in query_terms}
for toks in tokenized:
for term in set(toks) & query_terms:
df[term] += 1
idf = {
term: math.log((n_docs - df[term] + 0.5) / (df[term] + 0.5) + 1.0) for term in query_terms
}
scores = []
for toks, dl in zip(tokenized, doc_lens):
if dl == 0:
scores.append(0.0)
continue
tf: dict[str, int] = {}
for token in toks:
if token in query_terms:
tf[token] = tf.get(token, 0) + 1
score = 0.0
for term, freq in tf.items():
num = freq * (k1 + 1)
den = freq + k1 * (1 - b + b * dl / avgdl)
score += idf[term] * num / den
scores.append(score)
return scores
def _coerce_metadata_value(value: Any) -> Any:
if isinstance(value, bool):
return int(value)
return value
def _compare_metadata(actual: Any, op: str, expected: Any) -> bool:
actual = _coerce_metadata_value(actual)
expected = _coerce_metadata_value(expected)
if op == "$eq":
return actual == expected
if op == "$ne":
return actual != expected
if op == "$in":
return actual in (expected or [])
if op == "$nin":
return actual not in (expected or [])
if op == "$contains":
return str(expected) in str(actual or "")
try:
if op == "$gt":
return actual > expected
if op == "$gte":
return actual >= expected
if op == "$lt":
return actual < expected
if op == "$lte":
return actual <= expected
except TypeError:
return False
raise UnsupportedFilterError(f"operator {op!r} not supported by chroma backend")
def _matches_where(meta: dict, where: Optional[dict]) -> bool:
if not where:
return True
if not isinstance(where, dict):
return False
for key, expected in where.items():
if key == "$and":
if not all(_matches_where(meta, clause) for clause in expected or []):
return False
continue
if key == "$or":
if not any(_matches_where(meta, clause) for clause in expected or []):
return False
continue
if key.startswith("$"):
raise UnsupportedFilterError(f"operator {key!r} not supported by chroma backend")
actual = meta.get(key)
if isinstance(expected, dict):
for op, operand in expected.items():
if not _compare_metadata(actual, op, operand):
return False
elif actual != expected:
return False
return True
def _metadata_cell_value(sval, ival, fval, bval):
if sval is not None:
return sval
if ival is not None:
return ival
if fval is not None:
return fval
if bval is not None:
return bool(bval)
return None
def _segment_appears_healthy(seg_dir: str) -> bool:
"""Return True if a chromadb HNSW segment dir looks intact.
Sniff-tests the chromadb-written segment metadata file
(``index_metadata.pickle``) for its expected format bytes without
parsing it. ChromaDB writes that file after a successful HNSW flush;
a complete write starts with byte ``0x80`` and ends with byte
``0x2e`` (the protocol/terminator byte sequence chromadb serializes
with).
When metadata is missing, the segment is either *never-persisted*
(sub-threshold: fewer records than ``batch_size``, so chromadb never
triggered ``_persist()``) or *partially flushed* (persist started but
crashed). The two are distinguished by ``link_lists.bin``: chromadb
writes link data during persist, so an empty/absent ``link_lists.bin``
together with absent metadata means no persist was ever attempted.
Note: ``data_level0.bin`` is pre-allocated at index creation and its
size does not indicate actual record count.
Deliberately format-sniffs only; never deserializes. Deserialization
can execute arbitrary code, and the byte-sniff is sufficient to
distinguish a complete write from truncation, zero-fill, or
partial-flush corruption.
Assumes pickle protocol >= 2 (``0x80`` PROTO marker). Matches what
chromadb writes today; if a future chromadb version emits protocol
0/1 segments, this check would start returning False on healthy
files and quarantine_stale_hnsw would conservatively rename them
out of the way.
"""
meta_path = os.path.join(seg_dir, "index_metadata.pickle")
if not os.path.isfile(meta_path):
link_path = os.path.join(seg_dir, "link_lists.bin")
try:
link_has_data = os.path.isfile(link_path) and os.path.getsize(link_path) > 0
except OSError:
return False
# Both absent → sub-threshold, never persisted.
# link_lists written but metadata not → interrupted persist.
return not link_has_data
if not _hnsw_payload_appears_sane(seg_dir):
return False
return _hnsw_metadata_marker_intact(seg_dir)
def quarantine_stale_hnsw(palace_path: str, stale_seconds: float = 300.0) -> list[str]:
"""Rename HNSW segment dirs that look unsafe to open.
This catches two classes of HNSW corruption before ChromaDB opens the
native segment reader:
1. stale-by-mtime segments whose ``index_metadata.pickle`` fails the
existing format sniff-test;
2. structurally impossible HNSW payloads where ``link_lists.bin`` is much
larger than ``data_level0.bin``.
The second check is intentionally not gated by mtime. A segment with a
300x link/data ratio is unsafe regardless of whether its mtime is recent;
letting Chroma open it can SIGSEGV before Python fallback code runs.
The original directory is renamed, not deleted, so recovery remains
possible if the heuristic ever misfires.
"""
db_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return []
try:
sqlite_mtime = os.path.getmtime(db_path)
except OSError:
return []
moved: list[str] = []
try:
entries = os.listdir(palace_path)
except OSError:
return []
for name in entries:
if "-" not in name or name.startswith(".") or ".drift-" in name:
continue
seg_dir = os.path.join(palace_path, name)
if not os.path.isdir(seg_dir):
continue
hnsw_bin = os.path.join(seg_dir, "data_level0.bin")
if not os.path.isfile(hnsw_bin):
continue
try:
hnsw_mtime = os.path.getmtime(hnsw_bin)
except OSError:
continue
payload_ratio = _hnsw_link_to_data_ratio(seg_dir)
payload_corrupt = payload_ratio is not None and payload_ratio > _HNSW_LINK_TO_DATA_MAX_RATIO
if not payload_corrupt and sqlite_mtime - hnsw_mtime < stale_seconds:
continue
# Stage 2: integrity gate. Mtime drift alone is not corruption because
# Chroma flushes HNSW asynchronously. A healthy metadata file proves the
# ordinary stale-by-mtime case is just flush lag.
if not payload_corrupt and _segment_appears_healthy(seg_dir):
logger.info(
"HNSW mtime gap %.0fs on %s exceeds threshold but segment "
"metadata and payload size are intact — flush-lag, not "
"corruption. Leaving in place.",
sqlite_mtime - hnsw_mtime,
seg_dir,
)
continue
stamp = _dt.datetime.now().strftime("%Y%m%d-%H%M%S")
target = f"{seg_dir}.drift-{stamp}"
if payload_corrupt:
reason = (
f"link_lists.bin/data_level0.bin ratio {payload_ratio:.1f}x "
f"exceeds {_HNSW_LINK_TO_DATA_MAX_RATIO:.1f}x"
)
else:
reason = (
f"sqlite {sqlite_mtime - hnsw_mtime:.0f}s newer than HNSW "
"and integrity check failed"
)
try:
os.rename(seg_dir, target)
moved.append(target)
logger.warning(
"Quarantined corrupt HNSW segment %s (%s); renamed to %s",
seg_dir,
reason,
target,
)
except OSError:
logger.exception("Failed to quarantine corrupt HNSW segment %s", seg_dir)
return moved
def _vector_segment_id(palace_path: str, collection_name: str) -> Optional[str]:
"""Return the VECTOR segment UUID for ``collection_name`` or ``None``.
Reads ``chroma.sqlite3`` directly so we never have to load a segment
that may segfault on open (#1222 is exactly this case).
"""
db_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return None
try:
conn = sqlite3.connect(sqlite_read_uri(db_path), uri=True)
try:
row = conn.execute(
"""
SELECT s.id
FROM segments s
JOIN collections c ON s.collection = c.id
WHERE c.name = ? AND s.scope = 'VECTOR'
LIMIT 1
""",
(collection_name,),
).fetchone()
return row[0] if row else None
finally:
conn.close()
except sqlite3.Error:
return None
class _PersistentDataStub:
"""Minimal stand-in for chromadb's ``PersistentData`` during safe unpickling.
Accepts any constructor args so pickle's REDUCE opcode succeeds,
captures ``__setstate__`` into ``__dict__``. Only used by
:func:`_hnsw_element_count` — never persisted, never re-pickled.
"""
def __init__(self, *args, **kwargs):
# Some chromadb versions pickle PersistentData by passing init args
# positionally via REDUCE. We don't care about reconstructing the
# object faithfully — we only need the id_to_label dict — so swallow
# all positional args and re-expose the relevant attributes via
# __setstate__ or __dict__ population further down.
pass
def __setstate__(self, state):
if isinstance(state, dict):
self.__dict__.update(state)
elif isinstance(state, tuple) and len(state) == 2 and isinstance(state[1], dict):
# (slot_state, dict_state) two-tuple form — only the dict part has
# the named attributes we care about.
self.__dict__.update(state[1])
class _SafePersistentDataUnpickler:
"""Whitelist-only unpickler for ``index_metadata.pickle``.
Allows only ``PersistentData`` from chromadb's HNSW module; everything
else raises ``UnpicklingError``. Standard container types (dict, list,
tuple, str, int, float) are handled by the pickle machinery itself
and don't need allowlisting via ``find_class`` — only constructed
classes do.
This is the same trust model chromadb uses (it pickles its own files),
but with a tight class allowlist so a tampered file can't instantiate
arbitrary classes during deserialization.
"""
_ALLOWED = frozenset(
{
(
"chromadb.segment.impl.vector.local_persistent_hnsw",
"PersistentData",
),
}
)
@classmethod
def load(cls, path: str):
import pickle
class _Restricted(pickle.Unpickler):
def find_class(self, module: str, name: str):
if (module, name) in cls._ALLOWED:
return _PersistentDataStub
raise pickle.UnpicklingError(f"disallowed class: {module}.{name}")
with open(path, "rb") as f:
return _Restricted(f).load()
def _hnsw_element_count(palace_path: str, segment_id: str) -> Optional[int]:
"""Return the element count chromadb thinks the HNSW segment holds.
Reads ``index_metadata.pickle`` via a tight-allowlist unpickler and
counts ``id_to_label`` entries. This is the count chromadb consults
when sizing/loading the HNSW index on next open — distinct from
hnswlib's internal ``cur_element_count`` in the binary files. For
#1222's divergence check this is the number that matters, because it
is what gets compared against ``count() * resize_factor`` when
chromadb decides whether to resize HNSW on load.
Uses :class:`_SafePersistentDataUnpickler` rather than chromadb's own
``PersistentData.load_from_file`` so the probe works even when
``hnswlib`` is not installed (chromadb's persistent_hnsw module
imports hnswlib at module load — a probe that requires hnswlib would
refuse to run in environments where the segfault risk is moot
anyway). The allowlisted unpickler is also the safer default: the
pickle file is owned by the same user, but a tighter trust boundary
costs us nothing.
Returns ``None`` when the file is absent (fresh / never-flushed
segment) or the unpickle fails. Callers treat ``None`` as "unknown".
"""
pickle_path = os.path.join(palace_path, segment_id, "index_metadata.pickle")
if not os.path.isfile(pickle_path):
return None
try:
pd = _SafePersistentDataUnpickler.load(pickle_path)
# ChromaDB serializes PersistentData differently across versions:
# 1.5.x writes a plain dict via ``__reduce_ex__``; older versions
# pickled the class instance and rely on ``__setstate__`` to
# populate ``__dict__``. Handle both shapes.
if isinstance(pd, dict):
id_to_label = pd.get("id_to_label")
else:
id_to_label = getattr(pd, "id_to_label", None)
if isinstance(id_to_label, dict):
return len(id_to_label)
return None
except Exception:
logger.debug("_hnsw_element_count failed for %s", pickle_path, exc_info=True)
return None
# Divergence threshold: chromadb's HNSW flushes asynchronously, so HNSW
# typically lags sqlite by up to ``sync_threshold`` records under active
# write load — that's the *brute-force batch* that hasn't been compacted
# into HNSW yet, plus the un-persisted tail beyond the last sync. Two
# synchronization windows worth (2 × sync_threshold) is a safe steady-
# state ceiling; anything past that is real divergence, not flush-lag.
#
# The threshold floor scales with whatever ``hnsw:sync_threshold`` the
# collection was created with (read via :func:`_read_sync_threshold`).
# ``_HNSW_DIVERGENCE_FALLBACK_FLOOR`` is the floor used when we can't
# read the collection metadata (older palaces missing the row, sqlite
# unreadable). 2000 = 2 × chromadb's default sync_threshold of 1000.
#
# Why dynamic: legacy palaces may still carry ``sync_threshold = 50_000``
# (the pre-#1579 guard), so flush-lag can grow up to 50K on those palaces.
# New palaces use sync_threshold=2 (#1579) and flush almost immediately.
# A fixed 2000 floor would flag actively-written legacy palaces as
# DIVERGED the moment their queue exceeded 10% of sqlite_count, even
# though chromadb is behaving correctly. The floor must scale with the
# per-collection sync_threshold to distinguish real corruption (#1222 was
# 176 613 missing of 192 997, orders of magnitude past any reasonable
# sync_threshold) from expected steady-state lag.
_HNSW_DIVERGENCE_FALLBACK_FLOOR = 2000
_HNSW_DIVERGENCE_FRACTION = 0.10
_HNSW_PERSISTENT_DIVERGENCE_GRACE_SECONDS = 300.0
def _read_sync_threshold(palace_path: str, collection_name: str) -> int:
"""Return the ``hnsw:sync_threshold`` for a collection, or 1000 default.
The configured sync_threshold drives chromadb's HNSW flush cadence —
larger values mean fewer, bigger flushes (less index-bloat risk per
PR #1191) but also larger steady-state lag between
``index_metadata.pickle`` and the live sqlite count. The divergence
probe scales its tolerance to ``2 × sync_threshold`` so that lag is
not mistaken for corruption.
Falls back to 1000 (chromadb's own default) if the collection has no
explicit setting — matches what older mempalace palaces were created
with before PR #1191.
"""
db_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return 1000
try:
conn = sqlite3.connect(sqlite_read_uri(db_path), uri=True)
try:
cur = conn.cursor()
cur.execute(
"""
SELECT cm.int_value
FROM collection_metadata cm
JOIN collections c ON cm.collection_id = c.id
WHERE c.name = ? AND cm.key = 'hnsw:sync_threshold'
""",
(collection_name,),
)
row = cur.fetchone()
if row and row[0] is not None:
return int(row[0])
return 1000
finally:
conn.close()
except Exception:
logger.debug("_read_sync_threshold failed", exc_info=True)
return 1000
def _collection_has_sync_threshold_metadata(palace_path: str, collection_name: str) -> bool:
"""Return True when the collection explicitly stores hnsw:sync_threshold."""
db_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return False
try:
conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
try:
row = conn.execute(
"""
SELECT 1
FROM collection_metadata cm
JOIN collections c ON cm.collection_id = c.id
WHERE c.name = ?
AND cm.key = 'hnsw:sync_threshold'
LIMIT 1
""",
(collection_name,),
).fetchone()
return row is not None
finally:
conn.close()
except Exception:
logger.debug("_collection_has_sync_threshold_metadata failed", exc_info=True)
return False
def _hnsw_metadata_age_seconds(palace_path: str, segment_id: str) -> Optional[float]:
"""Return index_metadata.pickle age in seconds, or None when unreadable."""
pickle_path = os.path.join(palace_path, segment_id, "index_metadata.pickle")
try:
return max(0.0, time.time() - os.path.getmtime(pickle_path))
except OSError:
return None
def hnsw_capacity_status(palace_path: str, collection_name: str = "mempalace_drawers") -> dict:
"""Compare sqlite embedding count against HNSW element count.
The #1222 failure mode: ``max_elements`` froze at 16 384 while sqlite
accumulated 192 997 embeddings. Every subsequent tool call segfaulted
when chromadb tried to load the undersized HNSW. This probe runs
*before* anything touches the segment so we can warn (or fall back to
BM25) instead of crashing.
Returns a dict with:
* ``segment_id`` — VECTOR segment UUID, or ``None`` if no palace
* ``sqlite_count`` — embeddings present in chroma.sqlite3
* ``hnsw_count`` — elements chromadb's pickle knows about
* ``divergence`` — ``sqlite_count - hnsw_count`` when both known
* ``diverged`` — True when divergence exceeds the threshold
* ``status`` — ``"ok"`` | ``"diverged"`` | ``"unknown"``
* ``message`` — human-readable summary
Never raises — a probe that throws would defeat the point.
"""
out: dict[str, Any] = {
"segment_id": None,
"sqlite_count": None,
"hnsw_count": None,
"divergence": None,
"diverged": False,
"status": "unknown",
"message": "",
}
try:
seg_id = _vector_segment_id(palace_path, collection_name)
out["segment_id"] = seg_id
sqlite_count = _sqlite_embedding_count(palace_path, collection_name)
out["sqlite_count"] = sqlite_count
if seg_id is None or sqlite_count is None:
out["message"] = "palace state unreadable; skipping HNSW capacity check"
return out
hnsw_count = _hnsw_element_count(palace_path, seg_id)
out["hnsw_count"] = hnsw_count
sync_threshold = _read_sync_threshold(palace_path, collection_name)
has_explicit_sync_threshold = _collection_has_sync_threshold_metadata(
palace_path,
collection_name,
)
metadata_age_seconds = (
_hnsw_metadata_age_seconds(palace_path, seg_id) if hnsw_count is not None else None
)
out["hnsw_metadata_age_seconds"] = metadata_age_seconds
if hnsw_count is None:
# No pickle yet, so this probe cannot measure HNSW capacity.
# Chroma 1.5.x can have binary HNSW files without a flushed
# metadata pickle; absence of the pickle alone is not proof that
# vector search is unusable or dangerous. Keep the status unknown
# so MCP does not globally disable vectors on an inconclusive
# signal. Corrupt/invalid metadata, when present, is handled by
# quarantine_invalid_hnsw_metadata before Chroma opens.
out["message"] = (
"HNSW capacity unavailable: metadata has not been flushed; "
"leaving vector search enabled"
)
return out
divergence = sqlite_count - hnsw_count
out["divergence"] = divergence
# Newer palaces explicitly store mempalace's low sync threshold
# (currently 2), so a gap of dozens of rows is far beyond ordinary
# flush lag. Older palaces may lack the metadata row; keep the
# historical floor for fresh lag there, but do not let a stale pickle
# sit below the floor forever (#1816).
if has_explicit_sync_threshold:
threshold = max(0, 2 * sync_threshold)
else:
divergence_floor = max(_HNSW_DIVERGENCE_FALLBACK_FLOOR, 2 * sync_threshold)
threshold = max(
divergence_floor,
int(sqlite_count * _HNSW_DIVERGENCE_FRACTION),
)
out["threshold"] = threshold
stale_below_threshold = (
not has_explicit_sync_threshold
and divergence > 0
and metadata_age_seconds is not None
and metadata_age_seconds >= _HNSW_PERSISTENT_DIVERGENCE_GRACE_SECONDS
)
if divergence > threshold or stale_below_threshold:
out["status"] = "diverged"
out["diverged"] = True
pct = 100.0 * divergence / max(sqlite_count, 1)
if divergence > threshold:
reason = f"exceeds threshold {threshold:,}"
else:
age = metadata_age_seconds or 0.0
reason = f"persisted below the old flush-lag floor for {age:.0f}s"
out["message"] = (
f"HNSW index holds {hnsw_count:,} elements but sqlite has "
f"{sqlite_count:,} embeddings - {divergence:,} drawers "
f"({pct:.0f}%) are missing from the flushed HNSW index "
f"({reason}). Vector reads are disabled until "
"`mempalace repair` rebuilds it."
)
else:
out["status"] = "ok"
out["message"] = (
f"HNSW {hnsw_count:,} / sqlite {sqlite_count:,} (within flush-lag tolerance)"
)
if divergence < 0:
out["message"] += " (HNSW has extra flushed elements; treating as safe)"
except Exception:
logger.debug("hnsw_capacity_status failed", exc_info=True)
out["message"] = "HNSW capacity probe raised; skipping"
return out
def _sqlite_embedding_count(palace_path: str, collection_name: str) -> Optional[int]:
"""Count rows in chroma.sqlite3.embeddings for ``collection_name``.
Mirrors :func:`mempalace.repair.sqlite_drawer_count` but kept in this
module so the backend probe doesn't pull in the repair CLI module.
"""
db_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return None
try:
conn = sqlite3.connect(sqlite_read_uri(db_path), uri=True)
try:
row = conn.execute(
"""
SELECT COUNT(*)
FROM embeddings e
JOIN segments s ON e.segment_id = s.id
JOIN collections c ON s.collection = c.id
WHERE c.name = ?
""",
(collection_name,),
).fetchone()
return int(row[0]) if row and row[0] is not None else None
finally:
conn.close()
except sqlite3.Error:
return None
def _sqlite_wing_room_counts(
palace_path: str, collection_name: str
) -> Optional[tuple[int, dict[str, dict[str, int]]]]:
"""Tally drawers by wing/room straight from ``chroma.sqlite3``.
Returns ``(total, {wing: {room: count}})`` or ``None`` when the read
cannot be trusted — missing DB file, the collection has not been
bootstrapped, or any sqlite error (including a sustained writer lock).
``None`` signals the caller to fall back to the ChromaDB client path
(which also emits the right state-specific guidance for absent/empty
palaces).
The point of reading sqlite directly is to count drawers **without opening
the collection**, because opening it cold-loads the HNSW vector index. On
large palaces that load costs tens of seconds of CPU per call — a steep,
pointless tax for an inspection command that only needs metadata the
relational tables already hold (#1681). Wings/rooms live in plain
``embedding_metadata`` rows, joined to ``embeddings`` on the
``(id, key)`` primary key, so the tally is a bounded scan of the metadata
segment: sub-second warm, a few seconds cold on a multi-GB DB — versus the
~60s the vector-index load costs.
Sibling readers that count the same way: :func:`_sqlite_embedding_count`
(total only) and ``mcp_server._tool_status_via_sqlite`` (independent
wing/room histograms for the #1222 fallback). This one cross-tabulates
wing→room to match the ChromaDB ``status()`` output shape.
Notes:
- ``busy_timeout`` lets a transient checkpoint lock resolve instead of
instantly demoting to the slow HNSW path; a *sustained* lock still
raises and falls back (slow but correct).
- ``s.scope = 'METADATA'`` makes the single-segment join explicit so a
future ChromaDB that also stored per-vector-segment rows could not
silently double every count.
- ``COALESCE`` over ``string_value``/``int_value``/``float_value`` matches
the ChromaDB path, which surfaces a numeric wing/room natively rather
than dropping it to ``"?"``.
"""
db_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return None
try:
conn = sqlite3.connect(sqlite_read_uri(db_path), uri=True)
try:
# Wait out a transient writer/checkpoint lock rather than falling
# straight back to the expensive vector-index path (#1681).
conn.execute("PRAGMA busy_timeout = 3000")
# Distinguish "collection never bootstrapped" (-> None, so the
# caller can show the 'initialized but empty' message) from
# "collection exists with zero drawers" (-> a real 0 tally).
if (
conn.execute(
"SELECT 1 FROM collections WHERE name = ?", (collection_name,)
).fetchone()
is None
):
return None
rows = conn.execute(
"""
SELECT COALESCE(wm.string_value, CAST(wm.int_value AS TEXT),
CAST(wm.float_value AS TEXT), '?') AS wing,
COALESCE(rm.string_value, CAST(rm.int_value AS TEXT),
CAST(rm.float_value AS TEXT), '?') AS room,
COUNT(*) AS n
FROM embeddings e
JOIN segments s ON e.segment_id = s.id AND s.scope = 'METADATA'
JOIN collections c ON s.collection = c.id
LEFT JOIN embedding_metadata wm ON wm.id = e.id AND wm.key = 'wing'
LEFT JOIN embedding_metadata rm ON rm.id = e.id AND rm.key = 'room'
WHERE c.name = ?
GROUP BY wing, room
""",
(collection_name,),
).fetchall()
finally:
conn.close()
except sqlite3.Error:
return None
total = 0
wing_rooms: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
for wing, room, n in rows:
wing_rooms[wing][room] += int(n)
total += int(n)
return total, wing_rooms
def _pin_hnsw_threads(collection) -> None:
"""Best-effort retrofit: pin ``hnsw:num_threads=1`` on an existing collection.
Fresh collections set this via ``metadata=`` at creation. Legacy palaces
built before that change keep the default (parallel insert) and can hit
the HNSW race described in #974/#965. ChromaDB's
``collection.modify(configuration=...)`` lets us re-apply ``num_threads=1``
in memory at load time so every new process is protected.
Note: in chromadb 1.5.x the modified ``configuration_json["hnsw"]`` does
not persist to disk across ``PersistentClient`` reopens, so this must
run on every ``get_collection`` call, not just once.
"""
try:
from chromadb.api.collection_configuration import (
UpdateCollectionConfiguration,
UpdateHNSWConfiguration,
)
except ImportError:
logger.debug("_pin_hnsw_threads skipped: chromadb too old", exc_info=True)
return
try:
collection.modify(
configuration=UpdateCollectionConfiguration(hnsw=UpdateHNSWConfiguration(num_threads=1))
)
except Exception:
logger.debug("_pin_hnsw_threads modify failed", exc_info=True)
_BLOB_FIX_MARKER = ".blob_seq_ids_migrated"
_COLLECTION_TYPE_MARKER = ".collection_type_fixed"
def _valid_dimensionality(value: object) -> bool:
return isinstance(value, Integral) and not isinstance(value, bool) and int(value) > 0
def _persisted_metadata_value(obj: object, name: str) -> object:
if isinstance(obj, dict):
return obj.get(name)
return getattr(obj, name, None)
def _persisted_metadata_fields(obj: object) -> tuple[object, object]:
return _persisted_metadata_value(obj, "dimensionality"), _persisted_metadata_value(
obj, "id_to_label"
)
def _missing_dimensionality_appears_recoverable(
persisted: object, id_to_label: dict, seg_dir: str
) -> bool:
total = _persisted_metadata_value(persisted, "total_elements_added")
label_to_id = _persisted_metadata_value(persisted, "label_to_id")
data_path = os.path.join(seg_dir, "data_level0.bin")
link_path = os.path.join(seg_dir, "link_lists.bin")
if not isinstance(total, Integral) or isinstance(total, bool):
return False
if not isinstance(label_to_id, dict):
return False
try:
if not (
os.path.isfile(data_path)
and os.path.isfile(link_path)
and os.path.getsize(data_path) > _HNSW_MISSING_METADATA_DATA_FLOOR
):
return False
except OSError:
return False
if not _hnsw_payload_appears_sane(seg_dir):
return False
label_count = len(id_to_label)
# total_elements_added is monotonic across every add, while id_to_label and
# label_to_id hold only live elements, so a segment that has had deletions
# carries total_elements_added > label_count. Require >= (not ==), otherwise
# every post-deletion dim-None segment is wrongly quarantined (#1710); the
# label-map size and bijection checks still reject inconsistent label maps.
if int(total) < label_count or len(label_to_id) != label_count:
return False
try:
return all(label_to_id.get(label) == item_id for item_id, label in id_to_label.items())
except TypeError:
return False
def quarantine_invalid_hnsw_metadata(palace_path: str) -> list[str]:
"""Quarantine segment dirs whose ``index_metadata.pickle`` is unreadable or invalid.
Chroma's persisted HNSW metadata is untrusted disk state. If a segment has
labels but invalid or partial metadata, current Chroma versions can accept
the pickle and crash later in the Rust loader. We rename the entire segment
out of the way before ``PersistentClient`` opens so Chroma can rebuild
cleanly instead of touching known-bad metadata.
"""
try:
entries = os.listdir(palace_path)
except OSError:
return []
moved: list[str] = []
for name in entries:
if "-" not in name or name.startswith(".") or ".drift-" in name or ".corrupt-" in name:
continue
seg_dir = os.path.join(palace_path, name)
if not os.path.isdir(seg_dir):
continue
meta_path = os.path.join(seg_dir, "index_metadata.pickle")
if not os.path.isfile(meta_path):
continue
reason = None
try:
persisted = _SafePersistentDataUnpickler.load(meta_path)
except (EOFError, OSError):
logger.debug(
"Skipping invalid-HNSW quarantine for transient metadata read in %s",
meta_path,
exc_info=True,
)
continue
except pickle.UnpicklingError as exc:
if "truncated" in str(exc).lower() or "ran out of input" in str(exc).lower():
logger.debug(
"Skipping invalid-HNSW quarantine for transient metadata read in %s",
meta_path,
exc_info=True,
)
continue
reason = f"invalid index_metadata.pickle: {exc}"
except Exception as exc:
reason = f"invalid index_metadata.pickle: {exc}"
else:
if not isinstance(persisted, dict) and not (
hasattr(persisted, "dimensionality") or hasattr(persisted, "id_to_label")
):
reason = f"unrecognized index_metadata.pickle payload: {type(persisted).__name__}"
else:
dimensionality, id_to_label = _persisted_metadata_fields(persisted)
if id_to_label is not None and not isinstance(id_to_label, dict):
reason = f"invalid id_to_label type {type(id_to_label).__name__}"
else:
has_labels = bool(id_to_label)
if (
has_labels
and dimensionality is None
and not _missing_dimensionality_appears_recoverable(
persisted, id_to_label, seg_dir
)
):
reason = (
"labels present but dimensionality is missing or invalid "
f"({dimensionality!r})"
)
elif (
has_labels
and dimensionality is not None
and not _valid_dimensionality(dimensionality)
):
reason = (
"labels present but dimensionality is missing or invalid "
f"({dimensionality!r})"
)
elif dimensionality is not None and not _valid_dimensionality(dimensionality):
reason = f"invalid dimensionality {dimensionality!r}"
if reason is None:
continue
stamp = _dt.datetime.now().strftime("%Y%m%d-%H%M%S")
target = f"{seg_dir}.corrupt-{stamp}"
try:
os.rename(seg_dir, target)
moved.append(target)
logger.warning("Quarantined invalid HNSW metadata in %s: %s", seg_dir, reason)
except OSError:
logger.exception("Failed to quarantine invalid HNSW metadata in %s", seg_dir)
return moved
def _fix_blob_seq_ids(palace_path: str) -> None:
"""Fix ChromaDB 0.6.x -> 1.5.x migration bug: BLOB seq_ids -> INTEGER.
ChromaDB 0.6.x stored seq_id as big-endian 8-byte BLOBs. ChromaDB 1.5.x
expects INTEGER. The auto-migration doesn't convert existing rows, causing
the Rust compactor to crash with "mismatched types; Rust type u64 (as SQL
type INTEGER) is not compatible with SQL type BLOB".
Scoped to the ``embeddings`` table only. The ``max_seq_id`` table used
to be included in this loop, but chromadb 1.5.x writes its own BLOB
format there (``b'\\x11\\x11'`` + 6 ASCII digits). Misinterpreting that
format via ``int.from_bytes(..., 'big')`` yields a ~1.23e18 integer
that silently suppresses every subsequent write for the affected
segment (``embeddings_queue`` filters on ``seq_id > start``). chromadb
owns the ``max_seq_id`` column — we leave it alone. Palaces already
poisoned by the old behaviour can be repaired via
``mempalace repair --mode max-seq-id``.
Defense-in-depth: rows with the sysdb-10 ``b'\\x11\\x11'`` prefix in
``embeddings`` are skipped rather than converted. Real 0.6.x BLOBs are
pure big-endian u64 with no text prefix, so the prefix check is a
no-op for genuine legacy data.
Must run BEFORE PersistentClient is created (the compactor fires on init).
Opening a Python sqlite3 connection against a ChromaDB 1.5.x WAL-mode
database leaves state that segfaults the next PersistentClient call. After
the migration has run once successfully, a marker file is written so
subsequent opens skip the sqlite connection entirely. Already-migrated
palaces can touch the marker manually to opt into the fast path.
"""
db_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return
marker = os.path.join(palace_path, _BLOB_FIX_MARKER)
if os.path.isfile(marker):
return
try:
with contextlib.closing(sqlite3.connect(db_path)) as conn:
try:
rows = conn.execute(
"SELECT rowid, seq_id FROM embeddings WHERE typeof(seq_id) = 'blob'"
).fetchall()
except sqlite3.OperationalError:
return
safe_rows = [(rowid, blob) for rowid, blob in rows if not blob.startswith(b"\x11\x11")]
skipped = len(rows) - len(safe_rows)
if skipped:
logger.warning(
"Skipped %d sysdb-10-format BLOB seq_id(s) in embeddings (not converting)",
skipped,
)
if safe_rows:
updates = [
(int.from_bytes(blob, byteorder="big"), rowid) for rowid, blob in safe_rows
]
conn.executemany("UPDATE embeddings SET seq_id = ? WHERE rowid = ?", updates)
logger.info("Fixed %d BLOB seq_ids in embeddings", len(updates))
conn.commit()
except Exception:
logger.exception("Could not fix BLOB seq_ids in %s", db_path)
return
# Write marker whether or not rows needed migration — the palace is now
# confirmed to be in the INTEGER-seq_id state and future opens can skip the
# sqlite3.connect() entirely.
try:
Path(marker).touch()
except OSError:
logger.exception("Could not write migration marker %s", marker)
def _fix_missing_collection_type(palace_path: str) -> None:
"""Add ``_type`` to ``collections.config_json_str`` where absent.
chromadb <= 1.5.8 writes ``config_json_str = '{}'`` (empty JSON) when
creating collections. chromadb 1.5.9 switched from the permissive
``load_collection_configuration_from_json_str`` to
``CollectionConfigurationInternal.from_json`` which requires a ``_type``
key — its absence raises ``KeyError: '_type'`` on palace open.
This migration adds the missing marker so both old and new chromadb
versions can load the collection. The value
``"CollectionConfigurationInternal"`` matches what ``to_json()`` writes
for freshly-created collections.
Same lifecycle constraints as :func:`_fix_blob_seq_ids`: must run
BEFORE ``PersistentClient`` is created.
"""
db_path = os.path.join(palace_path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return
marker = os.path.join(palace_path, _COLLECTION_TYPE_MARKER)
if os.path.isfile(marker):
return
conn = sqlite3.connect(db_path)
try:
try:
rows = conn.execute("SELECT id, config_json_str FROM collections").fetchall()
except sqlite3.OperationalError:
return
updates = []
for coll_id, config_str in rows:
if not config_str:
config_str = "{}"
try:
config = json.loads(config_str)
except (json.JSONDecodeError, TypeError):
continue
if not isinstance(config, dict):
continue
if "_type" not in config:
config["_type"] = "CollectionConfigurationInternal"
updates.append((json.dumps(config), coll_id))
if updates:
conn.executemany(
"UPDATE collections SET config_json_str = ? WHERE id = ?",
updates,
)
conn.commit()
logger.info(
"Fixed %d collection(s) missing _type in config_json_str",
len(updates),
)
except Exception:
logger.exception("Could not fix collection config_json_str in %s", db_path)
return
finally:
conn.close()
try:
Path(marker).touch()
except OSError:
logger.exception("Could not write migration marker %s", marker)
# ---------------------------------------------------------------------------
# Collection adapter
# ---------------------------------------------------------------------------
def _as_list(v: Any) -> list:
"""Coerce possibly-None scalar-or-list into a list (defensive for chroma nulls)."""
if v is None:
return []
if isinstance(v, list):
return v
return [v]
def _close_client(client) -> None:
"""Call ``PersistentClient.close()`` if available, swallow otherwise.
chromadb 1.5.x exposes ``Client.close()`` to release rust-side SQLite
file locks; older versions relied on GC. Try/except keeps forward-compat.
"""
if client is None:
return
try:
client.close()
except Exception:
logger.debug("client.close() unavailable or failed", exc_info=True)
class ChromaCollection(BaseCollection):
"""Thin adapter translating ChromaDB dict returns into typed results.
When ``palace_path`` is set, all write methods (``add``, ``upsert``,
``update``, ``delete``) acquire ``mine_palace_lock(palace_path)`` for the
duration of the underlying chromadb call. This serializes MCP and other
direct-backend writers against ``mempalace mine`` and against each other,
closing the race between concurrent writers that triggers ChromaDB's
multi-threaded HNSW corruption (#974/#965).
The lock is the same primitive used by ``miner.mine()`` so re-entrant
acquisition from inside the mine pipeline (mine -> _mine_body ->
collection.upsert) is short-circuited by the per-thread guard inside
``mine_palace_lock`` — no self-deadlock.
``palace_path=None`` disables the wrapping, preserving the legacy
no-lock behaviour for callers that construct a ``ChromaCollection``
directly without going through ``ChromaBackend``.
"""
def __init__(self, collection, palace_path: Optional[str] = None):
self._collection = collection
self._palace_path = palace_path
@contextlib.contextmanager
def _write_lock(self):
"""Acquire ``mine_palace_lock`` for the configured palace, if any.
No-op (yields immediately) when ``self._palace_path`` is None.
"""
if self._palace_path is None:
yield
return
# Late import — palace.py imports ChromaBackend from this module.
from ..palace import mine_palace_lock
with mine_palace_lock(self._palace_path):
yield
# ------------------------------------------------------------------
# Writes
# ------------------------------------------------------------------
@staticmethod
def _sanitize_metadatas_for_chromadb(metadatas):
"""chromadb 1.5.x rejects None and empty-dict entries in the metadatas
list (ValueError: Expected metadata to be a non-empty dict, got 0
metadata attributes in add). Coerce any such entry to a sentinel so
the write succeeds. Operators can later locate coerced drawers via
``where={"_repaired_empty_meta": True}``.
This is the chokepoint catch-all: even if a caller's own sanitizer
misses a case (or skips for performance), reaching the chromadb
client always goes through here first.
"""
if metadatas is None:
return None
return [
m if (isinstance(m, dict) and len(m) > 0) else {"_repaired_empty_meta": True}
for m in metadatas
]
@staticmethod
def _sanitize_documents_for_chromadb(documents):
"""Strip lone UTF-16 surrogates and embedded NUL characters from every
document before it reaches the chromadb client.
A single lone surrogate (U+D800–U+DFFF) raises ``UnicodeEncodeError``
inside chromadb's encode path and aborts the *entire* add/upsert batch
with a ``-32000`` Internal Error, silently dropping every other row in
the same batch (#1235).
#1235 fixed this for the MCP write tools via ``sanitize_content``, but
the bulk ingest paths (miner, convo_miner, sweeper, diary_ingest) build
documents without routing through that helper and reach this backend
directly. Sanitising here makes the chokepoint catch-all complete: the
sibling :meth:`_sanitize_metadatas_for_chromadb` already guarantees this
for metadata one method over; documents get the same guarantee.
A document containing an embedded NUL (U+0000) — routine in mined
Bash tool output — is well-formed UTF-8, unlike a lone surrogate, but
can corrupt the FTS5 inverted index for the whole collection rather
than just failing to store that one row (a ChromaDB-side bug; tracked
upstream). Stripping it here is the same defense-in-depth already
applied to surrogates: sanitize input we don't control before it
reaches a datastore we don't control.
"""
if documents is None:
return None
from ..config import strip_lone_surrogates, strip_nul_bytes
def _sanitize(d):
return strip_nul_bytes(strip_lone_surrogates(d))
# chromadb accepts OneOrMany[Document]: a bare str is a single document,
# not an iterable of characters. Handle it explicitly so we don't split
# it into per-character documents — that would be exactly the kind of
# silent corruption this method exists to prevent.
if isinstance(documents, str):
return _sanitize(documents)
return [_sanitize(d) if isinstance(d, str) else d for d in documents]
def add(self, *, documents, ids, metadatas=None, embeddings=None):
kwargs: dict[str, Any] = {
"documents": self._sanitize_documents_for_chromadb(documents),
"ids": ids,
}
sanitized = self._sanitize_metadatas_for_chromadb(metadatas)
if sanitized is not None:
kwargs["metadatas"] = sanitized
if embeddings is not None:
kwargs["embeddings"] = embeddings
with self._write_lock():
self._collection.add(**kwargs)
def upsert(self, *, documents, ids, metadatas=None, embeddings=None):
kwargs: dict[str, Any] = {
"documents": self._sanitize_documents_for_chromadb(documents),
"ids": ids,
}
sanitized = self._sanitize_metadatas_for_chromadb(metadatas)
if sanitized is not None:
kwargs["metadatas"] = sanitized
if embeddings is not None:
kwargs["embeddings"] = embeddings
with self._write_lock():
self._collection.upsert(**kwargs)
def update(
self,
*,
ids,
documents=None,
metadatas=None,
embeddings=None,
):
if documents is None and metadatas is None and embeddings is None:
raise ValueError("update requires at least one of documents, metadatas, embeddings")
kwargs: dict[str, Any] = {"ids": ids}
if documents is not None:
kwargs["documents"] = self._sanitize_documents_for_chromadb(documents)
if metadatas is not None:
kwargs["metadatas"] = metadatas
if embeddings is not None:
kwargs["embeddings"] = embeddings
with self._write_lock():
self._collection.update(**kwargs)
# ------------------------------------------------------------------
# Reads
# ------------------------------------------------------------------
def query(
self,
*,
query_texts=None,
query_embeddings=None,
n_results=10,
where=None,
where_document=None,
include=None,
) -> QueryResult:
_validate_where(where)
_validate_where(where_document)
if (query_texts is None) == (query_embeddings is None):
raise ValueError("query requires exactly one of query_texts or query_embeddings")
chosen = query_texts if query_texts is not None else query_embeddings
if not chosen:
raise ValueError("query input must be a non-empty list")
spec = _IncludeSpec.resolve(include, default_distances=True)
chroma_include: list[str] = []
if spec.documents:
chroma_include.append("documents")
if spec.metadatas:
chroma_include.append("metadatas")
if spec.distances:
chroma_include.append("distances")
if spec.embeddings:
chroma_include.append("embeddings")
kwargs: dict[str, Any] = {
"n_results": n_results,
"include": chroma_include,
}
if query_texts is not None:
kwargs["query_texts"] = query_texts
if query_embeddings is not None:
kwargs["query_embeddings"] = query_embeddings
if where is not None:
kwargs["where"] = where
if where_document is not None:
kwargs["where_document"] = where_document
raw = self._collection.query(**kwargs)
num_queries = (
len(query_texts)
if query_texts is not None
else (len(query_embeddings) if query_embeddings is not None else 1)
)
ids = raw.get("ids") or []
if not ids:
return QueryResult.empty(
num_queries=num_queries,
embeddings_requested=spec.embeddings,
)
documents = raw.get("documents") or [[] for _ in ids]
metadatas = raw.get("metadatas") or [[] for _ in ids]
distances = raw.get("distances") or [[] for _ in ids]
embeddings_raw = raw.get("embeddings") if spec.embeddings else None
def _none_list_to_empty(outer):
return [(inner or []) for inner in outer]
return QueryResult(
ids=_none_list_to_empty(ids),
documents=_none_list_to_empty(documents),
metadatas=_none_list_to_empty(metadatas),
distances=_none_list_to_empty(distances),
embeddings=(
[list(inner) for inner in embeddings_raw]
if spec.embeddings and embeddings_raw is not None
else None
),
)
def get(
self,
*,
ids=None,
where=None,
where_document=None,
limit=None,
offset=None,
include=None,
) -> GetResult:
_validate_where(where)
_validate_where(where_document)
spec = _IncludeSpec.resolve(include, default_distances=False)
chroma_include: list[str] = []
if spec.documents:
chroma_include.append("documents")
if spec.metadatas:
chroma_include.append("metadatas")
if spec.embeddings:
chroma_include.append("embeddings")
kwargs: dict[str, Any] = {"include": chroma_include}
if ids is not None:
kwargs["ids"] = ids
if where is not None:
kwargs["where"] = where
if where_document is not None:
kwargs["where_document"] = where_document
if limit is not None:
kwargs["limit"] = limit
if offset is not None:
kwargs["offset"] = offset
raw = self._collection.get(**kwargs)
out_ids = list(raw.get("ids") or [])
out_docs = list(raw.get("documents") or []) if spec.documents else []
out_metas = list(raw.get("metadatas") or []) if spec.metadatas else []
out_embeds = raw.get("embeddings") if spec.embeddings else None
# Pad doc/meta lists to match ids so downstream zipping is safe.
if spec.documents and len(out_docs) < len(out_ids):
out_docs = out_docs + [""] * (len(out_ids) - len(out_docs))
if spec.metadatas and len(out_metas) < len(out_ids):
out_metas = out_metas + [{}] * (len(out_ids) - len(out_metas))
return GetResult(
ids=out_ids,
documents=out_docs,
metadatas=out_metas,
embeddings=[list(v) for v in out_embeds] if out_embeds is not None else None,
)
def delete(self, *, ids=None, where=None):
_validate_where(where)
kwargs: dict[str, Any] = {}
if ids is not None:
kwargs["ids"] = ids
if where is not None:
kwargs["where"] = where
with self._write_lock():
self._collection.delete(**kwargs)
def count(self):
return self._collection.count()
def lexical_search(
self,
*,
query: str,
n_results: int = 10,
where: Optional[dict] = None,
) -> LexicalResult:
"""Return lexical BM25 candidates for this collection.
This is the normal healthy-Chroma implementation behind the optional
backend capability. The HNSW-disabled fallback in ``searcher.py`` still
reads ``chroma.sqlite3`` directly and remains Chroma-only.
"""
_validate_where(where)
sqlite_hits = self._lexical_search_via_sqlite(query=query, n_results=n_results, where=where)
if sqlite_hits is not None:
return LexicalResult(hits=sqlite_hits)
# Directly-constructed ChromaCollection test doubles may not carry a
# palace path. Keep lexical_search usable in that shape, but normal
# MemPalace paths above use Chroma's FTS table instead of scanning every
# drawer through the Python client.
total = self.count()
docs: list[str] = []
metas: list[dict] = []
ids: list[str] = []
offset = 0
batch_size = 1000
while offset < total:
kwargs: dict[str, Any] = {
"include": ["documents", "metadatas"],
"limit": batch_size,
"offset": offset,
}
if where:
kwargs["where"] = where
batch = self.get(**kwargs)
if not batch.ids:
break
ids.extend(batch.ids)
docs.extend(doc or "" for doc in batch.documents)
metas.extend(meta or {} for meta in batch.metadatas)
offset += len(batch.ids)
scores = _bm25_scores(query, docs)
hits = [
LexicalHit(id=doc_id, document=doc, metadata=meta, score=float(score))
for doc_id, doc, meta, score in zip(ids, docs, metas, scores)
if score > 0
]
hits.sort(key=lambda hit: hit.score, reverse=True)
return LexicalResult(hits=hits[:n_results])
def _collection_name(self) -> Optional[str]:
name = getattr(self._collection, "name", None)
if callable(name):
try:
name = name()
except TypeError:
name = None
return str(name) if name else None
def _lexical_search_via_sqlite(
self,
*,
query: str,
n_results: int,
where: Optional[dict],
max_candidates: int = 500,
) -> Optional[list[LexicalHit]]:
if not self._palace_path:
return None
db_path = os.path.join(self._palace_path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return []
collection_name = self._collection_name()
if not collection_name:
return []
tokens = [t for t in _tokenize(query) if len(t) >= 3]
use_recency_fallback = not tokens
candidate_ids: list[int] = []
# Map internal embeddings.id (rowid, used to join embedding_metadata)
# to the public embeddings.embedding_id so returned LexicalHit.id values
# round-trip through get(ids=...). The two differ: id is the integer
# rowid, embedding_id is the user-facing drawer id.
public_ids: dict[int, str] = {}
try:
conn = sqlite3.connect(sqlite_read_uri(db_path), uri=True)
conn.row_factory = sqlite3.Row
except sqlite3.Error:
logger.debug("Chroma lexical sqlite open failed", exc_info=True)
return []
try:
if tokens:
fts_query = " OR ".join(tokens)
# If a metadata filter is present, do not cap before filtering:
# otherwise a common term can fill the window with wrong-scope
# rows and hide valid scoped hits later in the FTS result set.
limit_sql = "" if where else "LIMIT ?"
params = [fts_query, collection_name]
if not where:
params.append(max(max_candidates, n_results))
try:
rows = conn.execute(
f"""
SELECT e.id, e.embedding_id
FROM embedding_fulltext_search
JOIN embeddings e ON e.id = embedding_fulltext_search.rowid
JOIN segments s ON e.segment_id = s.id
JOIN collections c ON s.collection = c.id
WHERE embedding_fulltext_search MATCH ?
AND c.name = ?
{limit_sql}
""",
params,
).fetchall()
candidate_ids = [int(row[0]) for row in rows]
public_ids.update({int(row[0]): str(row[1]) for row in rows})
except sqlite3.Error:
logger.debug(
"Chroma lexical FTS query failed; using recency fallback", exc_info=True
)
use_recency_fallback = True
if not candidate_ids and use_recency_fallback:
order_expr = "e.created_at DESC"
try:
rows = conn.execute(
f"""
SELECT e.id, e.embedding_id
FROM embeddings e
JOIN segments s ON e.segment_id = s.id
JOIN collections c ON s.collection = c.id
WHERE c.name = ?
ORDER BY {order_expr}
LIMIT ?
""",
(collection_name, max(max_candidates, n_results)),
).fetchall()
except sqlite3.Error:
logger.debug(
"Chroma lexical recency fallback failed; ordering by id", exc_info=True
)
rows = conn.execute(
"""
SELECT e.id, e.embedding_id
FROM embeddings e
JOIN segments s ON e.segment_id = s.id
JOIN collections c ON s.collection = c.id
WHERE c.name = ?
ORDER BY e.id DESC
LIMIT ?
""",
(collection_name, max(max_candidates, n_results)),
).fetchall()
candidate_ids = [int(row[0]) for row in rows]
public_ids.update({int(row[0]): str(row[1]) for row in rows})
if not candidate_ids:
return []
meta_columns = {
row["name"]
for row in conn.execute("PRAGMA table_info(embedding_metadata)").fetchall()
}
value_columns = [
col
for col in ("string_value", "int_value", "float_value", "bool_value")
if col in meta_columns
]
if not value_columns:
return []
meta_rows = []
for start in range(0, len(candidate_ids), 900):
chunk_ids = candidate_ids[start : start + 900]
placeholders = ",".join("?" for _ in chunk_ids)
meta_rows.extend(
conn.execute(
f"""
SELECT id, key, {", ".join(value_columns)}
FROM embedding_metadata
WHERE id IN ({placeholders})
""",
chunk_ids,
).fetchall()
)
except sqlite3.Error:
logger.debug("Chroma lexical sqlite read failed", exc_info=True)
return []
finally:
conn.close()
drawers: dict[int, dict] = {}
for row in meta_rows:
emb_id = int(row["id"])
key = row["key"]
values = {col: row[col] if col in row.keys() else None for col in value_columns}
value = _metadata_cell_value(
values.get("string_value"),
values.get("int_value"),
values.get("float_value"),
values.get("bool_value"),
)
drawer = drawers.setdefault(emb_id, {"metadata": {}, "document": ""})
if key == "chroma:document":
drawer["document"] = str(value or "")
else:
drawer["metadata"][key] = value
ordered = []
for emb_id in candidate_ids:
drawer = drawers.get(emb_id)
if drawer is None:
continue
meta = drawer["metadata"]
if not _matches_where(meta, where):
continue
ordered.append((emb_id, drawer["document"], meta))
docs = [doc for _, doc, _ in ordered]
scores = _bm25_scores(query, docs)
hits = [
LexicalHit(
id=public_ids.get(emb_id, str(emb_id)),
document=doc,
metadata=meta,
score=float(score),
)
for (emb_id, doc, meta), score in zip(ordered, scores)
if score > 0
]
hits.sort(key=lambda hit: hit.score, reverse=True)
return hits[:n_results]
@property
def metadata(self) -> dict:
"""Pass-through to the underlying ChromaDB collection's metadata.
Used by the searcher to detect legacy palaces that were created
without ``hnsw:space=cosine`` and therefore silently use L2
distance, which breaks cosine-based similarity interpretation.
Returns ``{}`` when metadata is absent so callers can do a plain
``.get("hnsw:space")`` without None-checks.
"""
return self._collection.metadata or {}
@property
def distance_metric(self) -> str:
"""Report this collection's actual space from ``hnsw:space``.
MemPalace sets ``hnsw:space=cosine`` on every creation path, so a
healthy palace reports ``"cosine"``. When the key is absent, empty, or
an unrecognized value, the collection is genuinely using Chroma's HNSW
default — **L2** (Euclidean) — because cosine was never set on it. We
report ``"l2"`` in that case so core ranking maps the distances
correctly; reporting ``"cosine"`` here would reintroduce the
floor-every-result-to-zero misranking this property exists to fix.
"""
space = str(self.metadata.get("hnsw:space", "") or "").lower()
if space in ("cosine", "l2", "ip"):
return space
return "l2"
# ------------------------------------------------------------------
# Embedder identity (RFC 001)
#
# Stored in a small sidecar JSON in the palace dir rather than the Chroma
# collection metadata: ``collection.modify(metadata=...)`` replaces the
# whole dict and some Chroma versions reject re-passing the immutable
# ``hnsw:*`` construction keys, so mutating it on every open is fragile.
# The sidecar is keyed by collection name (a palace may hold several).
# This is complementary to Chroma's own embedding-function-name check —
# the core check runs at open time and yields the clean cross-backend
# error before Chroma's read-time rejection fires.
# ------------------------------------------------------------------
def _embedder_sidecar_path(self) -> Optional[str]:
if not self._palace_path:
return None
return os.path.join(self._palace_path, EMBEDDER_SIDECAR_FILENAME)
def get_stored_embedder_identity(self):
return read_embedder_sidecar(self._embedder_sidecar_path(), self._collection_name())
def set_embedder_identity(self, identity) -> None:
write_embedder_sidecar(self._embedder_sidecar_path(), self._collection_name(), identity)
# ---------------------------------------------------------------------------
# Backend
# ---------------------------------------------------------------------------
class ChromaBackend(BaseBackend):
"""MemPalace's default ChromaDB backend.
Maintains two caches:
* ``self._clients`` — ``palace_path -> PersistentClient`` for callers
using the ``PalaceRef`` / :meth:`get_collection` path.
* An inode+mtime freshness check absorbed from ``mcp_server._get_client``
(merged via #757) ensuring a palace rebuild on disk is detected on the
next :meth:`get_collection` call.
"""
name = "chroma"
capabilities = frozenset(
{
"supports_embeddings_in",
"supports_embeddings_passthrough",
"supports_embeddings_out",
"supports_metadata_filters",
"supports_contains_fast",
"supports_lexical_search",
"local_mode",
}
)
def __init__(self):
# palace_path -> PersistentClient
self._clients: dict[str, Any] = {}
# palace_path -> (inode, mtime) of chroma.sqlite3 at cache time.
self._freshness: dict[str, tuple[int, float]] = {}
self._closed = False
@staticmethod
def _resolve_embedding_function():
"""Return the EF for the user's ``embedding_device`` setting.
Both ``get_collection`` and ``get_or_create_collection`` must receive
the EF explicitly — ChromaDB 1.x does not persist it with the
collection, so a reader that omits the argument silently gets the
library default and its queries won't match the writer's vectors.
"""
try:
from ..embedding import get_embedding_function
return get_embedding_function()
except Exception:
logger.exception("Failed to build embedding function; using chromadb default")
return None
@staticmethod
def _explain_ef_mismatch(error: Exception, palace_path: str) -> Optional[str]:
"""If ``error`` looks like a ChromaDB EF-name mismatch, return a
user-friendly explanation. Otherwise return None so the caller can
re-raise unchanged.
Triggered when ``MEMPALACE_EMBEDDING_MODEL`` is switched on an
existing palace — ChromaDB persists the EF name on the collection
and refuses reads with a different one. The bare ValueError
ChromaDB raises doesn't mention rebuild-index or the env var, so
users hit it and don't know how to recover.
"""
msg = str(error)
if "Embedding function conflict" not in msg and "embedding function" not in msg.lower():
return None
try:
from ..config import MempalaceConfig
current_model = MempalaceConfig().embedding_model
except Exception:
current_model = "unknown"
rebuild_cmd = f"mempalace --palace {shlex.quote(palace_path)} repair rebuild-index"
return (
f"Embedding model mismatch reading palace at {palace_path!r}.\n"
f" Underlying ChromaDB error: {msg}\n"
f" Current MEMPALACE_EMBEDDING_MODEL={current_model!r}.\n"
f" The palace was built with a different embedding model. Either:\n"
f" (a) revert the model: unset MEMPALACE_EMBEDDING_MODEL (or set "
f"the previous value), or\n"
f" (b) re-embed in place: `{rebuild_cmd}` "
f"(writes new vectors with the current model)."
)
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
@staticmethod
def _db_stat(palace_path: str) -> tuple[int, float]:
"""Return ``(inode, mtime)`` of ``chroma.sqlite3`` or ``(0, 0.0)`` if absent."""
db_path = os.path.join(palace_path, "chroma.sqlite3")
try:
st = os.stat(db_path)
return (st.st_ino, st.st_mtime)
except OSError:
return (0, 0.0)
def _client(self, palace_path: str):
"""Return a cached ``PersistentClient``, rebuilding on inode/mtime change.
Handles the palace-rebuild case (repair/nuke/purge) by invalidating the
cache when ``chroma.sqlite3`` changes on disk. Mirrors the semantics of
``mcp_server._get_client`` (merged via #757):
* DB file missing while we hold a cached client → drop the cache so we
do not serve stale data after a rebuild that has not yet re-created
the DB.
* Transition 0 → nonzero stat (DB created after cache) counts as a
change, so the cached client is replaced with one that sees the DB.
* FAT/exFAT filesystems return inode 0; we never fire inode comparisons
when either side is 0 (safe fallback) but still honor mtime.
* Mtime change uses an epsilon (0.01 s) to tolerate FS timestamp
granularity without thrashing.
"""
if self._closed:
from .base import BackendClosedError # late import avoids cycles at module load
raise BackendClosedError("ChromaBackend has been closed")
cached = self._clients.get(palace_path)
cached_inode, cached_mtime = self._freshness.get(palace_path, (0, 0.0))
current_inode, current_mtime = self._db_stat(palace_path)
db_path = os.path.join(palace_path, "chroma.sqlite3")
# DB was present when cache was built but is now missing → invalidate.
if cached is not None and not os.path.isfile(db_path):
_close_client(self._clients.pop(palace_path, None))
self._freshness.pop(palace_path, None)
cached = None
cached_inode, cached_mtime = 0, 0.0
inode_changed = current_inode != 0 and cached_inode != 0 and current_inode != cached_inode
# Transition from no-stat (0.0) to a real stat counts as a change so we
# pick up a DB that was created after the cache was built.
mtime_appeared = cached_mtime == 0.0 and current_mtime != 0.0
mtime_changed = (
current_mtime != 0.0
and cached_mtime != 0.0
and abs(current_mtime - cached_mtime) > 0.01
)
if cached is None or inode_changed or mtime_changed or mtime_appeared:
# Drop the per-process quarantine gate so the HNSW pre-checks
# run again against the new disk state. An inode swap means a
# different physical DB (post-restore, fresh palace at the same
# path); an mtime/appearance change means an external in-place
# write (closet_llm, mine, compress) that may have drifted the
# HNSW index while this process was running.
if (
inode_changed
or mtime_changed
or (mtime_appeared and palace_path in self._freshness)
):
ChromaBackend._quarantined_paths.discard(palace_path)
ChromaBackend._prepare_palace_for_open(palace_path)
cached = chromadb.PersistentClient(path=palace_path)
self._clients[palace_path] = cached
# Re-stat after the client constructor runs: chromadb creates
# chroma.sqlite3 lazily, so the stat captured before the call
# may still be (0, 0.0) on first open.
self._freshness[palace_path] = self._db_stat(palace_path)
return cached
# ------------------------------------------------------------------
# Public static helpers (legacy; prefer :meth:`get_collection`)
# ------------------------------------------------------------------
# Per-process record of palaces that have already had the cold-start
# quarantine invoked at least once. The proactive HNSW checks are a
# *cold-start* protection -- they catch segments that arrive stale relative
# to ``chroma.sqlite3`` or invalid on disk (e.g. cross-machine replication,
# partial restore, crashed-mid-write). The gate is cleared whenever the
# palace changes on disk (inode swap, mtime bump, or file appearance), so
# external writes that drift HNSW segments are caught on the next open
# without requiring a full process restart.
#
# Thread-safety: this set is mutated without a lock. Two concurrent
# ``make_client()`` calls for the same palace can both pass the
# membership check and both invoke the cold-start quarantine. That's
# safe because the functions are idempotent (mtime checks + timestamped
# rename of distinct directories), so the worst-case race produces one
# redundant rename attempt that no-ops. Idempotency is the safety
# property; locking would add cost without correctness gain.
_quarantined_paths: set[str] = set()
@staticmethod
def _prepare_palace_for_open(palace_path: str) -> None:
"""Run the pre-open safety pass shared by :meth:`make_client` and
:meth:`_client`.
Four steps, all required before constructing a ``PersistentClient``:
1. ``_fix_missing_collection_type`` — adds the ``_type`` marker to
``collections.config_json_str`` that chromadb 1.5.9+ requires
but <= 1.5.8 never wrote (#1611).
2. ``_fix_blob_seq_ids`` — repairs the BLOB seq_id quirk that bites
certain chromadb migrations.
3. ``quarantine_invalid_hnsw_metadata`` — renames aside any HNSW
``index_metadata.pickle`` that fails to load, so chromadb opens
against an empty index instead of crashing on the unloadable
pickle (#1266 / PR #1285).
4. ``quarantine_stale_hnsw`` -- gated by :attr:`_quarantined_paths`
so it fires once per palace until the gate is re-armed by a
disk change. This is the SIGSEGV prevention path for stale
HNSW segments (see #1121, #1132, #1263); wiring it through
this helper means CLI mining, search, repair, and status all
benefit, not just the legacy ``make_client`` callers.
Idempotent: safe to call from any code path that is about to open or
re-open a palace. The ``_quarantined_paths`` gate prevents thrash on
hot paths (e.g. ``_client()`` is called on every backend operation).
"""
_fix_missing_collection_type(palace_path)
_fix_blob_seq_ids(palace_path)
if palace_path not in ChromaBackend._quarantined_paths:
quarantine_invalid_hnsw_metadata(palace_path)
quarantine_stale_hnsw(palace_path)
ChromaBackend._quarantined_paths.add(palace_path)
@staticmethod
def make_client(palace_path: str):
"""Create a fresh ``PersistentClient`` (runs pre-open safety pass first).
Deprecated-ish: exposed for legacy long-lived callers that manage their
own client cache. New code should obtain a collection through
:meth:`get_collection` which manages caching internally.
Quarantines HNSW segments on first open and after any detected
disk change. See :attr:`_quarantined_paths` for the gate logic.
"""
ChromaBackend._prepare_palace_for_open(palace_path)
return chromadb.PersistentClient(path=palace_path)
@staticmethod
def backend_version() -> str:
"""Return the installed chromadb package version string."""
return chromadb.__version__
# ------------------------------------------------------------------
# BaseBackend surface
# ------------------------------------------------------------------
def get_collection(
self,
*args,
**kwargs,
) -> ChromaCollection:
"""Obtain a collection for a palace.
Supports two calling conventions during the RFC 001 transition:
* New (preferred): ``get_collection(palace=PalaceRef, collection_name=...,
create=False, options=None)``.
* Legacy: ``get_collection(palace_path, collection_name, create=False)``
— still used by callers not yet migrated.
"""
palace_ref, collection_name, create, options = _normalize_get_collection_args(args, kwargs)
palace_path = palace_ref.local_path
if palace_path is None:
raise PalaceNotFoundError("ChromaBackend requires PalaceRef.local_path")
if not create and not os.path.isdir(palace_path):
raise PalaceNotFoundError(palace_path)
if create:
os.makedirs(palace_path, exist_ok=True)
try:
os.chmod(palace_path, 0o700)
except (OSError, NotImplementedError):
pass
client = self._client(palace_path)
hnsw_space = "cosine"
if options and isinstance(options, dict):
hnsw_space = options.get("hnsw_space", hnsw_space)
ef = self._resolve_embedding_function()
ef_kwargs = {"embedding_function": ef} if ef is not None else {}
if create:
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,
)
except ValueError as e:
explanation = self._explain_ef_mismatch(e, palace_path)
if explanation:
raise ValueError(explanation) from e
raise
else:
try:
collection = client.get_collection(collection_name, **ef_kwargs)
except _ChromaNotFoundError as e:
raise CollectionNotInitializedError(palace_path) from e
except ValueError as e:
explanation = self._explain_ef_mismatch(e, palace_path)
if explanation:
raise ValueError(explanation) from e
raise
_pin_hnsw_threads(collection)
return ChromaCollection(collection, palace_path=palace_path)
def close_palace(self, palace) -> None:
"""Drop cached handles for ``palace`` and release its SQLite file lock.
Accepts ``PalaceRef`` or legacy path str. chromadb's rust-side file
lock is held until ``PersistentClient.close()`` is called, so plain
dict eviction would leave the palace path unreopenable and
unremovable in the same process.
"""
path = palace.local_path if isinstance(palace, PalaceRef) else palace
if path is None:
return
_close_client(self._clients.pop(path, None))
self._freshness.pop(path, None)
def close(self) -> None:
for client in self._clients.values():
_close_client(client)
self._clients.clear()
self._freshness.clear()
self._closed = True
def health(self, palace: Optional[PalaceRef] = None) -> HealthStatus:
if self._closed:
return HealthStatus.unhealthy("backend closed")
return HealthStatus.healthy()
@classmethod
def detect(cls, path: str) -> bool:
"""Return True when ``path`` looks like a chroma palace.
Verifies the SQLite magic header rather than file presence alone.
Bare ``sqlite3.connect()`` against a missing path leaves a 0-byte
file behind (the SQLite header is written on the first statement,
not on connection), so file-presence alone treats those artifacts
as real chroma palaces and breaks multi-backend resolution. The
16-byte ``SQLite format 3\\x00`` magic prefix is written as soon
as chromadb's ``PersistentClient`` does any work, so this check
accepts every real chroma palace while rejecting empty / garbage
files. See #1893.
"""
db_path = os.path.join(path, "chroma.sqlite3")
if not os.path.isfile(db_path):
return False
try:
with open(db_path, "rb") as f:
return f.read(16) == b"SQLite format 3\x00"
except OSError:
return False
# ------------------------------------------------------------------
# Legacy (pre-RFC 001) surface — retained while callers migrate.
# ------------------------------------------------------------------
def get_or_create_collection(self, palace_path: str, collection_name: str) -> ChromaCollection:
"""Legacy shim for ``get_collection(..., create=True)`` by path string."""
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 the given HNSW space."""
ef = self._resolve_embedding_function()
ef_kwargs = {"embedding_function": ef} if ef is not None else {}
collection = self._client(palace_path).create_collection(
collection_name,
metadata={
"hnsw:space": hnsw_space,
"hnsw:num_threads": 1,
**_HNSW_BLOAT_GUARD,
},
**ef_kwargs,
)
return ChromaCollection(collection, palace_path=palace_path)
def _normalize_get_collection_args(args, kwargs):
"""Unify legacy positional ``(palace_path, collection_name, create)`` calls
with the new kwargs-only ``(palace=PalaceRef, collection_name=..., create=...)``.
Returns ``(PalaceRef, collection_name, create, options)``.
"""
# New-style: palace= kwarg with a PalaceRef (spec path).
if "palace" in kwargs:
palace_ref = kwargs.pop("palace")
if not isinstance(palace_ref, PalaceRef):
raise TypeError("palace= must be a PalaceRef instance")
collection_name = kwargs.pop("collection_name")
create = kwargs.pop("create", False)
options = kwargs.pop("options", None)
if kwargs:
raise TypeError(f"unexpected kwargs: {sorted(kwargs)}")
if args:
raise TypeError("positional args not allowed with palace= kwarg")
return palace_ref, collection_name, create, options
# Legacy: first positional is a path string.
if args:
palace_path = args[0]
rest = list(args[1:])
collection_name = kwargs.pop("collection_name", None) or (rest.pop(0) if rest else None)
if collection_name is None:
raise TypeError("collection_name is required")
create = kwargs.pop("create", False)
if rest:
create = rest.pop(0)
if rest:
raise TypeError(f"unexpected positional args: {rest!r}")
if kwargs:
raise TypeError(f"unexpected kwargs: {sorted(kwargs)}")
return (
PalaceRef(id=palace_path, local_path=palace_path),
collection_name,
bool(create),
None,
)
# Legacy kwargs-only (palace_path=..., collection_name=..., create=...)
if "palace_path" in kwargs:
palace_path = kwargs.pop("palace_path")
collection_name = kwargs.pop("collection_name")
create = kwargs.pop("create", False)
if kwargs:
raise TypeError(f"unexpected kwargs: {sorted(kwargs)}")
return (
PalaceRef(id=palace_path, local_path=palace_path),
collection_name,
bool(create),
None,
)
raise TypeError("get_collection requires palace= or a positional palace_path")