380 lines
9.6 KiB
Go
380 lines
9.6 KiB
Go
package plugins
|
||
|
||
import (
|
||
"database/sql"
|
||
"errors"
|
||
"fmt"
|
||
"kurumibot/database/mdb"
|
||
"kurumibot/database/psql"
|
||
"kurumibot/database/red"
|
||
"kurumibot/laniakea"
|
||
"kurumibot/utils/ai"
|
||
"strconv"
|
||
"strings"
|
||
|
||
"github.com/google/uuid"
|
||
)
|
||
|
||
func RegisterRP(bot *laniakea.Bot) {
|
||
rp := laniakea.NewPlugin("RP")
|
||
rp = rp.Command(selectWaifu, "rpwaifu", "рпвайфу")
|
||
rp = rp.Payload(selectWaifu, "rp.selwaifu")
|
||
rp = rp.Command(rpPresetsList, "rpplist")
|
||
rp = rp.Command(rpPresetSet, "rppset")
|
||
rp = rp.Command(rpScenarioList, "rpscenlist")
|
||
rp = rp.Command(newChat, "newchat")
|
||
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())
|
||
}
|
||
|
||
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
|
||
}
|
||
rpRep := red.NewRPRepository(db)
|
||
err = rpRep.SetSelectedWaifu(ctx.FromID, waifuId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
ctx.Answer(fmt.Sprintf("Была выбрана вайфу %d", waifuId))
|
||
}
|
||
|
||
func rpPresetsList(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
rep := psql.NewRPRepository(db.PostgresSQL)
|
||
presets, err := rep.GetAllPresets()
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
out := make([]string, len(presets))
|
||
for i, preset := range presets {
|
||
out[i] = fmt.Sprintf(
|
||
"%s) *%s*\n%s",
|
||
preset.ID, laniakea.EscapeMarkdown(preset.Name), preset.Description,
|
||
)
|
||
}
|
||
ctx.Answer(strings.Join(out, "\n"))
|
||
}
|
||
func rpPresetSet(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
if len(ctx.Args) == 0 || ctx.Args[0] == "" {
|
||
return
|
||
}
|
||
presetId := ctx.Args[0]
|
||
rep := psql.NewRPRepository(db.PostgresSQL)
|
||
user, err := rep.GetOrCreateUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
preset, err := rep.UpdateUserPreset(user, presetId)
|
||
if err != nil {
|
||
if errors.Is(err, sql.ErrNoRows) {
|
||
ctx.Answer("Данный пресет не найден")
|
||
} else {
|
||
ctx.Error(err)
|
||
}
|
||
return
|
||
}
|
||
ctx.Answer(fmt.Sprintf("Был выбран пресет %s", preset.Name))
|
||
}
|
||
|
||
func rpScenarioList(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
rep := psql.NewRPRepository(db.PostgresSQL)
|
||
scenarios, err := rep.GetAllScenarios()
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
out := make([]string, len(scenarios))
|
||
for i, scenario := range scenarios {
|
||
out[i] = fmt.Sprintf("%d) *%s*", scenario.ID, scenario.Name)
|
||
}
|
||
ctx.Answer(strings.Join(out, "\n"))
|
||
}
|
||
func newChat(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
redisRpRep := red.NewRPRepository(db)
|
||
waifuId := redisRpRep.GetSelectedWaifu(ctx.FromID)
|
||
if waifuId == 0 {
|
||
ctx.Answer("Не выбрана вайфу")
|
||
return
|
||
}
|
||
chatId := uuid.New()
|
||
err := redisRpRep.SetChatId(ctx.FromID, waifuId, chatId.String())
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
scenarioId, err := strconv.Atoi(ctx.Args[0])
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
rep := psql.NewRPRepository(db.PostgresSQL)
|
||
scenario, err := rep.GetScenario(scenarioId)
|
||
if err != nil {
|
||
if errors.Is(err, sql.ErrNoRows) {
|
||
ctx.Answerf("Сценарий %d не найден", scenarioId)
|
||
} else {
|
||
ctx.Error(err)
|
||
}
|
||
return
|
||
}
|
||
err = redisRpRep.SetChatPrompt(ctx.FromID, waifuId, scenario.Prompt)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
err = redisRpRep.SetCounter(ctx.FromID, waifuId, 0)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
ctx.Answer("Был создан новый чат")
|
||
}
|
||
|
||
func rpUserPromptGet(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
rep := psql.NewRPRepository(db.PostgresSQL)
|
||
user, err := rep.GetOrCreateUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
if user.UserPrompt == "" {
|
||
ctx.Answer("У тебя нет описания")
|
||
return
|
||
}
|
||
ctx.Answerf("Вот твое РП описание пользователя\n%s", user.UserPrompt)
|
||
}
|
||
func rpUserPromptSet(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
if len(ctx.Args) == 0 || ctx.Args[0] == "" {
|
||
return
|
||
}
|
||
prompt := strings.Join(ctx.Args[1:], " ")
|
||
rep := psql.NewRPRepository(db.PostgresSQL)
|
||
user, err := rep.GetOrCreateUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
user.UserPrompt = prompt
|
||
err = rep.UpdateUser(user)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
ctx.Answer("Описание пользователя было обновлено.")
|
||
}
|
||
|
||
func generate(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
|
||
}
|
||
}
|
||
waifu, err := psql.GetWaifuById(waifuId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
rpRep := psql.NewRPRepository(db.PostgresSQL)
|
||
rpUser, err := rpRep.GetUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
preset, err := rpRep.GetUserPreset(rpUser)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
userPrompt := ""
|
||
if rpUser.UserPrompt != "" {
|
||
userPrompt = fmt.Sprintf("Вот описание моего персонажа %s ", rpUser.UserPrompt)
|
||
}
|
||
beforeHistory := ai.Message{
|
||
Role: "system",
|
||
Content: fmt.Sprintf(
|
||
"%s %s %s %s",
|
||
ai.FormatPrompt(preset.PreHistory, waifu.Name, ctx.Msg.From.FirstName),
|
||
fmt.Sprintf("Вот краткое описание твоего персонажа: %s", waifu.RpPrompt),
|
||
userPrompt,
|
||
redisRpRep.GetChatPrompt(ctx.FromID, waifuId),
|
||
),
|
||
}
|
||
afterHistory := ai.Message{
|
||
Role: "system",
|
||
Content: ai.FormatPrompt(preset.PostHistory, waifu.Name, ctx.Msg.From.FirstName),
|
||
}
|
||
|
||
history, err := mdb.GetChatHistory(db, chatId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
messages := []ai.Message{beforeHistory}
|
||
for _, m := range history {
|
||
messages = append(messages, ai.Message{
|
||
Role: m.Role,
|
||
Content: m.Message,
|
||
})
|
||
}
|
||
|
||
userMessage := strings.Join(ctx.Args, " ")
|
||
messages = append(messages, afterHistory, ai.Message{
|
||
Role: "user",
|
||
Content: userMessage,
|
||
})
|
||
err = mdb.UpdateChatHistory(db, chatId, "user", userMessage)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
m := ctx.Answer("Генерация запущена...")
|
||
api := ai.NewOpenAIAPI(ai.GPTBaseUrl, "", "deepseek-ai/deepseek-v3.1-terminus")
|
||
res, err := api.CreateCompletion(ai.CreateCompletionReq{
|
||
Messages: messages,
|
||
Verbosity: "low",
|
||
Temperature: 1.0,
|
||
})
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
if len(res.Choices) == 0 {
|
||
m.Edit("Не удалось сгенерировать ответ. Попробуйте снова позже")
|
||
return
|
||
}
|
||
|
||
agentAnswer := res.Choices[0].Message
|
||
err = mdb.UpdateChatHistory(db, chatId, agentAnswer.Role, agentAnswer.Content)
|
||
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)
|
||
}
|
||
|
||
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 {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
messages = make([]ai.Message, 0)
|
||
for _, m := range history {
|
||
messages = append(messages, ai.Message{
|
||
Role: m.Role,
|
||
Content: m.Message,
|
||
})
|
||
}
|
||
res, err = api.CompressChat(messages)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
if len(res.Choices) == 0 {
|
||
m.Edit("Не удалось сжать диалог")
|
||
return
|
||
}
|
||
|
||
compressedHistory := res.Choices[0].Message.Content
|
||
|
||
chatId = uuid.New().String()
|
||
err := redisRpRep.SetChatId(ctx.FromID, waifuId, chatId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
err = mdb.UpdateChatHistory(db, chatId, "assistant", compressedHistory)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
err = mdb.UpdateChatHistory(db, chatId, agentAnswer.Role, agentAnswer.Content)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
counter = 0
|
||
m.Edit("Сжатие завершено")
|
||
}
|
||
err = redisRpRep.SetCounter(ctx.FromID, waifuId, counter)
|
||
if err != nil {
|
||
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)
|
||
}
|