feat(llm): add configurable request/token rate limiting for LLM client
Introduce process-local and Redis-backed global rate limiting for the shared LLM client, plus a dedicated generation consumer prefetch knob. - Add request (RPM) and token (TPM) rate limits with burst controls, process or global scope, and fail-closed/process-degraded failover - Wire the LLM client through NewWithRedis so limits can be shared across replicas via a Redis token bucket - Add generation_consumer_prefetch to tune RabbitMQ generation consumers independently of the default prefetch - Bump generation worker_concurrency to 24 and promote golang.org/x/time to a direct dependency - Drop t.Parallel() from a few middleware/bootstrap tests Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -7,6 +7,8 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/config"
|
||||
)
|
||||
|
||||
@@ -53,15 +55,25 @@ type Client interface {
|
||||
}
|
||||
|
||||
func New(cfg config.LLMConfig) Client {
|
||||
return newClient(cfg, nil)
|
||||
}
|
||||
|
||||
func NewWithRedis(cfg config.LLMConfig, redisClient goredis.Cmdable) Client {
|
||||
return newClient(cfg, redisClient)
|
||||
}
|
||||
|
||||
func newClient(cfg config.LLMConfig, redisClient goredis.Cmdable) Client {
|
||||
provider := strings.ToLower(strings.TrimSpace(cfg.Provider))
|
||||
var client Client
|
||||
switch provider {
|
||||
case "", "disabled":
|
||||
return disabledClient{reason: "llm provider is disabled"}
|
||||
client = disabledClient{reason: "llm provider is disabled"}
|
||||
case "ark":
|
||||
return NewArkClient(cfg)
|
||||
client = NewArkClient(cfg)
|
||||
default:
|
||||
return disabledClient{reason: fmt.Sprintf("unsupported llm provider %q", cfg.Provider)}
|
||||
client = disabledClient{reason: fmt.Sprintf("unsupported llm provider %q", cfg.Provider)}
|
||||
}
|
||||
return NewRateLimitedClientWithRedis(client, cfg, redisClient)
|
||||
}
|
||||
|
||||
type disabledClient struct {
|
||||
|
||||
@@ -0,0 +1,472 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
"golang.org/x/time/rate"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/config"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultRateLimitMaxOutputTokens int64 = 4000
|
||||
redisRateLimitFailoverCooldown = 5 * time.Second
|
||||
)
|
||||
|
||||
var redisTokenBucketScript = goredis.NewScript(`
|
||||
local tokens_key = KEYS[1]
|
||||
local timestamp_key = KEYS[2]
|
||||
local cost = tonumber(ARGV[1])
|
||||
local rate = tonumber(ARGV[2])
|
||||
local capacity = tonumber(ARGV[3])
|
||||
local ttl = tonumber(ARGV[4])
|
||||
|
||||
if cost <= 0 or rate <= 0 or capacity <= 0 then
|
||||
return 0
|
||||
end
|
||||
|
||||
local current_time = redis.call("TIME")
|
||||
local now = tonumber(current_time[1]) * 1000 + math.floor(tonumber(current_time[2]) / 1000)
|
||||
local tokens = tonumber(redis.call("GET", tokens_key))
|
||||
if tokens == nil then
|
||||
tokens = capacity
|
||||
end
|
||||
local timestamp = tonumber(redis.call("GET", timestamp_key))
|
||||
if timestamp == nil then
|
||||
timestamp = now
|
||||
end
|
||||
if now < timestamp then
|
||||
timestamp = now
|
||||
end
|
||||
|
||||
local elapsed = now - timestamp
|
||||
if elapsed > 0 then
|
||||
tokens = math.min(capacity, tokens + elapsed * rate / 1000)
|
||||
end
|
||||
|
||||
if tokens >= cost then
|
||||
tokens = tokens - cost
|
||||
redis.call("SET", tokens_key, tokens, "PX", ttl)
|
||||
redis.call("SET", timestamp_key, now, "PX", ttl)
|
||||
return 0
|
||||
end
|
||||
|
||||
local wait = math.ceil((cost - tokens) / rate * 1000)
|
||||
redis.call("SET", tokens_key, tokens, "PX", ttl)
|
||||
redis.call("SET", timestamp_key, now, "PX", ttl)
|
||||
return wait
|
||||
`)
|
||||
|
||||
type rateLimitedClient struct {
|
||||
next Client
|
||||
requestLimiter *rate.Limiter
|
||||
tokenLimiter *rate.Limiter
|
||||
redisLimiter *redisRateLimiter
|
||||
redisFailoverRequests *rate.Limiter
|
||||
redisFailoverTokens *rate.Limiter
|
||||
redisUnavailableUntil int64
|
||||
rateLimitKey string
|
||||
requestRateLimitRPM int
|
||||
requestRateLimitBurst int
|
||||
tokenRateLimitTPM int64
|
||||
tokenRateLimitBurst int
|
||||
maxOutputTokens int64
|
||||
estimatedCharsPerToken float64
|
||||
}
|
||||
|
||||
func NewRateLimitedClient(next Client, cfg config.LLMConfig) Client {
|
||||
return NewRateLimitedClientWithRedis(next, cfg, nil)
|
||||
}
|
||||
|
||||
func NewRateLimitedClientWithRedis(next Client, cfg config.LLMConfig, redisClient goredis.Cmdable) Client {
|
||||
if next == nil {
|
||||
next = disabledClient{reason: "llm client is nil"}
|
||||
}
|
||||
|
||||
config.NormalizeLLMConfig(&cfg)
|
||||
if cfg.RequestRateLimitRPM <= 0 && cfg.TokenRateLimitTPM <= 0 {
|
||||
return next
|
||||
}
|
||||
|
||||
wrapped := &rateLimitedClient{
|
||||
next: next,
|
||||
rateLimitKey: llmRateLimitKey(cfg),
|
||||
requestRateLimitRPM: cfg.RequestRateLimitRPM,
|
||||
requestRateLimitBurst: cfg.RequestRateLimitBurst,
|
||||
tokenRateLimitTPM: cfg.TokenRateLimitTPM,
|
||||
tokenRateLimitBurst: cfg.TokenRateLimitBurst,
|
||||
maxOutputTokens: cfg.MaxOutputTokens,
|
||||
estimatedCharsPerToken: cfg.EstimatedCharsPerToken,
|
||||
}
|
||||
if wrapped.maxOutputTokens <= 0 {
|
||||
wrapped.maxOutputTokens = defaultRateLimitMaxOutputTokens
|
||||
}
|
||||
|
||||
if cfg.RateLimitScope == "global" && redisClient != nil {
|
||||
wrapped.redisLimiter = &redisRateLimiter{redis: redisClient}
|
||||
wrapped.configureRedisFailover(cfg)
|
||||
return wrapped
|
||||
}
|
||||
|
||||
if wrapped.requestRateLimitRPM > 0 {
|
||||
wrapped.requestLimiter = rate.NewLimiter(
|
||||
rate.Limit(float64(wrapped.requestRateLimitRPM)/60.0),
|
||||
rateLimitBurst(wrapped.requestRateLimitRPM, wrapped.requestRateLimitBurst),
|
||||
)
|
||||
}
|
||||
if wrapped.tokenRateLimitTPM > 0 {
|
||||
wrapped.tokenLimiter = rate.NewLimiter(
|
||||
rate.Limit(float64(wrapped.tokenRateLimitTPM)/60.0),
|
||||
rateLimitInt64Burst(wrapped.tokenRateLimitTPM, wrapped.tokenRateLimitBurst),
|
||||
)
|
||||
}
|
||||
|
||||
return wrapped
|
||||
}
|
||||
|
||||
func (c *rateLimitedClient) configureRedisFailover(cfg config.LLMConfig) {
|
||||
if c == nil || cfg.RateLimitFailoverMode != "process_degraded" {
|
||||
return
|
||||
}
|
||||
replicas := cfg.RateLimitFailoverReplicas
|
||||
if replicas <= 0 {
|
||||
replicas = 1
|
||||
}
|
||||
if c.requestRateLimitRPM > 0 {
|
||||
limit := divideLimit(c.requestRateLimitRPM, replicas)
|
||||
burst := divideLimit(rateLimitBurst(c.requestRateLimitRPM, c.requestRateLimitBurst), replicas)
|
||||
c.redisFailoverRequests = rate.NewLimiter(rate.Limit(float64(limit)/60.0), burst)
|
||||
}
|
||||
if c.tokenRateLimitTPM > 0 {
|
||||
limit := divideLimit64(c.tokenRateLimitTPM, replicas)
|
||||
burst := divideLimit(rateLimitInt64Burst(c.tokenRateLimitTPM, c.tokenRateLimitBurst), replicas)
|
||||
c.redisFailoverTokens = rate.NewLimiter(rate.Limit(float64(limit)/60.0), burst)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *rateLimitedClient) Validate() error {
|
||||
if c == nil || c.next == nil {
|
||||
return disabledClient{reason: "llm client is nil"}.Validate()
|
||||
}
|
||||
return c.next.Validate()
|
||||
}
|
||||
|
||||
func (c *rateLimitedClient) Generate(ctx context.Context, req GenerateRequest, onDelta func(string)) (*GenerateResult, error) {
|
||||
if err := c.Validate(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tokens := estimateRateLimitTokens(req, c.maxOutputTokens, c.estimatedCharsPerToken)
|
||||
if c.redisLimiter != nil {
|
||||
if err := c.waitGlobalRateLimit(ctx, "requests", 1, int64(c.requestRateLimitRPM), c.requestRateLimitBurst); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := c.waitGlobalRateLimit(ctx, "tokens", tokens, c.tokenRateLimitTPM, c.tokenRateLimitBurst); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return c.next.Generate(ctx, req, onDelta)
|
||||
}
|
||||
if c.requestLimiter != nil {
|
||||
if err := c.requestLimiter.Wait(ctx); err != nil {
|
||||
return nil, fmt.Errorf("llm request rate limit wait: %w", err)
|
||||
}
|
||||
}
|
||||
if c.tokenLimiter != nil {
|
||||
if err := waitLimiterN(ctx, c.tokenLimiter, tokens); err != nil {
|
||||
return nil, fmt.Errorf("llm token rate limit wait: %w", err)
|
||||
}
|
||||
}
|
||||
return c.next.Generate(ctx, req, onDelta)
|
||||
}
|
||||
|
||||
func (c *rateLimitedClient) waitGlobalRateLimit(ctx context.Context, bucket string, cost int, limitPerMinute int64, burst int) error {
|
||||
if c == nil || c.redisLimiter == nil || cost <= 0 || limitPerMinute <= 0 {
|
||||
return nil
|
||||
}
|
||||
if c.redisCircuitOpen(time.Now()) {
|
||||
return c.waitRedisFailoverRateLimit(ctx, bucket, cost)
|
||||
}
|
||||
if err := c.waitRedisRateLimit(ctx, bucket, cost, limitPerMinute, burst); err != nil {
|
||||
if !c.canUseRedisFailover(err) {
|
||||
return err
|
||||
}
|
||||
c.markRedisUnavailable(time.Now())
|
||||
if fallbackErr := c.waitRedisFailoverRateLimit(ctx, bucket, cost); fallbackErr != nil {
|
||||
return fmt.Errorf("llm %s global rate limit failed and process-degraded failover also failed: %w", bucket, errors.Join(err, fallbackErr))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *rateLimitedClient) canUseRedisFailover(err error) bool {
|
||||
if c == nil || err == nil {
|
||||
return false
|
||||
}
|
||||
if c.redisFailoverRequests == nil && c.redisFailoverTokens == nil {
|
||||
return false
|
||||
}
|
||||
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (c *rateLimitedClient) waitRedisFailoverRateLimit(ctx context.Context, bucket string, cost int) error {
|
||||
switch bucket {
|
||||
case "requests":
|
||||
if c.redisFailoverRequests == nil {
|
||||
return nil
|
||||
}
|
||||
if err := c.redisFailoverRequests.Wait(ctx); err != nil {
|
||||
return fmt.Errorf("llm %s process-degraded rate limit wait: %w", bucket, err)
|
||||
}
|
||||
case "tokens":
|
||||
if c.redisFailoverTokens == nil {
|
||||
return nil
|
||||
}
|
||||
if err := waitLimiterN(ctx, c.redisFailoverTokens, cost); err != nil {
|
||||
return fmt.Errorf("llm %s process-degraded rate limit wait: %w", bucket, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *rateLimitedClient) redisCircuitOpen(now time.Time) bool {
|
||||
if c == nil {
|
||||
return false
|
||||
}
|
||||
until := atomic.LoadInt64(&c.redisUnavailableUntil)
|
||||
return until > now.UnixNano()
|
||||
}
|
||||
|
||||
func (c *rateLimitedClient) markRedisUnavailable(now time.Time) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
atomic.StoreInt64(&c.redisUnavailableUntil, now.Add(redisRateLimitFailoverCooldown).UnixNano())
|
||||
}
|
||||
|
||||
func (c *rateLimitedClient) waitRedisRateLimit(ctx context.Context, bucket string, cost int, limitPerMinute int64, burst int) error {
|
||||
if c == nil || c.redisLimiter == nil || cost <= 0 || limitPerMinute <= 0 {
|
||||
return nil
|
||||
}
|
||||
capacity := rateLimitInt64Burst(limitPerMinute, burst)
|
||||
if capacity <= 0 {
|
||||
return nil
|
||||
}
|
||||
remaining := cost
|
||||
for remaining > 0 {
|
||||
chunk := remaining
|
||||
if chunk > capacity {
|
||||
chunk = capacity
|
||||
}
|
||||
if err := c.redisLimiter.Wait(ctx, c.rateLimitKey, bucket, chunk, limitPerMinute, capacity); err != nil {
|
||||
return fmt.Errorf("llm %s global rate limit wait: %w", bucket, err)
|
||||
}
|
||||
remaining -= chunk
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func estimateRateLimitTokens(req GenerateRequest, fallbackMaxOutputTokens int64, charsPerToken float64) int {
|
||||
if charsPerToken <= 0 {
|
||||
charsPerToken = 1
|
||||
}
|
||||
promptChars := utf8.RuneCountInString(strings.TrimSpace(req.Prompt))
|
||||
promptTokens := int64(math.Ceil(float64(promptChars) / charsPerToken))
|
||||
outputTokens := req.MaxOutputTokens
|
||||
if outputTokens <= 0 {
|
||||
outputTokens = fallbackMaxOutputTokens
|
||||
}
|
||||
if outputTokens <= 0 {
|
||||
outputTokens = defaultRateLimitMaxOutputTokens
|
||||
}
|
||||
|
||||
total := promptTokens + outputTokens
|
||||
if total <= 0 {
|
||||
return 1
|
||||
}
|
||||
if total > int64(maxInt()) {
|
||||
return maxInt()
|
||||
}
|
||||
return int(total)
|
||||
}
|
||||
|
||||
func waitLimiterN(ctx context.Context, limiter *rate.Limiter, n int) error {
|
||||
if limiter == nil || n <= 0 {
|
||||
return nil
|
||||
}
|
||||
burst := limiter.Burst()
|
||||
if burst <= 0 {
|
||||
return nil
|
||||
}
|
||||
for n > 0 {
|
||||
chunk := n
|
||||
if chunk > burst {
|
||||
chunk = burst
|
||||
}
|
||||
if err := limiter.WaitN(ctx, chunk); err != nil {
|
||||
return err
|
||||
}
|
||||
n -= chunk
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func rateLimitBurst(limit, configured int) int {
|
||||
if configured > 0 {
|
||||
return configured
|
||||
}
|
||||
burst := limit / 10
|
||||
if burst < 1 {
|
||||
return 1
|
||||
}
|
||||
return burst
|
||||
}
|
||||
|
||||
func rateLimitInt64Burst(limit int64, configured int) int {
|
||||
if configured > 0 {
|
||||
return configured
|
||||
}
|
||||
burst := limit / 10
|
||||
if burst < 1 {
|
||||
return 1
|
||||
}
|
||||
if burst > int64(maxInt()) {
|
||||
return maxInt()
|
||||
}
|
||||
return int(burst)
|
||||
}
|
||||
|
||||
func divideLimit(limit, divisor int) int {
|
||||
if limit <= 0 {
|
||||
return 0
|
||||
}
|
||||
if divisor <= 1 {
|
||||
return limit
|
||||
}
|
||||
quotient := limit / divisor
|
||||
if quotient < 1 {
|
||||
return 1
|
||||
}
|
||||
return quotient
|
||||
}
|
||||
|
||||
func divideLimit64(limit int64, divisor int) int64 {
|
||||
if limit <= 0 {
|
||||
return 0
|
||||
}
|
||||
if divisor <= 1 {
|
||||
return limit
|
||||
}
|
||||
quotient := limit / int64(divisor)
|
||||
if quotient < 1 {
|
||||
return 1
|
||||
}
|
||||
return quotient
|
||||
}
|
||||
|
||||
func maxInt() int {
|
||||
return int(^uint(0) >> 1)
|
||||
}
|
||||
|
||||
type redisRateLimiter struct {
|
||||
redis goredis.Cmdable
|
||||
}
|
||||
|
||||
func (l *redisRateLimiter) Wait(ctx context.Context, group, bucket string, cost int, limitPerMinute int64, capacity int) error {
|
||||
if l == nil || l.redis == nil || cost <= 0 || limitPerMinute <= 0 || capacity <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
ratePerSecond := float64(limitPerMinute) / 60.0
|
||||
if ratePerSecond <= 0 {
|
||||
return nil
|
||||
}
|
||||
ttl := redisBucketTTL(ratePerSecond, capacity)
|
||||
keys := redisBucketKeys(group, bucket)
|
||||
|
||||
for {
|
||||
waitMillis, err := redisTokenBucketScript.Run(
|
||||
ctx,
|
||||
l.redis,
|
||||
keys,
|
||||
cost,
|
||||
ratePerSecond,
|
||||
capacity,
|
||||
ttl.Milliseconds(),
|
||||
).Int64()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if waitMillis <= 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
timer := time.NewTimer(time.Duration(waitMillis) * time.Millisecond)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
timer.Stop()
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func redisBucketTTL(ratePerSecond float64, capacity int) time.Duration {
|
||||
if ratePerSecond <= 0 || capacity <= 0 {
|
||||
return 2 * time.Minute
|
||||
}
|
||||
refill := time.Duration(math.Ceil(float64(capacity)/ratePerSecond)) * time.Second
|
||||
if refill < time.Minute {
|
||||
return 2 * time.Minute
|
||||
}
|
||||
return refill * 2
|
||||
}
|
||||
|
||||
func redisBucketKeys(group, bucket string) []string {
|
||||
escapedGroup := sanitizeRateLimitKey(group)
|
||||
escapedBucket := sanitizeRateLimitKey(bucket)
|
||||
prefix := fmt.Sprintf("geo:llm:rate:{%s}:%s", escapedGroup, escapedBucket)
|
||||
return []string{prefix + ":tokens", prefix + ":ts"}
|
||||
}
|
||||
|
||||
func llmRateLimitKey(cfg config.LLMConfig) string {
|
||||
if key := strings.TrimSpace(cfg.RateLimitKey); key != "" {
|
||||
return key
|
||||
}
|
||||
parts := []string{
|
||||
strings.TrimSpace(cfg.Provider),
|
||||
strings.TrimSpace(cfg.BaseURL),
|
||||
strings.TrimSpace(cfg.Model),
|
||||
}
|
||||
key := strings.Join(parts, ":")
|
||||
if strings.Trim(key, ": ") == "" {
|
||||
return "default"
|
||||
}
|
||||
return key
|
||||
}
|
||||
|
||||
var rateLimitKeySanitizer = regexp.MustCompile(`[^a-zA-Z0-9_.:-]+`)
|
||||
|
||||
func sanitizeRateLimitKey(value string) string {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return "default"
|
||||
}
|
||||
value = rateLimitKeySanitizer.ReplaceAllString(value, "_")
|
||||
value = strings.Trim(value, "_")
|
||||
if value == "" {
|
||||
return "default"
|
||||
}
|
||||
return value
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/geo-platform/tenant-api/internal/shared/config"
|
||||
)
|
||||
|
||||
type stubLLMClient struct {
|
||||
req GenerateRequest
|
||||
calls int
|
||||
}
|
||||
|
||||
func (c *stubLLMClient) Validate() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *stubLLMClient) Generate(_ context.Context, req GenerateRequest, _ func(string)) (*GenerateResult, error) {
|
||||
c.req = req
|
||||
c.calls++
|
||||
return &GenerateResult{Content: "ok", Model: req.Model}, nil
|
||||
}
|
||||
|
||||
func TestNewRateLimitedClientReturnsBaseWhenLimitsDisabled(t *testing.T) {
|
||||
base := &stubLLMClient{}
|
||||
got := NewRateLimitedClient(base, config.LLMConfig{})
|
||||
if got != base {
|
||||
t.Fatalf("expected base client when rate limits are disabled")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimitedClientDelegatesGenerate(t *testing.T) {
|
||||
base := &stubLLMClient{}
|
||||
client := NewRateLimitedClient(base, config.LLMConfig{
|
||||
RequestRateLimitRPM: 60000,
|
||||
TokenRateLimitTPM: 600000,
|
||||
MaxOutputTokens: 100,
|
||||
EstimatedCharsPerToken: 1,
|
||||
})
|
||||
|
||||
result, err := client.Generate(context.Background(), GenerateRequest{
|
||||
Prompt: "hello",
|
||||
MaxOutputTokens: 10,
|
||||
}, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("Generate() error = %v", err)
|
||||
}
|
||||
if result.Content != "ok" {
|
||||
t.Fatalf("Generate() content = %q, want ok", result.Content)
|
||||
}
|
||||
if base.req.Prompt != "hello" {
|
||||
t.Fatalf("delegated prompt = %q, want hello", base.req.Prompt)
|
||||
}
|
||||
if base.calls != 1 {
|
||||
t.Fatalf("delegated calls = %d, want 1", base.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateRateLimitTokensIncludesFallbackOutput(t *testing.T) {
|
||||
got := estimateRateLimitTokens(GenerateRequest{Prompt: " 你好世界 "}, 16000, 1)
|
||||
if got != 16004 {
|
||||
t.Fatalf("estimateRateLimitTokens() = %d, want 16004", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateRateLimitTokensUsesCharsPerToken(t *testing.T) {
|
||||
got := estimateRateLimitTokens(GenerateRequest{
|
||||
Prompt: "abcd",
|
||||
MaxOutputTokens: 6,
|
||||
}, 16000, 2)
|
||||
if got != 8 {
|
||||
t.Fatalf("estimateRateLimitTokens() = %d, want 8", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimitedClientStopsBeforeDelegateWhenContextCanceled(t *testing.T) {
|
||||
base := &stubLLMClient{}
|
||||
client := NewRateLimitedClient(base, config.LLMConfig{
|
||||
RequestRateLimitRPM: 1,
|
||||
RequestRateLimitBurst: 1,
|
||||
TokenRateLimitTPM: 1,
|
||||
TokenRateLimitBurst: 1,
|
||||
MaxOutputTokens: 100,
|
||||
EstimatedCharsPerToken: 1,
|
||||
})
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
if _, err := client.Generate(ctx, GenerateRequest{Prompt: "hello"}, nil); err == nil {
|
||||
t.Fatalf("Generate() expected context cancellation error")
|
||||
}
|
||||
if base.calls != 0 {
|
||||
t.Fatalf("delegated calls = %d, want 0", base.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGlobalRateLimitedClientsShareRedisBudget(t *testing.T) {
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := goredis.NewClient(&goredis.Options{Addr: mr.Addr()})
|
||||
defer rdb.Close()
|
||||
|
||||
cfg := config.LLMConfig{
|
||||
RateLimitScope: "global",
|
||||
RateLimitKey: "ark-shared",
|
||||
RequestRateLimitRPM: 60,
|
||||
RequestRateLimitBurst: 1,
|
||||
}
|
||||
firstBase := &stubLLMClient{}
|
||||
first := NewRateLimitedClientWithRedis(firstBase, cfg, rdb)
|
||||
if _, err := first.Generate(context.Background(), GenerateRequest{Prompt: "first"}, nil); err != nil {
|
||||
t.Fatalf("first Generate() error = %v", err)
|
||||
}
|
||||
if firstBase.calls != 1 {
|
||||
t.Fatalf("first delegated calls = %d, want 1", firstBase.calls)
|
||||
}
|
||||
|
||||
secondBase := &stubLLMClient{}
|
||||
second := NewRateLimitedClientWithRedis(secondBase, cfg, rdb)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
if _, err := second.Generate(ctx, GenerateRequest{Prompt: "second"}, nil); err == nil {
|
||||
t.Fatalf("second Generate() expected global rate limit wait error")
|
||||
}
|
||||
if secondBase.calls != 0 {
|
||||
t.Fatalf("second delegated calls = %d, want 0", secondBase.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGlobalRateLimitedClientsShareRedisTokenBudget(t *testing.T) {
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := goredis.NewClient(&goredis.Options{Addr: mr.Addr()})
|
||||
defer rdb.Close()
|
||||
|
||||
cfg := config.LLMConfig{
|
||||
RateLimitScope: "global",
|
||||
RateLimitKey: "ark-token-shared",
|
||||
TokenRateLimitTPM: 60,
|
||||
TokenRateLimitBurst: 10,
|
||||
MaxOutputTokens: 5,
|
||||
EstimatedCharsPerToken: 1,
|
||||
}
|
||||
firstBase := &stubLLMClient{}
|
||||
first := NewRateLimitedClientWithRedis(firstBase, cfg, rdb)
|
||||
if _, err := first.Generate(context.Background(), GenerateRequest{
|
||||
Prompt: "12345",
|
||||
MaxOutputTokens: 5,
|
||||
}, nil); err != nil {
|
||||
t.Fatalf("first Generate() error = %v", err)
|
||||
}
|
||||
if firstBase.calls != 1 {
|
||||
t.Fatalf("first delegated calls = %d, want 1", firstBase.calls)
|
||||
}
|
||||
|
||||
secondBase := &stubLLMClient{}
|
||||
second := NewRateLimitedClientWithRedis(secondBase, cfg, rdb)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
if _, err := second.Generate(ctx, GenerateRequest{
|
||||
Prompt: "12345",
|
||||
MaxOutputTokens: 5,
|
||||
}, nil); err == nil {
|
||||
t.Fatalf("second Generate() expected global token rate limit wait error")
|
||||
}
|
||||
if secondBase.calls != 0 {
|
||||
t.Fatalf("second delegated calls = %d, want 0", secondBase.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGlobalRateLimitedClientFailsClosedWhenRedisUnavailable(t *testing.T) {
|
||||
mr := miniredis.RunT(t)
|
||||
addr := mr.Addr()
|
||||
mr.Close()
|
||||
|
||||
rdb := goredis.NewClient(&goredis.Options{
|
||||
Addr: addr,
|
||||
DialTimeout: 10 * time.Millisecond,
|
||||
ReadTimeout: 10 * time.Millisecond,
|
||||
WriteTimeout: 10 * time.Millisecond,
|
||||
})
|
||||
defer rdb.Close()
|
||||
|
||||
cfg := config.LLMConfig{
|
||||
RateLimitScope: "global",
|
||||
RateLimitKey: "ark-fail-closed",
|
||||
RequestRateLimitRPM: 60,
|
||||
RequestRateLimitBurst: 1,
|
||||
}
|
||||
base := &stubLLMClient{}
|
||||
client := NewRateLimitedClientWithRedis(base, cfg, rdb)
|
||||
|
||||
if _, err := client.Generate(context.Background(), GenerateRequest{Prompt: "blocked"}, nil); err == nil {
|
||||
t.Fatalf("Generate() expected redis rate limit error")
|
||||
}
|
||||
if base.calls != 0 {
|
||||
t.Fatalf("delegated calls = %d, want 0", base.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGlobalRateLimitedClientUsesProcessDegradedFailoverWhenRedisUnavailable(t *testing.T) {
|
||||
mr := miniredis.RunT(t)
|
||||
addr := mr.Addr()
|
||||
mr.Close()
|
||||
|
||||
rdb := goredis.NewClient(&goredis.Options{
|
||||
Addr: addr,
|
||||
DialTimeout: 10 * time.Millisecond,
|
||||
ReadTimeout: 10 * time.Millisecond,
|
||||
WriteTimeout: 10 * time.Millisecond,
|
||||
})
|
||||
defer rdb.Close()
|
||||
|
||||
cfg := config.LLMConfig{
|
||||
RateLimitScope: "global",
|
||||
RateLimitKey: "ark-process-degraded",
|
||||
RateLimitFailoverMode: "process_degraded",
|
||||
RateLimitFailoverReplicas: 3,
|
||||
RequestRateLimitRPM: 60,
|
||||
RequestRateLimitBurst: 3,
|
||||
TokenRateLimitTPM: 60,
|
||||
TokenRateLimitBurst: 3,
|
||||
MaxOutputTokens: 1,
|
||||
EstimatedCharsPerToken: 1,
|
||||
}
|
||||
base := &stubLLMClient{}
|
||||
client := NewRateLimitedClientWithRedis(base, cfg, rdb)
|
||||
|
||||
if _, err := client.Generate(context.Background(), GenerateRequest{
|
||||
Prompt: "",
|
||||
MaxOutputTokens: 1,
|
||||
}, nil); err != nil {
|
||||
t.Fatalf("first Generate() error = %v", err)
|
||||
}
|
||||
if base.calls != 1 {
|
||||
t.Fatalf("delegated calls = %d, want 1", base.calls)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
|
||||
defer cancel()
|
||||
if _, err := client.Generate(ctx, GenerateRequest{
|
||||
Prompt: "",
|
||||
MaxOutputTokens: 1,
|
||||
}, nil); err == nil {
|
||||
t.Fatalf("second Generate() expected process-degraded rate limit wait error")
|
||||
}
|
||||
if base.calls != 1 {
|
||||
t.Fatalf("delegated calls after blocked call = %d, want 1", base.calls)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user