From 406247b14985ab57902e8e42898dc8cb8960290d Mon Sep 17 00:00:00 2001 From: Peter Stone Date: Sun, 8 Mar 2026 20:50:21 +0000 Subject: feat(executor): implement Gemini-based task classification and load balancing - Add Classifier using gemini-2.0-flash-lite to automatically select agent/model. - Update Pool to track per-agent active tasks and rate limit status. - Enable classification for all tasks (top-level and subtasks). - Refine SystemStatus to be dynamic across all supported agents. - Add unit tests for the classifier and updated pool logic. - Minor UI improvements for project selection and 'Start Next' action. --- internal/executor/classifier.go | 109 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 internal/executor/classifier.go (limited to 'internal/executor/classifier.go') diff --git a/internal/executor/classifier.go b/internal/executor/classifier.go new file mode 100644 index 0000000..79ebc27 --- /dev/null +++ b/internal/executor/classifier.go @@ -0,0 +1,109 @@ +package executor + +import ( + "context" + "encoding/json" + "fmt" + "os/exec" + "strings" +) + +type Classification struct { + AgentType string `json:"agent_type"` + Model string `json:"model"` + Reason string `json:"reason"` +} + +type SystemStatus struct { + ActiveTasks map[string]int + RateLimited map[string]bool +} + +type Classifier struct { + GeminiBinaryPath string +} + +const classificationPrompt = ` +You are a task classifier for Claudomator. +Given a task description and system status, select the best agent (claude or gemini) and model to use. + +Agent Types: +- claude: Best for complex coding, reasoning, and tool use. +- gemini: Best for large context, fast reasoning, and multimodal tasks. + +Available Models: +Claude: +- claude-3-5-sonnet-latest (balanced) +- claude-3-5-sonnet-20241022 (stable) +- claude-3-opus-20240229 (most powerful, expensive) +- claude-3-5-haiku-20241022 (fast, cheap) + +Gemini: +- gemini-2.0-flash-lite (fastest, most efficient, best for simple tasks) +- gemini-2.0-flash (fast, multimodal) +- gemini-1.5-flash (fast, balanced) +- gemini-1.5-pro (more powerful, larger context) + +Selection Criteria: +- Agent: Prefer the one with least running tasks and no active rate limit. +- Model: Select based on task complexity. Use powerful models (opus, pro) for complex reasoning/coding, flash-lite/flash/haiku for simple tasks. + +Task: +Name: %s +Instructions: %s + +System Status: +%s + +Respond with ONLY a JSON object: +{ + "agent_type": "claude" | "gemini", + "model": "model-name", + "reason": "brief reason" +} +` + +func (c *Classifier) Classify(ctx context.Context, taskName, instructions string, status SystemStatus) (*Classification, error) { + statusStr := "" + for agent, active := range status.ActiveTasks { + statusStr += fmt.Sprintf("- Agent %s: %d active tasks, Rate Limited: %t\n", agent, active, status.RateLimited[agent]) + } + + prompt := fmt.Sprintf(classificationPrompt, + taskName, instructions, statusStr, + ) + + binary := c.GeminiBinaryPath + if binary == "" { + binary = "gemini" + } + + // Use a minimal model for classification to be fast and cheap. + args := []string{ + "--prompt", prompt, + "--model", "gemini-2.0-flash-lite", + "--output-format", "json", + } + + cmd := exec.CommandContext(ctx, binary, args...) + out, err := cmd.Output() + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + return nil, fmt.Errorf("classifier failed (%v): %s", err, string(exitErr.Stderr)) + } + return nil, fmt.Errorf("classifier failed: %w", err) + } + + var cls Classification + // Gemini might wrap the JSON in markdown code blocks. + cleanOut := strings.TrimSpace(string(out)) + cleanOut = strings.TrimPrefix(cleanOut, "```json") + cleanOut = strings.TrimSuffix(cleanOut, "```") + cleanOut = strings.TrimSpace(cleanOut) + + if err := json.Unmarshal([]byte(cleanOut), &cls); err != nil { + return nil, fmt.Errorf("failed to parse classification JSON: %w\nOutput: %s", err, cleanOut) + } + + return &cls, nil +} -- cgit v1.2.3