refactor(provider): optimize cooldown strategy resolution
This commit is contained in:
parent
9c55a5a506
commit
7577f952cd
5 changed files with 100 additions and 32 deletions
|
|
@ -234,57 +234,53 @@ func applyCooldownKeys(cfg *config.Config, candidates []providers.FallbackCandid
|
|||
return candidates
|
||||
}
|
||||
|
||||
strategyLookup := buildCooldownStrategyLookup(cfg)
|
||||
resolved := make([]providers.FallbackCandidate, len(candidates))
|
||||
copy(resolved, candidates)
|
||||
|
||||
for i := range resolved {
|
||||
resolved[i].CooldownKey = resolveCooldownKey(cfg, resolved[i])
|
||||
resolved[i].CooldownKey = resolveCooldownKey(strategyLookup, resolved[i])
|
||||
}
|
||||
|
||||
return resolved
|
||||
}
|
||||
|
||||
func resolveCooldownKey(cfg *config.Config, candidate providers.FallbackCandidate) string {
|
||||
if candidate.Provider == "" {
|
||||
return ""
|
||||
func buildCooldownStrategyLookup(cfg *config.Config) map[string]string {
|
||||
if cfg == nil || len(cfg.ModelList) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
if cfg != nil && candidateUsesPerModelCooldown(cfg, candidate) {
|
||||
return providers.ModelKey(candidate.Provider, candidate.Model)
|
||||
}
|
||||
|
||||
return candidate.Provider
|
||||
}
|
||||
|
||||
func candidateUsesPerModelCooldown(cfg *config.Config, candidate providers.FallbackCandidate) bool {
|
||||
if cfg == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
candidateKey := providers.ModelKey(candidate.Provider, candidate.Model)
|
||||
strategyLookup := make(map[string]string, len(cfg.ModelList))
|
||||
for i := range cfg.ModelList {
|
||||
ref := providers.ParseModelRef(cfg.ModelList[i].Model, "openai")
|
||||
if ref == nil {
|
||||
continue
|
||||
}
|
||||
if providers.ModelKey(ref.Provider, ref.Model) != candidateKey {
|
||||
|
||||
strategy := config.NormalizeCooldownStrategy(cfg.ModelList[i].CooldownStrategy)
|
||||
if strategy == "" {
|
||||
continue
|
||||
}
|
||||
if normalizeCooldownStrategy(cfg.ModelList[i].CooldownStrategy) == "model" {
|
||||
return true
|
||||
}
|
||||
strategyLookup[providers.ModelKey(ref.Provider, ref.Model)] = strategy
|
||||
}
|
||||
|
||||
return false
|
||||
return strategyLookup
|
||||
}
|
||||
|
||||
func normalizeCooldownStrategy(strategy string) string {
|
||||
switch strings.ReplaceAll(strings.ToLower(strings.TrimSpace(strategy)), "_", "-") {
|
||||
case "model", "per-model":
|
||||
return "model"
|
||||
default:
|
||||
return "provider"
|
||||
func resolveCooldownKey(strategyLookup map[string]string, candidate providers.FallbackCandidate) string {
|
||||
if strings.TrimSpace(candidate.Provider) == "" {
|
||||
if strings.TrimSpace(candidate.Model) == "" {
|
||||
return ""
|
||||
}
|
||||
return providers.ModelKey(candidate.Provider, candidate.Model)
|
||||
}
|
||||
|
||||
candidateKey := providers.ModelKey(candidate.Provider, candidate.Model)
|
||||
if strategyLookup[candidateKey] == "model" {
|
||||
return candidateKey
|
||||
}
|
||||
|
||||
return candidate.Provider
|
||||
}
|
||||
|
||||
// resolveAgentWorkspace determines the workspace directory for an agent.
|
||||
|
|
|
|||
|
|
@ -211,6 +211,50 @@ func TestNewAgentInstance_ResolvePerModelCooldownKeys(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestNewAgentInstance_ResolveLightModelPerModelCooldownKeys(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
Model: "primary-model",
|
||||
Routing: &config.RoutingConfig{
|
||||
Enabled: true,
|
||||
LightModel: "light-model",
|
||||
Threshold: 0.5,
|
||||
},
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
{
|
||||
ModelName: "primary-model",
|
||||
Model: "litellm/openai/gpt-4o-mini",
|
||||
},
|
||||
{
|
||||
ModelName: "light-model",
|
||||
Model: " LiteLLM/OpenAI/GPT-4O ",
|
||||
CooldownStrategy: "per_model",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
|
||||
|
||||
if agent.Router == nil {
|
||||
t.Fatal("expected Router to be initialized for light model routing")
|
||||
}
|
||||
if len(agent.LightCandidates) != 1 {
|
||||
t.Fatalf("len(LightCandidates) = %d, want 1", len(agent.LightCandidates))
|
||||
}
|
||||
if got := agent.LightCandidates[0].Provider; got != "litellm" {
|
||||
t.Fatalf("light candidate provider = %q, want %q", got, "litellm")
|
||||
}
|
||||
if got := agent.LightCandidates[0].CooldownKey; got != "litellm/openai/gpt-4o" {
|
||||
t.Fatalf("light candidate cooldown key = %q, want %q", got, "litellm/openai/gpt-4o")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
mediaDir := media.TempDir()
|
||||
|
|
|
|||
|
|
@ -696,6 +696,7 @@ func (c *ModelConfig) APIKey() string {
|
|||
func (c *ModelConfig) IsVirtual() bool {
|
||||
return c.isVirtual
|
||||
}
|
||||
}
|
||||
|
||||
// Validate checks if the ModelConfig has all required fields.
|
||||
func (c *ModelConfig) Validate() error {
|
||||
|
|
@ -736,7 +737,6 @@ func (c *ModelConfig) SetAPIKey(value string) {
|
|||
c.APIKeys = append(c.APIKeys, NewSecureString(value))
|
||||
}
|
||||
}
|
||||
|
||||
type GatewayConfig struct {
|
||||
Host string `json:"host" env:"PICOCLAW_GATEWAY_HOST"`
|
||||
Port int `json:"port" env:"PICOCLAW_GATEWAY_PORT"`
|
||||
|
|
|
|||
|
|
@ -243,6 +243,29 @@ func TestModelConfig_Validate(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestNormalizeCooldownStrategy(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{input: "", want: "provider"},
|
||||
{input: "provider", want: "provider"},
|
||||
{input: "model", want: "model"},
|
||||
{input: "per-model", want: "model"},
|
||||
{input: "per_model", want: "model"},
|
||||
{input: " Per_Model ", want: "model"},
|
||||
{input: "backend", want: ""},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.input, func(t *testing.T) {
|
||||
if got := NormalizeCooldownStrategy(tt.input); got != tt.want {
|
||||
t.Fatalf("NormalizeCooldownStrategy(%q) = %q, want %q", tt.input, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfig_ValidateModelList(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
|
|
|||
|
|
@ -46,7 +46,13 @@ func (c FallbackCandidate) cooldownKey() string {
|
|||
if strings.TrimSpace(c.CooldownKey) != "" {
|
||||
return c.CooldownKey
|
||||
}
|
||||
return ModelKey(c.Provider, c.Model)
|
||||
if strings.TrimSpace(c.Provider) == "" {
|
||||
if strings.TrimSpace(c.Model) == "" {
|
||||
return ""
|
||||
}
|
||||
return ModelKey(c.Provider, c.Model)
|
||||
}
|
||||
return c.Provider
|
||||
}
|
||||
|
||||
// ResolveCandidates parses model config into a deduplicated candidate list.
|
||||
|
|
@ -127,8 +133,7 @@ func (fc *FallbackChain) Execute(
|
|||
return nil, context.Canceled
|
||||
}
|
||||
|
||||
// Check cooldown (per provider/model, not just provider).
|
||||
// This allows multi-key failover where different keys use different model names.
|
||||
// Check cooldown using the resolved provider- or model-scoped key.
|
||||
if !fc.cooldown.IsAvailable(cooldownKey) {
|
||||
remaining := fc.cooldown.CooldownRemaining(cooldownKey)
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue