tests/test_chat_compact.py

import pytest

from app.websockets import handlers


class FakeConnectionManager:
    def __init__(self):
        self.sent = []

    async def send(self, project_id, event_type, payload):
        self.sent.append((project_id, event_type, payload))


class FakeCerebrasClient:
    def __init__(self, tokens):
        self.tokens = tokens
        self.last_messages = None

    async def stream_complete(self, messages, model=None, **kwargs):
        self.last_messages = messages
        for token in self.tokens:
            yield token


class FakeTutor:
    def __init__(self, client):
        self._client = client


@pytest.mark.asyncio
async def test_chat_compact_sends_distilled_summary(monkeypatch):
    cm = FakeConnectionManager()
    client = FakeCerebrasClient(["The student ", "discussed attention", " mechanisms."])
    monkeypatch.setattr(handlers, "_cm", cm)
    monkeypatch.setattr(handlers, "_get_tutor", lambda: FakeTutor(client))

    await handlers.handle_event("project-1", "CHAT_COMPACT", {
        "message_id": "compact-1",
        "history": [
            {"role": "user", "content": "What is self-attention?"},
            {"role": "assistant", "content": "It relates each token to every other token."},
        ],
    })

    assert cm.sent == [("project-1", "CHAT_COMPACTED", {
        "message_id": "compact-1",
        "summary": "The student discussed attention mechanisms.",
    })]
    # The transcript handed to the model must actually contain both turns.
    assert "self-attention" in client.last_messages[0]["content"]
    assert "every other token" in client.last_messages[0]["content"]


@pytest.mark.asyncio
async def test_chat_compact_with_empty_history_short_circuits_without_calling_the_model(monkeypatch):
    cm = FakeConnectionManager()
    client = FakeCerebrasClient([])
    monkeypatch.setattr(handlers, "_cm", cm)
    monkeypatch.setattr(handlers, "_get_tutor", lambda: FakeTutor(client))

    await handlers.handle_event("project-1", "CHAT_COMPACT", {
        "message_id": "compact-2",
        "history": [],
    })

    assert cm.sent == [("project-1", "CHAT_COMPACTED", {"message_id": "compact-2", "summary": ""})]
    assert client.last_messages is None


@pytest.mark.asyncio
async def test_chat_compact_truncates_an_oversized_transcript_before_calling_the_model(monkeypatch):
    cm = FakeConnectionManager()
    client = FakeCerebrasClient(["ok"])
    monkeypatch.setattr(handlers, "_cm", cm)
    monkeypatch.setattr(handlers, "_get_tutor", lambda: FakeTutor(client))

    huge_message = {"role": "user", "content": "x" * 50000}
    await handlers.handle_event("project-1", "CHAT_COMPACT", {
        "message_id": "compact-3",
        "history": [huge_message],
    })

    sent_transcript = client.last_messages[0]["content"]
    # MAX_COMPACT_TRANSCRIPT_CHARS caps the transcript itself; the whole system prompt
    # (instructions + transcript) should stay well short of the raw 50k-char input.
    assert len(sent_transcript) < 25000