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

"""Semantic search quality benchmarks across backend×provider combinations.

Runs identical query suites (lexical + paraphrase) against five configurations:

  sqlite-fts          SQLite FTS5, no embeddings
  sqlite-fastembed    SQLite + FastEmbed (384-d ONNX)
  postgres-fts        Postgres tsvector, no embeddings
  postgres-fastembed  Postgres + FastEmbed (384-d)
  postgres-openai     Postgres + OpenAI (1536-d, needs OPENAI_API_KEY)

Quality is measured via hit@1, recall@5, and MRR@10.  A comparison table
is printed at the end, and JSON-lines artifacts are written when
``BASIC_MEMORY_BENCHMARK_OUTPUT`` is set.
"""

from __future__ import annotations

import time

import pytest

from basic_memory.schemas.search import SearchItemType, SearchQuery, SearchRetrievalMode

from semantic.conftest import (
    ALL_COMBOS,
    SearchCombo,
    create_search_service,
    skip_if_needed,
    _create_fastembed_provider,
    _create_openai_provider,
)
from semantic.corpus import QUERY_SUITES, seed_benchmark_notes
from semantic.metrics import (
    QualityMetrics,
    first_relevant_rank,
    format_comparison_table,
    write_benchmark_artifact,
)


# --- Thresholds (conservative, tighten over time) ---
# Keys: (combo.name, suite_name, mode)
# Only combos/suites we have strong expectations for are listed.

RECALL_AT_5_THRESHOLDS: dict[tuple[str, str, str], float] = {
    # FTS-only: realistic corpus makes pure keyword matching harder
    ("sqlite-fts", "lexical", "fts"): 0.25,
    ("postgres-fts", "lexical", "fts"): 0.25,
    # FastEmbed hybrid should improve on FTS for both suites
    ("sqlite-fastembed", "lexical", "hybrid"): 0.37,
    ("sqlite-fastembed", "paraphrase", "hybrid"): 0.25,
    ("postgres-fastembed", "lexical", "hybrid"): 0.37,
    ("postgres-fastembed", "paraphrase", "hybrid"): 0.25,
    # OpenAI hybrid should handle paraphrases better than FastEmbed.
    ("postgres-openai", "lexical", "hybrid"): 0.37,
    ("postgres-openai", "paraphrase", "hybrid"): 0.25,
}


def _resolve_provider(combo: SearchCombo):
    """Build the embedding provider for a combo, or None for FTS-only."""
    if combo.provider_name == "fastembed":
        return _create_fastembed_provider()
    if combo.provider_name == "openai":
        return _create_openai_provider()
    return None


def _retrieval_modes(combo: SearchCombo) -> list[SearchRetrievalMode]:
    """Return the retrieval modes to test for a given combo.

    FTS-only combos: [FTS]
    Semantic combos: [FTS, VECTOR, HYBRID]
    """
    if combo.provider_name is None:
        return [SearchRetrievalMode.FTS]
    return [SearchRetrievalMode.FTS, SearchRetrievalMode.VECTOR, SearchRetrievalMode.HYBRID]


# --- Parameterized test ---


@pytest.mark.asyncio
@pytest.mark.semantic
@pytest.mark.benchmark
@pytest.mark.parametrize("combo", ALL_COMBOS, ids=[c.name for c in ALL_COMBOS])
async def test_semantic_quality(
    combo: SearchCombo,
    sqlite_engine_factory,
    postgres_engine_factory,
    tmp_path,
):
    """Benchmark search quality for a single backend×provider combo."""
    skip_if_needed(combo)

    # Pick the right engine factory
    from basic_memory.config import DatabaseBackend

    if combo.backend == DatabaseBackend.SQLITE:
        engine_factory_result = sqlite_engine_factory
    else:
        if postgres_engine_factory is None:
            pytest.skip("Postgres engine not available")
        engine_factory_result = postgres_engine_factory

    provider = _resolve_provider(combo)
    search_service = await create_search_service(
        engine_factory_result, combo, tmp_path, embedding_provider=provider
    )

    # Seed corpus
    entities = await seed_benchmark_notes(search_service, note_count=240)
    assert len(entities) == 240

    # Collect metrics for each (suite, mode)
    all_metrics: list[QualityMetrics] = []

    for suite_name, cases in QUERY_SUITES.items():
        for mode in _retrieval_modes(combo):
            metrics = QualityMetrics(
                combo=combo.name,
                suite=suite_name,
                mode=mode.value,
            )

            for case in cases:
                t0 = time.perf_counter()
                results = await search_service.search(
                    SearchQuery(
                        text=case.text,
                        retrieval_mode=mode,
                        entity_types=[SearchItemType.ENTITY],
                    ),
                    limit=10,
                )
                latency = time.perf_counter() - t0

                rank = first_relevant_rank(results, case.expected_topic, k=10) if results else None
                metrics.record(case.text, case.expected_topic, rank, latency=latency)

            all_metrics.append(metrics)

    # Print comparison table
    table = format_comparison_table(all_metrics)
    print(f"\n{table}")

    # Write JSON artifact
    write_benchmark_artifact(all_metrics)

    # Enforce thresholds
    for m in all_metrics:
        threshold_key = (m.combo, m.suite, m.mode)
        threshold = RECALL_AT_5_THRESHOLDS.get(threshold_key)
        if threshold is not None:
            assert m.recall_at_5 >= threshold, (
                f"recall@5 for {m.combo}/{m.suite}/{m.mode}: {m.recall_at_5:.3f} < {threshold:.3f}"
            )