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:
+12
-1
@@ -78,6 +78,7 @@ rabbitmq:
|
||||
compliance_review_dlq: compliance.review.run.dlq
|
||||
compliance_review_dlq_routing_key: compliance.review.run.dlq
|
||||
consumer_prefetch: 10
|
||||
generation_consumer_prefetch: 1
|
||||
publish_channel_pool_size: 32
|
||||
|
||||
scheduler:
|
||||
@@ -218,6 +219,16 @@ llm:
|
||||
temperature: 0.7
|
||||
reasoning_effort: minimal
|
||||
web_search_limit: 3
|
||||
# Global Redis-backed budget shared by all replicas.
|
||||
rate_limit_scope: global
|
||||
rate_limit_key: ark:doubao-seed-2-0
|
||||
rate_limit_failover_mode: process_degraded
|
||||
rate_limit_failover_replicas: 10
|
||||
request_rate_limit_rpm: 30000
|
||||
request_rate_limit_burst: 3000
|
||||
token_rate_limit_tpm: 5000000
|
||||
token_rate_limit_burst: 500000
|
||||
estimated_chars_per_token: 1
|
||||
|
||||
compliance:
|
||||
enabled: true
|
||||
@@ -246,7 +257,7 @@ retrieval:
|
||||
|
||||
generation:
|
||||
queue_size: 128
|
||||
worker_concurrency: 8
|
||||
worker_concurrency: 24
|
||||
stream_enabled: false
|
||||
article_timeout: 8m
|
||||
task_lease_ttl: 10m
|
||||
|
||||
@@ -206,7 +206,7 @@ metadata:
|
||||
labels:
|
||||
app.kubernetes.io/name: worker-generate
|
||||
spec:
|
||||
replicas: 1
|
||||
replicas: 3
|
||||
selector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/name: worker-generate
|
||||
|
||||
@@ -83,6 +83,7 @@ rabbitmq:
|
||||
compliance_review_dlq: compliance.review.run.dlq
|
||||
compliance_review_dlq_routing_key: compliance.review.run.dlq
|
||||
consumer_prefetch: 10
|
||||
generation_consumer_prefetch: 1
|
||||
publish_channel_pool_size: 32
|
||||
|
||||
scheduler:
|
||||
@@ -200,6 +201,16 @@ llm:
|
||||
temperature: 0.7
|
||||
reasoning_effort: minimal
|
||||
web_search_limit: 3
|
||||
# Global Redis-backed budget shared by all replicas.
|
||||
rate_limit_scope: global
|
||||
rate_limit_key: ark:doubao-seed-2-0
|
||||
rate_limit_failover_mode: process_degraded
|
||||
rate_limit_failover_replicas: 10
|
||||
request_rate_limit_rpm: 30000
|
||||
request_rate_limit_burst: 3000
|
||||
token_rate_limit_tpm: 5000000
|
||||
token_rate_limit_burst: 500000
|
||||
estimated_chars_per_token: 1
|
||||
|
||||
compliance:
|
||||
enabled: true
|
||||
@@ -228,7 +239,7 @@ retrieval:
|
||||
|
||||
generation:
|
||||
queue_size: 128
|
||||
worker_concurrency: 8
|
||||
worker_concurrency: 24
|
||||
stream_enabled: false
|
||||
article_timeout: 8m
|
||||
task_lease_ttl: 10m
|
||||
|
||||
@@ -70,6 +70,7 @@ rabbitmq:
|
||||
compliance_review_dlq: compliance.review.run.dlq
|
||||
compliance_review_dlq_routing_key: compliance.review.run.dlq
|
||||
consumer_prefetch: 10
|
||||
generation_consumer_prefetch: 1
|
||||
publish_channel_pool_size: 32
|
||||
# Prefer environment override when needed:
|
||||
# export RABBITMQ_URL=amqp://geo:geo_dev@localhost:5672/geo
|
||||
@@ -198,6 +199,16 @@ llm:
|
||||
temperature: 0.7
|
||||
reasoning_effort: minimal
|
||||
web_search_limit: 3
|
||||
# Global Redis-backed budget shared by all replicas.
|
||||
rate_limit_scope: global
|
||||
rate_limit_key: ark:doubao-seed-2-0
|
||||
rate_limit_failover_mode: process_degraded
|
||||
rate_limit_failover_replicas: 10
|
||||
request_rate_limit_rpm: 30000
|
||||
request_rate_limit_burst: 3000
|
||||
token_rate_limit_tpm: 5000000
|
||||
token_rate_limit_burst: 500000
|
||||
estimated_chars_per_token: 1
|
||||
# top_p: 0.95
|
||||
# Prefer environment override for secrets:
|
||||
# export LLM_API_KEY=your-api-key
|
||||
@@ -233,7 +244,7 @@ retrieval:
|
||||
|
||||
generation:
|
||||
queue_size: 128
|
||||
worker_concurrency: 8
|
||||
worker_concurrency: 24
|
||||
stream_enabled: false
|
||||
article_timeout: 8m
|
||||
|
||||
|
||||
@@ -83,6 +83,7 @@ rabbitmq:
|
||||
compliance_review_dlq: compliance.review.run.dlq
|
||||
compliance_review_dlq_routing_key: compliance.review.run.dlq
|
||||
consumer_prefetch: 10
|
||||
generation_consumer_prefetch: 1
|
||||
publish_channel_pool_size: 32
|
||||
|
||||
scheduler:
|
||||
@@ -204,6 +205,16 @@ llm:
|
||||
temperature: 0.7
|
||||
reasoning_effort: minimal
|
||||
web_search_limit: 3
|
||||
# Global Redis-backed budget shared by all replicas.
|
||||
rate_limit_scope: global
|
||||
rate_limit_key: ark:doubao-seed-2-0
|
||||
rate_limit_failover_mode: process_degraded
|
||||
rate_limit_failover_replicas: 10
|
||||
request_rate_limit_rpm: 30000
|
||||
request_rate_limit_burst: 3000
|
||||
token_rate_limit_tpm: 5000000
|
||||
token_rate_limit_burst: 500000
|
||||
estimated_chars_per_token: 1
|
||||
|
||||
compliance:
|
||||
enabled: true
|
||||
@@ -232,7 +243,7 @@ retrieval:
|
||||
|
||||
generation:
|
||||
queue_size: 128
|
||||
worker_concurrency: 8
|
||||
worker_concurrency: 24
|
||||
stream_enabled: false
|
||||
article_timeout: 8m
|
||||
task_lease_ttl: 10m
|
||||
|
||||
+1
-1
@@ -33,6 +33,7 @@ require (
|
||||
golang.org/x/image v0.38.0
|
||||
golang.org/x/net v0.48.0
|
||||
golang.org/x/sync v0.20.0
|
||||
golang.org/x/time v0.5.0
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
|
||||
@@ -83,7 +84,6 @@ require (
|
||||
golang.org/x/arch v0.3.0 // indirect
|
||||
golang.org/x/sys v0.39.0 // indirect
|
||||
golang.org/x/text v0.35.0 // indirect
|
||||
golang.org/x/time v0.5.0 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect
|
||||
google.golang.org/grpc v1.79.3 // indirect
|
||||
google.golang.org/protobuf v1.36.10 // indirect
|
||||
|
||||
@@ -135,7 +135,7 @@ func New(configPath string) (*App, error) {
|
||||
logger.Warn("tenant login guard disabled by config; brute-force / rate-limit protection is off")
|
||||
}
|
||||
loginGuard := auth.NewLoginGuard(rdb, loginGuardCfg)
|
||||
llmClient := llm.NewReloadableClient(llm.New(cfg.LLM))
|
||||
llmClient := llm.NewReloadableClient(llm.NewWithRedis(cfg.LLM, rdb))
|
||||
retrievalProvider := retrieval.NewReloadableProvider(retrieval.NewProvider(cfg.Retrieval, logger))
|
||||
vectorStore := retrieval.NewReloadableVectorStore(retrieval.NewQdrantStore(cfg.Qdrant, logger))
|
||||
objectStorageClient := objectstorage.NewReloadableClient(objectstorage.New(cfg.ObjectStorage, logger))
|
||||
@@ -294,7 +294,7 @@ func (a *App) applyReloadedConfig(event config.ReloadEvent) {
|
||||
|
||||
a.JWT.Update(event.Current.JWT.Secret, event.Current.JWT.AccessTTL, event.Current.JWT.RefreshTTL)
|
||||
if a.ReloadableLLM != nil {
|
||||
a.ReloadableLLM.Update(llm.New(event.Current.LLM))
|
||||
a.ReloadableLLM.Update(llm.NewWithRedis(event.Current.LLM, a.Redis))
|
||||
}
|
||||
if a.ReloadableRetrieval != nil {
|
||||
a.ReloadableRetrieval.Update(retrieval.NewProvider(event.Current.Retrieval, a.Logger))
|
||||
|
||||
@@ -12,8 +12,6 @@ import (
|
||||
)
|
||||
|
||||
func TestNewEngineUsesForwardedClientIPFromTrustedProxy(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
engine, err := NewEngine(config.ServerConfig{
|
||||
@@ -38,8 +36,6 @@ func TestNewEngineUsesForwardedClientIPFromTrustedProxy(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestNewEngineRejectsInvalidTrustedProxy(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
_, err := NewEngine(config.ServerConfig{
|
||||
|
||||
@@ -158,6 +158,7 @@ type RabbitMQConfig struct {
|
||||
ComplianceReviewDLQ string `mapstructure:"compliance_review_dlq"`
|
||||
ComplianceReviewDLQRouteKey string `mapstructure:"compliance_review_dlq_routing_key"`
|
||||
ConsumerPrefetch int `mapstructure:"consumer_prefetch"`
|
||||
GenerationConsumerPrefetch int `mapstructure:"generation_consumer_prefetch"`
|
||||
PublishChannelPoolSize int `mapstructure:"publish_channel_pool_size"`
|
||||
}
|
||||
|
||||
@@ -340,6 +341,15 @@ type LLMConfig struct {
|
||||
TopP float64 `mapstructure:"top_p"`
|
||||
ReasoningEffort string `mapstructure:"reasoning_effort"`
|
||||
WebSearchLimit int32 `mapstructure:"web_search_limit"`
|
||||
RequestRateLimitRPM int `mapstructure:"request_rate_limit_rpm"`
|
||||
RequestRateLimitBurst int `mapstructure:"request_rate_limit_burst"`
|
||||
TokenRateLimitTPM int64 `mapstructure:"token_rate_limit_tpm"`
|
||||
TokenRateLimitBurst int `mapstructure:"token_rate_limit_burst"`
|
||||
EstimatedCharsPerToken float64 `mapstructure:"estimated_chars_per_token"`
|
||||
RateLimitScope string `mapstructure:"rate_limit_scope"`
|
||||
RateLimitKey string `mapstructure:"rate_limit_key"`
|
||||
RateLimitFailoverMode string `mapstructure:"rate_limit_failover_mode"`
|
||||
RateLimitFailoverReplicas int `mapstructure:"rate_limit_failover_replicas"`
|
||||
}
|
||||
|
||||
type ComplianceConfig struct {
|
||||
@@ -480,6 +490,7 @@ func loadWithFiles(configPath string) (*Config, []string, error) {
|
||||
normalizeMonitoringDispatchConfig(&cfg.MonitoringDispatch)
|
||||
normalizeMembershipConfig(&cfg.Membership)
|
||||
normalizeBrandLibraryConfig(&cfg.BrandLibrary)
|
||||
NormalizeLLMConfig(&cfg.LLM)
|
||||
NormalizeComplianceConfig(&cfg.Compliance)
|
||||
NormalizeGenerationConfig(&cfg.Generation)
|
||||
NormalizeBrowserFetchConfig(&cfg.BrowserFetch)
|
||||
@@ -924,6 +935,47 @@ func NormalizeComplianceConfig(cfg *ComplianceConfig) {
|
||||
}
|
||||
}
|
||||
|
||||
func NormalizeLLMConfig(cfg *LLMConfig) {
|
||||
if cfg == nil {
|
||||
return
|
||||
}
|
||||
if cfg.RequestRateLimitRPM < 0 {
|
||||
cfg.RequestRateLimitRPM = 0
|
||||
}
|
||||
if cfg.RequestRateLimitBurst < 0 {
|
||||
cfg.RequestRateLimitBurst = 0
|
||||
}
|
||||
if cfg.TokenRateLimitTPM < 0 {
|
||||
cfg.TokenRateLimitTPM = 0
|
||||
}
|
||||
if cfg.TokenRateLimitBurst < 0 {
|
||||
cfg.TokenRateLimitBurst = 0
|
||||
}
|
||||
if cfg.EstimatedCharsPerToken <= 0 {
|
||||
cfg.EstimatedCharsPerToken = 1
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(cfg.RateLimitScope)) {
|
||||
case "", "process", "local":
|
||||
cfg.RateLimitScope = "process"
|
||||
case "global", "redis", "distributed":
|
||||
cfg.RateLimitScope = "global"
|
||||
default:
|
||||
cfg.RateLimitScope = "process"
|
||||
}
|
||||
cfg.RateLimitKey = strings.TrimSpace(cfg.RateLimitKey)
|
||||
switch strings.ToLower(strings.TrimSpace(cfg.RateLimitFailoverMode)) {
|
||||
case "", "fail_closed", "closed":
|
||||
cfg.RateLimitFailoverMode = "fail_closed"
|
||||
case "process_degraded", "local_degraded", "process", "local":
|
||||
cfg.RateLimitFailoverMode = "process_degraded"
|
||||
default:
|
||||
cfg.RateLimitFailoverMode = "fail_closed"
|
||||
}
|
||||
if cfg.RateLimitFailoverReplicas < 0 {
|
||||
cfg.RateLimitFailoverReplicas = 0
|
||||
}
|
||||
}
|
||||
|
||||
func candidateConfigPaths(configPath string, local bool) []string {
|
||||
trimmed := strings.TrimSpace(configPath)
|
||||
if trimmed == "" {
|
||||
@@ -992,6 +1044,33 @@ func applyEnvOverrides(cfg *Config) {
|
||||
if model, ok := lookupNonEmptyEnv("LLM_MONITORING_ANSWER_PARSE_MODEL"); ok {
|
||||
cfg.LLM.MonitoringAnswerParseModel = model
|
||||
}
|
||||
if n, ok := lookupIntEnv("LLM_REQUEST_RATE_LIMIT_RPM"); ok {
|
||||
cfg.LLM.RequestRateLimitRPM = n
|
||||
}
|
||||
if n, ok := lookupIntEnv("LLM_REQUEST_RATE_LIMIT_BURST"); ok {
|
||||
cfg.LLM.RequestRateLimitBurst = n
|
||||
}
|
||||
if n, ok := lookupInt64Env("LLM_TOKEN_RATE_LIMIT_TPM"); ok {
|
||||
cfg.LLM.TokenRateLimitTPM = n
|
||||
}
|
||||
if n, ok := lookupIntEnv("LLM_TOKEN_RATE_LIMIT_BURST"); ok {
|
||||
cfg.LLM.TokenRateLimitBurst = n
|
||||
}
|
||||
if n, ok := lookupFloatEnv("LLM_ESTIMATED_CHARS_PER_TOKEN"); ok {
|
||||
cfg.LLM.EstimatedCharsPerToken = n
|
||||
}
|
||||
if scope, ok := lookupNonEmptyEnv("LLM_RATE_LIMIT_SCOPE"); ok {
|
||||
cfg.LLM.RateLimitScope = scope
|
||||
}
|
||||
if key, ok := lookupNonEmptyEnv("LLM_RATE_LIMIT_KEY"); ok {
|
||||
cfg.LLM.RateLimitKey = key
|
||||
}
|
||||
if mode, ok := lookupNonEmptyEnv("LLM_RATE_LIMIT_FAILOVER_MODE"); ok {
|
||||
cfg.LLM.RateLimitFailoverMode = mode
|
||||
}
|
||||
if n, ok := lookupIntEnv("LLM_RATE_LIMIT_FAILOVER_REPLICAS"); ok {
|
||||
cfg.LLM.RateLimitFailoverReplicas = n
|
||||
}
|
||||
if enabled, ok := lookupBoolEnv("COMPLIANCE_ENABLED"); ok {
|
||||
cfg.Compliance.Enabled = enabled
|
||||
}
|
||||
@@ -1136,6 +1215,12 @@ func applyEnvOverrides(cfg *Config) {
|
||||
if rabbitmqURL, ok := lookupNonEmptyEnv("RABBITMQ_URL"); ok {
|
||||
cfg.RabbitMQ.URL = rabbitmqURL
|
||||
}
|
||||
if n, ok := lookupIntEnv("RABBITMQ_CONSUMER_PREFETCH"); ok {
|
||||
cfg.RabbitMQ.ConsumerPrefetch = n
|
||||
}
|
||||
if n, ok := lookupIntEnv("RABBITMQ_GENERATION_CONSUMER_PREFETCH"); ok {
|
||||
cfg.RabbitMQ.GenerationConsumerPrefetch = n
|
||||
}
|
||||
if url, ok := lookupNonEmptyEnv("QDRANT_URL"); ok {
|
||||
cfg.Qdrant.URL = url
|
||||
}
|
||||
@@ -1391,6 +1476,12 @@ func normalizeRabbitMQConfig(cfg *RabbitMQConfig) {
|
||||
if strings.TrimSpace(cfg.ComplianceReviewDLQRouteKey) == "" {
|
||||
cfg.ComplianceReviewDLQRouteKey = "compliance.review.run.dlq"
|
||||
}
|
||||
if cfg.ConsumerPrefetch < 0 {
|
||||
cfg.ConsumerPrefetch = 0
|
||||
}
|
||||
if cfg.GenerationConsumerPrefetch < 0 {
|
||||
cfg.GenerationConsumerPrefetch = 0
|
||||
}
|
||||
if cfg.PublishChannelPoolSize <= 0 {
|
||||
cfg.PublishChannelPoolSize = 16
|
||||
}
|
||||
|
||||
@@ -51,6 +51,125 @@ compliance:
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadAppliesLLMRateLimitConfig(t *testing.T) {
|
||||
configPath := writeTestConfig(t, `
|
||||
llm:
|
||||
rate_limit_scope: global
|
||||
rate_limit_key: ark:shared
|
||||
request_rate_limit_rpm: 30000
|
||||
request_rate_limit_burst: 100
|
||||
token_rate_limit_tpm: 5000000
|
||||
token_rate_limit_burst: 166000
|
||||
estimated_chars_per_token: 1.25
|
||||
rate_limit_failover_mode: process_degraded
|
||||
rate_limit_failover_replicas: 10
|
||||
`)
|
||||
|
||||
cfg, err := Load(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Load() error = %v", err)
|
||||
}
|
||||
|
||||
if cfg.LLM.RateLimitScope != "global" ||
|
||||
cfg.LLM.RateLimitKey != "ark:shared" ||
|
||||
cfg.LLM.RequestRateLimitRPM != 30000 ||
|
||||
cfg.LLM.RequestRateLimitBurst != 100 ||
|
||||
cfg.LLM.TokenRateLimitTPM != 5000000 ||
|
||||
cfg.LLM.TokenRateLimitBurst != 166000 ||
|
||||
cfg.LLM.EstimatedCharsPerToken != 1.25 ||
|
||||
cfg.LLM.RateLimitFailoverMode != "process_degraded" ||
|
||||
cfg.LLM.RateLimitFailoverReplicas != 10 {
|
||||
t.Fatalf("unexpected llm rate limit config: %#v", cfg.LLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadAppliesLLMRateLimitEnvOverrides(t *testing.T) {
|
||||
t.Setenv("LLM_RATE_LIMIT_SCOPE", "global")
|
||||
t.Setenv("LLM_RATE_LIMIT_KEY", "env-ark-shared")
|
||||
t.Setenv("LLM_REQUEST_RATE_LIMIT_RPM", "9000")
|
||||
t.Setenv("LLM_REQUEST_RATE_LIMIT_BURST", "90")
|
||||
t.Setenv("LLM_TOKEN_RATE_LIMIT_TPM", "1500000")
|
||||
t.Setenv("LLM_TOKEN_RATE_LIMIT_BURST", "150000")
|
||||
t.Setenv("LLM_ESTIMATED_CHARS_PER_TOKEN", "1.5")
|
||||
t.Setenv("LLM_RATE_LIMIT_FAILOVER_MODE", "local")
|
||||
t.Setenv("LLM_RATE_LIMIT_FAILOVER_REPLICAS", "12")
|
||||
|
||||
configPath := writeTestConfig(t, `
|
||||
llm:
|
||||
rate_limit_scope: process
|
||||
request_rate_limit_rpm: 10000
|
||||
token_rate_limit_tpm: 1660000
|
||||
`)
|
||||
|
||||
cfg, err := Load(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Load() error = %v", err)
|
||||
}
|
||||
|
||||
if cfg.LLM.RateLimitScope != "global" ||
|
||||
cfg.LLM.RateLimitKey != "env-ark-shared" ||
|
||||
cfg.LLM.RequestRateLimitRPM != 9000 ||
|
||||
cfg.LLM.RequestRateLimitBurst != 90 ||
|
||||
cfg.LLM.TokenRateLimitTPM != 1500000 ||
|
||||
cfg.LLM.TokenRateLimitBurst != 150000 ||
|
||||
cfg.LLM.EstimatedCharsPerToken != 1.5 ||
|
||||
cfg.LLM.RateLimitFailoverMode != "process_degraded" ||
|
||||
cfg.LLM.RateLimitFailoverReplicas != 12 {
|
||||
t.Fatalf("unexpected llm env rate limit config: %#v", cfg.LLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadNormalizesLLMRateLimitConfig(t *testing.T) {
|
||||
configPath := writeTestConfig(t, `
|
||||
llm:
|
||||
request_rate_limit_rpm: -1
|
||||
request_rate_limit_burst: -1
|
||||
token_rate_limit_tpm: -1
|
||||
token_rate_limit_burst: -1
|
||||
estimated_chars_per_token: 0
|
||||
rate_limit_scope: broken
|
||||
rate_limit_key: " shared-key "
|
||||
rate_limit_failover_mode: broken
|
||||
rate_limit_failover_replicas: -1
|
||||
`)
|
||||
|
||||
cfg, err := Load(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Load() error = %v", err)
|
||||
}
|
||||
|
||||
if cfg.LLM.RequestRateLimitRPM != 0 ||
|
||||
cfg.LLM.RequestRateLimitBurst != 0 ||
|
||||
cfg.LLM.TokenRateLimitTPM != 0 ||
|
||||
cfg.LLM.TokenRateLimitBurst != 0 ||
|
||||
cfg.LLM.EstimatedCharsPerToken != 1 ||
|
||||
cfg.LLM.RateLimitScope != "process" ||
|
||||
cfg.LLM.RateLimitKey != "shared-key" ||
|
||||
cfg.LLM.RateLimitFailoverMode != "fail_closed" ||
|
||||
cfg.LLM.RateLimitFailoverReplicas != 0 {
|
||||
t.Fatalf("unexpected normalized llm rate limit config: %#v", cfg.LLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadAppliesRabbitMQGenerationPrefetchEnvOverride(t *testing.T) {
|
||||
t.Setenv("RABBITMQ_GENERATION_CONSUMER_PREFETCH", "1")
|
||||
|
||||
configPath := writeTestConfig(t, `
|
||||
rabbitmq:
|
||||
consumer_prefetch: 10
|
||||
generation_consumer_prefetch: 5
|
||||
`)
|
||||
|
||||
cfg, err := Load(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Load() error = %v", err)
|
||||
}
|
||||
|
||||
if cfg.RabbitMQ.ConsumerPrefetch != 10 || cfg.RabbitMQ.GenerationConsumerPrefetch != 1 {
|
||||
t.Fatalf("unexpected rabbitmq prefetch config: %#v", cfg.RabbitMQ)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDatabaseConfigDSNEscapesCredentials(t *testing.T) {
|
||||
cfg := DatabaseConfig{
|
||||
Host: "db.internal",
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -280,7 +280,11 @@ func (c *Client) ConsumeMonitorProjectionRebuild(consumerName string) (<-chan am
|
||||
}
|
||||
|
||||
func (c *Client) ConsumeGenerationTask(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) {
|
||||
return c.consume(consumerName, c.cfg.GenerationQueue)
|
||||
prefetch := c.cfg.GenerationConsumerPrefetch
|
||||
if prefetch <= 0 {
|
||||
prefetch = c.cfg.ConsumerPrefetch
|
||||
}
|
||||
return c.consumeWithPrefetch(consumerName, c.cfg.GenerationQueue, prefetch)
|
||||
}
|
||||
|
||||
func (c *Client) ConsumeGenerationEvents(consumerName string) (<-chan amqp.Delivery, *amqp.Channel, error) {
|
||||
@@ -312,14 +316,18 @@ func (c *Client) ConsumeComplianceReviewTask(consumerName string) (<-chan amqp.D
|
||||
}
|
||||
|
||||
func (c *Client) consume(consumerName, queue string) (<-chan amqp.Delivery, *amqp.Channel, error) {
|
||||
return c.consumeWithPrefetch(consumerName, queue, c.cfg.ConsumerPrefetch)
|
||||
}
|
||||
|
||||
func (c *Client) consumeWithPrefetch(consumerName, queue string, prefetch int) (<-chan amqp.Delivery, *amqp.Channel, error) {
|
||||
for attempt := 0; attempt < 2; attempt++ {
|
||||
conn, ch, err := c.openChannel()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
if c.cfg.ConsumerPrefetch > 0 {
|
||||
if err := ch.Qos(c.cfg.ConsumerPrefetch, 0, false); err != nil {
|
||||
if prefetch > 0 {
|
||||
if err := ch.Qos(prefetch, 0, false); err != nil {
|
||||
_ = ch.Close()
|
||||
c.invalidateConnection(conn)
|
||||
if attempt == 1 {
|
||||
|
||||
@@ -18,8 +18,6 @@ import (
|
||||
)
|
||||
|
||||
func TestLoggerIncludesDetailedAuthFailureContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
var logBuffer bytes.Buffer
|
||||
@@ -74,8 +72,6 @@ func TestLoggerIncludesDetailedAuthFailureContext(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestLoggerUsesForwardedClientIPFromTrustedProxy(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
var logBuffer bytes.Buffer
|
||||
@@ -107,8 +103,6 @@ func TestLoggerUsesForwardedClientIPFromTrustedProxy(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestLoggerRedactsSensitiveQueryValues(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
var logBuffer bytes.Buffer
|
||||
|
||||
@@ -10,8 +10,6 @@ import (
|
||||
)
|
||||
|
||||
func TestSecurityHeadersAddsNoStoreForSensitiveAuthPaths(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(SecurityHeaders())
|
||||
@@ -34,8 +32,6 @@ func TestSecurityHeadersAddsNoStoreForSensitiveAuthPaths(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestCORSWithConfigRejectsUnknownPreflightOrigin(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(CORSWithConfig(CORSConfig{AllowedOrigins: []string{"https://admin.example.com"}}))
|
||||
@@ -54,8 +50,6 @@ func TestCORSWithConfigRejectsUnknownPreflightOrigin(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestCORSWithConfigAllowsConfiguredOrigin(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(CORSWithConfig(CORSConfig{AllowedOrigins: []string{"https://admin.example.com"}}))
|
||||
|
||||
Reference in New Issue
Block a user