Newer
Older
navi-1 / tests / unit / core / test_compression_breadcrumbs.py
@Eugene Sukhodolskiy Eugene Sukhodolskiy 21 hours ago 7 KB core: split the compressor into a package
"""Compression breadcrumbs: the KEYS line and the raw-chunk index it feeds."""

from types import SimpleNamespace

import navi.config as _config
import navi.core.compression.compressor as _cm
from navi.config import Settings
from navi.core.compressor import ContextCompressor, compress_context
from navi.core.session import InMemorySessionStore
from navi.llm.base import Message
from tests.conftest_factory import FakeLLMBackend


def _conversation() -> list[Message]:
    return [
        Message(role="system", content="sys"),
        Message(role="user", content="use `uv run pytest` in navi/core/compressor.py", sequence_number=1),
        Message(role="assistant", content="ok, port 8099", sequence_number=2),
        Message(role="user", content="next", sequence_number=3),
        Message(role="assistant", content="next answer", sequence_number=4),
        Message(role="user", content="latest", sequence_number=5),
        Message(role="assistant", content="latest answer", sequence_number=6),
    ]


class TestKeysLine:
    async def test_the_keys_line_is_stripped_and_stored_on_the_summary(self):
        backend = FakeLLMBackend(responses=[
            "## Goal\ndo it\n\nKEYS: `uv run pytest`, navi/core/compressor.py, port 8099"
        ])
        new_context, summary = await compress_context(
            context=_conversation(), llm=backend, model="test", temperature=0.3, keep_recent=2,
        )
        assert "KEYS" not in summary
        assert summary.startswith("## Goal")
        summary_msg = next(m for m in new_context if m.is_summary)
        assert summary_msg.metadata["breadcrumbs"] == [
            "uv run pytest", "navi/core/compressor.py", "port 8099",
        ]

    async def test_a_summary_without_keys_keeps_empty_metadata(self):
        backend = FakeLLMBackend(responses=["plain summary"])
        new_context, summary = await compress_context(
            context=_conversation(), llm=backend, model="test", temperature=0.3, keep_recent=2,
        )
        assert summary == "plain summary"
        assert next(m for m in new_context if m.is_summary).metadata == {}

    async def test_a_keys_only_response_does_not_become_an_empty_summary(self):
        backend = FakeLLMBackend(responses=["KEYS: alpha, beta"])
        _, summary = await compress_context(
            context=_conversation(), llm=backend, model="test", temperature=0.3, keep_recent=2,
        )
        assert summary == "(summary unavailable)"


class FakeChunkStore:
    def __init__(self):
        self.calls = []

    async def add_chunks(self, chunks, user_id=None, session_id=None):
        self.calls.append((chunks, user_id, session_id))
        return len(chunks)


class TestIndexDemoted:
    def _demoted(self) -> list[Message]:
        return [
            Message(role="user", content="run `navi_search --port 8099`", sequence_number=1),
            Message(role="assistant", content="done, see navi/memory/_chunks.py", sequence_number=2),
            Message(role="user", content="old summary text", sequence_number=3, is_summary=True),
        ]

    async def test_chunks_are_written_with_the_summary_keys_and_extracted_names(self):
        store = FakeChunkStore()
        summary_msg = Message(
            role="user", content="summary", is_summary=True,
            metadata={"breadcrumbs": ["from_summary_key"]},
        )
        session = SimpleNamespace(user_id="u1")
        compressor = ContextCompressor()

        await compressor._index_demoted(store, session, "s1", self._demoted(), summary_msg)

        chunks, user_id, session_id = store.calls[0]
        assert (user_id, session_id) == ("u1", "s1")
        assert len(chunks) == 1
        keys = chunks[0]["breadcrumbs"]
        assert "from_summary_key" in keys
        assert "navi/memory/_chunks.py" in keys
        assert "navi_search --port 8099" in keys
        assert "old summary text" not in chunks[0]["content"]

    async def test_the_hard_truncate_path_falls_back_to_the_extractor(self):
        store = FakeChunkStore()
        await ContextCompressor()._index_demoted(
            store, SimpleNamespace(user_id=None), "s1", self._demoted(), None,
        )
        keys = store.calls[0][0][0]["breadcrumbs"]
        assert "navi/memory/_chunks.py" in keys

    async def test_indexing_can_be_switched_off(self, monkeypatch):
        new = Settings(_env_file=None, memory_search_enabled=False)
        monkeypatch.setattr(_config, "settings", new)
        monkeypatch.setattr(_cm, "settings", new)
        store = FakeChunkStore()
        await ContextCompressor()._index_demoted(
            store, SimpleNamespace(user_id="u1"), "s1", self._demoted(), None,
        )
        assert store.calls == []

    async def test_a_broken_store_never_breaks_compression(self):
        class Broken(FakeChunkStore):
            async def add_chunks(self, chunks, user_id=None, session_id=None):
                raise RuntimeError("db down")

        # No exception escapes: losing an index entry is cheaper than losing a turn.
        await ContextCompressor()._index_demoted(
            Broken(), SimpleNamespace(user_id="u1"), "s1", self._demoted(), None,
        )

    async def test_no_messages_means_no_call(self):
        store = FakeChunkStore()
        await ContextCompressor()._index_demoted(store, SimpleNamespace(user_id="u1"), "s1", [], None)
        assert store.calls == []


class TestDemotedTurnsAreIndexedOnce:
    """`session.messages` and `session.context` share Message objects (that is how
    `PgSessionStore.get` builds them), and a turn dropped by one compression stays
    in the display history forever. Demoting the union indexed every dropped turn
    twice in the same pass and again on every later pass, under a new chunk
    boundary each time — duplicate rows and repeated embedding work."""

    @staticmethod
    def _turns(start: int, count: int) -> list[Message]:
        out: list[Message] = []
        for i in range(count):
            out.append(Message(role="user", content=f"turn-{start + i} question",
                               sequence_number=(start + i) * 2))
            out.append(Message(role="assistant", content=f"turn-{start + i} answer",
                               sequence_number=(start + i) * 2 + 1))
        return out

    async def test_a_dropped_turn_is_written_once_and_never_reindexed(self):
        store = FakeChunkStore()
        sessions = InMemorySessionStore()
        session = await sessions.create(profile_id="test")
        session.context = self._turns(0, 8)
        session.messages = list(session.context)
        compressor = ContextCompressor()
        backend = FakeLLMBackend(responses=["summary one", "summary two"])

        async def compress():
            return await compressor.compress_and_save_session(
                session=session,
                session_store=sessions,
                llm=backend,
                model="m",
                temperature=0.3,
                session_id=session.id,
                reason="pytest",
                keep_recent=1,
                window=65_536,
                memory_store=store,
            )

        await compress()
        first_text = "\n".join(c["content"] for c in store.calls[0][0])
        assert first_text.count("turn-0 question") == 1

        added = self._turns(8, 4)
        session.context.extend(added)
        session.messages.extend(added)
        await compress()
        second_text = "\n".join(c["content"] for c in store.calls[1][0])
        assert "turn-8 question" in second_text
        assert "turn-0 question" not in second_text
        assert second_text.count("turn-8 question") == 1