"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, 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.
+// 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 log and the tree once to fill the queue again. After that, only a
+// new message or a new node adds work.
type Memory struct {
- mutex sync.Mutex
- root string
- log *Log
- tree *Tree
- chat *View
- cview *View
+ 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.
Compaction viewState `json:"compaction"`
}
-// 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. 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,
+ 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())
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 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.
+// 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 {
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 {
+ if m.tree.Built(0, id) {
+ continue
+ }
+
+ leaf, free := freeLeaf(m.log.Message(id))
+ if !free {
+ m.queue.AddLeaf(id)
+ continue
+ }
+
+ err := m.build(leaf)
+ if err != nil {
+ return err
+ }
+ }
+
+ for _, node := range m.tree.Nodes() {
+ if node.Index%2 != 0 {
+ continue
+ }
+
+ err := m.settle(node)
+ 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))
+ err := m.addMessage(m.log.Message(m.chat.end))
if err != nil {
return err
}
return m.save()
}
-// Close closes the files of the chat and releases its lock.
+// 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()
return errors.Join(err, m.log.Close())
}
-// Append logs a new message and adds it to the tree and the view.
+// 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()
return Message{}, err
}
+ m.queue.Retry()
+ m.dispatch()
+
return message, nil
}
-// addMessage builds the leaf of message when it needs no model call and appends
-// it to the view.
+// 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.
func (m *Memory) addMessage(message Message) error {
- _, built := m.tree.Node(0, message.ID)
leaf, free := freeLeaf(message)
- if free && !built {
+ if free {
err := m.build(leaf)
if err != nil {
return err
}
+ } else {
+ m.queue.AddLeaf(message.ID)
}
merged := m.chat.Append(m.tree)
return nil
}
-// build adds node to the tree, and then each parent above it that needs no
-// model call.
+// build adds node to the tree and settles its parent.
func (m *Memory) build(node Node) error {
- for {
- err := m.tree.Add(node)
- if err != nil {
- return err
- }
+ err := m.tree.Add(node)
+ if err != nil {
+ return err
+ }
+
+ return m.settle(node)
+}
+// settle looks at the parent of node 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(node Node) error {
+ for {
sibling, built := m.tree.Node(node.Level, node.Index^1)
- if !built {
+ if !built || m.tree.Built(node.Level+1, node.Index/2) {
return nil
}
parent, free := freeMerge(left, right)
if !free {
+ m.queue.AddMerge(nodeKey{node.Level + 1, node.Index / 2})
return nil
}
+ err := m.tree.Add(parent)
+ if err != nil {
+ return err
+ }
+
node = parent
}
}
+// dispatch starts the builds that are ready, up to the limit.
+func (m *Memory) dispatch() {
+ for m.ctx.Err() == nil {
+ key, ready := m.queue.Next()
+ if !ready {
+ return
+ }
+
+ task := m.task(key)
+ m.workers.Add(1)
+ go m.run(key, task)
+ }
+}
+
+// task returns the build task of the node key.
+func (m *Memory) task(key nodeKey) Task {
+ task := Task{
+ Level: key.level,
+ Index: key.index,
+ Context: m.context(key.level, key.index),
+ }
+
+ if key.level == 0 {
+ task.Input = []string{leafSource(m.log.Message(key.index))}
+ return task
+ }
+
+ for _, index := range []int{2 * key.index, 2*key.index + 1} {
+ child, _ := m.tree.Node(key.level-1, index)
+ task.Input = append(task.Input, strings.ReplaceAll(child.Text, "\n", " "))
+ }
+
+ return task
+}
+
+// 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{
package main
import (
+ "context"
+ "errors"
"fmt"
+ "maps"
"os"
"path/filepath"
+ "slices"
+ "strings"
+ "sync"
"testing"
+ "unicode/utf8"
"github.com/stretchr/testify/require"
)
func TestMemoryRestartLoadsTheSameView(t *testing.T) {
root := t.TempDir()
clock := newTestClock()
- memory, err := openMemory(t.Context(), root, clock.Now)
+ memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
require.NoError(t, err)
memory.chat.high = 200
require.Less(t, len(memory.chat.lines), 100)
require.NoError(t, memory.Close())
- memory, err = openMemory(t.Context(), root, clock.Now)
+ memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
require.NoError(t, err)
defer func() {
}
func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) {
- memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now)
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
require.NoError(t, err)
defer func() {
}
}
+func TestMemoryRunsAtMostEightBuilds(t *testing.T) {
+ summarizer := newFakeSummarizer()
+ summarizer.gate = make(chan struct{})
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ for n := range 12 {
+ _, err = memory.Append("echo", longText(n))
+ require.NoError(t, err)
+ }
+
+ memory.mutex.Lock()
+ running := slices.Collect(maps.Keys(memory.queue.running))
+ memory.mutex.Unlock()
+ slices.SortFunc(running, func(a, b nodeKey) int {
+ return a.index - b.index
+ })
+ require.Equal(t, []nodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5}, {0, 6}, {0, 7}}, running)
+
+ close(summarizer.gate)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.Equal(t, maxBuilds, summarizer.peak)
+
+ for level, count := range []int{12, 6, 3, 1} {
+ for index := range count {
+ require.True(t, memory.tree.Built(level, index), "node (%d, %d)", level, index)
+ }
+ }
+}
+
+func TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) {
+ summarizer := newFakeSummarizer()
+ summarizer.SetFailing(true)
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ for n := range 9 {
+ _, err = memory.Append("echo", longText(n))
+ require.NoError(t, err)
+ }
+
+ require.NoError(t, memory.Wait(t.Context()))
+ for _, task := range summarizer.Tasks() {
+ require.Less(t, task.Index, 8)
+ }
+
+ require.False(t, memory.tree.Built(0, 8))
+ summarizer.SetFailing(false)
+
+ _, err = memory.Append("user", "retry")
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.True(t, memory.tree.Built(0, 8))
+ require.True(t, memory.tree.Built(3, 0))
+}
+
+func TestMemoryMergesAfterBothHalves(t *testing.T) {
+ summarizer := newFakeSummarizer()
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ _, err = memory.Append("echo", longText(0))
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+
+ _, err = memory.Append("echo", longText(1))
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.Len(t, summarizer.Tasks(), 3)
+
+ left, _ := memory.tree.Node(0, 0)
+ right, _ := memory.tree.Node(0, 1)
+ merge := summarizer.Tasks()[2]
+ require.Equal(t, []string{left.Text, right.Text}, merge.Input)
+ require.Equal(t, "<chat>\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n</chat>", merge.Context)
+ require.True(t, memory.tree.Built(1, 0))
+}
+
+func TestMemoryRestartRestoresTheQueue(t *testing.T) {
+ root := t.TempDir()
+ clock := newTestClock()
+ summarizer := newFakeSummarizer()
+ summarizer.SetFailing(true)
+ memory, err := openMemory(t.Context(), root, clock.Now, summarizer)
+ require.NoError(t, err)
+
+ for n := range 4 {
+ _, err = memory.Append("echo", longText(n))
+ require.NoError(t, err)
+ }
+
+ require.NoError(t, memory.Wait(t.Context()))
+ require.NoError(t, memory.Close())
+
+ memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ require.NoError(t, memory.Wait(t.Context()))
+ require.True(t, memory.tree.Built(2, 0))
+}
+
func TestMemoryBuildsFreeParents(t *testing.T) {
- memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now)
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer())
require.NoError(t, err)
defer func() {
}
func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {
- memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now)
+ summarizer := newFakeSummarizer()
+ summarizer.SetFailing(true)
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer)
require.NoError(t, err)
defer func() {
require.NoError(t, memory.Close())
}()
- _, err = memory.Append("echo", string(make([]byte, nodeBytes)))
+ _, err = memory.Append("echo", longText(0))
require.NoError(t, err)
-
+ require.NoError(t, memory.Wait(t.Context()))
require.Equal(t, "<chat>\n0+1|"+pending+"\n</chat>", memory.View())
+
+ summarizer.SetFailing(false)
+
+ _, err = memory.Append("user", "again")
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.Equal(t, "<chat>\n0+1|"+fakeLine(memory.task(nodeKey{0, 0}))+"\n1+1|user: again\n</chat>", memory.View())
}
func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
root := t.TempDir()
clock := newTestClock()
- memory, err := openMemory(t.Context(), root, clock.Now)
+ memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
require.NoError(t, err)
_, err = memory.Append("user", "a")
require.NoError(t, err)
require.NoError(t, memory.Close())
- memory, err = openMemory(t.Context(), root, clock.Now)
+ memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
require.NoError(t, err)
defer func() {
func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
root := t.TempDir()
clock := newTestClock()
- memory, err := openMemory(t.Context(), root, clock.Now)
+ memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
require.NoError(t, err)
_, err = memory.Append("user", "a")
require.NoError(t, memory.Close())
require.NoError(t, os.Remove(filepath.Join(root, "view.json")))
- _, err = openMemory(t.Context(), root, clock.Now)
+ _, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
require.ErrorIs(t, err, errNoView)
}
+
+// fakeSummarizer builds each node with a line of fakeLine. It fails each task
+// while failing is true, and each task waits for gate when gate is not nil.
+type fakeSummarizer struct {
+ mutex sync.Mutex
+ gate chan struct{}
+ failing bool
+ active int
+ peak int
+ tasks []Task
+}
+
+func newFakeSummarizer() *fakeSummarizer {
+ return &fakeSummarizer{}
+}
+
+func (f *fakeSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
+ f.mutex.Lock()
+ f.tasks = append(f.tasks, task)
+ f.active++
+ f.peak = max(f.peak, f.active)
+ fail := f.failing
+ f.mutex.Unlock()
+
+ defer func() {
+ f.mutex.Lock()
+ f.active--
+ f.mutex.Unlock()
+ }()
+
+ if f.gate != nil {
+ select {
+ case <-f.gate:
+ case <-ctx.Done():
+ return "", ctx.Err()
+ }
+ }
+
+ if fail {
+ return "", errors.New("fake failure")
+ }
+
+ return fakeLine(task), nil
+}
+
+// fakeLine returns a line of nodeBytes bytes or less that names the node and
+// starts with the input. A real line of this size makes each merge a model
+// call.
+func fakeLine(task Task) string {
+ count := 1 << task.Level
+ line := fmt.Sprintf("summary %d+%d: %s", task.Index*count, count, strings.Join(task.Input, " "))
+ line = strings.ReplaceAll(line, "\n", " ")
+ for len(line) > nodeBytes-12 {
+ _, size := utf8.DecodeLastRuneInString(line)
+ line = line[:len(line)-size]
+ }
+
+ return line
+}
+
+// longText returns a text that is too long for its own node.
+func longText(n int) string {
+ return fmt.Sprintf("long %d ", n) + strings.Repeat("x", nodeBytes)
+}
+
+// SetFailing makes the next tasks fail or not.
+func (f *fakeSummarizer) SetFailing(failing bool) {
+ f.mutex.Lock()
+ defer f.mutex.Unlock()
+
+ f.failing = failing
+}
+
+// Tasks returns the tasks so far.
+func (f *fakeSummarizer) Tasks() []Task {
+ f.mutex.Lock()
+ defer f.mutex.Unlock()
+
+ return slices.Clone(f.tasks)
+}
--- /dev/null
+package main
+
+import (
+ "slices"
+)
+
+const (
+ // maxBuilds is the number of model builds that can run at once.
+ maxBuilds = 8
+
+ // leafWindow is the number of unbuilt leaves that can be before a leaf
+ // whose build starts.
+ leafWindow = 8
+)
+
+// queue holds the nodes that wait for a model build. A leaf joins it when its
+// message is too long to be its own node. A merge joins it when both of its
+// children are built and their texts do not fit in one node. So the core never
+// scans the tree for work.
+//
+// A leaf can start when fewer than leafWindow leaves before it are unbuilt.
+// Leaves start before merges, oldest first, because a turn waits for the leaves
+// of all earlier messages. Merges start in the order in which they became
+// ready. A failed build waits until the next message.
+type queue struct {
+ leaves []int
+ merges []nodeKey
+ queued map[nodeKey]bool
+ running map[nodeKey]bool
+ failed map[nodeKey]bool
+}
+
+func newQueue() *queue {
+ return &queue{
+ queued: make(map[nodeKey]bool),
+ running: make(map[nodeKey]bool),
+ failed: make(map[nodeKey]bool),
+ }
+}
+
+// AddLeaf adds the leaf of message id. The IDs must come in increasing order.
+func (q *queue) AddLeaf(id int) {
+ q.leaves = append(q.leaves, id)
+}
+
+// AddMerge adds a merge whose children are built. A merge that is in the queue
+// already stays as it is.
+func (q *queue) AddMerge(key nodeKey) {
+ if q.queued[key] {
+ return
+ }
+
+ q.queued[key] = true
+ q.merges = append(q.merges, key)
+}
+
+// Next returns the next node that can start, and marks it as running.
+func (q *queue) Next() (nodeKey, bool) {
+ if len(q.running) >= maxBuilds {
+ return nodeKey{}, false
+ }
+
+ for _, id := range q.leaves[:min(len(q.leaves), leafWindow)] {
+ key := nodeKey{0, id}
+ if q.startable(key) {
+ q.running[key] = true
+ return key, true
+ }
+ }
+
+ for _, key := range q.merges {
+ if q.startable(key) {
+ q.running[key] = true
+ return key, true
+ }
+ }
+
+ return nodeKey{}, false
+}
+
+func (q *queue) startable(key nodeKey) bool {
+ return !q.running[key] && !q.failed[key]
+}
+
+// Done removes a built node from the queue.
+func (q *queue) Done(key nodeKey) {
+ delete(q.running, key)
+ if key.level == 0 {
+ q.leaves = slices.DeleteFunc(q.leaves, func(id int) bool {
+ return id == key.index
+ })
+
+ return
+ }
+
+ delete(q.queued, key)
+ q.merges = slices.DeleteFunc(q.merges, func(other nodeKey) bool {
+ return other == key
+ })
+}
+
+// Fail keeps a node in the queue, but it cannot start until Retry.
+func (q *queue) Fail(key nodeKey) {
+ delete(q.running, key)
+ q.failed[key] = true
+}
+
+// Retry lets the failed nodes start again.
+func (q *queue) Retry() {
+ clear(q.failed)
+}
+
+// Running returns the number of builds that run.
+func (q *queue) Running() int {
+ return len(q.running)
+}
projects / onechat / commitdiff