From 4a47d5d1eccf8807ed03540501826e143ae39bf2 Mon Sep 17 00:00:00 2001
From: liushuang <liushuangbill@outlook.com>
Date: Tue, 4 Jun 2024 15:40:44 +0800
Subject: [PATCH] feat: support rate limits headers

---
 client.go            |  4 ++++
 message_test.go      | 55 ++++++++++++++++++++++++++++++++++++++++++++
 ratelimit_headers.go | 46 ++++++++++++++++++++++++++++++++++++
 3 files changed, 105 insertions(+)
 create mode 100644 ratelimit_headers.go

diff --git a/client.go b/client.go
index 3ea47b3..f156998 100644
--- a/client.go
+++ b/client.go
@@ -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{
diff --git a/message_test.go b/message_test.go
index 434bdf5..d1a9dec 100644
--- a/message_test.go
+++ b/message_test.go
@@ -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)
@@ -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
@@ -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)
 }
 
diff --git a/ratelimit_headers.go b/ratelimit_headers.go
new file mode 100644
index 0000000..181132c
--- /dev/null
+++ b/ratelimit_headers.go
@@ -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,
+	}
+}