136 lines
3.3 KiB
Go
136 lines
3.3 KiB
Go
package plugins
|
||
|
||
import (
|
||
"fmt"
|
||
"kurumibot/database/mdb"
|
||
"kurumibot/database/psql"
|
||
"kurumibot/database/red"
|
||
"kurumibot/laniakea"
|
||
"kurumibot/utils/ai"
|
||
"os"
|
||
"strconv"
|
||
"strings"
|
||
|
||
"github.com/google/uuid"
|
||
)
|
||
|
||
func RegisterTestRP(bot *laniakea.Bot) {
|
||
rp := laniakea.NewPlugin("RP")
|
||
rp = rp.Command(selectWaifu, "rpwaifu", "рпвайфу")
|
||
rp = rp.Payload(selectWaifu, "rp.selwaifu")
|
||
rp = rp.Command(newChat, "newchat")
|
||
rp = rp.Command(generate, "g", "gen", "г")
|
||
|
||
bot.AddPlugins(rp.Build())
|
||
}
|
||
|
||
func selectWaifu(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
waifuId, err := strconv.Atoi(ctx.Args[0])
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
err = red.RPSetSelectedWaifu(db, ctx.FromID, waifuId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
ctx.Answer(fmt.Sprintf("Была выбрана вайфу %d", waifuId))
|
||
}
|
||
|
||
func newChat(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
waifuId := red.RPGetSelectedWaifu(db, ctx.FromID)
|
||
if waifuId == 0 {
|
||
ctx.Answer("Не выбрана вайфу")
|
||
return
|
||
}
|
||
chatId := uuid.New()
|
||
err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId.String())
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
|
||
prompt := strings.Join(ctx.Args[1:], " ")
|
||
err = red.RPSetChatPrompt(db, ctx.FromID, waifuId, prompt)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
ctx.Answer("Был создан новый чат.")
|
||
}
|
||
|
||
func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
waifuId := red.RPGetSelectedWaifu(db, ctx.FromID)
|
||
if waifuId == 0 {
|
||
ctx.Answer("Не выбрана вайфу")
|
||
return
|
||
}
|
||
chatId := red.RPGetChatId(db, ctx.FromID, waifuId)
|
||
if chatId == "" {
|
||
err := red.RPSetChatId(db, ctx.FromID, waifuId, uuid.New().String())
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
}
|
||
waifu, err := psql.GetWaifuById(waifuId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
messages := []ai.Message{
|
||
{
|
||
Role: "system",
|
||
Content: fmt.Sprintf(
|
||
"%s %s %s",
|
||
ai.FormatPrompt(ai.PreHistoryPrompt, waifu.Name, ctx.Msg.From.FirstName),
|
||
fmt.Sprintf("Вот краткое описание твоего персонажа: %s", waifu.RpPrompt),
|
||
red.RPGetChatPrompt(db, ctx.FromID, waifuId),
|
||
),
|
||
},
|
||
}
|
||
|
||
api := ai.NewOpenAIAPI(ai.CosmoRPUrl, os.Getenv("PAWAN_KEY"), "cosmorp-2.5")
|
||
history, err := mdb.GetChatHistory(db, chatId)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
for _, m := range history {
|
||
messages = append(messages, ai.Message{
|
||
Role: m.Role,
|
||
Content: m.Message,
|
||
})
|
||
}
|
||
userMessage := strings.Join(ctx.Args, " ")
|
||
messages = append(messages, ai.Message{
|
||
Role: "system", Content: ai.FormatPrompt(ai.PostHistoryPrompt, waifu.Name, ctx.Msg.From.FirstName),
|
||
}, ai.Message{
|
||
Role: "user", Content: userMessage,
|
||
})
|
||
err = mdb.UpdateChatHistory(db, chatId, "user", userMessage)
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
m := ctx.Answer("Генерация запущена...")
|
||
res, err := api.CreateCompletion(ai.CreateCompletionReq{
|
||
Messages: messages,
|
||
Verbosity: "low",
|
||
Temperature: 0.7,
|
||
})
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
response := make([]string, 0)
|
||
for _, choice := range res.Choices {
|
||
m := choice.Message
|
||
response = append(response, m.Content)
|
||
err = mdb.UpdateChatHistory(db, chatId, m.Role, m.Content)
|
||
}
|
||
m.Edit(laniakea.EscapeMarkdown(strings.Join(response, "\n")))
|
||
}
|