mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
refactor(proxy): MemoryRanker + ImageCompressionDecision + branch-aware version-sync
Three independent contract-pattern follow-ons bundled into one PR. Same frozen-dataclass + factory + apply_to_tags + Rust-portable shape that PR #473 / #477 / #483 established. ## (1) MemoryRanker + 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. * ``MemoryRanker`` Protocol — pluggable re-ranker; future PRs add source-weight + access-count rankers behind the same interface. * ``RecencyBoostRanker`` — first concrete impl. Final score is ``cosine × exp(-age_days / decay_days)``. Default decay 30 days (half-life ~21 days; 60-day-old factor 0.135, 90-day-old 0.050). * ``MemoryCandidate`` — backend-agnostic frozen value type that flows through the ranker. ``MemoryCandidate.from_backend_result`` adapter converts the existing ``MemoryResult`` shape (with nested ``memory.created_at``) into the ranker's flatter form. * Wired into ``memory_handler.search_and_format_context`` as an optional ``ranker=`` kwarg — backwards-compat: ``None`` (default) preserves the pure-cosine path identically. Defensive: * ``created_at=None`` → factor 1.0 (recency-neutral, back-compat with legacy rows / migrating backends) * Negative age (clock skew) → clamped to factor 1.0 (a future-dated row can't outrank a real fresh memory) * Sort is stable on ties — same input → same output every turn, so consecutive turns inject memories in the same order (prefix-cache friendly) Performance: O(N) over candidates where N=top_k≈10. One ``math.exp`` per candidate. Sub-microsecond. Zero new I/O. ## (2) ImageCompressionDecision Mirror of :class:`CompressionDecision` for image compression. Two sites today (``openai.py:1203``, ``anthropic.py:868``) gate inline; both already respect bypass (no Gemini-class drift bug like text compression had), but consolidating into a value type: * Locks bypass-respect via AST contract test — future sites can't drift on it * Surfaces ``image_skip_reason`` in ``RequestOutcome.tags`` for dashboard slicing (same observability surface as ``passthrough_reason`` and ``memory_skip_reason``) * Same Rust-port shape as the other decision types Precedence: ``bypass_header`` > ``image_optimize_disabled`` > ``no_messages`` > ``should_compress=True``. Anthropic's extra ``is_cache_mode`` check stays inline because it's Anthropic-specific (openai/gemini don't have it). Documented in a code comment. ## (3) Branch-aware sync-plugin-versions hook Pre-this-fix the pre-commit ``sync-plugin-versions`` hook ran on every commit and bumped manifests to the predicted-next-release version. Every PR ended up carrying the prediction as collateral ("Why are we bumping ``.claude-plugin/marketplace.json`` — we should not, right??" - user, on PR #483). Fix: the hook is now a NO-OP unless EITHER: * We're on the ``main`` branch, OR * ``HEADROOM_SYNC_VERSIONS=1`` is set explicitly (release workflow) On feature branches the hook prints a single line explaining the skip and exits cleanly. The release workflow opts in via the env var; behaviour on main / at release time is unchanged. ## Test coverage * 16 new tests on ``MemoryRanker`` / ``RecencyBoostRanker`` (frozen, equal cosine wins by recency, decay configurable, NULL timestamp neutral, no-mutation contract, Rust-port shape) * 17 new tests on ``ImageCompressionDecision`` (frozen, all 3 skip reasons, precedence, observability fields, apply_to_tags) * 1 new AST invariant test (extends ``test_handler_outcome_tag_invariant.py``) — locks "no raw ``if self.config.image_optimize and messages and not _bypass:`` conjunction in any handler" All existing memory + cache-stability + handler tests pass (203 ✓). ``make ci-precheck`` clean. ## Rust portability All three new value types port cleanly to frozen Rust structs + pure functions. Same migration pattern as ``CompressionDecision`` (already locked in for the SmartCrusher Rust port). ## Zero-regression contract * Default ``ranker=None`` → memory_handler behaves identically to pre-this-PR (pure cosine; no perf change) * Image decision migration is identity at the bypass/optimize/messages gate — no behaviour change, just contract consolidation * Hook fix is no-op on feature branches (less churn) and unchanged on main (release flow preserved)
This commit is contained in:
parent
10580439bb
commit
a7b197c6ec
9 changed files with 982 additions and 24 deletions
|
|
@ -864,12 +864,18 @@ class AnthropicHandlerMixin:
|
|||
# turn becomes historical on the next request, so even "latest turn only"
|
||||
# rewrites can invalidate the next cache read when the client resends the
|
||||
# original transcript.
|
||||
if (
|
||||
self.config.image_optimize
|
||||
and messages
|
||||
and not _bypass
|
||||
and not is_cache_mode(self.config.mode)
|
||||
):
|
||||
#
|
||||
# Bypass / image_optimize / messages gating routes through
|
||||
# ImageCompressionDecision for uniformity with CompressionDecision +
|
||||
# MemoryDecision. The cache_mode check stays inline because it's
|
||||
# Anthropic-specific (sites in openai.py / gemini.py don't have it).
|
||||
from headroom.proxy.image_compression_decision import ImageCompressionDecision
|
||||
|
||||
_image_decision = ImageCompressionDecision.decide(
|
||||
headers=request.headers, config=self.config, messages=messages
|
||||
)
|
||||
_image_decision.apply_to_tags(tags)
|
||||
if _image_decision.should_compress and not is_cache_mode(self.config.mode):
|
||||
compressor = None
|
||||
try:
|
||||
compressor = _get_image_compressor()
|
||||
|
|
|
|||
|
|
@ -1200,8 +1200,19 @@ class OpenAIHandlerMixin:
|
|||
if _bypass:
|
||||
logger.info(f"[{request_id}] Bypass: skipping compression (header)")
|
||||
|
||||
# Image compression: tile alignment + ML-based technique routing
|
||||
if self.config.image_optimize and messages and not _bypass:
|
||||
# Image compression: tile alignment + ML-based technique routing.
|
||||
# Gated on ImageCompressionDecision — same value-type pattern
|
||||
# as CompressionDecision + MemoryDecision; locks bypass-respect
|
||||
# in tests so a future site can't drift.
|
||||
from headroom.proxy.image_compression_decision import ImageCompressionDecision
|
||||
|
||||
_image_decision = ImageCompressionDecision.decide(
|
||||
headers=request.headers, config=self.config, messages=messages
|
||||
)
|
||||
# tags is populated downstream at L1229 — defer apply_to_tags
|
||||
# to where the tags dict exists. The decision is captured here
|
||||
# so the conditional is uniform with the other gates.
|
||||
if _image_decision.should_compress:
|
||||
from headroom.proxy.helpers import _get_image_compressor
|
||||
|
||||
compressor = None
|
||||
|
|
@ -1229,6 +1240,10 @@ class OpenAIHandlerMixin:
|
|||
headers.pop("accept-encoding", None)
|
||||
tags = extract_tags(headers)
|
||||
client = classify_client(headers)
|
||||
# Surface the image-compression decision (computed earlier) into
|
||||
# tags now that the tags dict exists. Same observability pattern
|
||||
# the funnel uses for passthrough_reason + memory_skip_reason.
|
||||
_image_decision.apply_to_tags(tags)
|
||||
# PR-A5 (P5-49): strip internal x-headroom-* from upstream-bound
|
||||
# headers AFTER `_extract_tags` reads them. Inbound bypass gating
|
||||
# uses `request.headers.get(...)` above; memory user-id reads
|
||||
|
|
|
|||
125
headroom/proxy/image_compression_decision.py
Normal file
125
headroom/proxy/image_compression_decision.py
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
"""``ImageCompressionDecision``: canonical "should this request have
|
||||
images compressed?" gate.
|
||||
|
||||
Mirror of :class:`CompressionDecision` (text compression) and
|
||||
:class:`MemoryDecision`. Pre-this-PR image compression was gated at
|
||||
two sites (``openai.py:1203``, ``anthropic.py:868``) by inline
|
||||
conjunctions. Both already checked ``_bypass`` (no drift bug like
|
||||
the text-compression Gemini bypass-misses) — but consolidating into
|
||||
a value type still pays off:
|
||||
|
||||
* Locks the contract via tests so a future site can't drift
|
||||
* ``apply_to_tags()`` surfaces ``image_skip_reason`` to
|
||||
:class:`RequestOutcome.tags` for dashboard slicing
|
||||
* Rust-portable shape, same as the other decision types
|
||||
|
||||
Precedence (highest first):
|
||||
|
||||
1. ``bypass_header`` — user's explicit opt-out
|
||||
2. ``image_optimize_disabled`` — operator ``config.image_optimize=False``
|
||||
3. ``no_messages`` — empty / missing messages
|
||||
4. otherwise → ``should_compress=True``
|
||||
|
||||
Distinct from text :class:`CompressionDecision`'s
|
||||
``compression_disabled`` reason: operators can enable text + disable
|
||||
image (or vice versa) independently. Same shape, different gate.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from headroom.proxy.helpers import _headroom_bypass_enabled
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImageCompressionDecision:
|
||||
"""Immutable, value-equal snapshot of the image-compression gate.
|
||||
|
||||
Construction policy: use :meth:`decide`. Direct construction is
|
||||
allowed for tests but unusual — handlers always go through
|
||||
``decide``. The constituent observability booleans
|
||||
(``bypass_header_set`` etc.) MUST match the inputs ``decide`` saw;
|
||||
the factory enforces that invariant, and the dataclass being
|
||||
frozen means downstream code can't violate it.
|
||||
"""
|
||||
|
||||
should_compress: bool
|
||||
# When ``should_compress`` is False, the canonical reason surfaced
|
||||
# in ``RequestOutcome.tags["image_skip_reason"]`` so the dashboard
|
||||
# can slice image-skipped traffic by cause. One of:
|
||||
# * "bypass_header" — user set x-headroom-bypass/mode
|
||||
# * "image_optimize_disabled" — operator config off
|
||||
# * "no_messages" — empty / missing messages
|
||||
# When ``should_compress`` is True, this is None.
|
||||
passthrough_reason: str | None
|
||||
|
||||
# Observability: every constituent boolean exposed so debug tools
|
||||
# answer "what did the decision see?" without re-running it.
|
||||
bypass_header_set: bool
|
||||
image_optimize_enabled: bool
|
||||
has_messages: bool
|
||||
|
||||
@classmethod
|
||||
def decide(
|
||||
cls,
|
||||
*,
|
||||
headers: Any,
|
||||
config: Any,
|
||||
messages: Sequence[Any] | None,
|
||||
) -> ImageCompressionDecision:
|
||||
"""Compute the canonical image-compression decision.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
headers
|
||||
Inbound request headers. Accepts any object with a
|
||||
``.get(key)`` method (dict, starlette Headers, mapping).
|
||||
Bypass detected via ``_headroom_bypass_enabled``.
|
||||
config
|
||||
``HeadroomConfig``-shaped object; only ``image_optimize``
|
||||
is read.
|
||||
messages
|
||||
Request messages. ``None`` and ``[]`` are equivalent.
|
||||
"""
|
||||
bypass = _headroom_bypass_enabled(headers)
|
||||
image_ok = bool(getattr(config, "image_optimize", False))
|
||||
has_msgs = bool(messages)
|
||||
|
||||
if bypass:
|
||||
reason: str | None = "bypass_header"
|
||||
should = False
|
||||
elif not image_ok:
|
||||
reason = "image_optimize_disabled"
|
||||
should = False
|
||||
elif not has_msgs:
|
||||
reason = "no_messages"
|
||||
should = False
|
||||
else:
|
||||
reason = None
|
||||
should = True
|
||||
|
||||
return cls(
|
||||
should_compress=should,
|
||||
passthrough_reason=reason,
|
||||
bypass_header_set=bypass,
|
||||
image_optimize_enabled=image_ok,
|
||||
has_messages=has_msgs,
|
||||
)
|
||||
|
||||
def apply_to_tags(self, tags: dict[str, str]) -> None:
|
||||
"""Stamp the skip reason into a tags dict for dashboard slicing.
|
||||
|
||||
Mutates ``tags`` in place. No-op when ``should_compress=True``
|
||||
— absence vs presence is the signal.
|
||||
|
||||
Mirror of :meth:`CompressionDecision.apply_to_tags` and
|
||||
:meth:`MemoryDecision.apply_to_tags`. Multiple decision tags
|
||||
coexist in the same dict (``passthrough_reason``,
|
||||
``memory_skip_reason``, ``image_skip_reason``) for full
|
||||
dashboard slicing.
|
||||
"""
|
||||
if self.passthrough_reason is not None:
|
||||
tags["image_skip_reason"] = self.passthrough_reason
|
||||
|
|
@ -622,6 +622,8 @@ class MemoryHandler:
|
|||
user_id: str,
|
||||
messages: list[dict[str, Any]],
|
||||
request_context: RequestContext | None = None,
|
||||
*,
|
||||
ranker: Any | None = None,
|
||||
) -> str | None:
|
||||
"""Search memories and format as context injection.
|
||||
|
||||
|
|
@ -636,6 +638,14 @@ class MemoryHandler:
|
|||
omitted, behaves as before this fix — single-bucket search
|
||||
against the legacy backend. Production handlers always
|
||||
pass it; tests / mocks can keep the simpler call shape.
|
||||
ranker: Optional :class:`~headroom.proxy.memory_ranker.MemoryRanker`
|
||||
— re-ranks the backend's cosine-only candidates by an
|
||||
additional signal (recency, source, access count, …).
|
||||
When ``None`` (default), behaviour is identical to pre-
|
||||
this-PR: pure cosine + min_similarity floor. When
|
||||
provided, candidates are adapted to
|
||||
:class:`MemoryCandidate`, re-ranked, then re-filtered by
|
||||
``min_similarity`` on the boosted score.
|
||||
|
||||
Returns:
|
||||
Formatted context string, or None if no relevant memories.
|
||||
|
|
@ -689,23 +699,48 @@ class MemoryHandler:
|
|||
)
|
||||
return None
|
||||
|
||||
# Filter by minimum similarity
|
||||
filtered_results = [r for r in results if r.score >= self.config.min_similarity]
|
||||
# Optional re-rank: when a MemoryRanker is provided, adapt
|
||||
# results to MemoryCandidate, re-rank, then filter by
|
||||
# min_similarity on the BOOSTED score. The re-rank can
|
||||
# promote a fresh weak-cosine memory above a stale strong-
|
||||
# cosine one (RecencyBoostRanker default behaviour).
|
||||
if ranker is not None:
|
||||
from headroom.proxy.memory_ranker import MemoryCandidate
|
||||
|
||||
if not filtered_results:
|
||||
logger.debug(
|
||||
f"Memory: {len(results)} memories found but none above threshold "
|
||||
f"{self.config.min_similarity}"
|
||||
)
|
||||
return None
|
||||
candidates = [MemoryCandidate.from_backend_result(r) for r in results]
|
||||
ranked = ranker.rank(candidates)
|
||||
# Filter on the post-rank score (the ranker may have
|
||||
# boosted or attenuated original cosine values).
|
||||
ranked = [c for c in ranked if c.score >= self.config.min_similarity]
|
||||
if not ranked:
|
||||
logger.debug(
|
||||
f"Memory: {len(results)} memories found but none above threshold "
|
||||
f"{self.config.min_similarity} after re-rank"
|
||||
)
|
||||
return None
|
||||
memory_lines = []
|
||||
for i, candidate in enumerate(ranked, 1):
|
||||
memory_lines.append(f"{i}. {candidate.content}")
|
||||
if candidate.related_entities:
|
||||
entities_str = ", ".join(candidate.related_entities[:3])
|
||||
memory_lines.append(f" (Related: {entities_str})")
|
||||
else:
|
||||
# Pre-PR-this behaviour: pure cosine, no boost.
|
||||
filtered_results = [r for r in results if r.score >= self.config.min_similarity]
|
||||
|
||||
# Format as context
|
||||
memory_lines = []
|
||||
for i, result in enumerate(filtered_results, 1):
|
||||
memory_lines.append(f"{i}. {result.memory.content}")
|
||||
if hasattr(result, "related_entities") and result.related_entities:
|
||||
entities_str = ", ".join(result.related_entities[:3])
|
||||
memory_lines.append(f" (Related: {entities_str})")
|
||||
if not filtered_results:
|
||||
logger.debug(
|
||||
f"Memory: {len(results)} memories found but none above threshold "
|
||||
f"{self.config.min_similarity}"
|
||||
)
|
||||
return None
|
||||
|
||||
memory_lines = []
|
||||
for i, result in enumerate(filtered_results, 1):
|
||||
memory_lines.append(f"{i}. {result.memory.content}")
|
||||
if hasattr(result, "related_entities") and result.related_entities:
|
||||
entities_str = ", ".join(result.related_entities[:3])
|
||||
memory_lines.append(f" (Related: {entities_str})")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Memory: Search failed for user {effective_user_id}: {e}")
|
||||
|
|
|
|||
198
headroom/proxy/memory_ranker.py
Normal file
198
headroom/proxy/memory_ranker.py
Normal file
|
|
@ -0,0 +1,198 @@
|
|||
"""``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
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Protocol
|
||||
|
||||
# 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)
|
||||
|
||||
@classmethod
|
||||
def from_backend_result(cls, result: object) -> MemoryCandidate:
|
||||
"""Adapter from backend ``MemoryResult`` shape to ``MemoryCandidate``.
|
||||
|
||||
The backend returns objects with ``.score``, ``.memory.content``,
|
||||
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).
|
||||
"""
|
||||
score = float(getattr(result, "score", 0.0))
|
||||
memory = getattr(result, "memory", None)
|
||||
content = str(getattr(memory, "content", "")) if memory is not None else ""
|
||||
raw_dt = getattr(memory, "created_at", None) if memory is not None else None
|
||||
created_at = _parse_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,
|
||||
)
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, datetime):
|
||||
return value if value.tzinfo is not None else value.replace(tzinfo=_UTC)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
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):
|
||||
factor = self._recency_factor(now, c.created_at)
|
||||
boosted.append((idx, c, c.score * factor))
|
||||
|
||||
# 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,
|
||||
)
|
||||
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)``.
|
||||
"""
|
||||
if created_at is None:
|
||||
return 1.0
|
||||
# Normalize to UTC-aware for safe subtraction.
|
||||
if created_at.tzinfo is None:
|
||||
created_at = created_at.replace(tzinfo=_UTC)
|
||||
delta = now - created_at
|
||||
age_days = delta.total_seconds() / 86400.0
|
||||
if age_days <= 0:
|
||||
# Future-dated row (clock skew). Clamp to neutral.
|
||||
return 1.0
|
||||
return math.exp(-age_days / self.decay_days)
|
||||
|
|
@ -1,7 +1,22 @@
|
|||
"""Sync plugin manifest versions to the repo's computed release semver."""
|
||||
"""Sync plugin manifest versions to the repo's computed release semver.
|
||||
|
||||
Branch-aware: by default this script is a NO-OP on feature branches.
|
||||
Pre-this-fix it ran on every commit and bumped the manifests to the
|
||||
PREDICTED next release version, which polluted every PR with version-
|
||||
bump noise (the prediction advanced as commits landed; each PR ended
|
||||
up carrying the bump as collateral).
|
||||
|
||||
Sync now only runs when EITHER:
|
||||
* We're on the ``main`` branch, OR
|
||||
* ``HEADROOM_SYNC_VERSIONS=1`` is set explicitly (release workflow)
|
||||
|
||||
Result: feature-branch PRs no longer carry manifest bumps; the
|
||||
release workflow still gets a canonical sync at publish time.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
|
@ -33,8 +48,50 @@ def compute_repo_semver(root: Path) -> str:
|
|||
return info.npm_version
|
||||
|
||||
|
||||
def _current_branch(root: Path) -> str | None:
|
||||
"""Return the current git branch name, or None if git isn't usable."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "rev-parse", "--abbrev-ref", "HEAD"],
|
||||
cwd=root,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
except (FileNotFoundError, OSError):
|
||||
return None
|
||||
if result.returncode != 0:
|
||||
return None
|
||||
return result.stdout.strip() or None
|
||||
|
||||
|
||||
def _should_sync(root: Path) -> bool:
|
||||
"""Decide whether to actually run the sync.
|
||||
|
||||
Release workflow opts in via ``HEADROOM_SYNC_VERSIONS=1``; otherwise
|
||||
we only sync on ``main`` (where the next-release prediction
|
||||
legitimately lives). On feature branches we no-op — the prediction
|
||||
would just create PR-level noise.
|
||||
"""
|
||||
if os.environ.get("HEADROOM_SYNC_VERSIONS") == "1":
|
||||
return True
|
||||
branch = _current_branch(root)
|
||||
if branch is None:
|
||||
# Git unavailable or detached HEAD — safest default is no-op.
|
||||
return False
|
||||
return branch == "main"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
root = ROOT
|
||||
if not _should_sync(root):
|
||||
# Quiet no-op on feature branches. Print a single line so
|
||||
# pre-commit users see the reason if they look.
|
||||
branch = _current_branch(root) or "<unknown>"
|
||||
print(
|
||||
f"sync-plugin-versions: skipping on branch '{branch}' (set HEADROOM_SYNC_VERSIONS=1 to force)"
|
||||
)
|
||||
return
|
||||
version = compute_repo_semver(root)
|
||||
subprocess.run(
|
||||
[
|
||||
|
|
|
|||
|
|
@ -116,3 +116,41 @@ def test_outcome_call_sites_pass_client_kwarg() -> None:
|
|||
"`client = classify_client(headers)` and thread `client=client` "
|
||||
"into the outcome construction. See PR #473 for the pattern."
|
||||
)
|
||||
|
||||
|
||||
# ── Invariant: image-compression must route through ImageCompressionDecision ──
|
||||
|
||||
|
||||
import re # noqa: E402 -- only used by the image-decision invariant below
|
||||
|
||||
|
||||
def test_no_raw_image_optimize_gate_in_handlers() -> None:
|
||||
"""Locks the post-this-PR contract: image compression must be
|
||||
gated by :class:`ImageCompressionDecision`, not by an inline
|
||||
``if self.config.image_optimize and messages and not _bypass:``
|
||||
conjunction. Pre-PR-this both sites used the raw conjunction;
|
||||
consolidating into a value type means a future site (e.g., new
|
||||
provider handler) can't drift on bypass-respect or skip-reason
|
||||
observability.
|
||||
|
||||
Allowed forms after this PR:
|
||||
* ``if _image_decision.should_compress``
|
||||
* ``if _image_decision.should_compress and ...``
|
||||
"""
|
||||
pattern = re.compile(r"^\s*if\s*\(?\s*self\.config\.image_optimize\s+and\s+messages\b")
|
||||
offenders: list[tuple[str, int, str]] = []
|
||||
for f in HANDLER_FILES:
|
||||
text = f.read_text(encoding="utf-8")
|
||||
for i, line in enumerate(text.splitlines(), start=1):
|
||||
if pattern.match(line):
|
||||
offenders.append((f.name, i, line.rstrip()))
|
||||
if offenders:
|
||||
formatted = "\n".join(f" {f}:{ln} {src!r}" for f, ln, src in offenders)
|
||||
pytest.fail(
|
||||
f"{len(offenders)} handler site(s) use the pre-PR raw image "
|
||||
"gate `if self.config.image_optimize and messages [and ...]`:\n"
|
||||
f"{formatted}\n\n"
|
||||
"Replace with `ImageCompressionDecision.decide(...)` + "
|
||||
"`if _image_decision.should_compress:`. See "
|
||||
"headroom/proxy/image_compression_decision.py for the pattern."
|
||||
)
|
||||
|
|
|
|||
241
tests/test_image_compression_decision.py
Normal file
241
tests/test_image_compression_decision.py
Normal file
|
|
@ -0,0 +1,241 @@
|
|||
"""Tests for :class:`headroom.proxy.image_compression_decision.ImageCompressionDecision`.
|
||||
|
||||
Image compression today is gated at two sites (``openai.py:1203`` +
|
||||
``anthropic.py:868``) by inline conjunctions. Both already check
|
||||
``_bypass`` (the drift problem CompressionDecision fixed is NOT
|
||||
present here), but consolidating into a value type still pays off:
|
||||
|
||||
* test-lockable contract (no future site can forget bypass)
|
||||
* ``apply_to_tags()`` surfaces ``image_skip_reason`` in
|
||||
``RequestOutcome.tags`` — dashboards can slice image-skipped
|
||||
traffic by cause (same observability surface as
|
||||
``passthrough_reason`` and ``memory_skip_reason``)
|
||||
* Rust-portable shape mirrors :class:`CompressionDecision` exactly
|
||||
|
||||
Precedence (highest first):
|
||||
1. ``bypass_header`` — user opt-out
|
||||
2. ``image_optimize_disabled`` — operator config off
|
||||
3. ``no_messages`` — nothing to inspect
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import FrozenInstanceError
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
from headroom.proxy.image_compression_decision import ImageCompressionDecision
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _config(*, image_optimize: bool = True) -> Any:
|
||||
"""Minimal stand-in for ``HeadroomConfig`` — only the field the
|
||||
decision reads."""
|
||||
return SimpleNamespace(image_optimize=image_optimize)
|
||||
|
||||
|
||||
def _msgs(n: int = 1) -> list[dict[str, str]]:
|
||||
return [{"role": "user", "content": f"hi-{i}"} for i in range(n)]
|
||||
|
||||
|
||||
# ── Value-type contract ───────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_decision_is_frozen() -> None:
|
||||
d = ImageCompressionDecision.decide(headers={}, config=_config(), messages=_msgs())
|
||||
try:
|
||||
d.should_compress = False # type: ignore[misc]
|
||||
except FrozenInstanceError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError("ImageCompressionDecision must be frozen")
|
||||
|
||||
|
||||
def test_decision_is_value_equal() -> None:
|
||||
a = ImageCompressionDecision.decide(headers={}, config=_config(), messages=_msgs())
|
||||
b = ImageCompressionDecision.decide(headers={}, config=_config(), messages=_msgs())
|
||||
assert a == b
|
||||
|
||||
|
||||
# ── Precedence ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_compresses_when_every_gate_open() -> None:
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={}, config=_config(image_optimize=True), messages=_msgs()
|
||||
)
|
||||
assert d.should_compress is True
|
||||
assert d.passthrough_reason is None
|
||||
|
||||
|
||||
def test_bypass_header_wins() -> None:
|
||||
"""Bypass is the user's explicit "don't touch my bytes" signal —
|
||||
image compression mutates bytes (tile-aligns / re-encodes), so
|
||||
bypass must skip it. Mirror of CompressionDecision."""
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={"x-headroom-bypass": "true"},
|
||||
config=_config(image_optimize=True),
|
||||
messages=_msgs(),
|
||||
)
|
||||
assert d.should_compress is False
|
||||
assert d.passthrough_reason == "bypass_header"
|
||||
|
||||
|
||||
def test_passthrough_mode_header_also_triggers_bypass_skip() -> None:
|
||||
"""``x-headroom-mode: passthrough`` alt spelling — mirrors
|
||||
``_headroom_bypass_enabled`` semantics across all proxy gates."""
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={"x-headroom-mode": "passthrough"},
|
||||
config=_config(image_optimize=True),
|
||||
messages=_msgs(),
|
||||
)
|
||||
assert d.should_compress is False
|
||||
assert d.passthrough_reason == "bypass_header"
|
||||
|
||||
|
||||
def test_image_optimize_disabled_is_skip() -> None:
|
||||
"""Operator config — ``config.image_optimize = False``. Distinct
|
||||
reason from ``compression_disabled`` (text compression's gate);
|
||||
operators can enable text + disable image independently."""
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={}, config=_config(image_optimize=False), messages=_msgs()
|
||||
)
|
||||
assert d.should_compress is False
|
||||
assert d.passthrough_reason == "image_optimize_disabled"
|
||||
|
||||
|
||||
def test_no_messages_is_skip() -> None:
|
||||
"""Empty or missing messages — nothing to look at. Same shape as
|
||||
CompressionDecision's no_messages reason."""
|
||||
d = ImageCompressionDecision.decide(headers={}, config=_config(), messages=[])
|
||||
assert d.should_compress is False
|
||||
assert d.passthrough_reason == "no_messages"
|
||||
|
||||
|
||||
def test_messages_none_is_skip() -> None:
|
||||
"""``messages=None`` is treated identically to empty list."""
|
||||
d = ImageCompressionDecision.decide(headers={}, config=_config(), messages=None)
|
||||
assert d.should_compress is False
|
||||
assert d.passthrough_reason == "no_messages"
|
||||
|
||||
|
||||
# ── Precedence ordering ──────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_bypass_beats_image_optimize_disabled() -> None:
|
||||
"""User signal beats operator signal — bypass is the more
|
||||
informative dashboard slice."""
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={"x-headroom-bypass": "true"},
|
||||
config=_config(image_optimize=False),
|
||||
messages=_msgs(),
|
||||
)
|
||||
assert d.passthrough_reason == "bypass_header"
|
||||
|
||||
|
||||
def test_bypass_beats_no_messages() -> None:
|
||||
"""Bypass+no-messages surfaces bypass — user opted out, the
|
||||
empty body is incidental."""
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={"x-headroom-bypass": "true"},
|
||||
config=_config(),
|
||||
messages=[],
|
||||
)
|
||||
assert d.passthrough_reason == "bypass_header"
|
||||
|
||||
|
||||
def test_image_optimize_disabled_beats_no_messages() -> None:
|
||||
"""When config is off AND messages empty, surface the operator
|
||||
decision — more meaningful for dashboards."""
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={}, config=_config(image_optimize=False), messages=[]
|
||||
)
|
||||
assert d.passthrough_reason == "image_optimize_disabled"
|
||||
|
||||
|
||||
# ── Observability fields ─────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_observability_booleans_populated_when_compressing() -> None:
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={}, config=_config(image_optimize=True), messages=_msgs(2)
|
||||
)
|
||||
assert d.bypass_header_set is False
|
||||
assert d.image_optimize_enabled is True
|
||||
assert d.has_messages is True
|
||||
|
||||
|
||||
def test_observability_booleans_populated_when_passthrough() -> None:
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={"x-headroom-bypass": "true"},
|
||||
config=_config(image_optimize=True),
|
||||
messages=_msgs(),
|
||||
)
|
||||
assert d.bypass_header_set is True
|
||||
assert d.image_optimize_enabled is True
|
||||
assert d.has_messages is True
|
||||
|
||||
|
||||
# ── apply_to_tags ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_apply_to_tags_stamps_reason_when_passthrough() -> None:
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={"x-headroom-bypass": "true"}, config=_config(), messages=_msgs()
|
||||
)
|
||||
tags: dict[str, str] = {}
|
||||
d.apply_to_tags(tags)
|
||||
assert tags == {"image_skip_reason": "bypass_header"}
|
||||
|
||||
|
||||
def test_apply_to_tags_is_a_noop_when_compressing() -> None:
|
||||
d = ImageCompressionDecision.decide(headers={}, config=_config(), messages=_msgs())
|
||||
tags: dict[str, str] = {"client": "codex"}
|
||||
d.apply_to_tags(tags)
|
||||
assert tags == {"client": "codex"}
|
||||
|
||||
|
||||
def test_apply_to_tags_preserves_pre_existing_entries() -> None:
|
||||
"""Image skip reason coexists with other slicing tags (client,
|
||||
passthrough_reason, memory_skip_reason) — they all live in the
|
||||
same RequestOutcome.tags dict."""
|
||||
d = ImageCompressionDecision.decide(
|
||||
headers={}, config=_config(image_optimize=False), messages=_msgs()
|
||||
)
|
||||
tags: dict[str, str] = {
|
||||
"client": "claude-code",
|
||||
"passthrough_reason": "compression_disabled",
|
||||
"memory_skip_reason": "no_user_id",
|
||||
}
|
||||
d.apply_to_tags(tags)
|
||||
assert tags["client"] == "claude-code"
|
||||
assert tags["passthrough_reason"] == "compression_disabled"
|
||||
assert tags["memory_skip_reason"] == "no_user_id"
|
||||
assert tags["image_skip_reason"] == "image_optimize_disabled"
|
||||
|
||||
|
||||
def test_apply_to_tags_for_every_skip_reason() -> None:
|
||||
"""Every reason name round-trips cleanly into the tag dict."""
|
||||
cases: dict[str, dict[str, Any]] = {
|
||||
"bypass_header": {
|
||||
"headers": {"x-headroom-bypass": "true"},
|
||||
"config": _config(),
|
||||
"messages": _msgs(),
|
||||
},
|
||||
"image_optimize_disabled": {
|
||||
"headers": {},
|
||||
"config": _config(image_optimize=False),
|
||||
"messages": _msgs(),
|
||||
},
|
||||
"no_messages": {
|
||||
"headers": {},
|
||||
"config": _config(),
|
||||
"messages": [],
|
||||
},
|
||||
}
|
||||
for expected_reason, kwargs in cases.items():
|
||||
d = ImageCompressionDecision.decide(**kwargs)
|
||||
tags: dict[str, str] = {}
|
||||
d.apply_to_tags(tags)
|
||||
assert tags.get("image_skip_reason") == expected_reason, expected_reason
|
||||
243
tests/test_memory_ranker.py
Normal file
243
tests/test_memory_ranker.py
Normal file
|
|
@ -0,0 +1,243 @@
|
|||
"""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")
|
||||
|
||||
|
||||
# ── 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue