fix(provider): preserve multikey cooldown isolation
This commit is contained in:
parent
7577f952cd
commit
01d907c9c5
2 changed files with 50 additions and 10 deletions
|
|
@ -3,6 +3,7 @@ package providers
|
|||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
|
@ -42,7 +43,7 @@ func NewFallbackChain(cooldown *CooldownTracker) *FallbackChain {
|
|||
return &FallbackChain{cooldown: cooldown}
|
||||
}
|
||||
|
||||
func (c FallbackCandidate) cooldownKey() string {
|
||||
func (c FallbackCandidate) cooldownKey(candidates []FallbackCandidate) string {
|
||||
if strings.TrimSpace(c.CooldownKey) != "" {
|
||||
return c.CooldownKey
|
||||
}
|
||||
|
|
@ -52,9 +53,49 @@ func (c FallbackCandidate) cooldownKey() string {
|
|||
}
|
||||
return ModelKey(c.Provider, c.Model)
|
||||
}
|
||||
if belongsToMultiKeySet(c, candidates) {
|
||||
return ModelKey(c.Provider, c.Model)
|
||||
}
|
||||
return c.Provider
|
||||
}
|
||||
|
||||
func belongsToMultiKeySet(candidate FallbackCandidate, candidates []FallbackCandidate) bool {
|
||||
provider := NormalizeProvider(candidate.Provider)
|
||||
model := strings.ToLower(strings.TrimSpace(candidate.Model))
|
||||
if provider == "" || model == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
baseModel, isReplica := multiKeyBaseModel(model)
|
||||
if isReplica {
|
||||
return true
|
||||
}
|
||||
|
||||
for _, other := range candidates {
|
||||
if NormalizeProvider(other.Provider) != provider {
|
||||
continue
|
||||
}
|
||||
otherBase, otherIsReplica := multiKeyBaseModel(other.Model)
|
||||
if otherIsReplica && otherBase == baseModel {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
func multiKeyBaseModel(model string) (string, bool) {
|
||||
normalized := strings.ToLower(strings.TrimSpace(model))
|
||||
idx := strings.LastIndex(normalized, "__key_")
|
||||
if idx <= 0 {
|
||||
return normalized, false
|
||||
}
|
||||
if _, err := strconv.Atoi(normalized[idx+len("__key_"):]); err != nil {
|
||||
return normalized, false
|
||||
}
|
||||
return normalized[:idx], true
|
||||
}
|
||||
|
||||
// ResolveCandidates parses model config into a deduplicated candidate list.
|
||||
func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandidate {
|
||||
return ResolveCandidatesWithLookup(cfg, defaultProvider, nil)
|
||||
|
|
@ -126,7 +167,7 @@ func (fc *FallbackChain) Execute(
|
|||
}
|
||||
|
||||
for i, candidate := range candidates {
|
||||
cooldownKey := candidate.cooldownKey()
|
||||
cooldownKey := candidate.cooldownKey(candidates)
|
||||
|
||||
// Check context before each attempt.
|
||||
if ctx.Err() == context.Canceled {
|
||||
|
|
|
|||
|
|
@ -165,8 +165,8 @@ func TestFallback_CooldownSkip(t *testing.T) {
|
|||
ct, _ := newTestTracker(now)
|
||||
fc := NewFallbackChain(ct)
|
||||
|
||||
// Put openai/gpt-4 in cooldown (using ModelKey now)
|
||||
ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit)
|
||||
// Put openai in cooldown
|
||||
ct.MarkFailure("openai", FailoverRateLimit)
|
||||
|
||||
candidates := []FallbackCandidate{
|
||||
makeCandidate("openai", "gpt-4"),
|
||||
|
|
@ -258,9 +258,9 @@ func TestFallback_AllInCooldown(t *testing.T) {
|
|||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct)
|
||||
|
||||
// Put all models in cooldown (using ModelKey now)
|
||||
ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit)
|
||||
ct.MarkFailure(ModelKey("anthropic", "claude"), FailoverBilling)
|
||||
// Put all providers in cooldown
|
||||
ct.MarkFailure("openai", FailoverRateLimit)
|
||||
ct.MarkFailure("anthropic", FailoverBilling)
|
||||
|
||||
candidates := []FallbackCandidate{
|
||||
makeCandidate("openai", "gpt-4"),
|
||||
|
|
@ -336,13 +336,12 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
|
|||
fc := NewFallbackChain(ct)
|
||||
|
||||
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
|
||||
modelKey := ModelKey("openai", "gpt-4")
|
||||
|
||||
attempt := 0
|
||||
run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
|
||||
attempt++
|
||||
if attempt == 1 {
|
||||
ct.MarkFailure(modelKey, FailoverRateLimit) // simulate failure tracked elsewhere
|
||||
ct.MarkFailure("openai", FailoverRateLimit) // simulate failure tracked elsewhere
|
||||
}
|
||||
return &LLMResponse{Content: "ok", FinishReason: "stop"}, nil
|
||||
}
|
||||
|
|
@ -351,7 +350,7 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if !ct.IsAvailable(modelKey) {
|
||||
if !ct.IsAvailable("openai") {
|
||||
t.Error("success should reset cooldown")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue