fix: spawn tool now respects target agent's model config
Fixes #1322 Previously, when spawn tool was called with agent_id parameter, the subagent would use the caller agent's model instead of the target agent's configured model. Changes: - Add modelResolver field to SubagentManager - Add SetModelResolver() method to inject model resolution function - In runTask(), use target agent's model when agent_id is specified - Add test case for model resolver functionality Test Plan: - All existing tests pass - New TestSubagentManager_ModelResolver test passes
This commit is contained in:
parent
30584f04cb
commit
cdbd09bf2f
3 changed files with 66 additions and 1 deletions
|
|
@ -224,6 +224,13 @@ func registerSharedTools(
|
|||
if cfg.Tools.IsToolEnabled("subagent") {
|
||||
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Workspace)
|
||||
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
|
||||
// Set model resolver so spawn can use target agent's model
|
||||
subagentManager.SetModelResolver(func(targetAgentID string) string {
|
||||
if targetAgent, ok := registry.GetAgent(targetAgentID); ok {
|
||||
return targetAgent.Model
|
||||
}
|
||||
return ""
|
||||
})
|
||||
spawnTool := tools.NewSpawnTool(subagentManager)
|
||||
currentAgentID := agentID
|
||||
spawnTool.SetAllowlistChecker(func(targetAgentID string) bool {
|
||||
|
|
|
|||
|
|
@ -77,3 +77,38 @@ func TestSpawnTool_Execute_NilManager(t *testing.T) {
|
|||
t.Errorf("Error message should mention manager not configured, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubagentManager_ModelResolver(t *testing.T) {
|
||||
provider := &MockLLMProvider{}
|
||||
manager := NewSubagentManager(provider, "default-model", "/tmp/test")
|
||||
|
||||
// Set up model resolver
|
||||
resolvedAgentID := ""
|
||||
manager.SetModelResolver(func(agentID string) string {
|
||||
resolvedAgentID = agentID
|
||||
if agentID == "premium-agent" {
|
||||
return "gpt-4"
|
||||
}
|
||||
return ""
|
||||
})
|
||||
|
||||
// Verify resolver is set
|
||||
if manager.modelResolver == nil {
|
||||
t.Fatal("Model resolver should be set")
|
||||
}
|
||||
|
||||
// Test resolver is called with correct agent ID
|
||||
result := manager.modelResolver("premium-agent")
|
||||
if resolvedAgentID != "premium-agent" {
|
||||
t.Errorf("Expected resolver to be called with 'premium-agent', got '%s'", resolvedAgentID)
|
||||
}
|
||||
if result != "gpt-4" {
|
||||
t.Errorf("Expected 'gpt-4', got '%s'", result)
|
||||
}
|
||||
|
||||
// Test fallback for unknown agent
|
||||
result = manager.modelResolver("unknown-agent")
|
||||
if result != "" {
|
||||
t.Errorf("Expected empty string for unknown agent, got '%s'", result)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -34,6 +34,9 @@ type SubagentManager struct {
|
|||
hasMaxTokens bool
|
||||
hasTemperature bool
|
||||
nextID int
|
||||
// modelResolver resolves agentID to model name.
|
||||
// Returns empty string if agent not found (falls back to defaultModel).
|
||||
modelResolver func(agentID string) string
|
||||
}
|
||||
|
||||
func NewSubagentManager(
|
||||
|
|
@ -61,6 +64,16 @@ func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) {
|
|||
sm.hasTemperature = true
|
||||
}
|
||||
|
||||
// SetModelResolver sets a function to resolve agentID to model name.
|
||||
// When spawn is called with agent_id, this resolver is used to get the
|
||||
// target agent's configured model. If the resolver returns empty string
|
||||
// or is not set, falls back to the defaultModel.
|
||||
func (sm *SubagentManager) SetModelResolver(resolver func(agentID string) string) {
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
sm.modelResolver = resolver
|
||||
}
|
||||
|
||||
// SetTools sets the tool registry for subagent execution.
|
||||
// If not set, subagent will have access to the provided tools.
|
||||
func (sm *SubagentManager) SetTools(tools *ToolRegistry) {
|
||||
|
|
@ -147,8 +160,18 @@ After completing the task, provide a clear summary of what was done.`
|
|||
temperature := sm.temperature
|
||||
hasMaxTokens := sm.hasMaxTokens
|
||||
hasTemperature := sm.hasTemperature
|
||||
modelResolver := sm.modelResolver
|
||||
defaultModel := sm.defaultModel
|
||||
sm.mu.RUnlock()
|
||||
|
||||
// Resolve target agent model if agentID is specified
|
||||
model := defaultModel
|
||||
if task.AgentID != "" && modelResolver != nil {
|
||||
if resolvedModel := modelResolver(task.AgentID); resolvedModel != "" {
|
||||
model = resolvedModel
|
||||
}
|
||||
}
|
||||
|
||||
var llmOptions map[string]any
|
||||
if hasMaxTokens || hasTemperature {
|
||||
llmOptions = map[string]any{}
|
||||
|
|
@ -162,7 +185,7 @@ After completing the task, provide a clear summary of what was done.`
|
|||
|
||||
loopResult, err := RunToolLoop(ctx, ToolLoopConfig{
|
||||
Provider: sm.provider,
|
||||
Model: sm.defaultModel,
|
||||
Model: model,
|
||||
Tools: tools,
|
||||
MaxIterations: maxIter,
|
||||
LLMOptions: llmOptions,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue