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

   1 package main
   2
   3 import (
   4         "fmt"
   5         "math/rand/v2"
   6         "os"
   7         "path/filepath"
   8         "runtime"
   9         "strings"
  10         "testing"
  11         "time"
  12
  13         "github.com/stretchr/testify/require"
  14 )
  15
  16 const (
  17         // replayMessages is the length of the synthetic log.
  18         replayMessages = 100000
  19
  20         // replayShortMessages is the length of the synthetic log with -short.
  21         // It is long enough for batches of the chat view.
  22         replayShortMessages = 5000
  23 )
  24
  25 // TestReplaySyntheticLog runs a synthetic log through the memory core with the
  26 // fake summarizer. After each message it waits for the builds, as a turn does.
  27 // It checks that each view stays within its sawtooth, that no view line holds a
  28 // whole long message or a part of an unbuilt one, and that a restart in the
  29 // middle and at the end loads the same views byte for byte. With -short, the
  30 // log is shorter, and the test still requires a batch of the chat view.
  31 func TestReplaySyntheticLog(t *testing.T) {
  32         start := time.Now()
  33         messages := replayMessages
  34         if testing.Short() {
  35                 messages = replayShortMessages
  36         }
  37
  38         root := t.TempDir()
  39         clock := newTestClock()
  40         random := rand.New(rand.NewPCG(1, 2))
  41         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
  42         require.NoError(t, err)
  43
  44         defer func() {
  45                 require.NoError(t, memory.Close())
  46         }()
  47
  48         batches := 0
  49
  50         low := viewChatHigh
  51         long := make([]bool, messages)
  52         for id := range messages {
  53                 if id == messages/2 {
  54                         memory = replayRequireRestart(t, memory, root, clock)
  55                 }
  56
  57                 clock.now = clock.now.Add(time.Minute)
  58                 kind, text := replayMessage(random, id)
  59                 long[id] = len(kind)+len(": ")+len(text) > nodeBytes
  60                 lines := replayViewLines(memory)
  61                 _, err = memory.Append(kind, text)
  62                 require.NoError(t, err)
  63
  64                 size, merged := replayRequireViews(t, memory, lines, long, id%1000 == 0)
  65                 if merged {
  66                         batches++
  67                         low = min(low, size)
  68                 }
  69
  70                 require.NoError(t, memory.Wait(t.Context()))
  71         }
  72
  73         memory = replayRequireRestart(t, memory, root, clock)
  74         require.Equal(t, messages, memory.log.Len())
  75         require.Positive(t, batches)
  76
  77         require.NoError(t, memory.Wait(t.Context()))
  78         t.Logf("%d messages, %d nodes, %d batches in %s", messages, memory.tree.Count(), batches, time.Since(start).Round(time.Second))
  79         t.Logf("chat view: %d lines, %d bytes, lowest after a batch %d bytes", len(memory.chat.lines), memory.chat.Size(), low)
  80         t.Logf("compaction view: %d lines, %d bytes", len(memory.cview.lines), memory.cview.Size())
  81
  82         runtime.GC()
  83
  84         var stats runtime.MemStats
  85
  86         runtime.ReadMemStats(&stats)
  87         t.Logf("heap in use: %d KB", stats.HeapInuse/1024)
  88 }
  89
  90 // replayRequireRestart closes memory, opens it again, and checks that the views
  91 // and view.json stay the same byte for byte.
  92 func replayRequireRestart(t *testing.T, memory *Memory, root string, clock *TestClock) *Memory {
  93         t.Helper()
  94
  95         require.NoError(t, memory.Wait(t.Context()))
  96
  97         view := memory.View()
  98         context := memory.Context(0, memory.log.Len())
  99         saved, err := os.ReadFile(filepath.Join(root, "view.json"))
 100         require.NoError(t, err)
 101         require.NoError(t, memory.Close())
 102
 103         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 104         require.NoError(t, err)
 105         require.NoError(t, memory.Wait(t.Context()))
 106         require.Equal(t, view, memory.View())
 107         require.Equal(t, context, memory.Context(0, memory.log.Len()))
 108
 109         loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
 110         require.NoError(t, err)
 111         require.Equal(t, saved, loaded)
 112
 113         return memory
 114 }
 115
 116 // replayMessage returns the kind and the text of message id. Most messages are
 117 // tool calls and their output. Users write short lines. One output in a hundred
 118 // is over 30 KB.
 119 func replayMessage(random *rand.Rand, id int) (string, string) {
 120         roll := random.IntN(100)
 121         switch {
 122         case roll < 5:
 123                 return "user", replayText(id, 10+random.IntN(200))
 124         case roll < 10:
 125                 return "reply", replayText(id, 100+random.IntN(1500))
 126         case roll < 12:
 127                 return "note", replayText(id, 50+random.IntN(600))
 128         case roll < 13:
 129                 return "work", replayText(id, 300+random.IntN(3000))
 130         case roll < 56:
 131                 return "tool", replayText(id, 60+random.IntN(400))
 132         case roll < 57:
 133                 return "echo", replayText(id, 30000+random.IntN(30000))
 134         default:
 135                 return "echo", replayText(id, 20+random.IntN(4000))
 136         }
 137 }
 138
 139 // replayViewLines returns the number of lines in the chat view of memory.
 140 func replayViewLines(memory *Memory) int {
 141         memory.mutex.Lock()
 142         defer memory.mutex.Unlock()
 143
 144         return len(memory.chat.lines)
 145 }
 146
 147 // replayRequireViews checks the views after a new message, and with full
 148 // compares them with the tree. Each view is at most its high limit, and a batch
 149 // merges the chat view to its low limit. The leaf of the new message can be
 150 // built after the append, so the sizes count it as the pending mark, as the
 151 // append did. It returns the size of the chat view and if a batch merged it,
 152 // from the number of lines before the message.
 153 func replayRequireViews(t *testing.T, memory *Memory, lines int, long []bool, full bool) (int, bool) {
 154         t.Helper()
 155
 156         memory.mutex.Lock()
 157         defer memory.mutex.Unlock()
 158
 159         size := replayAppendSize(memory.chat)
 160         require.LessOrEqual(t, size, viewChatHigh)
 161         require.LessOrEqual(t, replayAppendSize(memory.cview), viewCompactionHigh)
 162
 163         merged := len(memory.chat.lines) < lines
 164         if merged {
 165                 require.LessOrEqual(t, size, viewChatLow)
 166         }
 167
 168         replayRequireOnlySummaries(t, memory, long, full)
 169
 170         return size, merged
 171 }
 172
 173 // replayText returns a text of about size bytes that starts with its message ID
 174 // and holds some non-ASCII characters.
 175 func replayText(id, size int) string {
 176         head := fmt.Sprintf("m%d ", id)
 177         words := "ação data line café "
 178
 179         return head + strings.Repeat(words, max(1, (size-len(head))/len(words)))
 180 }
 181
 182 // replayAppendSize returns the size of view with its last line as the pending
 183 // mark.
 184 func replayAppendSize(view *View) int {
 185         last := view.lines[len(view.lines)-1]
 186         return view.Size() - view.lineSize(last) + len(viewPending)
 187 }
 188
 189 // replayRequireOnlySummaries checks each line of the chat view as the model
 190 // reads it. A built line holds its node, a long message has a summary of the
 191 // fake summarizer and not its text, and an unbuilt line holds only the pending
 192 // mark. With full, it also checks that the view keeps the node of each built
 193 // line. The caller holds the lock of memory.
 194 func replayRequireOnlySummaries(t *testing.T, memory *Memory, long []bool, full bool) {
 195         t.Helper()
 196
 197         rendered := strings.Split(memory.chat.Render(), "\n")
 198         require.Len(t, rendered, len(memory.chat.lines)+2)
 199
 200         for n, line := range memory.chat.lines {
 201                 head := fmt.Sprintf("%d+%d|", line.index<<line.level, 1<<line.level)
 202                 node, built := memory.chat.nodes[line]
 203                 if full {
 204                         require.Equal(t, treeIsBuilt(t, memory.tree, line.level, line.index), built)
 205                 }
 206
 207                 if !built {
 208                         require.Equal(t, 0, line.level)
 209                         require.Equal(t, head+viewPending, rendered[n+1])
 210
 211                         continue
 212                 }
 213
 214                 require.Equal(t, head+strings.ReplaceAll(node.Text, "\n", " "), rendered[n+1])
 215
 216                 if line.level == 0 && long[line.index] {
 217                         require.True(t, strings.HasPrefix(node.Text, "summary "), "message %d", line.index)
 218                 }
 219         }
 220 }