feat(teams): implement Agent Teams architecture
- Added 'team' and 'spawn_sub_agent' tools to support Coordinator-Worker patterns. - Implemented execution strategies: Sequential, Parallel, DAG, and Evaluator-Optimizer. - Added global token budget tracking via 'RemainingTokenBudget' for team cost control. - Designed 'ConcurrencyUpgradeable' interface and 'ConcurrentFS' wrapper for opt-in, thread-safe file operations during parallel/DAG runs. - Tested file locking mechanisms preventing race conditions during high concurrency.
This commit is contained in:
Administrator
parent
2f4f45080b
commit
e7de26ee1b
10 changed files with 1015 additions and 28 deletions
|
|
@ -155,6 +155,12 @@ 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)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
|
||||
// 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 sysFs.WriteFile(path, newContent)
|
||||
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)
|
||||
}
|
||||
|
|
|
|||
105
pkg/tools/spawn_sub_agent.go
Normal file
105
pkg/tools/spawn_sub_agent.go
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
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 {
|
||||
return "Directly delegate a specific task to a new, isolated sub-agent. You (the main agent) should autonomously determine the appropriate expert role and specific task based on the user's high-level request. It will execute independently and return the final result."
|
||||
}
|
||||
|
||||
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').",
|
||||
},
|
||||
},
|
||||
"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)
|
||||
|
||||
// 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,
|
||||
}
|
||||
}
|
||||
|
|
@ -248,6 +248,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.
|
||||
|
|
|
|||
571
pkg/tools/team.go
Normal file
571
pkg/tools/team.go
Normal file
|
|
@ -0,0 +1,571 @@
|
|||
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
|
||||
DependsOn []string // List of member IDs this member depends on
|
||||
}
|
||||
|
||||
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 {
|
||||
return "Compose and execute a team of distinct sub-agents. You (the main agent) should autonomously analyze the user's request, determine the necessary specialized roles, break down the work into sub-tasks, and assign them. Execute sequentially (passing output from one to the next) or concurrently in parallel."
|
||||
}
|
||||
|
||||
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.",
|
||||
},
|
||||
"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"},
|
||||
},
|
||||
},
|
||||
"required": []string{"role", "task"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": []string{"strategy", "members"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *TeamTool) SetContext(channel, chatID string) {
|
||||
t.originChannel = channel
|
||||
t.originChatID = chatID
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
members = append(members, TeamMember{
|
||||
ID: id,
|
||||
Role: role,
|
||||
Task: task,
|
||||
DependsOn: dependsOn,
|
||||
})
|
||||
}
|
||||
|
||||
// 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":
|
||||
return t.executeSequential(teamCtx, baseConfig, members)
|
||||
case "dag":
|
||||
return t.executeDAG(teamCtx, baseConfig, members)
|
||||
case "evaluator_optimizer":
|
||||
return t.executeEvaluatorOptimizer(teamCtx, baseConfig, members)
|
||||
}
|
||||
return t.executeParallel(teamCtx, baseConfig, members)
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
func (t *TeamTool) executeSequential(ctx context.Context, config 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},
|
||||
}
|
||||
|
||||
loopResult, err := RunToolLoop(ctx, config, 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, config 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, role, task string) {
|
||||
defer wg.Done()
|
||||
|
||||
messages := []providers.Message{
|
||||
{Role: "system", Content: role},
|
||||
{Role: "user", Content: task},
|
||||
}
|
||||
|
||||
loopResult, err := RunToolLoop(ctx, config, messages, t.originChannel, t.originChatID)
|
||||
|
||||
if err != nil {
|
||||
resultsChan <- workResult{index: index, role: role, err: err}
|
||||
return
|
||||
}
|
||||
resultsChan <- workResult{index: index, role: role, res: loopResult.Content}
|
||||
}(i, m.Role, m.Task)
|
||||
}
|
||||
|
||||
// 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, config 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!)
|
||||
workerResult, err := RunToolLoop(ctx, config, 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},
|
||||
}
|
||||
|
||||
evalResult, err := RunToolLoop(ctx, config, 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, config 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},
|
||||
}
|
||||
|
||||
loopResult, err := RunToolLoop(ctx, config, 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"
|
||||
|
|
@ -23,12 +24,14 @@ type ToolLoopConfig struct {
|
|||
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,18 @@ func RunToolLoop(
|
|||
return nil, fmt.Errorf("LLM call failed: %w", err)
|
||||
}
|
||||
|
||||
// 3.5 Token Budget Enforcement
|
||||
if response.Usage != nil && config.RemainingTokenBudget != nil {
|
||||
newBudget := config.RemainingTokenBudget.Add(-int64(response.Usage.TotalTokens))
|
||||
if newBudget < 0 {
|
||||
logger.ErrorCF("toolloop", "Token budget exceeded", map[string]any{
|
||||
"used_iteration": response.Usage.TotalTokens,
|
||||
"deficit": newBudget,
|
||||
})
|
||||
return nil, fmt.Errorf("Token Budget Exceeded: Team budget completely consumed")
|
||||
}
|
||||
}
|
||||
|
||||
// 4. If no tool calls, we're done
|
||||
if len(response.ToolCalls) == 0 {
|
||||
finalContent = response.Content
|
||||
|
|
@ -158,5 +173,6 @@ func RunToolLoop(
|
|||
return &ToolLoopResult{
|
||||
Content: finalContent,
|
||||
Iterations: iteration,
|
||||
Messages: messages,
|
||||
}, nil
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue