okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

Aly Raffauffeat(api)!: migrate daemon and clients to protocol 361c43ca

main
branchesmainall refs →
3.6 KiB108 linesraw
1"""The model's own notes survive each compaction and lead every later request."""
2
3import json
4import unittest
5
6from harness import Albedo, Provider, error, text
7
8MARKER = "keeping notes for yourself"
9EVICTED = "<newly-evicted-history>"
10
11
12def last_user(request):
13    return next(
14        (
15            str(item.get("content", ""))
16            for item in reversed(request["input"])
17            if item.get("role") == "user"
18        ),
19        "",
20    )
21
22
23class NotesTests(unittest.TestCase):
24    def setUp(self):
25        self.written = []
26        self.notes_fail = False
27
28        def reply(request):
29            if MARKER in request.get("instructions", ""):
30                if self.notes_fail:
31                    return error(500)
32                return text(self.written.pop(0))
33            if EVICTED in last_user(request):
34                return text("rolling summary")
35            return text("done")
36
37        self.provider = Provider(reply)
38        self.addCleanup(self.provider.close)
39        self.app = Albedo(self.provider, protocol="responses")
40        self.app.__enter__()
41        self.addCleanup(self.app.__exit__, None, None, None)
42        self.session = self.app.session()
43
44    def say(self, *prompts):
45        for prompt in prompts:
46            self.app.prompt(self.session, prompt).close()
47            self.app.idle(self.session)
48
49    def compact(self):
50        with self.app.api(f"/sessions/{self.session}/compaction", {}) as response:
51            self.assertEqual(json.load(response)["state"], "compacted")
52        self.app.idle(self.session)
53
54    def note_calls(self):
55        return [
56            item["request"]
57            for item in self.provider.requests
58            if MARKER in item["request"].get("instructions", "")
59        ]
60
61    def head_of_next_request(self):
62        self.say("next")
63        return self.provider.requests[-1]["request"]["input"][0]["content"]
64
65    def test_notes_are_rewritten_from_the_previous_ones_and_lead_each_request(self):
66        self.written = ["NOTES-ONE commit abc123", "NOTES-TWO"]
67        self.say("first question", "second question", "third question")
68        self.compact()
69
70        (first,) = self.note_calls()
71        self.assertIn("at most 2000 tokens", first["instructions"])
72        self.assertIn("(none)", last_user(first))
73        self.assertIn("first question", last_user(first))
74        self.assertIn("[assistant]\ndone", last_user(first))
75        self.assertNotIn("provider item", last_user(first))
76        head = self.head_of_next_request()
77        self.assertIn("compaction notes", head)
78        self.assertIn("NOTES-ONE commit abc123", head)
79
80        self.say("fourth question", "fifth question")
81        self.compact()
82
83        second = self.note_calls()[1]
84        previous, evicted = last_user(second).split(EVICTED)
85        self.assertIn("NOTES-ONE commit abc123", previous)
86        self.assertIn("fourth question", evicted)
87        self.assertNotIn("first question", evicted)
88        head = self.head_of_next_request()
89        self.assertIn("NOTES-TWO", head)
90        self.assertNotIn("NOTES-ONE", head)
91
92    def test_a_failed_rewrite_keeps_the_old_notes_and_the_compaction(self):
93        self.written = ["NOTES-ONE"]
94        self.say("first question", "second question", "third question")
95        self.compact()
96        self.say("fourth question", "fifth question")
97        self.notes_fail = True
98        self.compact()
99
100        self.assertEqual(len(self.note_calls()), 2)
101        head = self.head_of_next_request()
102        self.assertIn("NOTES-ONE", head)
103        request = self.provider.requests[-1]["request"]
104        self.assertIn("rolling summary", json.dumps(request["input"]))
105
106
107if __name__ == "__main__":
108    unittest.main()