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"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"regexp"
|
"regexp"
|
||||||
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
|
|
@ -239,7 +240,7 @@ func applyCooldownKeys(cfg *config.Config, candidates []providers.FallbackCandid
|
||||||
copy(resolved, candidates)
|
copy(resolved, candidates)
|
||||||
|
|
||||||
for i := range resolved {
|
for i := range resolved {
|
||||||
resolved[i].CooldownKey = resolveCooldownKey(strategyLookup, resolved[i])
|
resolved[i].CooldownKey = resolveCooldownKey(strategyLookup, resolved[i], resolved)
|
||||||
}
|
}
|
||||||
|
|
||||||
return resolved
|
return resolved
|
||||||
|
|
@ -267,7 +268,11 @@ func buildCooldownStrategyLookup(cfg *config.Config) map[string]string {
|
||||||
return strategyLookup
|
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.Provider) == "" {
|
||||||
if strings.TrimSpace(candidate.Model) == "" {
|
if strings.TrimSpace(candidate.Model) == "" {
|
||||||
return ""
|
return ""
|
||||||
|
|
@ -276,6 +281,9 @@ func resolveCooldownKey(strategyLookup map[string]string, candidate providers.Fa
|
||||||
}
|
}
|
||||||
|
|
||||||
candidateKey := providers.ModelKey(candidate.Provider, candidate.Model)
|
candidateKey := providers.ModelKey(candidate.Provider, candidate.Model)
|
||||||
|
if belongsToMultiKeySet(candidate, candidates) {
|
||||||
|
return candidateKey
|
||||||
|
}
|
||||||
if strategyLookup[candidateKey] == "model" {
|
if strategyLookup[candidateKey] == "model" {
|
||||||
return candidateKey
|
return candidateKey
|
||||||
}
|
}
|
||||||
|
|
@ -283,6 +291,46 @@ func resolveCooldownKey(strategyLookup map[string]string, candidate providers.Fa
|
||||||
return candidate.Provider
|
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.
|
// resolveAgentWorkspace determines the workspace directory for an agent.
|
||||||
func resolveAgentWorkspace(agentCfg *config.AgentConfig, defaults *config.AgentDefaults) string {
|
func resolveAgentWorkspace(agentCfg *config.AgentConfig, defaults *config.AgentDefaults) string {
|
||||||
if agentCfg != nil && strings.TrimSpace(agentCfg.Workspace) != "" {
|
if agentCfg != nil && strings.TrimSpace(agentCfg.Workspace) != "" {
|
||||||
|
|
|
||||||
|
|
@ -172,11 +172,11 @@ func TestNewAgentInstance_ResolvePerModelCooldownKeys(t *testing.T) {
|
||||||
Agents: config.AgentsConfig{
|
Agents: config.AgentsConfig{
|
||||||
Defaults: config.AgentDefaults{
|
Defaults: config.AgentDefaults{
|
||||||
Workspace: tmpDir,
|
Workspace: tmpDir,
|
||||||
Model: "router-a",
|
ModelName: "router-a",
|
||||||
ModelFallbacks: []string{"router-b", "shared-provider"},
|
ModelFallbacks: []string{"router-b", "shared-provider"},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
ModelList: []config.ModelConfig{
|
ModelList: []*config.ModelConfig{
|
||||||
{
|
{
|
||||||
ModelName: "router-a",
|
ModelName: "router-a",
|
||||||
Model: "litellm/openai/gpt-4o-mini",
|
Model: "litellm/openai/gpt-4o-mini",
|
||||||
|
|
@ -218,7 +218,7 @@ func TestNewAgentInstance_ResolveLightModelPerModelCooldownKeys(t *testing.T) {
|
||||||
Agents: config.AgentsConfig{
|
Agents: config.AgentsConfig{
|
||||||
Defaults: config.AgentDefaults{
|
Defaults: config.AgentDefaults{
|
||||||
Workspace: tmpDir,
|
Workspace: tmpDir,
|
||||||
Model: "primary-model",
|
ModelName: "primary-model",
|
||||||
Routing: &config.RoutingConfig{
|
Routing: &config.RoutingConfig{
|
||||||
Enabled: true,
|
Enabled: true,
|
||||||
LightModel: "light-model",
|
LightModel: "light-model",
|
||||||
|
|
@ -226,7 +226,7 @@ func TestNewAgentInstance_ResolveLightModelPerModelCooldownKeys(t *testing.T) {
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
ModelList: []config.ModelConfig{
|
ModelList: []*config.ModelConfig{
|
||||||
{
|
{
|
||||||
ModelName: "primary-model",
|
ModelName: "primary-model",
|
||||||
Model: "litellm/openai/gpt-4o-mini",
|
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) {
|
func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
|
||||||
workspace := t.TempDir()
|
workspace := t.TempDir()
|
||||||
mediaDir := media.TempDir()
|
mediaDir := media.TempDir()
|
||||||
|
|
|
||||||
|
|
@ -667,10 +667,10 @@ type ModelConfig struct {
|
||||||
Workspace string `json:"workspace,omitempty"` // Workspace path for CLI-based providers
|
Workspace string `json:"workspace,omitempty"` // Workspace path for CLI-based providers
|
||||||
|
|
||||||
// Optional optimizations
|
// Optional optimizations
|
||||||
RPM int `json:"rpm,omitempty"` // Requests per minute limit
|
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")
|
MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens")
|
||||||
RequestTimeout int `json:"request_timeout,omitempty"`
|
RequestTimeout int `json:"request_timeout,omitempty"`
|
||||||
ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive
|
ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive
|
||||||
// CooldownStrategy controls the scope of cooldown tracking for this model.
|
// CooldownStrategy controls the scope of cooldown tracking for this model.
|
||||||
// - "provider" (default): cooldown is shared across all models in the provider
|
// - "provider" (default): cooldown is shared across all models in the provider
|
||||||
// - "model" or "per-model": cooldown is isolated to this specific model
|
// - "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 {
|
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 {
|
||||||
|
|
@ -737,6 +736,7 @@ 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"`
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue