basicmachines-co-basic-memory/tests/repository/test_vector_threshold.py

"""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,
        allow_relaxed: bool = False,
    ) -> 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