feature: add client rate limit support

You can specify your own rate limiter to limit the number of requests sent to the remote API. The easiest way is probably to use the off the shelf Limiter implemented by https://pkg.go.dev/golang.org/x/time/rate#Limiter We had to refactor the API a bit so we can wrap the http client with the provided limiter and also make room for adding retries. Signed-off-by: Milos Gajdos <[email protected]>
milosgajdos · Mar 22, 2024 · e77ae2f · e77ae2f
1 parent da75163
commit e77ae2f
Show file tree

Hide file tree

Showing 8 changed files with 91 additions and 19 deletions.
diff --git a/client/client.go b/client/client.go
@@ -0,0 +1,70 @@
+package client
+
+import (
+	"context"
+	"net/http"
+)
+
+// HTTP is a HTTP client.
+type HTTP struct {
+	c  *http.Client
+	rl Limiter
+}
+
+// Options are client options
+type Options struct {
+	HTTPClient *http.Client
+	Limiter    Limiter
+}
+
+// Option is functional graph option.
+type Option func(*Options)
+
+// Limiter is used to apply rate limits.
+type Limiter interface {
+	// Wait must block until limiter
+	// permits another request to proceed.
+	Wait(context.Context) error
+}
+
+// NewHTTP creates a new HTTP client and returns it.
+func NewHTTP(opts ...Option) *HTTP {
+	options := Options{
+		HTTPClient: &http.Client{},
+	}
+	for _, apply := range opts {
+		apply(&options)
+	}
+
+	return &HTTP{
+		c:  options.HTTPClient,
+		rl: options.Limiter,
+	}
+}
+
+// Do dispatches the HTTP request to the network
+func (c *HTTP) Do(req *http.Request) (*http.Response, error) {
+	err := c.rl.Wait(req.Context()) // This is a blocking call. Honors the rate limit
+	if err != nil {
+		return nil, err
+	}
+	resp, err := c.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	return resp, nil
+}
+
+// WithHTTPClient sets the HTTP client.
+func WithHTTPClient(c *http.Client) Option {
+	return func(o *Options) {
+		o.HTTPClient = c
+	}
+}
+
+// WithLimiter sets the http rate limiter.
+func WithLimiter(l Limiter) Option {
+	return func(o *Options) {
+		o.Limiter = l
+	}
+}
diff --git a/cohere/client.go b/cohere/client.go
@@ -1,10 +1,10 @@
 package cohere
 
 import (
-	"net/http"
 	"os"
 
 	"github.com/milosgajdos/go-embeddings"
+	"github.com/milosgajdos/go-embeddings/client"
 )
 
 const (
@@ -24,7 +24,7 @@ type Options struct {
 	APIKey     string
 	BaseURL    string
 	Version    string
-	HTTPClient *http.Client
+	HTTPClient *client.HTTP
 }
 
 // Option is functional graph option.
@@ -39,7 +39,7 @@ func NewClient(opts ...Option) *Client {
 		APIKey:     os.Getenv("COHERE_API_KEY"),
 		BaseURL:    BaseURL,
 		Version:    EmbedAPIVersion,
-		HTTPClient: &http.Client{},
+		HTTPClient: client.NewHTTP(),
 	}
 
 	for _, apply := range opts {
@@ -78,7 +78,7 @@ func WithVersion(version string) Option {
 }
 
 // WithHTTPClient sets the HTTP client.
-func WithHTTPClient(httpClient *http.Client) Option {
+func WithHTTPClient(httpClient *client.HTTP) Option {
 	return func(o *Options) {
 		o.HTTPClient = httpClient
 	}

diff --git a/cohere/client_test.go b/cohere/client_test.go
@@ -1,9 +1,9 @@
 package cohere
 
 import (
-	"net/http"
 	"testing"
 
+	"github.com/milosgajdos/go-embeddings/client"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -45,7 +45,7 @@ func TestClient(t *testing.T) {
 		c := NewClient()
 		assert.NotNil(t, c.opts.HTTPClient)
 
-		testVal := &http.Client{}
+		testVal := client.NewHTTP()
 		c = NewClient(WithHTTPClient(testVal))
 		assert.NotNil(t, c.opts.HTTPClient)
 	})

diff --git a/openai/client.go b/openai/client.go
@@ -1,10 +1,10 @@
 package openai
 
 import (
-	"net/http"
 	"os"
 
 	"github.com/milosgajdos/go-embeddings"
+	"github.com/milosgajdos/go-embeddings/client"
 )
 
 const (
@@ -26,7 +26,7 @@ type Options struct {
 	BaseURL    string
 	Version    string
 	OrgID      string
-	HTTPClient *http.Client
+	HTTPClient *client.HTTP
 }
 
 // Option is functional graph option.
@@ -41,7 +41,7 @@ func NewClient(opts ...Option) *Client {
 		APIKey:     os.Getenv("OPENAI_API_KEY"),
 		BaseURL:    BaseURL,
 		Version:    EmbedAPIVersion,
-		HTTPClient: &http.Client{},
+		HTTPClient: client.NewHTTP(),
 	}
 
 	for _, apply := range opts {
@@ -87,7 +87,7 @@ func WithOrgID(orgID string) Option {
 }
 
 // WithHTTPClient sets the HTTP client.
-func WithHTTPClient(httpClient *http.Client) Option {
+func WithHTTPClient(httpClient *client.HTTP) Option {
 	return func(o *Options) {
 		o.HTTPClient = httpClient
 	}

diff --git a/openai/client_test.go b/openai/client_test.go
@@ -1,9 +1,9 @@
 package openai
 
 import (
-	"net/http"
 	"testing"
 
+	"github.com/milosgajdos/go-embeddings/client"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -54,7 +54,7 @@ func TestClient(t *testing.T) {
 		c := NewClient()
 		assert.NotNil(t, c.opts.HTTPClient)
 
-		testVal := &http.Client{}
+		testVal := client.NewHTTP()
 		c = NewClient(WithHTTPClient(testVal))
 		assert.NotNil(t, c.opts.HTTPClient)
 	})

diff --git a/request/request.go b/request/request.go
@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+
+	"github.com/milosgajdos/go-embeddings/client"
 )
 
 // NewHTTP creates a new HTTP request from the provided parameters  and returns it.
@@ -42,7 +44,7 @@ func NewHTTP(ctx context.Context, method, url string, body io.Reader, opts ...Op
 }
 
 // Do sends the HTTP request req using the client and returns the response.
-func Do[T error](client *http.Client, req *http.Request) (*http.Response, error) {
+func Do[T error](client *client.HTTP, req *http.Request) (*http.Response, error) {
 	resp, err := client.Do(req)
 	if err != nil {
 		return nil, err

diff --git a/vertexai/client.go b/vertexai/client.go
@@ -1,10 +1,10 @@
 package vertexai
 
 import (
-	"net/http"
 	"os"
 
 	"github.com/milosgajdos/go-embeddings"
+	"github.com/milosgajdos/go-embeddings/client"
 	"golang.org/x/oauth2"
 )
 
@@ -29,7 +29,7 @@ type Options struct {
 	ProjectID  string
 	ModelID    string
 	BaseURL    string
-	HTTPClient *http.Client
+	HTTPClient *client.HTTP
 }
 
 // Option is functional graph option.
@@ -50,7 +50,7 @@ func NewClient(opts ...Option) *Client {
 		ModelID:    os.Getenv("VERTEXAI_MODEL_ID"),
 		ProjectID:  os.Getenv("GOOGLE_PROJECT_ID"),
 		BaseURL:    BaseURL,
-		HTTPClient: &http.Client{},
+		HTTPClient: client.NewHTTP(),
 	}
 
 	for _, apply := range opts {
@@ -106,7 +106,7 @@ func WithBaseURL(baseURL string) Option {
 }
 
 // WithHTTPClient sets the HTTP client.
-func WithHTTPClient(httpClient *http.Client) Option {
+func WithHTTPClient(httpClient *client.HTTP) Option {
 	return func(o *Options) {
 		o.HTTPClient = httpClient
 	}

diff --git a/vertexai/client_test.go b/vertexai/client_test.go
@@ -1,9 +1,9 @@
 package vertexai
 
 import (
-	"net/http"
 	"testing"
 
+	"github.com/milosgajdos/go-embeddings/client"
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/oauth2"
 )
@@ -76,7 +76,7 @@ func TestClient(t *testing.T) {
 		c := NewClient()
 		assert.NotNil(t, c.opts.HTTPClient)
 
-		testVal := &http.Client{}
+		testVal := client.NewHTTP()
 		c = NewClient(WithHTTPClient(testVal))
 		assert.NotNil(t, c.opts.HTTPClient)
 	})