diff --git a/database/red/rp_chats.go b/database/red/rp_chats.go index 21d8856..585f534 100644 --- a/database/red/rp_chats.go +++ b/database/red/rp_chats.go @@ -4,17 +4,27 @@ import ( "context" "fmt" "kurumibot/laniakea" + + "github.com/redis/go-redis/v9" ) var ctx = context.Background() -func RPSetSelectedWaifu(db *laniakea.DatabaseContext, userId, waifuId int) error { - key := fmt.Sprintf("ai.chats.rp.%d", userId) - return db.Redis.Set(ctx, key, waifuId, 0).Err() +type RPRepository struct { + client *redis.Client } -func RPGetSelectedWaifu(db *laniakea.DatabaseContext, userId int) int { + +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) - res := db.Redis.Get(ctx, key) + 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 } @@ -22,39 +32,39 @@ func RPGetSelectedWaifu(db *laniakea.DatabaseContext, userId int) int { return i } -func RPSetChatId(db *laniakea.DatabaseContext, userId, waifuId int, chatId string) error { +func (rep *RPRepository) SetChatId(userId, waifuId int, chatId string) error { key := fmt.Sprintf("ai.chats.rp.%d.%d.chat", userId, waifuId) - return db.Redis.Set(context.Background(), key, chatId, 0).Err() + return rep.client.Set(context.Background(), key, chatId, 0).Err() } -func RPGetChatId(db *laniakea.DatabaseContext, userId, waifuId int) string { +func (rep *RPRepository) GetChatId(userId, waifuId int) string { key := fmt.Sprintf("ai.chats.rp.%d.%d.chat", userId, waifuId) - res := db.Redis.Get(ctx, key) + res := rep.client.Get(ctx, key) if res.Err() != nil { return "" } return res.Val() } -func RPSetChatPrompt(db *laniakea.DatabaseContext, userId, waifuId int, prompt string) error { +func (rep *RPRepository) SetChatPrompt(userId, waifuId int, prompt string) error { key := fmt.Sprintf("ai.chats.rp.%d.%d.prompt", userId, waifuId) - return db.Redis.Set(context.Background(), key, prompt, 0).Err() + return rep.client.Set(context.Background(), key, prompt, 0).Err() } -func RPGetChatPrompt(db *laniakea.DatabaseContext, userId, waifuId int) string { +func (rep *RPRepository) GetChatPrompt(userId, waifuId int) string { key := fmt.Sprintf("ai.chats.rp.%d.%d.prompt", userId, waifuId) - res := db.Redis.Get(ctx, key) + res := rep.client.Get(ctx, key) if res.Err() != nil { return "" } return res.Val() } -func RPSetCounter(db *laniakea.DatabaseContext, userId, waifuId, counter int) error { +func (rep *RPRepository) SetCounter(userId, waifuId, counter int) error { key := fmt.Sprintf("ai.chats.rp.%d.%d.counter", userId, waifuId) - return db.Redis.Set(ctx, key, counter, 0).Err() + return rep.client.Set(ctx, key, counter, 0).Err() } -func RPGetCounter(db *laniakea.DatabaseContext, userId, waifuId int) int { +func (rep *RPRepository) GetCounter(userId, waifuId int) int { key := fmt.Sprintf("ai.chats.rp.%d.%d.counter", userId, waifuId) - res := db.Redis.Get(ctx, key) + res := rep.client.Get(ctx, key) if res.Err() != nil { return 0 } @@ -62,20 +72,16 @@ func RPGetCounter(db *laniakea.DatabaseContext, userId, waifuId int) int { return i } -func RPSetIndex(db *laniakea.DatabaseContext, userId, waifuId, index int) error { - key := fmt.Sprintf("ai.chats.rp.%d.%d.index", userId, waifuId) - return db.Redis.Set(ctx, key, index, 0).Err() +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 RPGetIndex(db *laniakea.DatabaseContext, userId, waifuId int) int { - key := fmt.Sprintf("ai.chats.rp.%d.%d.index", userId, waifuId) - res := db.Redis.Get(ctx, key) +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, err := res.Int() - if err != nil { - return 0 - } + i, _ := res.Int() return i - } diff --git a/db.docker-compose.yml b/db.docker-compose.yml index 62d6e85..4b543ed 100644 --- a/db.docker-compose.yml +++ b/db.docker-compose.yml @@ -3,7 +3,6 @@ services: image: postgres:18.1-alpine networks: - bot - restart: always ports: - "5432:5432" environment: @@ -22,7 +21,6 @@ services: image: mongo:noble networks: - bot - restart: always ports: - "27017:27017" environment: @@ -36,7 +34,6 @@ services: image: redis:8.4-alpine networks: - bot - restart: always ports: - "6379:6379" volumes: diff --git a/laniakea b/laniakea index 2162378..ce13b19 160000 --- a/laniakea +++ b/laniakea @@ -1 +1 @@ -Subproject commit 21623788c610d089cac54fd0c72ddc4a7931e525 +Subproject commit ce13b1967658012da013965d60f44f0cd4a6ccf5 diff --git a/plugins/rp.go b/plugins/rp.go index f771094..6899e1c 100644 --- a/plugins/rp.go +++ b/plugins/rp.go @@ -26,17 +26,26 @@ func RegisterRP(bot *laniakea.Bot) { rp = rp.Command(rpUserPromptGet, "rpuserpget") rp = rp.Command(rpUserPromptSet, "rpuserpset") rp = rp.Command(generate, "g", "gen", "г") + rp = rp.Payload(debugTokens, "rp.tokens") bot.AddPlugins(rp.Build()) } +func debugTokens(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { + rep := red.NewRPRepository(db) + waifuId := rep.GetSelectedWaifu(ctx.FromID) + tokens := rep.GetChatTokens(ctx.FromID, waifuId) + ctx.Answerf("%d", tokens) +} + func selectWaifu(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { waifuId, err := strconv.Atoi(ctx.Args[0]) if err != nil { ctx.Error(err) return } - err = red.RPSetSelectedWaifu(db, ctx.FromID, waifuId) + rpRep := red.NewRPRepository(db) + err = rpRep.SetSelectedWaifu(ctx.FromID, waifuId) if err != nil { ctx.Error(err) return @@ -98,13 +107,14 @@ func rpScenarioList(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { ctx.Answer(strings.Join(out, "\n")) } func newChat(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { - waifuId := red.RPGetSelectedWaifu(db, ctx.FromID) + redisRpRep := red.NewRPRepository(db) + waifuId := redisRpRep.GetSelectedWaifu(ctx.FromID) if waifuId == 0 { ctx.Answer("Не выбрана вайфу") return } chatId := uuid.New() - err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId.String()) + err := redisRpRep.SetChatId(ctx.FromID, waifuId, chatId.String()) if err != nil { ctx.Error(err) return @@ -125,12 +135,12 @@ func newChat(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { } return } - err = red.RPSetChatPrompt(db, ctx.FromID, waifuId, scenario.Prompt) + err = redisRpRep.SetChatPrompt(ctx.FromID, waifuId, scenario.Prompt) if err != nil { ctx.Error(err) return } - err = red.RPSetCounter(db, ctx.FromID, waifuId, 0) + err = redisRpRep.SetCounter(ctx.FromID, waifuId, 0) if err != nil { ctx.Error(err) return @@ -172,15 +182,16 @@ func rpUserPromptSet(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { } func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { - waifuId := red.RPGetSelectedWaifu(db, ctx.FromID) + redisRpRep := red.NewRPRepository(db) + waifuId := redisRpRep.GetSelectedWaifu(ctx.FromID) if waifuId == 0 { ctx.Answer("Не выбрана вайфу") return } - chatId := red.RPGetChatId(db, ctx.FromID, waifuId) + chatId := redisRpRep.GetChatId(ctx.FromID, waifuId) if chatId == "" { chatId = uuid.New().String() - err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId) + err := redisRpRep.SetChatId(ctx.FromID, waifuId, chatId) if err != nil { ctx.Error(err) return @@ -214,7 +225,7 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { ai.FormatPrompt(preset.PreHistory, waifu.Name, ctx.Msg.From.FirstName), fmt.Sprintf("Вот краткое описание твоего персонажа: %s", waifu.RpPrompt), userPrompt, - red.RPGetChatPrompt(db, ctx.FromID, waifuId), + redisRpRep.GetChatPrompt(ctx.FromID, waifuId), ), } afterHistory := ai.Message{ @@ -267,17 +278,28 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { if err != nil { ctx.Error(err) } + rpUser.UsedTokens = rpUser.UsedTokens + res.Usage.TotalTokens + tokens := redisRpRep.GetChatTokens(ctx.FromID, waifuId) + tokens += int(res.Usage.CompletionTokens) + err = redisRpRep.SetChatTokens(ctx.FromID, waifuId, tokens) + if err != nil { + ctx.Error(err) + } err = rpRep.UpdateUser(rpUser) if err != nil { ctx.Error(err) - return } - m.Delete() - ctx.Answer(laniakea.EscapeMarkdown(agentAnswer.Content)) - counter := red.RPGetCounter(db, ctx.FromID, waifuId) + 1 - if counter == 10 { + m.Delete() + //kb := laniakea.NewInlineKeyboard() + //kb.AddCallbackButton("Перегенерировать", laniakea.CallbackData{Command: "rp.tokens"}) + //ctx.Keyboard("Test", kb) + ctx.Answer(laniakea.EscapeMarkdown(agentAnswer.Content)) + return + + counter := redisRpRep.GetCounter(ctx.FromID, waifuId) + 1 + if counter == 20 { m := ctx.Answer("Запущено сжатие чата.") history, err = mdb.GetChatHistory(db, chatId) if err != nil { @@ -304,7 +326,7 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { compressedHistory := res.Choices[0].Message.Content chatId = uuid.New().String() - err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId) + err := redisRpRep.SetChatId(ctx.FromID, waifuId, chatId) if err != nil { ctx.Error(err) return @@ -323,7 +345,7 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { counter = 0 m.Edit("Сжатие завершено") } - err = red.RPSetCounter(db, ctx.FromID, waifuId, counter) + err = redisRpRep.SetCounter(ctx.FromID, waifuId, counter) if err != nil { ctx.Error(err) }