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