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