[onechat] / memory.go

package main

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "os"
        "path/filepath"
        "strings"
        "sync"
        "time"
)

// errNoView tells that a chat has messages but no saved view.
var errNoView = errors.New("the chat has messages but no saved view")

// Summarizer builds the text of one node with a model.
type Summarizer interface {
        Summarize(ctx context.Context, task Task) (string, error)
}

// Task is the build of the node (Level, Index). Context is the compaction view
// for the build. Input holds the source "kind: text" of a leaf, or the two
// child lines of a merge with their newlines as spaces.
type Task struct {
        Level   int
        Index   int
        Context string
        Input   []string
}

// Memory is the memory of one chat: the message log, the summary tree, the chat
// view, the compaction view, and the queue of the builds. The compaction view
// is the chat view merged further, at the limits of 16 to 32 KB. It takes the
// same new lines, and it merges again in one batch when it is too large or when
// the chat view merges. So the compactions share its prefix in the prompt
// cache. It saves the views to view.json after each change and loads them at
// start. It never builds the views again from the log, because a new view
// differs from the live one and so loses the prompt cache.
//
// The summarizer builds the nodes in the queue in the background. The start
// reads the indexes once to fill the queue again. After that, only a new
// message or a new node adds work. Memory keeps the nodes of the views and the
// tasks of the running builds. It reads all other messages and nodes from their
// day files through the indexes.
type Memory struct {
        mutex      sync.Mutex
        root       string
        log        *Log
        tree       *Tree
        chat       *View
        cview      *View
        queue      *queue
        summarizer Summarizer
        ctx        context.Context
        cancel     context.CancelFunc
        workers    sync.WaitGroup
        changed    chan struct{}
}

// viewFile is the content of view.json.
type viewFile struct {
        Chat       viewState `json:"chat"`
        Compaction viewState `json:"compaction"`
}

// openMemory opens the chat in root and starts the builds that are ready. The
// clock gives the dates of new messages and the days of new nodes. The builds
// stop when ctx ends or at Close.
func openMemory(ctx context.Context, root string, clock func() time.Time, summarizer Summarizer) (*Memory, error) {
        log, err := openLog(ctx, root, clock)
        if err != nil {
                return nil, err
        }

        m := &Memory{
                root:       root,
                log:        log,
                queue:      newQueue(),
                summarizer: summarizer,
                changed:    make(chan struct{}),
        }

        m.ctx, m.cancel = context.WithCancel(ctx)

        m.tree, err = openTree(ctx, root, log.Len(), clock)
        if err != nil {
                return nil, errors.Join(err, m.Close())
        }

        err = m.load()
        if err != nil {
                return nil, errors.Join(fmt.Errorf("load view: %w", err), m.Close())
        }

        err = m.restoreQueue()
        if err != nil {
                return nil, errors.Join(fmt.Errorf("restore build queue: %w", err), m.Close())
        }

        err = m.catchUp()
        if err != nil {
                return nil, errors.Join(err, m.Close())
        }

        m.mutex.Lock()
        m.dispatch()
        m.mutex.Unlock()

        return m, nil
}

// load reads the saved views.
func (m *Memory) load() error {
        data, err := os.ReadFile(m.viewPath())
        if errors.Is(err, os.ErrNotExist) && m.log.Len() == 0 {
                m.chat = newView(chatHigh, chatLow)
                m.cview = newView(compactionHigh, compactionLow)

                return nil
        }

        if errors.Is(err, os.ErrNotExist) {
                return errNoView
        }

        if err != nil {
                return fmt.Errorf("read view: %w", err)
        }

        var file viewFile
        err = json.Unmarshal(data, &file)
        if err != nil {
                return fmt.Errorf("decode view: %w", err)
        }

        m.chat, err = restoreView(file.Chat, chatHigh, chatLow, m.tree)
        if err != nil {
                return fmt.Errorf("chat view: %w", err)
        }

        m.cview, err = restoreView(file.Compaction, compactionHigh, compactionLow, m.tree)
        if err != nil {
                return fmt.Errorf("compaction view: %w", err)
        }

        if m.cview.end != m.chat.end {
                return fmt.Errorf("compaction view covers %d messages, but the chat view %d", m.cview.end, m.chat.end)
        }

        if m.chat.end > m.log.Len() {
                return fmt.Errorf("view covers %d messages, but the log has %d", m.chat.end, m.log.Len())
        }

        return nil
}

// restoreQueue adds the work of the messages in the view to the queue. It
// builds each node that needs no model call and that a crash left out.
func (m *Memory) restoreQueue() error {
        for id := range m.chat.end {
                built, err := m.tree.Built(0, id)
                if err != nil {
                        return err
                }

                if built {
                        continue
                }

                message, err := m.log.Message(id)
                if err != nil {
                        return err
                }

                leaf, free := freeLeaf(message)
                if !free {
                        m.queue.AddLeaf(id)
                        continue
                }

                err = m.build(leaf)
                if err != nil {
                        return err
                }
        }

        keys, err := m.tree.Keys()
        if err != nil {
                return err
        }

        for _, key := range keys {
                if key.index%2 != 0 {
                        continue
                }

                err = m.settle(key)
                if err != nil {
                        return err
                }
        }

        return nil
}

// catchUp adds the messages that the view does not cover yet. A crash can stop
// an append after the log write and before the view save.
func (m *Memory) catchUp() error {
        if m.chat.end == m.log.Len() {
                return nil
        }

        for m.chat.end < m.log.Len() {
                message, err := m.log.Message(m.chat.end)
                if err != nil {
                        return err
                }

                err = m.addMessage(message)
                if err != nil {
                        return err
                }
        }

        return m.save()
}

// Close stops the builds, waits for them, closes the files of the chat, and
// releases its lock.
func (m *Memory) Close() error {
        m.cancel()
        m.workers.Wait()

        m.mutex.Lock()
        defer m.mutex.Unlock()

        var err error
        if m.tree != nil {
                err = m.tree.Close()
        }

        return errors.Join(err, m.log.Close())
}

// Append logs a new message, adds it to the tree and the views, and starts the
// builds that are ready, also the failed ones.
func (m *Memory) Append(kind, text string) (Message, error) {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        message, err := m.log.Append(kind, text)
        if err != nil {
                return Message{}, err
        }

        err = m.addMessage(message)
        if err != nil {
                return Message{}, err
        }

        err = m.save()
        if err != nil {
                return Message{}, err
        }

        m.queue.Retry()
        m.dispatch()

        return message, nil
}

// addMessage builds the leaf of message when it needs no model call, or else
// adds the leaf to the queue. Then it appends the leaf to the views. A crash can
// stop an append after the leaf build and before the view save, so a leaf that
// is built already stays as it is, and restoreQueue settles its parent.
func (m *Memory) addMessage(message Message) error {
        built, err := m.tree.Built(0, message.ID)
        if err != nil {
                return err
        }

        leaf, free := freeLeaf(message)
        switch {
        case built:
        case free:
                err = m.build(leaf)
                if err != nil {
                        return err
                }
        default:
                m.queue.AddLeaf(message.ID)
        }

        merged, err := m.chat.Append(m.tree)
        if err != nil {
                return fmt.Errorf("append to chat view: %w", err)
        }

        _, err = m.cview.Append(m.tree)
        if err == nil && merged {
                err = m.cview.Coarsen(m.chat, m.tree)
        }

        if err != nil {
                return fmt.Errorf("append to compaction view: %w", err)
        }

        return nil
}

// build adds node to the tree, keeps it in the views when it is a pending line,
// and settles its parent.
func (m *Memory) build(node Node) error {
        err := m.tree.Add(node)
        if err != nil {
                return err
        }

        m.chat.Refresh(node)
        m.cview.Refresh(node)

        return m.settle(nodeKey{node.Level, node.Index})
}

// settle looks at the parent of the built node key when both children are
// built. It builds each parent that needs no model call, up the tree, and adds
// the first other parent to the queue.
func (m *Memory) settle(key nodeKey) error {
        for {
                parent := nodeKey{key.level + 1, key.index / 2}
                ready, err := m.ready(parent)
                if err != nil || !ready {
                        return err
                }

                left, right, err := m.children(parent)
                if err != nil {
                        return err
                }

                node, free := freeMerge(left, right)
                if !free {
                        m.queue.AddMerge(parent)
                        return nil
                }

                err = m.tree.Add(node)
                if err != nil {
                        return err
                }

                key = parent
        }
}

// ready tells if both children of the node key are built and the node is not.
func (m *Memory) ready(key nodeKey) (bool, error) {
        for _, index := range []int{2 * key.index, 2*key.index + 1} {
                built, err := m.tree.Built(key.level-1, index)
                if err != nil || !built {
                        return false, err
                }
        }

        built, err := m.tree.Built(key.level, key.index)

        return !built, err
}

// children reads the two children of the node key.
func (m *Memory) children(key nodeKey) (Node, Node, error) {
        var children [2]Node
        for n := range children {
                child, built, err := m.tree.Node(key.level-1, 2*key.index+n)
                if err != nil {
                        return Node{}, Node{}, err
                }

                if !built {
                        return Node{}, Node{}, fmt.Errorf("node (%d, %d) is not built", key.level-1, 2*key.index+n)
                }

                children[n] = child
        }

        return children[0], children[1], nil
}

// dispatch starts the builds that are ready, up to the limit. A build whose
// input cannot be read fails and waits for the next message.
func (m *Memory) dispatch() {
        for m.ctx.Err() == nil {
                key, ready := m.queue.Next()
                if !ready {
                        return
                }

                task, err := m.task(key)
                if err != nil {
                        m.queue.Fail(key)
                        continue
                }

                m.workers.Add(1)
                go m.run(key, task)
        }
}

// task returns the build task of the node key, with its input read from the log
// or the tree.
func (m *Memory) task(key nodeKey) (Task, error) {
        task := Task{
                Level:   key.level,
                Index:   key.index,
                Context: m.context(key.level, key.index),
        }

        if key.level == 0 {
                message, err := m.log.Message(key.index)
                if err != nil {
                        return Task{}, err
                }

                task.Input = []string{leafSource(message)}

                return task, nil
        }

        left, right, err := m.children(key)
        if err != nil {
                return Task{}, err
        }

        for _, child := range []Node{left, right} {
                task.Input = append(task.Input, strings.ReplaceAll(child.Text, "\n", " "))
        }

        return task, nil
}

// run builds one node with the summarizer. A failed build stays in the queue
// until the next message.
func (m *Memory) run(key nodeKey, task Task) {
        defer m.workers.Done()

        text, err := m.summarizer.Summarize(m.ctx, task)

        m.mutex.Lock()
        defer m.mutex.Unlock()

        if err == nil {
                err = m.build(Node{
                        Level: key.level,
                        Index: key.index,
                        Text:  text,
                        Size:  len(text),
                })
        }

        if err != nil {
                m.queue.Fail(key)
        } else {
                m.queue.Done(key)
        }

        close(m.changed)
        m.changed = make(chan struct{})
        m.dispatch()
}

// Wait returns when no build runs. The failed builds wait for the next message.
func (m *Memory) Wait(ctx context.Context) error {
        for {
                m.mutex.Lock()
                running := m.queue.Running()
                changed := m.changed
                m.mutex.Unlock()

                if running == 0 {
                        return nil
                }

                select {
                case <-changed:
                case <-ctx.Done():
                        return ctx.Err()
                }
        }
}

// save writes the views to view.json.
func (m *Memory) save() error {
        data, err := json.Marshal(viewFile{
                Chat:       m.chat.State(),
                Compaction: m.cview.State(),
        })
        if err != nil {
                return fmt.Errorf("encode view: %w", err)
        }

        err = writeFileAtomic(m.viewPath(), append(data, '\n'))
        if err != nil {
                return fmt.Errorf("save view: %w", err)
        }

        return nil
}

func (m *Memory) viewPath() string {
        return filepath.Join(m.root, "view.json")
}

// View returns the chat view as the model reads it.
func (m *Memory) View() string {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        return m.chat.Render()
}

// Context returns the compaction view for the build of node (l, i). It ends
// before the message of a leaf, or after the last message of a merge, and it
// stops at the first line that is not built.
func (m *Memory) Context(level, index int) string {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        return m.context(level, index)
}

func (m *Memory) context(level, index int) string {
        end := index
        if level > 0 {
                end = nodeEnd(level, index)
        }

        return m.cview.Context(end)
}