tests/test_memory_variants.py

"""Tests for the Part C memory/context variants and their wiring.

Covers the per-call-view middlewares (sliding window, observation truncation,
prompt compression), preset resolution + middleware assembly, the generalized
compaction gate, and the per-request retrieval-budget override. All offline:
no model client, no API keys, no vector DB.
"""

from __future__ import annotations

import unittest
from types import SimpleNamespace

import tiktoken
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage

from app.chat_service import (
    InContextHistoryRetrievalMiddleware,
    ObservationTruncationMiddleware,
    PromptCompressionMiddleware,
    SlidingWindowMiddleware,
    build_agent_middleware,
)
from app.chroma_rag import LocalChromaRetriever
from app.memory_presets import resolve_memory_preset
from app.telemetry import COMPACTION_SIGNAL_NAMES, pop_turn_signals, reset_turn_signals
from evals.common import COMPACTION_SIGNAL_KEYS, compaction_active


def make_request(messages: list, turn_id: str = "t1") -> SimpleNamespace:
    """A stand-in for langchain's ModelRequest with .messages/.runtime/.override."""
    runtime = SimpleNamespace(context=SimpleNamespace(kb_session_id=turn_id))
    req = SimpleNamespace(messages=messages, runtime=runtime)
    req.override = lambda messages=None, **_: make_request(
        req.messages if messages is None else messages, turn_id
    )
    return req


class SlidingWindowTests(unittest.TestCase):
    def test_no_trim_within_window(self) -> None:
        reset_turn_signals("t1")
        out = SlidingWindowMiddleware(keep=5)._trim(
            make_request([HumanMessage("q1"), AIMessage("a1")])
        )
        self.assertEqual(len(out.messages), 2)
        self.assertEqual(pop_turn_signals("t1"), {})

    def test_cut_lands_on_user_boundary(self) -> None:
        msgs = [
            HumanMessage("q1"),
            AIMessage("a1"),
            HumanMessage("q2"),
            AIMessage(""),
            ToolMessage(content="r", tool_call_id="c1"),
            AIMessage("a2"),
            HumanMessage("q3"),
        ]
        reset_turn_signals("t1")
        out = SlidingWindowMiddleware(keep=5)._trim(make_request(msgs))
        self.assertEqual(out.messages[0].type, "human")
        self.assertEqual(pop_turn_signals("t1")["dropped_messages"], 2)

    def test_advances_past_tool_to_avoid_orphan(self) -> None:
        # A naive last-3 cut would start on the tool result; the window must
        # advance to the next user message so no tool result is orphaned.
        msgs = [
            HumanMessage("q1"),
            AIMessage("a1"),
            HumanMessage("q2"),
            AIMessage(""),
            ToolMessage(content="r", tool_call_id="c1"),
            AIMessage("a2"),
            HumanMessage("q3"),
        ]
        reset_turn_signals("t1")
        out = SlidingWindowMiddleware(keep=3)._trim(make_request(msgs))
        self.assertNotEqual(out.messages[0].type, "tool")
        self.assertEqual(out.messages[0].type, "human")

    def test_trims_long_tool_turn_instead_of_no_op(self) -> None:
        # The current turn's tool loop is longer than keep, so the naive cut
        # lands inside it (no later user message). Must still drop the prior
        # turn and keep the current turn intact, not return the full list.
        msgs = [
            HumanMessage("old"),
            AIMessage("oldA"),  # prior turn -> droppable
            HumanMessage("current"),
            AIMessage(""),
            ToolMessage(content="r1", tool_call_id="c1"),
            AIMessage(""),
            ToolMessage(content="r2", tool_call_id="c2"),
            AIMessage("final"),
        ]
        reset_turn_signals("t1")
        out = SlidingWindowMiddleware(keep=3)._trim(make_request(msgs))
        self.assertEqual(out.messages[0].content, "current")  # current turn kept
        self.assertEqual(pop_turn_signals("t1")["dropped_messages"], 2)


class ObservationTruncationTests(unittest.TestCase):
    def _mw(self) -> ObservationTruncationMiddleware:
        return ObservationTruncationMiddleware(
            head_chars=10, tail_chars=5, trigger_chars=30
        )

    def test_truncates_large_tool_output(self) -> None:
        msgs = [HumanMessage("q"), ToolMessage(content="X" * 100, tool_call_id="c1")]
        reset_turn_signals("t1")
        out = self._mw()._truncate(make_request(msgs))
        tool_msg = out.messages[1]
        self.assertEqual(tool_msg.type, "tool")
        self.assertLess(len(tool_msg.content), 100)
        self.assertTrue(tool_msg.content.startswith("X" * 10))
        signals = pop_turn_signals("t1")
        self.assertEqual(signals["truncated_tool_outputs"], 1)
        self.assertGreater(signals["chars_saved"], 0)

    def test_small_output_and_non_tool_untouched(self) -> None:
        msgs = [HumanMessage("Y" * 100), ToolMessage(content="tiny", tool_call_id="c1")]
        reset_turn_signals("t1")
        out = self._mw()._truncate(make_request(msgs))
        self.assertEqual(out.messages[0].content, "Y" * 100)  # human untouched
        self.assertEqual(out.messages[1].content, "tiny")  # below trigger
        self.assertEqual(pop_turn_signals("t1"), {})

    def test_just_over_trigger_never_inflates(self) -> None:
        # Content barely over the trigger: the marker boilerplate would make the
        # "truncated" copy longer, so it must be left untouched (no inflation,
        # no negative chars_saved).
        mw = ObservationTruncationMiddleware(
            head_chars=20, tail_chars=5, trigger_chars=26
        )
        content = "Z" * 27  # > trigger, but head+tail+marker >> 27
        reset_turn_signals("t1")
        out = mw._truncate(
            make_request([ToolMessage(content=content, tool_call_id="c")])
        )
        self.assertEqual(out.messages[0].content, content)
        self.assertEqual(pop_turn_signals("t1"), {})


class PromptCompressionTests(unittest.TestCase):
    def test_collapses_whitespace_and_reports(self) -> None:
        reset_turn_signals("t1")
        out = PromptCompressionMiddleware()._compress(
            make_request([HumanMessage("a   b\n\n\n\nc   ")])
        )
        self.assertEqual(out.messages[0].content, "a b\n\nc")
        signals = pop_turn_signals("t1")
        self.assertEqual(signals["compressed_messages"], 1)
        self.assertGreater(signals["chars_saved"], 0)

    def test_already_compact_untouched(self) -> None:
        reset_turn_signals("t1")
        PromptCompressionMiddleware()._compress(make_request([HumanMessage("a b\nc")]))
        self.assertEqual(pop_turn_signals("t1"), {})


def _stub_embed(texts: list) -> list:
    """Deterministic embedder: blocks containing 'KEEPME' point one way."""
    return [[1.0, 0.0] if "KEEPME" in t else [0.0, 1.0] for t in texts]


class InContextHistoryRetrievalTests(unittest.TestCase):
    def _msgs(self) -> list:
        return [
            HumanMessage("blah one"),
            AIMessage("a1"),
            HumanMessage("the KEEPME fact is important"),
            AIMessage("a2"),
            HumanMessage("blah three"),
            AIMessage("a3"),
            HumanMessage("recall the KEEPME thing"),  # current turn
        ]

    def test_retrieves_relevant_older_block_drops_rest(self) -> None:
        mw = InContextHistoryRetrievalMiddleware(
            keep_recent=1, top_k=1, embed_fn=_stub_embed
        )
        reset_turn_signals("t1")
        out = mw._select(make_request(self._msgs()))
        contents = [m.content for m in out.messages]
        self.assertIn("the KEEPME fact is important", contents)  # retrieved
        self.assertIn("recall the KEEPME thing", contents)  # current turn
        self.assertNotIn("blah one", contents)  # irrelevant, dropped
        self.assertNotIn("blah three", contents)
        self.assertEqual(out.messages[0].type, "human")  # no orphaned tool/ai
        signals = pop_turn_signals("t1")
        self.assertEqual(signals["history_retrievals"], 1)
        self.assertGreater(signals["dropped_messages"], 0)

    def test_no_op_when_nothing_old_enough(self) -> None:
        mw = InContextHistoryRetrievalMiddleware(
            keep_recent=2, top_k=3, embed_fn=_stub_embed
        )
        reset_turn_signals("t1")
        msgs = [HumanMessage("q1"), AIMessage("a1"), HumanMessage("q2")]
        out = mw._select(make_request(msgs))
        self.assertEqual(len(out.messages), 3)
        self.assertEqual(pop_turn_signals("t1"), {})


class PresetResolutionTests(unittest.TestCase):
    NEW_PRESETS = (
        "observation_truncation",
        "sliding_window",
        "prompt_compression",
        "selective_retention",
        "context_reset",
        "clear_retrieval_kb",
        "incontext_history_retrieval",
    )

    def test_all_new_presets_resolve(self) -> None:
        for name in self.NEW_PRESETS:
            self.assertEqual(resolve_memory_preset(name).name, name)

    def test_axis_a_alternatives_disable_summarization(self) -> None:
        self.assertFalse(resolve_memory_preset("sliding_window").summarization)
        self.assertFalse(resolve_memory_preset("prompt_compression").summarization)

    def test_summary_prompt_variants_are_valid_templates(self) -> None:
        for name in ("selective_retention", "context_reset"):
            cfg = resolve_memory_preset(name)
            self.assertTrue(cfg.summarization)
            self.assertIsNotNone(cfg.summary_prompt)
            self.assertIn("{messages}", cfg.summary_prompt)
            # Only {messages} is a field, so .format must not raise.
            cfg.summary_prompt.format(messages="X")


class BuildMiddlewareTests(unittest.TestCase):
    def test_sliding_window_stack(self) -> None:
        mws = build_agent_middleware(
            model=None, memory_config=resolve_memory_preset("sliding_window")
        )
        names = {type(m).__name__ for m in mws}
        self.assertIn("SlidingWindowMiddleware", names)
        self.assertNotIn("SummarizationMiddleware", names)

    def test_prompt_compression_stack(self) -> None:
        mws = build_agent_middleware(
            model=None, memory_config=resolve_memory_preset("prompt_compression")
        )
        self.assertIn("PromptCompressionMiddleware", {type(m).__name__ for m in mws})

    def test_incontext_history_retrieval_stack(self) -> None:
        mws = build_agent_middleware(
            model=None,
            memory_config=resolve_memory_preset("incontext_history_retrieval"),
        )
        names = {type(m).__name__ for m in mws}
        self.assertIn("InContextHistoryRetrievalMiddleware", names)
        self.assertNotIn("SummarizationMiddleware", names)


class CompactionGateTests(unittest.TestCase):
    def test_recognizes_old_and_new_signals(self) -> None:
        self.assertTrue(compaction_active({"summary_messages": 2}))
        self.assertTrue(compaction_active({"cleared_tool_outputs": 1}))
        self.assertTrue(compaction_active({"dropped_messages": 3}))
        self.assertTrue(compaction_active({"truncated_tool_outputs": 1}))

    def test_no_signal_means_inactive(self) -> None:
        self.assertFalse(
            compaction_active({"summary_messages": 0, "dropped_messages": 0})
        )
        self.assertFalse(compaction_active({}))
        self.assertFalse(compaction_active(None))

    def test_app_and_eval_signal_lists_stay_in_sync(self) -> None:
        # Bidirectional: the per-call-view turn signals (app) must equal the eval
        # gate keys minus the two checkpoint-only markers. Pins both lists and
        # flags any phantom signal name that has no producer.
        checkpoint_markers = {"summary_messages", "cleared_tool_outputs"}
        self.assertEqual(
            set(COMPACTION_SIGNAL_NAMES),
            set(COMPACTION_SIGNAL_KEYS) - checkpoint_markers,
        )

    def test_signal_names_disjoint_from_reserved_stats(self) -> None:
        # Turn signals spread last into the context_stats event, so a name that
        # collided with a real metric would silently clobber it.
        reserved = {
            "input_tokens",
            "output_tokens",
            "total_tokens",
            "est_cost_usd",
            "llm_calls",
            "ttft_ms",
            "total_ms",
            "summary_messages",
            "cleared_tool_outputs",
            "context_messages",
            "context_tokens_approx",
        }
        self.assertEqual(set(COMPACTION_SIGNAL_NAMES) & reserved, set())


class RetrievalBudgetTests(unittest.TestCase):
    def _retriever(self, default_budget: int) -> LocalChromaRetriever:
        retriever = LocalChromaRetriever.__new__(LocalChromaRetriever)
        retriever._encoding = tiktoken.get_encoding("cl100k_base")
        retriever._token_budget = default_budget
        return retriever

    def test_override_caps_results_below_default(self) -> None:
        retriever = self._retriever(10_000)
        results = [SimpleNamespace(score=0.9, content="word " * 50) for _ in range(3)]
        self.assertEqual(len(retriever._apply_token_budget(results)), 3)
        self.assertEqual(
            len(retriever._apply_token_budget(results, token_budget=60)), 1
        )

    def test_low_score_filtered(self) -> None:
        retriever = self._retriever(10_000)
        results = [SimpleNamespace(score=0.05, content="x")]
        self.assertEqual(retriever._apply_token_budget(results), [])


if __name__ == "__main__":
    unittest.main()