支持 TOKEN 设置最大调用次数

server/db.go

@@ -43,6 +43,16 @@ func PutToken(token types.Token) error {
 	return db.Put(key, token)
 }
 
+func GetToken(name string) (types.Token, error) {
+	key := TokenPrefix + name
+	token, err := db.Get(key)
+	if err != nil {
+		return types.Token{}, err
+	}
+
+	return token.(types.Token), nil
+}
+
 func RemoveToken(token string) error {
 	key := TokenPrefix + token
 	return db.Delete(key)
@@ -51,7 +61,22 @@ func RemoveToken(token string) error {
 // GetChatRoles 获取聊天角色
 // chat/roles
 func GetChatRoles() map[string]types.ChatRole {
-	return nil
+	items := db.Search(ChatRolePrefix)
+	var roles = make(map[string]types.ChatRole)
+	for _, v := range items {
+		var role types.ChatRole
+		err := json.Unmarshal([]byte(v), &role)
+		if err != nil {
+			continue
+		}
+		roles[role.Key] = role
+	}
+	return roles
+}
+
+func PutChatRole(role types.ChatRole) error {
+	key := ChatRolePrefix + role.Key
+	return db.Put(key, role)
 }
 
 // GetChatHistory 获取聊天历史记录