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