fix(agent): preserve per-model identity in rate limiting and fallback
This commit is contained in:
parent
3f7cf1a0b6
commit
51d61ac970
10 changed files with 385 additions and 94 deletions
|
|
@ -74,7 +74,9 @@ model_list:
|
||||||
- gpt-4o-mini # must also be in model_list; its own rpm applies
|
- gpt-4o-mini # must also be in model_list; its own rpm applies
|
||||||
```
|
```
|
||||||
|
|
||||||
If the primary candidate's bucket is empty the call blocks until a token is available (or context is cancelled). If context deadline is hit while waiting, the error propagates and the agent retries normally.
|
If the current candidate's bucket is empty and there are more candidates available, PicoClaw skips the locally saturated candidate and tries the next fallback immediately. Only the last remaining candidate waits for a token to refill. If the context deadline is hit while waiting on that last candidate, the wait error propagates.
|
||||||
|
|
||||||
|
For `model_list` aliases that resolve to the same underlying provider/model, rate limiting is keyed by the stable config identity (for example `model_name`) rather than the resolved runtime model string. This preserves distinct RPM settings for multi-key and alias-based configurations.
|
||||||
|
|
||||||
### Burst behaviour
|
### Burst behaviour
|
||||||
|
|
||||||
|
|
@ -89,5 +91,5 @@ To reduce burstiness for strict APIs, set a lower `rpm` and rely on the steady-s
|
||||||
| `pkg/providers/ratelimiter.go` | `RateLimiter` (token bucket) + `RateLimiterRegistry` |
|
| `pkg/providers/ratelimiter.go` | `RateLimiter` (token bucket) + `RateLimiterRegistry` |
|
||||||
| `pkg/providers/ratelimiter_test.go` | Unit tests for limiter and registry |
|
| `pkg/providers/ratelimiter_test.go` | Unit tests for limiter and registry |
|
||||||
| `pkg/providers/fallback.go` | `FallbackCandidate.RPM` field; `FallbackChain.rl`; `Wait()` call in `Execute`/`ExecuteImage` |
|
| `pkg/providers/fallback.go` | `FallbackCandidate.RPM` field; `FallbackChain.rl`; `Wait()` call in `Execute`/`ExecuteImage` |
|
||||||
| `pkg/agent/model_resolution.go` | `lookupModelConfigByProtocolModel`; propagates `RPM` into `FallbackCandidate` |
|
| `pkg/agent/model_resolution.go` | Resolves candidates from `model_list`, preserving stable config identity and propagating `RPM` into `FallbackCandidate` |
|
||||||
| `pkg/agent/loop.go` | Build `RateLimiterRegistry`, register all agents' candidates, pass to `NewFallbackChain` |
|
| `pkg/agent/loop.go` | Build `RateLimiterRegistry`, register all agents' candidates, pass to `NewFallbackChain` |
|
||||||
|
|
|
||||||
|
|
@ -165,6 +165,58 @@ func TestNewAgentInstance_ResolveCandidatesFromModelListAlias(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestNewAgentInstance_PreservesDistinctLimiterIdentityForSharedResolvedModel(t *testing.T) {
|
||||||
|
tmpDir := t.TempDir()
|
||||||
|
|
||||||
|
cfg := &config.Config{
|
||||||
|
Agents: config.AgentsConfig{
|
||||||
|
Defaults: config.AgentDefaults{
|
||||||
|
Workspace: tmpDir,
|
||||||
|
ModelName: "glm-4.7",
|
||||||
|
ModelFallbacks: []string{"glm-4.7__key_1"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
ModelList: []*config.ModelConfig{
|
||||||
|
{
|
||||||
|
ModelName: "glm-4.7",
|
||||||
|
Model: "zhipu/glm-4.7",
|
||||||
|
RPM: 1,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
ModelName: "glm-4.7__key_1",
|
||||||
|
Model: "zhipu/glm-4.7",
|
||||||
|
RPM: 3,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
|
||||||
|
if len(agent.Candidates) != 2 {
|
||||||
|
t.Fatalf("len(Candidates) = %d, want 2", len(agent.Candidates))
|
||||||
|
}
|
||||||
|
|
||||||
|
first := agent.Candidates[0]
|
||||||
|
second := agent.Candidates[1]
|
||||||
|
if first.Provider != "zhipu" || first.Model != "glm-4.7" {
|
||||||
|
t.Fatalf("first candidate = %s/%s, want zhipu/glm-4.7", first.Provider, first.Model)
|
||||||
|
}
|
||||||
|
if second.Provider != "zhipu" || second.Model != "glm-4.7" {
|
||||||
|
t.Fatalf("second candidate = %s/%s, want zhipu/glm-4.7", second.Provider, second.Model)
|
||||||
|
}
|
||||||
|
if first.IdentityKey != "model_name:glm-4.7" {
|
||||||
|
t.Fatalf("first identity key = %q, want %q", first.IdentityKey, "model_name:glm-4.7")
|
||||||
|
}
|
||||||
|
if second.IdentityKey != "model_name:glm-4.7__key_1" {
|
||||||
|
t.Fatalf("second identity key = %q, want %q", second.IdentityKey, "model_name:glm-4.7__key_1")
|
||||||
|
}
|
||||||
|
if first.RPM != 1 {
|
||||||
|
t.Fatalf("first RPM = %d, want 1", first.RPM)
|
||||||
|
}
|
||||||
|
if second.RPM != 3 {
|
||||||
|
t.Fatalf("second RPM = %d, want 3", second.RPM)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
|
func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
|
||||||
workspace := t.TempDir()
|
workspace := t.TempDir()
|
||||||
mediaDir := media.TempDir()
|
mediaDir := media.TempDir()
|
||||||
|
|
|
||||||
|
|
@ -3453,7 +3453,7 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt
|
||||||
return "", fmt.Errorf("failed to initialize model %q: %w", value, err)
|
return "", fmt.Errorf("failed to initialize model %q: %w", value, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
nextCandidates := resolveModelCandidates(cfg, cfg.Agents.Defaults.Provider, modelCfg.Model, agent.Fallbacks)
|
nextCandidates := resolveModelCandidates(cfg, cfg.Agents.Defaults.Provider, value, agent.Fallbacks)
|
||||||
if len(nextCandidates) == 0 {
|
if len(nextCandidates) == 0 {
|
||||||
return "", fmt.Errorf("model %q did not resolve to any provider candidates", value)
|
return "", fmt.Errorf("model %q did not resolve to any provider candidates", value)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -8,44 +8,102 @@ import (
|
||||||
"github.com/sipeed/picoclaw/pkg/providers"
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
)
|
)
|
||||||
|
|
||||||
func buildModelListResolver(cfg *config.Config) func(raw string) (string, bool) {
|
func ensureProtocolModel(model string) string {
|
||||||
ensureProtocol := func(model string) string {
|
model = strings.TrimSpace(model)
|
||||||
model = strings.TrimSpace(model)
|
if model == "" {
|
||||||
if model == "" {
|
return ""
|
||||||
return ""
|
}
|
||||||
}
|
if strings.Contains(model, "/") {
|
||||||
if strings.Contains(model, "/") {
|
return model
|
||||||
return model
|
}
|
||||||
}
|
return "openai/" + model
|
||||||
return "openai/" + model
|
}
|
||||||
|
|
||||||
|
func modelConfigIdentityKey(mc *config.ModelConfig) string {
|
||||||
|
if mc == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
if name := strings.TrimSpace(mc.ModelName); name != "" {
|
||||||
|
return "model_name:" + name
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
func candidateFromModelConfig(
|
||||||
|
defaultProvider string,
|
||||||
|
mc *config.ModelConfig,
|
||||||
|
) (providers.FallbackCandidate, bool) {
|
||||||
|
if mc == nil {
|
||||||
|
return providers.FallbackCandidate{}, false
|
||||||
}
|
}
|
||||||
|
|
||||||
return func(raw string) (string, bool) {
|
ref := providers.ParseModelRef(ensureProtocolModel(mc.Model), defaultProvider)
|
||||||
raw = strings.TrimSpace(raw)
|
if ref == nil {
|
||||||
if raw == "" || cfg == nil {
|
return providers.FallbackCandidate{}, false
|
||||||
return "", false
|
|
||||||
}
|
|
||||||
|
|
||||||
if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" {
|
|
||||||
return ensureProtocol(mc.Model), true
|
|
||||||
}
|
|
||||||
|
|
||||||
for i := range cfg.ModelList {
|
|
||||||
fullModel := strings.TrimSpace(cfg.ModelList[i].Model)
|
|
||||||
if fullModel == "" {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
if fullModel == raw {
|
|
||||||
return ensureProtocol(fullModel), true
|
|
||||||
}
|
|
||||||
_, modelID := providers.ExtractProtocol(fullModel)
|
|
||||||
if modelID == raw {
|
|
||||||
return ensureProtocol(fullModel), true
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return "", false
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
return providers.FallbackCandidate{
|
||||||
|
Provider: ref.Provider,
|
||||||
|
Model: ref.Model,
|
||||||
|
RPM: mc.RPM,
|
||||||
|
IdentityKey: modelConfigIdentityKey(mc),
|
||||||
|
}, true
|
||||||
|
}
|
||||||
|
|
||||||
|
func lookupModelConfigByRef(cfg *config.Config, raw string) *config.ModelConfig {
|
||||||
|
raw = strings.TrimSpace(raw)
|
||||||
|
if raw == "" || cfg == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" {
|
||||||
|
return mc
|
||||||
|
}
|
||||||
|
|
||||||
|
for i := range cfg.ModelList {
|
||||||
|
mc := cfg.ModelList[i]
|
||||||
|
if mc == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
fullModel := strings.TrimSpace(mc.Model)
|
||||||
|
if fullModel == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if fullModel == raw {
|
||||||
|
return mc
|
||||||
|
}
|
||||||
|
_, modelID := providers.ExtractProtocol(fullModel)
|
||||||
|
if modelID == raw {
|
||||||
|
return mc
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func resolveModelCandidate(
|
||||||
|
cfg *config.Config,
|
||||||
|
defaultProvider string,
|
||||||
|
raw string,
|
||||||
|
) (providers.FallbackCandidate, bool) {
|
||||||
|
raw = strings.TrimSpace(raw)
|
||||||
|
if raw == "" {
|
||||||
|
return providers.FallbackCandidate{}, false
|
||||||
|
}
|
||||||
|
|
||||||
|
if mc := lookupModelConfigByRef(cfg, raw); mc != nil {
|
||||||
|
return candidateFromModelConfig(defaultProvider, mc)
|
||||||
|
}
|
||||||
|
|
||||||
|
ref := providers.ParseModelRef(raw, defaultProvider)
|
||||||
|
if ref == nil {
|
||||||
|
return providers.FallbackCandidate{}, false
|
||||||
|
}
|
||||||
|
|
||||||
|
return providers.FallbackCandidate{
|
||||||
|
Provider: ref.Provider,
|
||||||
|
Model: ref.Model,
|
||||||
|
}, true
|
||||||
}
|
}
|
||||||
|
|
||||||
func resolveModelCandidates(
|
func resolveModelCandidates(
|
||||||
|
|
@ -54,45 +112,31 @@ func resolveModelCandidates(
|
||||||
primary string,
|
primary string,
|
||||||
fallbacks []string,
|
fallbacks []string,
|
||||||
) []providers.FallbackCandidate {
|
) []providers.FallbackCandidate {
|
||||||
candidates := providers.ResolveCandidatesWithLookup(
|
seen := make(map[string]bool)
|
||||||
providers.ModelConfig{
|
candidates := make([]providers.FallbackCandidate, 0, 1+len(fallbacks))
|
||||||
Primary: primary,
|
|
||||||
Fallbacks: fallbacks,
|
|
||||||
},
|
|
||||||
defaultProvider,
|
|
||||||
buildModelListResolver(cfg),
|
|
||||||
)
|
|
||||||
|
|
||||||
// Propagate RPM from ModelConfig to each candidate.
|
addCandidate := func(raw string) {
|
||||||
if cfg != nil {
|
candidate, ok := resolveModelCandidate(cfg, defaultProvider, raw)
|
||||||
for i := range candidates {
|
if !ok {
|
||||||
mc := lookupModelConfigByProtocolModel(cfg, candidates[i].Provider, candidates[i].Model)
|
return
|
||||||
if mc != nil && mc.RPM > 0 {
|
|
||||||
candidates[i].RPM = mc.RPM
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
key := candidate.StableKey()
|
||||||
|
if seen[key] {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
seen[key] = true
|
||||||
|
candidates = append(candidates, candidate)
|
||||||
|
}
|
||||||
|
|
||||||
|
addCandidate(primary)
|
||||||
|
for _, fallback := range fallbacks {
|
||||||
|
addCandidate(fallback)
|
||||||
}
|
}
|
||||||
|
|
||||||
return candidates
|
return candidates
|
||||||
}
|
}
|
||||||
|
|
||||||
// lookupModelConfigByProtocolModel finds a ModelConfig whose resolved provider/model matches.
|
|
||||||
func lookupModelConfigByProtocolModel(cfg *config.Config, provider, model string) *config.ModelConfig {
|
|
||||||
for _, mc := range cfg.ModelList {
|
|
||||||
if mc == nil || strings.TrimSpace(mc.Model) == "" {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
p, m := providers.ExtractProtocol(mc.Model)
|
|
||||||
if p == "" {
|
|
||||||
p = "openai"
|
|
||||||
}
|
|
||||||
if p == provider && m == model {
|
|
||||||
return mc
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func resolvedCandidateModel(candidates []providers.FallbackCandidate, fallback string) string {
|
func resolvedCandidateModel(candidates []providers.FallbackCandidate, fallback string) string {
|
||||||
if len(candidates) > 0 && strings.TrimSpace(candidates[0].Model) != "" {
|
if len(candidates) > 0 && strings.TrimSpace(candidates[0].Model) != "" {
|
||||||
return candidates[0].Model
|
return candidates[0].Model
|
||||||
|
|
|
||||||
|
|
@ -208,7 +208,10 @@ func (p *Provider) Chat(
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// Check for SSO token expiration errors and provide actionable guidance
|
// Check for SSO token expiration errors and provide actionable guidance
|
||||||
if isSSOTokenError(err) {
|
if isSSOTokenError(err) {
|
||||||
return nil, fmt.Errorf("bedrock converse: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w", err)
|
return nil, fmt.Errorf(
|
||||||
|
"bedrock converse: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w",
|
||||||
|
err,
|
||||||
|
)
|
||||||
}
|
}
|
||||||
return nil, fmt.Errorf("bedrock converse: %w", err)
|
return nil, fmt.Errorf("bedrock converse: %w", err)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -583,13 +583,17 @@ func TestIsSSOTokenError(t *testing.T) {
|
||||||
expected: true,
|
expected: true,
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
name: "full SSO error message",
|
name: "full SSO error message",
|
||||||
err: fmt.Errorf("get identity: get credentials: failed to refresh cached credentials, refresh cached SSO token failed, unable to refresh SSO token"),
|
err: fmt.Errorf(
|
||||||
|
"get identity: get credentials: failed to refresh cached credentials, refresh cached SSO token failed, unable to refresh SSO token",
|
||||||
|
),
|
||||||
expected: true,
|
expected: true,
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
name: "SSO token file missing",
|
name: "SSO token file missing",
|
||||||
err: fmt.Errorf("get identity: get credentials: failed to refresh cached credentials, failed to read cached SSO token file, open ~/.aws/sso/cache/abc123.json: no such file or directory"),
|
err: fmt.Errorf(
|
||||||
|
"get identity: get credentials: failed to refresh cached credentials, failed to read cached SSO token file, open ~/.aws/sso/cache/abc123.json: no such file or directory",
|
||||||
|
),
|
||||||
expected: true,
|
expected: true,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -15,9 +15,19 @@ type FallbackChain struct {
|
||||||
|
|
||||||
// FallbackCandidate represents one model/provider to try.
|
// FallbackCandidate represents one model/provider to try.
|
||||||
type FallbackCandidate struct {
|
type FallbackCandidate struct {
|
||||||
Provider string
|
Provider string
|
||||||
Model string
|
Model string
|
||||||
RPM int // requests per minute; 0 means unrestricted
|
RPM int // requests per minute; 0 means unrestricted
|
||||||
|
IdentityKey string // optional stable config identity for cooldown/rate limiting
|
||||||
|
}
|
||||||
|
|
||||||
|
// StableKey returns the candidate's config-level identity when available,
|
||||||
|
// otherwise it falls back to the runtime provider/model key.
|
||||||
|
func (c FallbackCandidate) StableKey() string {
|
||||||
|
if key := strings.TrimSpace(c.IdentityKey); key != "" {
|
||||||
|
return key
|
||||||
|
}
|
||||||
|
return ModelKey(c.Provider, c.Model)
|
||||||
}
|
}
|
||||||
|
|
||||||
// FallbackResult contains the successful response and metadata about all attempts.
|
// FallbackResult contains the successful response and metadata about all attempts.
|
||||||
|
|
@ -120,9 +130,9 @@ func (fc *FallbackChain) Execute(
|
||||||
return nil, context.Canceled
|
return nil, context.Canceled
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check cooldown (per provider/model, not just provider).
|
// Check cooldown per stable candidate identity, not just provider/model.
|
||||||
// This allows multi-key failover where different keys use different model names.
|
// This allows aliases and multi-key configs to fail over independently.
|
||||||
cooldownKey := ModelKey(candidate.Provider, candidate.Model)
|
cooldownKey := candidate.StableKey()
|
||||||
if !fc.cooldown.IsAvailable(cooldownKey) {
|
if !fc.cooldown.IsAvailable(cooldownKey) {
|
||||||
remaining := fc.cooldown.CooldownRemaining(cooldownKey)
|
remaining := fc.cooldown.CooldownRemaining(cooldownKey)
|
||||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||||
|
|
@ -140,9 +150,29 @@ func (fc *FallbackChain) Execute(
|
||||||
}
|
}
|
||||||
|
|
||||||
// Enforce per-candidate rate limit before calling the provider.
|
// Enforce per-candidate rate limit before calling the provider.
|
||||||
|
// If this candidate is locally saturated, try other candidates first.
|
||||||
if fc.rl != nil {
|
if fc.rl != nil {
|
||||||
if waitErr := fc.rl.Wait(ctx, cooldownKey); waitErr != nil {
|
if !fc.rl.TryAcquire(cooldownKey) {
|
||||||
return nil, waitErr
|
if i < len(candidates)-1 {
|
||||||
|
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||||
|
Provider: candidate.Provider,
|
||||||
|
Model: candidate.Model,
|
||||||
|
Skipped: true,
|
||||||
|
Reason: FailoverRateLimit,
|
||||||
|
Error: fmt.Errorf("%s waiting for local rate limit token", cooldownKey),
|
||||||
|
})
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if waitErr := fc.rl.Wait(ctx, cooldownKey); waitErr != nil {
|
||||||
|
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||||
|
Provider: candidate.Provider,
|
||||||
|
Model: candidate.Model,
|
||||||
|
Skipped: true,
|
||||||
|
Reason: FailoverRateLimit,
|
||||||
|
Error: waitErr,
|
||||||
|
})
|
||||||
|
return nil, waitErr
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -240,10 +270,30 @@ func (fc *FallbackChain) ExecuteImage(
|
||||||
}
|
}
|
||||||
|
|
||||||
// Enforce per-candidate rate limit before calling the provider.
|
// Enforce per-candidate rate limit before calling the provider.
|
||||||
imageKey := ModelKey(candidate.Provider, candidate.Model)
|
// If this candidate is locally saturated, try other candidates first.
|
||||||
|
imageKey := candidate.StableKey()
|
||||||
if fc.rl != nil {
|
if fc.rl != nil {
|
||||||
if waitErr := fc.rl.Wait(ctx, imageKey); waitErr != nil {
|
if !fc.rl.TryAcquire(imageKey) {
|
||||||
return nil, waitErr
|
if i < len(candidates)-1 {
|
||||||
|
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||||
|
Provider: candidate.Provider,
|
||||||
|
Model: candidate.Model,
|
||||||
|
Skipped: true,
|
||||||
|
Reason: FailoverRateLimit,
|
||||||
|
Error: fmt.Errorf("%s waiting for local rate limit token", imageKey),
|
||||||
|
})
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if waitErr := fc.rl.Wait(ctx, imageKey); waitErr != nil {
|
||||||
|
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||||
|
Provider: candidate.Provider,
|
||||||
|
Model: candidate.Model,
|
||||||
|
Skipped: true,
|
||||||
|
Reason: FailoverRateLimit,
|
||||||
|
Error: waitErr,
|
||||||
|
})
|
||||||
|
return nil, waitErr
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -293,6 +293,42 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestFallback_LocalRateLimitSkipsToHealthyFallback(t *testing.T) {
|
||||||
|
ct := NewCooldownTracker()
|
||||||
|
rl := NewRateLimiterRegistry()
|
||||||
|
rl.Register("model_name:primary", 1)
|
||||||
|
if err := rl.Wait(context.Background(), "model_name:primary"); err != nil {
|
||||||
|
t.Fatalf("failed to pre-drain primary limiter: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
fc := NewFallbackChain(ct, rl)
|
||||||
|
candidates := []FallbackCandidate{
|
||||||
|
{Provider: "openai", Model: "gpt-4o", IdentityKey: "model_name:primary"},
|
||||||
|
{Provider: "anthropic", Model: "claude", IdentityKey: "model_name:fallback"},
|
||||||
|
}
|
||||||
|
|
||||||
|
run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
|
||||||
|
if provider != "anthropic" || model != "claude" {
|
||||||
|
t.Fatalf("expected fallback candidate to run, got %s/%s", provider, model)
|
||||||
|
}
|
||||||
|
return &LLMResponse{Content: "fallback ok", FinishReason: "stop"}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
result, err := fc.Execute(ctx, candidates, run)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("expected fallback success, got error: %v", err)
|
||||||
|
}
|
||||||
|
if result.Provider != "anthropic" || result.Model != "claude" {
|
||||||
|
t.Fatalf("result = %s/%s, want anthropic/claude", result.Provider, result.Model)
|
||||||
|
}
|
||||||
|
if len(result.Attempts) != 1 || !result.Attempts[0].Skipped {
|
||||||
|
t.Fatalf("expected one skipped primary attempt, got %+v", result.Attempts)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// --- Image Fallback Tests ---
|
// --- Image Fallback Tests ---
|
||||||
|
|
||||||
func TestImageFallback_Success(t *testing.T) {
|
func TestImageFallback_Success(t *testing.T) {
|
||||||
|
|
@ -384,6 +420,42 @@ func TestImageFallback_RetryOnOtherErrors(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestImageFallback_LocalRateLimitSkipsToHealthyFallback(t *testing.T) {
|
||||||
|
ct := NewCooldownTracker()
|
||||||
|
rl := NewRateLimiterRegistry()
|
||||||
|
rl.Register("model_name:primary-image", 1)
|
||||||
|
if err := rl.Wait(context.Background(), "model_name:primary-image"); err != nil {
|
||||||
|
t.Fatalf("failed to pre-drain primary image limiter: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
fc := NewFallbackChain(ct, rl)
|
||||||
|
candidates := []FallbackCandidate{
|
||||||
|
{Provider: "openai", Model: "gpt-4o", IdentityKey: "model_name:primary-image"},
|
||||||
|
{Provider: "anthropic", Model: "claude-sonnet", IdentityKey: "model_name:fallback-image"},
|
||||||
|
}
|
||||||
|
|
||||||
|
run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
|
||||||
|
if provider != "anthropic" || model != "claude-sonnet" {
|
||||||
|
t.Fatalf("expected image fallback candidate to run, got %s/%s", provider, model)
|
||||||
|
}
|
||||||
|
return &LLMResponse{Content: "image fallback ok", FinishReason: "stop"}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
result, err := fc.ExecuteImage(ctx, candidates, run)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("expected image fallback success, got error: %v", err)
|
||||||
|
}
|
||||||
|
if result.Provider != "anthropic" || result.Model != "claude-sonnet" {
|
||||||
|
t.Fatalf("result = %s/%s, want anthropic/claude-sonnet", result.Provider, result.Model)
|
||||||
|
}
|
||||||
|
if len(result.Attempts) != 1 || !result.Attempts[0].Skipped {
|
||||||
|
t.Fatalf("expected one skipped primary attempt, got %+v", result.Attempts)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestImageFallback_NoCandidates(t *testing.T) {
|
func TestImageFallback_NoCandidates(t *testing.T) {
|
||||||
ct := NewCooldownTracker()
|
ct := NewCooldownTracker()
|
||||||
fc := NewFallbackChain(ct, nil)
|
fc := NewFallbackChain(ct, nil)
|
||||||
|
|
|
||||||
|
|
@ -18,6 +18,15 @@ type RateLimiter struct {
|
||||||
nowFunc func() time.Time // for testing
|
nowFunc func() time.Time // for testing
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (rl *RateLimiter) refillLocked(now time.Time) {
|
||||||
|
elapsed := now.Sub(rl.lastTick).Seconds()
|
||||||
|
rl.lastTick = now
|
||||||
|
|
||||||
|
// Refill tokens proportional to elapsed time.
|
||||||
|
refill := elapsed * float64(rl.rpm) / 60.0
|
||||||
|
rl.tokens = min(rl.maxBurst, rl.tokens+refill)
|
||||||
|
}
|
||||||
|
|
||||||
// newRateLimiter creates a RateLimiter that allows rpm requests/minute.
|
// newRateLimiter creates a RateLimiter that allows rpm requests/minute.
|
||||||
func newRateLimiter(rpm int) *RateLimiter {
|
func newRateLimiter(rpm int) *RateLimiter {
|
||||||
return &RateLimiter{
|
return &RateLimiter{
|
||||||
|
|
@ -35,12 +44,7 @@ func (rl *RateLimiter) Wait(ctx context.Context) error {
|
||||||
for {
|
for {
|
||||||
rl.mu.Lock()
|
rl.mu.Lock()
|
||||||
now := rl.nowFunc()
|
now := rl.nowFunc()
|
||||||
elapsed := now.Sub(rl.lastTick).Seconds()
|
rl.refillLocked(now)
|
||||||
rl.lastTick = now
|
|
||||||
|
|
||||||
// Refill tokens proportional to elapsed time.
|
|
||||||
refill := elapsed * float64(rl.rpm) / 60.0
|
|
||||||
rl.tokens = min(rl.maxBurst, rl.tokens+refill)
|
|
||||||
|
|
||||||
if rl.tokens >= 1.0 {
|
if rl.tokens >= 1.0 {
|
||||||
rl.tokens--
|
rl.tokens--
|
||||||
|
|
@ -53,15 +57,32 @@ func (rl *RateLimiter) Wait(ctx context.Context) error {
|
||||||
waitSec := deficit / (float64(rl.rpm) / 60.0)
|
waitSec := deficit / (float64(rl.rpm) / 60.0)
|
||||||
rl.mu.Unlock()
|
rl.mu.Unlock()
|
||||||
|
|
||||||
|
timer := time.NewTimer(time.Duration(waitSec * float64(time.Second)))
|
||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
|
if !timer.Stop() {
|
||||||
|
<-timer.C
|
||||||
|
}
|
||||||
return ctx.Err()
|
return ctx.Err()
|
||||||
case <-time.After(time.Duration(waitSec * float64(time.Second))):
|
case <-timer.C:
|
||||||
// Loop to re-check (another goroutine may have consumed the token).
|
// Loop to re-check (another goroutine may have consumed the token).
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TryAcquire attempts to consume a token without blocking.
|
||||||
|
func (rl *RateLimiter) TryAcquire() bool {
|
||||||
|
rl.mu.Lock()
|
||||||
|
defer rl.mu.Unlock()
|
||||||
|
|
||||||
|
rl.refillLocked(rl.nowFunc())
|
||||||
|
if rl.tokens < 1.0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
rl.tokens--
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
// RateLimiterRegistry holds per-candidate rate limiters.
|
// RateLimiterRegistry holds per-candidate rate limiters.
|
||||||
// Candidates with RPM=0 are unrestricted.
|
// Candidates with RPM=0 are unrestricted.
|
||||||
// Thread-safe for concurrent reads/writes.
|
// Thread-safe for concurrent reads/writes.
|
||||||
|
|
@ -100,12 +121,24 @@ func (r *RateLimiterRegistry) Wait(ctx context.Context, key string) error {
|
||||||
return rl.Wait(ctx)
|
return rl.Wait(ctx)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TryAcquire attempts to consume a token for the given key without blocking.
|
||||||
|
// If no limiter is registered for key, it returns true.
|
||||||
|
func (r *RateLimiterRegistry) TryAcquire(key string) bool {
|
||||||
|
r.mu.RLock()
|
||||||
|
rl := r.limiters[key]
|
||||||
|
r.mu.RUnlock()
|
||||||
|
if rl == nil {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return rl.TryAcquire()
|
||||||
|
}
|
||||||
|
|
||||||
// RegisterCandidates registers rate limiters for all candidates that have RPM > 0.
|
// RegisterCandidates registers rate limiters for all candidates that have RPM > 0.
|
||||||
// Candidates with RPM == 0 are ignored (no restriction).
|
// Candidates with RPM == 0 are ignored (no restriction).
|
||||||
func (r *RateLimiterRegistry) RegisterCandidates(candidates []FallbackCandidate) {
|
func (r *RateLimiterRegistry) RegisterCandidates(candidates []FallbackCandidate) {
|
||||||
for _, c := range candidates {
|
for _, c := range candidates {
|
||||||
if c.RPM > 0 {
|
if c.RPM > 0 {
|
||||||
r.Register(ModelKey(c.Provider, c.Model), c.RPM)
|
r.Register(c.StableKey(), c.RPM)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -149,6 +149,37 @@ func TestRateLimiterRegistry_RegisterCandidates(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestRateLimiterRegistry_RegisterCandidatesUsesStableIdentity(t *testing.T) {
|
||||||
|
r := NewRateLimiterRegistry()
|
||||||
|
candidates := []FallbackCandidate{
|
||||||
|
{Provider: "openai", Model: "gpt-4o", RPM: 1, IdentityKey: "model_name:primary"},
|
||||||
|
{Provider: "openai", Model: "gpt-4o", RPM: 2, IdentityKey: "model_name:fallback"},
|
||||||
|
}
|
||||||
|
r.RegisterCandidates(candidates)
|
||||||
|
|
||||||
|
if err := r.Wait(context.Background(), "model_name:primary"); err != nil {
|
||||||
|
t.Fatalf("primary first call should pass: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.Wait(context.Background(), "model_name:fallback"); err != nil {
|
||||||
|
t.Fatalf("fallback first call should pass: %v", err)
|
||||||
|
}
|
||||||
|
if err := r.Wait(context.Background(), "model_name:fallback"); err != nil {
|
||||||
|
t.Fatalf("fallback second call should pass: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
ctxPrimary, cancelPrimary := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||||
|
defer cancelPrimary()
|
||||||
|
if err := r.Wait(ctxPrimary, "model_name:primary"); err == nil {
|
||||||
|
t.Fatal("primary second call should have been limited")
|
||||||
|
}
|
||||||
|
|
||||||
|
ctxFallback, cancelFallback := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||||
|
defer cancelFallback()
|
||||||
|
if err := r.Wait(ctxFallback, "model_name:fallback"); err == nil {
|
||||||
|
t.Fatal("fallback third call should have been limited")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// TestRateLimiter_Concurrency verifies thread safety under concurrent access.
|
// TestRateLimiter_Concurrency verifies thread safety under concurrent access.
|
||||||
func TestRateLimiter_Concurrency(t *testing.T) {
|
func TestRateLimiter_Concurrency(t *testing.T) {
|
||||||
rpm := 20
|
rpm := 20
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue