chopratejas--headroom/tests/test_memory_ranker.py

"""Tests for :class:`headroom.proxy.memory_ranker.MemoryRanker` +
:class:`RecencyBoostRanker`.

Pre-this-PR Headroom ranked memory candidates by pure cosine
similarity. Every other memory system we surveyed (Letta, Mem0,
Cognee, Supermemory) re-ranks beyond cosine — recency / source /
access-count / decay are table-stakes. The pure-cosine baseline
returns 6-month-old memories with 0.9 similarity ahead of fresh
memories with 0.5 — wrong for most use cases.

``RecencyBoostRanker`` is the first ranker we ship: a pure-function
``score = cosine × exp(-age_days / decay_days)`` re-ranker. Default
``decay_days=30`` (half-life ~21 days). Other rankers (source-weight,
access-count) plug into the same :class:`MemoryRanker` protocol in
follow-on PRs.

Performance: O(N) over candidates where N = top_k = ~10. One ``exp()``
per candidate. Sub-microsecond per request — no embedding compute, no
I/O. The ranker is pure and Rust-portable.
"""

from __future__ import annotations

from dataclasses import FrozenInstanceError
from datetime import datetime, timedelta, timezone

from headroom.proxy.memory_ranker import (
    MemoryCandidate,
    RecencyBoostRanker,
)

_UTC = timezone.utc

# ── Helpers ───────────────────────────────────────────────────────────


def _candidate(content: str, score: float, age_days: float = 0.0) -> MemoryCandidate:
    """Build a MemoryCandidate at the given cosine score and age."""
    created = datetime.now(_UTC) - timedelta(days=age_days)
    return MemoryCandidate(content=content, score=score, created_at=created)


def _candidate_no_timestamp(content: str, score: float) -> MemoryCandidate:
    """Build a MemoryCandidate without a created_at (back-compat shape)."""
    return MemoryCandidate(content=content, score=score, created_at=None)


# ── MemoryCandidate value-type contract ──────────────────────────────


def test_candidate_is_frozen() -> None:
    """Frozen so a ranker can't mutate a candidate's score and lie about
    which candidates it returned."""
    c = _candidate("x", 0.9, age_days=0)
    try:
        c.score = 0.1  # type: ignore[misc]
    except FrozenInstanceError:
        pass
    else:
        raise AssertionError("MemoryCandidate must be frozen")


def test_from_backend_result_preserves_memory_id() -> None:
    """The adapter must carry ``memory.id`` through to MemoryCandidate.id
    so the auto-tail block can render it as the bracketed handle the
    model uses for memory_update / memory_delete. Pre-this-fix the
    adapter dropped the ID, which silently regressed the [id] auto-tail
    format on the ranker path."""

    class _Mem:
        id = "mem_abc_123"
        content = "User prefers Python."
        created_at = None
        metadata = {"source": "memory_save"}

    class _Result:
        memory = _Mem()
        score = 0.91
        related_entities = ("python",)

    cand = MemoryCandidate.from_backend_result(_Result())
    assert cand.id == "mem_abc_123"
    assert cand.content == "User prefers Python."
    assert cand.score == 0.91


def test_rank_preserves_memory_id() -> None:
    """The ranker must not drop the backend ID when rebuilding candidates."""
    cand = MemoryCandidate(content="User prefers Python.", score=0.91, id="mem_abc_123")
    out = RecencyBoostRanker().rank([cand])
    assert out[0].id == "mem_abc_123"


def test_from_backend_result_handles_missing_id() -> None:
    """Defensive: legacy backend rows without an ID become ``id=""``;
    the auto-tail formatter renders ``[?]`` for those rows, no crash."""

    class _Mem:
        # no .id attribute
        content = "legacy row"
        created_at = None
        metadata = {}

    class _Result:
        memory = _Mem()
        score = 0.5
        related_entities = ()

    cand = MemoryCandidate.from_backend_result(_Result())
    assert cand.id == ""


# ── RecencyBoostRanker contract ──────────────────────────────────────


def test_ranker_is_frozen() -> None:
    """The ranker config is itself immutable — operators set
    ``decay_days`` at construction; runtime cannot edit it."""
    r = RecencyBoostRanker()
    try:
        r.decay_days = 99  # type: ignore[misc]
    except FrozenInstanceError:
        pass
    else:
        raise AssertionError("RecencyBoostRanker must be frozen")


def test_ranker_default_decay_is_thirty_days() -> None:
    """30-day decay is the conservative default. At 30 days, factor is
    ~0.37 (e^{-1}); at 90 days, ~0.05. Tuned so a fresh memory with
    weak cosine doesn't dominate, but a 6-month-old strong-cosine
    can't dominate either."""
    assert RecencyBoostRanker().decay_days == 30.0


def test_ranker_default_decay_is_configurable() -> None:
    """Operators can tune decay; e.g., 7 days for an aggressive
    recency bias on rapidly-evolving codebases."""
    r = RecencyBoostRanker(decay_days=7.0)
    assert r.decay_days == 7.0


def test_ranker_returns_list_preserving_shape() -> None:
    """Output is a list of candidates (re-ranked). Length matches
    input length — the ranker does NOT filter, only re-orders. The
    budget filters; the ranker ranks."""
    candidates = [_candidate("a", 0.9), _candidate("b", 0.5)]
    out = RecencyBoostRanker().rank(candidates)
    assert len(out) == 2
    assert {c.content for c in out} == {"a", "b"}


# ── Recency boost behaviour ──────────────────────────────────────────


def test_equal_cosine_younger_wins() -> None:
    """Two candidates with identical cosine score — the younger one
    wins because its recency factor is closer to 1.0."""
    fresh = _candidate("fresh", 0.5, age_days=0)
    old = _candidate("old", 0.5, age_days=60)
    out = RecencyBoostRanker().rank([old, fresh])
    assert out[0].content == "fresh"
    assert out[1].content == "old"


def test_old_strong_cosine_can_still_beat_young_weak_cosine() -> None:
    """The boost is multiplicative, not absolute — a 60-day-old memory
    with 0.9 cosine (0.9 × 0.135 ≈ 0.12) still loses to a 0-day-old
    memory with 0.5 cosine (0.5 × 1.0 = 0.5). But a 5-day-old memory
    with 0.9 (0.9 × 0.847 ≈ 0.76) beats a 0-day-old with 0.5."""
    very_old_strong = _candidate("old_strong", 0.9, age_days=60)
    fresh_weak = _candidate("fresh_weak", 0.5, age_days=0)
    out = RecencyBoostRanker().rank([very_old_strong, fresh_weak])
    # fresh_weak should win because 60-day decay flattens the strong cosine
    assert out[0].content == "fresh_weak"

    # Versus: slightly-old strong beats fresh weak
    slightly_old_strong = _candidate("slightly_old_strong", 0.9, age_days=5)
    fresh_weak2 = _candidate("fresh_weak2", 0.5, age_days=0)
    out2 = RecencyBoostRanker().rank([fresh_weak2, slightly_old_strong])
    assert out2[0].content == "slightly_old_strong"


def test_decay_rate_changes_winner() -> None:
    """An aggressive decay_days=7 makes a 30-day-old memory much
    weaker than a default decay_days=30. Locks the configurability
    contract."""
    old_strong = _candidate("old_strong", 0.9, age_days=30)
    fresh_weak = _candidate("fresh_weak", 0.6, age_days=0)

    # decay_days=30: old × e^{-1} ≈ 0.331; fresh = 0.6 → fresh wins
    r_default = RecencyBoostRanker(decay_days=30.0)
    out_default = r_default.rank([old_strong, fresh_weak])
    assert out_default[0].content == "fresh_weak"

    # decay_days=120 (loose): old × e^{-0.25} ≈ 0.701; fresh = 0.6 → old wins
    r_loose = RecencyBoostRanker(decay_days=120.0)
    out_loose = r_loose.rank([old_strong, fresh_weak])
    assert out_loose[0].content == "old_strong"


def test_zero_age_memory_keeps_full_cosine() -> None:
    """At age=0 days, the recency factor is e^0 = 1.0 — the boosted
    score equals the original cosine. Fresh memories see no penalty."""
    fresh = _candidate("fresh", 0.7, age_days=0)
    out = RecencyBoostRanker().rank([fresh])
    # Compare with tolerance — datetime.now() drift between
    # _candidate() and rank() is microseconds, so factor ~ 1.0.
    assert out[0].score == 0.7 or abs(out[0].score - 0.7) < 1e-3


def test_candidate_without_timestamp_keeps_pure_cosine() -> None:
    """Backwards-compat: pre-this-PR candidates may not have a
    ``created_at`` (older rows / older backends). NULL timestamp
    means "treat as recency-neutral" — factor 1.0. Pure cosine."""
    no_ts = _candidate_no_timestamp("legacy", 0.8)
    out = RecencyBoostRanker().rank([no_ts])
    assert out[0].score == 0.8


def test_mixed_with_and_without_timestamps() -> None:
    """A backend that returns SOME candidates with timestamps and
    SOME without (e.g., during a migration) must still produce a
    sensible ranking. NULL-timestamp candidates get factor 1.0,
    timestamped ones get their decay."""
    fresh_ts = _candidate("fresh_ts", 0.6, age_days=0)
    old_ts = _candidate("old_ts", 0.6, age_days=60)
    no_ts_neutral = _candidate_no_timestamp("no_ts", 0.6)
    out = RecencyBoostRanker().rank([old_ts, no_ts_neutral, fresh_ts])
    # fresh_ts (~0.6) and no_ts (=0.6) tied at top — old_ts decayed.
    assert out[-1].content == "old_ts"
    assert {out[0].content, out[1].content} == {"fresh_ts", "no_ts"}


# ── Stability + edge cases ───────────────────────────────────────────


def test_empty_input_returns_empty_output() -> None:
    """No candidates → no candidates. Boundary case."""
    assert RecencyBoostRanker().rank([]) == []


def test_ranking_is_stable_for_identical_candidates() -> None:
    """Two candidates with identical content + score + age → stable
    order (no spurious reshuffling). Important for prefix-cache
    stability: a deterministic ranker means consecutive turns inject
    the same memory in the same order, preserving byte-stable
    output."""
    a = _candidate("same", 0.5, age_days=10)
    b = _candidate("same", 0.5, age_days=10)
    out = RecencyBoostRanker().rank([a, b])
    assert len(out) == 2


def test_negative_age_treated_as_zero() -> None:
    """Defensive: a candidate with a future ``created_at`` (clock
    skew) shouldn't crash or give a > 1.0 factor. ``exp(-age/decay)``
    with negative age gives > 1; we clamp to 1.0 so a clock-skewed
    candidate can't outrank a real fresh one with score=1.0
    artifically."""
    future = _candidate("future", 0.5, age_days=-10)  # 10 days in future
    fresh = _candidate("fresh", 0.5, age_days=0)
    out = RecencyBoostRanker().rank([future, fresh])
    # Both should have factor 1.0 (clamped) — score equal → stable order
    assert {out[0].content, out[1].content} == {"future", "fresh"}
    assert out[0].score == 0.5 or abs(out[0].score - 0.5) < 1e-3


# ── Rust-port shape ─────────────────────────────────────────────────


def test_ranker_is_pure_no_side_effects() -> None:
    """Calling rank() twice with the same input gives the same
    output. No state on the ranker; no I/O. Rust-portable."""
    candidates = [_candidate("a", 0.7, age_days=5), _candidate("b", 0.5, age_days=20)]
    r = RecencyBoostRanker()
    out1 = r.rank(candidates)
    out2 = r.rank(candidates)
    assert [c.content for c in out1] == [c.content for c in out2]
    # Inputs preserved — ranker did not mutate
    assert candidates[0].content == "a"
    assert candidates[1].content == "b"


def test_ranker_does_not_mutate_input_list() -> None:
    """Defence-in-depth: the input list and its elements must be
    unchanged after ranking. Frozen candidates make element mutation
    impossible; the list order itself must also be preserved."""
    a = _candidate("a", 0.5, age_days=20)
    b = _candidate("b", 0.5, age_days=5)
    candidates = [a, b]
    RecencyBoostRanker().rank(candidates)
    assert candidates == [a, b]  # original list order preserved