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
|
// spawn_status which are added below — preventing recursive
|
||||||
// subagent spawning.
|
// subagent spawning.
|
||||||
subagentManager.SetTools(agent.Tools.Clone())
|
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 {
|
if spawnEnabled {
|
||||||
spawnTool := tools.NewSpawnTool(subagentManager)
|
spawnTool := tools.NewSpawnTool(subagentManager)
|
||||||
currentAgentID := agentID
|
currentAgentID := agentID
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,7 @@ import (
|
||||||
"context"
|
"context"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
|
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)
|
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
|
mu sync.RWMutex
|
||||||
provider providers.LLMProvider
|
provider providers.LLMProvider
|
||||||
defaultModel string
|
defaultModel string
|
||||||
|
resolveModel func(agentID string) (string, bool)
|
||||||
workspace string
|
workspace string
|
||||||
tools *ToolRegistry
|
tools *ToolRegistry
|
||||||
maxIterations int
|
maxIterations int
|
||||||
|
|
@ -61,6 +62,13 @@ func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) {
|
||||||
sm.hasTemperature = true
|
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.
|
// 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) {
|
||||||
|
|
@ -144,8 +152,16 @@ After completing the task, provide a clear summary of what was done.`
|
||||||
temperature := sm.temperature
|
temperature := sm.temperature
|
||||||
hasMaxTokens := sm.hasMaxTokens
|
hasMaxTokens := sm.hasMaxTokens
|
||||||
hasTemperature := sm.hasTemperature
|
hasTemperature := sm.hasTemperature
|
||||||
|
resolveModel := sm.resolveModel
|
||||||
sm.mu.RUnlock()
|
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
|
var llmOptions map[string]any
|
||||||
if hasMaxTokens || hasTemperature {
|
if hasMaxTokens || hasTemperature {
|
||||||
llmOptions = map[string]any{}
|
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{
|
loopResult, err := RunToolLoop(ctx, ToolLoopConfig{
|
||||||
Provider: sm.provider,
|
Provider: sm.provider,
|
||||||
Model: sm.defaultModel,
|
Model: model,
|
||||||
Tools: tools,
|
Tools: tools,
|
||||||
MaxIterations: maxIter,
|
MaxIterations: maxIter,
|
||||||
LLMOptions: llmOptions,
|
LLMOptions: llmOptions,
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,7 @@ import (
|
||||||
"context"
|
"context"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/providers"
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
)
|
)
|
||||||
|
|
@ -11,6 +12,7 @@ import (
|
||||||
// MockLLMProvider is a test implementation of LLMProvider
|
// MockLLMProvider is a test implementation of LLMProvider
|
||||||
type MockLLMProvider struct {
|
type MockLLMProvider struct {
|
||||||
lastOptions map[string]any
|
lastOptions map[string]any
|
||||||
|
lastModel string
|
||||||
}
|
}
|
||||||
|
|
||||||
func (m *MockLLMProvider) Chat(
|
func (m *MockLLMProvider) Chat(
|
||||||
|
|
@ -21,6 +23,7 @@ func (m *MockLLMProvider) Chat(
|
||||||
options map[string]any,
|
options map[string]any,
|
||||||
) (*providers.LLMResponse, error) {
|
) (*providers.LLMResponse, error) {
|
||||||
m.lastOptions = options
|
m.lastOptions = options
|
||||||
|
m.lastModel = model
|
||||||
// Find the last user message to generate a response
|
// Find the last user message to generate a response
|
||||||
for i := len(messages) - 1; i >= 0; i-- {
|
for i := len(messages) - 1; i >= 0; i-- {
|
||||||
if messages[i].Role == "user" {
|
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
|
// TestSubagentTool_Name verifies tool name
|
||||||
func TestSubagentTool_Name(t *testing.T) {
|
func TestSubagentTool_Name(t *testing.T) {
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue