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

"""Quality metric computation and reporting for semantic search benchmarks.

Computes hit@1, recall@5, and MRR@10 from search results, plus timing
data for performance comparison across backends and providers.
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from statistics import mean

from basic_memory.repository.search_index_row import SearchIndexRow


# --- Relevance helpers ---


def first_relevant_rank(results: list[SearchIndexRow], expected_topic: str, k: int) -> int | None:
    """Return the 1-based rank of the first result matching ``expected_topic``, or None."""
    expected_prefix = f"bench/{expected_topic}-"
    for rank, row in enumerate(results[:k], start=1):
        if (row.permalink or "").startswith(expected_prefix):
            return rank
    return None


# --- Metric dataclass ---


@dataclass
class QualityMetrics:
    """Aggregated quality metrics for a (combo, suite, mode) triple."""

    combo: str
    suite: str
    mode: str
    cases: int = 0
    hits_at_1: int = 0
    hits_at_5: int = 0
    reciprocal_rank_sum: float = 0.0
    per_query: list[dict] = field(default_factory=list)
    # Timing data: per-query latencies in seconds
    latencies: list[float] = field(default_factory=list)

    def record(
        self, query_text: str, expected_topic: str, rank: int | None, latency: float = 0.0
    ) -> None:
        self.cases += 1
        entry = {"query": query_text, "expected": expected_topic, "rank": rank}
        self.per_query.append(entry)
        self.latencies.append(latency)
        if rank is None:
            return
        self.reciprocal_rank_sum += 1.0 / rank
        if rank == 1:
            self.hits_at_1 += 1
        if rank <= 5:
            self.hits_at_5 += 1

    @property
    def hit_at_1(self) -> float:
        return self.hits_at_1 / self.cases if self.cases else 0.0

    @property
    def recall_at_5(self) -> float:
        return self.hits_at_5 / self.cases if self.cases else 0.0

    @property
    def mrr_at_10(self) -> float:
        return self.reciprocal_rank_sum / self.cases if self.cases else 0.0

    @property
    def total_time_ms(self) -> float:
        return sum(self.latencies) * 1000

    @property
    def avg_latency_ms(self) -> float:
        return mean(self.latencies) * 1000 if self.latencies else 0.0

    def as_dict(self) -> dict:
        return {
            "combo": self.combo,
            "suite": self.suite,
            "mode": self.mode,
            "cases": self.cases,
            "hit_at_1": round(self.hit_at_1, 4),
            "recall_at_5": round(self.recall_at_5, 4),
            "mrr_at_10": round(self.mrr_at_10, 4),
            "avg_latency_ms": round(self.avg_latency_ms, 2),
            "total_time_ms": round(self.total_time_ms, 2),
        }


# --- Comparison table ---


def format_comparison_table(all_metrics: list[QualityMetrics]) -> str:
    """Format a list of QualityMetrics into an ASCII comparison table."""
    header = (
        f"{'Combo':<25} {'Suite':<12} {'Mode':<8} "
        f"{'hit@1':>6} {'R@5':>6} {'MRR@10':>7} {'avg_ms':>8} {'total_ms':>9}"
    )
    separator = "-" * len(header)
    lines = [separator, header, separator]

    for m in sorted(all_metrics, key=lambda x: (x.suite, x.combo, x.mode)):
        lines.append(
            f"{m.combo:<25} {m.suite:<12} {m.mode:<8} "
            f"{m.hit_at_1:>6.3f} {m.recall_at_5:>6.3f} {m.mrr_at_10:>7.3f} "
            f"{m.avg_latency_ms:>8.1f} {m.total_time_ms:>9.1f}"
        )

    lines.append(separator)
    return "\n".join(lines)


# --- Artifact output ---


def write_benchmark_artifact(all_metrics: list[QualityMetrics]) -> None:
    """Append JSON-lines benchmark artifact if BASIC_MEMORY_BENCHMARK_OUTPUT is set."""
    output_path = os.getenv("BASIC_MEMORY_BENCHMARK_OUTPUT")
    if not output_path:
        return

    artifact_path = Path(output_path).expanduser()
    artifact_path.parent.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now(timezone.utc).isoformat()

    with artifact_path.open("a", encoding="utf-8") as f:
        for m in all_metrics:
            payload = {
                "benchmark": f"semantic-quality-{m.combo}-{m.suite}-{m.mode}",
                "timestamp_utc": timestamp,
                "metrics": m.as_dict(),
            }
            f.write(json.dumps(payload, sort_keys=True) + "\n")