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