tests/test_memory_presets.py

from __future__ import annotations

import types
import unittest
from unittest.mock import patch

from app.chat_service import (
    DeepSeekCacheIsolationMiddleware,
    InstrumentedSummarizationMiddleware,
    PrefixPreservingCompactionMiddleware,
    SourcePreferenceMiddleware,
    StableToolOutputCapMiddleware,
    StudentProfileMiddleware,
    build_agent,
    build_agent_middleware,
)
from app.memory_presets import (
    DEFAULT_MEMORY_PRESET,
    MEMORY_PRESETS,
    resolve_memory_preset,
)
from langchain.agents.middleware import (
    ContextEditingMiddleware,
    SummarizationMiddleware,
)


def _stub_model():
    # SummarizationMiddleware reads model._llm_type at init.
    return types.SimpleNamespace(_llm_type="fake-chat-model")


class MemoryPresetResolutionTests(unittest.TestCase):
    def test_prod_preset_matches_production_constants(self) -> None:
        prod = MEMORY_PRESETS["prod"]
        self.assertTrue(prod.summarization)
        self.assertEqual(prod.summarization_trigger_tokens, 30_000)
        self.assertEqual(prod.summarization_keep_messages, 20)
        self.assertTrue(prod.context_editing)
        self.assertEqual(prod.context_editing_trigger_tokens, 5_000)
        self.assertEqual(prod.context_editing_keep, 5)
        self.assertFalse(prod.longterm_memory)

    def test_default_resolution(self) -> None:
        with patch.dict("os.environ", {}, clear=False):
            import os

            os.environ.pop("AI_TUTOR_MEMORY_PRESET", None)
            self.assertEqual(resolve_memory_preset("").name, DEFAULT_MEMORY_PRESET)
            self.assertEqual(resolve_memory_preset(None).name, DEFAULT_MEMORY_PRESET)

    def test_env_var_default_and_explicit_override(self) -> None:
        with patch.dict("os.environ", {"AI_TUTOR_MEMORY_PRESET": "full_history"}):
            self.assertEqual(resolve_memory_preset("").name, "full_history")
            # An explicit request value beats the env default.
            self.assertEqual(resolve_memory_preset("aggressive").name, "aggressive")

    def test_unknown_preset_raises(self) -> None:
        with self.assertRaises(ValueError):
            resolve_memory_preset("does_not_exist")
        with patch.dict("os.environ", {"AI_TUTOR_MEMORY_PRESET": "typo"}):
            with self.assertRaises(ValueError):
                resolve_memory_preset("")

    def test_deepseek_experiment_arms_are_single_axis_configs(self) -> None:
        raw = MEMORY_PRESETS["exp_fh_raw"]
        capped = MEMORY_PRESETS["exp_fh_cap10k"]
        compact = MEMORY_PRESETS["exp_c200_cap10k"]
        self.assertTrue(raw.experiment_mode)
        self.assertFalse(raw.summarization)
        self.assertFalse(raw.context_editing)
        self.assertEqual(capped.tool_output_cap_bytes, 40_000)
        self.assertEqual(compact.summarization_trigger_tokens, 200_000)
        self.assertEqual(compact.summarization_keep_tokens, 50_000)
        self.assertIsNone(compact.summarization_trim_tokens)
        self.assertEqual(compact.summarization_input_guard_tokens, 900_000)
        self.assertEqual(compact.experiment_request_guard_tokens, 990_000)
        self.assertFalse(compact.context_editing)
        structured = MEMORY_PRESETS["exp_c200_cap10k_structured"]
        self.assertEqual(structured.summarization_strategy, "structured_prefix")
        self.assertEqual(structured.summarization_trigger_tokens, 200_000)
        self.assertEqual(structured.summarization_keep_tokens, 50_000)
        self.assertEqual(structured.tool_output_cap_bytes, 40_000)


class MiddlewareAssemblyTests(unittest.TestCase):
    def test_full_history_disables_compaction(self) -> None:
        middleware = build_agent_middleware(
            _stub_model(), MEMORY_PRESETS["full_history"]
        )
        self.assertEqual(len(middleware), 1)
        self.assertIsInstance(middleware[0], SourcePreferenceMiddleware)

    def test_prod_has_editing_then_summarization(self) -> None:
        middleware = build_agent_middleware(_stub_model(), MEMORY_PRESETS["prod"])
        self.assertIsInstance(middleware[0], ContextEditingMiddleware)
        self.assertIsInstance(middleware[1], SummarizationMiddleware)
        self.assertIsInstance(middleware[-1], SourcePreferenceMiddleware)
        self.assertEqual(len(middleware), 3)

    def test_single_technique_presets(self) -> None:
        summarization_only = build_agent_middleware(
            _stub_model(), MEMORY_PRESETS["summarization_only"]
        )
        self.assertFalse(
            any(isinstance(m, ContextEditingMiddleware) for m in summarization_only)
        )
        self.assertTrue(
            any(isinstance(m, SummarizationMiddleware) for m in summarization_only)
        )
        editing_only = build_agent_middleware(
            _stub_model(), MEMORY_PRESETS["editing_only"]
        )
        self.assertTrue(
            any(isinstance(m, ContextEditingMiddleware) for m in editing_only)
        )
        self.assertFalse(
            any(isinstance(m, SummarizationMiddleware) for m in editing_only)
        )

    def test_profile_memory_adds_student_profile_middleware(self) -> None:
        middleware = build_agent_middleware(
            _stub_model(), MEMORY_PRESETS["profile_memory"]
        )
        self.assertTrue(
            any(isinstance(m, StudentProfileMiddleware) for m in middleware)
        )

    def test_experiment_stack_has_isolation_cap_and_full_input_summary(self) -> None:
        middleware = build_agent_middleware(
            _stub_model(), MEMORY_PRESETS["exp_c200_cap10k"]
        )
        self.assertTrue(
            any(isinstance(m, DeepSeekCacheIsolationMiddleware) for m in middleware)
        )
        self.assertTrue(
            any(isinstance(m, StableToolOutputCapMiddleware) for m in middleware)
        )
        summary = next(
            m for m in middleware if isinstance(m, InstrumentedSummarizationMiddleware)
        )
        self.assertEqual(summary.keep, ("tokens", 50_000))
        self.assertIsNone(summary.trim_tokens_to_summarize)
        self.assertEqual(summary.summary_input_guard_tokens, 900_000)
        self.assertFalse(
            any(isinstance(m, ContextEditingMiddleware) for m in middleware)
        )

    def test_structured_compactor_runs_after_request_shaping_middleware(self) -> None:
        middleware = build_agent_middleware(
            _stub_model(), MEMORY_PRESETS["exp_c200_cap10k_structured"]
        )
        self.assertIsInstance(middleware[0], DeepSeekCacheIsolationMiddleware)
        self.assertTrue(
            any(isinstance(m, StableToolOutputCapMiddleware) for m in middleware)
        )
        self.assertIsInstance(middleware[-2], SourcePreferenceMiddleware)
        self.assertIsInstance(middleware[-1], PrefixPreservingCompactionMiddleware)
        self.assertFalse(
            any(type(m) is InstrumentedSummarizationMiddleware for m in middleware)
        )

    def test_build_agent_cache_keys_include_memory_config(self) -> None:
        build_agent.cache_clear()
        created = []

        def fake_create_agent(**kwargs):
            agent = types.SimpleNamespace(kwargs=kwargs)
            created.append(agent)
            return agent

        try:
            with (
                patch(
                    "app.chat_service.build_chat_model",
                    return_value=_stub_model(),
                ),
                patch("app.chat_service.build_system_prompt", return_value="prompt"),
                patch("app.chat_service.create_agent", side_effect=fake_create_agent),
            ):
                prod = build_agent(
                    "google-genai:gemini-3.5-flash",
                    memory_config=MEMORY_PRESETS["prod"],
                )
                full_history = build_agent(
                    "google-genai:gemini-3.5-flash",
                    memory_config=MEMORY_PRESETS["full_history"],
                )
                full_history_again = build_agent(
                    "google-genai:gemini-3.5-flash",
                    memory_config=MEMORY_PRESETS["full_history"],
                )
        finally:
            build_agent.cache_clear()

        self.assertIsNot(prod, full_history)
        self.assertIs(full_history, full_history_again)
        self.assertEqual(len(created), 2)
        # Long-term memory needs the store wired into every agent.
        self.assertIn("store", created[0].kwargs)


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