127 lines
3.6 KiB
Go
127 lines
3.6 KiB
Go
package red
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"kurumibot/utils"
|
|
"strings"
|
|
|
|
"git.nix13.pw/scuroneko/laniakea"
|
|
"github.com/google/uuid"
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
var ctx = context.Background()
|
|
|
|
type RPRepository struct {
|
|
client *redis.Client
|
|
}
|
|
|
|
func NewRPRepository(db *laniakea.DatabaseContext) *RPRepository {
|
|
return &RPRepository{client: db.Redis}
|
|
}
|
|
|
|
func (rep *RPRepository) SetSelectedWaifu(userId, waifuId int) error {
|
|
key := fmt.Sprintf("ai.chats.rp.%d", userId)
|
|
return rep.client.Set(ctx, key, waifuId, 0).Err()
|
|
}
|
|
func (rep *RPRepository) GetSelectedWaifu(userId int) int {
|
|
key := fmt.Sprintf("ai.chats.rp.%d", userId)
|
|
res := rep.client.Get(ctx, key)
|
|
if res.Err() != nil {
|
|
return 0
|
|
}
|
|
i, _ := res.Int()
|
|
return i
|
|
}
|
|
|
|
func (rep *RPRepository) SetChatId(userId, waifuId int, chatId string) error {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.chat", userId, waifuId)
|
|
return rep.client.Set(context.Background(), key, chatId, 0).Err()
|
|
}
|
|
func (rep *RPRepository) GetChatId(userId, waifuId int) string {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.chat", userId, waifuId)
|
|
res := rep.client.Get(ctx, key)
|
|
if res.Err() != nil {
|
|
return ""
|
|
}
|
|
return res.Val()
|
|
}
|
|
func (rep *RPRepository) GetOrCreateChatId(userId, waifuId int) (string, error) {
|
|
chatId := rep.GetChatId(userId, waifuId)
|
|
if chatId == "" {
|
|
chatId = uuid.New().String()
|
|
}
|
|
err := rep.SetChatId(userId, waifuId, chatId)
|
|
return chatId, err
|
|
}
|
|
|
|
func (rep *RPRepository) SetChatPrompt(userId, waifuId int, prompt string) error {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.prompt", userId, waifuId)
|
|
return rep.client.Set(context.Background(), key, prompt, 0).Err()
|
|
}
|
|
func (rep *RPRepository) GetChatPrompt(userId, waifuId int) string {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.prompt", userId, waifuId)
|
|
res := rep.client.Get(ctx, key)
|
|
if res.Err() != nil {
|
|
return ""
|
|
}
|
|
return res.Val()
|
|
}
|
|
|
|
func (rep *RPRepository) SetCounter(userId, waifuId, counter int) error {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.counter", userId, waifuId)
|
|
return rep.client.Set(ctx, key, counter, 0).Err()
|
|
}
|
|
func (rep *RPRepository) GetCounter(userId, waifuId int) int {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.counter", userId, waifuId)
|
|
res := rep.client.Get(ctx, key)
|
|
if res.Err() != nil {
|
|
return 0
|
|
}
|
|
i, _ := res.Int()
|
|
return i
|
|
}
|
|
|
|
func (rep *RPRepository) SetChatTokens(userId, waifuId, tokens int) error {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.tokens", userId, waifuId)
|
|
return rep.client.Set(ctx, key, tokens, 0).Err()
|
|
}
|
|
func (rep *RPRepository) GetChatTokens(userId, waifuId int) int {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.tokens", userId, waifuId)
|
|
res := rep.client.Get(ctx, key)
|
|
if res.Err() != nil {
|
|
return 0
|
|
}
|
|
i, _ := res.Int()
|
|
return i
|
|
}
|
|
|
|
func (rep *RPRepository) SetChatSettingID(userId, waifuId, settingId int) error {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.setting_id", userId, waifuId)
|
|
return rep.client.Set(ctx, key, settingId, 0).Err()
|
|
}
|
|
func (rep *RPRepository) GetChatSettingID(userId, waifuId int) int {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.setting_id", userId, waifuId)
|
|
res := rep.client.Get(ctx, key)
|
|
if res.Err() != nil {
|
|
return 0
|
|
}
|
|
i, _ := res.Int()
|
|
return i
|
|
}
|
|
|
|
func (rep *RPRepository) SetChatScenariosIDs(userId, waifuId int, scenarioIds string) error {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.scenario_id", userId, waifuId)
|
|
return rep.client.Set(ctx, key, scenarioIds, 0).Err()
|
|
}
|
|
func (rep *RPRepository) GetChatScenariosIDs(userId, waifuId int) []int {
|
|
key := fmt.Sprintf("ai.chats.rp.%d.%d.scenario_id", userId, waifuId)
|
|
res := rep.client.Get(ctx, key)
|
|
if res.Err() != nil {
|
|
return []int{0}
|
|
}
|
|
ids := strings.Split(res.Val(), ",")
|
|
return utils.Map(ids, utils.StringToInt)
|
|
}
|