|
from pr.config import RECENT_MESSAGES_TO_KEEP
|
|
from pr.core.context import compress_context, should_compress_context
|
|
|
|
|
|
class TestContextManagement:
|
|
|
|
def test_should_compress_context_below_threshold(self):
|
|
messages = [{"role": "user", "content": "test"}] * 10
|
|
assert should_compress_context(messages) is False
|
|
|
|
def test_should_compress_context_above_threshold(self):
|
|
messages = [{"role": "user", "content": "test"}] * 35
|
|
assert should_compress_context(messages) is True
|
|
|
|
def test_compress_context_preserves_system_message(self):
|
|
messages = [
|
|
{"role": "system", "content": "System prompt"},
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi"},
|
|
] * 40 # Ensure compression
|
|
compressed = compress_context(messages)
|
|
assert compressed[0]["role"] == "system"
|
|
assert "System prompt" in compressed[0]["content"]
|
|
|
|
def test_compress_context_keeps_recent_messages(self):
|
|
messages = [{"role": "user", "content": f"msg{i}"} for i in range(40)]
|
|
compressed = compress_context(messages)
|
|
# Should keep recent messages
|
|
recent = compressed[-RECENT_MESSAGES_TO_KEEP:]
|
|
assert len(recent) == RECENT_MESSAGES_TO_KEEP
|
|
# Check that the messages are the most recent ones
|
|
for i, msg in enumerate(recent):
|
|
expected_index = 40 - RECENT_MESSAGES_TO_KEEP + i
|
|
assert msg["content"] == f"msg{expected_index}"
|