nisse/nisse-core/tests/test_compact.py

"""Compaction (L6, L7): a long driver run's middle gives way to the model's
own notes, in one step now and then, so the run stays inside the window and
the start of its prompt stays the same between compactions."""

import io
import json
from itertools import pairwise
from pathlib import Path
from typing import Any

import pytest

from nisse_core.log.canonical import (
    AssistantText,
    CanonicalEvent,
    Compaction,
    Fact,
    StoredEvent,
    ToolCall,
    ToolResult,
    UserMessage,
)
from nisse_core.loop.driver import Driver, ToolOutcome
from nisse_core.loop.tools import tool
from nisse_core.models.compact import CompactionTuning, Compactor
from nisse_core.models.context import BudgetedClient
from nisse_core.models.interruptible import Interruptible, Stopped
from nisse_core.oplog import OpLog
from nisse_core.state.store import StateStore

READ = tool("read_file", "Read a file.", path="Its path.")
POST = tool("post_review", "Post the review. Ends your turn.", body="The review.")
TOOLS = [READ, POST]
BRIEF = "BRIEF: review PR #1"
FILE_CHARS = 6000  # about 2000 tokens a read


class Model:
    """Reads one file a turn, then posts; writes notes when asked for them."""

    def __init__(self, reads: int, notes: Any = "NOTES: all fine so far") -> None:
        self.reads = reads
        self.notes = notes
        self.requests: list[dict[str, Any]] = []
        self.turn = 0
        self.on_notes: Any = None

    def complete(
        self,
        events: list[StoredEvent],
        tools: Any,
        model: str,
        system: str | None = None,
        *,
        max_tokens: int | None = None,
        require_tool_call: bool = False,
    ) -> list[CanonicalEvent]:
        asked = _asks_for_notes(events)
        self.requests.append(
            {
                "events": list(events),
                "tools": tools,
                "system": system,
                "require_tool_call": require_tool_call,
                "notes": asked,
            }
        )
        if asked:
            if self.on_notes is not None:
                self.on_notes()
            notes = self.notes(events) if callable(self.notes) else self.notes
            return [AssistantText(text=notes, model=model)]
        self.turn += 1
        if self.turn > self.reads:
            return [ToolCall(call_id="c-post", tool="post_review", arguments={})]
        return [
            AssistantText(text=f"reading f{self.turn}", model=model),
            ToolCall(
                call_id=f"c-{self.turn}",
                tool="read_file",
                arguments={"path": f"f{self.turn}.py"},
            ),
        ]

    def turns_sent(self) -> list[list[StoredEvent]]:
        return [r["events"] for r in self.requests if not r["notes"]]


def _asks_for_notes(events: list[StoredEvent]) -> bool:
    last = events[-1].event if events else None
    return isinstance(last, UserMessage) and last.source == "compaction"


class Files:
    def dispatch(self, name: str, arguments: dict[str, Any]) -> ToolOutcome:
        path = arguments["path"]
        return ToolOutcome(f"{path}\n" + "x" * FILE_CHARS, True)


class World:
    def __init__(
        self,
        tmp_path: Path,
        model: Model,
        *,
        tuning: CompactionTuning | None = None,
        interruptible: bool = False,
    ) -> None:
        self.stream = io.StringIO()
        oplog = OpLog("t", self.stream)
        store = StateStore(tmp_path / "s.db", oplog=OpLog("t", io.StringIO()))
        self.log = store.transcript("x", "driver")
        self.log.as_actor("orchestrator").append(Fact(text=BRIEF))
        self.model = model
        # A 20k window less the margin and output ceiling: 17,000 tokens for
        # the prompt, so compaction starts at about 10,200.
        budgeted = BudgetedClient(
            model,
            lambda m: (20_000, None),
            max_output_tokens=1000,
            oplog=oplog,
            enforces_tool_choice=lambda m: True,
        )
        self.client: Any = Interruptible(budgeted) if interruptible else budgeted
        self.driver = Driver(
            self.client,
            "m",
            max_turns=60,
            compactor=Compactor(self.client, tuning, oplog=oplog),
        )

    def run(self) -> Any:
        return self.driver.run(
            self.log,
            system="you review",
            tools=TOOLS,
            waiting={"post_review"},
            workspace=Files(),
        )

    def oplog(self, event: str) -> list[dict[str, Any]]:
        records = [json.loads(line) for line in self.stream.getvalue().splitlines()]
        return [r for r in records if r["event"] == event]

    def compactions(self) -> list[Compaction]:
        return [e.event for e in self.log.read() if isinstance(e.event, Compaction)]


def test_a_long_run_is_compacted_keeping_its_brief_and_recent_turns(
    tmp_path: Path,
) -> None:
    world = World(tmp_path, Model(reads=7))
    assert world.run().kind == "wait"

    [compaction] = world.compactions()
    [logged] = world.oplog("context.compacted")
    assert logged["tokens_after"] < logged["tokens_before"]
    assert logged["work_item"] == "x"
    assert world.oplog("context.trimmed") == []  # fitting never had to cut

    last = world.model.turns_sent()[-1]
    # The brief first, word for word; then the notes; then the recent turns.
    assert last[0].event.text == BRIEF
    assert isinstance(last[1].event, Compaction)
    assert last[1].event.summary.startswith("NOTES: all fine so far")
    assert any(
        isinstance(e.event, ToolResult) and e.event.result.startswith("f7.py")
        for e in last
    )
    assert not any(
        isinstance(e.event, ToolResult) and e.event.result.startswith("f1.py")
        for e in last
    )
    # The turn just before the compaction is still there, word for word.
    sent = world.model.turns_sent()
    after = next(
        i for i, r in enumerate(sent) if any(isinstance(e.event, Compaction) for e in r)
    )
    latest = [e for e in sent[after - 1] if isinstance(e.event, ToolResult)][-1]
    assert latest in sent[after]
    # The calls it replaced are listed by code, so the model can make them again.
    assert '- read_file {"path": "f1.py"}' in compaction.summary


def test_the_notes_are_asked_for_over_the_runs_own_request(tmp_path: Path) -> None:
    world = World(tmp_path, Model(reads=8))
    world.run()
    requests = world.model.requests
    asked = next(i for i, r in enumerate(requests) if r["notes"])
    before, notes = requests[asked - 1], requests[asked]
    # Same system prompt and tools, and the last request plus the model's
    # reply and its result before it: only the request for notes is new,
    # so the server reuses its cache.
    assert notes["system"] == before["system"]
    assert notes["tools"] == before["tools"]
    assert before["require_tool_call"] is True
    assert notes["require_tool_call"] is False
    ids = [e.id for e in notes["events"][:-1]]
    assert ids[: len(before["events"])] == [e.id for e in before["events"]]


def test_the_prompt_start_stays_the_same_between_compactions(tmp_path: Path) -> None:
    world = World(tmp_path, Model(reads=20))
    world.run()
    assert len(world.compactions()) >= 2
    sent = world.model.turns_sent()
    unchanged = 0
    for earlier, later in pairwise(sent):
        compacted = any(isinstance(e.event, Compaction) for e in later) and not any(
            isinstance(e.event, Compaction) and e.event == c.event
            for e in later
            for c in earlier
            if isinstance(c.event, Compaction)
        )
        if compacted:
            continue
        # Between compactions, each request extends the one before it.
        assert [e.id for e in later[: len(earlier)]] == [e.id for e in earlier]
        unchanged += 1
    assert unchanged > len(sent) // 2


def test_a_second_compaction_folds_in_the_first(tmp_path: Path) -> None:
    world = World(tmp_path, Model(reads=20))
    world.run()
    first, second, *_ = world.compactions()
    assert second.replaces[0] == first.replaces[0]
    assert second.replaces[1] > first.replaces[1]
    # The first one's list of calls is carried into the second.
    assert '- read_file {"path": "f1.py"}' in second.summary
    last = world.model.turns_sent()[-1]
    assert sum(isinstance(e.event, Compaction) for e in last) == 1
    assert last[0].event.text == BRIEF


def test_tool_calls_are_never_parted_from_their_results(tmp_path: Path) -> None:
    world = World(tmp_path, Model(reads=20))
    world.run()
    for sent in world.model.turns_sent():
        calls = {e.event.call_id for e in sent if isinstance(e.event, ToolCall)}
        answered = {e.event.call_id for e in sent if isinstance(e.event, ToolResult)}
        assert answered <= calls


def test_a_short_run_is_never_compacted(tmp_path: Path) -> None:
    world = World(tmp_path, Model(reads=2))
    world.run()
    assert world.compactions() == []
    assert not any(r["notes"] for r in world.model.requests)


def test_empty_notes_fail_without_failing_the_run(tmp_path: Path) -> None:
    world = World(tmp_path, Model(reads=8, notes=""))
    assert world.run().kind == "wait"
    assert world.compactions() == []
    failed = world.oplog("context.compaction_failed")
    assert failed and failed[0]["cause"] == "the notes were empty"
    # Not tried again every turn: only after keep_turns more.
    asked = [i for i, r in enumerate(world.model.requests) if r["notes"]]
    assert all(b - a > 3 for a, b in pairwise(asked))


def test_with_no_compaction_fitting_cuts_and_says_so(tmp_path: Path) -> None:
    world = World(tmp_path, Model(reads=12, notes=""))
    world.run()
    [first, *_] = world.oplog("context.trimmed")
    assert first["results_cut"] > 0
    assert first["tokens_after"] <= first["limit"] < first["tokens_before"]


def test_notes_that_leave_it_too_long_are_not_kept(tmp_path: Path) -> None:
    world = World(tmp_path, Model(reads=8, notes="n" * 40_000))
    world.run()
    assert world.compactions() == []
    [failed, *_] = world.oplog("context.compaction_failed")
    assert "would still be" in failed["cause"]


def test_a_stop_while_writing_notes_breaks_off_cleanly(tmp_path: Path) -> None:
    model = Model(reads=8)
    world = World(tmp_path, model, interruptible=True)
    model.on_notes = world.client.interrupt
    with pytest.raises(Stopped):
        world.run()
    assert world.compactions() == []


def test_compact_at_zero_turns_it_off(tmp_path: Path) -> None:
    world = World(tmp_path, Model(reads=10), tuning=CompactionTuning(compact_at=0))
    world.run()
    assert world.compactions() == []
    assert world.oplog("context.trimmed")


def test_a_view_from_a_start_hides_compactions_of_what_came_before(
    tmp_path: Path,
) -> None:
    world = World(tmp_path, Model(reads=0))
    log = world.log
    log.as_actor("orchestrator").append(Fact(text="old work"))
    start = log.read()[-1].id
    log.as_actor("orchestrator").append(Fact(text="BRIEF again"))
    log.append(Compaction(summary="stale", model="m", replaces=[1, start]))
    log.append(
        Compaction(summary="this run's", model="m", replaces=[start + 1, start + 1])
    )
    shown = log.since(start).view()
    assert [e.event.summary for e in shown if isinstance(e.event, Compaction)] == [
        "this run's"
    ]