git.giacomolli.uk - onechat/blob - view_layout.go

   1 package main
   2
   3 import (
   4         "crypto/sha256"
   5         "encoding/hex"
   6         "encoding/json"
   7         "errors"
   8         "fmt"
   9         "os"
  10         "strings"
  11 )
  12
  13 // viewLayoutFileName is the file of the view layout in the directory of an
  14 // agent.
  15 const viewLayoutFileName = "layout.json"
  16
  17 // ViewLayout splits the view of a turn request into user messages that only
  18 // grow between batches. The view goes out as the same messages as in the
  19 // request before, while their lines give the same text, and the new lines
  20 // follow as one new message. So the end of the view of each earlier turn stays
  21 // a message end. On the ChatGPT backend, the first request of a turn with this
  22 // layout read the view of the turn before from the prompt cache in each of 10
  23 // turns, and with the view in one message it missed once in 5.
  24 //
  25 // The first message starts with "<chat>\n". The closing tag goes into the
  26 // state message after the view, so that no view message changes when lines
  27 // come. The layout keeps the line count and a hash of the text of each message
  28 // in layout.json, so it outlives a restart.
  29 type ViewLayout struct {
  30         path     string
  31         Messages []ViewLayoutMessage `json:"messages"`
  32 }
  33
  34 // ViewLayoutMessage is one view message: its number of lines and the hash of
  35 // its text.
  36 type ViewLayoutMessage struct {
  37         Lines int    `json:"lines"`
  38         Hash  string `json:"hash"`
  39 }
  40
  41 // loadViewLayout reads the layout at path. A missing file gives an empty
  42 // layout.
  43 func loadViewLayout(path string) (*ViewLayout, error) {
  44         layout := &ViewLayout{
  45                 path: path,
  46         }
  47
  48         data, err := os.ReadFile(path)
  49         if errors.Is(err, os.ErrNotExist) {
  50                 return layout, nil
  51         }
  52
  53         if err != nil {
  54                 return nil, fmt.Errorf("read view layout: %w", err)
  55         }
  56
  57         err = json.Unmarshal(data, layout)
  58         if err != nil {
  59                 return nil, fmt.Errorf("decode view layout: %w", err)
  60         }
  61
  62         return layout, nil
  63 }
  64
  65 // Split returns the view messages for lines and saves the layout. It keeps the
  66 // earlier messages up to the first one whose text differs, and puts the rest
  67 // of the lines into one new message.
  68 func (l *ViewLayout) Split(lines []string) ([]string, error) {
  69         var (
  70                 texts []string
  71                 kept  []ViewLayoutMessage
  72         )
  73
  74         start := 0
  75         for _, message := range l.Messages {
  76                 if start+message.Lines > len(lines) {
  77                         break
  78                 }
  79
  80                 text := viewMessage(len(texts) == 0, lines[start:start+message.Lines])
  81                 if viewTextHash(text) != message.Hash {
  82                         break
  83                 }
  84
  85                 texts = append(texts, text)
  86                 kept = append(kept, message)
  87                 start += message.Lines
  88         }
  89
  90         if start < len(lines) || len(texts) == 0 {
  91                 text := viewMessage(len(texts) == 0, lines[start:])
  92                 texts = append(texts, text)
  93                 kept = append(kept, ViewLayoutMessage{
  94                         Lines: len(lines) - start,
  95                         Hash:  viewTextHash(text),
  96                 })
  97         }
  98
  99         l.Messages = kept
 100         data, err := json.Marshal(l)
 101         if err != nil {
 102                 return nil, fmt.Errorf("encode view layout: %w", err)
 103         }
 104
 105         err = writeFileAtomic(l.path, append(data, '\n'))
 106         if err != nil {
 107                 return nil, fmt.Errorf("save view layout: %w", err)
 108         }
 109
 110         return texts, nil
 111 }
 112
 113 // viewMessage returns the text of a view message with lines. The first one
 114 // opens the chat.
 115 func viewMessage(first bool, lines []string) string {
 116         text := strings.Join(lines, "")
 117         if first {
 118                 text = "<chat>\n" + text
 119         }
 120
 121         return text
 122 }
 123
 124 // viewTextHash returns a short hash of text.
 125 func viewTextHash(text string) string {
 126         sum := sha256.Sum256([]byte(text))
 127         return hex.EncodeToString(sum[:16])
 128 }