byre/internal/store/definition.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"

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

func insertDefinition(ctx context.Context, tx *sql.Tx, def api.WorkflowDefinition) (string, error) {
	id := ids.New(ids.WorkflowDefinition)
	_, err := tx.ExecContext(ctx, `INSERT INTO workflow_definitions (id, description) VALUES (?, ?)`, id, def.Description)
	if err != nil {
		return "", err
	}
	next := 0
	var insert func(steps []api.WorkflowStep, parent *int) error
	insert = func(steps []api.WorkflowStep, parent *int) error {
		for pos, s := range steps {
			stepID := next
			next++
			kind, children := "stage", (*[]api.WorkflowStep)(nil)
			switch {
			case s.Order != nil:
				kind, children = "order", s.Order
			case s.Parallel != nil:
				kind, children = "parallel", s.Parallel
			}
			_, err := tx.ExecContext(ctx, `INSERT INTO workflow_steps (definition_id, id, parent, position, kind, stage, model, prompt) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
				id, stepID, parent, pos, kind, s.Stage, s.Model, s.Prompt)
			if err != nil {
				return err
			}
			if s.Agents != nil {
				for i, a := range *s.Agents {
					_, err = tx.ExecContext(ctx, `INSERT INTO workflow_agents (definition_id, step_id, position, name, model, prompt) VALUES (?, ?, ?, ?, ?, ?)`,
						id, stepID, i, a.Name, a.Model, a.Prompt)
					if err != nil {
						return err
					}
				}
			}
			if children != nil {
				err = insert(*children, &stepID)
				if err != nil {
					return err
				}
			}
		}
		return nil
	}
	return id, insert(def.Steps, nil)
}

type stepRow struct {
	id   int
	kind string
	step api.WorkflowStep
}

func loadDefinition(ctx context.Context, q querier, id string) (api.WorkflowDefinition, error) {
	def := api.WorkflowDefinition{}
	err := q.QueryRowContext(ctx, `SELECT description FROM workflow_definitions WHERE id = ?`, id).Scan(&def.Description)
	if errors.Is(err, sql.ErrNoRows) {
		return def, fmt.Errorf("%w: no workflow definition %s", ErrNotFound, id)
	}
	if err != nil {
		return def, err
	}

	agents := map[int][]api.TeamMember{}
	rows, err := q.QueryContext(ctx, `SELECT step_id, name, model, prompt FROM workflow_agents WHERE definition_id = ? ORDER BY step_id, position`, id)
	if err != nil {
		return def, err
	}
	for rows.Next() {
		step := 0
		a := api.TeamMember{}
		err = rows.Scan(&step, &a.Name, &a.Model, &a.Prompt)
		if err != nil {
			rows.Close()
			return def, err
		}
		agents[step] = append(agents[step], a)
	}
	err = errors.Join(rows.Err(), rows.Close())
	if err != nil {
		return def, err
	}

	children := map[int][]stepRow{}
	rows, err = q.QueryContext(ctx, `SELECT id, coalesce(parent, -1), kind, stage, model, prompt FROM workflow_steps WHERE definition_id = ? ORDER BY coalesce(parent, -1), position`, id)
	if err != nil {
		return def, err
	}
	for rows.Next() {
		r := stepRow{}
		parent := 0
		err = rows.Scan(&r.id, &parent, &r.kind, &r.step.Stage, &r.step.Model, &r.step.Prompt)
		if err != nil {
			rows.Close()
			return def, err
		}
		if a, found := agents[r.id]; found {
			r.step.Agents = &a
		}
		children[parent] = append(children[parent], r)
	}
	err = errors.Join(rows.Err(), rows.Close())
	if err != nil {
		return def, err
	}

	var build func(parent int) []api.WorkflowStep
	build = func(parent int) []api.WorkflowStep {
		steps := []api.WorkflowStep{}
		for _, r := range children[parent] {
			switch r.kind {
			case "order":
				sub := build(r.id)
				r.step.Order = &sub
			case "parallel":
				sub := build(r.id)
				r.step.Parallel = &sub
			}
			steps = append(steps, r.step)
		}
		return steps
	}
	def.Steps = build(-1)
	return def, nil
}