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 (
  63                 texts []string
  64                 kept  []ViewLayoutMessage
  65         )
  66
  67         start := 0
  68         for _, message := range l.Messages {
  69                 if start+message.Lines > len(lines) {
  70                         break
  71                 }
  72
  73                 text := viewMessage(len(texts) == 0, lines[start:start+message.Lines])
  74                 if viewTextHash(text) != message.Hash {
  75                         break
  76                 }
  77
  78                 texts = append(texts, text)
  79                 kept = append(kept, message)
  80                 start += message.Lines
  81         }
  82
  83         if start < len(lines) || len(texts) == 0 {
  84                 text := viewMessage(len(texts) == 0, lines[start:])
  85                 texts = append(texts, text)
  86                 kept = append(kept, ViewLayoutMessage{
  87                         Lines: len(lines) - start,
  88                         Hash:  viewTextHash(text),
  89                 })
  90         }
  91
  92         l.Messages = kept
  93         data, err := json.Marshal(l)
  94         if err != nil {
  95                 return nil, fmt.Errorf("encode view layout: %w", err)
  96         }
  97
  98         err = writeFileAtomic(l.path, append(data, '\n'))
  99         if err != nil {
 100                 return nil, fmt.Errorf("save view layout: %w", err)
 101         }
 102
 103         return texts, nil
 104 }
 105
 106 // ViewLayoutMessage is one view message: its number of lines and the hash of
 107 // its text.
 108 type ViewLayoutMessage struct {
 109         Lines int    `json:"lines"`
 110         Hash  string `json:"hash"`
 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 }