X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/ba9faee0cd6c35020997794ef7bf134e8db12e96..e0ab6220b6fd35f7e92a04e4bf149940185713ab:/responses.go?ds=inline diff --git a/responses.go b/responses.go index 096439f..19ab6a1 100644 --- 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,25 @@ 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. Turn names the request of a turn: +// first, next, or warmup. Node and Turn go only to the usage file, never to the +// endpoint. A call with StopAtOutput ends at the first output item, after the +// endpoint has read the whole input, and gives no output and no usage. type Request struct { Model string Effort string Instructions string Tools []Tool Input []any + Node *nodeKey + Turn string + StopAtOutput bool } // Tool is the definition of a function tool. @@ -188,8 +201,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,26 +252,34 @@ 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, } - response, err := c.send(httpRequest, &record) + if request.Node != nil { + record.Node = &[2]int{request.Node.level, request.Node.index} + } + + record.Turn = request.Turn + response, err := c.send(httpRequest, &record, request.StopAtOutput) 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)) } // send makes the HTTP call and reads its stream. It puts the status and the // quota headers into record. -func (c *Client) send(httpRequest *http.Request, record *usageRecord) (response Response, err error) { +func (c *Client) send(httpRequest *http.Request, record *usageRecord, stopAtOutput bool) (response Response, err error) { httpResponse, err := c.http.Do(httpRequest) if err != nil { return Response{}, fmt.Errorf("send model call: %w", err) @@ -262,7 +298,7 @@ func (c *Client) send(httpRequest *http.Request, record *usageRecord) (response return Response{}, errors.Join(err, readErr) } - return readStream(httpResponse.Body) + return readStream(httpResponse.Body, stopAtOutput) } // usageLog appends one JSON line for each model call to usage.jsonl. Calls can @@ -272,16 +308,22 @@ 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. Turn names the request of a turn. 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"` + Turn string `json:"turn,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"` } @@ -364,8 +406,9 @@ func quotaHeaders(header http.Header) map[string]string { // readStream reads the server-sent events of a call until the response // completes. The output items come from the output_item.done events, or from -// the completed response when the stream sent none. -func readStream(stream io.Reader) (Response, error) { +// the completed response when the stream sent none. With stopAtOutput it ends +// at the first output item. +func readStream(stream io.Reader, stopAtOutput bool) (Response, error) { var response Response reader := bufio.NewReader(stream) var data bytes.Buffer @@ -383,7 +426,7 @@ func readStream(stream io.Reader) (Response, error) { data.WriteByte('\n') } } else if data.Len() > 0 { - done, eventErr := readEvent(&response, data.Bytes()) + done, eventErr := readEvent(&response, data.Bytes(), stopAtOutput) data.Reset() if eventErr != nil || done { return response, eventErr @@ -398,7 +441,7 @@ func readStream(stream io.Reader) (Response, error) { // readEvent adds one stream event to response, and tells if it ends the // stream. -func readEvent(response *Response, data []byte) (bool, error) { +func readEvent(response *Response, data []byte, stopAtOutput bool) (bool, error) { var event streamEvent err := json.Unmarshal(data, &event) if err != nil { @@ -406,6 +449,8 @@ func readEvent(response *Response, data []byte) (bool, error) { } switch event.Type { + case "response.output_item.added": + return stopAtOutput, nil case "response.output_item.done": response.Output = append(response.Output, event.Item) case "response.completed":