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
|
return candidates
|
||||||
}
|
}
|
||||||
|
|
||||||
|
strategyLookup := buildCooldownStrategyLookup(cfg)
|
||||||
resolved := make([]providers.FallbackCandidate, len(candidates))
|
resolved := make([]providers.FallbackCandidate, len(candidates))
|
||||||
copy(resolved, candidates)
|
copy(resolved, candidates)
|
||||||
|
|
||||||
for i := range resolved {
|
for i := range resolved {
|
||||||
resolved[i].CooldownKey = resolveCooldownKey(cfg, resolved[i])
|
resolved[i].CooldownKey = resolveCooldownKey(strategyLookup, resolved[i])
|
||||||
}
|
}
|
||||||
|
|
||||||
return resolved
|
return resolved
|
||||||
}
|
}
|
||||||
|
|
||||||
func resolveCooldownKey(cfg *config.Config, candidate providers.FallbackCandidate) string {
|
func buildCooldownStrategyLookup(cfg *config.Config) map[string]string {
|
||||||
if candidate.Provider == "" {
|
if cfg == nil || len(cfg.ModelList) == 0 {
|
||||||
return ""
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
if cfg != nil && candidateUsesPerModelCooldown(cfg, candidate) {
|
strategyLookup := make(map[string]string, len(cfg.ModelList))
|
||||||
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)
|
|
||||||
for i := range cfg.ModelList {
|
for i := range cfg.ModelList {
|
||||||
ref := providers.ParseModelRef(cfg.ModelList[i].Model, "openai")
|
ref := providers.ParseModelRef(cfg.ModelList[i].Model, "openai")
|
||||||
if ref == nil {
|
if ref == nil {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if providers.ModelKey(ref.Provider, ref.Model) != candidateKey {
|
|
||||||
|
strategy := config.NormalizeCooldownStrategy(cfg.ModelList[i].CooldownStrategy)
|
||||||
|
if strategy == "" {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if normalizeCooldownStrategy(cfg.ModelList[i].CooldownStrategy) == "model" {
|
strategyLookup[providers.ModelKey(ref.Provider, ref.Model)] = strategy
|
||||||
return true
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return false
|
return strategyLookup
|
||||||
}
|
}
|
||||||
|
|
||||||
func normalizeCooldownStrategy(strategy string) string {
|
func resolveCooldownKey(strategyLookup map[string]string, candidate providers.FallbackCandidate) string {
|
||||||
switch strings.ReplaceAll(strings.ToLower(strings.TrimSpace(strategy)), "_", "-") {
|
if strings.TrimSpace(candidate.Provider) == "" {
|
||||||
case "model", "per-model":
|
if strings.TrimSpace(candidate.Model) == "" {
|
||||||
return "model"
|
return ""
|
||||||
default:
|
|
||||||
return "provider"
|
|
||||||
}
|
}
|
||||||
|
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.
|
// 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) {
|
func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
|
||||||
workspace := t.TempDir()
|
workspace := t.TempDir()
|
||||||
mediaDir := media.TempDir()
|
mediaDir := media.TempDir()
|
||||||
|
|
|
||||||
|
|
@ -696,6 +696,7 @@ func (c *ModelConfig) APIKey() string {
|
||||||
func (c *ModelConfig) IsVirtual() bool {
|
func (c *ModelConfig) IsVirtual() bool {
|
||||||
return c.isVirtual
|
return c.isVirtual
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Validate checks if the ModelConfig has all required fields.
|
// Validate checks if the ModelConfig has all required fields.
|
||||||
func (c *ModelConfig) Validate() error {
|
func (c *ModelConfig) Validate() error {
|
||||||
|
|
@ -736,7 +737,6 @@ func (c *ModelConfig) SetAPIKey(value string) {
|
||||||
c.APIKeys = append(c.APIKeys, NewSecureString(value))
|
c.APIKeys = append(c.APIKeys, NewSecureString(value))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
type GatewayConfig struct {
|
type GatewayConfig struct {
|
||||||
Host string `json:"host" env:"PICOCLAW_GATEWAY_HOST"`
|
Host string `json:"host" env:"PICOCLAW_GATEWAY_HOST"`
|
||||||
Port int `json:"port" env:"PICOCLAW_GATEWAY_PORT"`
|
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) {
|
func TestConfig_ValidateModelList(t *testing.T) {
|
||||||
tests := []struct {
|
tests := []struct {
|
||||||
name string
|
name string
|
||||||
|
|
|
||||||
|
|
@ -46,7 +46,13 @@ func (c FallbackCandidate) cooldownKey() string {
|
||||||
if strings.TrimSpace(c.CooldownKey) != "" {
|
if strings.TrimSpace(c.CooldownKey) != "" {
|
||||||
return c.CooldownKey
|
return c.CooldownKey
|
||||||
}
|
}
|
||||||
|
if strings.TrimSpace(c.Provider) == "" {
|
||||||
|
if strings.TrimSpace(c.Model) == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
return ModelKey(c.Provider, c.Model)
|
return ModelKey(c.Provider, c.Model)
|
||||||
|
}
|
||||||
|
return c.Provider
|
||||||
}
|
}
|
||||||
|
|
||||||
// ResolveCandidates parses model config into a deduplicated candidate list.
|
// ResolveCandidates parses model config into a deduplicated candidate list.
|
||||||
|
|
@ -127,8 +133,7 @@ func (fc *FallbackChain) Execute(
|
||||||
return nil, context.Canceled
|
return nil, context.Canceled
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check cooldown (per provider/model, not just provider).
|
// Check cooldown using the resolved provider- or model-scoped key.
|
||||||
// This allows multi-key failover where different keys use different model names.
|
|
||||||
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{
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue