package red import ( "context" "database/sql" "errors" "fmt" "strings" "ymgb/database" "ymgb/database/psql" "ymgb/utils" "github.com/google/uuid" "github.com/redis/go-redis/v9" ) var ctx = context.Background() type RPRepository struct { client *redis.Client db *database.Context } type RPChat struct { ID uuid.UUID UserID int64 WaifuID int64 Prompt string Counter int ChatTokens int64 SettingID int Setting *psql.RPSetting ScenariosIDs []int Scenarios []psql.RPScenario } func NewRPRepository(db *database.Context) RPRepository { return RPRepository{db.Redis, db} } func (rep RPRepository) SetSelectedWaifu(userId, waifuId int64) error { key := fmt.Sprintf("ai.chats.rp.%d", userId) return rep.client.Set(ctx, key, waifuId, 0).Err() } func (rep RPRepository) GetSelectedWaifu(userId int64) int64 { key := fmt.Sprintf("ai.chats.rp.%d", userId) res := rep.client.Get(ctx, key) if res.Err() != nil { return 0 } i, _ := res.Int() return int64(i) } func (rep RPRepository) SetChatId(userId, waifuId int64, 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 int64) 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 int64) (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) GetChat(userId int64) (RPChat, error) { var chat RPChat waifuId := rep.GetSelectedWaifu(userId) if waifuId == 0 { return chat, errors.New("no WaifuID") } chatId, err := rep.GetOrCreateChatId(userId, waifuId) if err != nil { return chat, err } chat.ID = uuid.MustParse(chatId) chat.UserID = userId chat.WaifuID = waifuId chat.Prompt = rep.GetChatPrompt(userId, waifuId) chat.Counter = rep.GetCounter(userId, waifuId) chat.ChatTokens = int64(rep.GetChatTokens(userId, waifuId)) chat.SettingID = rep.GetChatSettingID(userId, waifuId) psqlRep := psql.NewRPRepository(rep.db) if chat.SettingID > 0 { setting, err := psqlRep.GetSetting(chat.SettingID) if err != nil { if errors.Is(err, sql.ErrNoRows) { return chat, fmt.Errorf("setting %d not found", chat.SettingID) } return chat, err } chat.Setting = &setting } chat.ScenariosIDs = rep.GetChatScenariosIDs(userId, waifuId) chat.Scenarios = make([]psql.RPScenario, len(chat.ScenariosIDs)) for i, id := range chat.ScenariosIDs { if id <= 0 { continue } chat.Scenarios[i], err = psqlRep.GetScenario(id) if err != nil { if errors.Is(err, sql.ErrNoRows) { return chat, fmt.Errorf("scenario %d not found", id) } return chat, err } } return chat, nil } func (rep RPRepository) SaveChat(chat RPChat) error { waifuId := chat.WaifuID userId := chat.UserID var err error if err = rep.SetChatPrompt(userId, waifuId, chat.Prompt); err != nil { return err } if err = rep.SetCounter(userId, waifuId, chat.Counter); err != nil { return err } if err = rep.SetChatTokens(userId, waifuId, chat.ChatTokens); err != nil { return err } var settingId int if chat.Setting != nil { settingId = chat.Setting.ID } else { settingId = chat.SettingID } if err = rep.SetChatSettingID(userId, waifuId, settingId); err != nil { return err } var ids []int if len(chat.Scenarios) > 0 { ids = utils.Map(chat.Scenarios, func(s psql.RPScenario) int { return s.ID }) } else { ids = chat.ScenariosIDs } if err = rep.SetChatScenariosIDs(userId, waifuId, ids); err != nil { return err } return nil } func (rep RPRepository) SetChatPrompt(userId, waifuId int64, 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 int64) 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 int64, 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 int64) 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 int64) 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 int64) int64 { 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 int64(i) } func (rep RPRepository) SetChatSettingID(userId, waifuId int64, 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 int64) 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 int64, scenarioIds []int) error { key := fmt.Sprintf("ai.chats.rp.%d.%d.scenario_id", userId, waifuId) ids := strings.Join(utils.Map(scenarioIds, utils.AnyToString), ",") return rep.client.Set(ctx, key, ids, 0).Err() } func (rep RPRepository) GetChatScenariosIDs(userId, waifuId int64) []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) }