admin/basicmachines-co-basic-memory
1
0
Fork 0
mirror of https://github.com/basicmachines-co/basic-memory synced 2026-06-21 13:47:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(core): use (type, id) keys in vector search hydration to prevent id collisions

Browse Source 
Root cause: entity, observation, and relation rows in search_index carry
ids from independent auto-increment sequences, so rows of different types
routinely share the same numeric id (guaranteed in young databases).
_search_vector_only parsed each vector hit's chunk_key (e.g. 'entity:4:0')
but discarded the type, and _fetch_search_index_rows_by_ids keyed its
result dict by bare row.id with no type discrimination. Whichever row the
database returned last clobbered the other in the dict; the clobbered hit
then hydrated against the wrong row or found None and was silently
dropped from results. The FTS-filter branch already guarded this with
(id, type) tuples, but the primary vector lookup path and the hybrid
fusion maps missed the same treatment.

Fix: introduce a SearchIndexKey = tuple[str, int] alias and key every map
in the vector/hybrid retrieval path by (type, id) — the similarity and
chunk maps in _search_vector_only, the _fetch_search_index_rows_by_ids
result, the FTS-filter allowed keys, and the rows/fts/vec/fused score
maps in _search_hybrid. The SQL stays unchanged; bare ids are deduped
before the IN query and rows are discriminated by row.type when building
dict keys.

Tests: end-to-end SQLite regression test indexes an entity row and a
relation row sharing id 7, syncs vectors for both, and asserts vector
search returns both rows (and that the entity survives a
search_item_types filter); a hybrid fusion unit test asserts an entity
and relation sharing id 1 stay distinct with single-source scores. Both
fail without the fix. Existing mocked vector tests updated for tuple keys.

Fixes #982

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
phernandez
2026-06-11 23:54:41 -05:00
committed by Paul Hernandez
parent b3bdd5914f
commit 253e240d68
5 changed files with 181 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_memory/repository/search_repository_base.py
+53 -42
View File
@@ -40,6 +40,11 @@ BULLET_PATTERN = re.compile(r"^[\-\*]\s+")
OVERSIZED_ENTITY_VECTOR_SHARD_SIZE = 256
_SQLITE_MAX_PREPARE_WINDOW = 8
# Entity, observation, and relation rows in search_index carry ids from independent
# auto-increment sequences, so a bare id is ambiguous across row types. Every map in
# the vector/hybrid retrieval path must key rows by (type, id) to avoid collisions.
type SearchIndexKey = tuple[str, int]
@dataclass
class VectorSyncBatchResult:
@@ -1857,7 +1862,7 @@ class SearchRepositoryBase(ABC):
# ------------------------------------------------------------------
@staticmethod
def _parse_chunk_key(chunk_key: str) -> tuple[str, int]:
def _parse_chunk_key(chunk_key: str) -> SearchIndexKey:
"""Parse a chunk_key like 'observation:5:0' into (type, search_index_id)."""
parts = chunk_key.split(":")
return parts[0], int(parts[1])
@@ -1932,26 +1937,27 @@ class SearchRepositoryBase(ABC):
hydrate_start = time.perf_counter()
# Build per-search_index_row similarity scores from chunk-level results.
# Each chunk_key encodes the search_index row type and id.
# Each chunk_key encodes the search_index row type and id; keep both as the
# key because different row types can share the same numeric id (#982).
# Track the best similarity per row (for ranking) and all chunks (for context).
similarity_by_si_id: dict[int, float] = {}
chunks_by_si_id: dict[int, list[tuple[float, str]]] = {}
similarity_by_si_key: dict[SearchIndexKey, float] = {}
chunks_by_si_key: dict[SearchIndexKey, list[tuple[float, str]]] = {}
for row in vector_rows:
chunk_key = row.get("chunk_key", "")
distance = float(row["best_distance"])
similarity = self._distance_to_similarity(distance)
chunk_text = row.get("chunk_text", "")
try:
_, si_id = self._parse_chunk_key(chunk_key)
si_key = self._parse_chunk_key(chunk_key)
except (ValueError, IndexError):
# Fallback: group by entity_id for chunks without parseable keys
continue
current = similarity_by_si_id.get(si_id)
current = similarity_by_si_key.get(si_key)
if current is None or similarity > current:
similarity_by_si_id[si_id] = similarity
chunks_by_si_id.setdefault(si_id, []).append((similarity, chunk_text))
similarity_by_si_key[si_key] = similarity
chunks_by_si_key.setdefault(si_key, []).append((similarity, chunk_text))
if not similarity_by_si_id:
if not similarity_by_si_key:
hydrate_ms = (time.perf_counter() - hydrate_start) * 1000
_log_vector_summary()
return []
@@ -1962,16 +1968,17 @@ class SearchRepositoryBase(ABC):
min_similarity if min_similarity is not None else self._semantic_min_similarity
)
if effective_min_similarity > 0.0:
similarity_by_si_id = {
k: v for k, v in similarity_by_si_id.items() if v >= effective_min_similarity
similarity_by_si_key = {
k: v for k, v in similarity_by_si_key.items() if v >= effective_min_similarity
}
if not similarity_by_si_id:
if not similarity_by_si_key:
hydrate_ms = (time.perf_counter() - hydrate_start) * 1000
_log_vector_summary()
return []
# Fetch the actual search_index rows
si_ids = list(similarity_by_si_id.keys())
# Fetch the actual search_index rows. Colliding (type, id) keys share one
# bare id, so deduplicate while preserving first-seen order.
si_ids = list(dict.fromkeys(si_id for _, si_id in similarity_by_si_key))
search_index_rows = await self._fetch_search_index_rows_by_ids(si_ids)
# Apply optional filters if requested
@@ -2003,16 +2010,14 @@ class SearchRepositoryBase(ABC):
limit=VECTOR_FILTER_SCAN_LIMIT,
offset=0,
)
# Use (id, type) tuples to avoid collisions between different
# Use (type, id) tuples to avoid collisions between different
# search_index row types that share the same auto-increment id.
allowed_keys = {(row.id, row.type) for row in filtered_rows if row.id is not None}
search_index_rows = {
k: v for k, v in search_index_rows.items() if (v.id, v.type) in allowed_keys
}
allowed_keys = {(row.type, row.id) for row in filtered_rows if row.id is not None}
search_index_rows = {k: v for k, v in search_index_rows.items() if k in allowed_keys}
ranked_rows: list[SearchIndexRow] = []
for si_id, similarity in similarity_by_si_id.items():
row = search_index_rows.get(si_id)
for si_key, similarity in similarity_by_si_key.items():
row = search_index_rows.get(si_key)
if row is None:
continue
@@ -2022,7 +2027,7 @@ class SearchRepositoryBase(ABC):
if content_snippet and len(content_snippet) <= SMALL_NOTE_CONTENT_LIMIT:
matched_chunk_text = content_snippet
else:
si_chunks = chunks_by_si_id.get(si_id, [])
si_chunks = chunks_by_si_key.get(si_key, [])
si_chunks.sort(key=lambda c: c[0], reverse=True)
top_texts = [text for _, text in si_chunks[:TOP_CHUNKS_PER_RESULT]]
matched_chunk_text = "\n---\n".join(top_texts) if top_texts else None
@@ -2088,8 +2093,12 @@ class SearchRepositoryBase(ABC):
async def _fetch_search_index_rows_by_ids(
self, row_ids: list[int]
) -> dict[int, SearchIndexRow]:
"""Fetch search_index rows by their primary key (id), any type."""
) -> dict[SearchIndexKey, SearchIndexRow]:
"""Fetch search_index rows by id, keyed by (type, id) to disambiguate types.
A bare id can match one row per type (independent id sequences), so the
result must carry every matching row rather than letting one clobber another.
"""
if not row_ids:
return {}
placeholders = ",".join(f":id_{idx}" for idx in range(len(row_ids)))
@@ -2106,11 +2115,11 @@ class SearchRepositoryBase(ABC):
WHERE project_id = :project_id
AND id IN ({placeholders})
"""
result: dict[int, SearchIndexRow] = {}
result: dict[SearchIndexKey, SearchIndexRow] = {}
async with db.scoped_session(self.session_maker) as session:
row_result = await session.execute(text(sql), params)
for row in row_result.fetchall():
result[row.id] = SearchIndexRow(
result[(row.type, row.id)] = SearchIndexRow(
project_id=self.project_id,
id=row.id,
title=row.title,
@@ -2156,7 +2165,7 @@ class SearchRepositoryBase(ABC):
) -> List[SearchIndexRow]:
"""Fuse FTS and vector results using score-based fusion.
Uses search_index row id as the fusion key. The formula
Uses the search_index (type, id) pair as the fusion key. The formula
``max(vec, fts) + FUSION_BONUS * min(vec, fts)`` preserves
the dominant signal and rewards dual-source agreement.
"""
@@ -2199,17 +2208,19 @@ class SearchRepositoryBase(ABC):
vector_ms = (time.perf_counter() - vector_start) * 1000
fusion_start = time.perf_counter()
# --- Score-based fusion keyed on search_index row id ---
# --- Score-based fusion keyed on (type, id) ---
# A bare row id collides across row types (independent id sequences), so
# fusion must key on (type, id) or distinct rows would merge (#982).
# FTS scores are normalized to [0, 1] (BM25 is unbounded).
# Vector scores are used raw — already calibrated [0, 1] by _distance_to_similarity().
rows_by_id: dict[int, SearchIndexRow] = {}
rows_by_key: dict[SearchIndexKey, SearchIndexRow] = {}
# Normalize FTS scores to [0, 1] — handles both SQLite (negative bm25)
# and Postgres (positive ts_rank) by using absolute values
fts_abs = [abs(row.score or 0.0) for row in fts_results]
fts_max = max(fts_abs) if fts_abs else 1.0
fts_scores: dict[int, float] = {}
fts_scores: dict[SearchIndexKey, float] = {}
for row in fts_results:
if row.id is None:
continue
@@ -2217,32 +2228,32 @@ class SearchRepositoryBase(ABC):
# Gate: FTS scores below threshold contribute zero
if norm < FTS_GATE_THRESHOLD:
norm = 0.0
fts_scores[row.id] = norm
rows_by_id[row.id] = row
fts_scores[(row.type, row.id)] = norm
rows_by_key[(row.type, row.id)] = row
vec_scores: dict[int, float] = {}
vec_scores: dict[SearchIndexKey, float] = {}
for row in vector_results:
if row.id is None:
continue
# Trigger: no re-normalization by vec_max
# Why: vector similarity is already calibrated [0, 1]; re-normalizing
# inflates weak matches when the entire result set is mediocre
vec_scores[row.id] = row.score or 0.0
rows_by_id[row.id] = row
vec_scores[(row.type, row.id)] = row.score or 0.0
rows_by_key[(row.type, row.id)] = row
# Fuse: max(v, f) + FUSION_BONUS * min(v, f)
# Preserves the dominant signal; bonus rewards dual-source agreement.
# Output range: [0, 1.3] for dual-source, [0, 1.0] for single-source.
fused_scores: dict[int, float] = {}
for row_id in fts_scores.keys() | vec_scores.keys():
v = vec_scores.get(row_id, 0.0)
f = fts_scores.get(row_id, 0.0)
fused_scores[row_id] = max(v, f) + FUSION_BONUS * min(v, f)
fused_scores: dict[SearchIndexKey, float] = {}
for row_key in fts_scores.keys() | vec_scores.keys():
v = vec_scores.get(row_key, 0.0)
f = fts_scores.get(row_key, 0.0)
fused_scores[row_key] = max(v, f) + FUSION_BONUS * min(v, f)
ranked = sorted(fused_scores.items(), key=lambda item: item[1], reverse=True)
output: list[SearchIndexRow] = []
for row_id, fused_score in ranked[offset : offset + limit]:
row = rows_by_id[row_id]
for row_key, fused_score in ranked[offset : offset + limit]:
row = rows_by_key[row_key]
# Trigger: FTS-only results have no matched_chunk_text from vector search.
# Why: without chunk text, API falls back to truncated content, losing answer text.
# Outcome: FTS-only results get full content_snippet as matched_chunk.