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

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