Skip to content

Commit eacb758

Browse files
authored
feature: add client rate limit support (#20)
* 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]> --------- Signed-off-by: Milos Gajdos <[email protected]>
1 parent da75163 commit eacb758

File tree

8 files changed

+95
-19
lines changed

8 files changed

+95
-19
lines changed

client/client.go

+74
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
package client
2+
3+
import (
4+
"context"
5+
"net/http"
6+
)
7+
8+
// HTTP is a HTTP client.
9+
type HTTP struct {
10+
client *http.Client
11+
limiter Limiter
12+
}
13+
14+
// Options are client options
15+
type Options struct {
16+
HTTPClient *http.Client
17+
Limiter Limiter
18+
}
19+
20+
// Option is functional graph option.
21+
type Option func(*Options)
22+
23+
// Limiter is used to apply rate limits.
24+
// NOTE: you can use off the shelf limiter from
25+
// https://pkg.go.dev/golang.org/x/time/rate#Limiter
26+
type Limiter interface {
27+
// Wait must block until limiter
28+
// permits another request to proceed.
29+
Wait(context.Context) error
30+
}
31+
32+
// NewHTTP creates a new HTTP client and returns it.
33+
func NewHTTP(opts ...Option) *HTTP {
34+
options := Options{
35+
HTTPClient: &http.Client{},
36+
}
37+
for _, apply := range opts {
38+
apply(&options)
39+
}
40+
41+
return &HTTP{
42+
client: options.HTTPClient,
43+
limiter: options.Limiter,
44+
}
45+
}
46+
47+
// Do dispatches the HTTP request to the network
48+
func (h *HTTP) Do(req *http.Request) (*http.Response, error) {
49+
if h.limiter != nil {
50+
err := h.limiter.Wait(req.Context()) // This is a blocking call. Honors the rate limit
51+
if err != nil {
52+
return nil, err
53+
}
54+
}
55+
resp, err := h.client.Do(req)
56+
if err != nil {
57+
return nil, err
58+
}
59+
return resp, nil
60+
}
61+
62+
// WithHTTPClient sets the HTTP client.
63+
func WithHTTPClient(c *http.Client) Option {
64+
return func(o *Options) {
65+
o.HTTPClient = c
66+
}
67+
}
68+
69+
// WithLimiter sets the http rate limiter.
70+
func WithLimiter(l Limiter) Option {
71+
return func(o *Options) {
72+
o.Limiter = l
73+
}
74+
}

cohere/client.go

+4-4
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
package cohere
22

33
import (
4-
"net/http"
54
"os"
65

76
"github.com/milosgajdos/go-embeddings"
7+
"github.com/milosgajdos/go-embeddings/client"
88
)
99

1010
const (
@@ -24,7 +24,7 @@ type Options struct {
2424
APIKey string
2525
BaseURL string
2626
Version string
27-
HTTPClient *http.Client
27+
HTTPClient *client.HTTP
2828
}
2929

3030
// Option is functional graph option.
@@ -39,7 +39,7 @@ func NewClient(opts ...Option) *Client {
3939
APIKey: os.Getenv("COHERE_API_KEY"),
4040
BaseURL: BaseURL,
4141
Version: EmbedAPIVersion,
42-
HTTPClient: &http.Client{},
42+
HTTPClient: client.NewHTTP(),
4343
}
4444

4545
for _, apply := range opts {
@@ -78,7 +78,7 @@ func WithVersion(version string) Option {
7878
}
7979

8080
// WithHTTPClient sets the HTTP client.
81-
func WithHTTPClient(httpClient *http.Client) Option {
81+
func WithHTTPClient(httpClient *client.HTTP) Option {
8282
return func(o *Options) {
8383
o.HTTPClient = httpClient
8484
}

cohere/client_test.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
package cohere
22

33
import (
4-
"net/http"
54
"testing"
65

6+
"github.com/milosgajdos/go-embeddings/client"
77
"github.com/stretchr/testify/assert"
88
)
99

@@ -45,7 +45,7 @@ func TestClient(t *testing.T) {
4545
c := NewClient()
4646
assert.NotNil(t, c.opts.HTTPClient)
4747

48-
testVal := &http.Client{}
48+
testVal := client.NewHTTP()
4949
c = NewClient(WithHTTPClient(testVal))
5050
assert.NotNil(t, c.opts.HTTPClient)
5151
})

openai/client.go

+4-4
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
package openai
22

33
import (
4-
"net/http"
54
"os"
65

76
"github.com/milosgajdos/go-embeddings"
7+
"github.com/milosgajdos/go-embeddings/client"
88
)
99

1010
const (
@@ -26,7 +26,7 @@ type Options struct {
2626
BaseURL string
2727
Version string
2828
OrgID string
29-
HTTPClient *http.Client
29+
HTTPClient *client.HTTP
3030
}
3131

3232
// Option is functional graph option.
@@ -41,7 +41,7 @@ func NewClient(opts ...Option) *Client {
4141
APIKey: os.Getenv("OPENAI_API_KEY"),
4242
BaseURL: BaseURL,
4343
Version: EmbedAPIVersion,
44-
HTTPClient: &http.Client{},
44+
HTTPClient: client.NewHTTP(),
4545
}
4646

4747
for _, apply := range opts {
@@ -87,7 +87,7 @@ func WithOrgID(orgID string) Option {
8787
}
8888

8989
// WithHTTPClient sets the HTTP client.
90-
func WithHTTPClient(httpClient *http.Client) Option {
90+
func WithHTTPClient(httpClient *client.HTTP) Option {
9191
return func(o *Options) {
9292
o.HTTPClient = httpClient
9393
}

openai/client_test.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
package openai
22

33
import (
4-
"net/http"
54
"testing"
65

6+
"github.com/milosgajdos/go-embeddings/client"
77
"github.com/stretchr/testify/assert"
88
)
99

@@ -54,7 +54,7 @@ func TestClient(t *testing.T) {
5454
c := NewClient()
5555
assert.NotNil(t, c.opts.HTTPClient)
5656

57-
testVal := &http.Client{}
57+
testVal := client.NewHTTP()
5858
c = NewClient(WithHTTPClient(testVal))
5959
assert.NotNil(t, c.opts.HTTPClient)
6060
})

request/request.go

+3-1
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ import (
77
"fmt"
88
"io"
99
"net/http"
10+
11+
"github.com/milosgajdos/go-embeddings/client"
1012
)
1113

1214
// 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
4244
}
4345

4446
// Do sends the HTTP request req using the client and returns the response.
45-
func Do[T error](client *http.Client, req *http.Request) (*http.Response, error) {
47+
func Do[T error](client *client.HTTP, req *http.Request) (*http.Response, error) {
4648
resp, err := client.Do(req)
4749
if err != nil {
4850
return nil, err

vertexai/client.go

+4-4
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
package vertexai
22

33
import (
4-
"net/http"
54
"os"
65

76
"github.com/milosgajdos/go-embeddings"
7+
"github.com/milosgajdos/go-embeddings/client"
88
"golang.org/x/oauth2"
99
)
1010

@@ -29,7 +29,7 @@ type Options struct {
2929
ProjectID string
3030
ModelID string
3131
BaseURL string
32-
HTTPClient *http.Client
32+
HTTPClient *client.HTTP
3333
}
3434

3535
// Option is functional graph option.
@@ -50,7 +50,7 @@ func NewClient(opts ...Option) *Client {
5050
ModelID: os.Getenv("VERTEXAI_MODEL_ID"),
5151
ProjectID: os.Getenv("GOOGLE_PROJECT_ID"),
5252
BaseURL: BaseURL,
53-
HTTPClient: &http.Client{},
53+
HTTPClient: client.NewHTTP(),
5454
}
5555

5656
for _, apply := range opts {
@@ -106,7 +106,7 @@ func WithBaseURL(baseURL string) Option {
106106
}
107107

108108
// WithHTTPClient sets the HTTP client.
109-
func WithHTTPClient(httpClient *http.Client) Option {
109+
func WithHTTPClient(httpClient *client.HTTP) Option {
110110
return func(o *Options) {
111111
o.HTTPClient = httpClient
112112
}

vertexai/client_test.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
package vertexai
22

33
import (
4-
"net/http"
54
"testing"
65

6+
"github.com/milosgajdos/go-embeddings/client"
77
"github.com/stretchr/testify/assert"
88
"golang.org/x/oauth2"
99
)
@@ -76,7 +76,7 @@ func TestClient(t *testing.T) {
7676
c := NewClient()
7777
assert.NotNil(t, c.opts.HTTPClient)
7878

79-
testVal := &http.Client{}
79+
testVal := client.NewHTTP()
8080
c = NewClient(WithHTTPClient(testVal))
8181
assert.NotNil(t, c.opts.HTTPClient)
8282
})

0 commit comments

Comments
 (0)