35 lines
1.5 KiB
Python
35 lines
1.5 KiB
Python
|
import pytest
|
||
|
|
from pr.core.context import should_compress_context, compress_context
|
||
|
|
from pr.config import RECENT_MESSAGES_TO_KEEP
|
||
|
|
|
||
|
|
|
||
|
|
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}'
|