ScuroNeko/YaeMikoBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

some changes

Browse Source
This commit is contained in:
ScuroNeko
2026-01-21 15:58:30 +03:00
parent 02590a2a46
commit b8391e9c6b
4 changed files with 73 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
database/red/rp_chats.go
View File

@@ -4,17 +4,27 @@ import (
"context" "context"
"fmt" "fmt"
"kurumibot/laniakea" "kurumibot/laniakea"
"github.com/redis/go-redis/v9"
) )
var ctx = context.Background() var ctx = context.Background()
func RPSetSelectedWaifu(db *laniakea.DatabaseContext, userId, waifuId int) error { type RPRepository struct {
key := fmt.Sprintf("ai.chats.rp.%d", userId) client *redis.Client
return db.Redis.Set(ctx, key, waifuId, 0).Err()
} }
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) 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 { if res.Err() != nil {
return 0 return 0
} }
@@ -22,39 +32,39 @@ func RPGetSelectedWaifu(db *laniakea.DatabaseContext, userId int) int {
return i 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) 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) 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 { if res.Err() != nil {
return "" return ""
} }
return res.Val() 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) 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) 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 { if res.Err() != nil {
return "" return ""
} }
return res.Val() 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) 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) 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 { if res.Err() != nil {
return 0 return 0
} }
@@ -62,20 +72,16 @@ func RPGetCounter(db *laniakea.DatabaseContext, userId, waifuId int) int {
return i return i
} }
func RPSetIndex(db *laniakea.DatabaseContext, userId, waifuId, index int) error { func (rep *RPRepository) SetChatTokens(userId, waifuId, tokens int) error {
key := fmt.Sprintf("ai.chats.rp.%d.%d.index", userId, waifuId) key := fmt.Sprintf("ai.chats.rp.%d.%d.tokens", userId, waifuId)
return db.Redis.Set(ctx, key, index, 0).Err() return rep.client.Set(ctx, key, tokens, 0).Err()
} }
func RPGetIndex(db *laniakea.DatabaseContext, userId, waifuId int) int { func (rep *RPRepository) GetChatTokens(userId, waifuId int) int {
key := fmt.Sprintf("ai.chats.rp.%d.%d.index", userId, waifuId) key := fmt.Sprintf("ai.chats.rp.%d.%d.tokens", userId, waifuId)
res := db.Redis.Get(ctx, key) res := rep.client.Get(ctx, key)
if res.Err() != nil { if res.Err() != nil {
return 0 return 0
} }
i, err := res.Int() i, _ := res.Int()
if err != nil {
return 0
}
return i return i
} }

3
db.docker-compose.yml
View File

@@ -3,7 +3,6 @@ services:
image: postgres:18.1-alpine image: postgres:18.1-alpine
networks: networks:
- bot - bot
restart: always
ports: ports:
- "5432:5432" - "5432:5432"
environment: environment:
@@ -22,7 +21,6 @@ services:
image: mongo:noble image: mongo:noble
networks: networks:
- bot - bot
restart: always
ports: ports:
- "27017:27017" - "27017:27017"
environment: environment:
@@ -36,7 +34,6 @@ services:
image: redis:8.4-alpine image: redis:8.4-alpine
networks: networks:
- bot - bot
restart: always
ports: ports:
- "6379:6379" - "6379:6379"
volumes: volumes:

2
laniakea

Submodule laniakea updated: 21623788c6...ce13b19676

54
plugins/rp.go
View File

@@ -26,17 +26,26 @@ func RegisterRP(bot *laniakea.Bot) {
rp = rp.Command(rpUserPromptGet, "rpuserpget") rp = rp.Command(rpUserPromptGet, "rpuserpget")
rp = rp.Command(rpUserPromptSet, "rpuserpset") rp = rp.Command(rpUserPromptSet, "rpuserpset")
rp = rp.Command(generate, "g", "gen", "г") rp = rp.Command(generate, "g", "gen", "г")
rp = rp.Payload(debugTokens, "rp.tokens")
bot.AddPlugins(rp.Build()) 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) { func selectWaifu(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
waifuId, err := strconv.Atoi(ctx.Args[0]) waifuId, err := strconv.Atoi(ctx.Args[0])
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
return return
} }
err = red.RPSetSelectedWaifu(db, ctx.FromID, waifuId) rpRep := red.NewRPRepository(db)
err = rpRep.SetSelectedWaifu(ctx.FromID, waifuId)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
return return
@@ -98,13 +107,14 @@ func rpScenarioList(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
ctx.Answer(strings.Join(out, "\n")) ctx.Answer(strings.Join(out, "\n"))
} }
func newChat(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { 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 { if waifuId == 0 {
ctx.Answer("Не выбрана вайфу") ctx.Answer("Не выбрана вайфу")
return return
} }
chatId := uuid.New() chatId := uuid.New()
err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId.String()) err := redisRpRep.SetChatId(ctx.FromID, waifuId, chatId.String())
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
return return
@@ -125,12 +135,12 @@ func newChat(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
} }
return return
} }
err = red.RPSetChatPrompt(db, ctx.FromID, waifuId, scenario.Prompt) err = redisRpRep.SetChatPrompt(ctx.FromID, waifuId, scenario.Prompt)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
return return
} }
err = red.RPSetCounter(db, ctx.FromID, waifuId, 0) err = redisRpRep.SetCounter(ctx.FromID, waifuId, 0)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
return return
@@ -172,15 +182,16 @@ func rpUserPromptSet(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
} }
func generate(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 { if waifuId == 0 {
ctx.Answer("Не выбрана вайфу") ctx.Answer("Не выбрана вайфу")
return return
} }
chatId := red.RPGetChatId(db, ctx.FromID, waifuId) chatId := redisRpRep.GetChatId(ctx.FromID, waifuId)
if chatId == "" { if chatId == "" {
chatId = uuid.New().String() chatId = uuid.New().String()
err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId) err := redisRpRep.SetChatId(ctx.FromID, waifuId, chatId)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
return return
@@ -214,7 +225,7 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
ai.FormatPrompt(preset.PreHistory, waifu.Name, ctx.Msg.From.FirstName), ai.FormatPrompt(preset.PreHistory, waifu.Name, ctx.Msg.From.FirstName),
fmt.Sprintf("Вот краткое описание твоего персонажа: %s", waifu.RpPrompt), fmt.Sprintf("Вот краткое описание твоего персонажа: %s", waifu.RpPrompt),
userPrompt, userPrompt,
red.RPGetChatPrompt(db, ctx.FromID, waifuId), redisRpRep.GetChatPrompt(ctx.FromID, waifuId),
), ),
} }
afterHistory := ai.Message{ afterHistory := ai.Message{
@@ -267,17 +278,28 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
} }
rpUser.UsedTokens = rpUser.UsedTokens + res.Usage.TotalTokens 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) err = rpRep.UpdateUser(rpUser)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
return
} }
m.Delete()
ctx.Answer(laniakea.EscapeMarkdown(agentAnswer.Content))
counter := red.RPGetCounter(db, ctx.FromID, waifuId) + 1 m.Delete()
if counter == 10 { //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("Запущено сжатие чата.") m := ctx.Answer("Запущено сжатие чата.")
history, err = mdb.GetChatHistory(db, chatId) history, err = mdb.GetChatHistory(db, chatId)
if err != nil { if err != nil {
@@ -304,7 +326,7 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
compressedHistory := res.Choices[0].Message.Content compressedHistory := res.Choices[0].Message.Content
chatId = uuid.New().String() chatId = uuid.New().String()
err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId) err := redisRpRep.SetChatId(ctx.FromID, waifuId, chatId)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
return return
@@ -323,7 +345,7 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
counter = 0 counter = 0
m.Edit("Сжатие завершено") m.Edit("Сжатие завершено")
} }
err = red.RPSetCounter(db, ctx.FromID, waifuId, counter) err = redisRpRep.SetCounter(ctx.FromID, waifuId, counter)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
} }