Add weighted random workflow steps and a random choice agent tool

internal/store/definition.go

@@ -28,9 +28,11 @@ func insertDefinition(ctx context.Context, tx *sql.Tx, def api.WorkflowDefinitio
 				kind, children = "order", s.Order
 			case s.Parallel != nil:
 				kind, children = "parallel", s.Parallel
+			case s.Random != nil:
+				kind, children = "random", s.Random
 			}
-			_, 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)
+			_, err := tx.ExecContext(ctx, `INSERT INTO workflow_steps (definition_id, id, parent, position, kind, stage, model, prompt, weight) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+				id, stepID, parent, pos, kind, s.Stage, s.Model, s.Prompt, s.Weight)
 			if err != nil {
 				return err
 			}
@@ -92,14 +94,14 @@ func loadDefinition(ctx context.Context, q querier, id string) (api.WorkflowDefi
 	}
 
 	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)
+	rows, err = q.QueryContext(ctx, `SELECT id, coalesce(parent, -1), kind, stage, model, prompt, weight 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)
+		err = rows.Scan(&r.id, &parent, &r.kind, &r.step.Stage, &r.step.Model, &r.step.Prompt, &r.step.Weight)
 		if err != nil {
 			rows.Close()
 			return def, err
@@ -125,6 +127,9 @@ func loadDefinition(ctx context.Context, q querier, id string) (api.WorkflowDefi
 			case "parallel":
 				sub := build(r.id)
 				r.step.Parallel = &sub
+			case "random":
+				sub := build(r.id)
+				r.step.Random = &sub
 			}
 			steps = append(steps, r.step)
 		}

internal/store/migrations/105-add-random-workflow-steps.sql

@@ -0,0 +1,17 @@
+CREATE TABLE workflow_steps_next (
+	definition_id TEXT NOT NULL REFERENCES workflow_definitions(id),
+	id INTEGER NOT NULL,
+	parent INTEGER,
+	position INTEGER NOT NULL,
+	kind TEXT NOT NULL CHECK (kind IN ('stage', 'order', 'parallel', 'random')),
+	stage TEXT NOT NULL,
+	model TEXT NOT NULL,
+	prompt TEXT NOT NULL,
+	weight REAL NOT NULL DEFAULT 0,
+	PRIMARY KEY (definition_id, id),
+	UNIQUE (definition_id, parent, position)
+);
+INSERT INTO workflow_steps_next (definition_id, id, parent, position, kind, stage, model, prompt)
+	SELECT definition_id, id, parent, position, kind, stage, model, prompt FROM workflow_steps;
+DROP TABLE workflow_steps;
+ALTER TABLE workflow_steps_next RENAME TO workflow_steps;

internal/store/workflow.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 
 	"g.fc.run/ai/byre/internal/api"
+	"g.fc.run/ai/byre/internal/workflow"
 )
 
 type FrozenWorkflow struct {
@@ -38,7 +39,7 @@ func (s *Store) ListWorkflows(ctx context.Context) ([]api.WorkflowSummary, error
 }
 
 func (s *Store) GetWorkflow(ctx context.Context, name string) (api.WorkflowDefinition, error) {
-	f, err := s.FreezeWorkflow(ctx, name)
+	f, err := s.savedWorkflow(ctx, name)
 	return f.Definition, err
 }
 
@@ -192,6 +193,20 @@ func (s *Store) DeleteWorkflow(ctx context.Context, name, actor string) error {
 }
 
 func (s *Store) FreezeWorkflow(ctx context.Context, name string) (FrozenWorkflow, error) {
+	f, err := s.savedWorkflow(ctx, name)
+	if err != nil || !workflow.HasRandom(f.Definition) {
+		return f, err
+	}
+	f.Definition = workflow.Resolve(f.Definition)
+	err = s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		id, err := insertDefinition(ctx, tx, f.Definition)
+		f.DefinitionID = id
+		return err
+	})
+	return f, err
+}
+
+func (s *Store) savedWorkflow(ctx context.Context, name string) (FrozenWorkflow, error) {
 	f := FrozenWorkflow{Name: name}
 	err := s.db.QueryRowContext(ctx, `SELECT definition_id FROM workflows WHERE name = ?`, name).Scan(&f.DefinitionID)
 	if errors.Is(err, sql.ErrNoRows) {