Support get http header and x-ratelimit-* headers (#507)

* feat: add headers to http response

* feat: support rate limit headers

* fix: go lint

* fix: test coverage

* refactor streamReader

* refactor streamReader

* refactor: NewRateLimitHeaders to newRateLimitHeaders

* refactor: RateLimitHeaders Resets filed

* refactor: move RateLimitHeaders struct
This commit is contained in:
Liu Shuang
2023-10-10 10:29:41 -05:00
committed by GitHub
parent 8e165dc9aa
commit b77d01edca
5 changed files with 191 additions and 5 deletions

43
ratelimit.go Normal file
View File

@@ -0,0 +1,43 @@
package openai
import (
"net/http"
"strconv"
"time"
)
// RateLimitHeaders struct represents Openai rate limits headers.
type RateLimitHeaders struct {
LimitRequests int `json:"x-ratelimit-limit-requests"`
LimitTokens int `json:"x-ratelimit-limit-tokens"`
RemainingRequests int `json:"x-ratelimit-remaining-requests"`
RemainingTokens int `json:"x-ratelimit-remaining-tokens"`
ResetRequests ResetTime `json:"x-ratelimit-reset-requests"`
ResetTokens ResetTime `json:"x-ratelimit-reset-tokens"`
}
type ResetTime string
func (r ResetTime) String() string {
return string(r)
}
func (r ResetTime) Time() time.Time {
d, _ := time.ParseDuration(string(r))
return time.Now().Add(d)
}
func newRateLimitHeaders(h http.Header) RateLimitHeaders {
limitReq, _ := strconv.Atoi(h.Get("x-ratelimit-limit-requests"))
limitTokens, _ := strconv.Atoi(h.Get("x-ratelimit-limit-tokens"))
remainingReq, _ := strconv.Atoi(h.Get("x-ratelimit-remaining-requests"))
remainingTokens, _ := strconv.Atoi(h.Get("x-ratelimit-remaining-tokens"))
return RateLimitHeaders{
LimitRequests: limitReq,
LimitTokens: limitTokens,
RemainingRequests: remainingReq,
RemainingTokens: remainingTokens,
ResetRequests: ResetTime(h.Get("x-ratelimit-reset-requests")),
ResetTokens: ResetTime(h.Get("x-ratelimit-reset-tokens")),
}
}