Skip to content

Commit

Permalink
feat: support rate limits headers
Browse files Browse the repository at this point in the history
  • Loading branch information
liushuangls committed Jun 4, 2024
1 parent d291194 commit 4a47d5d
Show file tree
Hide file tree
Showing 3 changed files with 105 additions and 0 deletions.
4 changes: 4 additions & 0 deletions client.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,10 @@ func (h *httpHeader) Header() http.Header {
return http.Header(*h)
}

func (h *httpHeader) GetRateLimitHeaders() RateLimitHeaders {
return newRateLimitHeaders(h.Header())
}

// NewClient create new Anthropic API client
func NewClient(apikey string, opts ...ClientOption) *Client {
return &Client{
Expand Down
55 changes: 55 additions & 0 deletions message_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,16 @@ import (
//go:embed internal/test/sources/*
var sources embed.FS

var rateLimitHeaders = map[string]string{
"anthropic-ratelimit-requests-limit": "100",
"anthropic-ratelimit-requests-remaining": "99",
"anthropic-ratelimit-requests-reset": "2024-06-04T07:13:19Z",
"anthropic-ratelimit-tokens-limit": "10000",
"anthropic-ratelimit-tokens-remaining": "9900",
"anthropic-ratelimit-tokens-reset": "2024-06-04T07:13:19Z",
"retry-after": "100",
}

func TestMessages(t *testing.T) {
server := test.NewTestServer()
server.RegisterHandler("/v1/messages", handleMessagesEndpoint)
Expand Down Expand Up @@ -220,6 +230,48 @@ func TestMessagesToolUse(t *testing.T) {
}
}

func TestMessagesRateLimitHeaders(t *testing.T) {
server := test.NewTestServer()
server.RegisterHandler("/v1/messages", handleMessagesEndpoint)

ts := server.AnthropicTestServer()
ts.Start()
defer ts.Close()

baseUrl := ts.URL + "/v1"
client := anthropic.NewClient(
test.GetTestToken(),
anthropic.WithBaseURL(baseUrl),
)

resp, err := client.CreateMessages(context.Background(), anthropic.MessagesRequest{
Model: anthropic.ModelClaude3Haiku20240307,
Messages: []anthropic.Message{
anthropic.NewUserTextMessage("What is your name?"),
},
MaxTokens: 1000,
})
if err != nil {
t.Fatalf("CreateMessages error: %v", err)
}

rateLimitHeaders := resp.GetRateLimitHeaders()

bs, err := json.Marshal(rateLimitHeaders)
if err != nil {
t.Fatal(err)
}

bs2, err := json.Marshal(rateLimitHeaders)
if err != nil {
t.Fatal(err)
}

if string(bs) != string(bs2) {
t.Fatalf("rate limit headers mismatch. got %s, want %s", string(bs), string(bs2))
}
}

func handleMessagesEndpoint(w http.ResponseWriter, r *http.Request) {
var err error
var resBytes []byte
Expand Down Expand Up @@ -281,6 +333,9 @@ func handleMessagesEndpoint(w http.ResponseWriter, r *http.Request) {
}

resBytes, _ = json.Marshal(res)
for k, v := range rateLimitHeaders {
w.Header().Set(k, v)
}
_, _ = w.Write(resBytes)
}

Expand Down
46 changes: 46 additions & 0 deletions ratelimit_headers.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
package anthropic

import (
"net/http"
"strconv"
"time"
)

type RateLimitHeaders struct {
// The maximum number of requests allowed within the rate limit window.
RequestsLimit int `json:"anthropic-ratelimit-requests-limit"`
// The number of requests remaining within the current rate limit window.
RequestsRemaining int `json:"anthropic-ratelimit-requests-remaining"`
// The time when the request rate limit window will reset, provided in RFC 3339 format.
RequestsReset time.Time `json:"anthropic-ratelimit-requests-reset"`
// The maximum number of tokens allowed within the rate limit window.
TokensLimit int `json:"anthropic-ratelimit-tokens-limit"`
// The number of tokens remaining, rounded to the nearest thousand, within the current rate limit window.
TokensRemaining int `json:"anthropic-ratelimit-tokens-remaining"`
// The time when the token rate limit window will reset, provided in RFC 3339 format.
TokensReset time.Time `json:"anthropic-ratelimit-tokens-reset"`
// The number of seconds until the rate limit window resets.
RetryAfter int `json:"retry-after"`
}

func newRateLimitHeaders(h http.Header) RateLimitHeaders {
requestsLimit, _ := strconv.Atoi(h.Get("anthropic-ratelimit-requests-limit"))
requestsRemaining, _ := strconv.Atoi(h.Get("anthropic-ratelimit-requests-remaining"))
requestsReset, _ := time.Parse(time.RFC3339, h.Get("anthropic-ratelimit-requests-reset"))

tokensLimit, _ := strconv.Atoi(h.Get("anthropic-ratelimit-tokens-limit"))
tokensRemaining, _ := strconv.Atoi(h.Get("anthropic-ratelimit-tokens-remaining"))
tokensReset, _ := time.Parse(time.RFC3339, h.Get("anthropic-ratelimit-tokens-reset"))

retryAfter, _ := strconv.Atoi(h.Get("anthropic-ratelimit-retry-after"))

return RateLimitHeaders{
RequestsLimit: requestsLimit,
RequestsRemaining: requestsRemaining,
RequestsReset: requestsReset,
TokensLimit: tokensLimit,
TokensRemaining: tokensRemaining,
TokensReset: tokensReset,
RetryAfter: retryAfter,
}
}

0 comments on commit 4a47d5d

Please sign in to comment.