Name and order the probe code

[onechat] / probe.go

diff --git a/probe.go b/probe.go
index 32a3ba99b27f98997bfe113839f931cca61882a5..0b809e89f84d86ef7150b3f9dcfb18b50b6d77d4 100644 (file)
--- a/probe.go
+++ b/probe.go
@@ -40,10 +40,10 @@ const (
 // 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
 // reports the calls, the time per call, the calls at once, the cached share,
 // 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
 // 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
 // limits the builds that call the model at once in this run only. The core

-// still queues up to maxBuilds builds. The run holds the lock of the token
-// file, so it does not run while the daemon runs.
+// 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]")
 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]")


@@ -54,15 +54,15 @@ func probe(ctx context.Context, args []string, out io.Writer) (err error) {
                return fmt.Errorf("COUNT %q is not a positive number", args[1])
        }
 
                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 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 || 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
        }
        if err != nil {
                return err
        }


@@ -100,7 +100,7 @@ func probe(ctx context.Context, args []string, out io.Writer) (err error) {
        usagePath := filepath.Join(root, "usage.jsonl")
        client := newModelClient(address, auth, "onechat-"+probeName, usagePath)
        client.LimitCalls(probeCallLimit)
        usagePath := filepath.Join(root, "usage.jsonl")
        client := newModelClient(address, auth, "onechat-"+probeName, usagePath)
        client.LimitCalls(probeCallLimit)

-       summarizer := slotSummarizer{
+       summarizer := ProbeSlotSummarizer{

                inner: newCompactor(client, probePrompt, config.Compactor),
                slots: make(chan struct{}, builds),
        }
                inner: newCompactor(client, probePrompt, config.Compactor),
                slots: make(chan struct{}, builds),
        }


@@ -115,41 +115,63 @@ func probe(ctx context.Context, args []string, out io.Writer) (err error) {
                return errors.Join(err, memory.Close())
        }
 
                return errors.Join(err, memory.Close())
        }
 

-       run := probeRun{
+       run := ProbeRun{

                memory:    memory,
                client:    client,
                usagePath: usagePath,
        }
 
        appended, stop := run.Run(ctx, entries)
                memory:    memory,
                client:    client,
                usagePath: usagePath,
        }
 
        appended, stop := run.Run(ctx, entries)

-       sizes, err := modelLines(memory)
+       sizes, err := probeModelLines(memory)

        err = errors.Join(err, memory.Close())
        if err != nil {
                return err
        }
 
        err = errors.Join(err, memory.Close())
        if err != nil {
                return err
        }
 

-       records, err := readUsageFile(usagePath)
+       records, err := probeReadUsage(usagePath)

        if err != nil {
                return err
        }
 
        if err != nil {
                return err
        }
 

-       return writeProbeReport(out, appended, stop, records, sizes)
+       return probeWriteReport(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 {
+// 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 ProbeSlotSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
+       select {
+       case s.slots <- struct{}{}:
+       case <-ctx.Done():
+               return "", ctx.Err()
+       }
+
+       defer func() {
+               <-s.slots
+       }()
+
+       return s.inner.Summarize(ctx, task)
+}
+
+// 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    *ModelClient
        usagePath string
 }
 
 // Run appends the entries. It waits for the builds only before a user message,
        memory    *Memory
        client    *ModelClient
        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) {
+// as a turn does, so up to buildMax 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, ""
        appended, reason := r.append(ctx, entries)
        if reason == "" {
                return appended, ""


@@ -164,7 +186,7 @@ func (r *probeRun) Run(ctx context.Context, entries []probeEntry) (int, string)
 }
 
 // append appends the entries and stops at the first stop rule.
 }
 
 // append appends the entries and stops at the first stop rule.

-func (r *probeRun) append(ctx context.Context, entries []probeEntry) (int, string) {
+func (r *ProbeRun) append(ctx context.Context, entries []ProbeEntry) (int, string) {

        for n, entry := range entries {
                if entry.Kind == "user" {
                        reason := r.wait(ctx)
        for n, entry := range entries {
                if entry.Kind == "user" {
                        reason := r.wait(ctx)


@@ -189,7 +211,7 @@ func (r *probeRun) append(ctx context.Context, entries []probeEntry) (int, strin
 
 // 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.
 
 // 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 {
+func (r *ProbeRun) wait(ctx context.Context) string {

        for {
                reason := r.check()
                if reason != "" {
        for {
                reason := r.check()
                if reason != "" {


@@ -210,12 +232,12 @@ func (r *probeRun) wait(ctx context.Context) string {
 }
 
 // check returns why the run must stop, or "" when it can go on.
 }
 
 // check returns why the run must stop, or "" when it can go on.

-func (r *probeRun) check() string {
+func (r *ProbeRun) check() string {

        if r.client.LimitReached() {
                return fmt.Sprintf("call limit of %d reached", r.client.limit)
        }
 
        if r.client.LimitReached() {
                return fmt.Sprintf("call limit of %d reached", r.client.limit)
        }
 

-       records, err := readUsageFile(r.usagePath)
+       records, err := probeReadUsage(r.usagePath)

        if err != nil {
                return err.Error()
        }
        if err != nil {
                return err.Error()
        }


@@ -223,42 +245,14 @@ func (r *probeRun) check() string {
        return probeStop(records)
 }
 
        return probeStop(records)
 }
 

-// slotSummarizer lets at most cap(slots) builds call inner at once. A build
-// past the limit waits for a slot.
-type slotSummarizer 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) {
-       select {
-       case s.slots <- struct{}{}:
-       case <-ctx.Done():
-               return "", ctx.Err()
-       }
-
-       defer func() {
-               <-s.slots
-       }()
-
-       return s.inner.Summarize(ctx, task)
-}
-
-// probePrompt returns the system prompt of a probe run: the prompt of the agent
-// probeName, without instructions.
-func probePrompt() (string, error) {
-       return prompt(probeName, ""), nil
-}
-
-// probeEntry is one message of a probe log.
-type probeEntry struct {
+// ProbeEntry is one message of a probe log.
+type ProbeEntry struct {

        Kind string `json:"kind"`
        Text string `json:"text"`
 }
 
        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.
+func probeReadLog(path string, count int) (entries []ProbeEntry, err error) {

        file, err := os.Open(path)
        if err != nil {
                return nil, fmt.Errorf("open probe log: %w", err)
        file, err := os.Open(path)
        if err != nil {
                return nil, fmt.Errorf("open probe log: %w", err)


@@ -271,7 +265,7 @@ func readProbeLog(path string, count int) (entries []probeEntry, err error) {
        scanner := bufio.NewScanner(file)
        scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
        for len(entries) < count && scanner.Scan() {
        scanner := bufio.NewScanner(file)
        scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
        for len(entries) < count && scanner.Scan() {

-               var entry probeEntry
+               var entry ProbeEntry

                err = json.Unmarshal(scanner.Bytes(), &entry)
                if err != nil {
                        return nil, fmt.Errorf("probe log line %d: %w", len(entries)+1, err)
                err = json.Unmarshal(scanner.Bytes(), &entry)
                if err != nil {
                        return nil, fmt.Errorf("probe log line %d: %w", len(entries)+1, err)


@@ -292,34 +286,15 @@ func readProbeLog(path string, count int) (entries []probeEntry, err error) {
        return entries, nil
 }
 
        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 []ModelUsageRecord) 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 ""
+// probePrompt returns the system prompt of a probe run: the prompt of the agent
+// probeName, without instructions.
+func probePrompt() (string, error) {
+       return prompt(probeName, ""), nil

 }
 
 }
 

-// modelLines returns the sizes of the nodes that the model wrote, in tree
+// probeModelLines 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.
 // order. A node that fits its source with no model call is not one of them.

-func modelLines(memory *Memory) ([]int, error) {
+func probeModelLines(memory *Memory) ([]int, error) {

        memory.mutex.Lock()
        defer memory.mutex.Unlock()
 
        memory.mutex.Lock()
        defer memory.mutex.Unlock()
 


@@ -335,7 +310,7 @@ func modelLines(memory *Memory) ([]int, error) {
                        return nil, err
                }
 
                        return nil, err
                }
 

-               model, err := builtByModel(memory, key)
+               model, err := probeBuiltByModel(memory, key)

                if err != nil {
                        return nil, err
                }
                if err != nil {
                        return nil, err
                }


@@ -348,34 +323,10 @@ func modelLines(memory *Memory) ([]int, error) {
        return sizes, nil
 }
 
        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
+// probeReadUsage 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.
 // 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 []ModelUsageRecord, err error) {
+func probeReadUsage(path string) (records []ModelUsageRecord, err error) {

        file, err := os.Open(path)
        if errors.Is(err, os.ErrNotExist) {
                return nil, nil
        file, err := os.Open(path)
        if errors.Is(err, os.ErrNotExist) {
                return nil, nil


@@ -407,43 +358,14 @@ func readUsageFile(path string) (records []ModelUsageRecord, err error) {
        }
 
        slices.SortStableFunc(records, func(a, b ModelUsageRecord) int {
        }
 
        slices.SortStableFunc(records, func(a, b ModelUsageRecord) int {

-               return callStart(a).Compare(callStart(b))
+               return probeCallStart(a).Compare(probeCallStart(b))

        })
 
        return records, nil
 }
 
        })
 
        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 ModelUsageRecord) 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 ModelUsageRecord) 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 []ModelUsageRecord, sizes []int) error {
+// probeWriteReport writes the result of a probe run.
+func probeWriteReport(out io.Writer, appended int, stop string, records []ModelUsageRecord, sizes []int) error {

        var input, cached, output, failed int
        for _, record := range records {
                input += record.InputTokens
        var input, cached, output, failed int
        for _, record := range records {
                input += record.InputTokens


@@ -461,23 +383,23 @@ func writeProbeReport(out io.Writer, appended int, stop string, records []ModelU
 
        lines = append(lines,
                fmt.Sprintf("messages: %d", appended),
 
        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),
+               fmt.Sprintf("calls: %d (%d failed), %.2f per message", len(records), failed, probeRatio(len(records), appended)),
+               fmt.Sprintf("model lines: %d, retries: %d, %.2f calls per model line", len(sizes), len(records)-failed-len(sizes), probeRatio(len(records), len(sizes))),
+               fmt.Sprintf("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*probeRatio(cached, input), output),

        )
 
        if len(sizes) > 0 {
        )
 
        if len(sizes) > 0 {

-               lines = append(lines, sizeText(sizes))
+               lines = append(lines, probeSizeText(sizes))

        }
 
        if len(records) > 0 {
        }
 
        if len(records) > 0 {

-               lines = append(lines, timeText(records))
+               lines = append(lines, probeTimeText(records))

        }
 
        if len(records) > 0 {
                lines = append(lines,
        }
 
        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]),
+                       "quota at the first call: "+probeQuotaText(records[0]),
+                       "quota at the last call: "+probeQuotaText(records[len(records)-1]),

                )
        }
 
                )
        }
 


@@ -486,8 +408,77 @@ func writeProbeReport(out io.Writer, appended int, stop string, records []ModelU
        return err
 }
 
        return err
 }
 

-// sizeText describes the sizes of the model lines.
-func sizeText(sizes []int) string {
+// 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 []ModelUsageRecord) string {
+       for _, record := range records {
+               if record.Error != "" {
+                       return "call failed: " + record.Error
+               }
+       }
+
+       if len(records) == 0 {
+               return ""
+       }
+
+       first := probeUsedPercents(records[0])
+       for name, used := range probeUsedPercents(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 ""
+}
+
+// probeBuiltByModel 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 probeBuiltByModel(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
+}
+
+// probeCallStart returns the start of the call of record, or the zero time when
+// its date is not valid.
+func probeCallStart(record ModelUsageRecord) time.Time {
+       start, err := time.Parse(time.RFC3339Nano, record.Date)
+       if err != nil {
+               return time.Time{}
+       }
+
+       return start
+}
+
+// probeRatio returns a / b, or 0 when b is 0.
+func probeRatio(a, b int) float64 {
+       if b == 0 {
+               return 0
+       }
+
+       return float64(a) / float64(b)
+}
+
+// probeSizeText describes the sizes of the model lines.
+func probeSizeText(sizes []int) string {

        sorted := slices.Sorted(slices.Values(sizes))
        total, over := 0, 0
        for _, size := range sorted {
        sorted := slices.Sorted(slices.Values(sizes))
        total, over := 0, 0
        for _, size := range sorted {


@@ -497,12 +488,12 @@ func sizeText(sizes []int) string {
                }
        }
 
                }
        }
 

-       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)
+       return fmt.Sprintf("line bytes: min %d, median %d, mean %.0f, max %d, %d over %d", sorted[0], sorted[len(sorted)/2], probeRatio(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 []ModelUsageRecord) string {
+// probeTimeText describes the time per call and how many calls ran at once. A
+// call runs from its date for its seconds.
+func probeTimeText(records []ModelUsageRecord) string {

        type event struct {
                at    time.Time
                delta int
        type event struct {
                at    time.Time
                delta int


@@ -549,8 +540,8 @@ func timeText(records []ModelUsageRecord) string {
        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))
 }
 
        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 ModelUsageRecord) string {
+// probeQuotaText returns the quota headers of record, sorted by name.
+func probeQuotaText(record ModelUsageRecord) string {

        names := slices.Sorted(maps.Keys(record.Quota))
        parts := make([]string, len(names))
        for n, name := range names {
        names := slices.Sorted(maps.Keys(record.Quota))
        parts := make([]string, len(names))
        for n, name := range names {


@@ -560,11 +551,20 @@ func quotaText(record ModelUsageRecord) string {
        return strings.Join(parts, " ")
 }
 
        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
+// probeUsedPercents returns the used percent of each quota window of record, by
+// header name.
+func probeUsedPercents(record ModelUsageRecord) 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 float64(a) / float64(b)
+       return used

 }
 }