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         "maps"
  11         "os"
  12         "path/filepath"
  13         "slices"
  14         "strconv"
  15         "strings"
  16         "time"
  17 )
  18
  19 const (
  20         // probeName is the agent name of a probe run. The cache key comes from
  21         // it.
  22         probeName = "Probe"
  23
  24         // probeQuotaLimit is how many points the used percent of a quota window
  25         // can grow in a probe run before the run stops.
  26         probeQuotaLimit = 15.0
  27
  28         // probeLineBytes is the most bytes of one line of a probe log.
  29         probeLineBytes = 1 << 20
  30 )
  31
  32 // probe runs "onechat probe LOG COUNT". It appends the first COUNT messages of
  33 // LOG, one JSON object {kind, text} per line, to a new chat in probe/ in the
  34 // data directory, and builds its tree with the real compactor. After each
  35 // message it waits for the builds, as a turn does. It stops at the first
  36 // failed call, or when the used percent of a quota window grows by more than
  37 // probeQuotaLimit points. At the end it reports the calls, the cached share,
  38 // the quota, the retries, and the line sizes.
  39 func probe(ctx context.Context, args []string, out io.Writer) error {
  40         if len(args) != 2 {
  41                 return errors.New("usage: onechat probe LOG COUNT")
  42         }
  43
  44         count, err := strconv.Atoi(args[1])
  45         if err != nil || count <= 0 {
  46                 return fmt.Errorf("COUNT %q is not a positive number", args[1])
  47         }
  48
  49         entries, err := readProbeLog(args[0], count)
  50         if err != nil {
  51                 return err
  52         }
  53
  54         dir, err := dataDir()
  55         if err != nil {
  56                 return err
  57         }
  58
  59         config, err := loadConfig(dir)
  60         if err != nil {
  61                 return err
  62         }
  63
  64         address, auth, err := modelEndpoint(config, dir)
  65         if err != nil {
  66                 return err
  67         }
  68
  69         root := filepath.Join(dir, "probe", time.Now().Format("20060102-150405"))
  70         err = os.MkdirAll(root, 0o700)
  71         if err != nil {
  72                 return fmt.Errorf("create probe directory: %w", err)
  73         }
  74
  75         usagePath := filepath.Join(root, "usage.jsonl")
  76         client := newClient(address, auth, "onechat-"+probeName, usagePath)
  77         memory, err := openMemory(ctx, root, time.Now, newCompactor(client, probeName, config.Compactor))
  78         if err != nil {
  79                 return err
  80         }
  81
  82         _, err = fmt.Fprintf(out, "probe: %d messages into %s with %s at %s effort\n", len(entries), root, config.Compactor.Model, config.Compactor.Effort)
  83         if err != nil {
  84                 return errors.Join(err, memory.Close())
  85         }
  86
  87         appended, stop := runProbe(ctx, memory, entries, usagePath)
  88         sizes, err := modelLines(memory)
  89         err = errors.Join(err, memory.Close())
  90         if err != nil {
  91                 return err
  92         }
  93
  94         records, err := readUsageFile(usagePath)
  95         if err != nil {
  96                 return err
  97         }
  98
  99         return writeProbeReport(out, appended, stop, records, sizes)
 100 }
 101
 102 // runProbe appends the entries and waits for the builds after each one. It
 103 // returns the number of messages appended and why the run stopped early, or
 104 // "" when it did not.
 105 func runProbe(ctx context.Context, memory *Memory, entries []probeEntry, usagePath string) (int, string) {
 106         for n, entry := range entries {
 107                 _, err := memory.Append(entry.Kind, entry.Text)
 108                 if err == nil {
 109                         err = memory.Wait(ctx)
 110                 }
 111
 112                 if err != nil {
 113                         return n, err.Error()
 114                 }
 115
 116                 records, err := readUsageFile(usagePath)
 117                 if err != nil {
 118                         return n + 1, err.Error()
 119                 }
 120
 121                 reason := probeStop(records)
 122                 if reason != "" {
 123                         return n + 1, reason
 124                 }
 125         }
 126
 127         return len(entries), ""
 128 }
 129
 130 // probeEntry is one message of a probe log.
 131 type probeEntry struct {
 132         Kind string `json:"kind"`
 133         Text string `json:"text"`
 134 }
 135
 136 // readProbeLog reads the first count messages of the probe log at path.
 137 func readProbeLog(path string, count int) (entries []probeEntry, err error) {
 138         file, err := os.Open(path)
 139         if err != nil {
 140                 return nil, fmt.Errorf("open probe log: %w", err)
 141         }
 142
 143         defer func() {
 144                 err = errors.Join(err, file.Close())
 145         }()
 146
 147         scanner := bufio.NewScanner(file)
 148         scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
 149         for len(entries) < count && scanner.Scan() {
 150                 var entry probeEntry
 151                 err = json.Unmarshal(scanner.Bytes(), &entry)
 152                 if err != nil {
 153                         return nil, fmt.Errorf("probe log line %d: %w", len(entries)+1, err)
 154                 }
 155
 156                 entries = append(entries, entry)
 157         }
 158
 159         err = scanner.Err()
 160         if err != nil {
 161                 return nil, fmt.Errorf("read probe log: %w", err)
 162         }
 163
 164         if len(entries) < count {
 165                 return nil, fmt.Errorf("probe log has %d messages, fewer than %d", len(entries), count)
 166         }
 167
 168         return entries, nil
 169 }
 170
 171 // probeStop returns why a probe run must stop: a failed call, or a quota window
 172 // whose used percent grew by more than probeQuotaLimit points since the first
 173 // call. It returns "" when the run can go on.
 174 func probeStop(records []usageRecord) string {
 175         for _, record := range records {
 176                 if record.Error != "" {
 177                         return "call failed: " + record.Error
 178                 }
 179         }
 180
 181         if len(records) == 0 {
 182                 return ""
 183         }
 184
 185         first := usedPercents(records[0])
 186         for name, used := range usedPercents(records[len(records)-1]) {
 187                 start, found := first[name]
 188                 if found && used-start > probeQuotaLimit {
 189                         return fmt.Sprintf("%s grew from %g to %g", name, start, used)
 190                 }
 191         }
 192
 193         return ""
 194 }
 195
 196 // modelLines returns the sizes of the nodes that the model wrote, in tree
 197 // order. A node that fits its source with no model call is not one of them.
 198 func modelLines(memory *Memory) ([]int, error) {
 199         memory.mutex.Lock()
 200         defer memory.mutex.Unlock()
 201
 202         keys, err := memory.tree.Keys()
 203         if err != nil {
 204                 return nil, err
 205         }
 206
 207         var sizes []int
 208         for _, key := range keys {
 209                 node, _, err := memory.tree.Node(key.level, key.index)
 210                 if err != nil {
 211                         return nil, err
 212                 }
 213
 214                 model, err := builtByModel(memory, key)
 215                 if err != nil {
 216                         return nil, err
 217                 }
 218
 219                 if model {
 220                         sizes = append(sizes, node.Size)
 221                 }
 222         }
 223
 224         return sizes, nil
 225 }
 226
 227 // builtByModel tells if the node key needed a model call: a leaf whose source
 228 // is over nodeBytes, or a merge whose children do not fit in one node.
 229 func builtByModel(memory *Memory, key nodeKey) (bool, error) {
 230         if key.level == 0 {
 231                 message, err := memory.log.Message(key.index)
 232                 if err != nil {
 233                         return false, err
 234                 }
 235
 236                 _, free := freeLeaf(message)
 237
 238                 return !free, nil
 239         }
 240
 241         left, right, err := memory.children(key)
 242         if err != nil {
 243                 return false, err
 244         }
 245
 246         _, free := freeMerge(left, right)
 247
 248         return !free, nil
 249 }
 250
 251 // readUsageFile reads the records of usage.jsonl at path.
 252 func readUsageFile(path string) (records []usageRecord, err error) {
 253         file, err := os.Open(path)
 254         if errors.Is(err, os.ErrNotExist) {
 255                 return nil, nil
 256         }
 257
 258         if err != nil {
 259                 return nil, fmt.Errorf("open usage file: %w", err)
 260         }
 261
 262         defer func() {
 263                 err = errors.Join(err, file.Close())
 264         }()
 265
 266         scanner := bufio.NewScanner(file)
 267         scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
 268         for scanner.Scan() {
 269                 var record usageRecord
 270                 err = json.Unmarshal(scanner.Bytes(), &record)
 271                 if err != nil {
 272                         return nil, fmt.Errorf("decode usage: %w", err)
 273                 }
 274
 275                 records = append(records, record)
 276         }
 277
 278         err = scanner.Err()
 279         if err != nil {
 280                 return nil, fmt.Errorf("read usage file: %w", err)
 281         }
 282
 283         return records, nil
 284 }
 285
 286 // usedPercents returns the used percent of each quota window of record, by
 287 // header name.
 288 func usedPercents(record usageRecord) map[string]float64 {
 289         used := make(map[string]float64)
 290         for name, value := range record.Quota {
 291                 if !strings.HasSuffix(name, "-used-percent") {
 292                         continue
 293                 }
 294
 295                 percent, err := strconv.ParseFloat(value, 64)
 296                 if err == nil {
 297                         used[name] = percent
 298                 }
 299         }
 300
 301         return used
 302 }
 303
 304 // writeProbeReport writes the result of a probe run.
 305 func writeProbeReport(out io.Writer, appended int, stop string, records []usageRecord, sizes []int) error {
 306         var input, cached, output, failed int
 307         for _, record := range records {
 308                 input += record.InputTokens
 309                 cached += record.CachedTokens
 310                 output += record.OutputTokens
 311                 if record.Error != "" {
 312                         failed++
 313                 }
 314         }
 315
 316         var lines []string
 317         if stop != "" {
 318                 lines = append(lines, "stopped early: "+stop)
 319         }
 320
 321         lines = append(lines,
 322                 fmt.Sprintf("messages: %d", appended),
 323                 fmt.Sprintf("calls: %d (%d failed), %.2f per message", len(records), failed, ratio(len(records), appended)),
 324                 fmt.Sprintf("model lines: %d, retries: %d", len(sizes), len(records)-failed-len(sizes)),
 325                 fmt.Sprintf("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*ratio(cached, input), output),
 326         )
 327
 328         if len(sizes) > 0 {
 329                 lines = append(lines, sizeText(sizes))
 330         }
 331
 332         if len(records) > 0 {
 333                 lines = append(lines,
 334                         "quota at the first call: "+quotaText(records[0]),
 335                         "quota at the last call: "+quotaText(records[len(records)-1]),
 336                 )
 337         }
 338
 339         _, err := io.WriteString(out, strings.Join(lines, "\n")+"\n")
 340
 341         return err
 342 }
 343
 344 // sizeText describes the sizes of the model lines.
 345 func sizeText(sizes []int) string {
 346         sorted := slices.Sorted(slices.Values(sizes))
 347         total, over := 0, 0
 348         for _, size := range sorted {
 349                 total += size
 350                 if size > nodeBytes {
 351                         over++
 352                 }
 353         }
 354
 355         return fmt.Sprintf("line bytes: min %d, median %d, mean %.0f, max %d, %d over %d", sorted[0], sorted[len(sorted)/2], ratio(total, len(sorted)), sorted[len(sorted)-1], over, nodeBytes)
 356 }
 357
 358 // quotaText returns the quota headers of record, sorted by name.
 359 func quotaText(record usageRecord) string {
 360         names := slices.Sorted(maps.Keys(record.Quota))
 361         parts := make([]string, len(names))
 362         for n, name := range names {
 363                 parts[n] = strings.TrimPrefix(name, "x-codex-") + "=" + record.Quota[name]
 364         }
 365
 366         return strings.Join(parts, " ")
 367 }
 368
 369 // ratio returns a / b, or 0 when b is 0.
 370 func ratio(a, b int) float64 {
 371         if b == 0 {
 372                 return 0
 373         }
 374
 375         return float64(a) / float64(b)
 376 }