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

   1 package main
   2
   3 import (
   4         "bufio"
   5         "context"
   6         "encoding/json"
   7         "errors"
   8         "fmt"
   9         "io"
  10         "log/slog"
  11         "os"
  12         "path/filepath"
  13         "strconv"
  14         "time"
  15 )
  16
  17 const (
  18         // probeName is the agent name of a probe run. The cache key comes from
  19         // it.
  20         probeName = "Probe"
  21
  22         // probeCallLimit is the most model calls of a probe run.
  23         probeCallLimit = 660
  24 )
  25
  26 // ProbeSlotSummarizer lets at most cap(slots) builds call inner at once. A
  27 // build past the limit waits for a slot.
  28 type ProbeSlotSummarizer struct {
  29         inner Summarizer
  30         slots chan struct{}
  31 }
  32
  33 // Summarize waits for a slot and builds the node with inner.
  34 func (s ProbeSlotSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
  35         select {
  36         case s.slots <- struct{}{}:
  37         case <-ctx.Done():
  38                 return "", ctx.Err()
  39         }
  40
  41         defer func() {
  42                 <-s.slots
  43         }()
  44
  45         text, err := s.inner.Summarize(ctx, task)
  46         if err != nil {
  47                 return text, fmt.Errorf("summarize probe node: %w", err)
  48         }
  49
  50         return text, nil
  51 }
  52
  53 // probe runs "onechat probe LOG COUNT [BUILDS]". It appends the first COUNT
  54 // messages of LOG, one JSON object {kind, text} per line, to a new chat in
  55 // probe/ in the data directory, and builds its tree with the real compactor.
  56 // Before each user message it waits for the builds, as a turn does. It stops at
  57 // the first failed call, after probeCallLimit calls, or when the used percent
  58 // of a quota window grows by more than buildQuotaLimit points. At the end it
  59 // reports the calls, the time per call, the calls at once, the cached share,
  60 // the quota, the retries, and the line sizes. BUILDS, from 1 to buildMax,
  61 // limits the builds that call the model at once in this run only. The core
  62 // still queues up to buildMax builds. The run holds the lock of the token file,
  63 // so it does not run while the daemon runs.
  64 func probe(ctx context.Context, args []string, out io.Writer) (err error) {
  65         if len(args) != 2 && len(args) != 3 {
  66                 return errors.New("usage: onechat probe LOG COUNT [BUILDS]")
  67         }
  68
  69         count, err := strconv.Atoi(args[1])
  70         if err != nil {
  71                 return fmt.Errorf("parse probe count: %w", err)
  72         }
  73
  74         if count <= 0 {
  75                 return fmt.Errorf("COUNT %q is not a positive number", args[1])
  76         }
  77
  78         builds := buildMax
  79         if len(args) == 3 {
  80                 builds, err = strconv.Atoi(args[2])
  81                 if err != nil {
  82                         return fmt.Errorf("parse probe builds: %w", err)
  83                 }
  84
  85                 if builds < 1 || builds > buildMax {
  86                         return fmt.Errorf("BUILDS %q is not a number from 1 to %d", args[2], buildMax)
  87                 }
  88         }
  89
  90         entries, err := probeReadLog(args[0], count)
  91         if err != nil {
  92                 return fmt.Errorf("read probe input: %w", err)
  93         }
  94
  95         dir, err := dataDir()
  96         if err != nil {
  97                 return fmt.Errorf("find probe data directory: %w", err)
  98         }
  99
 100         config, err := loadConfig(dir)
 101         if err != nil {
 102                 return fmt.Errorf("load probe configuration: %w", err)
 103         }
 104
 105         tokens, err := lockTokens(dir)
 106         if err != nil {
 107                 return fmt.Errorf("lock probe tokens: %w", err)
 108         }
 109
 110         defer func() {
 111                 closeErr := tokens.Close()
 112                 if closeErr != nil {
 113                         err = errors.Join(err, fmt.Errorf("close probe token lock: %w", closeErr))
 114                 }
 115         }()
 116
 117         address, auth, err := modelEndpoint(config, dir)
 118         if err != nil {
 119                 return fmt.Errorf("select probe endpoint: %w", err)
 120         }
 121
 122         root := filepath.Join(dir, "probe", time.Now().Format("20060102-150405"))
 123         err = os.MkdirAll(root, 0o700)
 124         if err != nil {
 125                 return fmt.Errorf("create probe directory: %w", err)
 126         }
 127
 128         usagePath := filepath.Join(root, "usage.jsonl")
 129         client := newModelClient(address, auth, "onechat-"+probeName, usagePath)
 130         client.LimitCalls(probeCallLimit)
 131
 132         summarizer := ProbeSlotSummarizer{
 133                 inner: newCompactor(client, probePrompt, config.Compactor, promptTools(), &ImageStore{dir: filepath.Join(dir, imageDirName)}),
 134                 slots: make(chan struct{}, builds),
 135         }
 136
 137         // The report of the probe counts the failed calls, so the lines of the
 138         // failed builds are not needed.
 139         memory, err := openMemory(ctx, root, time.Now, summarizer, slog.New(slog.DiscardHandler))
 140         if err != nil {
 141                 return fmt.Errorf("open probe memory: %w", err)
 142         }
 143
 144         _, err = fmt.Fprintf(out, "probe: %d messages into %s with %s at %s effort, %d builds at once\n", len(entries), root, config.Compactor.Model, config.Compactor.Effort, builds)
 145         if err != nil {
 146                 return fmt.Errorf("write probe header: %w", errors.Join(err, memory.Close()))
 147         }
 148
 149         run := BuildRun{
 150                 memory:    memory,
 151                 client:    client,
 152                 usagePath: usagePath,
 153         }
 154
 155         appended, stop := run.Run(ctx, entries)
 156         sizes, err := memory.buildModelLines()
 157         err = errors.Join(err, memory.Close())
 158         if err != nil {
 159                 return fmt.Errorf("finish probe memory: %w", err)
 160         }
 161
 162         records, err := buildReadUsage(usagePath)
 163         if err != nil {
 164                 return fmt.Errorf("read probe usage: %w", err)
 165         }
 166
 167         err = buildWriteReport(out, appended, stop, records, sizes)
 168         if err != nil {
 169                 return fmt.Errorf("write probe report: %w", err)
 170         }
 171
 172         return nil
 173 }
 174
 175 // probeReadLog reads the first count messages of the probe log at path. Source
 176 // dates are ignored because a probe uses its own clock.
 177 func probeReadLog(path string, count int) (entries []DatedMessage, err error) {
 178         file, err := os.Open(path)
 179         if err != nil {
 180                 return nil, fmt.Errorf("open probe log: %w", err)
 181         }
 182
 183         defer func() {
 184                 closeErr := file.Close()
 185                 if closeErr != nil {
 186                         err = errors.Join(err, fmt.Errorf("close probe log: %w", closeErr))
 187                 }
 188         }()
 189
 190         scanner := bufio.NewScanner(file)
 191         scanner.Buffer(make([]byte, 64*1024), buildLineBytes)
 192
 193         for len(entries) < count && scanner.Scan() {
 194                 var entry DatedMessage
 195
 196                 err = json.Unmarshal(scanner.Bytes(), &entry)
 197                 if err != nil {
 198                         return nil, fmt.Errorf("probe log line %d: %w", len(entries)+1, err)
 199                 }
 200
 201                 entry.Date = ""
 202                 entries = append(entries, entry)
 203         }
 204
 205         err = scanner.Err()
 206         if err != nil {
 207                 return nil, fmt.Errorf("read probe log: %w", err)
 208         }
 209
 210         if len(entries) < count {
 211                 return nil, fmt.Errorf("probe log has %d messages, fewer than %d", len(entries), count)
 212         }
 213
 214         return entries, nil
 215 }
 216
 217 // probePrompt returns the system prompt of a probe run: the prompt of the agent
 218 // probeName, without instructions and with a fixed instructions file path.
 219 func probePrompt() (string, error) {
 220         return prompt(probeName, "/PROBE/instructions.md", ""), nil
 221 }