import (
"bufio"
- "cmp"
"context"
"encoding/json"
"errors"
"fmt"
"io"
- "maps"
+ "log/slog"
"os"
"path/filepath"
- "slices"
"strconv"
- "strings"
"time"
)
// it.
probeName = "Probe"
- // probeQuotaLimit is how many points the used percent of a quota window
- // can grow in a probe run before the run stops.
- probeQuotaLimit = 15.0
-
// probeCallLimit is the most model calls of a probe run.
probeCallLimit = 660
-
- // probeLineBytes is the most bytes of one line of a probe log.
- probeLineBytes = 1 << 20
)
// probe runs "onechat probe LOG COUNT [BUILDS]". It appends the first COUNT
// probe/ in the data directory, and builds its tree with the real compactor.
// Before each user message it waits for the builds, as a turn does. It stops at
// the first failed call, after probeCallLimit calls, or when the used percent
-// of a quota window grows by more than probeQuotaLimit points. At the end it
+// of a quota window grows by more than buildQuotaLimit points. At the end it
// reports the calls, the time per call, the calls at once, the cached share,
-// the quota, the retries, and the line sizes. BUILDS, from 1 to maxBuilds,
+// the quota, the retries, and the line sizes. BUILDS, from 1 to buildMax,
// limits the builds that call the model at once in this run only. The core
-// still queues up to maxBuilds builds.
-func probe(ctx context.Context, args []string, out io.Writer) error {
+// still queues up to buildMax builds. The run holds the lock of the token file,
+// so it does not run while the daemon runs.
+func probe(ctx context.Context, args []string, out io.Writer) (err error) {
if len(args) != 2 && len(args) != 3 {
return errors.New("usage: onechat probe LOG COUNT [BUILDS]")
}
count, err := strconv.Atoi(args[1])
- if err != nil || count <= 0 {
+ if err != nil {
+ return fmt.Errorf("parse probe count: %w", err)
+ }
+
+ if count <= 0 {
return fmt.Errorf("COUNT %q is not a positive number", args[1])
}
- builds := maxBuilds
+ builds := buildMax
if len(args) == 3 {
builds, err = strconv.Atoi(args[2])
- if err != nil || builds < 1 || builds > maxBuilds {
- return fmt.Errorf("BUILDS %q is not a number from 1 to %d", args[2], maxBuilds)
+ if err != nil {
+ return fmt.Errorf("parse probe builds: %w", err)
+ }
+
+ if builds < 1 || builds > buildMax {
+ return fmt.Errorf("BUILDS %q is not a number from 1 to %d", args[2], buildMax)
}
}
- entries, err := readProbeLog(args[0], count)
+ entries, err := probeReadLog(args[0], count)
if err != nil {
- return err
+ return fmt.Errorf("read probe input: %w", err)
}
dir, err := dataDir()
if err != nil {
- return err
+ return fmt.Errorf("find probe data directory: %w", err)
}
config, err := loadConfig(dir)
if err != nil {
- return err
+ return fmt.Errorf("load probe configuration: %w", err)
+ }
+
+ tokens, err := lockTokens(dir)
+ if err != nil {
+ return fmt.Errorf("lock probe tokens: %w", err)
}
+ defer func() {
+ closeErr := tokens.Close()
+ if closeErr != nil {
+ err = errors.Join(err, fmt.Errorf("close probe token lock: %w", closeErr))
+ }
+ }()
+
address, auth, err := modelEndpoint(config, dir)
if err != nil {
- return err
+ return fmt.Errorf("select probe endpoint: %w", err)
}
root := filepath.Join(dir, "probe", time.Now().Format("20060102-150405"))
}
usagePath := filepath.Join(root, "usage.jsonl")
- client := newClient(address, auth, "onechat-"+probeName, usagePath)
+ client := newModelClient(address, auth, "onechat-"+probeName, usagePath)
client.LimitCalls(probeCallLimit)
- summarizer := slotSummarizer{
- inner: newCompactor(client, probeName, config.Compactor),
+
+ summarizer := ProbeSlotSummarizer{
+ inner: newCompactor(client, probePrompt, config.Compactor),
slots: make(chan struct{}, builds),
}
- memory, err := openMemory(ctx, root, time.Now, summarizer)
+ // The report of the probe counts the failed calls, so the lines of the
+ // failed builds are not needed.
+ memory, err := openMemory(ctx, root, time.Now, summarizer, slog.New(slog.DiscardHandler))
if err != nil {
- return err
+ return fmt.Errorf("open probe memory: %w", err)
}
_, 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)
if err != nil {
- return errors.Join(err, memory.Close())
+ return fmt.Errorf("write probe header: %w", errors.Join(err, memory.Close()))
}
- run := probeRun{
+ run := BuildRun{
memory: memory,
client: client,
usagePath: usagePath,
}
appended, stop := run.Run(ctx, entries)
- sizes, err := modelLines(memory)
+ sizes, err := memory.buildModelLines()
err = errors.Join(err, memory.Close())
if err != nil {
- return err
+ return fmt.Errorf("finish probe memory: %w", err)
}
- records, err := readUsageFile(usagePath)
+ records, err := buildReadUsage(usagePath)
if err != nil {
- return err
- }
-
- return writeProbeReport(out, appended, stop, records, sizes)
-}
-
-// probeRun appends the messages of a probe run to its chat and stops the run
-// by the rules of probeStop.
-type probeRun struct {
- memory *Memory
- client *Client
- usagePath string
-}
-
-// Run appends the entries. It waits for the builds only before a user message,
-// as a turn does, so up to maxBuilds builds run at once in between, and it
-// waits for all builds at the end. It returns the number of messages appended
-// and why the run stopped early, or "" when it did not. After a stop it waits
-// for the calls in flight, so no sent call is cancelled.
-func (r *probeRun) Run(ctx context.Context, entries []probeEntry) (int, string) {
- appended, reason := r.append(ctx, entries)
- if reason == "" {
- return appended, ""
- }
-
- err := r.memory.Wait(ctx)
- if err != nil {
- reason += "; " + err.Error()
- }
-
- return appended, reason
-}
-
-// append appends the entries and stops at the first stop rule.
-func (r *probeRun) append(ctx context.Context, entries []probeEntry) (int, string) {
- for n, entry := range entries {
- if entry.Kind == "user" {
- reason := r.wait(ctx)
- if reason != "" {
- return n, reason
- }
- }
-
- _, err := r.memory.Append(entry.Kind, entry.Text)
- if err != nil {
- return n, err.Error()
- }
-
- reason := r.check()
- if reason != "" {
- return n + 1, reason
- }
- }
-
- return len(entries), r.wait(ctx)
-}
-
-// wait returns when no build runs, or when a stop rule ends the run. It checks
-// the rules each second, so a run stops soon also during a long wait.
-func (r *probeRun) wait(ctx context.Context) string {
- for {
- reason := r.check()
- if reason != "" {
- return reason
- }
-
- waitCtx, cancel := context.WithTimeout(ctx, time.Second)
- err := r.memory.Wait(waitCtx)
- cancel()
- if err == nil {
- return r.check()
- }
-
- if ctx.Err() != nil {
- return ctx.Err().Error()
- }
- }
-}
-
-// check returns why the run must stop, or "" when it can go on.
-func (r *probeRun) check() string {
- if r.client.LimitReached() {
- return fmt.Sprintf("call limit of %d reached", r.client.limit)
+ return fmt.Errorf("read probe usage: %w", err)
}
- records, err := readUsageFile(r.usagePath)
+ err = buildWriteReport(out, appended, stop, records, sizes)
if err != nil {
- return err.Error()
+ return fmt.Errorf("write probe report: %w", err)
}
- return probeStop(records)
+ return nil
}
-// slotSummarizer lets at most cap(slots) builds call inner at once. A build
-// past the limit waits for a slot.
-type slotSummarizer struct {
+// ProbeSlotSummarizer lets at most cap(slots) builds call inner at once. A
+// build past the limit waits for a slot.
+type ProbeSlotSummarizer struct {
inner Summarizer
slots chan struct{}
}
// Summarize waits for a slot and builds the node with inner.
-func (s slotSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
+func (s ProbeSlotSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
select {
case s.slots <- struct{}{}:
case <-ctx.Done():
return s.inner.Summarize(ctx, task)
}
-// probeEntry is one message of a probe log.
-type probeEntry struct {
- Kind string `json:"kind"`
- Text string `json:"text"`
-}
-
-// readProbeLog reads the first count messages of the probe log at path.
-func readProbeLog(path string, count int) (entries []probeEntry, err error) {
+// probeReadLog reads the first count messages of the probe log at path. Source
+// dates are ignored because a probe uses its own clock.
+func probeReadLog(path string, count int) (entries []DatedMessage, err error) {
file, err := os.Open(path)
if err != nil {
return nil, fmt.Errorf("open probe log: %w", err)
}
defer func() {
- err = errors.Join(err, file.Close())
+ closeErr := file.Close()
+ if closeErr != nil {
+ err = errors.Join(err, fmt.Errorf("close probe log: %w", closeErr))
+ }
}()
scanner := bufio.NewScanner(file)
- scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
+ scanner.Buffer(make([]byte, 64*1024), buildLineBytes)
+
for len(entries) < count && scanner.Scan() {
- var entry probeEntry
+ var entry DatedMessage
+
err = json.Unmarshal(scanner.Bytes(), &entry)
if err != nil {
return nil, fmt.Errorf("probe log line %d: %w", len(entries)+1, err)
}
+ entry.Date = ""
entries = append(entries, entry)
}
return entries, nil
}
-// probeStop returns why a probe run must stop: a failed call, or a quota window
-// whose used percent grew by more than probeQuotaLimit points since the first
-// call. It returns "" when the run can go on.
-func probeStop(records []usageRecord) string {
- for _, record := range records {
- if record.Error != "" {
- return "call failed: " + record.Error
- }
- }
-
- if len(records) == 0 {
- return ""
- }
-
- first := usedPercents(records[0])
- for name, used := range usedPercents(records[len(records)-1]) {
- start, found := first[name]
- if found && used-start > probeQuotaLimit {
- return fmt.Sprintf("%s grew from %g to %g", name, start, used)
- }
- }
-
- return ""
-}
-
-// modelLines returns the sizes of the nodes that the model wrote, in tree
-// order. A node that fits its source with no model call is not one of them.
-func modelLines(memory *Memory) ([]int, error) {
- memory.mutex.Lock()
- defer memory.mutex.Unlock()
-
- keys, err := memory.tree.Keys()
- if err != nil {
- return nil, err
- }
-
- var sizes []int
- for _, key := range keys {
- node, _, err := memory.tree.Node(key.level, key.index)
- if err != nil {
- return nil, err
- }
-
- model, err := builtByModel(memory, key)
- if err != nil {
- return nil, err
- }
-
- if model {
- sizes = append(sizes, node.Size)
- }
- }
-
- return sizes, nil
-}
-
-// builtByModel tells if the node key needed a model call: a leaf whose source
-// is over nodeBytes, or a merge whose children do not fit in one node.
-func builtByModel(memory *Memory, key nodeKey) (bool, error) {
- if key.level == 0 {
- message, err := memory.log.Message(key.index)
- if err != nil {
- return false, err
- }
-
- _, free := freeLeaf(message)
-
- return !free, nil
- }
-
- left, right, err := memory.children(key)
- if err != nil {
- return false, err
- }
-
- _, free := freeMerge(left, right)
-
- return !free, nil
-}
-
-// readUsageFile reads the records of usage.jsonl at path, in the order of the
-// start of their calls. The file has them in the order of the end of their
-// calls, which differs when calls run at once.
-func readUsageFile(path string) (records []usageRecord, err error) {
- file, err := os.Open(path)
- if errors.Is(err, os.ErrNotExist) {
- return nil, nil
- }
-
- if err != nil {
- return nil, fmt.Errorf("open usage file: %w", err)
- }
-
- defer func() {
- err = errors.Join(err, file.Close())
- }()
-
- scanner := bufio.NewScanner(file)
- scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
- for scanner.Scan() {
- var record usageRecord
- err = json.Unmarshal(scanner.Bytes(), &record)
- if err != nil {
- return nil, fmt.Errorf("decode usage: %w", err)
- }
-
- records = append(records, record)
- }
-
- err = scanner.Err()
- if err != nil {
- return nil, fmt.Errorf("read usage file: %w", err)
- }
-
- slices.SortStableFunc(records, func(a, b usageRecord) int {
- return callStart(a).Compare(callStart(b))
- })
-
- return records, nil
-}
-
-// callStart returns the start of the call of record, or the zero time when its
-// date is not valid.
-func callStart(record usageRecord) time.Time {
- start, err := time.Parse(time.RFC3339Nano, record.Date)
- if err != nil {
- return time.Time{}
- }
-
- return start
-}
-
-// usedPercents returns the used percent of each quota window of record, by
-// header name.
-func usedPercents(record usageRecord) map[string]float64 {
- used := make(map[string]float64)
- for name, value := range record.Quota {
- if !strings.HasSuffix(name, "-used-percent") {
- continue
- }
-
- percent, err := strconv.ParseFloat(value, 64)
- if err == nil {
- used[name] = percent
- }
- }
-
- return used
-}
-
-// writeProbeReport writes the result of a probe run.
-func writeProbeReport(out io.Writer, appended int, stop string, records []usageRecord, sizes []int) error {
- var input, cached, output, failed int
- for _, record := range records {
- input += record.InputTokens
- cached += record.CachedTokens
- output += record.OutputTokens
- if record.Error != "" {
- failed++
- }
- }
-
- var lines []string
- if stop != "" {
- lines = append(lines, "stopped early: "+stop)
- }
-
- lines = append(lines,
- fmt.Sprintf("messages: %d", appended),
- fmt.Sprintf("calls: %d (%d failed), %.2f per message", len(records), failed, ratio(len(records), appended)),
- fmt.Sprintf("model lines: %d, retries: %d, %.2f calls per model line", len(sizes), len(records)-failed-len(sizes), ratio(len(records), len(sizes))),
- fmt.Sprintf("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*ratio(cached, input), output),
- )
-
- if len(sizes) > 0 {
- lines = append(lines, sizeText(sizes))
- }
-
- if len(records) > 0 {
- lines = append(lines, timeText(records))
- }
-
- if len(records) > 0 {
- lines = append(lines,
- "quota at the first call: "+quotaText(records[0]),
- "quota at the last call: "+quotaText(records[len(records)-1]),
- )
- }
-
- _, err := io.WriteString(out, strings.Join(lines, "\n")+"\n")
-
- return err
-}
-
-// sizeText describes the sizes of the model lines.
-func sizeText(sizes []int) string {
- sorted := slices.Sorted(slices.Values(sizes))
- total, over := 0, 0
- for _, size := range sorted {
- total += size
- if size > nodeBytes {
- over++
- }
- }
-
- 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)
-}
-
-// timeText describes the time per call and how many calls ran at once. A call
-// runs from its date for its seconds.
-func timeText(records []usageRecord) string {
- type event struct {
- at time.Time
- delta int
- }
-
- var events []event
- seconds := make([]float64, 0, len(records))
- total := 0.0
- for _, record := range records {
- start, err := time.Parse(time.RFC3339Nano, record.Date)
- if err != nil {
- continue
- }
-
- length := time.Duration(record.Seconds * float64(time.Second))
- events = append(events, event{
- at: start,
- delta: 1,
- }, event{
- at: start.Add(length),
- delta: -1,
- })
- seconds = append(seconds, record.Seconds)
- total += record.Seconds
- }
-
- if len(events) == 0 {
- return "time: no dated calls"
- }
-
- slices.SortFunc(events, func(a, b event) int {
- return cmp.Or(a.at.Compare(b.at), cmp.Compare(a.delta, b.delta))
- })
-
- running, most := 0, 0
- for _, e := range events {
- running += e.delta
- most = max(most, running)
- }
-
- wall := events[len(events)-1].at.Sub(events[0].at).Seconds()
- slices.Sort(seconds)
-
- return fmt.Sprintf("time: %.0f s in all, %.1f calls per minute, %.1f s per call (median %.1f, max %.1f), calls at once: at most %d, %.1f on average", wall, 60*float64(len(seconds))/max(wall, 1e-9), total/float64(len(seconds)), seconds[len(seconds)/2], seconds[len(seconds)-1], most, total/max(wall, 1e-9))
-}
-
-// quotaText returns the quota headers of record, sorted by name.
-func quotaText(record usageRecord) string {
- names := slices.Sorted(maps.Keys(record.Quota))
- parts := make([]string, len(names))
- for n, name := range names {
- parts[n] = strings.TrimPrefix(name, "x-codex-") + "=" + record.Quota[name]
- }
-
- return strings.Join(parts, " ")
-}
-
-// ratio returns a / b, or 0 when b is 0.
-func ratio(a, b int) float64 {
- if b == 0 {
- return 0
- }
-
- return float64(a) / float64(b)
+// probePrompt returns the system prompt of a probe run: the prompt of the agent
+// probeName, without instructions and with a fixed instructions file path.
+func probePrompt() (string, error) {
+ return prompt(probeName, "/PROBE/instructions.md", ""), nil
}
projects / onechat / blobdiff