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.Command(rpInfo, "rp", "рп") rp.Payload(rpInfo, "rp.info") rp.Payload(rpWaifuList, "rp.waifu_list") rp.Payload(rpWaifuSet, "rp.waifu_set") rp.Payload(rpPresetsList, "rp.preset_list") rp.Payload(rpPresetSet, "rp.preset_set") rp.Payload(rpScenarioList, "rp.scenario_list") rp.Payload(chatStat, "rp.tokens") rp.Payload(newChat, "rp.new_chat") rp.Command(rpUserPromptGet, "rpuserpget") rp.Command(rpUserPromptSet, "rpuserpset") rp.Command(generate, "g", "gen", "г") rp.Payload(compress, "rp.compress_chat") rp.Payload(generalClose, "general.close") bot.AddPlugins(rp.Build()) } func rpInfo(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { rpRepRed := red.NewRPRepository(db) rpRepPsql := psql.NewRPRepository(db) waifuId := rpRepRed.GetSelectedWaifu(ctx.FromID) waifuRep := psql.NewWaifuRepository(db) waifu, err := waifuRep.GetById(waifuId) if err != nil { ctx.Error(err) } rpUser, err := rpRepPsql.GetUser(int64(ctx.FromID)) if err != nil { ctx.Error(err) } out := []string{ fmt.Sprintf("Привет, %s!", ctx.From.FirstName), fmt.Sprintf("*Выбранная вайфу*: %s", waifu.Name), fmt.Sprintf("*Выбранный пресет*: %s", laniakea.EscapeMarkdown(rpUser.Preset.Name)), fmt.Sprintf("*Твое описание персонажа*: %s", rpUser.UserPrompt), } kb := laniakea.NewInlineKeyboard(1) kb.AddCallbackButton("Статистика чата", "rp.tokens") kb.AddCallbackButton("Сменить пресет", "rp.preset_list") kb.AddCallbackButton("Сменить вайфу", "rp.waifu_list") kb.AddCallbackButton("Новый чат", "rp.scenario_list") if ctx.CallbackMsgId > 0 { ctx.EditCallback(strings.Join(out, "\n"), kb) } else { ctx.Keyboard(strings.Join(out, "\n"), kb) } } func rpWaifuList(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { waifuRep := psql.NewWaifuRepository(db) waifus, err := waifuRep.GetByUserId(ctx.FromID) if err != nil { ctx.Error(err) return } out := make([]string, len(waifus)) kb := laniakea.NewInlineKeyboard(2) for i, waifu := range waifus { out[i] = fmt.Sprintf("*%s* %d\\* из \"%s\"", waifu.Name, waifu.Rarity, waifu.Fandom) kb.AddCallbackButton(waifu.Name, "rp.waifu_set", waifu.ID) } kb.AddLine() kb.AddCallbackButton("На главную", "rp.info") ctx.EditCallback(strings.Join(out, "\n"), kb) } func rpWaifuSet(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { waifuId, err := strconv.Atoi(ctx.Args[0]) if err != nil { ctx.Error(err) return } rpRepRed := red.NewRPRepository(db) err = rpRepRed.SetSelectedWaifu(ctx.FromID, waifuId) if err != nil { ctx.Error(err) return } //rpRepPsql := psql.NewRPRepository(db) waifuRep := psql.NewWaifuRepository(db) waifu, err := waifuRep.GetById(waifuId) if err != nil { ctx.Error(err) } kb := laniakea.NewInlineKeyboard(1).AddCallbackButton("На главную", "rp.info") ctx.EditCallbackf("Была выбрана вайфу %s", kb, waifu.Name) } func rpPresetsList(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { rep := psql.NewRPRepository(db) presets, err := rep.GetAllPresets() if err != nil { ctx.Error(err) return } out := make([]string, len(presets)) kb := laniakea.NewInlineKeyboard(2) for i, preset := range presets { out[i] = fmt.Sprintf( "*%s*\n%s", laniakea.EscapeMarkdown(preset.Name), preset.Description, ) kb.AddCallbackButton(preset.Name, "rp.preset_set", preset.ID) } kb.AddLine() kb.AddCallbackButton("Назад", "rp.info") ctx.EditCallback(strings.Join(out, "\n"), kb) } 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) 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 } kb := laniakea.NewInlineKeyboard(1).AddCallbackButton("На главную", "rp.info") ctx.EditCallbackf("Был выбран пресет %s", kb, preset.Name) } func rpScenarioList(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { rep := psql.NewRPRepository(db) scenarios, err := rep.GetAllScenarios() if err != nil { ctx.Error(err) return } out := make([]string, len(scenarios)) kb := laniakea.NewInlineKeyboard(2) for i, scenario := range scenarios { out[i] = fmt.Sprintf("%d) *%s*\n%s\n", scenario.ID, scenario.Name, scenario.Description) kb.AddCallbackButton(scenario.Name, "rp.new_chat", scenario.ID) } kb.AddCallbackButton("Без сценария", "rp.new_chat", 0) kb.AddLine() kb.AddCallbackButton("На главную", "rp.info") ctx.EditCallback("Выбери сценарий\n"+strings.Join(out, "\n"), kb) } 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) kb := laniakea.NewInlineKeyboard(1).AddCallbackButton("На главную", "rp.info") out := []string{ "Статистика чата", fmt.Sprintf("ID: `%s`", chatId), fmt.Sprintf("Кол-во сообщений: %d", messageCount), fmt.Sprintf("Кол-во токенов: %d", tokens), } ctx.EditCallback(strings.Join(out, "\n"), kb) } 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 } if scenarioId > 0 { rep := psql.NewRPRepository(db) 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 } err = redisRpRep.SetChatTokens(ctx.FromID, waifuId, 0) if err != nil { ctx.Error(err) return } kb := laniakea.NewInlineKeyboard(2) kb.AddCallbackButton("На главную", "rp.info").AddCallbackButton("Закрыть", "general.close") ctx.EditCallback("Был создан новый чат", kb) } func rpUserPromptGet(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { rep := psql.NewRPRepository(db) 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) 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 } } waifuRep := psql.NewWaifuRepository(db) waifu, err := waifuRep.GetById(waifuId) if err != nil { ctx.Error(err) return } rpRep := psql.NewRPRepository(db) 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.From.FirstName), fmt.Sprintf("Вот краткое описание твоего персонажа: %s.", waifu.RpPrompt), redisRpRep.GetChatPrompt(ctx.FromID, waifuId), userPrompt, ), } afterHistory := ai.Message{ Role: "system", Content: ai.FormatPrompt(preset.PostHistory, waifu.Name, ctx.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(1).AddCallbackButton("Сжать чать", "rp.compress_chat") ctx.Keyboard(laniakea.EscapeMarkdown(agentAnswer.Content), kb) } func compress(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) { m := ctx.Answer("Запущено сжатие чата.") 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 } } 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, }) } api := ai.NewOpenAIAPI(ai.GPTBaseUrl, "", "deepseek-ai/deepseek-v3.1-terminus") 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 } m.Edit("Сжатие завершено") } func generalClose(ctx *laniakea.MsgContext, _ *laniakea.DatabaseContext) { ctx.CallbackDelete() }