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(rpInfo, "rp", "рп") rp = rp.Payload(rpInfo, "rp.info") rp = rp.Payload(rpWaifuList, "rp.waifu_list") rp = rp.Payload(rpWaifuSet, "rp.waifu_set") rp = rp.Payload(rpPresetsList, "rp.preset_list") rp = rp.Payload(rpPresetSet, "rp.preset_set") rp = rp.Payload(rpScenarioList, "rp.scenario_list") rp = rp.Payload(rpScenarioSet, "rp.scenario_set") rp = rp.Payload(chatStat, "rp.tokens") rp = rp.Payload(newChat, "rp.new_chat") rp = rp.Command(rpUserPromptGet, "rpuserpget") rp = rp.Command(rpUserPromptSet, "rpuserpset") rp = rp.Command(generate, "g", "gen", "г") rp = rp.Payload(compress, "rp.compress_chat") 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) waifu, err := psql.GetWaifuById(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", "_TODO_"), fmt.Sprintf("Выбранный сценарий: %s", "_TODO_"), fmt.Sprintf("Твое описание персонажа: %s", rpUser.UserPrompt), } kb := laniakea.NewInlineKeyboard(1) kb.AddCallbackButton("Статистика чата", "rp.tokens") kb.AddCallbackButton("Сменить пресет", "rp.preset_list") kb.AddCallbackButton("Выбрать сценарий", "rp.scenario_list") kb.AddCallbackButton("Сменить вайфу", "rp.waifu_list") kb.AddCallbackButton("Новый чат", "rp.new_chat") 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) { waifus, err := psql.GetUserWaifus(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) waifu, err := psql.GetWaifuById(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.scenario_set", scenario.ID) } kb.AddLine() kb.AddCallbackButton("На главную", "rp.info") ctx.EditCallback(strings.Join(out, "\n"), kb) } func rpScenarioSet(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {} 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 } 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 } ctx.Answer("Был создан новый чат") } 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 } } waifu, err := psql.GetWaifuById(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.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(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("Сжатие завершено") }