X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/6844f8c992ec4e276a3a0d93b544de0127eebbd2..d7f6fba630444f5065f59cd3729feab92c7798aa:/memory.go?ds=inline diff --git a/memory.go b/memory.go index c9fe1c5..49215ac 100644 --- a/memory.go +++ b/memory.go @@ -22,6 +22,64 @@ var ( errMemoryUnsettled = errors.New("a summary of the view failed, and it is tried again at the next message") ) +// 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{} +} + +// 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. Images holds the leaf +// image filenames and is empty for merges. +type Task struct { + Level int + Index int + Context string + Input []string + Images []string +} + +// MemoryViewFile is the content of view.json. +type MemoryViewFile 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. logger gets // one line for each failed build. The builds stop when ctx ends or at Close. @@ -75,46 +133,6 @@ func openMemory(ctx context.Context, root string, clock func() time.Time, summar return m, nil } -// 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 { @@ -137,169 +155,6 @@ func (m *Memory) Close() error { 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()) @@ -441,63 +296,6 @@ func (m *Memory) dispatch() { } } -// 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 -} - -// 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) - } - - 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") } @@ -532,20 +330,77 @@ func (m *Memory) settle(key NodeKey) error { return err } - node, free := freeMerge(left, right) - if !free { - m.queue.AddMerge(parent) + 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 + } +} + +// 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) + } - return nil - } + _, err = m.cview.Append(m.tree) + if err == nil && merged { + err = m.cview.Coarsen(m.chat, m.tree) + } - err = m.tree.Add(node) - if err != nil { - return err - } + if err != nil { + return fmt.Errorf("append to compaction view: %w", err) + } - key = parent + return nil +} + +// 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) + } + + err = writeFileAtomic(m.viewPath(), append(data, '\n')) + if err != nil { + return fmt.Errorf("save view: %w", err) } + + return nil } // task returns the build task of the node key, with its input read from the log @@ -564,6 +419,7 @@ func (m *Memory) task(key NodeKey) (Task, error) { } task.Input = []string{leafSource(message)} + task.Images = imageFiles(message.Text) return task, nil } @@ -580,6 +436,11 @@ func (m *Memory) task(key NodeKey) (Task, error) { return task, nil } +// 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) +} + // 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) { @@ -597,6 +458,7 @@ func (m *Memory) run(key NodeKey, task Task) { Text: text, Size: len(text), }) + m.syncer.Written() } @@ -616,11 +478,6 @@ func (m *Memory) run(key NodeKey, task Task) { m.dispatch() } -// 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} { @@ -635,6 +492,26 @@ func (m *Memory) ready(key NodeKey) (bool, error) { 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 +} + func (m *Memory) context(level, index int) string { end := index if level > 0 { @@ -644,6 +521,192 @@ func (m *Memory) context(level, index int) string { return m.cview.Context(end) } +// 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() + + return m.appendDated(kind, text, date) +} + +// appendDated appends under the memory lock and updates the views and builds. +func (m *Memory) appendDated(kind, text, date string) (Message, error) { + 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 +} + +// AppendOnce appends a queued message. prepare saves its log ID before the +// append, under the memory lock. A saved ID permits recovery after a stop +// between that save and the append, without a second log entry. prepare must +// not call Memory. The caller stops delivery after a failed append. +func (m *Memory) AppendOnce(kind, text string, id *int, prepare func(int) error) (Message, error) { + m.mutex.Lock() + defer m.mutex.Unlock() + + if id != nil && *id < m.log.Len() { + message, err := m.log.Message(*id) + if err != nil { + return Message{}, fmt.Errorf("read queued message: %w", err) + } + + if message.Kind != kind || message.Text != text { + return Message{}, errors.New("queued message does not match its log entry") + } + + return message, nil + } + + if id != nil && *id != m.log.Len() { + return Message{}, errors.New("queued message is beyond the log") + } + + err := prepare(m.log.Len()) + if err != nil { + return Message{}, fmt.Errorf("prepare queued message: %w", err) + } + + message, err := m.appendDated(kind, text, "") + if err != nil { + return Message{}, fmt.Errorf("append queued message: %w", err) + } + + return message, 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() + } + } +} + +// 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 +} + +// Flush returns when every write before the call is durable. +func (m *Memory) Flush(ctx context.Context) error { + return m.syncer.Flush(ctx) +} + +// 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) +} + // View returns the chat view as the model reads it. func (m *Memory) View() string { m.mutex.Lock() @@ -735,19 +798,3 @@ func (m *Memory) buildBuiltByModel(key NodeKey) (bool, error) { 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"` -}