geo/server/internal/shared/llm/rate_limited.go

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
}