113 lines
2.8 KiB
Go
113 lines
2.8 KiB
Go
package plugins
|
||
|
||
import (
|
||
"fmt"
|
||
"kurumibot/database/mdb"
|
||
"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, "selwaifu")
|
||
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, uint(ctx.FromID), uint(waifuId))
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
ctx.Answer(fmt.Sprintf("Была выбрана вайфу %d", waifuId))
|
||
}
|
||
|
||
func newChat(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
err := red.SaveChatId(db, uint(ctx.FromID), uuid.New().String())
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
ctx.Answer("Был создан новый чат.")
|
||
}
|
||
|
||
func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
|
||
chatId := red.GetChatId(db, uint(ctx.FromID))
|
||
if chatId == "0" {
|
||
ctx.Answer("Не выбрана вайфу")
|
||
return
|
||
}
|
||
if chatId == "" {
|
||
err := red.SaveChatId(db, uint(ctx.FromID), uuid.New().String())
|
||
if err != nil {
|
||
ctx.Error(err)
|
||
return
|
||
}
|
||
}
|
||
messages := []ai.Message{
|
||
{
|
||
Role: "system",
|
||
Content: ai.FormatPrompt(ai.PreHistoryPrompt, "Яэ Мико", ctx.Msg.From.FirstName),
|
||
},
|
||
{
|
||
Role: "system",
|
||
Content: "Вот краткое описание твоего персонажа: у неё розовые волосы, лисьи уши и фиолетовые глаза. Её рост 166 см",
|
||
},
|
||
}
|
||
|
||
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, "Яэ Мико", 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{
|
||
Model: "cosmorp-2.5",
|
||
Messages: messages,
|
||
})
|
||
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")))
|
||
}
|