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