byre/internal/workflow/file.go

package workflow

import (
	"bytes"
	"encoding/json"
	"fmt"
	"slices"
	"strings"

	"go.yaml.in/yaml/v3"

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

var keyOrder = []string{"description", "stages", "name", "after", "base", "model", "prompt"}

func Parse(data []byte) (api.WorkflowDefinition, error) {
	raw := any(nil)
	err := yaml.Unmarshal(data, &raw)
	if err != nil {
		return api.WorkflowDefinition{}, fmt.Errorf("%w: %v", ErrInvalid, err)
	}
	j, err := json.Marshal(raw)
	if err != nil {
		return api.WorkflowDefinition{}, fmt.Errorf("%w: %v", ErrInvalid, err)
	}

	def := api.WorkflowDefinition{}
	dec := json.NewDecoder(bytes.NewReader(j))
	dec.DisallowUnknownFields()
	err = dec.Decode(&def)
	if err != nil {
		return api.WorkflowDefinition{}, fmt.Errorf("%w: %v", ErrInvalid, err)
	}
	return def, Validate(def)
}

func Format(def api.WorkflowDefinition) ([]byte, error) {
	j, err := json.Marshal(def)
	if err != nil {
		return nil, err
	}
	doc := yaml.Node{}
	err = yaml.Unmarshal(j, &doc)
	if err != nil {
		return nil, err
	}
	restyle(&doc)

	out := &bytes.Buffer{}
	enc := yaml.NewEncoder(out)
	enc.SetIndent(2)
	err = enc.Encode(&doc)
	if err != nil {
		return nil, err
	}
	err = enc.Close()
	if err != nil {
		return nil, err
	}
	return out.Bytes(), nil
}

func Filename(name string) string {
	return strings.Map(func(r rune) rune {
		if strings.ContainsRune(`/\:*?"<>|`, r) {
			return '_'
		}
		return r
	}, name) + ".yaml"
}

func restyle(n *yaml.Node) {
	switch n.Kind {
	case yaml.ScalarNode:
		n.Style = 0
		if n.Tag == "!!str" && strings.Contains(n.Value, "\n") {
			n.Style = yaml.LiteralStyle
		}
	case yaml.MappingNode:
		n.Style = 0
		pairs := [][2]*yaml.Node{}
		for i := 0; i+1 < len(n.Content); i += 2 {
			pairs = append(pairs, [2]*yaml.Node{n.Content[i], n.Content[i+1]})
		}
		slices.SortStableFunc(pairs, func(a, b [2]*yaml.Node) int {
			return keyRank(a[0].Value) - keyRank(b[0].Value)
		})
		n.Content = n.Content[:0]
		for _, p := range pairs {
			n.Content = append(n.Content, p[0], p[1])
		}
	case yaml.SequenceNode:
		n.Style = 0
	}
	for _, c := range n.Content {
		restyle(c)
	}
}

func keyRank(key string) int {
	i := slices.Index(keyOrder, key)
	if i < 0 {
		return len(keyOrder)
	}
	return i
}