1050 lines
30 KiB
Go
1050 lines
30 KiB
Go
package plugins
|
||
|
||
import (
|
||
"database/sql"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
"ymgb/database"
|
||
"ymgb/database/mdb"
|
||
"ymgb/database/psql"
|
||
"ymgb/database/red"
|
||
"ymgb/openai"
|
||
"ymgb/utils"
|
||
"ymgb/utils/ai"
|
||
|
||
"git.nix13.pw/scuroneko/extypes"
|
||
"git.nix13.pw/scuroneko/laniakea"
|
||
"git.nix13.pw/scuroneko/laniakea/tgapi"
|
||
"github.com/google/uuid"
|
||
)
|
||
|
||
func RegisterRP() *laniakea.Plugin[database.Context] {
|
||
rp := laniakea.NewPlugin[database.Context]("RP")
|
||
rp.NewCommand(rpUserPromptSet, "rpuserpset")
|
||
rp.NewCommand(rpInfo, "rp").SetDescription("РП профиль пользователя")
|
||
rp.NewCommand(rpInfo, "рп")
|
||
|
||
rp.NewPayload(rpInfo, "rp.info")
|
||
rp.NewPayload(rpWaifuList, "rp.waifu_list")
|
||
rp.NewPayload(rpWaifuSet, "rp.waifu_set")
|
||
rp.NewPayload(rpPresetsList, "rp.preset_list")
|
||
rp.NewPayload(rpPresetSet, "rp.preset_set")
|
||
rp.NewPayload(rpModelList, "rp.model_list")
|
||
rp.NewPayload(rpModelSet, "rp.model_set")
|
||
rp.NewPayload(rpScenarioList, "rp.scenario_list")
|
||
rp.NewPayload(rpSettingList, "rp.setting_list")
|
||
rp.NewPayload(chatStat, "rp.chat_stat")
|
||
rp.NewPayload(newChatStage1, "rp.new_chat_s1")
|
||
rp.NewPayload(newChatStage2, "rp.new_chat_s2")
|
||
rp.NewPayload(newChat, "rp.new_chat")
|
||
rp.NewCommand(
|
||
generate, "g",
|
||
*laniakea.NewCommandArg("prompt").SetValueType(laniakea.CommandValueAnyType),
|
||
).SetDescription("Генерация РП")
|
||
rp.NewCommand(generate, "gen").SkipCommandAutoGen()
|
||
rp.NewCommand(generate, "г").SkipCommandAutoGen()
|
||
rp.NewPayload(compress, "rp.compress_chat")
|
||
rp.NewPayload(regenerateResponse, "rp.regenerate")
|
||
|
||
rp.NewPayload(compressSettingStage1, "rp.compress_setting_s1")
|
||
rp.NewPayload(compressSettingStage2, "rp.compress_setting_s2")
|
||
rp.NewPayload(compressSetting, "rp.compress_setting")
|
||
|
||
return rp
|
||
}
|
||
|
||
func rpInfo(ctx *laniakea.MsgContext, db *database.Context) {
|
||
userRep := psql.NewUserRepository(db)
|
||
_, err := userRep.GetOrCreate(ctx.FromID, ctx.From.FirstName)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
waifuRep := psql.NewWaifuRepository(db)
|
||
count, err := waifuRep.GetCountByUserId(ctx.FromID)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
if count == 0 {
|
||
kb := laniakea.NewInlineKeyboardJson(1).AddCallbackButton("Закрыть", "general.close")
|
||
ctx.Keyboard("Для доступа к этой команде нужно иметь хотя бы одну вайфу", kb)
|
||
return
|
||
}
|
||
|
||
rpRepRed := red.NewRPRepository(db)
|
||
rpRepPsql := psql.NewRPRepository(db)
|
||
rpUser, err := rpRepPsql.GetOrCreateUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
var waifu *psql.Waifu
|
||
waifuId := rpRepRed.GetSelectedWaifu(ctx.FromID)
|
||
if waifuId == 0 {
|
||
waifus, err := waifuRep.GetByUserId(ctx.FromID)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
waifu = waifus[0]
|
||
err = rpRepRed.SetSelectedWaifu(ctx.FromID, waifu.ID)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
} else {
|
||
waifu, err = waifuRep.GetById(waifuId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
}
|
||
compressText := "сообщений"
|
||
if rpUser.CompressMethod == "tokens" {
|
||
compressText = "токенов"
|
||
}
|
||
out := []string{
|
||
fmt.Sprintf("Привет, _%s_!", ctx.From.FirstName),
|
||
fmt.Sprintf("*Выбранная вайфу*: %s", waifu.Name),
|
||
fmt.Sprintf("*Выбранный пресет*: %s", laniakea.EscapeMarkdownV2(rpUser.Preset.Name)),
|
||
fmt.Sprintf("*Выбранная модель*: %s", rpUser.Model.Name),
|
||
fmt.Sprintf("*Использовано токенов*: %d", rpUser.UsedTokens),
|
||
fmt.Sprintf("*Настройки сжатия*: %d %s", rpUser.CompressLimit, compressText),
|
||
fmt.Sprintf("*Твоё описание персонажа*: `%s`", rpUser.UserPrompt),
|
||
"",
|
||
"Что бы установить описание персонажа, используй `/rpuserpset \"описание персонажа\"` без кавычек.",
|
||
}
|
||
|
||
kb := laniakea.NewInlineKeyboardJson(2)
|
||
kb.AddCallbackButton("Статистика чата", "rp.chat_stat")
|
||
kb.AddCallbackButton("Сменить вайфу", "rp.waifu_list")
|
||
kb.AddCallbackButton("Сменить пресет", "rp.preset_list")
|
||
kb.AddCallbackButton("Сменить модель", "rp.model_list")
|
||
kb.AddCallbackButton("Список сценариев", "rp.scenario_list")
|
||
kb.AddCallbackButton("Список сеттингов", "rp.setting_list")
|
||
kb.AddCallbackButton("Настройки сжатия", "rp.compress_setting_s1")
|
||
kb.AddCallbackButton("Новый чат", "rp.new_chat_s1")
|
||
kb.AddLine()
|
||
kb.AddCallbackButtonStyle("Закрыть", laniakea.ButtonStyleDanger, "general.close")
|
||
|
||
if ctx.CallbackMsgId > 0 {
|
||
ctx.EditCallbackMarkdown(laniakea.EscapePunctuation(strings.Join(out, "\n")), kb)
|
||
ctx.AnswerCbQuery()
|
||
} else {
|
||
ctx.KeyboardMarkdown(laniakea.EscapePunctuation(strings.Join(out, "\n")), kb)
|
||
}
|
||
}
|
||
|
||
func rpWaifuList(ctx *laniakea.MsgContext, db *database.Context) {
|
||
waifuRep := psql.NewWaifuRepository(db)
|
||
var waifus extypes.Slice[*psql.Waifu]
|
||
|
||
userRep := psql.NewUserRepository(db)
|
||
user, err := userRep.GetOrCreate(ctx.FromID, ctx.From.FirstName)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
if user.Group.IsPremium {
|
||
waifus, err = waifuRep.GetAll()
|
||
} else {
|
||
waifus, err = waifuRep.GetByUserId(ctx.FromID)
|
||
}
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
waifus = waifus.Filter(func(w *psql.Waifu) bool {
|
||
return len(w.RpPrompt) > 0
|
||
})
|
||
|
||
out := make([]string, len(waifus))
|
||
kb := laniakea.NewInlineKeyboardJson(2)
|
||
for i, waifu := range waifus {
|
||
owner := "нет"
|
||
if waifu.OwnerID.Valid && waifu.Owner != nil {
|
||
owner = waifu.Owner.Name
|
||
}
|
||
out[i] = fmt.Sprintf(
|
||
"*%s* %d☆ из \"%s\" \\(владелец: %s\\)",
|
||
waifu.Name, waifu.Rarity, waifu.Fandom, owner,
|
||
)
|
||
kb.AddCallbackButton(waifu.Name, "rp.waifu_set", waifu.ID)
|
||
}
|
||
kb.AddLine()
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStyleDanger, "rp.info")
|
||
ctx.EditCallbackMarkdown(strings.Join(out, "\n"), kb)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
func rpWaifuSet(ctx *laniakea.MsgContext, db *database.Context) {
|
||
waifuId, err := strconv.Atoi(ctx.Args[0])
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
rpRepRed := red.NewRPRepository(db)
|
||
err = rpRepRed.SetSelectedWaifu(ctx.FromID, int64(waifuId))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
//rpRepPsql := psql.NewRPRepository(db)
|
||
waifuRep := psql.NewWaifuRepository(db)
|
||
waifu, err := waifuRep.GetById(int64(waifuId))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
|
||
kb := laniakea.NewInlineKeyboardJson(1)
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStyleSuccess, "rp.info")
|
||
ctx.EditCallbackfMarkdown("Была выбрана вайфу *%s*", kb, waifu.Name)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
|
||
func rpPresetsList(ctx *laniakea.MsgContext, db *database.Context) {
|
||
rep := psql.NewRPRepository(db)
|
||
presets, err := rep.GetAllPresets()
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
out := make([]string, len(presets))
|
||
kb := laniakea.NewInlineKeyboardJson(2)
|
||
for i, preset := range presets {
|
||
out[i] = fmt.Sprintf(
|
||
"*%s* \\- %s",
|
||
laniakea.EscapeMarkdownV2(preset.Name), laniakea.EscapePunctuation(preset.Description),
|
||
)
|
||
kb.AddCallbackButton(preset.Name, "rp.preset_set", preset.ID)
|
||
}
|
||
kb.AddLine()
|
||
kb.AddCallbackButtonStyle("Назад", laniakea.ButtonStyleDanger, "rp.info")
|
||
ctx.EditCallbackMarkdown(strings.Join(out, "\n"), kb)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
func rpPresetSet(ctx *laniakea.MsgContext, db *database.Context) {
|
||
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.NewInlineKeyboardJson(1)
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStylePrimary, "rp.info")
|
||
ctx.EditCallbackf("Был выбран пресет %s", kb, preset.Name)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
|
||
func rpModelList(ctx *laniakea.MsgContext, db *database.Context) {
|
||
rep := psql.NewAIRepository(db)
|
||
models, err := rep.GetAllModels()
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
out := make([]string, len(models))
|
||
kb := laniakea.NewInlineKeyboardJson(2)
|
||
for i, model := range models {
|
||
out[i] = fmt.Sprintf("*%s* \\- размер контекста _%dK_", laniakea.EscapeMarkdownV2(model.Name), model.ContextSize)
|
||
kb.AddCallbackButton(model.Name, "rp.model_set", model.ID)
|
||
}
|
||
kb.AddLine()
|
||
kb.AddCallbackButtonStyle("Назад", laniakea.ButtonStyleDanger, "rp.info")
|
||
ctx.EditCallbackMarkdown(strings.Join(out, "\n"), kb)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
func rpModelSet(ctx *laniakea.MsgContext, db *database.Context) {
|
||
rep := psql.NewRPRepository(db)
|
||
user, err := rep.GetOrCreateUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
aiRep := psql.NewAIRepository(db)
|
||
model, err := aiRep.GetModel(ctx.Args[0])
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
user.SelectedModel = model.ID
|
||
err = rep.UpdateUser(user)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
kb := laniakea.NewInlineKeyboardJson(1)
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStyleSuccess, "rp.info")
|
||
ctx.EditCallback(fmt.Sprintf("Была установлена модель %s", model.Name), kb)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
|
||
func rpScenarioList(ctx *laniakea.MsgContext, db *database.Context) {
|
||
rep := psql.NewRPRepository(db)
|
||
scenarios, err := rep.GetAllScenarios()
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
out := make([]string, len(scenarios))
|
||
kb := laniakea.NewInlineKeyboardJson(1)
|
||
for i, scenario := range scenarios {
|
||
out[i] = fmt.Sprintf("*%s* \\- %s\n", scenario.Name, laniakea.EscapePunctuation(scenario.Description))
|
||
}
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStyleDanger, "rp.info")
|
||
ctx.EditCallbackMarkdown(strings.Join(out, "\n"), kb)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
func rpSettingList(ctx *laniakea.MsgContext, db *database.Context) {
|
||
rep := psql.NewRPRepository(db)
|
||
settings, err := rep.GetAllSettings()
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
out := make([]string, len(settings))
|
||
for i, setting := range settings {
|
||
out[i] = fmt.Sprintf("*%s* \\- %s\n", setting.Name, laniakea.EscapePunctuation(setting.Description))
|
||
}
|
||
kb := laniakea.NewInlineKeyboardJson(1)
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStyleSuccess, "rp.info")
|
||
ctx.EditCallbackMarkdown(strings.Join(out, "\n"), kb)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
|
||
func chatStat(ctx *laniakea.MsgContext, db *database.Context) {
|
||
redisRpRep := red.NewRPRepository(db)
|
||
chat, err := redisRpRep.GetChat(ctx.FromID)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
messageCount, err := mdb.GetRPChatHistorySize(db, chat.ID.String())
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
kb := laniakea.NewInlineKeyboardJson(1)
|
||
kb.AddCallbackButtonStyle("Сжать чат", laniakea.ButtonStyleSuccess, "rp.compress_chat")
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStyleDanger, "rp.info")
|
||
out := []string{
|
||
"Статистика чата",
|
||
fmt.Sprintf("*ID*: `%s`", chat.ID.String()),
|
||
fmt.Sprintf("*Кол-во сообщений*: %d", messageCount),
|
||
fmt.Sprintf("*Кол-во токенов*: %d", chat.ChatTokens),
|
||
}
|
||
|
||
if len(chat.Scenarios) > 0 && chat.Scenarios[0].ID > 0 {
|
||
scenarioNames := make([]string, len(chat.Scenarios))
|
||
for i, scenario := range chat.Scenarios {
|
||
if scenario.ID == 0 {
|
||
continue
|
||
}
|
||
scenarioNames[i] = fmt.Sprintf("%s \\(ID: %d\\)", scenario.Name, scenario.ID)
|
||
}
|
||
out = append(out, fmt.Sprintf("*Выбранные сценарии*: %s", strings.Join(scenarioNames, ", ")))
|
||
}
|
||
|
||
if chat.Setting != nil {
|
||
out = append(out, fmt.Sprintf("*Выбранный сеттинг*: %s (ID: %d)", chat.Setting.Name, chat.Setting.ID))
|
||
}
|
||
ctx.EditCallbackMarkdown(laniakea.EscapePunctuation(strings.Join(out, "\n")), kb)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
|
||
func newChatStage1(ctx *laniakea.MsgContext, db *database.Context) {
|
||
// Выбор сеттинга
|
||
rep := psql.NewRPRepository(db)
|
||
settings, err := rep.GetAllSettings()
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
out := []string{
|
||
"Выбери сеттинг для чата",
|
||
}
|
||
kb := laniakea.NewInlineKeyboardJson(2)
|
||
for _, setting := range settings {
|
||
out = append(out, fmt.Sprintf("*%s* \\- %s", setting.Name, laniakea.EscapePunctuation(setting.Description)))
|
||
kb.AddCallbackButton(setting.Name, "rp.new_chat_s2", setting.ID)
|
||
}
|
||
kb.AddCallbackButton("Без сеттинга", "rp.new_chat_s2", 0)
|
||
kb.AddLine()
|
||
kb.AddCallbackButtonStyle("Назад", laniakea.ButtonStyleDanger, "rp.info")
|
||
ctx.EditCallbackMarkdown(strings.Join(out, "\n"), kb)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
func newChatStage2(ctx *laniakea.MsgContext, db *database.Context) {
|
||
// Выбор сценария
|
||
if len(ctx.Args) == 0 {
|
||
ctx.Error(fmt.Errorf("zero args len"))
|
||
return
|
||
}
|
||
|
||
rep := psql.NewRPRepository(db)
|
||
scenarios, err := rep.GetAllScenarios()
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
settingId, err := strconv.Atoi(ctx.Args[0])
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
selectedScenariosIds := make(extypes.Slice[int], 0)
|
||
if len(ctx.Args) > 1 {
|
||
selectedScenariosIds = utils.Map(ctx.Args[1:], utils.StringToInt)
|
||
}
|
||
|
||
var out []string
|
||
kb := laniakea.NewInlineKeyboardJson(2)
|
||
var scenariosIds extypes.Slice[int]
|
||
for _, scenario := range scenarios {
|
||
isSelected := selectedScenariosIds.Index(scenario.ID) >= 0
|
||
prefix := ""
|
||
style := laniakea.ButtonStylePrimary
|
||
if isSelected {
|
||
prefix = "✅"
|
||
style = laniakea.ButtonStyleSuccess
|
||
}
|
||
out = append(out, fmt.Sprintf("%s*%s* \\- %s", prefix, scenario.Name, laniakea.EscapePunctuation(scenario.Description)))
|
||
if isSelected {
|
||
scenariosIds = selectedScenariosIds.Remove(scenario.ID)
|
||
} else {
|
||
scenariosIds = selectedScenariosIds.Push(scenario.ID)
|
||
}
|
||
|
||
kb.AddCallbackButtonStyle(
|
||
fmt.Sprintf("%s%s", prefix, scenario.Name), style,
|
||
"rp.new_chat_s2", utils.AppendToInt(settingId, scenariosIds)...,
|
||
)
|
||
}
|
||
//if len(scenariosIds) == 0 {
|
||
// scenariosIds = append(scenariosIds, 0)
|
||
//}
|
||
kb.AddCallbackButtonStyle("Создать", laniakea.ButtonStyleSuccess, "rp.new_chat", utils.AppendToInt(settingId, selectedScenariosIds)...)
|
||
kb.AddLine()
|
||
kb.AddCallbackButtonStyle("Назад", laniakea.ButtonStyleDanger, "rp.new_chat_s1")
|
||
ctx.EditCallbackMarkdown(strings.Join(out, "\n"), kb)
|
||
|
||
//setting, err := rep.GetSetting(settingId)
|
||
//if err != nil {
|
||
// ctx.Error(err)
|
||
//}
|
||
//ctx.AnswerCbQueryText(fmt.Sprintf("Ты выбрал сеттинг %s", setting.name))
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
func newChat(ctx *laniakea.MsgContext, db *database.Context) {
|
||
redisRpRep := red.NewRPRepository(db)
|
||
psqlRpRep := psql.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
|
||
}
|
||
|
||
chat, err := redisRpRep.GetChat(ctx.FromID)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
chatPrompt := ""
|
||
settingId, err := strconv.Atoi(ctx.Args[0])
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
if settingId > 0 {
|
||
setting, err := psqlRpRep.GetSetting(settingId)
|
||
if err != nil {
|
||
if errors.Is(err, sql.ErrNoRows) {
|
||
ctx.Answerf("Сеттинг №%d не найден", settingId)
|
||
} else {
|
||
ctx.Error(err)
|
||
}
|
||
return
|
||
}
|
||
chat.Setting = &setting
|
||
chat.SettingID = settingId
|
||
chatPrompt = "Вот краткое описание мира(сеттинг): " + setting.Prompt + "."
|
||
}
|
||
|
||
scenariosIds := utils.Map(ctx.Args[1:], utils.StringToInt)
|
||
chat.Scenarios = make([]psql.RPScenario, 0)
|
||
if len(scenariosIds) > 0 && scenariosIds[0] > 0 {
|
||
rep := psql.NewRPRepository(db)
|
||
var scenariosPrompt []string
|
||
|
||
for _, scenarioId := range scenariosIds {
|
||
scenario, err := rep.GetScenario(scenarioId)
|
||
if err != nil {
|
||
if errors.Is(err, sql.ErrNoRows) {
|
||
ctx.Answerf("Сценарий №%d не найден", scenarioId)
|
||
} else {
|
||
ctx.Error(err)
|
||
}
|
||
return
|
||
}
|
||
chat.Scenarios = append(chat.Scenarios, scenario)
|
||
scenariosPrompt = append(scenariosPrompt, scenario.Prompt)
|
||
}
|
||
|
||
chatPrompt += "Вот дополнительная информация: " + strings.Join(scenariosPrompt, ".")
|
||
}
|
||
chat.Prompt = chatPrompt
|
||
chat.Counter = 0
|
||
chat.ChatTokens = 0
|
||
err = redisRpRep.SaveChat(chat)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
kb := laniakea.NewInlineKeyboardJson(2)
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStyleSuccess, "rp.info")
|
||
kb.AddCallbackButtonStyle("Закрыть", laniakea.ButtonStyleDanger, "general.close")
|
||
ctx.EditCallbackMarkdown("Был создан новый чат\\. Для общения используй `/г промпт`", kb)
|
||
ctx.AnswerCbQuery()
|
||
}
|
||
|
||
func rpUserPromptSet(ctx *laniakea.MsgContext, db *database.Context) {
|
||
if len(ctx.Args) == 0 || ctx.Args[0] == "" {
|
||
return
|
||
}
|
||
prompt := strings.Join(ctx.Args, " ")
|
||
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 _getChatHistory(ctx *laniakea.MsgContext, db *database.Context) ([]openai.Message, error) {
|
||
redRep := red.NewRPRepository(db)
|
||
psqlRep := psql.NewRPRepository(db)
|
||
waifuRep := psql.NewWaifuRepository(db)
|
||
|
||
messages := make([]openai.Message, 0)
|
||
waifuId := redRep.GetSelectedWaifu(ctx.FromID)
|
||
chatId, err := redRep.GetOrCreateChatId(ctx.FromID, waifuId)
|
||
if err != nil {
|
||
return messages, err
|
||
}
|
||
|
||
waifu, err := waifuRep.GetById(waifuId)
|
||
if err != nil {
|
||
return messages, err
|
||
}
|
||
user, err := psqlRep.GetUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
return messages, err
|
||
}
|
||
preset, err := psqlRep.GetPreset(user.SelectedPreset)
|
||
if err != nil {
|
||
return messages, err
|
||
}
|
||
|
||
userPrompt := ""
|
||
if user.UserPrompt != "" {
|
||
userPrompt = fmt.Sprintf("Вот описание моего персонажа %s.", user.UserPrompt)
|
||
}
|
||
//waifuPrompt, err := psqlRep.GetWaifuPrompts(waifuId)
|
||
//if err != nil {
|
||
// return messages, err
|
||
//}
|
||
beforeHistory := openai.Message{
|
||
Role: "system",
|
||
Content: fmt.Sprintf(
|
||
"%s %s %s %s",
|
||
ai.FormatPrompt(preset.PreHistory, waifu.Name, ctx.From.FirstName),
|
||
fmt.Sprintf(
|
||
"Вот краткое описание твоего персонажа: %s.",
|
||
waifu.RpPrompt,
|
||
),
|
||
redRep.GetChatPrompt(ctx.FromID, waifuId),
|
||
userPrompt,
|
||
),
|
||
}
|
||
afterHistory := openai.Message{
|
||
Role: "system",
|
||
Content: ai.FormatPrompt(preset.PostHistory, waifu.Name, ctx.From.FirstName),
|
||
}
|
||
|
||
history, err := mdb.GetRPChatHistory(db, chatId)
|
||
if err != nil {
|
||
return messages, err
|
||
}
|
||
|
||
messages = append(messages, beforeHistory)
|
||
for _, m := range history {
|
||
messages = append(messages, openai.Message{Role: m.Role, Content: m.Message})
|
||
}
|
||
messages = append(messages, afterHistory)
|
||
return messages, nil
|
||
}
|
||
|
||
func generate(ctx *laniakea.MsgContext, db *database.Context) {
|
||
redisRpRep := red.NewRPRepository(db)
|
||
rpRep := psql.NewRPRepository(db)
|
||
waifuId := redisRpRep.GetSelectedWaifu(ctx.FromID)
|
||
if waifuId == 0 {
|
||
ctx.Answer("Не выбрана вайфу")
|
||
return
|
||
}
|
||
|
||
rpUser, err := rpRep.GetOrCreateUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
chatId, err := redisRpRep.GetOrCreateChatId(ctx.FromID, waifuId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
messages, err := _getChatHistory(ctx, db)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
userMessage := strings.TrimSpace(strings.Join(ctx.Args, " "))
|
||
|
||
kb := laniakea.NewInlineKeyboardJson(1)
|
||
kb.AddCallbackButtonStyle("Отменить", laniakea.ButtonStyleDanger, "rp.cancel")
|
||
m := ctx.Keyboard("Генерация запущена...", kb)
|
||
ctx.SendAction(tgapi.ChatActionTyping)
|
||
api := openai.NewOpenAIAPI(ai.GPTBaseUrl, "", rpUser.Model.Key)
|
||
res, err := api.CreateCompletionStream(messages, userMessage, 0.5)
|
||
if err != nil {
|
||
_ = api.Close()
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
answerContent := ""
|
||
draft := ctx.NewDraft()
|
||
var mu sync.Mutex
|
||
buffer := strings.Builder{} // накапливаем фрагменты
|
||
ticker := time.NewTicker(time.Second)
|
||
done := make(chan struct{}) // сигнал завершения горутины
|
||
|
||
// Горутина для периодической отправки
|
||
go func() {
|
||
flush := func() {
|
||
mu.Lock()
|
||
msg := buffer.String()
|
||
buffer.Reset()
|
||
mu.Unlock()
|
||
if msg != "" {
|
||
if err := draft.Push(msg); err != nil {
|
||
log.Println(err)
|
||
}
|
||
}
|
||
}
|
||
for {
|
||
select {
|
||
case <-ticker.C:
|
||
// Каждую секунду отправляем накопленное
|
||
flush()
|
||
ctx.SendAction(tgapi.ChatActionTyping)
|
||
case <-done:
|
||
// Завершаем работу: отправляем остаток и выходим
|
||
flush()
|
||
ticker.Stop()
|
||
return
|
||
}
|
||
}
|
||
}()
|
||
for r, err := range res {
|
||
if err != nil {
|
||
if errors.Is(err, io.EOF) {
|
||
break
|
||
}
|
||
ctx.Error(err)
|
||
close(done)
|
||
return
|
||
}
|
||
|
||
if len(r.Choices) == 0 {
|
||
continue
|
||
}
|
||
content := r.Choices[0].Delta.Content
|
||
if content == "" {
|
||
continue
|
||
}
|
||
|
||
if m != nil {
|
||
m.Delete()
|
||
m = nil
|
||
}
|
||
answerContent += content
|
||
|
||
mu.Lock()
|
||
buffer.WriteString(content)
|
||
mu.Unlock()
|
||
}
|
||
|
||
// Завершаем горутину и отправляем последнюю порцию
|
||
close(done)
|
||
_ = api.Close()
|
||
if answerContent == "" {
|
||
if m != nil {
|
||
m.Delete()
|
||
m = nil
|
||
}
|
||
ctx.Answer("Не удалось сгенерировать ответ. Попробуйте снова")
|
||
return
|
||
}
|
||
|
||
counter := redisRpRep.GetCounter(ctx.FromID, waifuId)
|
||
err = mdb.UpdateRPChatHistory(db, chatId, "user", userMessage, counter+1)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
err = mdb.UpdateRPChatHistory(db, chatId, "assistant", answerContent, counter+2)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
|
||
rpUser.UsedTokens = rpUser.UsedTokens + int64(len(userMessage)) + int64(len(answerContent))
|
||
err = rpRep.UpdateUser(rpUser)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
|
||
tokens := redisRpRep.GetChatTokens(ctx.FromID, waifuId)
|
||
tokens += int64(len(userMessage) + len(answerContent))
|
||
err = redisRpRep.SetChatTokens(ctx.FromID, waifuId, tokens)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
err = redisRpRep.SetCounter(ctx.FromID, waifuId, counter+2)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
|
||
kb = laniakea.NewInlineKeyboardJson(1)
|
||
kb.AddCallbackButtonStyle("🔄 Перегенерировать 🔄", laniakea.ButtonStyleSuccess, "rp.regenerate", counter+2)
|
||
ctx.Keyboard(answerContent, kb)
|
||
|
||
// Auto compress
|
||
compressMethod := rpUser.CompressMethod
|
||
switch compressMethod {
|
||
case "messages":
|
||
if counter+2 >= rpUser.CompressLimit {
|
||
m = ctx.Answer("Запущено сжатие чата…")
|
||
_compress(ctx, db)
|
||
m.Delete()
|
||
}
|
||
case "tokens":
|
||
if tokens >= int64(rpUser.CompressLimit*1000) {
|
||
m = ctx.Answer("Запущено сжатие чата…")
|
||
_compress(ctx, db)
|
||
m.Delete()
|
||
}
|
||
}
|
||
}
|
||
|
||
func regenerateResponse(ctx *laniakea.MsgContext, db *database.Context) {
|
||
ctx.AnswerCbQueryText("Запущена повторная генерация…")
|
||
|
||
redRep := red.NewRPRepository(db)
|
||
waifuId := redRep.GetSelectedWaifu(ctx.FromID)
|
||
|
||
count, err := strconv.Atoi(ctx.Args[0])
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
var messages extypes.Slice[openai.Message]
|
||
messages, err = _getChatHistory(ctx, db)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
chatId, err := redRep.GetOrCreateChatId(ctx.FromID, waifuId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
history, err := mdb.GetRPChatHistory(db, chatId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
//if messages.Len() == count {
|
||
// ctx.Bot.logger().Errorln("len(messages) == count")
|
||
// return
|
||
//}
|
||
|
||
// 0(system), 1, 2, ..., N-2(user, count-3), N-1(agent, count-2), N(system, count-1)
|
||
answerToDelete := history[count-1]
|
||
userReq := history[count-2]
|
||
|
||
d := fmt.Sprintf("\\[DEBUG]\n%s\n\n%s", answerToDelete.Message, userReq.Message)
|
||
ctx.Answer(d)
|
||
|
||
psqlRep := psql.NewRPRepository(db)
|
||
rpUser, err := psqlRep.GetOrCreateUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
api := openai.NewOpenAIAPI(ai.GPTBaseUrl, "", rpUser.Model.Key)
|
||
defer func(api *openai.API) {
|
||
_ = api.Close()
|
||
}(api)
|
||
|
||
messages = messages.Pop(count - 2)
|
||
messages = messages.Pop(count - 2)
|
||
|
||
ctx.EditCallback("Генерация запущена...", nil)
|
||
res, err := api.CreateCompletion(messages, userReq.Message, 0.5)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
err = mdb.DeleteRPChatEntry(db, answerToDelete)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
kb := laniakea.NewInlineKeyboardJson(1)
|
||
kb.AddCallbackButtonStyle("🔄 Перегенерировать 🔄", laniakea.ButtonStyleSuccess, "rp.regenerate", count)
|
||
ctx.EditCallback(res.Choices[0].Message.Content, kb)
|
||
}
|
||
|
||
func compress(ctx *laniakea.MsgContext, db *database.Context) {
|
||
ctx.AnswerCbQueryText("Запущено сжатие чата…")
|
||
m := ctx.Answer("Запущено сжатие чата…")
|
||
_compress(ctx, db)
|
||
kb := laniakea.NewInlineKeyboardJson(1)
|
||
kb.AddCallbackButton("Назад", "rp.chat_stat")
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStyleDanger, "rp.info")
|
||
m.Delete()
|
||
}
|
||
|
||
func _compress(ctx *laniakea.MsgContext, db *database.Context) {
|
||
redisRpRep := red.NewRPRepository(db)
|
||
waifuId := redisRpRep.GetSelectedWaifu(ctx.FromID)
|
||
if waifuId == 0 {
|
||
ctx.Answer("Не выбрана вайфу")
|
||
return
|
||
}
|
||
chatId, err := redisRpRep.GetOrCreateChatId(ctx.FromID, waifuId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
history, err := mdb.GetRPChatHistory(db, chatId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
messages := make([]openai.Message, 0)
|
||
for _, m := range history {
|
||
messages = append(messages, openai.Message{
|
||
Role: m.Role,
|
||
Content: m.Message,
|
||
})
|
||
}
|
||
|
||
psqlRpRep := psql.NewRPRepository(db)
|
||
user, err := psqlRpRep.GetUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
api := openai.NewOpenAIAPI(ai.GPTBaseUrl, "", user.Model.Key)
|
||
defer func(api *openai.API) {
|
||
_ = api.Close()
|
||
}(api)
|
||
res, err := ai.CompressChat(api, messages)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
if len(res.Choices) == 0 {
|
||
ctx.Answer("Не удалось сжать чат")
|
||
return
|
||
}
|
||
compressedHistory := strings.TrimSpace(res.Choices[0].Message.Content)
|
||
compressedHistory = strings.ReplaceAll(compressedHistory, "*", "")
|
||
|
||
ctx.Answer(compressedHistory)
|
||
tokens := len(compressedHistory)
|
||
|
||
chatId = uuid.New().String()
|
||
err = redisRpRep.SetChatId(ctx.FromID, waifuId, chatId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
err = mdb.UpdateRPChatHistory(db, chatId, "assistant", compressedHistory, 0)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
offset := 50
|
||
if user.CompressMethod == "messages" {
|
||
offset = user.CompressLimit / 2
|
||
}
|
||
offset = utils.Min(len(history), offset)
|
||
|
||
// Copy short history from prev chat
|
||
index := 0
|
||
for _, m := range history {
|
||
if m.Role == "assistant" {
|
||
tokens += len(m.Message)
|
||
_ = mdb.UpdateRPChatHistory(db, chatId, m.Role, m.Message, index)
|
||
index++
|
||
} else {
|
||
break
|
||
}
|
||
}
|
||
|
||
for _, m := range history[len(history)-offset:] {
|
||
tokens += len(m.Message)
|
||
err = mdb.UpdateRPChatHistory(db, chatId, m.Role, m.Message, index)
|
||
index++
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
}
|
||
|
||
err = redisRpRep.SetCounter(ctx.FromID, waifuId, index)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
err = redisRpRep.SetChatTokens(ctx.FromID, waifuId, int64(tokens))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
}
|
||
|
||
var messagesMethodCount = []int{
|
||
20, 50, 100, 200, 500,
|
||
}
|
||
var tokenMethodCount = []int{
|
||
16, 32, 64, 128, 256,
|
||
}
|
||
|
||
func compressSettingStage1(ctx *laniakea.MsgContext, _ *database.Context) {
|
||
kb := laniakea.NewInlineKeyboardJson(2)
|
||
kb.AddCallbackButton("По сообщениям", "rp.compress_setting_s2", "messages")
|
||
kb.AddCallbackButton("По токенам", "rp.compress_setting_s2", "tokens")
|
||
kb.AddCallbackButton("Отключить", "rp.compress_setting", "none", 0)
|
||
kb.AddLine().AddCallbackButton("На главную", "rp.info")
|
||
|
||
out := []string{
|
||
"Выбери метод для определения сжатия",
|
||
"*По сообщениям* \\— чат будет сжиматься после достижения определенного кол-ва сообщений в нём",
|
||
"*По токенам* \\— чат будет сжиматься после достижения определенного кол-ва токенов в нём",
|
||
}
|
||
|
||
ctx.EditCallbackMarkdown(strings.Join(out, "\n"), kb)
|
||
}
|
||
func compressSettingStage2(ctx *laniakea.MsgContext, _ *database.Context) {
|
||
if len(ctx.Args) == 0 {
|
||
return
|
||
}
|
||
method := ctx.Args[0]
|
||
out := []string{
|
||
"При выборе маленьких значений, чат будет сжиматься чаще, что замедляет вывод сообщений.",
|
||
"При выборе слишком больших значений, некоторые детали могут теряться.",
|
||
}
|
||
kb := laniakea.NewInlineKeyboardJson(3)
|
||
switch method {
|
||
case "messages":
|
||
out = append(out, "Выбери количество сообщений, после которых будет происходить сжатие чата")
|
||
for _, i := range messagesMethodCount {
|
||
kb.AddCallbackButton(strconv.Itoa(i), "rp.compress_setting", method, i)
|
||
}
|
||
case "tokens":
|
||
out = append(out, "Выбери количество токенов, после которых будет происходить сжатие чата")
|
||
for _, i := range tokenMethodCount {
|
||
kb.AddCallbackButton(strconv.Itoa(i), "rp.compress_setting", method, i)
|
||
}
|
||
}
|
||
kb.AddLine().AddCallbackButton("Назад", "rp.compress_setting_s1", method)
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStyleDanger, "rp.info")
|
||
ctx.EditCallback(strings.Join(out, "\n"), kb)
|
||
}
|
||
func compressSetting(ctx *laniakea.MsgContext, db *database.Context) {
|
||
rep := psql.NewRPRepository(db)
|
||
user, err := rep.GetUser(int64(ctx.FromID))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
if len(ctx.Args) != 2 {
|
||
return
|
||
}
|
||
method := ctx.Args[0]
|
||
count, err := strconv.Atoi(ctx.Args[1])
|
||
if err != nil {
|
||
return
|
||
}
|
||
|
||
user.CompressMethod = method
|
||
user.CompressLimit = count
|
||
_, err = rep.UpdateUserCompressSettings(user)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
}
|
||
kb := laniakea.NewInlineKeyboardJson(1)
|
||
kb.AddCallbackButtonStyle("На главную", laniakea.ButtonStylePrimary, "rp.info")
|
||
ctx.EditCallback("Настройки сжатия были обновлены", kb)
|
||
}
|