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 (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
|
@ -42,7 +43,7 @@ func NewFallbackChain(cooldown *CooldownTracker) *FallbackChain {
|
||||||
return &FallbackChain{cooldown: cooldown}
|
return &FallbackChain{cooldown: cooldown}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c FallbackCandidate) cooldownKey() string {
|
func (c FallbackCandidate) cooldownKey(candidates []FallbackCandidate) string {
|
||||||
if strings.TrimSpace(c.CooldownKey) != "" {
|
if strings.TrimSpace(c.CooldownKey) != "" {
|
||||||
return c.CooldownKey
|
return c.CooldownKey
|
||||||
}
|
}
|
||||||
|
|
@ -52,9 +53,49 @@ func (c FallbackCandidate) cooldownKey() string {
|
||||||
}
|
}
|
||||||
return ModelKey(c.Provider, c.Model)
|
return ModelKey(c.Provider, c.Model)
|
||||||
}
|
}
|
||||||
|
if belongsToMultiKeySet(c, candidates) {
|
||||||
|
return ModelKey(c.Provider, c.Model)
|
||||||
|
}
|
||||||
return c.Provider
|
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.
|
// ResolveCandidates parses model config into a deduplicated candidate list.
|
||||||
func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandidate {
|
func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandidate {
|
||||||
return ResolveCandidatesWithLookup(cfg, defaultProvider, nil)
|
return ResolveCandidatesWithLookup(cfg, defaultProvider, nil)
|
||||||
|
|
@ -126,7 +167,7 @@ func (fc *FallbackChain) Execute(
|
||||||
}
|
}
|
||||||
|
|
||||||
for i, candidate := range candidates {
|
for i, candidate := range candidates {
|
||||||
cooldownKey := candidate.cooldownKey()
|
cooldownKey := candidate.cooldownKey(candidates)
|
||||||
|
|
||||||
// Check context before each attempt.
|
// Check context before each attempt.
|
||||||
if ctx.Err() == context.Canceled {
|
if ctx.Err() == context.Canceled {
|
||||||
|
|
|
||||||
|
|
@ -165,8 +165,8 @@ func TestFallback_CooldownSkip(t *testing.T) {
|
||||||
ct, _ := newTestTracker(now)
|
ct, _ := newTestTracker(now)
|
||||||
fc := NewFallbackChain(ct)
|
fc := NewFallbackChain(ct)
|
||||||
|
|
||||||
// Put openai/gpt-4 in cooldown (using ModelKey now)
|
// Put openai in cooldown
|
||||||
ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit)
|
ct.MarkFailure("openai", FailoverRateLimit)
|
||||||
|
|
||||||
candidates := []FallbackCandidate{
|
candidates := []FallbackCandidate{
|
||||||
makeCandidate("openai", "gpt-4"),
|
makeCandidate("openai", "gpt-4"),
|
||||||
|
|
@ -258,9 +258,9 @@ func TestFallback_AllInCooldown(t *testing.T) {
|
||||||
ct := NewCooldownTracker()
|
ct := NewCooldownTracker()
|
||||||
fc := NewFallbackChain(ct)
|
fc := NewFallbackChain(ct)
|
||||||
|
|
||||||
// Put all models in cooldown (using ModelKey now)
|
// Put all providers in cooldown
|
||||||
ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit)
|
ct.MarkFailure("openai", FailoverRateLimit)
|
||||||
ct.MarkFailure(ModelKey("anthropic", "claude"), FailoverBilling)
|
ct.MarkFailure("anthropic", FailoverBilling)
|
||||||
|
|
||||||
candidates := []FallbackCandidate{
|
candidates := []FallbackCandidate{
|
||||||
makeCandidate("openai", "gpt-4"),
|
makeCandidate("openai", "gpt-4"),
|
||||||
|
|
@ -336,13 +336,12 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
|
||||||
fc := NewFallbackChain(ct)
|
fc := NewFallbackChain(ct)
|
||||||
|
|
||||||
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
|
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
|
||||||
modelKey := ModelKey("openai", "gpt-4")
|
|
||||||
|
|
||||||
attempt := 0
|
attempt := 0
|
||||||
run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
|
run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
|
||||||
attempt++
|
attempt++
|
||||||
if attempt == 1 {
|
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
|
return &LLMResponse{Content: "ok", FinishReason: "stop"}, nil
|
||||||
}
|
}
|
||||||
|
|
@ -351,7 +350,7 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("unexpected error: %v", err)
|
t.Fatalf("unexpected error: %v", err)
|
||||||
}
|
}
|
||||||
if !ct.IsAvailable(modelKey) {
|
if !ct.IsAvailable("openai") {
|
||||||
t.Error("success should reset cooldown")
|
t.Error("success should reset cooldown")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue