Add weighted random workflow steps and a random choice agent tool

internal/workflow/workflow.go

@@ -3,6 +3,8 @@ package workflow
 import (
 	"errors"
 	"fmt"
+	"math"
+	"math/rand/v2"
 	"slices"
 	"strings"
 	"unicode"
@@ -45,7 +47,7 @@ func Validate(def api.WorkflowDefinition) error {
 	}
 	seen := map[string]bool{}
 	for i, s := range def.Steps {
-		problems = append(problems, stepProblems(s, fmt.Sprintf("steps[%d]", i), seen)...)
+		problems = append(problems, stepProblems(s, fmt.Sprintf("steps[%d]", i), seen, false)...)
 	}
 	if len(problems) > 0 {
 		return fmt.Errorf("%w: %s", ErrInvalid, strings.Join(problems, "; "))
@@ -53,36 +55,52 @@ func Validate(def api.WorkflowDefinition) error {
 	return nil
 }
 
-func stepProblems(s api.WorkflowStep, where string, seen map[string]bool) []string {
+func stepProblems(s api.WorkflowStep, where string, seen map[string]bool, inRandom bool) []string {
 	kinds := 0
-	for _, given := range []bool{s.Stage != "", s.Order != nil, s.Parallel != nil} {
+	for _, given := range []bool{s.Stage != "", s.Order != nil, s.Parallel != nil, s.Random != nil} {
 		if given {
 			kinds++
 		}
 	}
 	if kinds != 1 {
-		return []string{fmt.Sprintf("%s: a step needs exactly one of stage, order and parallel", where)}
+		return []string{fmt.Sprintf("%s: a step needs exactly one of stage, order, parallel and random", where)}
+	}
+
+	problems := []string{}
+	if !inRandom && s.Weight != 0 {
+		problems = append(problems, fmt.Sprintf("%s: weight belongs only on a step directly in a random", where))
 	}
 
 	if s.Stage == "" {
 		if s.Model != "" || s.Prompt != "" || s.Agents != nil {
-			return []string{fmt.Sprintf("%s: model, prompt and agents belong on a stage, not on an order or parallel", where)}
+			return append(problems, fmt.Sprintf("%s: model, prompt and agents belong on a stage, not on an order, parallel or random", where))
 		}
 		list, kind := s.Order, "order"
-		if s.Parallel != nil {
+		switch {
+		case s.Parallel != nil:
 			list, kind = s.Parallel, "parallel"
+		case s.Random != nil:
+			list, kind = s.Random, "random"
 		}
 		if len(*list) == 0 {
-			return []string{fmt.Sprintf("%s: %s is empty", where, kind)}
+			return append(problems, fmt.Sprintf("%s: %s is empty", where, kind))
+		}
+		if s.Random != nil {
+			weights := []float64{}
+			for _, c := range *s.Random {
+				weights = append(weights, c.Weight)
+			}
+			p := WeightsProblem(weights)
+			if p != "" {
+				problems = append(problems, fmt.Sprintf("%s: %s", where, p))
+			}
 		}
-		problems := []string{}
 		for i, c := range *list {
-			problems = append(problems, stepProblems(c, fmt.Sprintf("%s.%s[%d]", where, kind, i), seen)...)
+			problems = append(problems, stepProblems(c, fmt.Sprintf("%s.%s[%d]", where, kind, i), seen, s.Random != nil)...)
 		}
 		return problems
 	}
 
-	problems := []string{}
 	where = fmt.Sprintf("%s (stage %q)", where, s.Stage)
 	p := nameProblem(s.Stage)
 	if p != "" {
@@ -135,6 +153,91 @@ func TeamProblems(members []api.TeamMember) []string {
 	return problems
 }
 
+func HasRandom(def api.WorkflowDefinition) bool {
+	var has func(steps []api.WorkflowStep) bool
+	has = func(steps []api.WorkflowStep) bool {
+		return slices.ContainsFunc(steps, func(s api.WorkflowStep) bool {
+			return s.Random != nil || (s.Order != nil && has(*s.Order)) || (s.Parallel != nil && has(*s.Parallel))
+		})
+	}
+	return has(def.Steps)
+}
+
+func Resolve(def api.WorkflowDefinition) api.WorkflowDefinition {
+	var resolve func(steps []api.WorkflowStep) []api.WorkflowStep
+	var pick func(s api.WorkflowStep) api.WorkflowStep
+	pick = func(s api.WorkflowStep) api.WorkflowStep {
+		switch {
+		case s.Order != nil:
+			sub := resolve(*s.Order)
+			s.Order = &sub
+		case s.Parallel != nil:
+			sub := resolve(*s.Parallel)
+			s.Parallel = &sub
+		case s.Random != nil:
+			weights := []float64{}
+			for _, c := range *s.Random {
+				weights = append(weights, c.Weight)
+			}
+			sum, _ := Weights(weights)
+			chosen := (*s.Random)[Pick(weights, sum)]
+			chosen.Weight = 0
+			return pick(chosen)
+		}
+		return s
+	}
+	resolve = func(steps []api.WorkflowStep) []api.WorkflowStep {
+		out := []api.WorkflowStep{}
+		for _, s := range steps {
+			out = append(out, pick(s))
+		}
+		return out
+	}
+	def.Steps = resolve(def.Steps)
+	return def
+}
+
+const weightSlack = 1e-6
+
+func Weights(weights []float64) (float64, string) {
+	sum := 0.0
+	for i, w := range weights {
+		if w < 0 || w > 1 {
+			return 0, fmt.Sprintf("weight %d is %g, outside 0 to 1", i+1, w)
+		}
+		sum += w
+	}
+	return sum, ""
+}
+
+func WeightsProblem(weights []float64) string {
+	sum, p := Weights(weights)
+	if p == "" && math.Abs(sum-1) > weightSlack {
+		p = fmt.Sprintf("the weights add up to %g, not 1", sum)
+	}
+	return p
+}
+
+func ChoiceWeightsProblem(weights []float64) string {
+	sum, p := Weights(weights)
+	if p == "" && sum > 1+weightSlack {
+		p = fmt.Sprintf("the weights add up to %g, more than 1", sum)
+	}
+	return p
+}
+
+func Pick(weights []float64, total float64) int {
+	n := rand.Float64() * total
+	cum := 0.0
+	for i, w := range weights {
+		cum += w
+		if n < cum {
+			return i
+		}
+	}
+	return -1
+}
+
 func Ends(def api.WorkflowDefinition) []string {
 	return build(def).ends
 }
@@ -152,9 +255,13 @@ func (g *graph) walk(s api.WorkflowStep, before []string) []string {
 			before = g.walk(c, before)
 		}
 		return before
-	case s.Parallel != nil:
+	case s.Parallel != nil || s.Random != nil:
+		branches := s.Parallel
+		if s.Random != nil {
+			branches = s.Random
+		}
 		out := []string{}
-		for _, b := range *s.Parallel {
+		for _, b := range *branches {
 			out = append(out, g.walk(b, before)...)
 		}
 		return out