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