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

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