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