X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/d647f33b3e70f9419f6c0740e1d01b210439a519..fe39a670d7af062af089bdba4dd07c3f0aee4048:/responses.go

diff --git a/responses.go b/responses.go
index fddec73..3055e70 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,6 +67,8 @@ 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
@@ -191,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,