Files
geo/server/internal/shared/llm/client.go
T
root 347cea1296 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>
2026-07-10 00:14:40 +08:00

93 lines
2.1 KiB
Go

package llm
import (
"context"
"errors"
"fmt"
"strings"
"time"
goredis "github.com/redis/go-redis/v9"
"github.com/geo-platform/tenant-api/internal/shared/config"
)
var ErrNotConfigured = errors.New("llm provider is not configured")
type GenerateRequest struct {
Model string
Prompt string
Timeout time.Duration
MaxOutputTokens int64
WebSearch *WebSearchOptions
ResponseFormat *ResponseFormat
}
type GenerateResult struct {
Content string
Model string
}
type ResponseFormat struct {
Type ResponseFormatType
Name string
Description string
SchemaJSON []byte
Strict bool
}
type ResponseFormatType string
const (
ResponseFormatTypeText ResponseFormatType = "text"
ResponseFormatTypeJSONObject ResponseFormatType = "json_object"
ResponseFormatTypeJSONSchema ResponseFormatType = "json_schema"
)
type WebSearchOptions struct {
Enabled bool
Limit int32
}
type Client interface {
Validate() error
Generate(ctx context.Context, req GenerateRequest, onDelta func(string)) (*GenerateResult, error)
}
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":
client = disabledClient{reason: "llm provider is disabled"}
case "ark":
client = NewArkClient(cfg)
default:
client = disabledClient{reason: fmt.Sprintf("unsupported llm provider %q", cfg.Provider)}
}
return NewRateLimitedClientWithRedis(client, cfg, redisClient)
}
type disabledClient struct {
reason string
}
func (c disabledClient) Validate() error {
if c.reason == "" {
c.reason = "missing LLM configuration"
}
return fmt.Errorf("%w: %s", ErrNotConfigured, c.reason)
}
func (c disabledClient) Generate(context.Context, GenerateRequest, func(string)) (*GenerateResult, error) {
return nil, c.Validate()
}