"""Tests for semantic_min_similarity threshold filtering in vector search."""
from contextlib import asynccontextmanager
from dataclasses import dataclass
from unittest.mock import AsyncMock, patch
import pytest
from basic_memory.repository.search_repository_base import (
SMALL_NOTE_CONTENT_LIMIT,
TOP_CHUNKS_PER_RESULT,
SearchRepositoryBase,
)
@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, **kwargs):
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
@asynccontextmanager
async def fake_scoped_session(session_maker):
"""Fake scoped_session that yields a mock session object."""
yield AsyncMock()
COMMON_SEARCH_KWARGS = dict(
search_text="test",
permalink=None,
permalink_match=None,
title=None,
note_types=None,
after_date=None,
search_item_types=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 = type("EP", (), {"embed_query": mock_embed, "dimensions": 384})()
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={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 = type("EP", (), {"embed_query": mock_embed, "dimensions": 384})()
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={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 = type("EP", (), {"embed_query": mock_embed, "dimensions": 384})()
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 = type("EP", (), {"embed_query": mock_embed, "dimensions": 384})()
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={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 = type("EP", (), {"embed_query": mock_embed, "dimensions": 384})()
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={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 = type("EP", (), {"embed_query": mock_embed, "dimensions": 384})()
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={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 = type("EP", (), {"embed_query": mock_embed, "dimensions": 384})()
# 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={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
# 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 = type("EP", (), {"embed_query": mock_embed, "dimensions": 384})()
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={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 = type("EP", (), {"embed_query": mock_embed, "dimensions": 384})()
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={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
