Spaces:
Sleeping
Sleeping
File size: 2,997 Bytes
2e818da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | 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
|