feat: yaml conversation trees more robust

1 files changed, 27 insertions, 1 deletions

diff --git a/internal/behavior/types.go b/internal/behavior/types.go
index 3c6e290..f8e65f7 100644
--- a/internal/behavior/types.go
+++ b/internal/behavior/types.go
@@ -1,6 +1,11 @@
 package behavior
 
-import "strings"
+import (
+	"math/rand"
+	"strings"
+
+	"gopkg.in/yaml.v3"
+)
 
 func WordPrefixMatch(input, name string) bool {
 	inputWords := strings.Fields(strings.ToLower(input))
@@ -186,3 +191,24 @@ type DropEntry struct {
 type DropTableDef struct {
 	Drops []DropEntry `yaml:"drops"`
 }
+
+type MessageList []string
+
+func (ml *MessageList) UnmarshalYAML(value *yaml.Node) error {
+	if value.Kind == yaml.ScalarNode {
+		var s string
+		if err := value.Decode(&s); err != nil {
+			return err
+		}
+		*ml = MessageList{s}
+		return nil
+	}
+	return value.Decode((*[]string)(ml))
+}
+
+func (ml MessageList) Pick() string {
+	if len(ml) == 0 {
+		return ""
+	}
+	return ml[rand.Intn(len(ml))]
+}