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:
@@ -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