Name package-level code for its feature

[onechat] / probe.go

diff --git a/probe.go b/probe.go
index bbce24f65541d455b05796a4d30c58fdf5903e6d..9d272ec69286596881b34a727fea09fe502683cf 100644 (file)
--- a/probe.go
+++ b/probe.go
@@ -42,8 +42,9 @@ const (
 // 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.
-func probe(ctx context.Context, args []string, out io.Writer) error {
+// 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]")
        }
@@ -76,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
@@ -91,7 +101,7 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
        client := newClient(address, auth, "onechat-"+probeName, usagePath)
        client.LimitCalls(probeCallLimit)
        summarizer := slotSummarizer{
-               inner: newCompactor(client, probeName, config.Compactor),
+               inner: newCompactor(client, probePrompt, config.Compactor),
                slots: make(chan struct{}, builds),
        }
 
@@ -235,6 +245,12 @@ func (s slotSummarizer) Summarize(ctx context.Context, task Task) (string, error
        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"`