basicmachines-co-basic-memory/tests/services/test_semantic_search.py

"""Semantic search service regression tests for local SQLite search."""

from datetime import datetime
from types import SimpleNamespace
from unittest.mock import AsyncMock

import pytest
from sqlalchemy import text

from basic_memory import db
from basic_memory.repository import EntityRepository
from basic_memory.repository.search_repository_base import VectorSyncBatchResult
from basic_memory.repository.semantic_errors import (
    SemanticDependenciesMissingError,
    SemanticSearchDisabledError,
)
from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
from basic_memory.schemas.search import SearchItemType, SearchQuery, SearchRetrievalMode


def _sqlite_repo(search_service) -> SQLiteSearchRepository:
    repository = search_service.repository
    if not isinstance(repository, SQLiteSearchRepository):
        pytest.skip("Semantic retrieval behavior is local SQLite-only in this phase.")
    return repository


@pytest.mark.asyncio
async def test_semantic_vector_search_fails_when_disabled(search_service, test_graph):
    """Vector mode should fail fast when semantic search is disabled."""
    repository = _sqlite_repo(search_service)
    repository._semantic_enabled = False

    with pytest.raises(SemanticSearchDisabledError):
        await search_service.search(
            SearchQuery(
                text="Connected Entity",
                retrieval_mode=SearchRetrievalMode.VECTOR,
            )
        )


@pytest.mark.asyncio
async def test_semantic_hybrid_search_fails_when_disabled(search_service, test_graph):
    """Hybrid mode should fail fast when semantic search is disabled."""
    repository = _sqlite_repo(search_service)
    repository._semantic_enabled = False

    with pytest.raises(SemanticSearchDisabledError):
        await search_service.search(
            SearchQuery(
                text="Root Entity",
                retrieval_mode=SearchRetrievalMode.HYBRID,
            )
        )


@pytest.mark.asyncio
async def test_semantic_vector_search_fails_when_provider_unavailable(search_service, test_graph):
    """Vector mode should fail fast when semantic provider is unavailable."""
    repository = _sqlite_repo(search_service)
    repository._semantic_enabled = True
    repository._embedding_provider = None
    repository._vector_tables_initialized = False

    with pytest.raises(SemanticDependenciesMissingError):
        await search_service.search(
            SearchQuery(
                text="Root Entity",
                retrieval_mode=SearchRetrievalMode.VECTOR,
            )
        )


@pytest.mark.asyncio
async def test_semantic_vector_mode_rejects_non_text_query(search_service, test_graph):
    """Vector mode should not silently fall back for title-only queries."""
    with pytest.raises(ValueError):
        await search_service.search(
            SearchQuery(
                title="Root",
                retrieval_mode=SearchRetrievalMode.VECTOR,
                entity_types=[SearchItemType.ENTITY],
            )
        )


@pytest.mark.asyncio
async def test_semantic_fts_mode_still_returns_observations(search_service, test_graph):
    """Explicit FTS mode should preserve existing mixed result behavior."""
    results = await search_service.search(
        SearchQuery(
            text="Root note 1",
            retrieval_mode=SearchRetrievalMode.FTS,
        )
    )

    assert results
    assert any(result.type == SearchItemType.OBSERVATION.value for result in results)


@pytest.mark.asyncio
async def test_semantic_vector_sync_skips_embed_opt_out_and_clears_vectors(
    search_service, monkeypatch
):
    """Embed opt-out should clear stale vectors instead of regenerating them."""
    repository = _sqlite_repo(search_service)
    repository._semantic_enabled = True

    monkeypatch.setattr(
        search_service.entity_repository,
        "find_by_id",
        AsyncMock(return_value=SimpleNamespace(id=42, entity_metadata={"embed": False})),
    )
    sync_vectors = AsyncMock()
    delete_entity_vectors = AsyncMock()
    monkeypatch.setattr(repository, "sync_entity_vectors", sync_vectors)
    monkeypatch.setattr(repository, "delete_entity_vector_rows", delete_entity_vectors)

    await search_service.sync_entity_vectors(42)

    sync_vectors.assert_not_awaited()
    delete_entity_vectors.assert_awaited_once_with(42)


@pytest.mark.asyncio
async def test_semantic_vector_sync_resumes_when_embed_opt_out_removed(search_service, monkeypatch):
    """Removing the opt-out should restore normal embedding sync."""
    repository = _sqlite_repo(search_service)
    repository._semantic_enabled = True

    monkeypatch.setattr(
        search_service.entity_repository,
        "find_by_id",
        AsyncMock(return_value=SimpleNamespace(id=42, entity_metadata={})),
    )
    sync_vectors = AsyncMock()
    execute_query = AsyncMock()
    monkeypatch.setattr(repository, "sync_entity_vectors", sync_vectors)
    monkeypatch.setattr(repository, "execute_query", execute_query)

    await search_service.sync_entity_vectors(42)

    sync_vectors.assert_awaited_once_with(42)
    execute_query.assert_not_awaited()


@pytest.mark.asyncio
async def test_semantic_vector_sync_batch_skips_embed_opt_out_and_reports_skips(
    search_service, monkeypatch
):
    """Batch vector sync should only embed eligible notes and report skipped opt-outs."""
    repository = _sqlite_repo(search_service)
    repository._semantic_enabled = True

    monkeypatch.setattr(
        search_service.entity_repository,
        "find_by_ids",
        AsyncMock(
            return_value=[
                SimpleNamespace(id=41, entity_metadata={"embed": False}),
                SimpleNamespace(id=42, entity_metadata={}),
            ]
        ),
    )
    sync_batch = AsyncMock(
        return_value=VectorSyncBatchResult(
            entities_total=1,
            entities_synced=1,
            entities_failed=0,
        )
    )
    delete_entity_vectors = AsyncMock()
    monkeypatch.setattr(repository, "sync_entity_vectors_batch", sync_batch)
    monkeypatch.setattr(repository, "delete_entity_vector_rows", delete_entity_vectors)

    result = await search_service.sync_entity_vectors_batch([41, 42])

    sync_batch.assert_awaited_once()
    sync_batch_args = sync_batch.await_args
    assert sync_batch_args is not None
    assert sync_batch_args.args[0] == [42]
    assert result.entities_total == 2
    assert result.entities_synced == 1
    assert result.entities_skipped == 1
    delete_entity_vectors.assert_awaited_once_with(41)


@pytest.mark.asyncio
async def test_embed_opt_out_note_still_participates_in_fts(
    search_service, session_maker, test_project
):
    """Per-note semantic opt-out should not remove the note from FTS search."""
    entity_repo = EntityRepository(session_maker, project_id=test_project.id)
    entity = await entity_repo.create(
        {
            "title": "FTS Opt Out",
            "note_type": "note",
            "entity_metadata": {"embed": False},
            "content_type": "text/markdown",
            "file_path": "test/fts-opt-out.md",
            "permalink": "test/fts-opt-out",
            "project_id": test_project.id,
            "created_at": datetime.now(),
            "updated_at": datetime.now(),
        }
    )

    await search_service.index_entity(
        entity,
        content="This note should stay searchable through full text indexing.",
    )

    results = await search_service.search(
        SearchQuery(
            text="stay searchable",
            retrieval_mode=SearchRetrievalMode.FTS,
        )
    )

    assert any(result.entity_id == entity.id for result in results)


@pytest.mark.asyncio
async def test_reindex_vectors_respects_embed_opt_out(search_service, monkeypatch):
    """Full vector reindex should route through the service-level opt-out filter."""
    monkeypatch.setattr(
        search_service.entity_repository,
        "find_all",
        AsyncMock(
            return_value=[
                SimpleNamespace(id=41, entity_metadata={"embed": False}),
                SimpleNamespace(id=42, entity_metadata={}),
            ]
        ),
    )
    purge_stale_rows = AsyncMock()
    sync_batch = AsyncMock(
        return_value=VectorSyncBatchResult(
            entities_total=2,
            entities_synced=1,
            entities_failed=0,
            entities_skipped=1,
        )
    )
    monkeypatch.setattr(search_service, "_purge_stale_search_rows", purge_stale_rows)
    monkeypatch.setattr(search_service, "sync_entity_vectors_batch", sync_batch)

    stats = await search_service.reindex_vectors()

    purge_stale_rows.assert_awaited_once()
    sync_batch.assert_awaited_once_with([41, 42], progress_callback=None)
    assert stats == {
        "total_entities": 2,
        "embedded": 1,
        "skipped": 1,
        "errors": 0,
    }


@pytest.mark.asyncio
async def test_reindex_vectors_purges_sqlite_vectors_before_sync(search_service, monkeypatch):
    """Regression for #829: stale vec0 rows must be purged before batch sync."""
    repository = _sqlite_repo(search_service)
    repository._semantic_enabled = True
    calls: list[str] = []

    monkeypatch.setattr(
        search_service.entity_repository,
        "find_all",
        AsyncMock(return_value=[SimpleNamespace(id=42, entity_metadata={})]),
    )

    async def delete_stale_vector_rows():
        calls.append("purge")

    async def sync_entity_vectors_batch(entity_ids, progress_callback=None):
        assert entity_ids == [42]
        assert progress_callback is None
        calls.append("sync")
        return VectorSyncBatchResult(entities_total=1, entities_synced=1, entities_failed=0)

    monkeypatch.setattr(repository, "delete_stale_vector_rows", delete_stale_vector_rows)
    monkeypatch.setattr(search_service, "sync_entity_vectors_batch", sync_entity_vectors_batch)

    stats = await search_service.reindex_vectors()

    assert calls == ["purge", "sync"]
    assert stats == {
        "total_entities": 1,
        "embedded": 1,
        "skipped": 0,
        "errors": 0,
    }


@pytest.mark.asyncio
async def test_reindex_all_uses_sqlite_vec_aware_drop(search_service, monkeypatch):
    """Full service reindex should not drop vec0 tables through a raw connection."""
    repository = _sqlite_repo(search_service)
    executed_sql: list[str] = []
    calls: list[str] = []

    async def execute_query(query, params=None):
        executed_sql.append(str(query))

    async def drop_vector_tables():
        calls.append("drop_vector_tables")

    monkeypatch.setattr(repository, "execute_query", execute_query)
    monkeypatch.setattr(repository, "drop_vector_tables", drop_vector_tables)
    monkeypatch.setattr(search_service, "init_search_index", AsyncMock())
    monkeypatch.setattr(search_service.entity_repository, "find_all", AsyncMock(return_value=[]))

    await search_service.reindex_all()

    assert calls == ["drop_vector_tables"]
    assert all("search_vector_embeddings" not in sql for sql in executed_sql)


@pytest.mark.asyncio
async def test_drop_vector_tables_skips_sqlite_vec_load_for_plain_table(
    search_service,
    monkeypatch,
):
    """Plain compatibility tables should not require sqlite-vec to be loaded."""
    repository = _sqlite_repo(search_service)
    await repository.drop_vector_tables()

    async with db.scoped_session(repository.session_maker) as session:
        await session.execute(
            text("CREATE TABLE search_vector_embeddings (rowid INTEGER PRIMARY KEY)")
        )
        await session.commit()

    async def fail_if_loaded(_session):
        raise AssertionError("plain tables should not load sqlite-vec")

    monkeypatch.setattr(repository, "_ensure_sqlite_vec_loaded", fail_if_loaded)

    await repository.drop_vector_tables()

    async with db.scoped_session(repository.session_maker) as session:
        result = await session.execute(
            text(
                "SELECT name FROM sqlite_master "
                "WHERE type = 'table' AND name = 'search_vector_embeddings'"
            )
        )
        assert result.scalar() is None


@pytest.mark.asyncio
async def test_reindex_vectors_force_full_clears_project_vectors_before_resync(
    search_service, monkeypatch
):
    """Force-full vector reindex should clear derived vectors before batch sync."""
    repository = _sqlite_repo(search_service)
    repository._semantic_enabled = True

    monkeypatch.setattr(
        search_service.entity_repository,
        "find_all",
        AsyncMock(
            return_value=[
                SimpleNamespace(id=41, entity_metadata={}),
                SimpleNamespace(id=42, entity_metadata={}),
            ]
        ),
    )
    purge_stale_rows = AsyncMock()
    delete_project_vectors = AsyncMock()
    sync_batch = AsyncMock(
        return_value=VectorSyncBatchResult(
            entities_total=2,
            entities_synced=2,
            entities_failed=0,
        )
    )
    monkeypatch.setattr(search_service, "_purge_stale_search_rows", purge_stale_rows)
    monkeypatch.setattr(repository, "delete_project_vector_rows", delete_project_vectors)
    monkeypatch.setattr(search_service, "sync_entity_vectors_batch", sync_batch)

    stats = await search_service.reindex_vectors(force_full=True)

    purge_stale_rows.assert_awaited_once()
    delete_project_vectors.assert_awaited_once()
    sync_batch.assert_awaited_once_with([41, 42], progress_callback=None)
    assert stats == {
        "total_entities": 2,
        "embedded": 2,
        "skipped": 0,
        "errors": 0,
    }


@pytest.mark.asyncio
async def test_semantic_vector_sync_batch_cleans_up_unknown_ids(search_service, monkeypatch):
    """Deleted entity IDs should still flow through repository cleanup instead of being dropped."""
    repository = _sqlite_repo(search_service)
    repository._semantic_enabled = True

    monkeypatch.setattr(
        search_service.entity_repository,
        "find_by_ids",
        AsyncMock(return_value=[SimpleNamespace(id=42, entity_metadata={})]),
    )
    sync_batch = AsyncMock(
        side_effect=[
            VectorSyncBatchResult(
                entities_total=1,
                entities_synced=1,
                entities_failed=0,
                entities_skipped=1,
            ),
            VectorSyncBatchResult(
                entities_total=1,
                entities_synced=1,
                entities_failed=0,
            ),
        ]
    )
    monkeypatch.setattr(repository, "sync_entity_vectors_batch", sync_batch)
    progress_callback = AsyncMock()

    result = await search_service.sync_entity_vectors_batch([41, 42], progress_callback)

    assert sync_batch.await_count == 2
    called_entity_ids = {tuple(call.args[0]) for call in sync_batch.await_args_list}
    assert called_entity_ids == {(41,), (42,)}
    progress_callback_calls = [
        call
        for call in sync_batch.await_args_list
        if call.kwargs.get("progress_callback") is not None
    ]
    assert len(progress_callback_calls) == 1
    assert progress_callback_calls[0].args[0] == [42]
    assert progress_callback_calls[0].kwargs["progress_callback"] is progress_callback
    assert result.entities_total == 2
    assert result.entities_synced == 2
    assert result.entities_failed == 0
    assert result.entities_skipped == 0