|
| 1 | +// Plan & Delegate integration harness. |
| 2 | +// |
| 3 | +// This runs the REAL go-micro stack end to end — real services, real |
| 4 | +// registry, real RPC, the real agent loop, real store, real delegate |
| 5 | +// routing — and mocks ONLY the LLM with a deterministic provider. It |
| 6 | +// proves the plumbing works without an API key, and it's reproducible. |
| 7 | +// |
| 8 | +// Swap MICRO_AI_PROVIDER/MICRO_AI_API_KEY (and remove --mock) to run the |
| 9 | +// exact same flow against a live model. |
| 10 | +// |
| 11 | +// Run: |
| 12 | +// |
| 13 | +// go run ./internal/harness/plan-delegate |
| 14 | +package main |
| 15 | + |
| 16 | +import ( |
| 17 | + "context" |
| 18 | + "encoding/json" |
| 19 | + "flag" |
| 20 | + "fmt" |
| 21 | + "os" |
| 22 | + "strings" |
| 23 | + "sync" |
| 24 | + "time" |
| 25 | + |
| 26 | + "go-micro.dev/v5" |
| 27 | + "go-micro.dev/v5/ai" |
| 28 | +) |
| 29 | + |
| 30 | +// --------------------------------------------------------------------------- |
| 31 | +// real services |
| 32 | +// --------------------------------------------------------------------------- |
| 33 | + |
| 34 | +type Task struct { |
| 35 | + ID string `json:"id"` |
| 36 | + Title string `json:"title"` |
| 37 | +} |
| 38 | + |
| 39 | +type AddRequest struct { |
| 40 | + Title string `json:"title" description:"Title of the task to add"` |
| 41 | +} |
| 42 | +type AddResponse struct { |
| 43 | + Task *Task `json:"task"` |
| 44 | +} |
| 45 | +type ListRequest struct{} |
| 46 | +type ListResponse struct { |
| 47 | + Tasks []*Task `json:"tasks"` |
| 48 | +} |
| 49 | + |
| 50 | +type TaskService struct { |
| 51 | + mu sync.Mutex |
| 52 | + tasks []*Task |
| 53 | + nextID int |
| 54 | +} |
| 55 | + |
| 56 | +// Add creates a new task with the given title. |
| 57 | +// @example {"title": "Design"} |
| 58 | +func (s *TaskService) Add(ctx context.Context, req *AddRequest, rsp *AddResponse) error { |
| 59 | + s.mu.Lock() |
| 60 | + defer s.mu.Unlock() |
| 61 | + s.nextID++ |
| 62 | + t := &Task{ID: fmt.Sprintf("task-%d", s.nextID), Title: req.Title} |
| 63 | + s.tasks = append(s.tasks, t) |
| 64 | + rsp.Task = t |
| 65 | + fmt.Printf(" \033[32m[task]\033[0m created %s %q\n", t.ID, t.Title) |
| 66 | + return nil |
| 67 | +} |
| 68 | + |
| 69 | +// List returns all tasks. |
| 70 | +// @example {} |
| 71 | +func (s *TaskService) List(ctx context.Context, req *ListRequest, rsp *ListResponse) error { |
| 72 | + s.mu.Lock() |
| 73 | + defer s.mu.Unlock() |
| 74 | + rsp.Tasks = append(rsp.Tasks, s.tasks...) |
| 75 | + return nil |
| 76 | +} |
| 77 | + |
| 78 | +func (s *TaskService) count() int { |
| 79 | + s.mu.Lock() |
| 80 | + defer s.mu.Unlock() |
| 81 | + return len(s.tasks) |
| 82 | +} |
| 83 | + |
| 84 | +type SendRequest struct { |
| 85 | + To string `json:"to" description:"Recipient address"` |
| 86 | + Message string `json:"message" description:"Message body"` |
| 87 | +} |
| 88 | +type SendResponse struct { |
| 89 | + Sent bool `json:"sent"` |
| 90 | +} |
| 91 | +type NotifyService struct { |
| 92 | + mu sync.Mutex |
| 93 | + sent int |
| 94 | +} |
| 95 | + |
| 96 | +// Send delivers a notification message to a recipient. |
| 97 | +// @example {"to": "owner@acme.com", "message": "ready"} |
| 98 | +func (s *NotifyService) Send(ctx context.Context, req *SendRequest, rsp *SendResponse) error { |
| 99 | + s.mu.Lock() |
| 100 | + s.sent++ |
| 101 | + s.mu.Unlock() |
| 102 | + fmt.Printf(" \033[35m[notify]\033[0m 📨 to=%s message=%q\n", req.To, req.Message) |
| 103 | + rsp.Sent = true |
| 104 | + return nil |
| 105 | +} |
| 106 | + |
| 107 | +func (s *NotifyService) count() int { |
| 108 | + s.mu.Lock() |
| 109 | + defer s.mu.Unlock() |
| 110 | + return s.sent |
| 111 | +} |
| 112 | + |
| 113 | +// --------------------------------------------------------------------------- |
| 114 | +// mock LLM provider — the ONLY fake. It "reasons" by simple heuristics |
| 115 | +// over the tools it's offered and the system prompt it's given, calling |
| 116 | +// the real tool handler exactly the way a real provider would. |
| 117 | +// --------------------------------------------------------------------------- |
| 118 | + |
| 119 | +type mockModel struct{ opts ai.Options } |
| 120 | + |
| 121 | +func newMock(opts ...ai.Option) ai.Model { |
| 122 | + m := &mockModel{} |
| 123 | + _ = m.Init(opts...) |
| 124 | + return m |
| 125 | +} |
| 126 | + |
| 127 | +func (m *mockModel) Init(opts ...ai.Option) error { |
| 128 | + for _, o := range opts { |
| 129 | + o(&m.opts) |
| 130 | + } |
| 131 | + return nil |
| 132 | +} |
| 133 | +func (m *mockModel) Options() ai.Options { return m.opts } |
| 134 | +func (m *mockModel) String() string { return "mock" } |
| 135 | +func (m *mockModel) Stream(ctx context.Context, req *ai.Request, _ ...ai.GenerateOption) (ai.Stream, error) { |
| 136 | + return nil, fmt.Errorf("stream not supported by mock") |
| 137 | +} |
| 138 | + |
| 139 | +// findTool returns the safe name of the first offered tool whose name |
| 140 | +// contains sub, or "" if none. |
| 141 | +func findTool(tools []ai.Tool, sub string) string { |
| 142 | + for _, t := range tools { |
| 143 | + if strings.Contains(t.Name, sub) { |
| 144 | + return t.Name |
| 145 | + } |
| 146 | + } |
| 147 | + return "" |
| 148 | +} |
| 149 | + |
| 150 | +func (m *mockModel) call(who, name string, input map[string]any) { |
| 151 | + args, _ := json.Marshal(input) |
| 152 | + fmt.Printf(" \033[33m[%s]\033[0m → %s(%s)\n", who, name, args) |
| 153 | + if m.opts.ToolHandler != nil { |
| 154 | + m.opts.ToolHandler(name, input) |
| 155 | + } |
| 156 | +} |
| 157 | + |
| 158 | +func (m *mockModel) Generate(ctx context.Context, req *ai.Request, _ ...ai.GenerateOption) (*ai.Response, error) { |
| 159 | + // Classify by the tools actually offered, not by prompt text: |
| 160 | + // the conductor has the task "Add" tool, comms has "Send". |
| 161 | + hasAdd := findTool(req.Tools, "Add") != "" |
| 162 | + hasSend := findTool(req.Tools, "Send") != "" |
| 163 | + |
| 164 | + switch { |
| 165 | + // comms agent: owns notify, has Send but not Add. |
| 166 | + case hasSend && !hasAdd: |
| 167 | + send := findTool(req.Tools, "Send") |
| 168 | + m.call("comms", send, map[string]any{ |
| 169 | + "to": "owner@acme.com", |
| 170 | + "message": "The launch plan is ready", |
| 171 | + }) |
| 172 | + return &ai.Response{Answer: "Notified owner@acme.com."}, nil |
| 173 | + |
| 174 | + // conductor: has the task Add tool — plan, create tasks, delegate. |
| 175 | + case hasAdd: |
| 176 | + if plan := findTool(req.Tools, "plan"); plan != "" { |
| 177 | + m.call("conductor", plan, map[string]any{ |
| 178 | + "steps": []any{ |
| 179 | + map[string]any{"task": "create Design task", "status": "pending"}, |
| 180 | + map[string]any{"task": "create Build task", "status": "pending"}, |
| 181 | + map[string]any{"task": "create Ship task", "status": "pending"}, |
| 182 | + map[string]any{"task": "notify owner via comms", "status": "pending"}, |
| 183 | + }, |
| 184 | + }) |
| 185 | + } |
| 186 | + if add := findTool(req.Tools, "Add"); add != "" { |
| 187 | + for _, title := range []string{"Design", "Build", "Ship"} { |
| 188 | + m.call("conductor", add, map[string]any{"title": title}) |
| 189 | + } |
| 190 | + } |
| 191 | + if del := findTool(req.Tools, "delegate"); del != "" { |
| 192 | + m.call("conductor", del, map[string]any{ |
| 193 | + "task": "Notify owner@acme.com that the launch plan is ready", |
| 194 | + "to": "comms", |
| 195 | + }) |
| 196 | + } |
| 197 | + return &ai.Response{Answer: "Created Design, Build and Ship, and had comms notify the owner."}, nil |
| 198 | + |
| 199 | + // ephemeral sub-agent or anything else. |
| 200 | + default: |
| 201 | + return &ai.Response{Reply: "subtask handled"}, nil |
| 202 | + } |
| 203 | +} |
| 204 | + |
| 205 | +func providerKey(provider string) string { |
| 206 | + if v := os.Getenv("MICRO_AI_API_KEY"); v != "" { |
| 207 | + return v |
| 208 | + } |
| 209 | + env := map[string]string{ |
| 210 | + "anthropic": "ANTHROPIC_API_KEY", |
| 211 | + "openai": "OPENAI_API_KEY", |
| 212 | + "gemini": "GEMINI_API_KEY", |
| 213 | + "groq": "GROQ_API_KEY", |
| 214 | + "mistral": "MISTRAL_API_KEY", |
| 215 | + "together": "TOGETHER_API_KEY", |
| 216 | + "atlascloud": "ATLASCLOUD_API_KEY", |
| 217 | + }[provider] |
| 218 | + return os.Getenv(env) |
| 219 | +} |
| 220 | + |
| 221 | +func main() { |
| 222 | + provider := flag.String("provider", "mock", "LLM provider: mock (default), anthropic, openai, gemini, groq, mistral, together, atlascloud") |
| 223 | + flag.Parse() |
| 224 | + |
| 225 | + apiKey := "" |
| 226 | + if *provider == "mock" { |
| 227 | + ai.Register("mock", newMock) |
| 228 | + } else { |
| 229 | + apiKey = providerKey(*provider) |
| 230 | + if apiKey == "" { |
| 231 | + fmt.Printf("no API key for provider %q — set MICRO_AI_API_KEY or the provider's key env\n", *provider) |
| 232 | + return |
| 233 | + } |
| 234 | + } |
| 235 | + |
| 236 | + fmt.Printf("\n\033[1mPlan & Delegate — live integration harness (provider: %s)\033[0m\n", *provider) |
| 237 | + fmt.Print("Real services, registry, RPC, agent loop, store, delegation.\n\n") |
| 238 | + |
| 239 | + // Real services. |
| 240 | + task := micro.New("task") |
| 241 | + task.Handle(new(TaskService)) |
| 242 | + go task.Run() |
| 243 | + |
| 244 | + notify := micro.New("notify") |
| 245 | + notify.Handle(new(NotifyService)) |
| 246 | + go notify.Run() |
| 247 | + |
| 248 | + // Real comms agent (owns notify), registered so delegate reaches it over RPC. |
| 249 | + comms := micro.NewAgent("comms", |
| 250 | + micro.AgentServices("notify"), |
| 251 | + micro.AgentPrompt("You handle outbound notifications. Use the notify service."), |
| 252 | + micro.AgentProvider(*provider), |
| 253 | + micro.AgentAPIKey(apiKey), |
| 254 | + ) |
| 255 | + go comms.Run() |
| 256 | + |
| 257 | + // Real conductor agent (owns task). |
| 258 | + conductor := micro.NewAgent("conductor", |
| 259 | + micro.AgentServices("task"), |
| 260 | + micro.AgentPrompt("You coordinate launch work. Plan first, create tasks, and delegate notifications to the \"comms\" agent."), |
| 261 | + micro.AgentProvider(*provider), |
| 262 | + micro.AgentAPIKey(apiKey), |
| 263 | + ) |
| 264 | + |
| 265 | + fmt.Println("waiting for services + comms agent to register...") |
| 266 | + time.Sleep(3 * time.Second) |
| 267 | + |
| 268 | + fmt.Print("\n\033[1m> prompt:\033[0m Create three launch tasks (Design, Build, Ship), then make sure owner@acme.com is notified.\n\n") |
| 269 | + |
| 270 | + resp, err := conductor.Ask(context.Background(), |
| 271 | + "Create three launch tasks: Design, Build, and Ship. Then make sure owner@acme.com is notified that the launch plan is ready.") |
| 272 | + if err != nil { |
| 273 | + fmt.Println("\033[31merror:\033[0m", err) |
| 274 | + os.Exit(1) |
| 275 | + } |
| 276 | + |
| 277 | + fmt.Println("\n\033[1m< conductor reply:\033[0m", resp.Reply) |
| 278 | + |
| 279 | + // Prove plan was persisted to the real store. |
| 280 | + if recs, _ := conductor.Options().Store.Read("agent/conductor/plan"); len(recs) > 0 { |
| 281 | + fmt.Printf("\n\033[1mstored plan (agent/conductor/plan):\033[0m %s\n", string(recs[0].Value)) |
| 282 | + } else { |
| 283 | + fmt.Println("\n\033[31m! plan was not persisted\033[0m") |
| 284 | + } |
| 285 | + |
| 286 | + fmt.Println("\n\033[32m✓ end-to-end flow complete\033[0m") |
| 287 | +} |
0 commit comments