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 // Each write goes to the operating system at once, and a syncer makes the
  43 // writes durable by group commit in the background. Flush waits until the
  44 // writes so far are durable.
  45 //
  46 // The summarizer builds the nodes in the queue in the background. The start
  47 // reads the indexes once to fill the queue again. After that, only a new
  48 // message or a new node adds work. Memory keeps the nodes of the views and the
  49 // tasks of the running builds. It reads all other messages and nodes from their
  50 // day files through the indexes.
  51 type Memory struct {
  52         mutex      sync.Mutex
  53         root       string
  54         log        *Log
  55         syncer     *syncer
  56         tree       *Tree
  57         chat       *View
  58         cview      *View
  59         queue      *queue
  60         summarizer Summarizer
  61         ctx        context.Context
  62         cancel     context.CancelFunc
  63         workers    sync.WaitGroup
  64         changed    chan struct{}
  65 }
  66
  67 // viewFile is the content of view.json.
  68 type viewFile struct {
  69         Chat       viewState `json:"chat"`
  70         Compaction viewState `json:"compaction"`
  71 }
  72
  73 // openMemory opens the chat in root and starts the builds that are ready. The
  74 // clock gives the dates of new messages and the days of new nodes. The builds
  75 // stop when ctx ends or at Close.
  76 func openMemory(ctx context.Context, root string, clock func() time.Time, summarizer Summarizer) (*Memory, error) {
  77         log, err := openLog(ctx, root, clock)
  78         if err != nil {
  79                 return nil, err
  80         }
  81
  82         m := &Memory{
  83                 root:       root,
  84                 log:        log,
  85                 queue:      newQueue(),
  86                 summarizer: summarizer,
  87                 changed:    make(chan struct{}),
  88         }
  89
  90         m.ctx, m.cancel = context.WithCancel(ctx)
  91
  92         m.syncer, err = openSyncer(root)
  93         if err != nil {
  94                 return nil, errors.Join(err, m.Close())
  95         }
  96
  97         m.tree, err = openTree(ctx, root, log.Len(), clock)
  98         if err != nil {
  99                 return nil, errors.Join(err, m.Close())
 100         }
 101
 102         err = m.load()
 103         if err != nil {
 104                 return nil, errors.Join(fmt.Errorf("load view: %w", err), m.Close())
 105         }
 106
 107         err = m.restoreQueue()
 108         if err != nil {
 109                 return nil, errors.Join(fmt.Errorf("restore build queue: %w", err), m.Close())
 110         }
 111
 112         err = m.catchUp()
 113         if err != nil {
 114                 return nil, errors.Join(err, m.Close())
 115         }
 116
 117         m.mutex.Lock()
 118         m.syncer.Written()
 119         m.dispatch()
 120         m.mutex.Unlock()
 121
 122         return m, nil
 123 }
 124
 125 // load reads the saved views.
 126 func (m *Memory) load() error {
 127         data, err := os.ReadFile(m.viewPath())
 128         if errors.Is(err, os.ErrNotExist) && m.log.Len() == 0 {
 129                 m.chat = newView(chatHigh, chatLow)
 130                 m.cview = newView(compactionHigh, compactionLow)
 131
 132                 return nil
 133         }
 134
 135         if errors.Is(err, os.ErrNotExist) {
 136                 return errNoView
 137         }
 138
 139         if err != nil {
 140                 return fmt.Errorf("read view: %w", err)
 141         }
 142
 143         var file viewFile
 144         err = json.Unmarshal(data, &file)
 145         if err != nil {
 146                 return fmt.Errorf("decode view: %w", err)
 147         }
 148
 149         m.chat, err = restoreView(file.Chat, chatHigh, chatLow, m.tree)
 150         if err != nil {
 151                 return fmt.Errorf("chat view: %w", err)
 152         }
 153
 154         m.cview, err = restoreView(file.Compaction, compactionHigh, compactionLow, m.tree)
 155         if err != nil {
 156                 return fmt.Errorf("compaction view: %w", err)
 157         }
 158
 159         if m.cview.end != m.chat.end {
 160                 return fmt.Errorf("compaction view covers %d messages, but the chat view %d", m.cview.end, m.chat.end)
 161         }
 162
 163         if m.chat.end > m.log.Len() {
 164                 return fmt.Errorf("view covers %d messages, but the log has %d", m.chat.end, m.log.Len())
 165         }
 166
 167         return nil
 168 }
 169
 170 // restoreQueue adds the work of the messages in the view to the queue. It
 171 // builds each node that needs no model call and that a crash left out.
 172 func (m *Memory) restoreQueue() error {
 173         for id := range m.chat.end {
 174                 built, err := m.tree.Built(0, id)
 175                 if err != nil {
 176                         return err
 177                 }
 178
 179                 if built {
 180                         continue
 181                 }
 182
 183                 message, err := m.log.Message(id)
 184                 if err != nil {
 185                         return err
 186                 }
 187
 188                 leaf, free := freeLeaf(message)
 189                 if !free {
 190                         m.queue.AddLeaf(id)
 191                         continue
 192                 }
 193
 194                 err = m.build(leaf)
 195                 if err != nil {
 196                         return err
 197                 }
 198         }
 199
 200         keys, err := m.tree.Keys()
 201         if err != nil {
 202                 return err
 203         }
 204
 205         for _, key := range keys {
 206                 if key.index%2 != 0 {
 207                         continue
 208                 }
 209
 210                 err = m.settle(key)
 211                 if err != nil {
 212                         return err
 213                 }
 214         }
 215
 216         return nil
 217 }
 218
 219 // catchUp adds the messages that the view does not cover yet. A crash can stop
 220 // an append after the log write and before the view save.
 221 func (m *Memory) catchUp() error {
 222         if m.chat.end == m.log.Len() {
 223                 return nil
 224         }
 225
 226         for m.chat.end < m.log.Len() {
 227                 message, err := m.log.Message(m.chat.end)
 228                 if err != nil {
 229                         return err
 230                 }
 231
 232                 err = m.addMessage(message)
 233                 if err != nil {
 234                         return err
 235                 }
 236         }
 237
 238         return m.save()
 239 }
 240
 241 // Close stops the builds, waits for them, makes the writes durable, closes the
 242 // files of the chat, and releases its lock.
 243 func (m *Memory) Close() error {
 244         m.cancel()
 245         m.workers.Wait()
 246
 247         m.mutex.Lock()
 248         defer m.mutex.Unlock()
 249
 250         var err error
 251         if m.syncer != nil {
 252                 err = m.syncer.Close()
 253         }
 254
 255         if m.tree != nil {
 256                 err = errors.Join(err, m.tree.Close())
 257         }
 258
 259         return errors.Join(err, m.log.Close())
 260 }
 261
 262 // Append logs a new message, adds it to the tree and the views, and starts the
 263 // builds that are ready, also the failed ones.
 264 func (m *Memory) Append(kind, text string) (Message, error) {
 265         m.mutex.Lock()
 266         defer m.mutex.Unlock()
 267
 268         message, err := m.log.Append(kind, text)
 269         if err != nil {
 270                 return Message{}, err
 271         }
 272
 273         defer m.syncer.Written()
 274
 275         err = m.addMessage(message)
 276         if err != nil {
 277                 return Message{}, err
 278         }
 279
 280         err = m.save()
 281         if err != nil {
 282                 return Message{}, err
 283         }
 284
 285         m.queue.Retry()
 286         m.dispatch()
 287
 288         return message, nil
 289 }
 290
 291 // addMessage builds the leaf of message when it needs no model call, or else
 292 // adds the leaf to the queue. Then it appends the leaf to the views. A crash can
 293 // stop an append after the leaf build and before the view save, so a leaf that
 294 // is built already stays as it is, and restoreQueue settles its parent.
 295 func (m *Memory) addMessage(message Message) error {
 296         built, err := m.tree.Built(0, message.ID)
 297         if err != nil {
 298                 return err
 299         }
 300
 301         leaf, free := freeLeaf(message)
 302         switch {
 303         case built:
 304         case free:
 305                 err = m.build(leaf)
 306                 if err != nil {
 307                         return err
 308                 }
 309         default:
 310                 m.queue.AddLeaf(message.ID)
 311         }
 312
 313         merged, err := m.chat.Append(m.tree)
 314         if err != nil {
 315                 return fmt.Errorf("append to chat view: %w", err)
 316         }
 317
 318         _, err = m.cview.Append(m.tree)
 319         if err == nil && merged {
 320                 err = m.cview.Coarsen(m.chat, m.tree)
 321         }
 322
 323         if err != nil {
 324                 return fmt.Errorf("append to compaction view: %w", err)
 325         }
 326
 327         return nil
 328 }
 329
 330 // build adds node to the tree, keeps it in the views when it is a pending line,
 331 // and settles its parent.
 332 func (m *Memory) build(node Node) error {
 333         err := m.tree.Add(node)
 334         if err != nil {
 335                 return err
 336         }
 337
 338         m.chat.Refresh(node)
 339         m.cview.Refresh(node)
 340
 341         return m.settle(nodeKey{node.Level, node.Index})
 342 }
 343
 344 // settle looks at the parent of the built node key when both children are
 345 // built. It builds each parent that needs no model call, up the tree, and adds
 346 // the first other parent to the queue.
 347 func (m *Memory) settle(key nodeKey) error {
 348         for {
 349                 parent := nodeKey{key.level + 1, key.index / 2}
 350                 ready, err := m.ready(parent)
 351                 if err != nil || !ready {
 352                         return err
 353                 }
 354
 355                 left, right, err := m.children(parent)
 356                 if err != nil {
 357                         return err
 358                 }
 359
 360                 node, free := freeMerge(left, right)
 361                 if !free {
 362                         m.queue.AddMerge(parent)
 363                         return nil
 364                 }
 365
 366                 err = m.tree.Add(node)
 367                 if err != nil {
 368                         return err
 369                 }
 370
 371                 key = parent
 372         }
 373 }
 374
 375 // ready tells if both children of the node key are built and the node is not.
 376 func (m *Memory) ready(key nodeKey) (bool, error) {
 377         for _, index := range []int{2 * key.index, 2*key.index + 1} {
 378                 built, err := m.tree.Built(key.level-1, index)
 379                 if err != nil || !built {
 380                         return false, err
 381                 }
 382         }
 383
 384         built, err := m.tree.Built(key.level, key.index)
 385
 386         return !built, err
 387 }
 388
 389 // children reads the two children of the node key.
 390 func (m *Memory) children(key nodeKey) (Node, Node, error) {
 391         var children [2]Node
 392         for n := range children {
 393                 child, built, err := m.tree.Node(key.level-1, 2*key.index+n)
 394                 if err != nil {
 395                         return Node{}, Node{}, err
 396                 }
 397
 398                 if !built {
 399                         return Node{}, Node{}, fmt.Errorf("node (%d, %d) is not built", key.level-1, 2*key.index+n)
 400                 }
 401
 402                 children[n] = child
 403         }
 404
 405         return children[0], children[1], nil
 406 }
 407
 408 // dispatch starts the builds that are ready, up to the limit. A build whose
 409 // input cannot be read fails and waits for the next message.
 410 func (m *Memory) dispatch() {
 411         for m.ctx.Err() == nil {
 412                 key, ready := m.queue.Next()
 413                 if !ready {
 414                         return
 415                 }
 416
 417                 task, err := m.task(key)
 418                 if err != nil {
 419                         m.queue.Fail(key)
 420                         continue
 421                 }
 422
 423                 m.workers.Add(1)
 424                 go m.run(key, task)
 425         }
 426 }
 427
 428 // task returns the build task of the node key, with its input read from the log
 429 // or the tree.
 430 func (m *Memory) task(key nodeKey) (Task, error) {
 431         task := Task{
 432                 Level:   key.level,
 433                 Index:   key.index,
 434                 Context: m.context(key.level, key.index),
 435         }
 436
 437         if key.level == 0 {
 438                 message, err := m.log.Message(key.index)
 439                 if err != nil {
 440                         return Task{}, err
 441                 }
 442
 443                 task.Input = []string{leafSource(message)}
 444
 445                 return task, nil
 446         }
 447
 448         left, right, err := m.children(key)
 449         if err != nil {
 450                 return Task{}, err
 451         }
 452
 453         for _, child := range []Node{left, right} {
 454                 task.Input = append(task.Input, strings.ReplaceAll(child.Text, "\n", " "))
 455         }
 456
 457         return task, nil
 458 }
 459
 460 // run builds one node with the summarizer. A failed build stays in the queue
 461 // until the next message.
 462 func (m *Memory) run(key nodeKey, task Task) {
 463         defer m.workers.Done()
 464
 465         text, err := m.summarizer.Summarize(m.ctx, task)
 466
 467         m.mutex.Lock()
 468         defer m.mutex.Unlock()
 469
 470         if err == nil {
 471                 err = m.build(Node{
 472                         Level: key.level,
 473                         Index: key.index,
 474                         Text:  text,
 475                         Size:  len(text),
 476                 })
 477                 m.syncer.Written()
 478         }
 479
 480         if err != nil {
 481                 m.queue.Fail(key)
 482         } else {
 483                 m.queue.Done(key)
 484         }
 485
 486         close(m.changed)
 487         m.changed = make(chan struct{})
 488         m.dispatch()
 489 }
 490
 491 // Wait returns when no build runs. The failed builds wait for the next message.
 492 func (m *Memory) Wait(ctx context.Context) error {
 493         for {
 494                 m.mutex.Lock()
 495                 running := m.queue.Running()
 496                 changed := m.changed
 497                 m.mutex.Unlock()
 498
 499                 if running == 0 {
 500                         return nil
 501                 }
 502
 503                 select {
 504                 case <-changed:
 505                 case <-ctx.Done():
 506                         return ctx.Err()
 507                 }
 508         }
 509 }
 510
 511 // Flush returns when every write before the call is durable.
 512 func (m *Memory) Flush(ctx context.Context) error {
 513         return m.syncer.Flush(ctx)
 514 }
 515
 516 // save writes the views to view.json.
 517 func (m *Memory) save() error {
 518         data, err := json.Marshal(viewFile{
 519                 Chat:       m.chat.State(),
 520                 Compaction: m.cview.State(),
 521         })
 522         if err != nil {
 523                 return fmt.Errorf("encode view: %w", err)
 524         }
 525
 526         err = writeFileAtomic(m.viewPath(), append(data, '\n'))
 527         if err != nil {
 528                 return fmt.Errorf("save view: %w", err)
 529         }
 530
 531         return nil
 532 }
 533
 534 func (m *Memory) viewPath() string {
 535         return filepath.Join(m.root, "view.json")
 536 }
 537
 538 // View returns the chat view as the model reads it.
 539 func (m *Memory) View() string {
 540         m.mutex.Lock()
 541         defer m.mutex.Unlock()
 542
 543         return m.chat.Render()
 544 }
 545
 546 // ViewLines returns the lines of the chat view, each with its newline, without
 547 // the leaf lines of the messages in skip.
 548 func (m *Memory) ViewLines(skip map[int]bool) []string {
 549         m.mutex.Lock()
 550         defer m.mutex.Unlock()
 551
 552         return m.chat.LinesWithout(skip)
 553 }
 554
 555 // Len returns the number of messages in the log.
 556 func (m *Memory) Len() int {
 557         m.mutex.Lock()
 558         defer m.mutex.Unlock()
 559
 560         return m.log.Len()
 561 }
 562
 563 // Message reads the message with the given ID.
 564 func (m *Memory) Message(id int) (Message, error) {
 565         m.mutex.Lock()
 566         defer m.mutex.Unlock()
 567
 568         return m.log.Message(id)
 569 }
 570
 571 // Node reads the node (l, i) if it is built.
 572 func (m *Memory) Node(level, index int) (Node, bool, error) {
 573         m.mutex.Lock()
 574         defer m.mutex.Unlock()
 575
 576         return m.tree.Node(level, index)
 577 }
 578
 579 // Context returns the compaction view for the build of node (l, i). It ends
 580 // before the message of a leaf, or after the last message of a merge, and it
 581 // stops at the first line that is not built.
 582 func (m *Memory) Context(level, index int) string {
 583         m.mutex.Lock()
 584         defer m.mutex.Unlock()
 585
 586         return m.context(level, index)
 587 }
 588
 589 func (m *Memory) context(level, index int) string {
 590         end := index
 591         if level > 0 {
 592                 end = nodeEnd(level, index)
 593         }
 594
 595         return m.cview.Context(end)
 596 }