byre/internal/workflow/workflow.go

package workflow

import (
	"errors"
	"fmt"
	"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)...)
	}
	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) []string {
	kinds := 0
	for _, given := range []bool{s.Stage != "", s.Order != nil, s.Parallel != nil} {
		if given {
			kinds++
		}
	}
	if kinds != 1 {
		return []string{fmt.Sprintf("%s: a step needs exactly one of stage, order and parallel", 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)}
		}
		list, kind := s.Order, "order"
		if s.Parallel != nil {
			list, kind = s.Parallel, "parallel"
		}
		if len(*list) == 0 {
			return []string{fmt.Sprintf("%s: %s is empty", where, kind)}
		}
		problems := []string{}
		for i, c := range *list {
			problems = append(problems, stepProblems(c, fmt.Sprintf("%s.%s[%d]", where, kind, i), seen)...)
		}
		return problems
	}

	problems := []string{}
	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) != "":
		problems = append(problems, fmt.Sprintf("%s: model 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 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:
		out := []string{}
		for _, b := range *s.Parallel {
			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
}