diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 65a1fe04d..097cb4efb 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -92,7 +92,7 @@ func NewAgentInstance( Primary: model, Fallbacks: fallbacks, } - resolveFromModelList := func(raw string) (string, bool) { + resolveFromModelList := func(raw string) (string, []string, bool) { ensureProtocol := func(model string) string { model = strings.TrimSpace(model) if model == "" { @@ -106,12 +106,12 @@ func NewAgentInstance( raw = strings.TrimSpace(raw) if raw == "" { - return "", false + return "", nil, false } if cfg != nil { if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" { - return ensureProtocol(mc.Model), true + return ensureProtocol(mc.Model), mc.Tags, true } for i := range cfg.ModelList { @@ -120,16 +120,16 @@ func NewAgentInstance( continue } if fullModel == raw { - return ensureProtocol(fullModel), true + return ensureProtocol(fullModel), cfg.ModelList[i].Tags, true } _, modelID := providers.ExtractProtocol(fullModel) if modelID == raw { - return ensureProtocol(fullModel), true + return ensureProtocol(fullModel), cfg.ModelList[i].Tags, true } } } - return "", false + return "", nil, false } candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 1b8654eb6..6e9d9d4c2 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -147,7 +147,7 @@ func registerSharedTools( agent.Tools.Register(tools.NewInstallSkillTool(registryMgr, agent.Workspace)) // Spawn tool with allowlist checker - subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Workspace, msgBus) + subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Candidates, agent.Workspace, msgBus) subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) spawnTool := tools.NewSpawnTool(subagentManager) currentAgentID := agentID diff --git a/pkg/config/config.go b/pkg/config/config.go index d84772d2b..80cd8253c 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -457,8 +457,9 @@ type OpenAIProviderConfig struct { // Default protocol is "openai" if no prefix is specified. type ModelConfig struct { // Required fields - ModelName string `json:"model_name"` // User-facing alias for the model - Model string `json:"model"` // Protocol/model-identifier (e.g., "openai/gpt-4o", "anthropic/claude-sonnet-4.6") + ModelName string `json:"model_name"` // User-facing alias for the model + Model string `json:"model"` // Protocol/model-identifier (e.g., "openai/gpt-4o", "anthropic/claude-sonnet-4.6") + Tags []string `json:"tags,omitempty"` // Model capability labels like 'vision' // HTTP-based providers APIBase string `json:"api_base,omitempty"` // API endpoint URL diff --git a/pkg/providers/fallback.go b/pkg/providers/fallback.go index 7ba563b66..2c0f52f21 100644 --- a/pkg/providers/fallback.go +++ b/pkg/providers/fallback.go @@ -16,6 +16,7 @@ type FallbackChain struct { type FallbackCandidate struct { Provider string Model string + Tags []string } // FallbackResult contains the successful response and metadata about all attempts. @@ -49,16 +50,19 @@ func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandid func ResolveCandidatesWithLookup( cfg ModelConfig, defaultProvider string, - lookup func(raw string) (resolved string, ok bool), + lookup func(raw string) (resolved string, tags []string, ok bool), ) []FallbackCandidate { seen := make(map[string]bool) var candidates []FallbackCandidate addCandidate := func(raw string) { candidateRaw := strings.TrimSpace(raw) + var modelTags []string + if lookup != nil { - if resolved, ok := lookup(candidateRaw); ok { + if resolved, tags, ok := lookup(candidateRaw); ok { candidateRaw = resolved + modelTags = tags } } @@ -74,6 +78,7 @@ func ResolveCandidatesWithLookup( candidates = append(candidates, FallbackCandidate{ Provider: ref.Provider, Model: ref.Model, + Tags: modelTags, }) } diff --git a/pkg/providers/fallback_test.go b/pkg/providers/fallback_test.go index 1783ebcb5..f9014db2d 100644 --- a/pkg/providers/fallback_test.go +++ b/pkg/providers/fallback_test.go @@ -459,11 +459,11 @@ func TestResolveCandidatesWithLookup_AliasResolvesToNestedModel(t *testing.T) { Fallbacks: nil, } - lookup := func(raw string) (string, bool) { + lookup := func(raw string) (string, []string, bool) { if raw == "step-3.5-flash" { - return "openrouter/stepfun/step-3.5-flash:free", true + return "openrouter/stepfun/step-3.5-flash:free", nil, true } - return "", false + return "", nil, false } candidates := ResolveCandidatesWithLookup(cfg, "", lookup) @@ -484,11 +484,11 @@ func TestResolveCandidatesWithLookup_DeduplicateAfterLookup(t *testing.T) { Fallbacks: []string{"openrouter/stepfun/step-3.5-flash:free"}, } - lookup := func(raw string) (string, bool) { + lookup := func(raw string) (string, []string, bool) { if raw == "step-3.5-flash" { - return "openrouter/stepfun/step-3.5-flash:free", true + return "openrouter/stepfun/step-3.5-flash:free", nil, true } - return "", false + return "", nil, false } candidates := ResolveCandidatesWithLookup(cfg, "", lookup) @@ -503,11 +503,11 @@ func TestResolveCandidatesWithLookup_AliasWithoutProtocolUsesDefaultProvider(t * Fallbacks: nil, } - lookup := func(raw string) (string, bool) { + lookup := func(raw string) (string, []string, bool) { if raw == "glm-5" { - return "glm-5", true + return "glm-5", nil, true } - return "", false + return "", nil, false } candidates := ResolveCandidatesWithLookup(cfg, "openai", lookup) diff --git a/pkg/tools/spawn_sub_agent.go b/pkg/tools/spawn_sub_agent.go index 7617f8a12..ecc71b69f 100644 --- a/pkg/tools/spawn_sub_agent.go +++ b/pkg/tools/spawn_sub_agent.go @@ -28,7 +28,13 @@ func (t *SpawnSubAgentTool) Name() string { } func (t *SpawnSubAgentTool) Description() string { - return "Directly delegate a specific task to a new, isolated sub-agent. You (the main agent) should autonomously determine the appropriate expert role and specific task based on the user's high-level request. It will execute independently and return the final result." + base := "Directly delegate a specific task to a new, isolated sub-agent. You (the main agent) should autonomously determine the appropriate expert role and specific task based on the user's high-level request. It will execute independently and return the final result." + if t.manager != nil { + if hint := t.manager.ModelCapabilityHint(); hint != "" { + return base + "\n\n" + hint + } + } + return base } func (t *SpawnSubAgentTool) Parameters() map[string]any { @@ -89,7 +95,11 @@ func (t *SpawnSubAgentTool) Execute(ctx context.Context, args map[string]any) *T // 2.1 Model Override (Heterogeneous Agents) if modelParam, ok := args["model"].(string); ok && strings.TrimSpace(modelParam) != "" { - config.Model = strings.TrimSpace(modelParam) + requestedModel := strings.TrimSpace(modelParam) + if !t.manager.IsModelAllowed(requestedModel) { + return ErrorResult(fmt.Sprintf("requested model '%s' is not in the allowed fallback candidates list for this agent workspace", requestedModel)).WithError(fmt.Errorf("model %s not allowed", requestedModel)) + } + config.Model = requestedModel } // Note: For MVP, we pass the current ToolRegistry unmodified. diff --git a/pkg/tools/spawn_test.go b/pkg/tools/spawn_test.go index 0646c82a9..b65a5e928 100644 --- a/pkg/tools/spawn_test.go +++ b/pkg/tools/spawn_test.go @@ -8,7 +8,7 @@ import ( func TestSpawnTool_Execute_EmptyTask(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSpawnTool(manager) ctx := context.Background() @@ -42,7 +42,7 @@ func TestSpawnTool_Execute_EmptyTask(t *testing.T) { func TestSpawnTool_Execute_ValidTask(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSpawnTool(manager) ctx := context.Background() diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 3e3551d19..03eeb8743 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -3,6 +3,7 @@ package tools import ( "context" "fmt" + "strings" "sync" "time" @@ -10,6 +11,26 @@ import ( "github.com/sipeed/picoclaw/pkg/providers" ) +// ModelTag constants define the recognized capability labels for models in config.json. +// These are set via `"tags": ["vision", "code"]` under each model in the model list. +const ( + ModelTagVision = "vision" // Supports image/screenshot input (multimodal) + ModelTagCode = "code" // Specialized for code generation and analysis + ModelTagFast = "fast" // Low-latency model, suited for lightweight tasks + ModelTagLongContext = "long-context" // Supports very long context windows (>100k tokens) + ModelTagReasoning = "reasoning" // Strong logical/math reasoning (e.g., o1, deepseek-r1) +) + +// modelTagDescriptions provides LLM-readable explanations of each known tag, +// injected at runtime into the tool description to guide model selection. +var modelTagDescriptions = map[string]string{ + ModelTagVision: "can analyze images and screenshots", + ModelTagCode: "specialized in code generation and debugging", + ModelTagFast: "fast and lightweight, ideal for simple or high-frequency tasks", + ModelTagLongContext: "handles very long inputs (>100k tokens)", + ModelTagReasoning: "excels at logical reasoning, math, and multi-step planning", +} + type SubagentTask struct { ID string Task string @@ -27,6 +48,7 @@ type SubagentManager struct { mu sync.RWMutex provider providers.LLMProvider defaultModel string + allowedModels []providers.FallbackCandidate bus *bus.MessageBus workspace string tools *ToolRegistry @@ -40,13 +62,16 @@ type SubagentManager struct { func NewSubagentManager( provider providers.LLMProvider, - defaultModel, workspace string, + defaultModel string, + candidates []providers.FallbackCandidate, + workspace string, bus *bus.MessageBus, ) *SubagentManager { return &SubagentManager{ tasks: make(map[string]*SubagentTask), provider: provider, defaultModel: defaultModel, + allowedModels: candidates, bus: bus, workspace: workspace, tools: NewToolRegistry(), @@ -55,6 +80,52 @@ func NewSubagentManager( } } +// IsModelAllowed checks if a specific requested model exists in the permitted candidates list. +func (sm *SubagentManager) IsModelAllowed(model string) bool { + // If the user requested the default model directly, that's automatically allowed + if model == sm.defaultModel { + return true + } + + // Otherwise, check against the resolved candidates (primary + fallbacks + explicitly configured) + for _, cand := range sm.allowedModels { + if cand.Model == model { + return true + } + } + return false +} + +// ModelCapabilityHint generates a human-readable summary of allowed models and their tags. +// This is injected into the coordinator's tool descriptions so the LLM can make better routing decisions. +func (sm *SubagentManager) ModelCapabilityHint() string { + if len(sm.allowedModels) == 0 { + return "" + } + + var modelLines []string + for _, cand := range sm.allowedModels { + if len(cand.Tags) == 0 { + modelLines = append(modelLines, fmt.Sprintf(" - %s (general purpose)", cand.Model)) + continue + } + var descs []string + for _, tag := range cand.Tags { + if desc, known := modelTagDescriptions[tag]; known { + descs = append(descs, fmt.Sprintf("%s (%s)", tag, desc)) + } else { + descs = append(descs, tag) + } + } + modelLines = append(modelLines, fmt.Sprintf(" - %s [%s]", cand.Model, strings.Join(descs, ", "))) + } + + hint := "When selecting a 'model' for sub-agents, use ONLY these configured models:\n" + hint += strings.Join(modelLines, "\n") + hint += "\nIf a task requires vision/image analysis, you MUST select a model with the 'vision' tag. If no suitable model is available, omit the 'model' field to use the default." + return hint +} + // SetLLMOptions sets max tokens and temperature for subagent LLM calls. func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) { sm.mu.Lock() diff --git a/pkg/tools/subagent_tool_test.go b/pkg/tools/subagent_tool_test.go index 59bfdffae..e1d87d10b 100644 --- a/pkg/tools/subagent_tool_test.go +++ b/pkg/tools/subagent_tool_test.go @@ -47,7 +47,7 @@ func (m *MockLLMProvider) GetContextWindow() int { func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) manager.SetLLMOptions(2048, 0.6) tool := NewSubagentTool(manager) tool.SetContext("cli", "direct") @@ -74,7 +74,7 @@ func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) { // TestSubagentTool_Name verifies tool name func TestSubagentTool_Name(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSubagentTool(manager) if tool.Name() != "subagent" { @@ -85,7 +85,7 @@ func TestSubagentTool_Name(t *testing.T) { // TestSubagentTool_Description verifies tool description func TestSubagentTool_Description(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSubagentTool(manager) desc := tool.Description() @@ -100,7 +100,7 @@ func TestSubagentTool_Description(t *testing.T) { // TestSubagentTool_Parameters verifies tool parameters schema func TestSubagentTool_Parameters(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSubagentTool(manager) params := tool.Parameters() @@ -150,7 +150,7 @@ func TestSubagentTool_Parameters(t *testing.T) { // TestSubagentTool_SetContext verifies context setting func TestSubagentTool_SetContext(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSubagentTool(manager) tool.SetContext("test-channel", "test-chat") @@ -164,7 +164,7 @@ func TestSubagentTool_SetContext(t *testing.T) { func TestSubagentTool_Execute_Success(t *testing.T) { provider := &MockLLMProvider{} msgBus := bus.NewMessageBus() - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus) tool := NewSubagentTool(manager) tool.SetContext("telegram", "chat-123") @@ -220,7 +220,7 @@ func TestSubagentTool_Execute_Success(t *testing.T) { func TestSubagentTool_Execute_NoLabel(t *testing.T) { provider := &MockLLMProvider{} msgBus := bus.NewMessageBus() - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus) tool := NewSubagentTool(manager) ctx := context.Background() @@ -243,7 +243,7 @@ func TestSubagentTool_Execute_NoLabel(t *testing.T) { // TestSubagentTool_Execute_MissingTask tests error handling for missing task func TestSubagentTool_Execute_MissingTask(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSubagentTool(manager) ctx := context.Background() @@ -294,7 +294,7 @@ func TestSubagentTool_Execute_NilManager(t *testing.T) { func TestSubagentTool_Execute_ContextPassing(t *testing.T) { provider := &MockLLMProvider{} msgBus := bus.NewMessageBus() - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus) tool := NewSubagentTool(manager) // Set context @@ -323,7 +323,7 @@ func TestSubagentTool_ForUserTruncation(t *testing.T) { // Create a mock provider that returns very long content provider := &MockLLMProvider{} msgBus := bus.NewMessageBus() - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus) tool := NewSubagentTool(manager) ctx := context.Background() diff --git a/pkg/tools/team.go b/pkg/tools/team.go index 4290487b0..f3c190aa2 100644 --- a/pkg/tools/team.go +++ b/pkg/tools/team.go @@ -38,7 +38,13 @@ func (t *TeamTool) Name() string { } func (t *TeamTool) Description() string { - return "Compose and execute a team of distinct sub-agents. You (the main agent) should autonomously analyze the user's request, determine the necessary specialized roles, break down the work into sub-tasks, and assign them. Execute sequentially (passing output from one to the next) or concurrently in parallel." + base := "Compose and execute a team of distinct sub-agents. You (the main agent) should autonomously analyze the user's request, determine the necessary specialized roles, break down the work into sub-tasks, and assign them. Execute sequentially (passing output from one to the next) or concurrently in parallel." + if t.manager != nil { + if hint := t.manager.ModelCapabilityHint(); hint != "" { + return base + "\n\n" + hint + } + } + return base } func (t *TeamTool) Parameters() map[string]any { @@ -210,14 +216,17 @@ func upgradeRegistryForConcurrency(original *ToolRegistry) *ToolRegistry { // buildWorkerConfig creates a ToolLoopConfig for a specific team member, // potentially overriding the model based on the member's definition. -func buildWorkerConfig(baseConfig ToolLoopConfig, registry *ToolRegistry, m TeamMember) ToolLoopConfig { +func buildWorkerConfig(baseConfig ToolLoopConfig, registry *ToolRegistry, m TeamMember, manager *SubagentManager) (ToolLoopConfig, error) { cfg := baseConfig cfg.Tools = registry // Heterogeneous Agents: Override model if this team member requested a specific one if m.Model != "" { + if !manager.IsModelAllowed(m.Model) { + return cfg, fmt.Errorf("requested model '%s' is not in the allowed fallback candidates list for this agent workspace", m.Model) + } cfg.Model = m.Model } - return cfg + return cfg, nil } func (t *TeamTool) executeSequential(ctx context.Context, baseConfig ToolLoopConfig, members []TeamMember) *ToolResult { @@ -238,7 +247,13 @@ func (t *TeamTool) executeSequential(ctx context.Context, baseConfig ToolLoopCon {Role: "user", Content: actualTask}, } - workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, m) + workerConfig, err := buildWorkerConfig(baseConfig, baseConfig.Tools, m, t.manager) + if err != nil { + errStr := fmt.Sprintf("Phase %d (Role: %s) configuration failed: %v", i+1, m.Role, err) + finalOutput.WriteString(errStr + "\n") + return ErrorResult(errStr).WithError(err) + } + loopResult, err := RunToolLoop(ctx, workerConfig, messages, t.originChannel, t.originChatID) if err != nil { errStr := fmt.Sprintf("Phase %d (Role: %s) failed: %v", i+1, m.Role, err) @@ -278,7 +293,12 @@ func (t *TeamTool) executeParallel(ctx context.Context, baseConfig ToolLoopConfi {Role: "user", Content: member.Task}, } - workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, member) + workerConfig, err := buildWorkerConfig(baseConfig, baseConfig.Tools, member, t.manager) + if err != nil { + resultsChan <- workResult{index: index, role: member.Role, err: err} + return + } + loopResult, err := RunToolLoop(ctx, workerConfig, messages, t.originChannel, t.originChatID) if err != nil { @@ -344,7 +364,13 @@ func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, baseConfig Too finalOutput.WriteString(fmt.Sprintf("## Attempt %d\n", attempt)) // 2. Trigger Worker (resumes from its exact previous state!) - workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, worker) + workerConfig, err := buildWorkerConfig(baseConfig, baseConfig.Tools, worker, t.manager) + if err != nil { + errStr := fmt.Sprintf("Worker configuration failed on attempt %d: %v", attempt, err) + finalOutput.WriteString(errStr + "\n") + return ErrorResult(errStr).WithError(err) + } + workerResult, err := RunToolLoop(ctx, workerConfig, workerMessages, t.originChannel, t.originChatID) if err != nil { errStr := fmt.Sprintf("Worker failed on attempt %d: %v", attempt, err) @@ -365,7 +391,13 @@ func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, baseConfig Too {Role: "user", Content: evalContext}, } - evalConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, evaluator) + evalConfig, err := buildWorkerConfig(baseConfig, baseConfig.Tools, evaluator, t.manager) + if err != nil { + errStr := fmt.Sprintf("Evaluator configuration failed on attempt %d: %v", attempt, err) + finalOutput.WriteString(errStr + "\n") + return ErrorResult(errStr).WithError(err) + } + evalResult, err := RunToolLoop(ctx, evalConfig, evalMessages, t.originChannel, t.originChatID) if err != nil { errStr := fmt.Sprintf("Evaluator failed on attempt %d: %v", attempt, err) @@ -514,7 +546,17 @@ func (t *TeamTool) executeDAG(ctx context.Context, baseConfig ToolLoopConfig, me {Role: "user", Content: actualTask}, } - workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, m) + workerConfig, err := buildWorkerConfig(baseConfig, baseConfig.Tools, m, t.manager) + if err != nil { + masterErrMu.Lock() + if masterErr == nil { + masterErr = err + } + masterErrMu.Unlock() + resultChan <- nodeResult{id: id, err: err} + return + } + loopResult, err := RunToolLoop(ctx, workerConfig, messages, t.originChannel, t.originChatID) if err != nil {