cubepi/tests/checkpointer/test_memory.py at f6a1f411aeb93e3c1eeac3f3ce617d3c317600df · cubeplexai/cubepi · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
from cubepi.checkpointer.memory import MemoryCheckpointer
from cubepi.providers.base import TextContent, UserMessage


class TestMemoryCheckpointer:
    async def test_load_empty_thread(self):
        cp = MemoryCheckpointer()
        data = await cp.load("thread-1")
        assert data is None

    async def test_append_and_load(self):
        cp = MemoryCheckpointer()
        msg1 = UserMessage(content=[TextContent(text="hello")])
        msg2 = UserMessage(content=[TextContent(text="world")])

        await cp.append("thread-1", [msg1])
        await cp.append("thread-1", [msg2])

        data = await cp.load("thread-1")
        assert data is not None
        assert len(data.messages) == 2
        assert data.messages[0].content[0].text == "hello"
        assert data.messages[1].content[0].text == "world"

    async def test_save_extra(self):
        cp = MemoryCheckpointer()
        await cp.append("thread-1", [UserMessage(content=[TextContent(text="hi")])])
        await cp.save_extra("thread-1", {"compaction_index": 5})

        data = await cp.load("thread-1")
        assert data is not None
        assert data.extra["compaction_index"] == 5

    async def test_save_extra_merges(self):
        cp = MemoryCheckpointer()
        await cp.append("thread-1", [UserMessage(content=[TextContent(text="hi")])])
        await cp.save_extra("thread-1", {"a": 1})
        await cp.save_extra("thread-1", {"b": 2})

        data = await cp.load("thread-1")
        assert data.extra == {"a": 1, "b": 2}

    async def test_multiple_threads(self):
        cp = MemoryCheckpointer()
        await cp.append("t1", [UserMessage(content=[TextContent(text="t1")])])
        await cp.append("t2", [UserMessage(content=[TextContent(text="t2")])])

        d1 = await cp.load("t1")
        d2 = await cp.load("t2")

        assert d1.messages[0].content[0].text == "t1"
        assert d2.messages[0].content[0].text == "t2"

    async def test_save_extra_creates_thread(self):
        """save_extra on a thread that has never been used should create an entry."""
        cp = MemoryCheckpointer()
        await cp.save_extra("new-thread", {"key": "value"})

        data = await cp.load("new-thread")
        assert data is not None
        assert data.messages == []
        assert data.extra == {"key": "value"}