"encoding/json"
"errors"
"fmt"
+ "log/slog"
"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")
+var (
+ // errMemoryNoView tells that a chat has messages but no saved view.
+ errMemoryNoView = errors.New("the chat has messages but no saved view")
-// Memory is the memory of one chat: the message log, the summary tree, the chat
-// view, and the compaction view. 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.
-type Memory struct {
- mutex sync.Mutex
- root string
- log *Log
- tree *Tree
- chat *View
- cview *View
-}
-
-// viewFile is the content of view.json.
-type viewFile struct {
- Chat viewState `json:"chat"`
- Compaction viewState `json:"compaction"`
-}
+ // errMemoryUnsettled tells that a line of the view stays pending,
+ // because its build failed.
+ errMemoryUnsettled = errors.New("a summary of the view failed, and it is tried again at the next message")
+)
-// openMemory opens the chat in root. The clock gives the dates of new messages
-// and the days of new nodes.
-func openMemory(ctx context.Context, root string, clock func() time.Time) (*Memory, error) {
+// 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. logger gets
+// one line for each failed build. The builds stop when ctx ends or at Close.
+func openMemory(ctx context.Context, root string, clock func() time.Time, summarizer Summarizer, logger *slog.Logger) (*Memory, error) {
log, err := openLog(ctx, root, clock)
if err != nil {
return nil, err
}
m := &Memory{
- root: root,
- log: log,
+ root: root,
+ log: log,
+ queue: newBuildQueue(),
+ summarizer: summarizer,
+ logger: logger,
+ changed: make(chan struct{}),
+ }
+
+ m.ctx, m.cancel = context.WithCancel(ctx)
+
+ m.syncer, err = openSyncer(root)
+ if err != nil {
+ return nil, errors.Join(err, m.Close())
}
m.tree, err = openTree(ctx, root, log.Len(), clock)
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.syncer.Written()
+ m.dispatch()
+ m.mutex.Unlock()
+
return m, nil
}
-// load reads the saved view. A crash can stop an append after the log write and
-// before the view save, so load continues the view with the messages that it
-// does not cover yet.
+// Summarizer builds the text of one node with a model.
+type Summarizer interface {
+ Summarize(ctx context.Context, task Task) (string, error)
+}
+
+// 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.
+//
+// Each write goes to the operating system at once, and a syncer makes the
+// writes durable by group commit in the background. Flush waits until the
+// writes so far are durable.
+//
+// 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
+ syncer *Syncer
+ tree *Tree
+ chat *View
+ cview *View
+ queue *BuildQueue
+ summarizer Summarizer
+ logger *slog.Logger
+ ctx context.Context
+ cancel context.CancelFunc
+ workers sync.WaitGroup
+ changed chan struct{}
+}
+
+// Close stops the builds, waits for them, makes the writes durable, 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.syncer != nil {
+ err = m.syncer.Close()
+ }
+
+ if m.tree != nil {
+ err = errors.Join(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) {
+ return m.AppendDated(kind, text, "")
+}
+
+// AppendDated appends a message with its RFC 3339 date and builds its nodes. An
+// empty date uses the clock, as Append does.
+func (m *Memory) AppendDated(kind, text, date string) (Message, error) {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ message, err := m.log.AppendDated(kind, text, date)
+ if err != nil {
+ return Message{}, fmt.Errorf("append dated log message: %w", err)
+ }
+
+ defer m.syncer.Written()
+
+ err = m.addMessage(message)
+ if err != nil {
+ return Message{}, fmt.Errorf("add dated message to memory: %w", err)
+ }
+
+ err = m.save()
+ if err != nil {
+ return Message{}, fmt.Errorf("save dated message view: %w", err)
+ }
+
+ m.queue.Retry()
+ m.dispatch()
+
+ return message, nil
+}
+
+// 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
+}
+
+// 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()
+ }
+ }
+}
+
+// Settle returns when the chat view has no pending line, so that every message
+// so far is summarized. A pending line with no build that runs belongs to a
+// failed build, which waits for the next message, so Settle returns
+// errMemoryUnsettled in place of a wait that never ends.
+func (m *Memory) Settle(ctx context.Context) error {
+ for {
+ m.mutex.Lock()
+
+ pending := m.chat.Pending()
+ running := m.queue.Running()
+ changed := m.changed
+ m.mutex.Unlock()
+
+ if !pending {
+ return nil
+ }
+
+ if running == 0 {
+ return errMemoryUnsettled
+ }
+
+ select {
+ case <-changed:
+ case <-ctx.Done():
+ return ctx.Err()
+ }
+ }
+}
+
+// Flush returns when every write before the call is durable.
+func (m *Memory) Flush(ctx context.Context) error {
+ return m.syncer.Flush(ctx)
+}
+
+// Snapshot flushes all memory writes and runs snapshot while no append or
+// summary write can change the log, indexes, tree, or views. snapshot must not
+// call Memory.
+func (m *Memory) Snapshot(ctx context.Context, snapshot func() error) error {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ err := m.Flush(ctx)
+ if err != nil {
+ return fmt.Errorf("flush memory snapshot: %w", err)
+ }
+
+ err = snapshot()
+ if err != nil {
+ return fmt.Errorf("copy memory snapshot: %w", err)
+ }
+
+ return nil
+}
+
+// ViewLines returns the lines of the chat view, each with its newline, without
+// the leaf lines of the messages in skip.
+func (m *Memory) ViewLines(skip map[int]bool) []string {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ return m.chat.LinesWithout(skip)
+}
+
+// Len returns the number of messages in the log.
+func (m *Memory) Len() int {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ return m.log.Len()
+}
+
+// Message reads the message with the given ID.
+func (m *Memory) Message(id int) (Message, error) {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ return m.log.Message(id)
+}
+
+// Node reads the node (l, i) if it is built.
+func (m *Memory) Node(level, index int) (Node, bool, error) {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ return m.tree.Node(level, index)
+}
+
+// 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)
+ m.chat = newView(viewChatHigh, viewChatLow)
+ m.cview = newView(viewCompactionHigh, viewCompactionLow)
return nil
}
if errors.Is(err, os.ErrNotExist) {
- return errNoView
+ return errMemoryNoView
}
if err != nil {
return fmt.Errorf("read view: %w", err)
}
- var file viewFile
+ var file MemoryViewFile
+
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)
+ m.chat, err = restoreView(file.Chat, viewChatHigh, viewChatLow, m.tree)
if err != nil {
return fmt.Errorf("chat view: %w", err)
}
- m.cview, err = restoreView(file.Compaction, compactionHigh, compactionLow, m.tree)
+ m.cview, err = restoreView(file.Compaction, viewCompactionHigh, viewCompactionLow, m.tree)
if err != nil {
return fmt.Errorf("compaction view: %w", err)
}
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() {
- err = m.addMessage(m.log.Message(m.chat.end))
+ 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 closes the files of the chat and releases its lock.
-func (m *Memory) Close() error {
- m.mutex.Lock()
- defer m.mutex.Unlock()
+// 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
+ }
- var err error
- if m.tree != nil {
- err = m.tree.Close()
- }
+ task, err := m.task(key)
+ if err != nil {
+ m.failed(key, fmt.Errorf("read its input: %w", err))
+ m.queue.Fail(key)
- return errors.Join(err, m.log.Close())
-}
+ continue
+ }
-// Append logs a new message and adds it to the tree and the view.
-func (m *Memory) Append(kind, text string) (Message, error) {
- m.mutex.Lock()
- defer m.mutex.Unlock()
+ m.workers.Add(1)
+
+ go m.run(key, task)
+ }
+}
- message, err := m.log.Append(kind, text)
+// 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 Message{}, err
+ return err
}
- err = m.addMessage(message)
+ 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 Message{}, err
+ 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)
}
- err = m.save()
if err != nil {
- return Message{}, err
+ return fmt.Errorf("append to compaction view: %w", err)
}
- return message, nil
+ return nil
}
-// addMessage builds the leaf of message when it needs no model call and appends
-// it to the view.
-func (m *Memory) addMessage(message Message) error {
- _, built := m.tree.Node(0, message.ID)
- leaf, free := freeLeaf(message)
- if free && !built {
- err := m.build(leaf)
- if err != nil {
- return err
- }
+// save writes the views to view.json.
+func (m *Memory) save() error {
+ data, err := json.Marshal(MemoryViewFile{
+ Chat: m.chat.State(),
+ Compaction: m.cview.State(),
+ })
+ if err != nil {
+ return fmt.Errorf("encode view: %w", err)
}
- merged := m.chat.Append(m.tree)
- m.cview.Append(m.tree)
- if merged {
- m.cview.Coarsen(m.chat, m.tree)
+ err = writeFileAtomic(m.viewPath(), append(data, '\n'))
+ if err != nil {
+ return fmt.Errorf("save view: %w", err)
}
return nil
}
-// build adds node to the tree, and then each parent above it that needs no
-// model call.
+func (m *Memory) viewPath() string {
+ return filepath.Join(m.root, "view.json")
+}
+
+// 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 {
- err := m.tree.Add(node)
+ 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
}
- sibling, built := m.tree.Node(node.Level, node.Index^1)
- if !built {
+ node, free := freeMerge(left, right)
+ if !free {
+ m.queue.AddMerge(parent)
+
return nil
}
- left, right := node, sibling
- if node.Index%2 != 0 {
- left, right = sibling, node
+ err = m.tree.Add(node)
+ if err != nil {
+ return err
}
- parent, free := freeMerge(left, right)
- if !free {
- return nil
+ key = parent
+ }
+}
+
+// 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
}
- node = parent
+ task.Input = []string{leafSource(message)}
+
+ return task, nil
}
-}
-// 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(),
- })
+ left, right, err := m.children(key)
if err != nil {
- return fmt.Errorf("encode view: %w", err)
+ 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),
+ })
+ m.syncer.Written()
+ }
+
+ if err != nil && m.ctx.Err() == nil {
+ m.failed(key, err)
}
- err = writeFileAtomic(m.viewPath(), append(data, '\n'))
if err != nil {
- return fmt.Errorf("save view: %w", err)
+ m.queue.Fail(key)
+ } else {
+ m.queue.Done(key)
}
- return nil
+ close(m.changed)
+
+ m.changed = make(chan struct{})
+ m.dispatch()
}
-func (m *Memory) viewPath() string {
- return filepath.Join(m.root, "view.json")
+// failed logs the failed build of the node key with its cause err.
+func (m *Memory) failed(key NodeKey, err error) {
+ m.logger.Warn("summary failed", "level", key.level, "index", key.index, "err", err)
+}
+
+// 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
+}
+
+func (m *Memory) context(level, index int) string {
+ end := index
+ if level > 0 {
+ end = nodeEnd(level, index)
+ }
+
+ return m.cview.Context(end)
}
// View returns the chat view as the model reads it.
m.mutex.Lock()
defer m.mutex.Unlock()
- return m.chat.Render(m.tree)
+ return m.chat.Render()
}
// Context returns the compaction view for the build of node (l, i). It ends
return m.context(level, index)
}
-func (m *Memory) context(level, index int) string {
- end := index
- if level > 0 {
- end = nodeEnd(level, index)
+// rebuildComplete checks the whole tree, not only the visible view. A compactor
+// can fail without an HTTP error, for example with an empty line. Such a failed
+// merge can be absent from the view and must prevent a swap.
+func (m *Memory) rebuildComplete() error {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ if len(m.queue.leaves) != 0 || len(m.queue.queued) != 0 {
+ return errors.New("rebuild failed: summary tree has unfinished nodes")
}
- return m.cview.Context(m.tree, end)
+ return nil
}
-// writeFileAtomic replaces the file at path with data. A reader or a crash sees
-// the old file or the new file, never a part of one.
-func writeFileAtomic(path string, data []byte) error {
- dir := filepath.Dir(path)
- file, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*")
+// buildModelLines returns the sizes of the nodes that the model wrote, in tree
+// order. A node that fits its source with no model call is not one of them.
+func (m *Memory) buildModelLines() ([]int, error) {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ keys, err := m.tree.Keys()
if err != nil {
- return fmt.Errorf("create temporary file: %w", err)
+ return nil, fmt.Errorf("list built nodes: %w", err)
}
- _, err = file.Write(data)
- if err == nil {
- err = file.Sync()
- }
+ var sizes []int
- err = errors.Join(err, file.Close())
- if err == nil {
- err = os.Rename(file.Name(), path)
- }
+ for _, key := range keys {
+ var node Node
- if err != nil {
- return errors.Join(fmt.Errorf("replace file: %w", err), os.Remove(file.Name()))
+ node, _, err = m.tree.Node(key.level, key.index)
+ if err != nil {
+ return nil, fmt.Errorf("read built node: %w", err)
+ }
+
+ var model bool
+
+ model, err = m.buildBuiltByModel(key)
+ if err != nil {
+ return nil, fmt.Errorf("check model source: %w", err)
+ }
+
+ if model {
+ sizes = append(sizes, node.Size)
+ }
}
- return syncDir(dir)
+ return sizes, nil
}
-// syncDir flushes the entries of dir to the disk.
-func syncDir(dir string) (err error) {
- handle, err := os.Open(dir)
- if err != nil {
- return fmt.Errorf("open directory: %w", err)
- }
+// buildBuiltByModel tells if the node key needed a model call: a leaf whose
+// source is over nodeBytes, or a merge whose children do not fit in one node.
+func (m *Memory) buildBuiltByModel(key NodeKey) (bool, error) {
+ if key.level == 0 {
+ message, err := m.log.Message(key.index)
+ if err != nil {
+ return false, fmt.Errorf("read leaf source: %w", err)
+ }
+
+ _, free := freeLeaf(message)
- defer func() {
- err = errors.Join(err, handle.Close())
- }()
+ return !free, nil
+ }
- err = handle.Sync()
+ left, right, err := m.children(key)
if err != nil {
- return fmt.Errorf("flush directory: %w", err)
+ return false, fmt.Errorf("read merge sources: %w", err)
}
- return nil
+ _, free := freeMerge(left, right)
+
+ return !free, nil
+}
+
+// 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
+}
+
+// MemoryViewFile is the content of view.json.
+type MemoryViewFile struct {
+ Chat ViewState `json:"chat"`
+ Compaction ViewState `json:"compaction"`
}
projects / onechat / blobdiff