]> git.giacomolli.uk - onechat/blobdiff - probe.go
Name and order the build queue code
[onechat] / probe.go
index bbce24f65541d455b05796a4d30c58fdf5903e6d..e12516d12b19559e532c067a59ea10ccf73e914e 100644 (file)
--- a/probe.go
+++ b/probe.go
@@ -40,10 +40,11 @@ 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.
-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]")
        }
        if len(args) != 2 && len(args) != 3 {
                return errors.New("usage: onechat probe LOG COUNT [BUILDS]")
        }
@@ -53,11 +54,11 @@ func probe(ctx context.Context, args []string, out io.Writer) 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)
                }
        }
 
                }
        }
 
@@ -76,6 +77,15 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
                return err
        }
 
                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
        address, auth, err := modelEndpoint(config, dir)
        if err != nil {
                return err
@@ -88,10 +98,10 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
        }
 
        usagePath := filepath.Join(root, "usage.jsonl")
        }
 
        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{
        client.LimitCalls(probeCallLimit)
        summarizer := slotSummarizer{
-               inner: newCompactor(client, probeName, config.Compactor),
+               inner: newCompactor(client, probePrompt, config.Compactor),
                slots: make(chan struct{}, builds),
        }
 
                slots: make(chan struct{}, builds),
        }
 
@@ -130,15 +140,15 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 // by the rules of probeStop.
 type probeRun struct {
        memory    *Memory
 // by the rules of probeStop.
 type probeRun struct {
        memory    *Memory
-       client    *Client
+       client    *ModelClient
        usagePath string
 }
 
 // Run appends the entries. It waits for the builds only before a user message,
        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.
+// 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 == "" {
 func (r *probeRun) Run(ctx context.Context, entries []probeEntry) (int, string) {
        appended, reason := r.append(ctx, entries)
        if reason == "" {
@@ -235,6 +245,12 @@ func (s slotSummarizer) Summarize(ctx context.Context, task Task) (string, error
        return s.inner.Summarize(ctx, task)
 }
 
        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"`
 // probeEntry is one message of a probe log.
 type probeEntry struct {
        Kind string `json:"kind"`
@@ -279,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.
 // 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
        for _, record := range records {
                if record.Error != "" {
                        return "call failed: " + record.Error
@@ -359,7 +375,7 @@ func builtByModel(memory *Memory, key nodeKey) (bool, 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.
 // 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) {
+func readUsageFile(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
@@ -376,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() {
        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)
                err = json.Unmarshal(scanner.Bytes(), &record)
                if err != nil {
                        return nil, fmt.Errorf("decode usage: %w", err)
@@ -390,7 +406,7 @@ func readUsageFile(path string) (records []usageRecord, err error) {
                return nil, fmt.Errorf("read usage file: %w", err)
        }
 
                return nil, fmt.Errorf("read usage file: %w", err)
        }
 
-       slices.SortStableFunc(records, func(a, b usageRecord) int {
+       slices.SortStableFunc(records, func(a, b ModelUsageRecord) int {
                return callStart(a).Compare(callStart(b))
        })
 
                return callStart(a).Compare(callStart(b))
        })
 
@@ -399,7 +415,7 @@ func readUsageFile(path string) (records []usageRecord, err error) {
 
 // callStart returns the start of the call of record, or the zero time when its
 // date is not valid.
 
 // 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 {
+func callStart(record ModelUsageRecord) time.Time {
        start, err := time.Parse(time.RFC3339Nano, record.Date)
        if err != nil {
                return time.Time{}
        start, err := time.Parse(time.RFC3339Nano, record.Date)
        if err != nil {
                return time.Time{}
@@ -410,7 +426,7 @@ func callStart(record usageRecord) time.Time {
 
 // usedPercents returns the used percent of each quota window of record, by
 // header name.
 
 // 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") {
        used := make(map[string]float64)
        for name, value := range record.Quota {
                if !strings.HasSuffix(name, "-used-percent") {
@@ -427,7 +443,7 @@ func usedPercents(record usageRecord) map[string]float64 {
 }
 
 // writeProbeReport writes the result of a probe run.
 }
 
 // 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
        var input, cached, output, failed int
        for _, record := range records {
                input += record.InputTokens
@@ -486,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.
 
 // 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
        type event struct {
                at    time.Time
                delta int
@@ -534,7 +550,7 @@ func timeText(records []usageRecord) string {
 }
 
 // quotaText returns the quota headers of record, sorted by name.
 }
 
 // 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 {
        names := slices.Sorted(maps.Keys(record.Quota))
        parts := make([]string, len(names))
        for n, name := range names {