1 files changed, 22 insertions, 1 deletions
diff --git a/models/models.go b/models/models.go
index 2e2ef34..918e35e 100644
--- a/models/models.go
+++ b/models/models.go
@@ -1,8 +1,8 @@
 package models
 
 import (
-	"gf-lt/config"
 	"fmt"
+	"gf-lt/config"
 	"strings"
 )
 
@@ -76,6 +76,27 @@ type ChatBody struct {
 	Messages []RoleMsg `json:"messages"`
 }
 
+func (cb *ChatBody) Rename(oldname, newname string) {
+	for i, m := range cb.Messages {
+		cb.Messages[i].Content = strings.ReplaceAll(m.Content, oldname, newname)
+		cb.Messages[i].Role = strings.ReplaceAll(m.Role, oldname, newname)
+	}
+}
+
+func (cb *ChatBody) ListRoles() []string {
+	namesMap := make(map[string]struct{})
+	for _, m := range cb.Messages {
+		namesMap[m.Role] = struct{}{}
+	}
+	resp := make([]string, len(namesMap))
+	i := 0
+	for k := range namesMap {
+		resp[i] = k
+		i++
+	}
+	return resp
+}
+
 type ChatToolsBody struct {
 	Model    string    `json:"model"`
 	Messages []RoleMsg `json:"messages"`