101 lines
3.6 KiB
Python
101 lines
3.6 KiB
Python
import unittest
|
|
|
|
from core.context import prepare_messages_for_llm, prepare_messages_with_stats
|
|
|
|
|
|
class ContextCompactionTests(unittest.TestCase):
|
|
def test_preserves_system_and_recent_messages(self) -> None:
|
|
messages = [
|
|
{"role": "system", "content": "rules"},
|
|
{"role": "user", "content": "old"},
|
|
{"role": "tool", "tool_call_id": "old-tool", "name": "shell", "content": "A" * 200},
|
|
{"role": "user", "content": "latest"},
|
|
{"role": "tool", "tool_call_id": "new-tool", "name": "shell", "content": "B" * 200},
|
|
]
|
|
|
|
prepared = prepare_messages_for_llm(
|
|
messages,
|
|
keep_recent=2,
|
|
old_tool_chars=40,
|
|
)
|
|
|
|
self.assertEqual(prepared[0], messages[0])
|
|
self.assertEqual(prepared[-2], messages[-2])
|
|
self.assertEqual(prepared[-1], messages[-1])
|
|
|
|
def test_compacts_old_tool_content_without_breaking_protocol_fields(self) -> None:
|
|
messages = [
|
|
{"role": "system", "content": "rules"},
|
|
{"role": "assistant", "tool_calls": [{"id": "tc1"}], "content": None},
|
|
{"role": "tool", "tool_call_id": "tc1", "name": "run_python", "content": "A" * 200},
|
|
{"role": "user", "content": "continue"},
|
|
]
|
|
|
|
prepared = prepare_messages_for_llm(
|
|
messages,
|
|
keep_recent=1,
|
|
old_tool_chars=40,
|
|
)
|
|
tool_msg = prepared[2]
|
|
|
|
self.assertEqual(tool_msg["role"], "tool")
|
|
self.assertEqual(tool_msg["tool_call_id"], "tc1")
|
|
self.assertEqual(tool_msg["name"], "run_python")
|
|
self.assertIn("[compacted old tool result", tool_msg["content"])
|
|
self.assertLess(len(tool_msg["content"]), 120)
|
|
|
|
def test_short_old_tool_content_is_left_unchanged(self) -> None:
|
|
messages = [
|
|
{"role": "system", "content": "rules"},
|
|
{"role": "tool", "tool_call_id": "tc1", "name": "grep", "content": "short"},
|
|
{"role": "user", "content": "next"},
|
|
]
|
|
|
|
prepared = prepare_messages_for_llm(
|
|
messages,
|
|
keep_recent=1,
|
|
old_tool_chars=40,
|
|
)
|
|
|
|
self.assertEqual(prepared[1]["content"], "short")
|
|
|
|
def test_compacts_old_load_skill_result_to_marker(self) -> None:
|
|
messages = [
|
|
{"role": "system", "content": "rules"},
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": "tc1",
|
|
"name": "load_skill",
|
|
"content": "[skill=proposal, dir=/sandbox/skills/proposal]\n" + "A" * 5000,
|
|
},
|
|
{"role": "user", "content": "next"},
|
|
]
|
|
|
|
prepared = prepare_messages_for_llm(messages, keep_recent=1)
|
|
|
|
self.assertIn("loaded skill: proposal", prepared[1]["content"])
|
|
self.assertIn("dir=/sandbox/skills/proposal", prepared[1]["content"])
|
|
self.assertNotIn("A" * 100, prepared[1]["content"])
|
|
|
|
def test_prepare_messages_reports_compaction_stats(self) -> None:
|
|
messages = [
|
|
{"role": "system", "content": "rules"},
|
|
{"role": "tool", "tool_call_id": "tc1", "name": "shell", "content": "A" * 200},
|
|
{"role": "user", "content": "next"},
|
|
]
|
|
|
|
prepared, stats = prepare_messages_with_stats(
|
|
messages,
|
|
keep_recent=1,
|
|
old_tool_chars=40,
|
|
)
|
|
|
|
self.assertLess(stats["sent_chars"], stats["original_chars"])
|
|
self.assertEqual(stats["compacted_tool_messages"], 1)
|
|
self.assertGreater(stats["saved_chars"], 0)
|
|
self.assertEqual(len(prepared), len(messages))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|