feat: add ResolveModel to SubagentManager with 3-level fallback
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
46201fb679
commit
2e2b2178a3
2 changed files with 102 additions and 2 deletions
|
|
@ -14,6 +14,7 @@ type SubagentTask struct {
|
||||||
ID string
|
ID string
|
||||||
Task string
|
Task string
|
||||||
Label string
|
Label string
|
||||||
|
Model string // model_name override for this task
|
||||||
AgentID string
|
AgentID string
|
||||||
OriginChannel string
|
OriginChannel string
|
||||||
OriginChatID string
|
OriginChatID string
|
||||||
|
|
@ -27,6 +28,8 @@ type SubagentManager struct {
|
||||||
mu sync.RWMutex
|
mu sync.RWMutex
|
||||||
provider providers.LLMProvider
|
provider providers.LLMProvider
|
||||||
defaultModel string
|
defaultModel string
|
||||||
|
subagentDefaultModel string // from SubagentsConfig.Model
|
||||||
|
modelValidator func(string) bool // validates model_name exists
|
||||||
bus *bus.MessageBus
|
bus *bus.MessageBus
|
||||||
workspace string
|
workspace string
|
||||||
tools *ToolRegistry
|
tools *ToolRegistry
|
||||||
|
|
@ -65,6 +68,41 @@ func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) {
|
||||||
sm.hasTemperature = true
|
sm.hasTemperature = true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetSubagentDefaultModel sets the default model for subagents from config.
|
||||||
|
func (sm *SubagentManager) SetSubagentDefaultModel(model string) {
|
||||||
|
sm.mu.Lock()
|
||||||
|
defer sm.mu.Unlock()
|
||||||
|
sm.subagentDefaultModel = model
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetModelValidator sets the function used to validate model names.
|
||||||
|
func (sm *SubagentManager) SetModelValidator(validator func(string) bool) {
|
||||||
|
sm.mu.Lock()
|
||||||
|
defer sm.mu.Unlock()
|
||||||
|
sm.modelValidator = validator
|
||||||
|
}
|
||||||
|
|
||||||
|
// ResolveModel resolves the model to use for a subagent task.
|
||||||
|
// Priority: requested > subagentDefaultModel > defaultModel (parent).
|
||||||
|
func (sm *SubagentManager) ResolveModel(requested string) (string, error) {
|
||||||
|
sm.mu.RLock()
|
||||||
|
validator := sm.modelValidator
|
||||||
|
subDefault := sm.subagentDefaultModel
|
||||||
|
parentModel := sm.defaultModel
|
||||||
|
sm.mu.RUnlock()
|
||||||
|
|
||||||
|
if requested != "" {
|
||||||
|
if validator != nil && !validator(requested) {
|
||||||
|
return "", fmt.Errorf("model %q not found in model_list", requested)
|
||||||
|
}
|
||||||
|
return requested, nil
|
||||||
|
}
|
||||||
|
if subDefault != "" {
|
||||||
|
return subDefault, nil
|
||||||
|
}
|
||||||
|
return parentModel, nil
|
||||||
|
}
|
||||||
|
|
||||||
// SetTools sets the tool registry for subagent execution.
|
// SetTools sets the tool registry for subagent execution.
|
||||||
// If not set, subagent will have access to the provided tools.
|
// If not set, subagent will have access to the provided tools.
|
||||||
func (sm *SubagentManager) SetTools(tools *ToolRegistry) {
|
func (sm *SubagentManager) SetTools(tools *ToolRegistry) {
|
||||||
|
|
|
||||||
|
|
@ -300,6 +300,68 @@ func TestSubagentTool_Execute_ContextPassing(t *testing.T) {
|
||||||
// but execution success indicates context was handled properly
|
// but execution success indicates context was handled properly
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestSubagentManager_ResolveModel(t *testing.T) {
|
||||||
|
provider := &MockLLMProvider{}
|
||||||
|
alwaysValid := func(name string) bool { return true }
|
||||||
|
|
||||||
|
t.Run("uses requested model when valid", func(t *testing.T) {
|
||||||
|
manager := NewSubagentManager(provider, "parent-model", "/tmp/test", nil)
|
||||||
|
manager.SetModelValidator(alwaysValid)
|
||||||
|
model, err := manager.ResolveModel("requested-model")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if model != "requested-model" {
|
||||||
|
t.Errorf("got %q, want %q", model, "requested-model")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("returns error for invalid requested model", func(t *testing.T) {
|
||||||
|
manager := NewSubagentManager(provider, "parent-model", "/tmp/test", nil)
|
||||||
|
manager.SetModelValidator(func(name string) bool { return name != "bad-model" })
|
||||||
|
_, err := manager.ResolveModel("bad-model")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for invalid model")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("falls back to subagent default model", func(t *testing.T) {
|
||||||
|
manager := NewSubagentManager(provider, "parent-model", "/tmp/test", nil)
|
||||||
|
manager.SetSubagentDefaultModel("subagent-default")
|
||||||
|
manager.SetModelValidator(alwaysValid)
|
||||||
|
model, err := manager.ResolveModel("")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if model != "subagent-default" {
|
||||||
|
t.Errorf("got %q, want %q", model, "subagent-default")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("falls back to parent model", func(t *testing.T) {
|
||||||
|
manager := NewSubagentManager(provider, "parent-model", "/tmp/test", nil)
|
||||||
|
manager.SetModelValidator(alwaysValid)
|
||||||
|
model, err := manager.ResolveModel("")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if model != "parent-model" {
|
||||||
|
t.Errorf("got %q, want %q", model, "parent-model")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("works without validator set (no validation)", func(t *testing.T) {
|
||||||
|
manager := NewSubagentManager(provider, "parent-model", "/tmp/test", nil)
|
||||||
|
model, err := manager.ResolveModel("any-model")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if model != "any-model" {
|
||||||
|
t.Errorf("got %q, want %q", model, "any-model")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
// TestSubagentTool_ForUserTruncation verifies long content is truncated for user
|
// TestSubagentTool_ForUserTruncation verifies long content is truncated for user
|
||||||
func TestSubagentTool_ForUserTruncation(t *testing.T) {
|
func TestSubagentTool_ForUserTruncation(t *testing.T) {
|
||||||
// Create a mock provider that returns very long content
|
// Create a mock provider that returns very long content
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue