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