Send the compaction view in blocks of four lines

[onechat] / responses.go

diff --git a/responses.go b/responses.go
index fddec7378650a7ca9c25e6ac1c960f5f53313f19..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,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,