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

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "errors"
   7         "fmt"
   8         "os"
   9         "path/filepath"
  10         "strings"
  11         "sync"
  12         "time"
  13 )
  14
  15 // errNoView tells that a chat has messages but no saved view.
  16 var errNoView = errors.New("the chat has messages but no saved view")
  17
  18 // Summarizer builds the text of one node with a model.
  19 type Summarizer interface {
  20         Summarize(ctx context.Context, task Task) (string, error)
  21 }
  22
  23 // Task is the build of the node (Level, Index). Context is the compaction view
  24 // for the build. Input holds the source "kind: text" of a leaf, or the two
  25 // child lines of a merge with their newlines as spaces.
  26 type Task struct {
  27         Level   int
  28         Index   int
  29         Context string
  30         Input   []string
  31 }
  32
  33 // Memory is the memory of one chat: the message log, the summary tree, the chat
  34 // view, the compaction view, and the queue of the builds. The compaction view
  35 // is the chat view merged further, at the limits of 16 to 32 KB. It takes the
  36 // same new lines, and it merges again in one batch when it is too large or when
  37 // the chat view merges. So the compactions share its prefix in the prompt
  38 // cache. It saves the views to view.json after each change and loads them at
  39 // start. It never builds the views again from the log, because a new view
  40 // differs from the live one and so loses the prompt cache.
  41 //
  42 // The summarizer builds the nodes in the queue in the background. The start
  43 // reads the log and the tree once to fill the queue again. After that, only a
  44 // new message or a new node adds work.
  45 type Memory struct {
  46         mutex      sync.Mutex
  47         root       string
  48         log        *Log
  49         tree       *Tree
  50         chat       *View
  51         cview      *View
  52         queue      *queue
  53         summarizer Summarizer
  54         ctx        context.Context
  55         cancel     context.CancelFunc
  56         workers    sync.WaitGroup
  57         changed    chan struct{}
  58 }
  59
  60 // viewFile is the content of view.json.
  61 type viewFile struct {
  62         Chat       viewState `json:"chat"`
  63         Compaction viewState `json:"compaction"`
  64 }
  65
  66 // openMemory opens the chat in root and starts the builds that are ready. The
  67 // clock gives the dates of new messages and the days of new nodes. The builds
  68 // stop when ctx ends or at Close.
  69 func openMemory(ctx context.Context, root string, clock func() time.Time, summarizer Summarizer) (*Memory, error) {
  70         log, err := openLog(ctx, root, clock)
  71         if err != nil {
  72                 return nil, err
  73         }
  74
  75         m := &Memory{
  76                 root:       root,
  77                 log:        log,
  78                 queue:      newQueue(),
  79                 summarizer: summarizer,
  80                 changed:    make(chan struct{}),
  81         }
  82
  83         m.ctx, m.cancel = context.WithCancel(ctx)
  84
  85         m.tree, err = openTree(ctx, root, log.Len(), clock)
  86         if err != nil {
  87                 return nil, errors.Join(err, m.Close())
  88         }
  89
  90         err = m.load()
  91         if err != nil {
  92                 return nil, errors.Join(fmt.Errorf("load view: %w", err), m.Close())
  93         }
  94
  95         err = m.restoreQueue()
  96         if err != nil {
  97                 return nil, errors.Join(fmt.Errorf("restore build queue: %w", err), m.Close())
  98         }
  99
 100         err = m.catchUp()
 101         if err != nil {
 102                 return nil, errors.Join(err, m.Close())
 103         }
 104
 105         m.mutex.Lock()
 106         m.dispatch()
 107         m.mutex.Unlock()
 108
 109         return m, nil
 110 }
 111
 112 // load reads the saved views.
 113 func (m *Memory) load() error {
 114         data, err := os.ReadFile(m.viewPath())
 115         if errors.Is(err, os.ErrNotExist) && m.log.Len() == 0 {
 116                 m.chat = newView(chatHigh, chatLow)
 117                 m.cview = newView(compactionHigh, compactionLow)
 118
 119                 return nil
 120         }
 121
 122         if errors.Is(err, os.ErrNotExist) {
 123                 return errNoView
 124         }
 125
 126         if err != nil {
 127                 return fmt.Errorf("read view: %w", err)
 128         }
 129
 130         var file viewFile
 131         err = json.Unmarshal(data, &file)
 132         if err != nil {
 133                 return fmt.Errorf("decode view: %w", err)
 134         }
 135
 136         m.chat, err = restoreView(file.Chat, chatHigh, chatLow, m.tree)
 137         if err != nil {
 138                 return fmt.Errorf("chat view: %w", err)
 139         }
 140
 141         m.cview, err = restoreView(file.Compaction, compactionHigh, compactionLow, m.tree)
 142         if err != nil {
 143                 return fmt.Errorf("compaction view: %w", err)
 144         }
 145
 146         if m.cview.end != m.chat.end {
 147                 return fmt.Errorf("compaction view covers %d messages, but the chat view %d", m.cview.end, m.chat.end)
 148         }
 149
 150         if m.chat.end > m.log.Len() {
 151                 return fmt.Errorf("view covers %d messages, but the log has %d", m.chat.end, m.log.Len())
 152         }
 153
 154         return nil
 155 }
 156
 157 // restoreQueue adds the work of the messages in the view to the queue. It
 158 // builds each node that needs no model call and that a crash left out.
 159 func (m *Memory) restoreQueue() error {
 160         for id := range m.chat.end {
 161                 if m.tree.Built(0, id) {
 162                         continue
 163                 }
 164
 165                 leaf, free := freeLeaf(m.log.Message(id))
 166                 if !free {
 167                         m.queue.AddLeaf(id)
 168                         continue
 169                 }
 170
 171                 err := m.build(leaf)
 172                 if err != nil {
 173                         return err
 174                 }
 175         }
 176
 177         for _, node := range m.tree.Nodes() {
 178                 if node.Index%2 != 0 {
 179                         continue
 180                 }
 181
 182                 err := m.settle(node)
 183                 if err != nil {
 184                         return err
 185                 }
 186         }
 187
 188         return nil
 189 }
 190
 191 // catchUp adds the messages that the view does not cover yet. A crash can stop
 192 // an append after the log write and before the view save.
 193 func (m *Memory) catchUp() error {
 194         if m.chat.end == m.log.Len() {
 195                 return nil
 196         }
 197
 198         for m.chat.end < m.log.Len() {
 199                 err := m.addMessage(m.log.Message(m.chat.end))
 200                 if err != nil {
 201                         return err
 202                 }
 203         }
 204
 205         return m.save()
 206 }
 207
 208 // Close stops the builds, waits for them, closes the files of the chat, and
 209 // releases its lock.
 210 func (m *Memory) Close() error {
 211         m.cancel()
 212         m.workers.Wait()
 213
 214         m.mutex.Lock()
 215         defer m.mutex.Unlock()
 216
 217         var err error
 218         if m.tree != nil {
 219                 err = m.tree.Close()
 220         }
 221
 222         return errors.Join(err, m.log.Close())
 223 }
 224
 225 // Append logs a new message, adds it to the tree and the views, and starts the
 226 // builds that are ready, also the failed ones.
 227 func (m *Memory) Append(kind, text string) (Message, error) {
 228         m.mutex.Lock()
 229         defer m.mutex.Unlock()
 230
 231         message, err := m.log.Append(kind, text)
 232         if err != nil {
 233                 return Message{}, err
 234         }
 235
 236         err = m.addMessage(message)
 237         if err != nil {
 238                 return Message{}, err
 239         }
 240
 241         err = m.save()
 242         if err != nil {
 243                 return Message{}, err
 244         }
 245
 246         m.queue.Retry()
 247         m.dispatch()
 248
 249         return message, nil
 250 }
 251
 252 // addMessage builds the leaf of message when it needs no model call, or else
 253 // adds the leaf to the queue. Then it appends the leaf to the views.
 254 func (m *Memory) addMessage(message Message) error {
 255         leaf, free := freeLeaf(message)
 256         if free {
 257                 err := m.build(leaf)
 258                 if err != nil {
 259                         return err
 260                 }
 261         } else {
 262                 m.queue.AddLeaf(message.ID)
 263         }
 264
 265         merged := m.chat.Append(m.tree)
 266         m.cview.Append(m.tree)
 267         if merged {
 268                 m.cview.Coarsen(m.chat, m.tree)
 269         }
 270
 271         return nil
 272 }
 273
 274 // build adds node to the tree and settles its parent.
 275 func (m *Memory) build(node Node) error {
 276         err := m.tree.Add(node)
 277         if err != nil {
 278                 return err
 279         }
 280
 281         return m.settle(node)
 282 }
 283
 284 // settle looks at the parent of node when both children are built. It builds
 285 // each parent that needs no model call, up the tree, and adds the first other
 286 // parent to the queue.
 287 func (m *Memory) settle(node Node) error {
 288         for {
 289                 sibling, built := m.tree.Node(node.Level, node.Index^1)
 290                 if !built || m.tree.Built(node.Level+1, node.Index/2) {
 291                         return nil
 292                 }
 293
 294                 left, right := node, sibling
 295                 if node.Index%2 != 0 {
 296                         left, right = sibling, node
 297                 }
 298
 299                 parent, free := freeMerge(left, right)
 300                 if !free {
 301                         m.queue.AddMerge(nodeKey{node.Level + 1, node.Index / 2})
 302                         return nil
 303                 }
 304
 305                 err := m.tree.Add(parent)
 306                 if err != nil {
 307                         return err
 308                 }
 309
 310                 node = parent
 311         }
 312 }
 313
 314 // dispatch starts the builds that are ready, up to the limit.
 315 func (m *Memory) dispatch() {
 316         for m.ctx.Err() == nil {
 317                 key, ready := m.queue.Next()
 318                 if !ready {
 319                         return
 320                 }
 321
 322                 task := m.task(key)
 323                 m.workers.Add(1)
 324                 go m.run(key, task)
 325         }
 326 }
 327
 328 // task returns the build task of the node key.
 329 func (m *Memory) task(key nodeKey) Task {
 330         task := Task{
 331                 Level:   key.level,
 332                 Index:   key.index,
 333                 Context: m.context(key.level, key.index),
 334         }
 335
 336         if key.level == 0 {
 337                 task.Input = []string{leafSource(m.log.Message(key.index))}
 338                 return task
 339         }
 340
 341         for _, index := range []int{2 * key.index, 2*key.index + 1} {
 342                 child, _ := m.tree.Node(key.level-1, index)
 343                 task.Input = append(task.Input, strings.ReplaceAll(child.Text, "\n", " "))
 344         }
 345
 346         return task
 347 }
 348
 349 // run builds one node with the summarizer. A failed build stays in the queue
 350 // until the next message.
 351 func (m *Memory) run(key nodeKey, task Task) {
 352         defer m.workers.Done()
 353
 354         text, err := m.summarizer.Summarize(m.ctx, task)
 355
 356         m.mutex.Lock()
 357         defer m.mutex.Unlock()
 358
 359         if err == nil {
 360                 err = m.build(Node{
 361                         Level: key.level,
 362                         Index: key.index,
 363                         Text:  text,
 364                         Size:  len(text),
 365                 })
 366         }
 367
 368         if err != nil {
 369                 m.queue.Fail(key)
 370         } else {
 371                 m.queue.Done(key)
 372         }
 373
 374         close(m.changed)
 375         m.changed = make(chan struct{})
 376         m.dispatch()
 377 }
 378
 379 // Wait returns when no build runs. The failed builds wait for the next message.
 380 func (m *Memory) Wait(ctx context.Context) error {
 381         for {
 382                 m.mutex.Lock()
 383                 running := m.queue.Running()
 384                 changed := m.changed
 385                 m.mutex.Unlock()
 386
 387                 if running == 0 {
 388                         return nil
 389                 }
 390
 391                 select {
 392                 case <-changed:
 393                 case <-ctx.Done():
 394                         return ctx.Err()
 395                 }
 396         }
 397 }
 398
 399 // save writes the views to view.json.
 400 func (m *Memory) save() error {
 401         data, err := json.Marshal(viewFile{
 402                 Chat:       m.chat.State(),
 403                 Compaction: m.cview.State(),
 404         })
 405         if err != nil {
 406                 return fmt.Errorf("encode view: %w", err)
 407         }
 408
 409         err = writeFileAtomic(m.viewPath(), append(data, '\n'))
 410         if err != nil {
 411                 return fmt.Errorf("save view: %w", err)
 412         }
 413
 414         return nil
 415 }
 416
 417 func (m *Memory) viewPath() string {
 418         return filepath.Join(m.root, "view.json")
 419 }
 420
 421 // View returns the chat view as the model reads it.
 422 func (m *Memory) View() string {
 423         m.mutex.Lock()
 424         defer m.mutex.Unlock()
 425
 426         return m.chat.Render(m.tree)
 427 }
 428
 429 // Context returns the compaction view for the build of node (l, i). It ends
 430 // before the message of a leaf, or after the last message of a merge, and it
 431 // stops at the first line that is not built.
 432 func (m *Memory) Context(level, index int) string {
 433         m.mutex.Lock()
 434         defer m.mutex.Unlock()
 435
 436         return m.context(level, index)
 437 }
 438
 439 func (m *Memory) context(level, index int) string {
 440         end := index
 441         if level > 0 {
 442                 end = nodeEnd(level, index)
 443         }
 444
 445         return m.cview.Context(m.tree, end)
 446 }
 447
 448 // writeFileAtomic replaces the file at path with data through a temporary file
 449 // and a rename. A reader or a crash of the process sees the old file or the new
 450 // file, never a part of one.
 451 func writeFileAtomic(path string, data []byte) error {
 452         file, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".*")
 453         if err != nil {
 454                 return fmt.Errorf("create temporary file: %w", err)
 455         }
 456
 457         _, err = file.Write(data)
 458         err = errors.Join(err, file.Close())
 459         if err == nil {
 460                 err = os.Rename(file.Name(), path)
 461         }
 462
 463         if err != nil {
 464                 return errors.Join(fmt.Errorf("replace file: %w", err), os.Remove(file.Name()))
 465         }
 466
 467         return nil
 468 }