fix(spawn): use target agent model
This commit is contained in:
parent
8a188cf7fc
commit
4d6ac5fc16
4 changed files with 104 additions and 1 deletions
|
|
@ -244,6 +244,13 @@ func registerSharedTools(
|
|||
// spawn_status which are added below — preventing recursive
|
||||
// subagent spawning.
|
||||
subagentManager.SetTools(agent.Tools.Clone())
|
||||
subagentManager.SetAgentModelResolver(func(targetAgentID string) (string, bool) {
|
||||
targetAgent, ok := registry.GetAgent(targetAgentID)
|
||||
if !ok || targetAgent == nil || targetAgent.Model == "" {
|
||||
return "", false
|
||||
}
|
||||
return targetAgent.Model, true
|
||||
})
|
||||
if spawnEnabled {
|
||||
spawnTool := tools.NewSpawnTool(subagentManager)
|
||||
currentAgentID := agentID
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
|
||||
|
|
@ -77,3 +78,39 @@ func TestSpawnTool_Execute_NilManager(t *testing.T) {
|
|||
t.Errorf("Error message should mention manager not configured, got: %s", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSpawnTool_ExecuteAsync_UsesTargetAgentModel(t *testing.T) {
|
||||
provider := &MockLLMProvider{}
|
||||
manager := NewSubagentManager(provider, "caller-model", "/tmp/test")
|
||||
manager.SetAgentModelResolver(func(agentID string) (string, bool) {
|
||||
if agentID == "analyst" {
|
||||
return "premium-model", true
|
||||
}
|
||||
return "", false
|
||||
})
|
||||
tool := NewSpawnTool(manager)
|
||||
|
||||
done := make(chan *ToolResult, 1)
|
||||
result := tool.ExecuteAsync(context.Background(), map[string]any{
|
||||
"task": "Investigate the issue",
|
||||
"agent_id": "analyst",
|
||||
}, func(_ context.Context, result *ToolResult) {
|
||||
done <- result
|
||||
})
|
||||
if result == nil || result.IsError {
|
||||
t.Fatalf("expected async spawn success, got: %+v", result)
|
||||
}
|
||||
|
||||
select {
|
||||
case callbackResult := <-done:
|
||||
if callbackResult == nil || callbackResult.IsError {
|
||||
t.Fatalf("expected successful callback result, got: %+v", callbackResult)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("timed out waiting for async spawn callback")
|
||||
}
|
||||
|
||||
if provider.lastModel != "premium-model" {
|
||||
t.Fatalf("model = %q, want %q", provider.lastModel, "premium-model")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ type SubagentManager struct {
|
|||
mu sync.RWMutex
|
||||
provider providers.LLMProvider
|
||||
defaultModel string
|
||||
resolveModel func(agentID string) (string, bool)
|
||||
workspace string
|
||||
tools *ToolRegistry
|
||||
maxIterations int
|
||||
|
|
@ -61,6 +62,13 @@ func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) {
|
|||
sm.hasTemperature = true
|
||||
}
|
||||
|
||||
// SetAgentModelResolver resolves the effective model for a target agent ID.
|
||||
func (sm *SubagentManager) SetAgentModelResolver(resolve func(agentID string) (string, bool)) {
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
sm.resolveModel = resolve
|
||||
}
|
||||
|
||||
// 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) {
|
||||
|
|
@ -144,8 +152,16 @@ After completing the task, provide a clear summary of what was done.`
|
|||
temperature := sm.temperature
|
||||
hasMaxTokens := sm.hasMaxTokens
|
||||
hasTemperature := sm.hasTemperature
|
||||
resolveModel := sm.resolveModel
|
||||
sm.mu.RUnlock()
|
||||
|
||||
model := sm.defaultModel
|
||||
if task.AgentID != "" && resolveModel != nil {
|
||||
if resolvedModel, ok := resolveModel(task.AgentID); ok && resolvedModel != "" {
|
||||
model = resolvedModel
|
||||
}
|
||||
}
|
||||
|
||||
var llmOptions map[string]any
|
||||
if hasMaxTokens || hasTemperature {
|
||||
llmOptions = map[string]any{}
|
||||
|
|
@ -159,7 +175,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,
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
|
@ -11,6 +12,7 @@ import (
|
|||
// MockLLMProvider is a test implementation of LLMProvider
|
||||
type MockLLMProvider struct {
|
||||
lastOptions map[string]any
|
||||
lastModel string
|
||||
}
|
||||
|
||||
func (m *MockLLMProvider) Chat(
|
||||
|
|
@ -21,6 +23,7 @@ func (m *MockLLMProvider) Chat(
|
|||
options map[string]any,
|
||||
) (*providers.LLMResponse, error) {
|
||||
m.lastOptions = options
|
||||
m.lastModel = model
|
||||
// Find the last user message to generate a response
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
if messages[i].Role == "user" {
|
||||
|
|
@ -69,6 +72,46 @@ func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestSubagentManager_RunTask_UsesResolvedTargetAgentModel(t *testing.T) {
|
||||
provider := &MockLLMProvider{}
|
||||
manager := NewSubagentManager(provider, "caller-model", "/tmp/test")
|
||||
manager.SetAgentModelResolver(func(agentID string) (string, bool) {
|
||||
if agentID == "analyst" {
|
||||
return "premium-model", true
|
||||
}
|
||||
return "", false
|
||||
})
|
||||
|
||||
done := make(chan *ToolResult, 1)
|
||||
_, err := manager.Spawn(
|
||||
context.Background(),
|
||||
"Investigate the issue",
|
||||
"analysis",
|
||||
"analyst",
|
||||
"cli",
|
||||
"direct",
|
||||
func(_ context.Context, result *ToolResult) {
|
||||
done <- result
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("Spawn() error: %v", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case result := <-done:
|
||||
if result == nil || result.IsError {
|
||||
t.Fatalf("expected successful async result, got: %+v", result)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("timed out waiting for subagent completion")
|
||||
}
|
||||
|
||||
if provider.lastModel != "premium-model" {
|
||||
t.Fatalf("model = %q, want %q", provider.lastModel, "premium-model")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSubagentTool_Name verifies tool name
|
||||
func TestSubagentTool_Name(t *testing.T) {
|
||||
provider := &MockLLMProvider{}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue