chopratejas--headroom/tests/test_memory_injection_budget.py

"""Tests for :class:`headroom.proxy.memory_injection.MemoryInjectionBudget`.

Pre-PR Headroom had **no token cap on injected memory**: top_k=10
candidates × ~400 tokens each = up to ~4000 tokens injected per
request. None of Letta/Mem0/Cognee/Supermemory ship a token-uncapped
injection path on the hot wire.

``MemoryInjectionBudget`` is the single configurable cap applied to
every injection site so all 5 sites are uniformly bounded — set the
budget once, apply at every handler, dashboards see the same shape.
"""

from __future__ import annotations

from dataclasses import FrozenInstanceError

from headroom.proxy.memory_injection import MemoryInjectionBudget

# ── Value-type contract ───────────────────────────────────────────────


def test_budget_is_frozen() -> None:
    b = MemoryInjectionBudget()
    try:
        b.max_tokens = 99  # type: ignore[misc]
    except FrozenInstanceError:
        pass
    else:
        raise AssertionError("MemoryInjectionBudget must be frozen")


def test_budget_defaults() -> None:
    """Default budget is conservative — 1024 tokens, 10 entries, 0.3
    similarity floor. Operators can override via config; the default
    is hard-set so a misconfiguration can't accidentally unbound
    injection."""
    b = MemoryInjectionBudget()
    assert b.max_tokens == 1024
    assert b.max_entries == 10
    assert b.min_similarity == 0.3


def test_budget_value_equal() -> None:
    a = MemoryInjectionBudget(max_tokens=512, max_entries=5, min_similarity=0.5)
    b = MemoryInjectionBudget(max_tokens=512, max_entries=5, min_similarity=0.5)
    assert a == b


# ── apply_to_text — bounding the formatted context block ─────────────


def test_apply_to_text_returns_input_when_under_budget() -> None:
    """Short context passes through unchanged — no spurious mutation."""
    b = MemoryInjectionBudget(max_tokens=1024)
    text = "## Relevant Memories\n1. small fact\n"
    out = b.apply_to_text(text)
    assert out == text


def test_apply_to_text_truncates_when_over_budget() -> None:
    """Large context is bounded — truncated at the budget. The
    truncation here is on the OUTPUT (the formatted injection block),
    NOT on the INPUT (which keeps full fidelity per MemoryQuery
    contract)."""
    # 4 tokens/char heuristic in our cap — make the input clearly
    # over even the most generous budget.
    b = MemoryInjectionBudget(max_tokens=128)  # ~512 chars at 4 char/token
    huge = "x" * 100000
    out = b.apply_to_text(huge)
    # Output should be substantially smaller than input.
    assert len(out) < len(huge)


def test_apply_to_text_preserves_full_lines() -> None:
    """When truncating, prefer cutting at line boundaries so the
    dashboard renders intact memory entries (no half-truncated bullet
    point)."""
    b = MemoryInjectionBudget(max_tokens=64)  # very tight
    text = "## Relevant Memories\n" + "".join(f"{i}. fact {i}\n" for i in range(100))
    out = b.apply_to_text(text)
    # No partial last line — every retained line ends in newline or is
    # the final line.
    if out and not out.endswith("\n"):
        # The last char is the closing of the final line; it must not
        # be in the middle of "fact " — easy heuristic: must not end
        # mid-word with a hanging digit-then-period.
        assert ". fact" not in out[-15:] or out.rstrip().endswith(("fact 0", "fact 1", "fact 2"))


def test_apply_to_text_handles_empty_input() -> None:
    """Empty input → empty output."""
    assert MemoryInjectionBudget().apply_to_text("") == ""


# ── apply_to_entries — bounding the list before formatting ───────────


def test_apply_to_entries_caps_entry_count() -> None:
    """Even if the backend returns 100 candidates, the budget caps
    entry count to ``max_entries``."""
    b = MemoryInjectionBudget(max_entries=3)
    entries = [{"content": f"entry {i}", "score": 0.9 - i * 0.01} for i in range(20)]
    out = b.apply_to_entries(entries)
    assert len(out) == 3


def test_apply_to_entries_preserves_order_of_input() -> None:
    """Budget doesn't re-rank — the backend's order is preserved. (The
    backend should already have ranked by score; budget just caps.)"""
    b = MemoryInjectionBudget(max_entries=2)
    entries = [
        {"content": "alpha", "score": 0.9},
        {"content": "beta", "score": 0.8},
        {"content": "gamma", "score": 0.7},
    ]
    out = b.apply_to_entries(entries)
    assert [e["content"] for e in out] == ["alpha", "beta"]


def test_apply_to_entries_filters_below_min_similarity() -> None:
    """Entries below ``min_similarity`` are dropped, regardless of
    entry-count budget remaining."""
    b = MemoryInjectionBudget(max_entries=10, min_similarity=0.5)
    entries = [
        {"content": "kept", "score": 0.9},
        {"content": "dropped", "score": 0.3},
        {"content": "kept2", "score": 0.55},
    ]
    out = b.apply_to_entries(entries)
    assert {e["content"] for e in out} == {"kept", "kept2"}