some changes

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
-
 }

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:

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)
 	}