fix(agent): keep multikey cooldown isolation in candidate keying
This commit is contained in:
parent
01d907c9c5
commit
83310d9933
3 changed files with 94 additions and 11 deletions
|
|
@ -6,6 +6,7 @@ import (
|
|||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
|
|
@ -239,7 +240,7 @@ func applyCooldownKeys(cfg *config.Config, candidates []providers.FallbackCandid
|
|||
copy(resolved, candidates)
|
||||
|
||||
for i := range resolved {
|
||||
resolved[i].CooldownKey = resolveCooldownKey(strategyLookup, resolved[i])
|
||||
resolved[i].CooldownKey = resolveCooldownKey(strategyLookup, resolved[i], resolved)
|
||||
}
|
||||
|
||||
return resolved
|
||||
|
|
@ -267,7 +268,11 @@ func buildCooldownStrategyLookup(cfg *config.Config) map[string]string {
|
|||
return strategyLookup
|
||||
}
|
||||
|
||||
func resolveCooldownKey(strategyLookup map[string]string, candidate providers.FallbackCandidate) string {
|
||||
func resolveCooldownKey(
|
||||
strategyLookup map[string]string,
|
||||
candidate providers.FallbackCandidate,
|
||||
candidates []providers.FallbackCandidate,
|
||||
) string {
|
||||
if strings.TrimSpace(candidate.Provider) == "" {
|
||||
if strings.TrimSpace(candidate.Model) == "" {
|
||||
return ""
|
||||
|
|
@ -276,6 +281,9 @@ func resolveCooldownKey(strategyLookup map[string]string, candidate providers.Fa
|
|||
}
|
||||
|
||||
candidateKey := providers.ModelKey(candidate.Provider, candidate.Model)
|
||||
if belongsToMultiKeySet(candidate, candidates) {
|
||||
return candidateKey
|
||||
}
|
||||
if strategyLookup[candidateKey] == "model" {
|
||||
return candidateKey
|
||||
}
|
||||
|
|
@ -283,6 +291,46 @@ func resolveCooldownKey(strategyLookup map[string]string, candidate providers.Fa
|
|||
return candidate.Provider
|
||||
}
|
||||
|
||||
func belongsToMultiKeySet(
|
||||
candidate providers.FallbackCandidate,
|
||||
candidates []providers.FallbackCandidate,
|
||||
) bool {
|
||||
provider := providers.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 providers.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
|
||||
}
|
||||
|
||||
// resolveAgentWorkspace determines the workspace directory for an agent.
|
||||
func resolveAgentWorkspace(agentCfg *config.AgentConfig, defaults *config.AgentDefaults) string {
|
||||
if agentCfg != nil && strings.TrimSpace(agentCfg.Workspace) != "" {
|
||||
|
|
|
|||
|
|
@ -172,11 +172,11 @@ func TestNewAgentInstance_ResolvePerModelCooldownKeys(t *testing.T) {
|
|||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
Model: "router-a",
|
||||
ModelName: "router-a",
|
||||
ModelFallbacks: []string{"router-b", "shared-provider"},
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
ModelList: []*config.ModelConfig{
|
||||
{
|
||||
ModelName: "router-a",
|
||||
Model: "litellm/openai/gpt-4o-mini",
|
||||
|
|
@ -218,7 +218,7 @@ func TestNewAgentInstance_ResolveLightModelPerModelCooldownKeys(t *testing.T) {
|
|||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
Model: "primary-model",
|
||||
ModelName: "primary-model",
|
||||
Routing: &config.RoutingConfig{
|
||||
Enabled: true,
|
||||
LightModel: "light-model",
|
||||
|
|
@ -226,7 +226,7 @@ func TestNewAgentInstance_ResolveLightModelPerModelCooldownKeys(t *testing.T) {
|
|||
},
|
||||
},
|
||||
},
|
||||
ModelList: []config.ModelConfig{
|
||||
ModelList: []*config.ModelConfig{
|
||||
{
|
||||
ModelName: "primary-model",
|
||||
Model: "litellm/openai/gpt-4o-mini",
|
||||
|
|
@ -255,6 +255,41 @@ func TestNewAgentInstance_ResolveLightModelPerModelCooldownKeys(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestNewAgentInstance_ResolveMultiKeyCooldownKeys(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
ModelName: "glm-main",
|
||||
ModelFallbacks: []string{"glm-replica"},
|
||||
},
|
||||
},
|
||||
ModelList: []*config.ModelConfig{
|
||||
{
|
||||
ModelName: "glm-main",
|
||||
Model: "zhipu/glm-4.7",
|
||||
},
|
||||
{
|
||||
ModelName: "glm-replica",
|
||||
Model: "zhipu/glm-4.7__key_1",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
|
||||
if len(agent.Candidates) != 2 {
|
||||
t.Fatalf("len(Candidates) = %d, want 2", len(agent.Candidates))
|
||||
}
|
||||
if got := agent.Candidates[0].CooldownKey; got != "zhipu/glm-4.7" {
|
||||
t.Fatalf("candidate[0] cooldown key = %q, want %q", got, "zhipu/glm-4.7")
|
||||
}
|
||||
if got := agent.Candidates[1].CooldownKey; got != "zhipu/glm-4.7__key_1" {
|
||||
t.Fatalf("candidate[1] cooldown key = %q, want %q", got, "zhipu/glm-4.7__key_1")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
mediaDir := media.TempDir()
|
||||
|
|
|
|||
|
|
@ -667,10 +667,10 @@ type ModelConfig struct {
|
|||
Workspace string `json:"workspace,omitempty"` // Workspace path for CLI-based providers
|
||||
|
||||
// Optional optimizations
|
||||
RPM int `json:"rpm,omitempty"` // Requests per minute limit
|
||||
MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens")
|
||||
RequestTimeout int `json:"request_timeout,omitempty"`
|
||||
ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive
|
||||
RPM int `json:"rpm,omitempty"` // Requests per minute limit
|
||||
MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens")
|
||||
RequestTimeout int `json:"request_timeout,omitempty"`
|
||||
ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive
|
||||
// CooldownStrategy controls the scope of cooldown tracking for this model.
|
||||
// - "provider" (default): cooldown is shared across all models in the provider
|
||||
// - "model" or "per-model": cooldown is isolated to this specific model
|
||||
|
|
@ -696,7 +696,6 @@ 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 {
|
||||
|
|
@ -737,6 +736,7 @@ 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"`
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue