mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
phernandez
253e240d68
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>
469 lines
17 KiB
Python
469 lines
17 KiB
Python
"""Tests for semantic_min_similarity threshold filtering in vector search."""
|
|
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from typing import Any, Optional, cast
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
|
|
from basic_memory.repository.embedding_provider import EmbeddingProvider
|
|
from basic_memory.repository.search_index_row import SearchIndexRow
|
|
from basic_memory.repository.search_repository_base import (
|
|
SMALL_NOTE_CONTENT_LIMIT,
|
|
TOP_CHUNKS_PER_RESULT,
|
|
SearchRepositoryBase,
|
|
)
|
|
from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode
|
|
|
|
|
|
@dataclass
|
|
class FakeRow:
|
|
"""Minimal stand-in for SearchIndexRow in threshold tests."""
|
|
|
|
id: int
|
|
type: str = "entity"
|
|
score: float = 0.0
|
|
matched_chunk_text: str | None = None
|
|
content_snippet: str | None = None
|
|
|
|
|
|
class ConcreteSearchRepo(SearchRepositoryBase):
|
|
"""Minimal concrete subclass for testing base class threshold logic."""
|
|
|
|
def __init__(self):
|
|
# Skip super().__init__ — we only need the attributes under test
|
|
self._semantic_enabled = True
|
|
self._semantic_vector_k = 100
|
|
self._semantic_min_similarity = 0.0
|
|
self._embedding_provider = None
|
|
self._vector_dimensions = 384
|
|
self._vector_tables_initialized = True
|
|
self.session_maker = None
|
|
self.project_id = 1
|
|
|
|
# --- Abstract method stubs (not exercised by these tests) ---
|
|
|
|
async def init_search_index(self):
|
|
pass # pragma: no cover
|
|
|
|
def _prepare_search_term(self, term, is_prefix=True):
|
|
return term # pragma: no cover
|
|
|
|
async def search(
|
|
self,
|
|
search_text: Optional[str] = None,
|
|
permalink: Optional[str] = None,
|
|
permalink_match: Optional[str] = None,
|
|
title: Optional[str] = None,
|
|
note_types: Optional[list[str]] = None,
|
|
after_date: Optional[datetime] = None,
|
|
search_item_types: Optional[list[SearchItemType]] = None,
|
|
categories: Optional[list[str]] = None,
|
|
metadata_filters: Optional[dict[str, Any]] = None,
|
|
retrieval_mode: SearchRetrievalMode = SearchRetrievalMode.FTS,
|
|
min_similarity: Optional[float] = None,
|
|
limit: int = 10,
|
|
offset: int = 0,
|
|
) -> list[SearchIndexRow]:
|
|
return [] # pragma: no cover
|
|
|
|
async def _ensure_vector_tables(self):
|
|
pass # pragma: no cover
|
|
|
|
async def _run_vector_query(self, session, query_embedding, candidate_limit):
|
|
return [] # pragma: no cover
|
|
|
|
async def _write_embeddings(self, session, jobs, embeddings):
|
|
pass # pragma: no cover
|
|
|
|
async def _delete_entity_chunks(self, session, entity_id):
|
|
pass # pragma: no cover
|
|
|
|
async def _delete_stale_chunks(self, session, stale_ids, entity_id):
|
|
pass # pragma: no cover
|
|
|
|
async def _update_timestamp_sql(self):
|
|
return "CURRENT_TIMESTAMP" # pragma: no cover
|
|
|
|
def _distance_to_similarity(self, distance: float) -> float:
|
|
return 1.0 / (1.0 + max(distance, 0.0))
|
|
|
|
|
|
def _make_vector_rows(scores: list[float]) -> list[dict]:
|
|
"""Build fake vector query rows with controlled distances.
|
|
|
|
Distance = (1/score) - 1 inverts the similarity formula:
|
|
similarity = 1 / (1 + distance)
|
|
"""
|
|
rows = []
|
|
for i, score in enumerate(scores):
|
|
distance = (1.0 / score) - 1.0
|
|
rows.append(
|
|
{
|
|
"chunk_key": f"entity:{i}:0",
|
|
"best_distance": distance,
|
|
"chunk_text": f"chunk text for entity:{i}:0",
|
|
}
|
|
)
|
|
return rows
|
|
|
|
|
|
def _fake_embedding_provider(mock_embed: AsyncMock) -> EmbeddingProvider:
|
|
return cast(
|
|
EmbeddingProvider,
|
|
type("EP", (), {"embed_query": mock_embed, "dimensions": 384})(),
|
|
)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def fake_scoped_session(session_maker):
|
|
"""Fake scoped_session that yields a mock session object."""
|
|
yield AsyncMock()
|
|
|
|
|
|
COMMON_SEARCH_KWARGS: dict[str, Any] = dict(
|
|
search_text="test",
|
|
permalink=None,
|
|
permalink_match=None,
|
|
title=None,
|
|
note_types=None,
|
|
after_date=None,
|
|
search_item_types=None,
|
|
categories=None,
|
|
metadata_filters=None,
|
|
limit=10,
|
|
offset=0,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_threshold_zero_returns_all():
|
|
"""With threshold=0.0 (default), all results pass through."""
|
|
repo = ConcreteSearchRepo()
|
|
repo._semantic_min_similarity = 0.0
|
|
|
|
fake_rows = _make_vector_rows([0.9, 0.5, 0.3])
|
|
|
|
mock_embed = AsyncMock(return_value=[0.0] * 384)
|
|
repo._embedding_provider = _fake_embedding_provider(mock_embed)
|
|
|
|
with (
|
|
patch(
|
|
"basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session
|
|
),
|
|
patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock),
|
|
patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock),
|
|
patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows),
|
|
patch.object(
|
|
repo,
|
|
"_fetch_search_index_rows_by_ids",
|
|
new_callable=AsyncMock,
|
|
return_value={("entity", i): FakeRow(id=i) for i in range(3)},
|
|
),
|
|
):
|
|
results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS)
|
|
|
|
assert len(results) == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_threshold_filters_low_scores():
|
|
"""Results below the threshold are excluded."""
|
|
repo = ConcreteSearchRepo()
|
|
repo._semantic_min_similarity = 0.6
|
|
|
|
# Scores: 0.9 (pass), 0.5 (fail), 0.3 (fail)
|
|
fake_rows = _make_vector_rows([0.9, 0.5, 0.3])
|
|
|
|
mock_embed = AsyncMock(return_value=[0.0] * 384)
|
|
repo._embedding_provider = _fake_embedding_provider(mock_embed)
|
|
|
|
with (
|
|
patch(
|
|
"basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session
|
|
),
|
|
patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock),
|
|
patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock),
|
|
patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows),
|
|
patch.object(
|
|
repo,
|
|
"_fetch_search_index_rows_by_ids",
|
|
new_callable=AsyncMock,
|
|
# Only entity_0 (score=0.9) passes the threshold; the fetch only gets id 0
|
|
return_value={("entity", 0): FakeRow(id=0)},
|
|
),
|
|
):
|
|
results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS)
|
|
|
|
# Only the 0.9 result passes the 0.6 threshold
|
|
assert len(results) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_threshold_returns_empty_when_all_below():
|
|
"""All results below threshold → empty list, no DB fetch."""
|
|
repo = ConcreteSearchRepo()
|
|
repo._semantic_min_similarity = 0.8
|
|
|
|
# All scores below 0.8
|
|
fake_rows = _make_vector_rows([0.5, 0.4, 0.3])
|
|
|
|
mock_embed = AsyncMock(return_value=[0.0] * 384)
|
|
repo._embedding_provider = _fake_embedding_provider(mock_embed)
|
|
|
|
mock_fetch = AsyncMock()
|
|
|
|
with (
|
|
patch(
|
|
"basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session
|
|
),
|
|
patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock),
|
|
patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock),
|
|
patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows),
|
|
patch.object(repo, "_fetch_search_index_rows_by_ids", mock_fetch),
|
|
):
|
|
results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS)
|
|
|
|
assert results == []
|
|
# Should short-circuit before fetching search_index rows
|
|
mock_fetch.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_per_query_min_similarity_overrides_instance_default():
|
|
"""Per-query min_similarity takes precedence over instance-level default."""
|
|
repo = ConcreteSearchRepo()
|
|
# Instance default would filter out 0.5 and 0.3
|
|
repo._semantic_min_similarity = 0.6
|
|
|
|
# Scores: 0.9, 0.5, 0.3
|
|
fake_rows = _make_vector_rows([0.9, 0.5, 0.3])
|
|
|
|
mock_embed = AsyncMock(return_value=[0.0] * 384)
|
|
repo._embedding_provider = _fake_embedding_provider(mock_embed)
|
|
|
|
with (
|
|
patch(
|
|
"basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session
|
|
),
|
|
patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock),
|
|
patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock),
|
|
patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows),
|
|
patch.object(
|
|
repo,
|
|
"_fetch_search_index_rows_by_ids",
|
|
new_callable=AsyncMock,
|
|
return_value={("entity", i): FakeRow(id=i) for i in range(3)},
|
|
),
|
|
):
|
|
# Override to 0.0 → all results pass through despite instance default of 0.6
|
|
results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS, min_similarity=0.0)
|
|
|
|
assert len(results) == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_per_query_min_similarity_tightens_threshold():
|
|
"""Per-query min_similarity=0.8 filters more aggressively than instance default."""
|
|
repo = ConcreteSearchRepo()
|
|
# Instance default is permissive
|
|
repo._semantic_min_similarity = 0.0
|
|
|
|
# Scores: 0.9, 0.5, 0.3
|
|
fake_rows = _make_vector_rows([0.9, 0.5, 0.3])
|
|
|
|
mock_embed = AsyncMock(return_value=[0.0] * 384)
|
|
repo._embedding_provider = _fake_embedding_provider(mock_embed)
|
|
|
|
with (
|
|
patch(
|
|
"basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session
|
|
),
|
|
patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock),
|
|
patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock),
|
|
patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows),
|
|
patch.object(
|
|
repo,
|
|
"_fetch_search_index_rows_by_ids",
|
|
new_callable=AsyncMock,
|
|
# Only id=0 (score=0.9) will be fetched after filtering
|
|
return_value={("entity", 0): FakeRow(id=0)},
|
|
),
|
|
):
|
|
# Override to 0.8 → only score=0.9 passes
|
|
results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS, min_similarity=0.8)
|
|
|
|
assert len(results) == 1
|
|
assert results[0].id == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_matched_chunk_text_populated_on_vector_results():
|
|
"""Vector search results carry the matched chunk text from the best-matching chunk."""
|
|
repo = ConcreteSearchRepo()
|
|
repo._semantic_min_similarity = 0.0
|
|
|
|
fake_rows = _make_vector_rows([0.9, 0.7])
|
|
|
|
mock_embed = AsyncMock(return_value=[0.0] * 384)
|
|
repo._embedding_provider = _fake_embedding_provider(mock_embed)
|
|
|
|
with (
|
|
patch(
|
|
"basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session
|
|
),
|
|
patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock),
|
|
patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock),
|
|
patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows),
|
|
patch.object(
|
|
repo,
|
|
"_fetch_search_index_rows_by_ids",
|
|
new_callable=AsyncMock,
|
|
return_value={("entity", i): FakeRow(id=i) for i in range(2)},
|
|
),
|
|
):
|
|
results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS)
|
|
|
|
assert len(results) == 2
|
|
# Results are sorted by score descending, so id=0 (0.9) first, id=1 (0.7) second
|
|
# Each entity has only 1 chunk and no content_snippet, so chunk text is used directly
|
|
assert results[0].matched_chunk_text == "chunk text for entity:0:0"
|
|
assert results[1].matched_chunk_text == "chunk text for entity:1:0"
|
|
|
|
|
|
def _make_multi_chunk_vector_rows(si_id: int, scores: list[float]) -> list[dict]:
|
|
"""Build multiple fake vector chunks for a single search_index row.
|
|
|
|
Each chunk gets a unique chunk_index within the same si_id.
|
|
Distance = (1/score) - 1 inverts the similarity formula.
|
|
"""
|
|
rows = []
|
|
for chunk_idx, score in enumerate(scores):
|
|
distance = (1.0 / score) - 1.0
|
|
rows.append(
|
|
{
|
|
"chunk_key": f"entity:{si_id}:{chunk_idx}",
|
|
"best_distance": distance,
|
|
"chunk_text": f"chunk-{chunk_idx} (sim={score})",
|
|
}
|
|
)
|
|
return rows
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_top_n_chunks_joined_in_matched_chunk_text():
|
|
"""Large note with 7 chunks: top 5 by similarity are joined with separator."""
|
|
repo = ConcreteSearchRepo()
|
|
repo._semantic_min_similarity = 0.0
|
|
|
|
# 7 chunks for entity 0, with varying similarities
|
|
chunk_scores = [0.6, 0.9, 0.4, 0.8, 0.75, 0.3, 0.85]
|
|
fake_rows = _make_multi_chunk_vector_rows(si_id=0, scores=chunk_scores)
|
|
|
|
mock_embed = AsyncMock(return_value=[0.0] * 384)
|
|
repo._embedding_provider = _fake_embedding_provider(mock_embed)
|
|
|
|
# content_snippet exceeds SMALL_NOTE_CONTENT_LIMIT → top-N chunks path
|
|
large_content = "x" * (SMALL_NOTE_CONTENT_LIMIT + 1)
|
|
|
|
with (
|
|
patch(
|
|
"basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session
|
|
),
|
|
patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock),
|
|
patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock),
|
|
patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows),
|
|
patch.object(
|
|
repo,
|
|
"_fetch_search_index_rows_by_ids",
|
|
new_callable=AsyncMock,
|
|
return_value={("entity", 0): FakeRow(id=0, content_snippet=large_content)},
|
|
),
|
|
):
|
|
results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS)
|
|
|
|
assert len(results) == 1
|
|
text = results[0].matched_chunk_text
|
|
assert text is not None
|
|
|
|
# Top 5 chunks by similarity: 0.9, 0.85, 0.8, 0.75, 0.6 (0.4 and 0.3 excluded)
|
|
parts = text.split("\n---\n")
|
|
assert len(parts) == TOP_CHUNKS_PER_RESULT
|
|
assert parts[0] == "chunk-1 (sim=0.9)"
|
|
assert parts[1] == "chunk-6 (sim=0.85)"
|
|
assert parts[2] == "chunk-3 (sim=0.8)"
|
|
assert parts[3] == "chunk-4 (sim=0.75)"
|
|
assert parts[4] == "chunk-0 (sim=0.6)"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_small_note_returns_full_content_as_matched_chunk():
|
|
"""Small note (content_snippet under limit) returns full content instead of chunks."""
|
|
repo = ConcreteSearchRepo()
|
|
repo._semantic_min_similarity = 0.0
|
|
|
|
fake_rows = _make_vector_rows([0.9])
|
|
|
|
mock_embed = AsyncMock(return_value=[0.0] * 384)
|
|
repo._embedding_provider = _fake_embedding_provider(mock_embed)
|
|
|
|
small_content = "This is a short note with all the important details."
|
|
assert len(small_content) <= SMALL_NOTE_CONTENT_LIMIT
|
|
|
|
with (
|
|
patch(
|
|
"basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session
|
|
),
|
|
patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock),
|
|
patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock),
|
|
patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows),
|
|
patch.object(
|
|
repo,
|
|
"_fetch_search_index_rows_by_ids",
|
|
new_callable=AsyncMock,
|
|
return_value={("entity", 0): FakeRow(id=0, content_snippet=small_content)},
|
|
),
|
|
):
|
|
results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS)
|
|
|
|
assert len(results) == 1
|
|
# Full content returned instead of the chunk text
|
|
assert results[0].matched_chunk_text == small_content
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_large_note_returns_chunks_not_full_content():
|
|
"""Large note (content_snippet over limit) returns top-N chunks, not full content."""
|
|
repo = ConcreteSearchRepo()
|
|
repo._semantic_min_similarity = 0.0
|
|
|
|
fake_rows = _make_vector_rows([0.9])
|
|
|
|
mock_embed = AsyncMock(return_value=[0.0] * 384)
|
|
repo._embedding_provider = _fake_embedding_provider(mock_embed)
|
|
|
|
large_content = "x" * (SMALL_NOTE_CONTENT_LIMIT + 500)
|
|
|
|
with (
|
|
patch(
|
|
"basic_memory.repository.search_repository_base.db.scoped_session", fake_scoped_session
|
|
),
|
|
patch.object(repo, "_ensure_vector_tables", new_callable=AsyncMock),
|
|
patch.object(repo, "_prepare_vector_session", new_callable=AsyncMock),
|
|
patch.object(repo, "_run_vector_query", new_callable=AsyncMock, return_value=fake_rows),
|
|
patch.object(
|
|
repo,
|
|
"_fetch_search_index_rows_by_ids",
|
|
new_callable=AsyncMock,
|
|
return_value={("entity", 0): FakeRow(id=0, content_snippet=large_content)},
|
|
),
|
|
):
|
|
results = await repo._search_vector_only(**COMMON_SEARCH_KWARGS)
|
|
|
|
assert len(results) == 1
|
|
# Should use chunk text, not the full content
|
|
assert results[0].matched_chunk_text == "chunk text for entity:0:0"
|
|
assert results[0].matched_chunk_text != large_content
|