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

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