56 lines
1.7 KiB
Go
56 lines
1.7 KiB
Go
package mdb
|
|
|
|
import (
|
|
"context"
|
|
"kurumibot/database"
|
|
"time"
|
|
|
|
"git.nix13.pw/scuroneko/laniakea"
|
|
"go.mongodb.org/mongo-driver/v2/bson"
|
|
)
|
|
|
|
type RPChatMessage struct {
|
|
Id bson.ObjectID `bson:"_id"`
|
|
ChatID string `bson:"chat_id"`
|
|
Role string `bson:"role"`
|
|
Message string `bson:"message"`
|
|
}
|
|
|
|
func GetChatHistory(db *laniakea.DatabaseContext, chatId string) ([]*RPChatMessage, error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
col := database.GetMongoCollection(db, "rp_chat_messages")
|
|
cursor, err := col.Find(ctx, bson.M{"chat_id": chatId})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
result := make([]*RPChatMessage, 0)
|
|
err = cursor.All(ctx, &result)
|
|
return result, err
|
|
}
|
|
func UpdateChatHistory(db *laniakea.DatabaseContext, chatId, role, message string) error {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
col := database.GetMongoCollection(db, "rp_chat_messages")
|
|
_, err := col.InsertOne(ctx, RPChatMessage{
|
|
bson.NewObjectID(),
|
|
chatId,
|
|
role,
|
|
message,
|
|
})
|
|
return err
|
|
}
|
|
func GetChatHistorySize(db *laniakea.DatabaseContext, chatId string) (int64, error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
col := database.GetMongoCollection(db, "rp_chat_messages")
|
|
return col.CountDocuments(ctx, bson.M{"chat_id": chatId})
|
|
}
|
|
func DeleteChatEntry(db *laniakea.DatabaseContext, entry *RPChatMessage) error {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
col := database.GetMongoCollection(db, "rp_chat_messages")
|
|
_, err := col.DeleteOne(ctx, bson.M{"chat_id": entry.ChatID})
|
|
return err
|
|
}
|