Send model calls to a Responses endpoint

[onechat] / memory.go

diff --git a/memory.go b/memory.go
index eb28341f69821e6158ad49e29fd2b644314ecff3..733673bf8607b50019e8417a80d4d0d9c6352151 100644 (file)
--- a/memory.go
+++ b/memory.go
@@ -7,6 +7,7 @@ import (
        "fmt"
        "os"
        "path/filepath"
        "fmt"
        "os"
        "path/filepath"

+       "strings"

        "sync"
        "time"
 )
        "sync"
        "time"
 )


@@ -14,21 +15,48 @@ import (
 // errNoView tells that a chat has messages but no saved view.
 var errNoView = errors.New("the chat has messages but no saved view")
 
 // 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
 // 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 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 {
 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.
 }
 
 // viewFile is the content of view.json.


@@ -37,19 +65,25 @@ type viewFile struct {
        Compaction viewState `json:"compaction"`
 }
 
        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{
        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())
        m.tree, err = openTree(ctx, root, log.Len(), clock)
        if err != nil {
                return nil, errors.Join(err, m.Close())


@@ -60,12 +94,24 @@ func openMemory(ctx context.Context, root string, clock func() time.Time) (*Memo
                return nil, errors.Join(fmt.Errorf("load view: %w", 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
 }
 
        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 {
 func (m *Memory) load() error {
        data, err := os.ReadFile(m.viewPath())
        if errors.Is(err, os.ErrNotExist) && m.log.Len() == 0 {


@@ -107,12 +153,72 @@ func (m *Memory) load() error {
                return fmt.Errorf("view covers %d messages, but the log has %d", 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() {
        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
                }
                if err != nil {
                        return err
                }


@@ -121,8 +227,12 @@ func (m *Memory) load() error {
        return m.save()
 }
 
        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 {
 func (m *Memory) Close() error {

+       m.cancel()
+       m.workers.Wait()
+

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


@@ -134,7 +244,8 @@ func (m *Memory) Close() error {
        return errors.Join(err, m.log.Close())
 }
 
        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()
 func (m *Memory) Append(kind, text string) (Message, error) {
        m.mutex.Lock()
        defer m.mutex.Unlock()


@@ -154,55 +265,219 @@ func (m *Memory) Append(kind, text string) (Message, error) {
                return Message{}, err
        }
 
                return Message{}, err
        }
 

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

        return message, nil
 }
 
        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 {
 func (m *Memory) addMessage(message Message) error {

-       _, built := m.tree.Node(0, message.ID)

        leaf, free := freeLeaf(message)
        leaf, free := freeLeaf(message)

-       if free && !built {
+       if free {

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

+       } else {
+               m.queue.AddLeaf(message.ID)

        }
 
        }
 

-       merged := m.chat.Append(m.tree)
-       m.cview.Append(m.tree)
-       if merged {
-               m.cview.Coarsen(m.chat, m.tree)
+       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
 }
 
        }
 
        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, keeps it in the views when it is a pending line,
+// and settles its parent.

 func (m *Memory) build(node Node) error {
 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 {
        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
                }
 
                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
                }
 
                        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
+               }
+
+               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

                }
                }

+       }

 
 

-               parent, free := freeMerge(left, right)
-               if !free {
+       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
                }
 
                        return nil
                }
 

-               node = parent
+               select {
+               case <-changed:
+               case <-ctx.Done():
+                       return ctx.Err()
+               }

        }
 }
 
        }
 }
 


@@ -233,7 +508,7 @@ func (m *Memory) View() string {
        m.mutex.Lock()
        defer m.mutex.Unlock()
 
        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
 }
 
 // Context returns the compaction view for the build of node (l, i). It ends


@@ -252,23 +527,19 @@ func (m *Memory) context(level, index int) string {
                end = nodeEnd(level, index)
        }
 
                end = nodeEnd(level, index)
        }
 

-       return m.cview.Context(m.tree, end)
+       return m.cview.Context(end)

 }
 
 }
 

-// 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.
+// writeFileAtomic replaces the file at path with data through a temporary file
+// and a rename. A reader or a crash of the process sees the old file or the new
+// file, never a part of one.

 func writeFileAtomic(path string, data []byte) error {
 func writeFileAtomic(path string, data []byte) error {

-       dir := filepath.Dir(path)
-       file, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*")
+       file, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".*")

        if err != nil {
                return fmt.Errorf("create temporary file: %w", err)
        }
 
        _, err = file.Write(data)
        if err != nil {
                return fmt.Errorf("create temporary file: %w", err)
        }
 
        _, err = file.Write(data)

-       if err == nil {
-               err = file.Sync()
-       }
-

        err = errors.Join(err, file.Close())
        if err == nil {
                err = os.Rename(file.Name(), path)
        err = errors.Join(err, file.Close())
        if err == nil {
                err = os.Rename(file.Name(), path)


@@ -278,24 +549,5 @@ func writeFileAtomic(path string, data []byte) error {
                return errors.Join(fmt.Errorf("replace file: %w", err), os.Remove(file.Name()))
        }
 
                return errors.Join(fmt.Errorf("replace file: %w", err), os.Remove(file.Name()))
        }
 

-       return syncDir(dir)
-}
-
-// 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)
-       }
-
-       defer func() {
-               err = errors.Join(err, handle.Close())
-       }()
-
-       err = handle.Sync()
-       if err != nil {
-               return fmt.Errorf("flush directory: %w", err)
-       }
-

        return nil
 }
        return nil
 }