X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/522d226230a2eae0a8a33ae0fbf32f96c6886b26..48d7bce28b18ecef4c52c98821675605464cecb9:/probe.go

diff --git a/probe.go b/probe.go
index acd4f6b..384d0f8 100644
--- a/probe.go
+++ b/probe.go
@@ -7,12 +7,10 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"maps"
+	"log/slog"
 	"os"
 	"path/filepath"
-	"slices"
 	"strconv"
-	"strings"
 	"time"
 )
 
@@ -21,49 +19,77 @@ const (
 	// 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
-
-	// probeLineBytes is the most bytes of one line of a probe log.
-	probeLineBytes = 1 << 20
+	// probeCallLimit is the most model calls of a probe run.
+	probeCallLimit = 660
 )
 
-// probe runs "onechat probe LOG COUNT". It appends the first COUNT messages of
-// LOG, one JSON object {kind, text} per line, to a new chat in probe/ in the
-// data directory, and builds its tree with the real compactor. After each
-// message it waits for the builds, as a turn does. It stops at the first
-// failed call, or when the used percent of a quota window grows by more than
-// probeQuotaLimit points. At the end it reports the calls, the cached share,
-// the quota, the retries, and the line sizes.
-func probe(ctx context.Context, args []string, out io.Writer) error {
-	if len(args) != 2 {
-		return errors.New("usage: onechat probe LOG COUNT")
+// probe runs "onechat probe LOG COUNT [BUILDS]". It appends the first COUNT
+// messages of LOG, one JSON object {kind, text} per line, to a new chat in
+// 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 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 buildMax,
+// limits the builds that call the model at once in this run only. The core
+// 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])
 	}
 
-	entries, err := readProbeLog(args[0], count)
+	builds := buildMax
+	if len(args) == 3 {
+		builds, err = strconv.Atoi(args[2])
+		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 := 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"))
@@ -73,86 +99,101 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 	}
 
 	usagePath := filepath.Join(root, "usage.jsonl")
-	client := newClient(address, auth, "onechat-"+probeName, usagePath)
-	memory, err := openMemory(ctx, root, time.Now, newCompactor(client, probeName, config.Compactor))
+	client := newModelClient(address, auth, "onechat-"+probeName, usagePath)
+	client.LimitCalls(probeCallLimit)
+
+	summarizer := ProbeSlotSummarizer{
+		inner: newCompactor(client, probePrompt, config.Compactor, promptTools(), &ImageStore{dir: filepath.Join(dir, imageDirName)}),
+		slots: make(chan struct{}, builds),
+	}
+
+	// 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\n", len(entries), root, config.Compactor.Model, config.Compactor.Effort)
+	_, 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 := BuildRun{
+		memory:    memory,
+		client:    client,
+		usagePath: usagePath,
 	}
 
-	appended, stop := runProbe(ctx, memory, entries, usagePath)
-	sizes, err := modelLines(memory)
+	appended, stop := run.Run(ctx, entries)
+	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 fmt.Errorf("read probe usage: %w", err)
 	}
 
-	return writeProbeReport(out, appended, stop, records, sizes)
-}
-
-// runProbe appends the entries and waits for the builds after each one. It
-// returns the number of messages appended and why the run stopped early, or
-// "" when it did not.
-func runProbe(ctx context.Context, memory *Memory, entries []probeEntry, usagePath string) (int, string) {
-	for n, entry := range entries {
-		_, err := memory.Append(entry.Kind, entry.Text)
-		if err == nil {
-			err = memory.Wait(ctx)
-		}
+	err = buildWriteReport(out, appended, stop, records, sizes)
+	if err != nil {
+		return fmt.Errorf("write probe report: %w", err)
+	}
 
-		if err != nil {
-			return n, err.Error()
-		}
+	return nil
+}
 
-		records, err := readUsageFile(usagePath)
-		if err != nil {
-			return n + 1, err.Error()
-		}
+// 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{}
+}
 
-		reason := probeStop(records)
-		if reason != "" {
-			return n + 1, reason
-		}
+// 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()
 	}
 
-	return len(entries), ""
-}
+	defer func() {
+		<-s.slots
+	}()
 
-// probeEntry is one message of a probe log.
-type probeEntry struct {
-	Kind string `json:"kind"`
-	Text string `json:"text"`
+	return s.inner.Summarize(ctx, task)
 }
 
-// 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)
 	}
 
@@ -168,209 +209,8 @@ func readProbeLog(path string, count int) (entries []probeEntry, err error) {
 	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.
-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)
-	}
-
-	return records, nil
-}
-
-// 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", len(sizes), len(records)-failed-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,
-			"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)
-}
-
-// 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
 }