Merge feat/team: Agent Teams architecture
This commit is contained in:
Administrator
commit
1548691851
18 changed files with 1377 additions and 64 deletions
|
|
@ -78,7 +78,9 @@ Your workspace is at: %s
|
|||
|
||||
3. **Memory** - When interacting with me if something seems memorable, update %s/memory/MEMORY.md
|
||||
|
||||
4. **Context summaries** - Conversation summaries provided as context are approximate references only. They may be incomplete or outdated. Always defer to explicit user instructions over summary content.`,
|
||||
4. **Context summaries** - Conversation summaries provided as context are approximate references only. They may be incomplete or outdated. Always defer to explicit user instructions over summary content.
|
||||
|
||||
5. **Team delegation** - For any task that is non-trivial, multi-step, or involves distinct concerns (e.g. "convert React to Vue", "build a feature", "analyze and report"), you MUST use the 'team' tool to delegate and parallelize. Do NOT attempt to handle complex tasks inline by calling tools one by one yourself. Decompose first, delegate second, then report the outcome.`,
|
||||
workspacePath, workspacePath, workspacePath, workspacePath, workspacePath)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -98,7 +98,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 == "" {
|
||||
|
|
@ -112,12 +112,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 {
|
||||
|
|
@ -126,16 +126,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)
|
||||
|
|
|
|||
|
|
@ -148,7 +148,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
|
||||
|
|
@ -156,6 +156,18 @@ func registerSharedTools(
|
|||
return registry.CanSpawnSubagent(currentAgentID, targetAgentID)
|
||||
})
|
||||
agent.Tools.Register(spawnTool)
|
||||
|
||||
teamTool := tools.NewTeamTool(subagentManager)
|
||||
agent.Tools.Register(teamTool)
|
||||
|
||||
spawnSubAgentTool := tools.NewSpawnSubAgentTool(subagentManager)
|
||||
agent.Tools.Register(spawnSubAgentTool)
|
||||
|
||||
// Direction 3: Hierarchical Decomposition.
|
||||
// Share the fully-built registry (which includes team, spawn_sub_agent, etc.) back
|
||||
// to the subagent manager so that all workers spawned by this agent also inherit
|
||||
// the full toolset — enabling sub-agents to recursively call 'team' themselves.
|
||||
subagentManager.SetTools(agent.Tools)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -235,6 +235,7 @@ func parseResponse(body []byte) (*LLMResponse, error) {
|
|||
|
||||
choice := apiResponse.Choices[0]
|
||||
toolCalls := make([]ToolCall, 0, len(choice.Message.ToolCalls))
|
||||
truncated := false
|
||||
for _, tc := range choice.Message.ToolCalls {
|
||||
arguments := make(map[string]any)
|
||||
name := ""
|
||||
|
|
@ -249,8 +250,10 @@ func parseResponse(body []byte) (*LLMResponse, error) {
|
|||
name = tc.Function.Name
|
||||
if tc.Function.Arguments != "" {
|
||||
if err := json.Unmarshal([]byte(tc.Function.Arguments), &arguments); err != nil {
|
||||
// JSON is malformed (likely truncated due to max_tokens). Log and signal truncation.
|
||||
log.Printf("openai_compat: failed to decode tool call arguments for %q: %v", name, err)
|
||||
arguments["raw"] = tc.Function.Arguments
|
||||
truncated = true
|
||||
continue // Skip this malformed tool call entirely
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -274,13 +277,19 @@ func parseResponse(body []byte) (*LLMResponse, error) {
|
|||
toolCalls = append(toolCalls, toolCall)
|
||||
}
|
||||
|
||||
finishReason := choice.FinishReason
|
||||
// Propagate truncation: if finish_reason is "length" or we detected bad JSON, mark as truncated.
|
||||
if truncated || finishReason == "length" {
|
||||
finishReason = "truncated"
|
||||
}
|
||||
|
||||
return &LLMResponse{
|
||||
Content: choice.Message.Content,
|
||||
ReasoningContent: choice.Message.ReasoningContent,
|
||||
Reasoning: choice.Message.Reasoning,
|
||||
ReasoningDetails: choice.Message.ReasoningDetails,
|
||||
ToolCalls: toolCalls,
|
||||
FinishReason: choice.FinishReason,
|
||||
FinishReason: finishReason,
|
||||
Usage: apiResponse.Usage,
|
||||
}, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,9 +2,7 @@ package tools
|
|||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"strings"
|
||||
)
|
||||
|
||||
|
|
@ -76,6 +74,12 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe
|
|||
return SilentResult(fmt.Sprintf("File edited: %s", path))
|
||||
}
|
||||
|
||||
func (t *EditFileTool) UpgradeToConcurrent() Tool {
|
||||
return &EditFileTool{
|
||||
fs: &ConcurrentFS{baseFS: t.fs},
|
||||
}
|
||||
}
|
||||
|
||||
type AppendFileTool struct {
|
||||
fs fileSystem
|
||||
}
|
||||
|
|
@ -132,31 +136,25 @@ func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *Tool
|
|||
return SilentResult(fmt.Sprintf("Appended to %s", path))
|
||||
}
|
||||
|
||||
// editFile reads the file via sysFs, performs the replacement, and writes back.
|
||||
// It uses a fileSystem interface, allowing the same logic for both restricted and unrestricted modes.
|
||||
func editFile(sysFs fileSystem, path, oldText, newText string) error {
|
||||
content, err := sysFs.ReadFile(path)
|
||||
if err != nil {
|
||||
return err
|
||||
func (t *AppendFileTool) UpgradeToConcurrent() Tool {
|
||||
return &AppendFileTool{
|
||||
fs: &ConcurrentFS{baseFS: t.fs},
|
||||
}
|
||||
|
||||
newContent, err := replaceEditContent(content, oldText, newText)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return sysFs.WriteFile(path, newContent)
|
||||
}
|
||||
|
||||
// appendFile reads the existing content (if any) via sysFs, appends new content, and writes back.
|
||||
func appendFile(sysFs fileSystem, path, appendContent string) error {
|
||||
content, err := sysFs.ReadFile(path)
|
||||
if err != nil && !errors.Is(err, fs.ErrNotExist) {
|
||||
return err
|
||||
}
|
||||
// editFile reads the file via sysFs, performs the replacement, and writes back atomically.
|
||||
func editFile(sysFs fileSystem, path, oldText, newText string) error {
|
||||
return sysFs.EditFile(path, func(content []byte) ([]byte, error) {
|
||||
return replaceEditContent(content, oldText, newText)
|
||||
})
|
||||
}
|
||||
|
||||
newContent := append(content, []byte(appendContent)...)
|
||||
return sysFs.WriteFile(path, newContent)
|
||||
// appendFile reads the existing content (if any) via sysFs, appends new content, and writes back atomically.
|
||||
func appendFile(sysFs fileSystem, path, appendContent string) error {
|
||||
return sysFs.EditFile(path, func(content []byte) ([]byte, error) {
|
||||
newContent := append(content, []byte(appendContent)...)
|
||||
return newContent, nil
|
||||
})
|
||||
}
|
||||
|
||||
// replaceEditContent handles the core logic of finding and replacing a single occurrence of oldText.
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import (
|
|||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/fileutil"
|
||||
|
|
@ -131,6 +132,12 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe
|
|||
return NewToolResult(string(content))
|
||||
}
|
||||
|
||||
func (t *ReadFileTool) UpgradeToConcurrent() Tool {
|
||||
return &ReadFileTool{
|
||||
fs: &ConcurrentFS{baseFS: t.fs},
|
||||
}
|
||||
}
|
||||
|
||||
type WriteFileTool struct {
|
||||
fs fileSystem
|
||||
}
|
||||
|
|
@ -188,6 +195,12 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]any) *ToolR
|
|||
return SilentResult(fmt.Sprintf("File written: %s", path))
|
||||
}
|
||||
|
||||
func (t *WriteFileTool) UpgradeToConcurrent() Tool {
|
||||
return &WriteFileTool{
|
||||
fs: &ConcurrentFS{baseFS: t.fs},
|
||||
}
|
||||
}
|
||||
|
||||
type ListDirTool struct {
|
||||
fs fileSystem
|
||||
}
|
||||
|
|
@ -253,6 +266,7 @@ func formatDirEntries(entries []os.DirEntry) *ToolResult {
|
|||
type fileSystem interface {
|
||||
ReadFile(path string) ([]byte, error)
|
||||
WriteFile(path string, data []byte) error
|
||||
EditFile(path string, editFn func([]byte) ([]byte, error)) error
|
||||
ReadDir(path string) ([]os.DirEntry, error)
|
||||
}
|
||||
|
||||
|
|
@ -273,6 +287,23 @@ func (h *hostFs) ReadFile(path string) ([]byte, error) {
|
|||
return content, nil
|
||||
}
|
||||
|
||||
func (h *hostFs) EditFile(path string, editFn func([]byte) ([]byte, error)) error {
|
||||
|
||||
content, err := os.ReadFile(path)
|
||||
if err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("failed to read file for editing: %w", err)
|
||||
}
|
||||
// If it doesn't exist, we pass an empty byte slice to the editFn.
|
||||
// This is important for "append" operations which might create new files.
|
||||
|
||||
newContent, err := editFn(content)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return fileutil.WriteFileAtomic(path, newContent, 0o600)
|
||||
}
|
||||
|
||||
func (h *hostFs) ReadDir(path string) ([]os.DirEntry, error) {
|
||||
return os.ReadDir(path)
|
||||
}
|
||||
|
|
@ -381,6 +412,63 @@ func (r *sandboxFs) WriteFile(path string, data []byte) error {
|
|||
})
|
||||
}
|
||||
|
||||
func (r *sandboxFs) EditFile(path string, editFn func([]byte) ([]byte, error)) error {
|
||||
return r.execute(path, func(root *os.Root, relPath string) error {
|
||||
// 1. Read
|
||||
content, err := root.ReadFile(relPath)
|
||||
if err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("failed to read file for editing: %w", err)
|
||||
}
|
||||
|
||||
// 2. Modify
|
||||
newContent, err := editFn(content)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 3. Write (reusing the atomic write logic)
|
||||
dir := filepath.Dir(relPath)
|
||||
if dir != "." && dir != "/" {
|
||||
if err := root.MkdirAll(dir, 0o755); err != nil {
|
||||
return fmt.Errorf("failed to create parent directories: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
tmpRelPath := fmt.Sprintf(".tmp-%d-%d", os.Getpid(), time.Now().UnixNano())
|
||||
tmpFile, err := root.OpenFile(tmpRelPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
|
||||
if err != nil {
|
||||
root.Remove(tmpRelPath)
|
||||
return fmt.Errorf("failed to open temp file: %w", err)
|
||||
}
|
||||
|
||||
if _, err := tmpFile.Write(newContent); err != nil {
|
||||
tmpFile.Close()
|
||||
root.Remove(tmpRelPath)
|
||||
return fmt.Errorf("failed to write temp file: %w", err)
|
||||
}
|
||||
if err := tmpFile.Sync(); err != nil {
|
||||
tmpFile.Close()
|
||||
root.Remove(tmpRelPath)
|
||||
return fmt.Errorf("failed to sync temp file: %w", err)
|
||||
}
|
||||
if err := tmpFile.Close(); err != nil {
|
||||
root.Remove(tmpRelPath)
|
||||
return fmt.Errorf("failed to close temp file: %w", err)
|
||||
}
|
||||
|
||||
if err := root.Rename(tmpRelPath, relPath); err != nil {
|
||||
root.Remove(tmpRelPath)
|
||||
return fmt.Errorf("failed to rename temp file over target: %w", err)
|
||||
}
|
||||
if dirFile, err := root.Open("."); err == nil {
|
||||
_ = dirFile.Sync()
|
||||
dirFile.Close()
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (r *sandboxFs) ReadDir(path string) ([]os.DirEntry, error) {
|
||||
var entries []os.DirEntry
|
||||
err := r.execute(path, func(root *os.Root, relPath string) error {
|
||||
|
|
@ -415,3 +503,50 @@ func getSafeRelPath(workspace, path string) (string, error) {
|
|||
|
||||
return rel, nil
|
||||
}
|
||||
|
||||
// ConcurrencyUpgradeable indicates a Tool operates on files and can be upgraded
|
||||
// to use a thread-safe locking proxy backend (`ConcurrentFS`) for Parallel or DAG agent teams.
|
||||
type ConcurrencyUpgradeable interface {
|
||||
UpgradeToConcurrent() Tool
|
||||
}
|
||||
|
||||
// Global file locks explicitly for concurrent agent strategies
|
||||
var globalFileLocks sync.Map // map[string]*sync.RWMutex
|
||||
|
||||
func getPathLock(path string) *sync.RWMutex {
|
||||
cleanPath := filepath.Clean(path)
|
||||
actual, _ := globalFileLocks.LoadOrStore(cleanPath, &sync.RWMutex{})
|
||||
return actual.(*sync.RWMutex)
|
||||
}
|
||||
|
||||
// ConcurrentFS is a lightweight proxy wrapper around any `fileSystem`.
|
||||
// It guarantees thread-safe, race-condition-free access by locking the absolute file path globally.
|
||||
type ConcurrentFS struct {
|
||||
baseFS fileSystem
|
||||
}
|
||||
|
||||
func (c *ConcurrentFS) ReadFile(path string) ([]byte, error) {
|
||||
lock := getPathLock(path)
|
||||
lock.RLock()
|
||||
defer lock.RUnlock()
|
||||
return c.baseFS.ReadFile(path)
|
||||
}
|
||||
|
||||
func (c *ConcurrentFS) WriteFile(path string, data []byte) error {
|
||||
lock := getPathLock(path)
|
||||
lock.Lock()
|
||||
defer lock.Unlock()
|
||||
return c.baseFS.WriteFile(path, data)
|
||||
}
|
||||
|
||||
func (c *ConcurrentFS) EditFile(path string, editFn func([]byte) ([]byte, error)) error {
|
||||
lock := getPathLock(path)
|
||||
lock.Lock()
|
||||
defer lock.Unlock()
|
||||
return c.baseFS.EditFile(path, editFn)
|
||||
}
|
||||
|
||||
func (c *ConcurrentFS) ReadDir(path string) ([]os.DirEntry, error) {
|
||||
// Directories rarely suffer from single-file corruption, but we delegate anyway.
|
||||
return c.baseFS.ReadDir(path)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -486,3 +486,61 @@ func TestRootRW_Write(t *testing.T) {
|
|||
assert.NoError(t, err)
|
||||
assert.Equal(t, newData, content)
|
||||
}
|
||||
|
||||
// TestConcurrentFS_RaceCondition simulates a high-concurrency environment
|
||||
// to prove that the ConcurrentFS proxy correctly prevents file corruption.
|
||||
func TestConcurrentFS_RaceCondition(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
testFile := filepath.Join(tmpDir, "race_test.txt")
|
||||
|
||||
// Pre-fill file with "0"
|
||||
err := os.WriteFile(testFile, []byte("0"), 0o600)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create a base FS (hostFs) and wrap it in ConcurrentFS
|
||||
baseFS := &hostFs{}
|
||||
concurrentFs := &ConcurrentFS{baseFS: baseFS}
|
||||
|
||||
numGoroutines := 100
|
||||
done := make(chan bool)
|
||||
|
||||
// Simulate 100 goroutines trying to append/edit simultaneously
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func() {
|
||||
_ = concurrentFs.EditFile(testFile, func(content []byte) ([]byte, error) {
|
||||
// Artificial parsing of a number to increment
|
||||
newContent := append(content, []byte("-x")...)
|
||||
return newContent, nil
|
||||
})
|
||||
done <- true
|
||||
}()
|
||||
}
|
||||
|
||||
// Wait for all to finish
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
<-done
|
||||
}
|
||||
|
||||
// Verify the file isn't corrupted and has exactly 100 "-x" additions
|
||||
finalData, err := os.ReadFile(testFile)
|
||||
assert.NoError(t, err)
|
||||
|
||||
finalStr := string(finalData)
|
||||
xCount := strings.Count(finalStr, "-x")
|
||||
assert.Equal(t, numGoroutines, xCount, "Race condition detected! The file was corrupted or missed writes.")
|
||||
}
|
||||
|
||||
func TestConcurrencyUpgradeable(t *testing.T) {
|
||||
// Verify that ReadFileTool implements the interface and upgrades correctly
|
||||
readTool := NewReadFileTool("", false)
|
||||
upgradable, ok := interface{}(readTool).(ConcurrencyUpgradeable)
|
||||
assert.True(t, ok, "ReadFileTool should implement ConcurrencyUpgradeable")
|
||||
|
||||
upgradedTool := upgradable.UpgradeToConcurrent()
|
||||
upgradedReadTool, ok := upgradedTool.(*ReadFileTool)
|
||||
assert.True(t, ok, "Upgraded tool should still be a *ReadFileTool")
|
||||
|
||||
// Ensure the internal fs is now a *ConcurrentFS
|
||||
_, isConcurrent := upgradedReadTool.fs.(*ConcurrentFS)
|
||||
assert.True(t, isConcurrent, "Internal fileSystem should be upgraded to *ConcurrentFS")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -35,6 +35,17 @@ func (r *ToolRegistry) Get(name string) (Tool, bool) {
|
|||
return tool, ok
|
||||
}
|
||||
|
||||
// ListTools returns a slice of all registered tool names.
|
||||
func (r *ToolRegistry) ListTools() []string {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
names := make([]string, 0, len(r.tools))
|
||||
for name := range r.tools {
|
||||
names = append(names, name)
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
func (r *ToolRegistry) Execute(ctx context.Context, name string, args map[string]any) *ToolResult {
|
||||
return r.ExecuteWithContext(ctx, name, args, "", "", nil)
|
||||
}
|
||||
|
|
|
|||
124
pkg/tools/spawn_sub_agent.go
Normal file
124
pkg/tools/spawn_sub_agent.go
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
// SpawnSubAgentTool executes a customized subagent task synchronously using Anthology-style single worker delegation.
|
||||
type SpawnSubAgentTool struct {
|
||||
manager *SubagentManager
|
||||
originChannel string
|
||||
originChatID string
|
||||
}
|
||||
|
||||
func NewSpawnSubAgentTool(manager *SubagentManager) *SpawnSubAgentTool {
|
||||
return &SpawnSubAgentTool{
|
||||
manager: manager,
|
||||
originChannel: "cli",
|
||||
originChatID: "direct",
|
||||
}
|
||||
}
|
||||
|
||||
func (t *SpawnSubAgentTool) Name() string {
|
||||
return "spawn_sub_agent"
|
||||
}
|
||||
|
||||
func (t *SpawnSubAgentTool) Description() string {
|
||||
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 {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"task": map[string]any{
|
||||
"type": "string",
|
||||
"description": "The specific task the sub-agent needs to accomplish.",
|
||||
},
|
||||
"role": map[string]any{
|
||||
"type": "string",
|
||||
"description": "The system prompt/role assignment for the sub-agent (e.g., 'You are an expert code reviewer').",
|
||||
},
|
||||
"model": map[string]any{
|
||||
"type": "string",
|
||||
"description": "Optional specific LLM model ID to route this task to (e.g., 'gpt-4o' for vision, 'claude-3-5-sonnet' for logic). If omitted, inherits the parent's model.",
|
||||
},
|
||||
},
|
||||
"required": []string{"task", "role"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *SpawnSubAgentTool) SetContext(channel, chatID string) {
|
||||
t.originChannel = channel
|
||||
t.originChatID = chatID
|
||||
}
|
||||
|
||||
func (t *SpawnSubAgentTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
|
||||
task, ok := args["task"].(string)
|
||||
if !ok || strings.TrimSpace(task) == "" {
|
||||
return ErrorResult("task is required").WithError(fmt.Errorf("task parameter is required"))
|
||||
}
|
||||
|
||||
role, ok := args["role"].(string)
|
||||
if !ok || strings.TrimSpace(role) == "" {
|
||||
return ErrorResult("role is required").WithError(fmt.Errorf("role parameter is required"))
|
||||
}
|
||||
|
||||
if t.manager == nil {
|
||||
return ErrorResult("Subagent manager not configured").WithError(fmt.Errorf("manager is nil"))
|
||||
}
|
||||
|
||||
// 1. Isolation: Each SubAgent gets a completely fresh message set
|
||||
messages := []providers.Message{
|
||||
{
|
||||
Role: "system",
|
||||
Content: role,
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Content: task,
|
||||
},
|
||||
}
|
||||
|
||||
// 2. Base Configuration (Timeout & LLM constraints)
|
||||
config := t.manager.BuildBaseWorkerConfig(ctx)
|
||||
|
||||
// 2.1 Model Override (Heterogeneous Agents)
|
||||
if modelParam, ok := args["model"].(string); ok && 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.
|
||||
// To enforce strict sandboxing later, we can construct a new ToolRegistry here based on args['allowed_tools'].
|
||||
|
||||
loopResult, err := RunToolLoop(ctx, config, messages, t.originChannel, t.originChatID)
|
||||
if err != nil {
|
||||
return ErrorResult(fmt.Sprintf("Subagent execution failed: %v", err)).WithError(err)
|
||||
}
|
||||
|
||||
// Return full details to LLM
|
||||
llmContent := fmt.Sprintf("Subagent (Role: %s) task completed:\nIterations: %d\nResult: %s",
|
||||
role, loopResult.Iterations, loopResult.Content)
|
||||
|
||||
return &ToolResult{
|
||||
ForLLM: llmContent,
|
||||
ForUser: "Sub-agent finished task.",
|
||||
Silent: false,
|
||||
IsError: false,
|
||||
Async: false,
|
||||
}
|
||||
}
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ package tools
|
|||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
|
|
@ -10,6 +11,28 @@ 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)
|
||||
ModelTagImageGen = "image-gen" // Supports image generation output (e.g. DALL-E, Stable Diffusion)
|
||||
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 (multimodal input)",
|
||||
ModelTagImageGen: "can generate images from text descriptions (e.g. DALL-E, Stable Diffusion)",
|
||||
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 +50,7 @@ type SubagentManager struct {
|
|||
mu sync.RWMutex
|
||||
provider providers.LLMProvider
|
||||
defaultModel string
|
||||
allowedModels []providers.FallbackCandidate
|
||||
bus *bus.MessageBus
|
||||
workspace string
|
||||
tools *ToolRegistry
|
||||
|
|
@ -40,13 +64,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 +82,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()
|
||||
|
|
@ -248,6 +321,31 @@ func (sm *SubagentManager) ListTasks() []*SubagentTask {
|
|||
return tasks
|
||||
}
|
||||
|
||||
// BuildBaseWorkerConfig returns a base ToolLoopConfig that can be customized for isolated workers.
|
||||
func (sm *SubagentManager) BuildBaseWorkerConfig(ctx context.Context) ToolLoopConfig {
|
||||
sm.mu.RLock()
|
||||
defer sm.mu.RUnlock()
|
||||
|
||||
var llmOptions map[string]any
|
||||
if sm.hasMaxTokens || sm.hasTemperature {
|
||||
llmOptions = map[string]any{}
|
||||
if sm.hasMaxTokens {
|
||||
llmOptions["max_tokens"] = sm.maxTokens
|
||||
}
|
||||
if sm.hasTemperature {
|
||||
llmOptions["temperature"] = sm.temperature
|
||||
}
|
||||
}
|
||||
|
||||
return ToolLoopConfig{
|
||||
Provider: sm.provider,
|
||||
Model: sm.defaultModel,
|
||||
Tools: sm.tools, // Note: Caller should replace this for isolated registry
|
||||
MaxIterations: sm.maxIterations,
|
||||
LLMOptions: llmOptions,
|
||||
}
|
||||
}
|
||||
|
||||
// SubagentTool executes a subagent task synchronously and returns the result.
|
||||
// Unlike SpawnTool which runs tasks asynchronously, SubagentTool waits for completion
|
||||
// and returns the result directly in the ToolResult.
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
740
pkg/tools/team.go
Normal file
740
pkg/tools/team.go
Normal file
|
|
@ -0,0 +1,740 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
type TeamTool struct {
|
||||
manager *SubagentManager
|
||||
originChannel string
|
||||
originChatID string
|
||||
}
|
||||
|
||||
type TeamMember struct {
|
||||
ID string
|
||||
Role string
|
||||
Task string
|
||||
Model string // Heterogeneous Agents: Optional specific model for this task
|
||||
DependsOn []string // List of member IDs this member depends on
|
||||
Produces string // Auto-reviewer: declares artifact type ("code", "data", "document")
|
||||
}
|
||||
|
||||
func NewTeamTool(manager *SubagentManager) *TeamTool {
|
||||
return &TeamTool{
|
||||
manager: manager,
|
||||
originChannel: "cli",
|
||||
originChatID: "direct",
|
||||
}
|
||||
}
|
||||
|
||||
func (t *TeamTool) Name() string {
|
||||
return "team"
|
||||
}
|
||||
|
||||
func (t *TeamTool) Description() string {
|
||||
base := `Compose and execute a team of specialized sub-agents to accomplish a complex task.
|
||||
|
||||
WHEN TO USE THIS TOOL (use proactively — do not attempt to handle these alone):
|
||||
- The task involves 2 or more distinct areas of concern (e.g. research + writing, coding + testing, data gathering + analysis).
|
||||
- The task would require more than 5 consecutive tool calls if done alone.
|
||||
- Any part of the task can be done in parallel to save time.
|
||||
- The task is large enough that a single agent would likely lose context or quality midway.
|
||||
- The user asks you to "build", "create", "generate", "analyze", or "convert" something non-trivial.
|
||||
When in doubt, prefer delegation over doing everything yourself.
|
||||
|
||||
CRITICAL RULES FOR TASK PLANNING:
|
||||
1. Think like a project manager: analyze the full task first, then design the team structure before spawning anyone.
|
||||
2. Decompose the task into the smallest independently-ownable units of work. A member should own exactly ONE distinct concern — not a broad compound goal.
|
||||
3. Identify dependencies between units: if one member's output is required by another, declare it via 'depends_on'. Independent units should run concurrently.
|
||||
4. Each member's 'task' must be precise and self-contained. Include relevant context (e.g. reference to outputs from dependencies) directly in the task description.
|
||||
5. Sub-agents are full agents with access to the same tools, including this 'team' tool. If a member's sub-task is itself complex, it may recursively form its own team.
|
||||
|
||||
Strategy guide:
|
||||
- sequential: each step depends on the full output of the previous step in a strict chain.
|
||||
- parallel: all tasks are fully independent with no shared inputs or outputs.
|
||||
- dag: most real-world tasks — some tasks depend on others, some can run concurrently.
|
||||
- evaluator_optimizer: the output needs iterative critique and revision cycles.`
|
||||
|
||||
if t.manager != nil {
|
||||
if hint := t.manager.ModelCapabilityHint(); hint != "" {
|
||||
return base + "\n\n" + hint
|
||||
}
|
||||
}
|
||||
return base
|
||||
}
|
||||
|
||||
func (t *TeamTool) Parameters() map[string]any {
|
||||
return map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"strategy": map[string]any{
|
||||
"type": "string",
|
||||
"enum": []string{"sequential", "parallel", "dag", "evaluator_optimizer"},
|
||||
"description": "How to run the team members. 'sequential': one after another. 'parallel': all at once. 'dag': execute based on declared dependencies. 'evaluator_optimizer': EXACTLY two members (worker & evaluator). The evaluator will check the worker's output; if it fails, the worker is revived with its FULL stateful memory intact and asked to fix it. Use this for complex generation tasks (like coding) requiring deep reasoning.",
|
||||
},
|
||||
"max_team_tokens": map[string]any{
|
||||
"type": "integer",
|
||||
"description": "The maximum combined LLM tokens (prompt + completion) this entire team is allowed to consume. Once exceeded, the team is instantly killed.",
|
||||
},
|
||||
"members": map[string]any{
|
||||
"type": "array",
|
||||
"description": "The list of sub-agents in the team.",
|
||||
"items": map[string]any{
|
||||
"type": "object",
|
||||
"properties": map[string]any{
|
||||
"id": map[string]any{
|
||||
"type": "string",
|
||||
"description": "Unique identifier for this member, used for dependencies in 'dag' strategy.",
|
||||
},
|
||||
"role": map[string]any{
|
||||
"type": "string",
|
||||
"description": "The system prompt/role assignment for the member.",
|
||||
},
|
||||
"task": map[string]any{
|
||||
"type": "string",
|
||||
"description": "The specific task this member needs to accomplish.",
|
||||
},
|
||||
"model": map[string]any{
|
||||
"type": "string",
|
||||
"description": "Optional specific LLM model ID to route this task to (e.g., 'gpt-4o' for vision, 'claude-3-5-sonnet' for logic). If omitted, inherits the parent's model.",
|
||||
},
|
||||
"depends_on": map[string]any{
|
||||
"type": "array",
|
||||
"description": "List of 'id' strings this member depends on. Only applicable for 'dag' strategy.",
|
||||
"items": map[string]any{"type": "string"},
|
||||
},
|
||||
"produces": map[string]any{
|
||||
"type": "string",
|
||||
"description": "Declares the type of artifact this member produces. Use 'code' for source code files, 'data' for structured data/JSON/CSV, 'document' for prose documents/reports. When set, the framework automatically appends a QA reviewer step after all workers finish to validate output correctness. Omit if no verification is needed.",
|
||||
},
|
||||
},
|
||||
"required": []string{"role", "task"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": []string{"strategy", "members"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *TeamTool) SetContext(channel, chatID string) {
|
||||
t.originChannel = channel
|
||||
t.originChatID = chatID
|
||||
}
|
||||
|
||||
// reviewerTaskTemplates maps a `produces` artifact type to the task prompt
|
||||
// that the auto-injected QA reviewer will receive.
|
||||
var reviewerTaskTemplates = map[string]string{
|
||||
"code": "You are a code quality reviewer. Read all code files in the workspace that were just written by your predecessors. Check for: syntax errors, incorrect or missing imports, broken logic, type mismatches, and any issues that would cause compilation or runtime failures. List every issue found with the filename and line number if possible. If everything looks correct, respond with 'REVIEW PASSED'.",
|
||||
"data": "You are a data validation reviewer. Read all output data files (JSON, CSV, YAML, etc.) in the workspace. Check for: invalid format, missing required fields, schema inconsistencies, and malformed values. List every issue found. If everything is valid, respond with 'REVIEW PASSED'.",
|
||||
"document": "You are a document quality reviewer. Read all output documents in the workspace. Check for: logical inconsistencies, incomplete sections, factual contradictions, and poor structure. List every issue found. If the documents are complete and correct, respond with 'REVIEW PASSED'.",
|
||||
}
|
||||
|
||||
// maybeRunAutoReviewer inspects TeamMembers for `produces` declarations.
|
||||
// If any member produced a verifiable artifact type, it runs an automatic
|
||||
// QA reviewer agent after all workers have completed.
|
||||
func (t *TeamTool) maybeRunAutoReviewer(
|
||||
ctx context.Context,
|
||||
members []TeamMember,
|
||||
baseConfig ToolLoopConfig,
|
||||
workerSummary string,
|
||||
) string {
|
||||
// Collect unique produces types from all members
|
||||
producedTypes := make(map[string]bool)
|
||||
for _, m := range members {
|
||||
if m.Produces != "" {
|
||||
producedTypes[m.Produces] = true
|
||||
}
|
||||
}
|
||||
if len(producedTypes) == 0 {
|
||||
return "" // No verifiable artifacts declared, skip review
|
||||
}
|
||||
|
||||
// Build reviewer task: combine templates for all declared artifact types
|
||||
var taskParts []string
|
||||
for artifactType := range producedTypes {
|
||||
if tmpl, ok := reviewerTaskTemplates[artifactType]; ok {
|
||||
taskParts = append(taskParts, tmpl)
|
||||
}
|
||||
}
|
||||
if len(taskParts) == 0 {
|
||||
return "" // Unknown produces types, skip
|
||||
}
|
||||
|
||||
reviewerTask := strings.Join(taskParts, "\n\n") +
|
||||
"\n\nContext from the workers that produced these artifacts:\n" + workerSummary
|
||||
|
||||
reviewerMessages := []providers.Message{
|
||||
{Role: "user", Content: reviewerTask},
|
||||
}
|
||||
|
||||
reviewerConfig := baseConfig
|
||||
loopResult, err := RunToolLoop(ctx, reviewerConfig, reviewerMessages, t.originChannel, t.originChatID)
|
||||
if err != nil {
|
||||
return fmt.Sprintf("[Auto-Reviewer] Failed to run: %v", err)
|
||||
}
|
||||
return "[Auto-Reviewer Result]\n" + loopResult.Content
|
||||
}
|
||||
|
||||
func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
|
||||
strategy, ok := args["strategy"].(string)
|
||||
if !ok || (strategy != "sequential" && strategy != "parallel" && strategy != "dag" && strategy != "evaluator_optimizer") {
|
||||
return ErrorResult("strategy must be 'sequential', 'parallel', 'dag', or 'evaluator_optimizer'")
|
||||
}
|
||||
|
||||
membersRaw, ok := args["members"].([]any)
|
||||
if !ok || len(membersRaw) == 0 {
|
||||
return ErrorResult("members map array is required and must not be empty")
|
||||
}
|
||||
|
||||
maxTokensFloat, ok := args["max_team_tokens"].(float64)
|
||||
var budget *atomic.Int64
|
||||
if ok && maxTokensFloat > 0 {
|
||||
budget = &atomic.Int64{}
|
||||
budget.Store(int64(maxTokensFloat))
|
||||
}
|
||||
|
||||
if t.manager == nil {
|
||||
return ErrorResult("Subagent manager not configured").WithError(fmt.Errorf("manager is nil"))
|
||||
}
|
||||
|
||||
var members []TeamMember
|
||||
for i, mRaw := range membersRaw {
|
||||
mMap, ok := mRaw.(map[string]any)
|
||||
if !ok {
|
||||
return ErrorResult(fmt.Sprintf("member at index %d is invalid", i))
|
||||
}
|
||||
|
||||
id, iOk := mMap["id"].(string)
|
||||
role, rOk := mMap["role"].(string)
|
||||
task, tOk := mMap["task"].(string)
|
||||
|
||||
if !rOk || !tOk || strings.TrimSpace(role) == "" || strings.TrimSpace(task) == "" {
|
||||
return ErrorResult(fmt.Sprintf("member at index %d is missing required 'role' or 'task'", i))
|
||||
}
|
||||
|
||||
// ID is highly recommended, generate one if missing for backwards compatibility
|
||||
if !iOk || strings.TrimSpace(id) == "" {
|
||||
id = fmt.Sprintf("member_%d", i)
|
||||
}
|
||||
|
||||
modelStr, _ := mMap["model"].(string)
|
||||
modelStr = strings.TrimSpace(modelStr)
|
||||
|
||||
var dependsOn []string
|
||||
if depRaw, dOk := mMap["depends_on"].([]any); dOk {
|
||||
for _, d := range depRaw {
|
||||
if dStr, dsOk := d.(string); dsOk {
|
||||
dependsOn = append(dependsOn, dStr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
producesStr, _ := mMap["produces"].(string)
|
||||
producesStr = strings.TrimSpace(producesStr)
|
||||
|
||||
members = append(members, TeamMember{
|
||||
ID: id,
|
||||
Role: role,
|
||||
Task: task,
|
||||
Model: modelStr,
|
||||
DependsOn: dependsOn,
|
||||
Produces: producesStr,
|
||||
})
|
||||
}
|
||||
|
||||
// Base struct setup
|
||||
baseConfig := t.manager.BuildBaseWorkerConfig(ctx)
|
||||
if budget != nil {
|
||||
baseConfig.RemainingTokenBudget = budget
|
||||
}
|
||||
|
||||
// Create a new master context for team bounding
|
||||
// In the future this could be overridden by an argument
|
||||
teamCtx, cancel := context.WithTimeout(ctx, 15*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
// If strategy is parallel or dag, we must upgrade the file tools to be concurrent-safe (locking)
|
||||
if strategy == "parallel" || strategy == "dag" {
|
||||
baseConfig.Tools = upgradeRegistryForConcurrency(baseConfig.Tools)
|
||||
}
|
||||
|
||||
switch strategy {
|
||||
case "sequential":
|
||||
result := t.executeSequential(teamCtx, baseConfig, members)
|
||||
if reviewNote := t.maybeRunAutoReviewer(teamCtx, members, baseConfig, result.ForLLM); reviewNote != "" {
|
||||
result.ForLLM += "\n\n" + reviewNote
|
||||
result.ForUser += "\n\n" + reviewNote
|
||||
}
|
||||
return result
|
||||
case "dag":
|
||||
result := t.executeDAG(teamCtx, baseConfig, members)
|
||||
if reviewNote := t.maybeRunAutoReviewer(teamCtx, members, baseConfig, result.ForLLM); reviewNote != "" {
|
||||
result.ForLLM += "\n\n" + reviewNote
|
||||
result.ForUser += "\n\n" + reviewNote
|
||||
}
|
||||
return result
|
||||
case "evaluator_optimizer":
|
||||
return t.executeEvaluatorOptimizer(teamCtx, baseConfig, members)
|
||||
}
|
||||
// parallel
|
||||
result := t.executeParallel(teamCtx, baseConfig, members)
|
||||
if reviewNote := t.maybeRunAutoReviewer(teamCtx, members, baseConfig, result.ForLLM); reviewNote != "" {
|
||||
result.ForLLM += "\n\n" + reviewNote
|
||||
result.ForUser += "\n\n" + reviewNote
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// upgradeRegistryForConcurrency takes an existing ToolRegistry, clones it,
|
||||
// and upgrades any tools that implement ConcurrencyUpgradeable to their locking counterparts.
|
||||
func upgradeRegistryForConcurrency(original *ToolRegistry) *ToolRegistry {
|
||||
if original == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
upgraded := NewToolRegistry()
|
||||
for _, name := range original.ListTools() {
|
||||
tool, ok := original.Get(name)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
if upgradeable, isUpgradeable := tool.(ConcurrencyUpgradeable); isUpgradeable {
|
||||
upgraded.Register(upgradeable.UpgradeToConcurrent())
|
||||
} else {
|
||||
upgraded.Register(tool)
|
||||
}
|
||||
}
|
||||
return upgraded
|
||||
}
|
||||
|
||||
// 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, 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, nil
|
||||
}
|
||||
|
||||
func (t *TeamTool) executeSequential(ctx context.Context, baseConfig ToolLoopConfig, members []TeamMember) *ToolResult {
|
||||
var finalOutput strings.Builder
|
||||
finalOutput.WriteString("Team Execution Summary (Sequential):\n\n")
|
||||
|
||||
var previousResult string
|
||||
|
||||
for i, m := range members {
|
||||
// If there is a previous result, we append it to the task so the new agent sees it.
|
||||
actualTask := m.Task
|
||||
if i > 0 && previousResult != "" {
|
||||
actualTask = fmt.Sprintf("%s\n\n--- Context from previous phase ---\n%s", m.Task, truncateContext(previousResult))
|
||||
}
|
||||
|
||||
messages := []providers.Message{
|
||||
{Role: "system", Content: m.Role},
|
||||
{Role: "user", Content: actualTask},
|
||||
}
|
||||
|
||||
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)
|
||||
finalOutput.WriteString(errStr + "\n")
|
||||
return ErrorResult(errStr).WithError(err) // Fail fast
|
||||
}
|
||||
|
||||
previousResult = loopResult.Content
|
||||
|
||||
finalOutput.WriteString(fmt.Sprintf("### Phase %d completed by Role: [%s]\n%s\n\n", i+1, m.Role, previousResult))
|
||||
}
|
||||
|
||||
return &ToolResult{
|
||||
ForLLM: finalOutput.String(),
|
||||
ForUser: "Team completed sequential execution successfully.",
|
||||
}
|
||||
}
|
||||
|
||||
func (t *TeamTool) executeParallel(ctx context.Context, baseConfig ToolLoopConfig, members []TeamMember) *ToolResult {
|
||||
var wg sync.WaitGroup
|
||||
type workResult struct {
|
||||
index int
|
||||
role string
|
||||
res string
|
||||
err error
|
||||
}
|
||||
|
||||
resultsChan := make(chan workResult, len(members))
|
||||
|
||||
for i, m := range members {
|
||||
wg.Add(1)
|
||||
go func(index int, member TeamMember) {
|
||||
defer wg.Done()
|
||||
|
||||
messages := []providers.Message{
|
||||
{Role: "system", Content: member.Role},
|
||||
{Role: "user", Content: member.Task},
|
||||
}
|
||||
|
||||
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 {
|
||||
resultsChan <- workResult{index: index, role: member.Role, err: err}
|
||||
return
|
||||
}
|
||||
resultsChan <- workResult{index: index, role: member.Role, res: loopResult.Content}
|
||||
}(i, m)
|
||||
}
|
||||
|
||||
// Wait for all goroutines to finish
|
||||
wg.Wait()
|
||||
close(resultsChan)
|
||||
|
||||
var finalOutput strings.Builder
|
||||
finalOutput.WriteString("Team Execution Summary (Parallel):\n\n")
|
||||
|
||||
// Pre-allocate to maintain order since channels don't guarantee arrival order
|
||||
orderedResults := make([]workResult, len(members))
|
||||
for res := range resultsChan {
|
||||
orderedResults[res.index] = res
|
||||
}
|
||||
|
||||
hasError := false
|
||||
for _, res := range orderedResults {
|
||||
if res.err != nil {
|
||||
hasError = true
|
||||
finalOutput.WriteString(fmt.Sprintf("### Worker [%s] FAILED:\n%v\n\n", res.role, res.err))
|
||||
} else {
|
||||
finalOutput.WriteString(fmt.Sprintf("### Worker [%s] Output:\n%s\n\n", res.role, res.res))
|
||||
}
|
||||
}
|
||||
|
||||
if hasError {
|
||||
return ErrorResult("One or more parallel workers failed.\n" + finalOutput.String())
|
||||
}
|
||||
|
||||
return &ToolResult{
|
||||
ForLLM: finalOutput.String(),
|
||||
ForUser: "Team completed parallel execution.",
|
||||
}
|
||||
}
|
||||
|
||||
func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, baseConfig ToolLoopConfig, members []TeamMember) *ToolResult {
|
||||
if len(members) != 2 {
|
||||
return ErrorResult("The evaluator_optimizer strategy requires exactly two members: [0] Worker, [1] Evaluator.")
|
||||
}
|
||||
|
||||
worker := members[0]
|
||||
evaluator := members[1]
|
||||
|
||||
var finalOutput strings.Builder
|
||||
finalOutput.WriteString("Team Execution Summary (Evaluator-Optimizer):\n\n")
|
||||
|
||||
// 1. Initialize the stateful memory for the worker
|
||||
workerMessages := []providers.Message{
|
||||
{Role: "system", Content: worker.Role},
|
||||
{Role: "user", Content: worker.Task},
|
||||
}
|
||||
|
||||
maxLoops := 5
|
||||
for attempt := 1; attempt <= maxLoops; attempt++ {
|
||||
finalOutput.WriteString(fmt.Sprintf("## Attempt %d\n", attempt))
|
||||
|
||||
// 2. Trigger Worker (resumes from its exact previous state!)
|
||||
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)
|
||||
finalOutput.WriteString(errStr + "\n")
|
||||
return ErrorResult(errStr).WithError(err)
|
||||
}
|
||||
|
||||
// Save the worker's cognitive state so it remembers its thought process for the next loop
|
||||
workerMessages = workerResult.Messages
|
||||
|
||||
finalOutput.WriteString(fmt.Sprintf("### Worker Output:\n%s\n\n", workerResult.Content))
|
||||
|
||||
// 3. Trigger Evaluator (Ephemeral, stateless evaluation)
|
||||
evalContext := fmt.Sprintf("%s\n\n--- Worker's Output to Evaluate ---\n%s\n\nIf the output is completely correct and fulfills the task, you MUST reply starting with strictly '[PASS]'. Otherwise, explain the issues in detail.", evaluator.Task, truncateContext(workerResult.Content))
|
||||
|
||||
evalMessages := []providers.Message{
|
||||
{Role: "system", Content: evaluator.Role},
|
||||
{Role: "user", Content: evalContext},
|
||||
}
|
||||
|
||||
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)
|
||||
finalOutput.WriteString(errStr + "\n")
|
||||
return ErrorResult(errStr).WithError(err)
|
||||
}
|
||||
|
||||
finalOutput.WriteString(fmt.Sprintf("### Evaluator Feedback:\n%s\n\n", evalResult.Content))
|
||||
|
||||
// 4. Check for PASS condition
|
||||
if strings.HasPrefix(strings.TrimSpace(evalResult.Content), "[PASS]") {
|
||||
finalOutput.WriteString("✅ Evaluation Passed! Loop finished successfully.\n")
|
||||
return &ToolResult{
|
||||
ForLLM: finalOutput.String(),
|
||||
ForUser: "Evaluator-Optimizer loop completed successfully.",
|
||||
}
|
||||
}
|
||||
|
||||
// 5. If not passed, and not the last attempt, inject feedback into Worker's stateful memory
|
||||
if attempt < maxLoops {
|
||||
injection := fmt.Sprintf("The evaluator rejected your previous attempt. Please fix the issues based on this feedback:\n\n%s", evalResult.Content)
|
||||
workerMessages = append(workerMessages, providers.Message{
|
||||
Role: "user",
|
||||
Content: injection,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
finalOutput.WriteString("❌ Maximum evaluation loops reached without a [PASS]. Returning current state.\n")
|
||||
return &ToolResult{
|
||||
ForLLM: finalOutput.String(),
|
||||
ForUser: "Evaluator-Optimizer loop exhausted maximum attempts.",
|
||||
}
|
||||
}
|
||||
|
||||
func (t *TeamTool) executeDAG(ctx context.Context, baseConfig ToolLoopConfig, members []TeamMember) *ToolResult {
|
||||
// 1. Build and VALIDATE dependency graph
|
||||
memberMap := make(map[string]TeamMember)
|
||||
inDegree := make(map[string]int)
|
||||
graph := make(map[string][]string) // node -> nodes that depend on it
|
||||
|
||||
// Register all valid members first
|
||||
for _, m := range members {
|
||||
memberMap[m.ID] = m
|
||||
inDegree[m.ID] = 0
|
||||
graph[m.ID] = []string{}
|
||||
}
|
||||
|
||||
// Build edges and check for ghost nodes
|
||||
for _, m := range members {
|
||||
for _, dep := range m.DependsOn {
|
||||
if _, exists := memberMap[dep]; !exists {
|
||||
return ErrorResult(fmt.Sprintf("DAG Validation Error: Member [%s] depends on undefined member [%s]", m.ID, dep))
|
||||
}
|
||||
graph[dep] = append(graph[dep], m.ID)
|
||||
inDegree[m.ID]++
|
||||
}
|
||||
}
|
||||
|
||||
// 1.5. Cycle Detection using Kahn's Algorithm
|
||||
var kahnQueue []string
|
||||
kahnInDegree := make(map[string]int)
|
||||
for k, v := range inDegree {
|
||||
kahnInDegree[k] = v
|
||||
if v == 0 {
|
||||
kahnQueue = append(kahnQueue, k)
|
||||
}
|
||||
}
|
||||
|
||||
processedCount := 0
|
||||
for len(kahnQueue) > 0 {
|
||||
curr := kahnQueue[0]
|
||||
kahnQueue = kahnQueue[1:]
|
||||
processedCount++
|
||||
|
||||
for _, dependent := range graph[curr] {
|
||||
kahnInDegree[dependent]--
|
||||
if kahnInDegree[dependent] == 0 {
|
||||
kahnQueue = append(kahnQueue, dependent)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if processedCount != len(members) {
|
||||
return ErrorResult("DAG Validation Error: Circular dependency (cycle) detected in the team layout. Please fix your 'depends_on' definitions.")
|
||||
}
|
||||
|
||||
// 2. Channels for coordination
|
||||
type nodeResult struct {
|
||||
id string
|
||||
res string
|
||||
err error
|
||||
}
|
||||
readyChan := make(chan string, len(members))
|
||||
resultChan := make(chan nodeResult, len(members))
|
||||
|
||||
// Channels specifically for passing context from dependencies to dependants
|
||||
contextMap := make(map[string]*strings.Builder)
|
||||
var contextMu sync.Mutex
|
||||
|
||||
// 3. Initialize queue with nodes having 0 in-degree
|
||||
nodesToProcess := len(members)
|
||||
for id, deg := range inDegree {
|
||||
if deg == 0 {
|
||||
readyChan <- id
|
||||
}
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var masterErr error
|
||||
var masterErrMu sync.Mutex
|
||||
|
||||
// Shared results store for the final output
|
||||
finalResults := make(map[string]string)
|
||||
var finalResultsMu sync.Mutex
|
||||
|
||||
// 4. DAG Execution Loop
|
||||
for i := 0; i < nodesToProcess; i++ {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ErrorResult("DAG execution timed out or cancelled")
|
||||
|
||||
case memberID := <-readyChan:
|
||||
wg.Add(1)
|
||||
go func(id string) {
|
||||
defer wg.Done()
|
||||
|
||||
m := memberMap[id]
|
||||
|
||||
// Construct the task with context from all dependencies
|
||||
actualTask := m.Task
|
||||
contextMu.Lock()
|
||||
b := contextMap[id]
|
||||
depsContext := ""
|
||||
if b != nil {
|
||||
depsContext = b.String()
|
||||
}
|
||||
contextMu.Unlock()
|
||||
|
||||
if depsContext != "" {
|
||||
actualTask = fmt.Sprintf("%s\n\n--- Context from dependencies ---\n%s", m.Task, truncateContext(depsContext))
|
||||
}
|
||||
|
||||
messages := []providers.Message{
|
||||
{Role: "system", Content: m.Role},
|
||||
{Role: "user", Content: actualTask},
|
||||
}
|
||||
|
||||
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 {
|
||||
masterErrMu.Lock()
|
||||
if masterErr == nil {
|
||||
masterErr = fmt.Errorf("worker [%s] failed: %v", m.ID, err)
|
||||
}
|
||||
masterErrMu.Unlock()
|
||||
resultChan <- nodeResult{id: id, err: err}
|
||||
return
|
||||
}
|
||||
|
||||
// Store result for final output
|
||||
finalResultsMu.Lock()
|
||||
finalResults[id] = loopResult.Content
|
||||
finalResultsMu.Unlock()
|
||||
|
||||
// Pass result to dependents
|
||||
resultChan <- nodeResult{id: id, res: loopResult.Content}
|
||||
}(memberID)
|
||||
|
||||
case res := <-resultChan:
|
||||
if res.err != nil {
|
||||
// Fast fail on first error
|
||||
return ErrorResult(res.err.Error())
|
||||
}
|
||||
|
||||
// Update dependents
|
||||
for _, dependentID := range graph[res.id] {
|
||||
contextMu.Lock()
|
||||
b := contextMap[dependentID]
|
||||
if b == nil {
|
||||
b = &strings.Builder{}
|
||||
}
|
||||
b.WriteString(fmt.Sprintf("--- Result from [%s] ---\n%s\n\n", res.id, res.res))
|
||||
contextMap[dependentID] = b
|
||||
contextMu.Unlock()
|
||||
|
||||
inDegree[dependentID]--
|
||||
if inDegree[dependentID] == 0 {
|
||||
readyChan <- dependentID
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for any remaining goroutines (though the select loop handles the exact count)
|
||||
wg.Wait()
|
||||
|
||||
if masterErr != nil {
|
||||
return ErrorResult(masterErr.Error())
|
||||
}
|
||||
|
||||
// 5. Format final output
|
||||
var finalOutput strings.Builder
|
||||
finalOutput.WriteString("Team Execution Summary (DAG):\n\n")
|
||||
|
||||
// Preserve original member order for final output readability
|
||||
for _, m := range members {
|
||||
if res, ok := finalResults[m.ID]; ok {
|
||||
finalOutput.WriteString(fmt.Sprintf("### Worker [%s] (Role: %s) Output:\n%s\n\n", m.ID, m.Role, res))
|
||||
}
|
||||
}
|
||||
|
||||
return &ToolResult{
|
||||
ForLLM: finalOutput.String(),
|
||||
ForUser: "Team completed DAG execution.",
|
||||
}
|
||||
}
|
||||
|
||||
// truncateContext prevents Context Window Explosion (Token Bombs)
|
||||
// by limiting the size of upstream results injected into downstream prompts.
|
||||
func truncateContext(ctx string) string {
|
||||
maxRunes := 8000
|
||||
runes := []rune(ctx)
|
||||
if len(runes) > maxRunes {
|
||||
return string(runes[:maxRunes]) + "\n...[Context truncated due to length]..."
|
||||
}
|
||||
return ctx
|
||||
}
|
||||
62
pkg/tools/team_test.go
Normal file
62
pkg/tools/team_test.go
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestUpgradeRegistryForConcurrency(t *testing.T) {
|
||||
// Create a standard tool registry
|
||||
original := NewToolRegistry()
|
||||
|
||||
// Register a mix of tools: some upgradeable, some not
|
||||
readTool := NewReadFileTool("", false)
|
||||
listTool := NewListDirTool("", false) // Not upgradeable
|
||||
writeTool := NewWriteFileTool("", false)
|
||||
|
||||
original.Register(readTool)
|
||||
original.Register(listTool)
|
||||
original.Register(writeTool)
|
||||
|
||||
// Perform the upgrade
|
||||
upgraded := upgradeRegistryForConcurrency(original)
|
||||
|
||||
// Verify count matches
|
||||
assert.Equal(t, len(original.ListTools()), len(upgraded.ListTools()), "Upgraded registry should have same number of tools")
|
||||
|
||||
// Verify ReadFileTool got upgraded
|
||||
actualReadTool, ok := upgraded.Get("read_file")
|
||||
assert.True(t, ok)
|
||||
if upgradedRead, isUpgraded := actualReadTool.(*ReadFileTool); isUpgraded {
|
||||
_, isConcurrent := upgradedRead.fs.(*ConcurrentFS)
|
||||
assert.True(t, isConcurrent, "read_file should have been upgraded to ConcurrentFS")
|
||||
}
|
||||
|
||||
// Verify WriteFileTool got upgraded
|
||||
actualWriteTool, ok := upgraded.Get("write_file")
|
||||
assert.True(t, ok)
|
||||
if upgradedWrite, isUpgraded := actualWriteTool.(*WriteFileTool); isUpgraded {
|
||||
_, isConcurrent := upgradedWrite.fs.(*ConcurrentFS)
|
||||
assert.True(t, isConcurrent, "write_file should have been upgraded to ConcurrentFS")
|
||||
}
|
||||
|
||||
// Verify ListDirTool remained the same
|
||||
actualListTool, ok := upgraded.Get("list_dir")
|
||||
assert.True(t, ok)
|
||||
_, isListDir := actualListTool.(*ListDirTool)
|
||||
assert.True(t, isListDir, "list_dir should still be ListDirTool")
|
||||
|
||||
// Double check list_dir doesn't randomly have ConcurrentFS injected
|
||||
if listImpl, ok := actualListTool.(*ListDirTool); ok {
|
||||
_, isConcurrent := listImpl.fs.(*ConcurrentFS)
|
||||
assert.False(t, isConcurrent, "list_dir should NOT have ConcurrentFS because it's not upgradeable")
|
||||
}
|
||||
|
||||
// Double check original registry was entirely unmodified
|
||||
origReadTool, _ := original.Get("read_file")
|
||||
if origRead, _ := origReadTool.(*ReadFileTool); origRead != nil {
|
||||
_, isConcurrent := origRead.fs.(*ConcurrentFS)
|
||||
assert.False(t, isConcurrent, "Original registry components MUST REMAIN completely lock-free")
|
||||
}
|
||||
}
|
||||
|
|
@ -10,6 +10,7 @@ import (
|
|||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
|
|
@ -18,17 +19,19 @@ import (
|
|||
|
||||
// ToolLoopConfig configures the tool execution loop.
|
||||
type ToolLoopConfig struct {
|
||||
Provider providers.LLMProvider
|
||||
Model string
|
||||
Tools *ToolRegistry
|
||||
MaxIterations int
|
||||
LLMOptions map[string]any
|
||||
Provider providers.LLMProvider
|
||||
Model string
|
||||
Tools *ToolRegistry
|
||||
MaxIterations int
|
||||
LLMOptions map[string]any
|
||||
RemainingTokenBudget *atomic.Int64
|
||||
}
|
||||
|
||||
// ToolLoopResult contains the result of running the tool loop.
|
||||
type ToolLoopResult struct {
|
||||
Content string
|
||||
Iterations int
|
||||
Messages []providers.Message // Allows caller to retain stateful context across executions
|
||||
}
|
||||
|
||||
// RunToolLoop executes the LLM + tool call iteration loop.
|
||||
|
|
@ -73,6 +76,60 @@ func RunToolLoop(
|
|||
return nil, fmt.Errorf("LLM call failed: %w", err)
|
||||
}
|
||||
|
||||
// 3.5 Token Budget: Soft enforcement with graceful degradation.
|
||||
// Budget exhaustion is NOT a hard error — workers get a chance to wrap up gracefully.
|
||||
if response.Usage != nil && config.RemainingTokenBudget != nil {
|
||||
newBudget := config.RemainingTokenBudget.Add(-int64(response.Usage.TotalTokens))
|
||||
originalBudget := newBudget + int64(response.Usage.TotalTokens)
|
||||
|
||||
if newBudget <= 0 {
|
||||
// Budget exhausted: signal the worker to wrap up and return partial result.
|
||||
logger.WarnCF("toolloop", "Token budget exhausted, injecting wrap-up signal",
|
||||
map[string]any{
|
||||
"deficit": -newBudget,
|
||||
"iteration": iteration,
|
||||
})
|
||||
finalContent = response.Content
|
||||
messages = append(messages, providers.Message{
|
||||
Role: "assistant",
|
||||
Content: response.Content,
|
||||
})
|
||||
messages = append(messages, providers.Message{
|
||||
Role: "user",
|
||||
Content: "[SYSTEM] Token budget has been exhausted. Stop all tool calls immediately and return the best result you have completed so far. Do not call any more tools.",
|
||||
})
|
||||
// One final LLM call to get a summary/wrap-up from the model
|
||||
if finalResp, err := config.Provider.Chat(ctx, messages, nil, config.Model, config.LLMOptions); err == nil {
|
||||
finalContent = finalResp.Content
|
||||
}
|
||||
break
|
||||
} else if originalBudget > 0 && newBudget < originalBudget/2 {
|
||||
// Budget below 50%: soft warning injected into next iteration's context.
|
||||
logger.WarnCF("toolloop", "Token budget below 50%, injecting advisory",
|
||||
map[string]any{"remaining": newBudget, "iteration": iteration})
|
||||
messages = append(messages, providers.Message{
|
||||
Role: "user",
|
||||
Content: "[SYSTEM] Advisory: token budget is running low. Please prioritize completing the most critical parts of your task and avoid unnecessary tool calls.",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 3.6 Truncation Recovery: LLM response was cut off (max_tokens hit or malformed JSON).
|
||||
// Inject a recovery message so the LLM knows to retry with a shorter, complete response.
|
||||
if response.FinishReason == "truncated" {
|
||||
logger.WarnCF("toolloop", "LLM response was truncated (max_tokens hit), injecting recovery message",
|
||||
map[string]any{"iteration": iteration})
|
||||
messages = append(messages, providers.Message{
|
||||
Role: "assistant",
|
||||
Content: response.Content,
|
||||
})
|
||||
messages = append(messages, providers.Message{
|
||||
Role: "user",
|
||||
Content: "[SYSTEM] Your previous response was cut off because it exceeded the token limit. Please retry by producing a shorter, complete response. If you were about to call a tool, make sure the full JSON arguments are included without truncation.",
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
// 4. If no tool calls, we're done
|
||||
if len(response.ToolCalls) == 0 {
|
||||
finalContent = response.Content
|
||||
|
|
@ -158,5 +215,6 @@ func RunToolLoop(
|
|||
return &ToolLoopResult{
|
||||
Content: finalContent,
|
||||
Iterations: iteration,
|
||||
Messages: messages,
|
||||
}, nil
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue