feat(team): centralize team tool config with allowed_models, limits, and toggles
This commit is contained in:
parent
55dd32267a
commit
f5f991686d
9 changed files with 173 additions and 64 deletions
|
|
@ -155,7 +155,7 @@ func NewAgentInstance(
|
||||||
Primary: model,
|
Primary: model,
|
||||||
Fallbacks: fallbacks,
|
Fallbacks: fallbacks,
|
||||||
}
|
}
|
||||||
resolveFromModelList := func(raw string) (string, []string, bool) {
|
resolveFromModelList := func(raw string) (string, bool) {
|
||||||
ensureProtocol := func(model string) string {
|
ensureProtocol := func(model string) string {
|
||||||
model = strings.TrimSpace(model)
|
model = strings.TrimSpace(model)
|
||||||
if model == "" {
|
if model == "" {
|
||||||
|
|
@ -169,12 +169,12 @@ func NewAgentInstance(
|
||||||
|
|
||||||
raw = strings.TrimSpace(raw)
|
raw = strings.TrimSpace(raw)
|
||||||
if raw == "" {
|
if raw == "" {
|
||||||
return "", nil, false
|
return "", false
|
||||||
}
|
}
|
||||||
|
|
||||||
if cfg != nil {
|
if cfg != nil {
|
||||||
if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" {
|
if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" {
|
||||||
return ensureProtocol(mc.Model), mc.Tags, true
|
return ensureProtocol(mc.Model), true
|
||||||
}
|
}
|
||||||
|
|
||||||
for i := range cfg.ModelList {
|
for i := range cfg.ModelList {
|
||||||
|
|
@ -183,16 +183,16 @@ func NewAgentInstance(
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if fullModel == raw {
|
if fullModel == raw {
|
||||||
return ensureProtocol(fullModel), cfg.ModelList[i].Tags, true
|
return ensureProtocol(fullModel), true
|
||||||
}
|
}
|
||||||
_, modelID := providers.ExtractProtocol(fullModel)
|
_, modelID := providers.ExtractProtocol(fullModel)
|
||||||
if modelID == raw {
|
if modelID == raw {
|
||||||
return ensureProtocol(fullModel), cfg.ModelList[i].Tags, true
|
return ensureProtocol(fullModel), true
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
return "", nil, false
|
return "", false
|
||||||
}
|
}
|
||||||
|
|
||||||
candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList)
|
candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList)
|
||||||
|
|
|
||||||
|
|
@ -220,7 +220,7 @@ func registerSharedTools(
|
||||||
}
|
}
|
||||||
|
|
||||||
// Spawn tool with allowlist checker
|
// Spawn tool with allowlist checker
|
||||||
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Candidates, agent.Workspace, msgBus)
|
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Candidates, agent.Workspace, cfg.Tools.Team, msgBus)
|
||||||
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
|
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
|
||||||
spawnTool := tools.NewSpawnTool(subagentManager)
|
spawnTool := tools.NewSpawnTool(subagentManager)
|
||||||
currentAgentID := agentID
|
currentAgentID := agentID
|
||||||
|
|
@ -230,10 +230,14 @@ func registerSharedTools(
|
||||||
agent.Tools.Register(spawnTool)
|
agent.Tools.Register(spawnTool)
|
||||||
|
|
||||||
teamTool := tools.NewTeamTool(subagentManager)
|
teamTool := tools.NewTeamTool(subagentManager)
|
||||||
|
if cfg.Tools.IsToolEnabled("team") {
|
||||||
agent.Tools.Register(teamTool)
|
agent.Tools.Register(teamTool)
|
||||||
|
}
|
||||||
|
|
||||||
spawnSubAgentTool := tools.NewSpawnSubAgentTool(subagentManager)
|
spawnSubAgentTool := tools.NewSpawnSubAgentTool(subagentManager)
|
||||||
|
if cfg.Tools.IsToolEnabled("spawn_sub_agent") {
|
||||||
agent.Tools.Register(spawnSubAgentTool)
|
agent.Tools.Register(spawnSubAgentTool)
|
||||||
|
}
|
||||||
|
|
||||||
// Direction 3: Hierarchical Decomposition.
|
// Direction 3: Hierarchical Decomposition.
|
||||||
// Share the fully-built registry (which includes team, spawn_sub_agent, etc.) back
|
// Share the fully-built registry (which includes team, spawn_sub_agent, etc.) back
|
||||||
|
|
@ -242,7 +246,7 @@ func registerSharedTools(
|
||||||
subagentManager.SetTools(agent.Tools)
|
subagentManager.SetTools(agent.Tools)
|
||||||
if cfg.Tools.IsToolEnabled("spawn") {
|
if cfg.Tools.IsToolEnabled("spawn") {
|
||||||
if cfg.Tools.IsToolEnabled("subagent") {
|
if cfg.Tools.IsToolEnabled("subagent") {
|
||||||
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Candidates, agent.Workspace, msgBus)
|
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Candidates, agent.Workspace, cfg.Tools.Team, msgBus)
|
||||||
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
|
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
|
||||||
spawnTool := tools.NewSpawnTool(subagentManager)
|
spawnTool := tools.NewSpawnTool(subagentManager)
|
||||||
currentAgentID := agentID
|
currentAgentID := agentID
|
||||||
|
|
|
||||||
|
|
@ -74,6 +74,22 @@ func (f *FlexibleStringSlice) UnmarshalText(text []byte) error {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type TeamModelConfig struct {
|
||||||
|
Name string `json:"name" yaml:"name"`
|
||||||
|
Tags []string `json:"tags,omitempty" yaml:"tags,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type TeamToolsConfig struct {
|
||||||
|
ToolConfig ` envPrefix:"PICOCLAW_TOOLS_TEAM_"`
|
||||||
|
MaxMembers int `json:"max_members" env:"PICOCLAW_TOOLS_TEAM_MAX_MEMBERS"`
|
||||||
|
MaxTeamTokens int `json:"max_team_tokens" env:"PICOCLAW_TOOLS_TEAM_MAX_TOKENS"`
|
||||||
|
MaxEvaluatorLoops int `json:"max_evaluator_loops" env:"PICOCLAW_TOOLS_TEAM_MAX_EVALUATOR_LOOPS"`
|
||||||
|
MaxTimeoutMinutes int `json:"max_timeout_minutes" env:"PICOCLAW_TOOLS_TEAM_MAX_TIMEOUT_MINUTES"`
|
||||||
|
DisableAutoReviewer bool `json:"disable_auto_reviewer" env:"PICOCLAW_TOOLS_TEAM_DISABLE_AUTO_REVIEWER"`
|
||||||
|
AllowedStrategies []string `json:"allowed_strategies" env:"PICOCLAW_TOOLS_TEAM_ALLOWED_STRATEGIES"`
|
||||||
|
AllowedModels []TeamModelConfig `json:"allowed_models" env:"-"`
|
||||||
|
}
|
||||||
|
|
||||||
type Config struct {
|
type Config struct {
|
||||||
Agents AgentsConfig `json:"agents"`
|
Agents AgentsConfig `json:"agents"`
|
||||||
Bindings []AgentBinding `json:"bindings,omitempty"`
|
Bindings []AgentBinding `json:"bindings,omitempty"`
|
||||||
|
|
@ -591,7 +607,6 @@ type ModelConfig struct {
|
||||||
// Required fields
|
// Required fields
|
||||||
ModelName string `json:"model_name"` // User-facing alias for the model
|
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")
|
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
|
// HTTP-based providers
|
||||||
APIBase string `json:"api_base,omitempty"` // API endpoint URL
|
APIBase string `json:"api_base,omitempty"` // API endpoint URL
|
||||||
|
|
@ -748,6 +763,8 @@ type ToolsConfig struct {
|
||||||
Spawn ToolConfig `json:"spawn" envPrefix:"PICOCLAW_TOOLS_SPAWN_"`
|
Spawn ToolConfig `json:"spawn" envPrefix:"PICOCLAW_TOOLS_SPAWN_"`
|
||||||
SPI ToolConfig `json:"spi" envPrefix:"PICOCLAW_TOOLS_SPI_"`
|
SPI ToolConfig `json:"spi" envPrefix:"PICOCLAW_TOOLS_SPI_"`
|
||||||
Subagent ToolConfig `json:"subagent" envPrefix:"PICOCLAW_TOOLS_SUBAGENT_"`
|
Subagent ToolConfig `json:"subagent" envPrefix:"PICOCLAW_TOOLS_SUBAGENT_"`
|
||||||
|
SpawnSubAgent ToolConfig `json:"spawn_sub_agent" envPrefix:"PICOCLAW_TOOLS_SPAWN_SUB_AGENT_"`
|
||||||
|
Team TeamToolsConfig `json:"team" envPrefix:"PICOCLAW_TOOLS_TEAM_"`
|
||||||
WebFetch ToolConfig `json:"web_fetch" envPrefix:"PICOCLAW_TOOLS_WEB_FETCH_"`
|
WebFetch ToolConfig `json:"web_fetch" envPrefix:"PICOCLAW_TOOLS_WEB_FETCH_"`
|
||||||
WriteFile ToolConfig `json:"write_file" envPrefix:"PICOCLAW_TOOLS_WRITE_FILE_"`
|
WriteFile ToolConfig `json:"write_file" envPrefix:"PICOCLAW_TOOLS_WRITE_FILE_"`
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,6 @@ type FallbackChain struct {
|
||||||
type FallbackCandidate struct {
|
type FallbackCandidate struct {
|
||||||
Provider string
|
Provider string
|
||||||
Model string
|
Model string
|
||||||
Tags []string
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// FallbackResult contains the successful response and metadata about all attempts.
|
// FallbackResult contains the successful response and metadata about all attempts.
|
||||||
|
|
@ -50,19 +49,17 @@ func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandid
|
||||||
func ResolveCandidatesWithLookup(
|
func ResolveCandidatesWithLookup(
|
||||||
cfg ModelConfig,
|
cfg ModelConfig,
|
||||||
defaultProvider string,
|
defaultProvider string,
|
||||||
lookup func(raw string) (resolved string, tags []string, ok bool),
|
lookup func(raw string) (resolved string, ok bool),
|
||||||
) []FallbackCandidate {
|
) []FallbackCandidate {
|
||||||
seen := make(map[string]bool)
|
seen := make(map[string]bool)
|
||||||
var candidates []FallbackCandidate
|
var candidates []FallbackCandidate
|
||||||
|
|
||||||
addCandidate := func(raw string) {
|
addCandidate := func(raw string) {
|
||||||
candidateRaw := strings.TrimSpace(raw)
|
candidateRaw := strings.TrimSpace(raw)
|
||||||
var modelTags []string
|
|
||||||
|
|
||||||
if lookup != nil {
|
if lookup != nil {
|
||||||
if resolved, tags, ok := lookup(candidateRaw); ok {
|
if resolved, ok := lookup(candidateRaw); ok {
|
||||||
candidateRaw = resolved
|
candidateRaw = resolved
|
||||||
modelTags = tags
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -78,7 +75,6 @@ func ResolveCandidatesWithLookup(
|
||||||
candidates = append(candidates, FallbackCandidate{
|
candidates = append(candidates, FallbackCandidate{
|
||||||
Provider: ref.Provider,
|
Provider: ref.Provider,
|
||||||
Model: ref.Model,
|
Model: ref.Model,
|
||||||
Tags: modelTags,
|
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -459,11 +459,11 @@ func TestResolveCandidatesWithLookup_AliasResolvesToNestedModel(t *testing.T) {
|
||||||
Fallbacks: nil,
|
Fallbacks: nil,
|
||||||
}
|
}
|
||||||
|
|
||||||
lookup := func(raw string) (string, []string, bool) {
|
lookup := func(raw string) (string, bool) {
|
||||||
if raw == "step-3.5-flash" {
|
if raw == "step-3.5-flash" {
|
||||||
return "openrouter/stepfun/step-3.5-flash:free", nil, true
|
return "openrouter/stepfun/step-3.5-flash:free", true
|
||||||
}
|
}
|
||||||
return "", nil, false
|
return "", false
|
||||||
}
|
}
|
||||||
|
|
||||||
candidates := ResolveCandidatesWithLookup(cfg, "", lookup)
|
candidates := ResolveCandidatesWithLookup(cfg, "", lookup)
|
||||||
|
|
@ -484,11 +484,11 @@ func TestResolveCandidatesWithLookup_DeduplicateAfterLookup(t *testing.T) {
|
||||||
Fallbacks: []string{"openrouter/stepfun/step-3.5-flash:free"},
|
Fallbacks: []string{"openrouter/stepfun/step-3.5-flash:free"},
|
||||||
}
|
}
|
||||||
|
|
||||||
lookup := func(raw string) (string, []string, bool) {
|
lookup := func(raw string) (string, bool) {
|
||||||
if raw == "step-3.5-flash" {
|
if raw == "step-3.5-flash" {
|
||||||
return "openrouter/stepfun/step-3.5-flash:free", nil, true
|
return "openrouter/stepfun/step-3.5-flash:free", true
|
||||||
}
|
}
|
||||||
return "", nil, false
|
return "", false
|
||||||
}
|
}
|
||||||
|
|
||||||
candidates := ResolveCandidatesWithLookup(cfg, "", lookup)
|
candidates := ResolveCandidatesWithLookup(cfg, "", lookup)
|
||||||
|
|
@ -503,11 +503,11 @@ func TestResolveCandidatesWithLookup_AliasWithoutProtocolUsesDefaultProvider(t *
|
||||||
Fallbacks: nil,
|
Fallbacks: nil,
|
||||||
}
|
}
|
||||||
|
|
||||||
lookup := func(raw string) (string, []string, bool) {
|
lookup := func(raw string) (string, bool) {
|
||||||
if raw == "glm-5" {
|
if raw == "glm-5" {
|
||||||
return "glm-5", nil, true
|
return "glm-5", true
|
||||||
}
|
}
|
||||||
return "", nil, false
|
return "", false
|
||||||
}
|
}
|
||||||
|
|
||||||
candidates := ResolveCandidatesWithLookup(cfg, "openai", lookup)
|
candidates := ResolveCandidatesWithLookup(cfg, "openai", lookup)
|
||||||
|
|
|
||||||
|
|
@ -4,11 +4,13 @@ import (
|
||||||
"context"
|
"context"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
|
func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
|
||||||
tool := NewSpawnTool(manager)
|
tool := NewSpawnTool(manager)
|
||||||
|
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
|
|
@ -42,7 +44,7 @@ func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
|
||||||
|
|
||||||
func TestSpawnTool_Execute_ValidTask(t *testing.T) {
|
func TestSpawnTool_Execute_ValidTask(t *testing.T) {
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
|
||||||
tool := NewSpawnTool(manager)
|
tool := NewSpawnTool(manager)
|
||||||
|
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
|
|
|
||||||
|
|
@ -8,6 +8,7 @@ import (
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/bus"
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
"github.com/sipeed/picoclaw/pkg/providers"
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -54,6 +55,7 @@ type SubagentManager struct {
|
||||||
bus *bus.MessageBus
|
bus *bus.MessageBus
|
||||||
workspace string
|
workspace string
|
||||||
tools *ToolRegistry
|
tools *ToolRegistry
|
||||||
|
teamConfig config.TeamToolsConfig
|
||||||
maxIterations int
|
maxIterations int
|
||||||
maxTokens int
|
maxTokens int
|
||||||
temperature float64
|
temperature float64
|
||||||
|
|
@ -67,6 +69,7 @@ func NewSubagentManager(
|
||||||
defaultModel string,
|
defaultModel string,
|
||||||
candidates []providers.FallbackCandidate,
|
candidates []providers.FallbackCandidate,
|
||||||
workspace string,
|
workspace string,
|
||||||
|
teamConfig config.TeamToolsConfig,
|
||||||
bus *bus.MessageBus,
|
bus *bus.MessageBus,
|
||||||
) *SubagentManager {
|
) *SubagentManager {
|
||||||
return &SubagentManager{
|
return &SubagentManager{
|
||||||
|
|
@ -74,6 +77,7 @@ func NewSubagentManager(
|
||||||
provider: provider,
|
provider: provider,
|
||||||
defaultModel: defaultModel,
|
defaultModel: defaultModel,
|
||||||
allowedModels: candidates,
|
allowedModels: candidates,
|
||||||
|
teamConfig: teamConfig,
|
||||||
bus: bus,
|
bus: bus,
|
||||||
workspace: workspace,
|
workspace: workspace,
|
||||||
tools: NewToolRegistry(),
|
tools: NewToolRegistry(),
|
||||||
|
|
@ -89,7 +93,21 @@ func (sm *SubagentManager) IsModelAllowed(model string) bool {
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
|
|
||||||
// Otherwise, check against the resolved candidates (primary + fallbacks + explicitly configured)
|
// 1. Check against explicitly allowed models in team config
|
||||||
|
for _, cand := range sm.teamConfig.AllowedModels {
|
||||||
|
if cand.Name == model {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Otherwise, check against the resolved candidates (primary + fallbacks + explicitly configured)
|
||||||
|
// If teamConfig.AllowedModels is set, we strictly enforce it and DO NOT fall back to candidates
|
||||||
|
// unless the candidate model has tags that overlap with AllowedTags. But since AllowedTags
|
||||||
|
// was not implemented yet, just check fallback for backwards compatibility if teamConfig is empty.
|
||||||
|
if len(sm.teamConfig.AllowedModels) > 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
for _, cand := range sm.allowedModels {
|
for _, cand := range sm.allowedModels {
|
||||||
if cand.Model == model {
|
if cand.Model == model {
|
||||||
return true
|
return true
|
||||||
|
|
@ -107,23 +125,22 @@ func (sm *SubagentManager) ModelCapabilityHint() string {
|
||||||
|
|
||||||
var modelLines []string
|
var modelLines []string
|
||||||
for _, cand := range sm.allowedModels {
|
for _, cand := range sm.allowedModels {
|
||||||
if len(cand.Tags) == 0 {
|
|
||||||
modelLines = append(modelLines, fmt.Sprintf(" - %s (general purpose)", cand.Model))
|
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 := "When selecting a 'model' for sub-agents, use ONLY these configured models:\n"
|
||||||
|
if len(sm.teamConfig.AllowedModels) > 0 {
|
||||||
|
for _, cand := range sm.teamConfig.AllowedModels {
|
||||||
|
tagsStr := ""
|
||||||
|
if len(cand.Tags) > 0 {
|
||||||
|
tagsStr = fmt.Sprintf(" [%s]", strings.Join(cand.Tags, ", "))
|
||||||
|
}
|
||||||
|
hint += fmt.Sprintf(" - %s%s\n", cand.Name, tagsStr)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
hint += strings.Join(modelLines, "\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."
|
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
|
return hint
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,7 @@ import (
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/bus"
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
"github.com/sipeed/picoclaw/pkg/providers"
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -47,7 +48,7 @@ func (m *MockLLMProvider) GetContextWindow() int {
|
||||||
|
|
||||||
func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) {
|
func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) {
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
|
||||||
manager.SetLLMOptions(2048, 0.6)
|
manager.SetLLMOptions(2048, 0.6)
|
||||||
tool := NewSubagentTool(manager)
|
tool := NewSubagentTool(manager)
|
||||||
|
|
||||||
|
|
@ -73,7 +74,7 @@ func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) {
|
||||||
// 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{}
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
|
||||||
tool := NewSubagentTool(manager)
|
tool := NewSubagentTool(manager)
|
||||||
|
|
||||||
if tool.Name() != "subagent" {
|
if tool.Name() != "subagent" {
|
||||||
|
|
@ -84,7 +85,7 @@ func TestSubagentTool_Name(t *testing.T) {
|
||||||
// TestSubagentTool_Description verifies tool description
|
// TestSubagentTool_Description verifies tool description
|
||||||
func TestSubagentTool_Description(t *testing.T) {
|
func TestSubagentTool_Description(t *testing.T) {
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
|
||||||
tool := NewSubagentTool(manager)
|
tool := NewSubagentTool(manager)
|
||||||
|
|
||||||
desc := tool.Description()
|
desc := tool.Description()
|
||||||
|
|
@ -99,7 +100,7 @@ func TestSubagentTool_Description(t *testing.T) {
|
||||||
// TestSubagentTool_Parameters verifies tool parameters schema
|
// TestSubagentTool_Parameters verifies tool parameters schema
|
||||||
func TestSubagentTool_Parameters(t *testing.T) {
|
func TestSubagentTool_Parameters(t *testing.T) {
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
|
||||||
tool := NewSubagentTool(manager)
|
tool := NewSubagentTool(manager)
|
||||||
|
|
||||||
params := tool.Parameters()
|
params := tool.Parameters()
|
||||||
|
|
@ -149,7 +150,7 @@ func TestSubagentTool_Parameters(t *testing.T) {
|
||||||
// TestSubagentTool_SetContext verifies context setting
|
// TestSubagentTool_SetContext verifies context setting
|
||||||
// func TestSubagentTool_SetContext(t *testing.T) {
|
// func TestSubagentTool_SetContext(t *testing.T) {
|
||||||
// provider := &MockLLMProvider{}
|
// provider := &MockLLMProvider{}
|
||||||
// manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
|
// manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
|
||||||
// tool := NewSubagentTool(manager)
|
// tool := NewSubagentTool(manager)
|
||||||
|
|
||||||
// tool.SetContext("test-channel", "test-chat")
|
// tool.SetContext("test-channel", "test-chat")
|
||||||
|
|
@ -163,7 +164,7 @@ func TestSubagentTool_Parameters(t *testing.T) {
|
||||||
func TestSubagentTool_Execute_Success(t *testing.T) {
|
func TestSubagentTool_Execute_Success(t *testing.T) {
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
msgBus := bus.NewMessageBus()
|
msgBus := bus.NewMessageBus()
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, msgBus)
|
||||||
tool := NewSubagentTool(manager)
|
tool := NewSubagentTool(manager)
|
||||||
|
|
||||||
ctx := WithToolContext(context.Background(), "telegram", "chat-123")
|
ctx := WithToolContext(context.Background(), "telegram", "chat-123")
|
||||||
|
|
@ -218,7 +219,7 @@ func TestSubagentTool_Execute_Success(t *testing.T) {
|
||||||
func TestSubagentTool_Execute_NoLabel(t *testing.T) {
|
func TestSubagentTool_Execute_NoLabel(t *testing.T) {
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
msgBus := bus.NewMessageBus()
|
msgBus := bus.NewMessageBus()
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, msgBus)
|
||||||
tool := NewSubagentTool(manager)
|
tool := NewSubagentTool(manager)
|
||||||
|
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
|
|
@ -241,7 +242,7 @@ func TestSubagentTool_Execute_NoLabel(t *testing.T) {
|
||||||
// TestSubagentTool_Execute_MissingTask tests error handling for missing task
|
// TestSubagentTool_Execute_MissingTask tests error handling for missing task
|
||||||
func TestSubagentTool_Execute_MissingTask(t *testing.T) {
|
func TestSubagentTool_Execute_MissingTask(t *testing.T) {
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
|
||||||
tool := NewSubagentTool(manager)
|
tool := NewSubagentTool(manager)
|
||||||
|
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
|
|
@ -292,7 +293,7 @@ func TestSubagentTool_Execute_NilManager(t *testing.T) {
|
||||||
func TestSubagentTool_Execute_ContextPassing(t *testing.T) {
|
func TestSubagentTool_Execute_ContextPassing(t *testing.T) {
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
msgBus := bus.NewMessageBus()
|
msgBus := bus.NewMessageBus()
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, msgBus)
|
||||||
tool := NewSubagentTool(manager)
|
tool := NewSubagentTool(manager)
|
||||||
|
|
||||||
channel := "test-channel"
|
channel := "test-channel"
|
||||||
|
|
@ -318,7 +319,7 @@ func TestSubagentTool_ForUserTruncation(t *testing.T) {
|
||||||
// Create a mock provider that returns very long content
|
// Create a mock provider that returns very long content
|
||||||
provider := &MockLLMProvider{}
|
provider := &MockLLMProvider{}
|
||||||
msgBus := bus.NewMessageBus()
|
msgBus := bus.NewMessageBus()
|
||||||
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus)
|
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, msgBus)
|
||||||
tool := NewSubagentTool(manager)
|
tool := NewSubagentTool(manager)
|
||||||
|
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
|
|
|
||||||
|
|
@ -167,6 +167,16 @@ func (t *TeamTool) maybeRunAutoReviewer(
|
||||||
return "" // Unknown produces types, skip
|
return "" // Unknown produces types, skip
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 3. Skip Auto-Reviewer if disabled by config
|
||||||
|
sm := t.manager
|
||||||
|
sm.mu.RLock()
|
||||||
|
disabled := sm.teamConfig.DisableAutoReviewer
|
||||||
|
sm.mu.RUnlock()
|
||||||
|
|
||||||
|
if disabled {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
reviewerTask := strings.Join(taskParts, "\n\n") +
|
reviewerTask := strings.Join(taskParts, "\n\n") +
|
||||||
"\n\nContext from the workers that produced these artifacts:\n" + workerSummary
|
"\n\nContext from the workers that produced these artifacts:\n" + workerSummary
|
||||||
|
|
||||||
|
|
@ -184,8 +194,37 @@ func (t *TeamTool) maybeRunAutoReviewer(
|
||||||
|
|
||||||
func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
|
func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
|
||||||
strategy, ok := args["strategy"].(string)
|
strategy, ok := args["strategy"].(string)
|
||||||
if !ok || (strategy != "sequential" && strategy != "parallel" && strategy != "dag" && strategy != "evaluator_optimizer") {
|
if !ok {
|
||||||
return ErrorResult("strategy must be 'sequential', 'parallel', 'dag', or 'evaluator_optimizer'")
|
return ErrorResult("strategy is required")
|
||||||
|
}
|
||||||
|
|
||||||
|
if t.manager == nil {
|
||||||
|
return ErrorResult("Subagent manager not configured").WithError(fmt.Errorf("manager is nil"))
|
||||||
|
}
|
||||||
|
|
||||||
|
sm := t.manager
|
||||||
|
sm.mu.RLock()
|
||||||
|
teamConfig := sm.teamConfig
|
||||||
|
sm.mu.RUnlock()
|
||||||
|
|
||||||
|
// 1. Validate Strategy
|
||||||
|
validStrategy := false
|
||||||
|
if len(teamConfig.AllowedStrategies) > 0 {
|
||||||
|
for _, s := range teamConfig.AllowedStrategies {
|
||||||
|
if strategy == s {
|
||||||
|
validStrategy = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Default allowed strategies if not configured
|
||||||
|
if strategy == "sequential" || strategy == "parallel" || strategy == "dag" || strategy == "evaluator_optimizer" {
|
||||||
|
validStrategy = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if !validStrategy {
|
||||||
|
return ErrorResult(fmt.Sprintf("strategy '%s' is not allowed by configuration", strategy))
|
||||||
}
|
}
|
||||||
|
|
||||||
membersRaw, ok := args["members"].([]any)
|
membersRaw, ok := args["members"].([]any)
|
||||||
|
|
@ -193,17 +232,38 @@ func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult
|
||||||
return ErrorResult("members map array is required and must not be empty")
|
return ErrorResult("members map array is required and must not be empty")
|
||||||
}
|
}
|
||||||
|
|
||||||
maxTokensFloat, ok := args["max_team_tokens"].(float64)
|
// 2. Validate Max Members
|
||||||
var budget *atomic.Int64
|
if teamConfig.MaxMembers > 0 && len(membersRaw) > teamConfig.MaxMembers {
|
||||||
if ok && maxTokensFloat > 0 {
|
return ErrorResult(fmt.Sprintf("Team exceeds maximum allowed members (%d). You requested %d members.", teamConfig.MaxMembers, len(membersRaw)))
|
||||||
budget = &atomic.Int64{}
|
|
||||||
budget.Store(int64(maxTokensFloat))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if t.manager == nil {
|
maxTokensFloat, ok := args["max_team_tokens"].(float64)
|
||||||
return ErrorResult("Subagent manager not configured").WithError(fmt.Errorf("manager is nil"))
|
|
||||||
|
// Enforce hard budgets from config
|
||||||
|
effectiveMaxTokens := int64(0)
|
||||||
|
if teamConfig.MaxTeamTokens > 0 {
|
||||||
|
effectiveMaxTokens = int64(teamConfig.MaxTeamTokens)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if ok && maxTokensFloat > 0 {
|
||||||
|
requestedTokens := int64(maxTokensFloat)
|
||||||
|
// If LLM requested tokens but config enforces a smaller hard limit, clamp it
|
||||||
|
if effectiveMaxTokens > 0 && requestedTokens > effectiveMaxTokens {
|
||||||
|
effectiveMaxTokens = requestedTokens // LLM asked for more, but we clamp to config
|
||||||
|
// Wait, the clamping logic: if requested > max_team_tokens, clamp effectively shrinks it to max_team_tokens.
|
||||||
|
effectiveMaxTokens = int64(teamConfig.MaxTeamTokens)
|
||||||
|
} else if effectiveMaxTokens == 0 || requestedTokens < effectiveMaxTokens {
|
||||||
|
effectiveMaxTokens = requestedTokens // LLM asked for less budget, let them be conservative
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
var budget *atomic.Int64
|
||||||
|
if effectiveMaxTokens > 0 {
|
||||||
|
budget = &atomic.Int64{}
|
||||||
|
budget.Store(effectiveMaxTokens)
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
var members []TeamMember
|
var members []TeamMember
|
||||||
for i, mRaw := range membersRaw {
|
for i, mRaw := range membersRaw {
|
||||||
mMap, ok := mRaw.(map[string]any)
|
mMap, ok := mRaw.(map[string]any)
|
||||||
|
|
@ -256,8 +316,11 @@ func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult
|
||||||
}
|
}
|
||||||
|
|
||||||
// Create a new master context for team bounding
|
// Create a new master context for team bounding
|
||||||
// In the future this could be overridden by an argument
|
timeoutDur := 15 * time.Minute
|
||||||
teamCtx, cancel := context.WithTimeout(ctx, 15*time.Minute)
|
if teamConfig.MaxTimeoutMinutes > 0 {
|
||||||
|
timeoutDur = time.Duration(teamConfig.MaxTimeoutMinutes) * time.Minute
|
||||||
|
}
|
||||||
|
teamCtx, cancel := context.WithTimeout(ctx, timeoutDur)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
|
||||||
// If strategy is parallel or dag, we must upgrade the file tools to be concurrent-safe (locking)
|
// If strategy is parallel or dag, we must upgrade the file tools to be concurrent-safe (locking)
|
||||||
|
|
@ -460,7 +523,16 @@ func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, baseConfig Too
|
||||||
{Role: "user", Content: worker.Task},
|
{Role: "user", Content: worker.Task},
|
||||||
}
|
}
|
||||||
|
|
||||||
|
sm := t.manager
|
||||||
|
sm.mu.RLock()
|
||||||
|
teamConfig := sm.teamConfig
|
||||||
|
sm.mu.RUnlock()
|
||||||
|
|
||||||
maxLoops := 5
|
maxLoops := 5
|
||||||
|
if teamConfig.MaxEvaluatorLoops > 0 {
|
||||||
|
maxLoops = teamConfig.MaxEvaluatorLoops
|
||||||
|
}
|
||||||
|
|
||||||
for attempt := 1; attempt <= maxLoops; attempt++ {
|
for attempt := 1; attempt <= maxLoops; attempt++ {
|
||||||
finalOutput.WriteString(fmt.Sprintf("## Attempt %d\n", attempt))
|
finalOutput.WriteString(fmt.Sprintf("## Attempt %d\n", attempt))
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue