Send the compaction view in blocks of four lines

[onechat] / responses.go

diff --git a/responses.go b/responses.go
index 096439f5cc9564d92f8fcfed7636ae60960c91a6..3055e706eb985c06f7b0a593877d3cd8f1aa3f24 100644 (file)
--- a/responses.go
+++ b/responses.go
@@ -13,6 +13,7 @@ import (
        "path/filepath"
        "strings"
        "sync"
+       "sync/atomic"
        "time"
 )
 
@@ -36,6 +37,9 @@ const (
        errorBytes = 4096
 )
 
+// errCallLimit tells that a client sent its limit of calls.
+var errCallLimit = errors.New("call limit reached")
+
 // authorizer adds the credentials of a model call to the headers of its
 // request.
 type authorizer interface {
@@ -63,16 +67,21 @@ type Client struct {
        cacheKey string
        usage    *usageLog
        http     *http.Client
+       limit    int64
+       calls    atomic.Int64
 }
 
 // Request is one model call. Input holds the input items in order, such as
-// userMessage values or output items of an earlier response.
+// userMessage values or output items of an earlier response. Node names the
+// node that a compaction builds, or is nil. It goes only to the usage file,
+// never to the endpoint.
 type Request struct {
        Model        string
        Effort       string
        Instructions string
        Tools        []Tool
        Input        []any
+       Node         *nodeKey
 }
 
 // Tool is the definition of a function tool.
@@ -188,8 +197,23 @@ func newClient(url string, auth authorizer, cacheKey, usagePath string) *Client
        }
 }
 
-// Create sends the call and reads its event stream to the end.
+// LimitCalls makes the client refuse each call after the first limit calls.
+func (c *Client) LimitCalls(limit int64) {
+       c.limit = limit
+}
+
+// LimitReached tells if the client refuses further calls.
+func (c *Client) LimitReached() bool {
+       return c.limit > 0 && c.calls.Load() >= c.limit
+}
+
+// Create sends the call and reads its event stream to the end. Over the call
+// limit it sends nothing and returns errCallLimit.
 func (c *Client) Create(ctx context.Context, request Request) (Response, error) {
+       if c.calls.Add(1) > c.limit && c.limit > 0 {
+               return Response{}, errCallLimit
+       }
+
        body, err := json.Marshal(requestBody{
                Model:        request.Model,
                Instructions: request.Instructions,
@@ -224,19 +248,26 @@ func (c *Client) Create(ctx context.Context, request Request) (Response, error)
                httpRequest.Header.Set("session_id", c.cacheKey)
        }
 
+       start := time.Now()
        record := usageRecord{
-               Date:  time.Now().Format(time.RFC3339Nano),
+               Date:  start.Format(time.RFC3339Nano),
                Model: request.Model,
        }
 
+       if request.Node != nil {
+               record.Node = &[2]int{request.Node.level, request.Node.index}
+       }
+
        response, err := c.send(httpRequest, &record)
        if err != nil {
                record.Error = err.Error()
        }
 
+       record.Seconds = time.Since(start).Seconds()
        record.InputTokens = response.Usage.Input
        record.CachedTokens = response.Usage.Cached
        record.OutputTokens = response.Usage.Output
+       record.TextBytes = len(response.Text)
 
        return response, errors.Join(err, c.usage.Write(record))
 }
@@ -272,16 +303,21 @@ type usageLog struct {
        path  string
 }
 
-// usageRecord is one line of usage.jsonl. Status is the HTTP status, or 0 when
-// no response came. Quota holds the x-codex-* headers of the response, without
-// the turn state.
+// usageRecord is one line of usage.jsonl. Date is the start of the call, and
+// Seconds its length. Node is the node [l, i] of a compaction. Status is the
+// HTTP status, or 0 when no response came. TextBytes is the size of the text of
+// the reply. Quota holds the x-codex-* headers of the response, without the
+// turn state.
 type usageRecord struct {
        Date         string            `json:"date"`
+       Seconds      float64           `json:"seconds"`
        Model        string            `json:"model"`
+       Node         *[2]int           `json:"node,omitempty"`
        Status       int               `json:"status"`
        InputTokens  int               `json:"input_tokens"`
        CachedTokens int               `json:"cached_tokens"`
        OutputTokens int               `json:"output_tokens"`
+       TextBytes    int               `json:"text_bytes"`
        Quota        map[string]string `json:"quota,omitempty"`
        Error        string            `json:"error,omitempty"`
 }