root/geo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(llm): add configurable request/token rate limiting for LLM client

Browse Source 
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:
liangxu
2026-07-10 00:14:40 +08:00
parent 0acd37918e
commit 347cea1296
16 changed files with 1016 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/internal/shared/config/config.go
+91
View File
@@ -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
}
server/internal/shared/config/config_test.go
+119
View File
@@ -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",