git.giacomolli.uk - onechat/blob - memory_test.go

   1 package main
   2
   3 import (
   4         "context"
   5         "errors"
   6         "fmt"
   7         "maps"
   8         "os"
   9         "path/filepath"
  10         "slices"
  11         "strings"
  12         "sync"
  13         "testing"
  14         "unicode/utf8"
  15
  16         "github.com/stretchr/testify/require"
  17 )
  18
  19 func TestMemoryRestartLoadsTheSameView(t *testing.T) {
  20         root := t.TempDir()
  21         clock := newTestClock()
  22         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
  23         require.NoError(t, err)
  24
  25         memory.chat.high = 200
  26         memory.chat.low = 100
  27         memory.cview.high = 80
  28         memory.cview.low = 40
  29         for n := range 100 {
  30                 _, err = memory.Append("user", fmt.Sprint(n))
  31                 require.NoError(t, err)
  32         }
  33
  34         view := memory.View()
  35         context := memory.Context(0, 100)
  36         saved, err := os.ReadFile(filepath.Join(root, "view.json"))
  37         require.NoError(t, err)
  38         require.Less(t, len(memory.chat.lines), 100)
  39         require.NoError(t, memory.Close())
  40
  41         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
  42         require.NoError(t, err)
  43
  44         defer func() {
  45                 require.NoError(t, memory.Close())
  46         }()
  47
  48         require.Equal(t, view, memory.View())
  49         require.Equal(t, context, memory.Context(0, 100))
  50
  51         loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
  52         require.NoError(t, err)
  53         require.Equal(t, saved, loaded)
  54 }
  55
  56 func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) {
  57         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
  58         require.NoError(t, err)
  59
  60         defer func() {
  61                 require.NoError(t, memory.Close())
  62         }()
  63
  64         memory.chat.high = 200
  65         memory.chat.low = 100
  66         memory.cview.high = 80
  67         memory.cview.low = 40
  68         for n := range 64 {
  69                 _, err = memory.Append("user", fmt.Sprint(n%10))
  70                 require.NoError(t, err)
  71
  72                 coarse := 0
  73                 for _, line := range memory.chat.lines {
  74                         outer := memory.cview.lines[coarse]
  75                         for nodeEnd(outer.level, outer.index) <= line.index<<line.level {
  76                                 coarse++
  77                                 outer = memory.cview.lines[coarse]
  78                         }
  79
  80                         require.GreaterOrEqual(t, outer.level, line.level, "message %d", n)
  81                 }
  82         }
  83 }
  84
  85 func TestMemoryRunsAtMostEightBuilds(t *testing.T) {
  86         summarizer := newFakeSummarizer()
  87         summarizer.gate = make(chan struct{})
  88         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
  89         require.NoError(t, err)
  90
  91         defer func() {
  92                 require.NoError(t, memory.Close())
  93         }()
  94
  95         for n := range 12 {
  96                 _, err = memory.Append("echo", longText(n))
  97                 require.NoError(t, err)
  98         }
  99
 100         memory.mutex.Lock()
 101         running := slices.Collect(maps.Keys(memory.queue.running))
 102         memory.mutex.Unlock()
 103         slices.SortFunc(running, func(a, b nodeKey) int {
 104                 return a.index - b.index
 105         })
 106         require.Equal(t, []nodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5}, {0, 6}, {0, 7}}, running)
 107
 108         close(summarizer.gate)
 109         require.NoError(t, memory.Wait(t.Context()))
 110         require.Equal(t, maxBuilds, summarizer.peak)
 111
 112         for level, count := range []int{12, 6, 3, 1} {
 113                 for index := range count {
 114                         require.True(t, memory.tree.Built(level, index), "node (%d, %d)", level, index)
 115                 }
 116         }
 117 }
 118
 119 func TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) {
 120         summarizer := newFakeSummarizer()
 121         summarizer.SetFailing(true)
 122         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
 123         require.NoError(t, err)
 124
 125         defer func() {
 126                 require.NoError(t, memory.Close())
 127         }()
 128
 129         for n := range 9 {
 130                 _, err = memory.Append("echo", longText(n))
 131                 require.NoError(t, err)
 132         }
 133
 134         require.NoError(t, memory.Wait(t.Context()))
 135         for _, task := range summarizer.Tasks() {
 136                 require.Less(t, task.Index, 8)
 137         }
 138
 139         require.False(t, memory.tree.Built(0, 8))
 140         summarizer.SetFailing(false)
 141
 142         _, err = memory.Append("user", "retry")
 143         require.NoError(t, err)
 144         require.NoError(t, memory.Wait(t.Context()))
 145         require.True(t, memory.tree.Built(0, 8))
 146         require.True(t, memory.tree.Built(3, 0))
 147 }
 148
 149 func TestMemoryMergesAfterBothHalves(t *testing.T) {
 150         summarizer := newFakeSummarizer()
 151         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
 152         require.NoError(t, err)
 153
 154         defer func() {
 155                 require.NoError(t, memory.Close())
 156         }()
 157
 158         _, err = memory.Append("echo", longText(0))
 159         require.NoError(t, err)
 160         require.NoError(t, memory.Wait(t.Context()))
 161
 162         _, err = memory.Append("echo", longText(1))
 163         require.NoError(t, err)
 164         require.NoError(t, memory.Wait(t.Context()))
 165         require.Len(t, summarizer.Tasks(), 3)
 166
 167         left, _ := memory.tree.Node(0, 0)
 168         right, _ := memory.tree.Node(0, 1)
 169         merge := summarizer.Tasks()[2]
 170         require.Equal(t, []string{left.Text, right.Text}, merge.Input)
 171         require.Equal(t, "<chat>\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n</chat>", merge.Context)
 172         require.True(t, memory.tree.Built(1, 0))
 173 }
 174
 175 func TestMemoryRestartRestoresTheQueue(t *testing.T) {
 176         root := t.TempDir()
 177         clock := newTestClock()
 178         summarizer := newFakeSummarizer()
 179         summarizer.SetFailing(true)
 180         memory, err := openMemory(t.Context(), root, clock.Now, summarizer)
 181         require.NoError(t, err)
 182
 183         for n := range 4 {
 184                 _, err = memory.Append("echo", longText(n))
 185                 require.NoError(t, err)
 186         }
 187
 188         require.NoError(t, memory.Wait(t.Context()))
 189         require.NoError(t, memory.Close())
 190
 191         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 192         require.NoError(t, err)
 193
 194         defer func() {
 195                 require.NoError(t, memory.Close())
 196         }()
 197
 198         require.NoError(t, memory.Wait(t.Context()))
 199         require.True(t, memory.tree.Built(2, 0))
 200 }
 201
 202 func TestMemoryBuildsFreeParents(t *testing.T) {
 203         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
 204         require.NoError(t, err)
 205
 206         defer func() {
 207                 require.NoError(t, memory.Close())
 208         }()
 209
 210         for _, text := range []string{"a", "b", "c", "d"} {
 211                 _, err = memory.Append("user", text)
 212                 require.NoError(t, err)
 213         }
 214
 215         node, built := memory.tree.Node(2, 0)
 216         require.True(t, built)
 217         require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
 218 }
 219
 220 func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {
 221         summarizer := newFakeSummarizer()
 222         summarizer.SetFailing(true)
 223         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
 224         require.NoError(t, err)
 225
 226         defer func() {
 227                 require.NoError(t, memory.Close())
 228         }()
 229
 230         _, err = memory.Append("echo", longText(0))
 231         require.NoError(t, err)
 232         require.NoError(t, memory.Wait(t.Context()))
 233         require.Equal(t, "<chat>\n0+1|"+pending+"\n</chat>", memory.View())
 234
 235         summarizer.SetFailing(false)
 236
 237         _, err = memory.Append("user", "again")
 238         require.NoError(t, err)
 239         require.NoError(t, memory.Wait(t.Context()))
 240         require.Equal(t, "<chat>\n0+1|"+fakeLine(memory.task(nodeKey{0, 0}))+"\n1+1|user: again\n</chat>", memory.View())
 241 }
 242
 243 func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
 244         root := t.TempDir()
 245         clock := newTestClock()
 246         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 247         require.NoError(t, err)
 248
 249         _, err = memory.Append("user", "a")
 250         require.NoError(t, err)
 251
 252         _, err = memory.log.Append("reply", "b")
 253         require.NoError(t, err)
 254         require.NoError(t, memory.Close())
 255
 256         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 257         require.NoError(t, err)
 258
 259         defer func() {
 260                 require.NoError(t, memory.Close())
 261         }()
 262
 263         require.Equal(t, "<chat>\n0+1|user: a\n1+1|reply: b\n</chat>", memory.View())
 264 }
 265
 266 func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
 267         root := t.TempDir()
 268         clock := newTestClock()
 269         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 270         require.NoError(t, err)
 271
 272         _, err = memory.Append("user", "a")
 273         require.NoError(t, err)
 274         require.NoError(t, memory.Close())
 275         require.NoError(t, os.Remove(filepath.Join(root, "view.json")))
 276
 277         _, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 278         require.ErrorIs(t, err, errNoView)
 279 }
 280
 281 // fakeSummarizer builds each node with a line of fakeLine. It fails each task
 282 // while failing is true, and each task waits for gate when gate is not nil.
 283 type fakeSummarizer struct {
 284         mutex   sync.Mutex
 285         gate    chan struct{}
 286         failing bool
 287         active  int
 288         peak    int
 289         tasks   []Task
 290 }
 291
 292 func newFakeSummarizer() *fakeSummarizer {
 293         return &fakeSummarizer{}
 294 }
 295
 296 func (f *fakeSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
 297         f.mutex.Lock()
 298         f.tasks = append(f.tasks, task)
 299         f.active++
 300         f.peak = max(f.peak, f.active)
 301         fail := f.failing
 302         f.mutex.Unlock()
 303
 304         defer func() {
 305                 f.mutex.Lock()
 306                 f.active--
 307                 f.mutex.Unlock()
 308         }()
 309
 310         if f.gate != nil {
 311                 select {
 312                 case <-f.gate:
 313                 case <-ctx.Done():
 314                         return "", ctx.Err()
 315                 }
 316         }
 317
 318         if fail {
 319                 return "", errors.New("fake failure")
 320         }
 321
 322         return fakeLine(task), nil
 323 }
 324
 325 // fakeLine returns a line of nodeBytes bytes or less that names the node and
 326 // starts with the input. A real line of this size makes each merge a model
 327 // call.
 328 func fakeLine(task Task) string {
 329         count := 1 << task.Level
 330         line := fmt.Sprintf("summary %d+%d: %s", task.Index*count, count, strings.Join(task.Input, " "))
 331         line = strings.ReplaceAll(line, "\n", " ")
 332         for len(line) > nodeBytes-12 {
 333                 _, size := utf8.DecodeLastRuneInString(line)
 334                 line = line[:len(line)-size]
 335         }
 336
 337         return line
 338 }
 339
 340 // longText returns a text that is too long for its own node.
 341 func longText(n int) string {
 342         return fmt.Sprintf("long %d ", n) + strings.Repeat("x", nodeBytes)
 343 }
 344
 345 // SetFailing makes the next tasks fail or not.
 346 func (f *fakeSummarizer) SetFailing(failing bool) {
 347         f.mutex.Lock()
 348         defer f.mutex.Unlock()
 349
 350         f.failing = failing
 351 }
 352
 353 // Tasks returns the tasks so far.
 354 func (f *fakeSummarizer) Tasks() []Task {
 355         f.mutex.Lock()
 356         defer f.mutex.Unlock()
 357
 358         return slices.Clone(f.tasks)
 359 }