byre/internal/workflow/workflow.go

package workflow

import (
	"errors"
	"fmt"
	"math"
	"math/rand/v2"
	"slices"
	"strings"
	"unicode"
	"unicode/utf8"

	"g.fc.run/ai/byre/internal/api"
)

var ErrInvalid = errors.New("invalid workflow")

const maxNameRunes = 100

func CleanName(name string) (string, error) {
	name = strings.TrimSpace(name)
	problem := nameProblem(name)
	if problem != "" {
		return "", fmt.Errorf("%w: workflow name %s", ErrInvalid, problem)
	}
	return name, nil
}

func nameProblem(name string) string {
	switch {
	case strings.TrimSpace(name) == "":
		return "is empty"
	case strings.TrimSpace(name) != name:
		return fmt.Sprintf("%q has leading or trailing spaces", name)
	case utf8.RuneCountInString(name) > maxNameRunes:
		return fmt.Sprintf("%q is longer than %d characters", name, maxNameRunes)
	case strings.ContainsFunc(name, unicode.IsControl):
		return fmt.Sprintf("%q isn't a single line", name)
	}
	return ""
}

func Validate(def api.WorkflowDefinition) error {
	problems := []string{}
	if len(def.Steps) == 0 {
		problems = append(problems, "it has no steps")
	}
	seen := map[string]bool{}
	for i, s := range def.Steps {
		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, "; "))
	}
	return nil
}

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, s.Random != nil} {
		if given {
			kinds++
		}
	}
	if kinds != 1 {
		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.Effort != "" || s.Prompt != "" || s.Agents != nil {
			return append(problems, fmt.Sprintf("%s: model, effort, prompt and agents belong on a stage, not on an order, parallel or random", where))
		}
		list, kind := s.Order, "order"
		switch {
		case s.Parallel != nil:
			list, kind = s.Parallel, "parallel"
		case s.Random != nil:
			list, kind = s.Random, "random"
		}
		if len(*list) == 0 {
			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))
			}
		}
		for i, c := range *list {
			problems = append(problems, stepProblems(c, fmt.Sprintf("%s.%s[%d]", where, kind, i), seen, s.Random != nil)...)
		}
		return problems
	}

	where = fmt.Sprintf("%s (stage %q)", where, s.Stage)
	p := nameProblem(s.Stage)
	if p != "" {
		problems = append(problems, fmt.Sprintf("%s: name %s", where, p))
	}
	if seen[s.Stage] {
		problems = append(problems, fmt.Sprintf("%s: another stage has the same name", where))
	}
	seen[s.Stage] = true

	switch {
	case s.Agents != nil && (strings.TrimSpace(s.Model) != "" || s.Effort != ""):
		problems = append(problems, fmt.Sprintf("%s: model and effort can't be given with agents, which name their own", where))
	case s.Agents != nil:
		for _, p := range TeamProblems(*s.Agents) {
			problems = append(problems, fmt.Sprintf("%s: %s", where, p))
		}
	case strings.TrimSpace(s.Model) == "":
		problems = append(problems, fmt.Sprintf("%s: model is required", where))
	}
	if strings.TrimSpace(s.Prompt) == "" {
		problems = append(problems, fmt.Sprintf("%s: prompt is empty", where))
	}
	return problems
}

func TeamProblems(members []api.TeamMember) []string {
	problems := []string{}
	if len(members) < 2 {
		problems = append(problems, "a team needs at least two agents")
	}
	seen := map[string]bool{}
	for i, mem := range members {
		where := fmt.Sprintf("agent %d (%q)", i+1, mem.Name)
		p := nameProblem(mem.Name)
		if p != "" {
			problems = append(problems, fmt.Sprintf("%s: name %s", where, p))
		}
		if seen[mem.Name] {
			problems = append(problems, fmt.Sprintf("%s: another agent has the same name", where))
		}
		seen[mem.Name] = true
		if strings.TrimSpace(mem.Model) == "" {
			problems = append(problems, fmt.Sprintf("%s: model is required", where))
		}
		if strings.TrimSpace(mem.Prompt) == "" {
			problems = append(problems, fmt.Sprintf("%s: prompt is empty", where))
		}
	}
	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
}

type graph struct {
	stages []api.WorkflowStep
	after  map[string][]string
	ends   []string
}

func (g *graph) walk(s api.WorkflowStep, before []string) []string {
	switch {
	case s.Order != nil:
		for _, c := range *s.Order {
			before = g.walk(c, before)
		}
		return before
	case s.Parallel != nil || s.Random != nil:
		branches := s.Parallel
		if s.Random != nil {
			branches = s.Random
		}
		out := []string{}
		for _, b := range *branches {
			out = append(out, g.walk(b, before)...)
		}
		return out
	}
	g.stages = append(g.stages, s)
	g.after[s.Stage] = before
	return []string{s.Stage}
}

func build(def api.WorkflowDefinition) graph {
	g := graph{after: map[string][]string{}}
	g.ends = g.walk(api.WorkflowStep{Order: &def.Steps}, []string{})
	return g
}

func Stages(def api.WorkflowDefinition) []api.WorkflowStep {
	return build(def).stages
}

func Get(def api.WorkflowDefinition, name string) api.WorkflowStep {
	stages := Stages(def)
	return stages[slices.IndexFunc(stages, func(s api.WorkflowStep) bool { return s.Stage == name })]
}

func After(def api.WorkflowDefinition, name string) []string {
	return build(def).after[name]
}

func Next(def api.WorkflowDefinition, name string) []string {
	g := build(def)
	next := []string{}
	for _, s := range g.stages {
		if slices.Contains(g.after[s.Stage], name) {
			next = append(next, s.Stage)
		}
	}
	return next
}

func Starts(def api.WorkflowDefinition) []string {
	g := build(def)
	starts := []string{}
	for _, s := range g.stages {
		if len(g.after[s.Stage]) == 0 {
			starts = append(starts, s.Stage)
		}
	}
	return starts
}