X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/63278e5c97fbad815c37c02f17d562367f4460a7..7b6d0a81e0f3c8e33b1b76d5ecbc5bdb25a87249:/probe.go

diff --git a/probe.go b/probe.go
index a1af30f..32a3ba9 100644
--- a/probe.go
+++ b/probe.go
@@ -33,17 +33,20 @@ const (
 	probeLineBytes = 1 << 20
 )
 
-// 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. 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 reports the
-// calls, the time per call, the calls at once, 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 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,
+// 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.
+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])
@@ -51,6 +54,14 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 		return fmt.Errorf("COUNT %q is not a positive number", args[1])
 	}
 
+	builds := maxBuilds
+	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)
+		}
+	}
+
 	entries, err := readProbeLog(args[0], count)
 	if err != nil {
 		return err
@@ -66,6 +77,15 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 		return err
 	}
 
+	tokens, err := lockTokens(dir)
+	if err != nil {
+		return err
+	}
+
+	defer func() {
+		err = errors.Join(err, tokens.Close())
+	}()
+
 	address, auth, err := modelEndpoint(config, dir)
 	if err != nil {
 		return err
@@ -78,14 +98,19 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 	}
 
 	usagePath := filepath.Join(root, "usage.jsonl")
-	client := newClient(address, auth, "onechat-"+probeName, usagePath)
+	client := newModelClient(address, auth, "onechat-"+probeName, usagePath)
 	client.LimitCalls(probeCallLimit)
-	memory, err := openMemory(ctx, root, time.Now, newCompactor(client, probeName, config.Compactor))
+	summarizer := slotSummarizer{
+		inner: newCompactor(client, probePrompt, config.Compactor),
+		slots: make(chan struct{}, builds),
+	}
+
+	memory, err := openMemory(ctx, root, time.Now, summarizer)
 	if err != nil {
 		return 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())
 	}
@@ -115,7 +140,7 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 // by the rules of probeStop.
 type probeRun struct {
 	memory    *Memory
-	client    *Client
+	client    *ModelClient
 	usagePath string
 }
 
@@ -198,6 +223,34 @@ func (r *probeRun) check() string {
 	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 {
 	Kind string `json:"kind"`
@@ -242,7 +295,7 @@ func readProbeLog(path string, count int) (entries []probeEntry, err error) {
 // 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 {
+func probeStop(records []ModelUsageRecord) string {
 	for _, record := range records {
 		if record.Error != "" {
 			return "call failed: " + record.Error
@@ -319,8 +372,10 @@ func builtByModel(memory *Memory, key nodeKey) (bool, error) {
 	return !free, nil
 }
 
-// readUsageFile reads the records of usage.jsonl at path.
-func readUsageFile(path string) (records []usageRecord, err error) {
+// 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 []ModelUsageRecord, err error) {
 	file, err := os.Open(path)
 	if errors.Is(err, os.ErrNotExist) {
 		return nil, nil
@@ -337,7 +392,7 @@ func readUsageFile(path string) (records []usageRecord, err error) {
 	scanner := bufio.NewScanner(file)
 	scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
 	for scanner.Scan() {
-		var record usageRecord
+		var record ModelUsageRecord
 		err = json.Unmarshal(scanner.Bytes(), &record)
 		if err != nil {
 			return nil, fmt.Errorf("decode usage: %w", err)
@@ -351,12 +406,27 @@ func readUsageFile(path string) (records []usageRecord, err error) {
 		return nil, fmt.Errorf("read usage file: %w", err)
 	}
 
+	slices.SortStableFunc(records, func(a, b ModelUsageRecord) 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 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 usageRecord) map[string]float64 {
+func usedPercents(record ModelUsageRecord) map[string]float64 {
 	used := make(map[string]float64)
 	for name, value := range record.Quota {
 		if !strings.HasSuffix(name, "-used-percent") {
@@ -373,7 +443,7 @@ func usedPercents(record usageRecord) map[string]float64 {
 }
 
 // writeProbeReport writes the result of a probe run.
-func writeProbeReport(out io.Writer, appended int, stop string, records []usageRecord, sizes []int) error {
+func writeProbeReport(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
@@ -432,7 +502,7 @@ func sizeText(sizes []int) string {
 
 // 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 {
+func timeText(records []ModelUsageRecord) string {
 	type event struct {
 		at    time.Time
 		delta int
@@ -476,11 +546,11 @@ func timeText(records []usageRecord) string {
 	wall := events[len(events)-1].at.Sub(events[0].at).Seconds()
 	slices.Sort(seconds)
 
-	return fmt.Sprintf("time: %.0f s in all, %.1f s per call (median %.1f, max %.1f), calls at once: at most %d, %.1f on average", wall, 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 usageRecord) string {
+func quotaText(record ModelUsageRecord) string {
 	names := slices.Sorted(maps.Keys(record.Quota))
 	parts := make([]string, len(names))
 	for n, name := range names {