From 7577f952cdb148dff5480ed65289dc462bb95110 Mon Sep 17 00:00:00 2001 From: Alix-007 <267018309+Alix-007@users.noreply.github.com> Date: Wed, 18 Mar 2026 10:36:40 +0800 Subject: [PATCH] refactor(provider): optimize cooldown strategy resolution --- pkg/agent/instance.go | 52 +++++++++++++++------------------ pkg/agent/instance_test.go | 44 ++++++++++++++++++++++++++++ pkg/config/config.go | 2 +- pkg/config/model_config_test.go | 23 +++++++++++++++ pkg/providers/fallback.go | 11 +++++-- 5 files changed, 100 insertions(+), 32 deletions(-) diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index a539b4885..c3453c45a 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -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. diff --git a/pkg/agent/instance_test.go b/pkg/agent/instance_test.go index 97bad2900..97c32e4f9 100644 --- a/pkg/agent/instance_test.go +++ b/pkg/agent/instance_test.go @@ -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() diff --git a/pkg/config/config.go b/pkg/config/config.go index 51c50421e..652651a26 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -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"` diff --git a/pkg/config/model_config_test.go b/pkg/config/model_config_test.go index a363e713c..a351cc39e 100644 --- a/pkg/config/model_config_test.go +++ b/pkg/config/model_config_test.go @@ -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 diff --git a/pkg/providers/fallback.go b/pkg/providers/fallback.go index fbeb9810d..8b3b94d94 100644 --- a/pkg/providers/fallback.go +++ b/pkg/providers/fallback.go @@ -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{