basicmachines-co-basic-memory/test-int/semantic/conftest.py

"""Fixtures for semantic search benchmark tests.

Provides a pgvector-enabled container, engine factories for both backends,
and a parameterized ``search_combo`` fixture that yields a configured
SearchService for each (backend, provider) combination.
"""

from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path

import pytest
import pytest_asyncio
from dotenv import load_dotenv
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool
from testcontainers.postgres import PostgresContainer

from basic_memory import db
from basic_memory.config import BasicMemoryConfig, DatabaseBackend
from basic_memory.db import DatabaseType, engine_session_factory
from basic_memory.markdown import EntityParser
from basic_memory.markdown.markdown_processor import MarkdownProcessor
from basic_memory.models.base import Base
from basic_memory.models.search import (
    CREATE_POSTGRES_SEARCH_INDEX_FTS,
    CREATE_POSTGRES_SEARCH_INDEX_METADATA,
    CREATE_POSTGRES_SEARCH_INDEX_PERMALINK,
    CREATE_POSTGRES_SEARCH_INDEX_TABLE,
    CREATE_SEARCH_INDEX,
)
from basic_memory.repository.embedding_provider import EmbeddingProvider
from basic_memory.repository.entity_repository import EntityRepository
from basic_memory.repository.project_repository import ProjectRepository
from basic_memory.repository.search_repository import SearchRepository
from basic_memory.services.file_service import FileService
from basic_memory.services.search_service import SearchService

# Load .env so OPENAI_API_KEY (and other keys) are available to providers
load_dotenv()


# --- Combo descriptor ---


@dataclass(frozen=True)
class SearchCombo:
    """Describes a (backend, provider) combination for benchmark parameterization."""

    name: str
    backend: DatabaseBackend
    provider_name: str | None  # None = FTS-only
    dimensions: int | None


# All combinations the suite covers
ALL_COMBOS = [
    SearchCombo("sqlite-fts", DatabaseBackend.SQLITE, None, None),
    SearchCombo("sqlite-fastembed", DatabaseBackend.SQLITE, "fastembed", 384),
    SearchCombo("postgres-fts", DatabaseBackend.POSTGRES, None, None),
    SearchCombo("postgres-fastembed", DatabaseBackend.POSTGRES, "fastembed", 384),
    SearchCombo("postgres-openai", DatabaseBackend.POSTGRES, "openai", 1536),
]


# --- Skip guards ---


def _docker_available() -> bool:
    """Check if Docker is available for testcontainers."""
    import shutil

    return shutil.which("docker") is not None


def _fastembed_available() -> bool:
    try:
        import fastembed  # noqa: F401

        return True
    except ImportError:
        return False


def _openai_key_available() -> bool:
    return bool(os.environ.get("OPENAI_API_KEY"))


def skip_if_needed(combo: SearchCombo) -> None:
    """Skip the current test if the combo's requirements aren't met."""
    if combo.backend == DatabaseBackend.POSTGRES and not _docker_available():
        pytest.skip("Docker not available for Postgres testcontainer")

    if combo.provider_name == "fastembed" and not _fastembed_available():
        pytest.skip("fastembed not installed (install/update basic-memory)")

    if combo.provider_name == "openai":
        if not _fastembed_available():
            pytest.skip("semantic dependencies not installed")
        if not _openai_key_available():
            pytest.skip("OPENAI_API_KEY not set")


# --- pgvector container (session-scoped, independent of main test suite) ---


@pytest.fixture(scope="session")
def pgvector_container():
    """Session-scoped pgvector container for semantic benchmarks.

    Uses pgvector/pgvector:pg16 image to get the vector extension.
    Only starts if Docker is available; yields None otherwise.
    """
    if not _docker_available():
        yield None
        return

    with PostgresContainer("pgvector/pgvector:pg16") as pg:
        yield pg


# --- Engine factories ---


@pytest_asyncio.fixture
async def sqlite_engine_factory(tmp_path):
    """Create a SQLite engine + session factory for benchmark use."""
    db_path = tmp_path / "bench.db"

    # Explicit config forces SQLite backend regardless of user's local config
    sqlite_config = BasicMemoryConfig(database_backend=DatabaseBackend.SQLITE)
    async with engine_session_factory(db_path, DatabaseType.FILESYSTEM, config=sqlite_config) as (
        engine,
        session_maker,
    ):
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        async with db.scoped_session(session_maker) as session:
            await session.execute(text("DROP TABLE IF EXISTS search_index"))
            await session.execute(CREATE_SEARCH_INDEX)
            await session.commit()

        yield engine, session_maker


@pytest_asyncio.fixture
async def postgres_engine_factory(pgvector_container):
    """Create a Postgres engine + session factory with pgvector extension."""
    if pgvector_container is None:
        yield None
        return

    sync_url = pgvector_container.get_connection_url()
    async_url = sync_url.replace("postgresql+psycopg2", "postgresql+asyncpg")

    engine = create_async_engine(async_url, echo=False, poolclass=NullPool)
    session_maker = async_sessionmaker(
        bind=engine,
        class_=AsyncSession,
        expire_on_commit=False,
        autoflush=False,
    )

    # Create schema from scratch for each test
    async with engine.begin() as conn:
        await conn.execute(text("DROP TABLE IF EXISTS search_vector_embeddings CASCADE"))
        await conn.execute(text("DROP TABLE IF EXISTS search_vector_chunks CASCADE"))
        await conn.execute(text("DROP TABLE IF EXISTS search_index CASCADE"))
        await conn.run_sync(Base.metadata.drop_all)
        await conn.run_sync(Base.metadata.create_all)
        await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_TABLE)
        await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_FTS)
        await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_METADATA)
        await conn.execute(CREATE_POSTGRES_SEARCH_INDEX_PERMALINK)

    yield engine, session_maker

    await engine.dispose()


# --- Embedding provider factories ---


def _create_fastembed_provider() -> EmbeddingProvider:
    from basic_memory.repository.fastembed_provider import FastEmbedEmbeddingProvider

    return FastEmbedEmbeddingProvider(model_name="bge-small-en-v1.5", batch_size=64)


def _create_openai_provider() -> EmbeddingProvider:
    from basic_memory.repository.openai_provider import OpenAIEmbeddingProvider

    return OpenAIEmbeddingProvider(model_name="text-embedding-3-small", dimensions=1536)


# --- Search service factory ---


async def create_search_service(
    engine_factory_result,
    combo: SearchCombo,
    tmp_path: Path,
    embedding_provider: EmbeddingProvider | None = None,
) -> SearchService:
    """Build a fully wired SearchService for a given combo."""
    engine, session_maker = engine_factory_result

    # Create test project
    project_repo = ProjectRepository(session_maker)
    project = await project_repo.create(
        {
            "name": "bench-project",
            "description": "Semantic benchmark project",
            "path": str(tmp_path),
            "is_active": True,
            "is_default": True,
        }
    )

    # Build app config
    semantic_enabled = combo.provider_name is not None
    app_config = BasicMemoryConfig(
        env="test",
        projects={"bench-project": str(tmp_path)},
        default_project="bench-project",
        database_backend=combo.backend,
        semantic_search_enabled=semantic_enabled,
    )

    # Create search repository (backend-specific)
    if combo.backend == DatabaseBackend.POSTGRES:
        from basic_memory.repository.postgres_search_repository import PostgresSearchRepository

        search_repo: SearchRepository = PostgresSearchRepository(
            session_maker,
            project_id=project.id,
            app_config=app_config,
            embedding_provider=embedding_provider,
        )
    else:
        from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository

        repo = SQLiteSearchRepository(
            session_maker,
            project_id=project.id,
            app_config=app_config,
        )
        # Inject provider directly for SQLite
        if embedding_provider is not None:
            repo._semantic_enabled = True
            repo._embedding_provider = embedding_provider
            repo._vector_dimensions = embedding_provider.dimensions
            repo._vector_tables_initialized = False
        search_repo = repo

    entity_repo = EntityRepository(session_maker, project_id=project.id)
    entity_parser = EntityParser(tmp_path)
    markdown_processor = MarkdownProcessor(entity_parser)
    file_service = FileService(tmp_path, markdown_processor)

    service = SearchService(search_repo, entity_repo, file_service)
    await service.init_search_index()
    return service