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, memory.tree.Built(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, memory.tree.Built(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, memory.tree.Built(0, 8))
 151         require.True(t, memory.tree.Built(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, _ := memory.tree.Node(0, 0)
 173         right, _ := memory.tree.Node(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, memory.tree.Built(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, memory.tree.Built(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, built := memory.tree.Node(2, 0)
 221         require.True(t, built)
 222         require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
 223 }
 224
 225 func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {
 226         summarizer := newFakeSummarizer()
 227         summarizer.SetFailing(true)
 228         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
 229         require.NoError(t, err)
 230
 231         defer func() {
 232                 require.NoError(t, memory.Close())
 233         }()
 234
 235         _, err = memory.Append("echo", longText(0))
 236         require.NoError(t, err)
 237         require.NoError(t, memory.Wait(t.Context()))
 238         require.Equal(t, "<chat>\n0+1|"+pending+"\n</chat>", memory.View())
 239
 240         summarizer.SetFailing(false)
 241
 242         _, err = memory.Append("user", "again")
 243         require.NoError(t, err)
 244         require.NoError(t, memory.Wait(t.Context()))
 245         require.Equal(t, "<chat>\n0+1|"+fakeLine(memory.task(nodeKey{0, 0}))+"\n1+1|user: again\n</chat>", memory.View())
 246 }
 247
 248 func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
 249         root := t.TempDir()
 250         clock := newTestClock()
 251         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 252         require.NoError(t, err)
 253
 254         _, err = memory.Append("user", "a")
 255         require.NoError(t, err)
 256
 257         _, err = memory.log.Append("reply", "b")
 258         require.NoError(t, err)
 259         require.NoError(t, memory.Close())
 260
 261         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 262         require.NoError(t, err)
 263
 264         defer func() {
 265                 require.NoError(t, memory.Close())
 266         }()
 267
 268         require.Equal(t, "<chat>\n0+1|user: a\n1+1|reply: b\n</chat>", memory.View())
 269 }
 270
 271 func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
 272         root := t.TempDir()
 273         clock := newTestClock()
 274         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 275         require.NoError(t, err)
 276
 277         _, err = memory.Append("user", "a")
 278         require.NoError(t, err)
 279         require.NoError(t, memory.Close())
 280         require.NoError(t, os.Remove(filepath.Join(root, "view.json")))
 281
 282         _, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 283         require.ErrorIs(t, err, errNoView)
 284 }
 285
 286 // fakeSummarizer builds each node with a line of fakeLine. It fails each task
 287 // while failing is true, and each task waits for gate when gate is not nil.
 288 type fakeSummarizer struct {
 289         mutex   sync.Mutex
 290         gate    chan struct{}
 291         failing bool
 292         active  int
 293         peak    int
 294         tasks   []Task
 295 }
 296
 297 func newFakeSummarizer() *fakeSummarizer {
 298         return &fakeSummarizer{}
 299 }
 300
 301 func (f *fakeSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
 302         f.mutex.Lock()
 303         f.tasks = append(f.tasks, task)
 304         f.active++
 305         f.peak = max(f.peak, f.active)
 306         fail := f.failing
 307         f.mutex.Unlock()
 308
 309         defer func() {
 310                 f.mutex.Lock()
 311                 f.active--
 312                 f.mutex.Unlock()
 313         }()
 314
 315         if f.gate != nil {
 316                 select {
 317                 case <-f.gate:
 318                 case <-ctx.Done():
 319                         return "", ctx.Err()
 320                 }
 321         }
 322
 323         if fail {
 324                 return "", errors.New("fake failure")
 325         }
 326
 327         return fakeLine(task), nil
 328 }
 329
 330 // fakeLine returns a line of nodeBytes bytes or less that names the node and
 331 // starts with the input. A real line of this size makes each merge a model
 332 // call.
 333 func fakeLine(task Task) string {
 334         count := 1 << task.Level
 335         line := fmt.Sprintf("summary %d+%d: %s", task.Index*count, count, strings.Join(task.Input, " "))
 336         line = strings.ReplaceAll(line, "\n", " ")
 337         for len(line) > nodeBytes-12 {
 338                 _, size := utf8.DecodeLastRuneInString(line)
 339                 line = line[:len(line)-size]
 340         }
 341
 342         return line
 343 }
 344
 345 // longText returns a text that is too long for its own node.
 346 func longText(n int) string {
 347         return fmt.Sprintf("long %d ", n) + strings.Repeat("x", nodeBytes)
 348 }
 349
 350 // SetFailing makes the next tasks fail or not.
 351 func (f *fakeSummarizer) SetFailing(failing bool) {
 352         f.mutex.Lock()
 353         defer f.mutex.Unlock()
 354
 355         f.failing = failing
 356 }
 357
 358 // Active returns the number of tasks that run now.
 359 func (f *fakeSummarizer) Active() int {
 360         f.mutex.Lock()
 361         defer f.mutex.Unlock()
 362
 363         return f.active
 364 }
 365
 366 // Tasks returns the tasks so far.
 367 func (f *fakeSummarizer) Tasks() []Task {
 368         f.mutex.Lock()
 369         defer f.mutex.Unlock()
 370
 371         return slices.Clone(f.tasks)
 372 }