"""The model's own notes survive each compaction and lead every later request."""
import json
import unittest
from harness import Albedo, Provider, error, text
MARKER = "keeping notes for yourself"
EVICTED = "<newly-evicted-history>"
def last_user(request):
return next(
(
str(item.get("content", ""))
for item in reversed(request["input"])
if item.get("role") == "user"
),
"",
)
class NotesTests(unittest.TestCase):
def setUp(self):
self.written = []
self.notes_fail = False
def reply(request):
if MARKER in request.get("instructions", ""):
if self.notes_fail:
return error(500)
return text(self.written.pop(0))
if EVICTED in last_user(request):
return text("rolling summary")
return text("done")
self.provider = Provider(reply)
self.addCleanup(self.provider.close)
self.app = Albedo(self.provider, protocol="responses")
self.app.__enter__()
self.addCleanup(self.app.__exit__, None, None, None)
self.session = self.app.session()
def say(self, *prompts):
for prompt in prompts:
self.app.prompt(self.session, prompt).close()
self.app.idle(self.session)
def compact(self):
with self.app.api(
f"/sessions/{self.session}/commands", {"name": "/compact"}
) as response:
self.assertTrue(json.load(response)["result"]["started"])
self.app.idle(self.session)
def note_calls(self):
return [
item["request"]
for item in self.provider.requests
if MARKER in item["request"].get("instructions", "")
]
def head_of_next_request(self):
self.say("next")
return self.provider.requests[-1]["request"]["input"][0]["content"]
def test_notes_are_rewritten_from_the_previous_ones_and_lead_each_request(self):
self.written = ["NOTES-ONE commit abc123", "NOTES-TWO"]
self.say("first question", "second question", "third question")
self.compact()
(first,) = self.note_calls()
self.assertIn("at most 2000 tokens", first["instructions"])
self.assertIn("(none)", last_user(first))
self.assertIn("first question", last_user(first))
self.assertIn("[assistant]\ndone", last_user(first))
self.assertNotIn("provider item", last_user(first))
head = self.head_of_next_request()
self.assertIn("compaction notes", head)
self.assertIn("NOTES-ONE commit abc123", head)
self.say("fourth question", "fifth question")
self.compact()
second = self.note_calls()[1]
previous, evicted = last_user(second).split(EVICTED)
self.assertIn("NOTES-ONE commit abc123", previous)
self.assertIn("fourth question", evicted)
self.assertNotIn("first question", evicted)
head = self.head_of_next_request()
self.assertIn("NOTES-TWO", head)
self.assertNotIn("NOTES-ONE", head)
def test_a_failed_rewrite_keeps_the_old_notes_and_the_compaction(self):
self.written = ["NOTES-ONE"]
self.say("first question", "second question", "third question")
self.compact()
self.say("fourth question", "fifth question")
self.notes_fail = True
self.compact()
self.assertEqual(len(self.note_calls()), 2)
head = self.head_of_next_request()
self.assertIn("NOTES-ONE", head)
request = self.provider.requests[-1]["request"]
self.assertIn("rolling summary", json.dumps(request["input"]))
if __name__ == "__main__":
unittest.main()