diff --git a/database/mdb/rp_chats.go b/database/mdb/rp_chats.go index 2ea07a3..389c93a 100644 --- a/database/mdb/rp_chats.go +++ b/database/mdb/rp_chats.go @@ -38,3 +38,9 @@ func UpdateChatHistory(db *laniakea.DatabaseContext, chatId, role, message strin }) return err } +func GetChatHistorySize(db *laniakea.DatabaseContext, chatId string) (int64, error) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + col := database.GetMongoCollection(db, "rp_chat_messages") + return col.CountDocuments(ctx, bson.M{"chat_id": chatId}) +} diff --git a/plugins/rp.go b/plugins/rp.go index 6899e1c..2222fa4 100644 --- a/plugins/rp.go +++ b/plugins/rp.go @@ -26,6 +26,7 @@ func RegisterRP(bot *laniakea.Bot) { rp = rp.Command(rpUserPromptGet, "rpuserpget") rp = rp.Command(rpUserPromptSet, "rpuserpset") rp = rp.Command(generate, "g", "gen", "г") + rp = rp.Command(chatStat, "chatstat") rp = rp.Payload(debugTokens, "rp.tokens") bot.AddPlugins(rp.Build()) @@ -350,3 +351,29 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { ctx.Error(err) } } + +func chatStat(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { + redisRpRep := red.NewRPRepository(db) + waifuId := redisRpRep.GetSelectedWaifu(ctx.FromID) + if waifuId == 0 { + ctx.Answer("Не выбрана вайфу") + return + } + chatId := redisRpRep.GetChatId(ctx.FromID, waifuId) + if chatId == "" { + chatId = uuid.New().String() + err := redisRpRep.SetChatId(ctx.FromID, waifuId, chatId) + if err != nil { + ctx.Error(err) + return + } + } + messageCount, err := mdb.GetChatHistorySize(db, chatId) + if err != nil { + ctx.Error(err) + return + } + + tokens := redisRpRep.GetChatTokens(ctx.FromID, waifuId) + ctx.Answerf("Кол-во сообщений: %d\nКол-во токенов: %d", messageCount, tokens) +}