X-Git-Url: http://git.giacomolli.uk/onechat/blobdiff_plain/d647f33b3e70f9419f6c0740e1d01b210439a519..a345df4e4bbf9dbaf3da1ce548896eef8f515de3:/responses.go?ds=sidebyside

diff --git a/responses.go b/responses.go
index fddec73..0ba4955 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,12 +67,15 @@ 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. Node names the
-// node that a compaction builds, or is nil. It goes only to the usage file,
-// never to the endpoint.
+// node that a compaction builds, or is nil. Turn names the request of a turn:
+// first or next. Node and Turn go only to the usage file, never to the
+// endpoint.
 type Request struct {
 	Model        string
 	Effort       string
@@ -76,6 +83,7 @@ type Request struct {
 	Tools        []Tool
 	Input        []any
 	Node         *nodeKey
+	Turn         string
 }
 
 // Tool is the definition of a function tool.
@@ -191,8 +199,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,
@@ -237,6 +260,7 @@ func (c *Client) Create(ctx context.Context, request Request) (Response, error)
 		record.Node = &[2]int{request.Node.level, request.Node.index}
 	}
 
+	record.Turn = request.Turn
 	response, err := c.send(httpRequest, &record)
 	if err != nil {
 		record.Error = err.Error()
@@ -285,13 +309,14 @@ type usageLog struct {
 // 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
+// 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"`