"""``MemoryRanker``: pluggable re-ranker for memory candidates. 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 weight, access count, and decay are table-stakes for not returning 6-month-old "winners" when fresh signal exists. This module ships the first ranker — :class:`RecencyBoostRanker` — plus the :class:`MemoryRanker` protocol that future rankers (source-weight, access-count) plug into. The ranker is **pure**: ``rank(candidates) -> ranked_candidates``, no I/O, no state, no mutation of inputs. Same Rust-port shape as ``CompressionDecision`` and ``MemoryDecision``. Performance: O(N) over candidates where N = top_k. One ``math.exp()`` per candidate. Sub-microsecond per request — no embedding compute, no network, no disk. """ from __future__ import annotations from dataclasses import dataclass, field from datetime import datetime, timezone from typing import Protocol from headroom.proxy.memory_rank_policy import ( boost_memory_score, memory_recency_factor, parse_memory_created_at, ) # Use ``timezone.utc`` (always available) instead of ``datetime.UTC`` # (Python 3.11+) so this module imports cleanly on older interpreters. _UTC = timezone.utc @dataclass(frozen=True) class MemoryCandidate: """Immutable retrieval candidate as it flows through the ranker. The shape is the **proxy-side internal contract** — the backend's return type (typically ``MemoryResult`` with nested ``MemoryResult.memory.created_at``) is adapted into this flatter shape at the ranker boundary so the ranker stays backend-agnostic. ``score`` is the cosine similarity as returned by the backend. Rankers MAY mutate ``score`` by returning a new candidate with an updated score (frozen dataclass means they cannot mutate in place). """ content: str score: float created_at: datetime | None = None source: str | None = None # e.g. "memory_save" | "traffic_learner" | "inline" related_entities: tuple[str, ...] = field(default_factory=tuple) # Backend memory ID. Empty string when not preserved (test fixtures, # legacy callers). Rendered as ``[id]`` in the auto-tail block so the # model can pass it to memory_update / memory_delete directly. id: str = "" @classmethod def from_backend_result(cls, result: object) -> MemoryCandidate: """Adapter from backend ``MemoryResult`` shape to ``MemoryCandidate``. The backend returns objects with ``.score``, ``.memory.content``, ``.memory.id``, and (optionally) ``.memory.created_at`` (str ISO timestamp) + ``.related_entities``. This adapter flattens that to the ranker's expected shape and parses the timestamp to ``datetime``. Missing / unparseable timestamps → ``None`` (recency-neutral). Missing IDs → ``""`` (rendered as ``[?]`` in the auto-tail block). """ score = float(getattr(result, "score", 0.0)) memory = getattr(result, "memory", None) content = str(getattr(memory, "content", "")) if memory is not None else "" memory_id = str(getattr(memory, "id", "") or "") if memory is not None else "" raw_dt = getattr(memory, "created_at", None) if memory is not None else None created_at = parse_memory_created_at(raw_dt) raw_related = getattr(result, "related_entities", None) or () related = tuple(str(x) for x in raw_related) source_meta = getattr(memory, "metadata", None) or {} source = source_meta.get("source") if isinstance(source_meta, dict) else None return cls( content=content, score=score, created_at=created_at, source=source, related_entities=related, id=memory_id, ) def _parse_created_at(value: object) -> datetime | None: """Best-effort parse of a timestamp into a UTC-aware datetime. Accepts ``datetime`` (returned as-is, UTC-normalized) or ISO-8601 string (with or without trailing ``Z``). Anything else → ``None`` so the ranker treats the candidate as recency-neutral. """ return parse_memory_created_at(value) class MemoryRanker(Protocol): """Re-ranks retrieval candidates. Pure function. Implementations MUST: * Not mutate the input list or its elements * Be deterministic (same input → same output) for prefix-cache stability across consecutive turns * Be backend-agnostic — work with any :class:`MemoryCandidate`, regardless of which backend produced it """ def rank(self, candidates: list[MemoryCandidate]) -> list[MemoryCandidate]: ... @dataclass(frozen=True) class RecencyBoostRanker: """Re-ranker applying an exponential recency decay to cosine scores. Final score: ``cosine × exp(-age_days / decay_days)``. At ``decay_days=30``: * age = 0 days → factor 1.000 * age = 15 days → factor 0.607 * age = 30 days → factor 0.368 * age = 60 days → factor 0.135 * age = 90 days → factor 0.050 Tuned so a fresh memory with weak cosine doesn't dominate (factor decays gradually), but a 6-month-old strong-cosine can't dominate either (factor approaches zero). Operators tune ``decay_days`` for their codebase's rate of change — 7 days for rapidly-evolving repos, 90 days for stable archival. Backwards-compat: candidates with ``created_at=None`` get factor 1.0 — treated as recency-neutral. Lets a backend during a migration return some rows with timestamps and some without without breaking the ranker. Defensive: negative ages (clock skew on future-timestamped rows) are clamped to factor 1.0 — a clock-skewed candidate cannot outrank a real fresh one with the same cosine. """ decay_days: float = 30.0 def rank(self, candidates: list[MemoryCandidate]) -> list[MemoryCandidate]: """Re-rank by ``score × recency_factor``. Pure; input unchanged. Sorts descending by boosted score. Ties broken by input order (Python's sort is stable) — deterministic output for prefix- cache stability across turns. """ if not candidates: return [] now = datetime.now(_UTC) boosted: list[tuple[int, MemoryCandidate, float]] = [] for idx, c in enumerate(candidates): new_score = boost_memory_score( score=c.score, now=now, created_at=c.created_at, decay_days=self.decay_days, ) boosted.append((idx, c, new_score)) # Sort descending by boosted score; stable on ties via the # captured idx — same input order preserved on ties for # deterministic output across turns. boosted.sort(key=lambda triple: (-triple[2], triple[0])) # Return new candidates with the boosted score so downstream # consumers (logging, budget filter) see the post-boost number. return [ MemoryCandidate( content=c.content, score=new_score, created_at=c.created_at, source=c.source, related_entities=c.related_entities, id=c.id, ) for _, c, new_score in boosted ] def _recency_factor(self, now: datetime, created_at: datetime | None) -> float: """Compute the recency multiplier for a single candidate. ``None`` timestamp → 1.0 (recency-neutral, backwards-compat). Future timestamps → 1.0 (clock-skew defence). Otherwise: ``exp(-age_days / decay_days)``. """ return memory_recency_factor( now=now, created_at=created_at, decay_days=self.decay_days, )