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.
 276 func (m *Memory) addMessage(message Message) error {
 277         leaf, free := freeLeaf(message)
 278         if free {
 279                 err := m.build(leaf)
 280                 if err != nil {
 281                         return err
 282                 }
 283         } else {
 284                 m.queue.AddLeaf(message.ID)
 285         }
 286
 287         merged, err := m.chat.Append(m.tree)
 288         if err != nil {
 289                 return fmt.Errorf("append to chat view: %w", err)
 290         }
 291
 292         _, err = m.cview.Append(m.tree)
 293         if err == nil && merged {
 294                 err = m.cview.Coarsen(m.chat, m.tree)
 295         }
 296
 297         if err != nil {
 298                 return fmt.Errorf("append to compaction view: %w", err)
 299         }
 300
 301         return nil
 302 }
 303
 304 // build adds node to the tree, keeps it in the views when it is a pending line,
 305 // and settles its parent.
 306 func (m *Memory) build(node Node) error {
 307         err := m.tree.Add(node)
 308         if err != nil {
 309                 return err
 310         }
 311
 312         m.chat.Refresh(node)
 313         m.cview.Refresh(node)
 314
 315         return m.settle(nodeKey{node.Level, node.Index})
 316 }
 317
 318 // settle looks at the parent of the built node key when both children are
 319 // built. It builds each parent that needs no model call, up the tree, and adds
 320 // the first other parent to the queue.
 321 func (m *Memory) settle(key nodeKey) error {
 322         for {
 323                 parent := nodeKey{key.level + 1, key.index / 2}
 324                 ready, err := m.ready(parent)
 325                 if err != nil || !ready {
 326                         return err
 327                 }
 328
 329                 left, right, err := m.children(parent)
 330                 if err != nil {
 331                         return err
 332                 }
 333
 334                 node, free := freeMerge(left, right)
 335                 if !free {
 336                         m.queue.AddMerge(parent)
 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 // ready tells if both children of the node key are built and the node is not.
 350 func (m *Memory) ready(key nodeKey) (bool, error) {
 351         for _, index := range []int{2 * key.index, 2*key.index + 1} {
 352                 built, err := m.tree.Built(key.level-1, index)
 353                 if err != nil || !built {
 354                         return false, err
 355                 }
 356         }
 357
 358         built, err := m.tree.Built(key.level, key.index)
 359
 360         return !built, err
 361 }
 362
 363 // children reads the two children of the node key.
 364 func (m *Memory) children(key nodeKey) (Node, Node, error) {
 365         var children [2]Node
 366         for n := range children {
 367                 child, built, err := m.tree.Node(key.level-1, 2*key.index+n)
 368                 if err != nil {
 369                         return Node{}, Node{}, err
 370                 }
 371
 372                 if !built {
 373                         return Node{}, Node{}, fmt.Errorf("node (%d, %d) is not built", key.level-1, 2*key.index+n)
 374                 }
 375
 376                 children[n] = child
 377         }
 378
 379         return children[0], children[1], nil
 380 }
 381
 382 // dispatch starts the builds that are ready, up to the limit. A build whose
 383 // input cannot be read fails and waits for the next message.
 384 func (m *Memory) dispatch() {
 385         for m.ctx.Err() == nil {
 386                 key, ready := m.queue.Next()
 387                 if !ready {
 388                         return
 389                 }
 390
 391                 task, err := m.task(key)
 392                 if err != nil {
 393                         m.queue.Fail(key)
 394                         continue
 395                 }
 396
 397                 m.workers.Add(1)
 398                 go m.run(key, task)
 399         }
 400 }
 401
 402 // task returns the build task of the node key, with its input read from the log
 403 // or the tree.
 404 func (m *Memory) task(key nodeKey) (Task, error) {
 405         task := Task{
 406                 Level:   key.level,
 407                 Index:   key.index,
 408                 Context: m.context(key.level, key.index),
 409         }
 410
 411         if key.level == 0 {
 412                 message, err := m.log.Message(key.index)
 413                 if err != nil {
 414                         return Task{}, err
 415                 }
 416
 417                 task.Input = []string{leafSource(message)}
 418
 419                 return task, nil
 420         }
 421
 422         left, right, err := m.children(key)
 423         if err != nil {
 424                 return Task{}, err
 425         }
 426
 427         for _, child := range []Node{left, right} {
 428                 task.Input = append(task.Input, strings.ReplaceAll(child.Text, "\n", " "))
 429         }
 430
 431         return task, nil
 432 }
 433
 434 // run builds one node with the summarizer. A failed build stays in the queue
 435 // until the next message.
 436 func (m *Memory) run(key nodeKey, task Task) {
 437         defer m.workers.Done()
 438
 439         text, err := m.summarizer.Summarize(m.ctx, task)
 440
 441         m.mutex.Lock()
 442         defer m.mutex.Unlock()
 443
 444         if err == nil {
 445                 err = m.build(Node{
 446                         Level: key.level,
 447                         Index: key.index,
 448                         Text:  text,
 449                         Size:  len(text),
 450                 })
 451         }
 452
 453         if err != nil {
 454                 m.queue.Fail(key)
 455         } else {
 456                 m.queue.Done(key)
 457         }
 458
 459         close(m.changed)
 460         m.changed = make(chan struct{})
 461         m.dispatch()
 462 }
 463
 464 // Wait returns when no build runs. The failed builds wait for the next message.
 465 func (m *Memory) Wait(ctx context.Context) error {
 466         for {
 467                 m.mutex.Lock()
 468                 running := m.queue.Running()
 469                 changed := m.changed
 470                 m.mutex.Unlock()
 471
 472                 if running == 0 {
 473                         return nil
 474                 }
 475
 476                 select {
 477                 case <-changed:
 478                 case <-ctx.Done():
 479                         return ctx.Err()
 480                 }
 481         }
 482 }
 483
 484 // save writes the views to view.json.
 485 func (m *Memory) save() error {
 486         data, err := json.Marshal(viewFile{
 487                 Chat:       m.chat.State(),
 488                 Compaction: m.cview.State(),
 489         })
 490         if err != nil {
 491                 return fmt.Errorf("encode view: %w", err)
 492         }
 493
 494         err = writeFileAtomic(m.viewPath(), append(data, '\n'))
 495         if err != nil {
 496                 return fmt.Errorf("save view: %w", err)
 497         }
 498
 499         return nil
 500 }
 501
 502 func (m *Memory) viewPath() string {
 503         return filepath.Join(m.root, "view.json")
 504 }
 505
 506 // View returns the chat view as the model reads it.
 507 func (m *Memory) View() string {
 508         m.mutex.Lock()
 509         defer m.mutex.Unlock()
 510
 511         return m.chat.Render()
 512 }
 513
 514 // Context returns the compaction view for the build of node (l, i). It ends
 515 // before the message of a leaf, or after the last message of a merge, and it
 516 // stops at the first line that is not built.
 517 func (m *Memory) Context(level, index int) string {
 518         m.mutex.Lock()
 519         defer m.mutex.Unlock()
 520
 521         return m.context(level, index)
 522 }
 523
 524 func (m *Memory) context(level, index int) string {
 525         end := index
 526         if level > 0 {
 527                 end = nodeEnd(level, index)
 528         }
 529
 530         return m.cview.Context(end)
 531 }
 532
 533 // writeFileAtomic replaces the file at path with data through a temporary file
 534 // and a rename. A reader or a crash of the process sees the old file or the new
 535 // file, never a part of one.
 536 func writeFileAtomic(path string, data []byte) error {
 537         file, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".*")
 538         if err != nil {
 539                 return fmt.Errorf("create temporary file: %w", err)
 540         }
 541
 542         _, err = file.Write(data)
 543         err = errors.Join(err, file.Close())
 544         if err == nil {
 545                 err = os.Rename(file.Name(), path)
 546         }
 547
 548         if err != nil {
 549                 return errors.Join(fmt.Errorf("replace file: %w", err), os.Remove(file.Name()))
 550         }
 551
 552         return nil
 553 }