129 lines
3.0 KiB
Go
129 lines
3.0 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/sipeed/picoclaw/pkg/bus"
|
|
"github.com/sipeed/picoclaw/pkg/providers"
|
|
)
|
|
|
|
type SubagentTask struct {
|
|
ID string
|
|
Task string
|
|
Label string
|
|
OriginChannel string
|
|
OriginChatID string
|
|
Status string
|
|
Result string
|
|
Created int64
|
|
}
|
|
|
|
type SubagentManager struct {
|
|
tasks map[string]*SubagentTask
|
|
mu sync.RWMutex
|
|
provider providers.LLMProvider
|
|
bus *bus.MessageBus
|
|
workspace string
|
|
nextID int
|
|
}
|
|
|
|
func NewSubagentManager(provider providers.LLMProvider, workspace string, bus *bus.MessageBus) *SubagentManager {
|
|
return &SubagentManager{
|
|
tasks: make(map[string]*SubagentTask),
|
|
provider: provider,
|
|
bus: bus,
|
|
workspace: workspace,
|
|
nextID: 1,
|
|
}
|
|
}
|
|
|
|
func (sm *SubagentManager) Spawn(ctx context.Context, task, label, originChannel, originChatID string) (string, error) {
|
|
sm.mu.Lock()
|
|
defer sm.mu.Unlock()
|
|
|
|
taskID := fmt.Sprintf("subagent-%d", sm.nextID)
|
|
sm.nextID++
|
|
|
|
subagentTask := &SubagentTask{
|
|
ID: taskID,
|
|
Task: task,
|
|
Label: label,
|
|
OriginChannel: originChannel,
|
|
OriginChatID: originChatID,
|
|
Status: "running",
|
|
Created: time.Now().UnixMilli(),
|
|
}
|
|
sm.tasks[taskID] = subagentTask
|
|
|
|
go sm.runTask(ctx, subagentTask)
|
|
|
|
if label != "" {
|
|
return fmt.Sprintf("Spawned subagent '%s' for task: %s", label, task), nil
|
|
}
|
|
return fmt.Sprintf("Spawned subagent for task: %s", task), nil
|
|
}
|
|
|
|
func (sm *SubagentManager) runTask(ctx context.Context, task *SubagentTask) {
|
|
task.Status = "running"
|
|
task.Created = time.Now().UnixMilli()
|
|
|
|
messages := []providers.Message{
|
|
{
|
|
Role: "system",
|
|
Content: "You are a subagent. Complete the given task independently and report the result.",
|
|
},
|
|
{
|
|
Role: "user",
|
|
Content: task.Task,
|
|
},
|
|
}
|
|
|
|
response, err := sm.provider.Chat(ctx, messages, nil, sm.provider.GetDefaultModel(), map[string]interface{}{
|
|
"max_tokens": 4096,
|
|
})
|
|
|
|
sm.mu.Lock()
|
|
defer sm.mu.Unlock()
|
|
|
|
if err != nil {
|
|
task.Status = "failed"
|
|
task.Result = fmt.Sprintf("Error: %v", err)
|
|
} else {
|
|
task.Status = "completed"
|
|
task.Result = response.Content
|
|
}
|
|
|
|
// Send announce message back to main agent
|
|
if sm.bus != nil {
|
|
announceContent := fmt.Sprintf("Task '%s' completed.\n\nResult:\n%s", task.Label, task.Result)
|
|
sm.bus.PublishInbound(bus.InboundMessage{
|
|
Channel: "system",
|
|
SenderID: fmt.Sprintf("subagent:%s", task.ID),
|
|
// Format: "original_channel:original_chat_id" for routing back
|
|
ChatID: fmt.Sprintf("%s:%s", task.OriginChannel, task.OriginChatID),
|
|
Content: announceContent,
|
|
})
|
|
}
|
|
}
|
|
|
|
func (sm *SubagentManager) GetTask(taskID string) (*SubagentTask, bool) {
|
|
sm.mu.RLock()
|
|
defer sm.mu.RUnlock()
|
|
task, ok := sm.tasks[taskID]
|
|
return task, ok
|
|
}
|
|
|
|
func (sm *SubagentManager) ListTasks() []*SubagentTask {
|
|
sm.mu.RLock()
|
|
defer sm.mu.RUnlock()
|
|
|
|
tasks := make([]*SubagentTask, 0, len(sm.tasks))
|
|
for _, task := range sm.tasks {
|
|
tasks = append(tasks, task)
|
|
}
|
|
return tasks
|
|
}
|