memory leak
This commit is contained in:
35
database/mdb/ai_chats.go
Normal file
35
database/mdb/ai_chats.go
Normal file
@@ -0,0 +1,35 @@
|
||||
package mdb
|
||||
|
||||
import (
|
||||
"context"
|
||||
"kurumibot/database"
|
||||
"time"
|
||||
|
||||
"git.nix13.pw/scuroneko/laniakea"
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
)
|
||||
|
||||
func GetGptChatHistory(db *laniakea.DatabaseContext, chatId string) ([]*AiChatMessage, error) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
col := database.GetMongoCollection(db, "gpt_chat_messages")
|
||||
cursor, err := col.Find(ctx, bson.M{"chatId": chatId})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make([]*AiChatMessage, 0)
|
||||
err = cursor.All(ctx, &result)
|
||||
return result, err
|
||||
}
|
||||
func UpdateGptChatHistory(db *laniakea.DatabaseContext, chatId, role, message string) error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
col := database.GetMongoCollection(db, "gpt_chat_messages")
|
||||
_, err := col.InsertOne(ctx, AiChatMessage{
|
||||
bson.NewObjectID(),
|
||||
chatId,
|
||||
role,
|
||||
message, 0,
|
||||
})
|
||||
return err
|
||||
}
|
||||
@@ -5,16 +5,19 @@ import (
|
||||
"kurumibot/database"
|
||||
"time"
|
||||
|
||||
"git.nix13.pw/scuroneko/extypes"
|
||||
"git.nix13.pw/scuroneko/laniakea"
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||||
)
|
||||
|
||||
type ConsoleLogEntry struct {
|
||||
Level string `bson:"level"`
|
||||
Prefix string `bson:"prefix"`
|
||||
Traceback string `bson:"traceback"`
|
||||
Message string `bson:"message"`
|
||||
Time time.Time `bson:"time"`
|
||||
TimeStamp int64 `bson:"timeStamp"`
|
||||
Level string `bson:"level" json:"level"`
|
||||
Prefix string `bson:"prefix" json:"prefix"`
|
||||
Traceback string `bson:"traceback" json:"traceback"`
|
||||
Message string `bson:"message" json:"message"`
|
||||
Time time.Time `bson:"time" json:"time"`
|
||||
TimeStamp int64 `bson:"timeStamp" json:"time_stamp"`
|
||||
}
|
||||
|
||||
func WriteConsoleLog(db *laniakea.DatabaseContext, e *ConsoleLogEntry) error {
|
||||
@@ -24,13 +27,29 @@ func WriteConsoleLog(db *laniakea.DatabaseContext, e *ConsoleLogEntry) error {
|
||||
_, err := col.InsertOne(ctx, e)
|
||||
return err
|
||||
}
|
||||
func GetConsoleLogs(db *laniakea.DatabaseContext) (extypes.Slice[*ConsoleLogEntry], error) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
col := database.GetMongoCollection(db, "logs")
|
||||
opts := options.Find()
|
||||
opts.SetLimit(5)
|
||||
opts.SetSort(bson.D{{"_id", 1}})
|
||||
cursor, err := col.Find(ctx, bson.D{}, opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer cursor.Close(ctx)
|
||||
result := make(extypes.Slice[*ConsoleLogEntry], 0)
|
||||
err = cursor.All(ctx, &result)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
type MessageLogEntry struct {
|
||||
MessageID int `bson:"messageId"`
|
||||
SenderID int `bson:"senderId"`
|
||||
ChatID int `bson:"chatId"`
|
||||
Text string `bson:"text"`
|
||||
TimeStamp int64 `bson:"timestamp"`
|
||||
MessageID int `bson:"messageId" json:"message_id"`
|
||||
SenderID int `bson:"senderId" json:"sender_id"`
|
||||
ChatID int `bson:"chatId" json:"chat_id"`
|
||||
Text string `bson:"text" json:"text"`
|
||||
TimeStamp int64 `bson:"timestamp" json:"timestamp"`
|
||||
}
|
||||
|
||||
func WriteMessageLog(db *laniakea.DatabaseContext, e *MessageLogEntry) error {
|
||||
@@ -40,3 +59,19 @@ func WriteMessageLog(db *laniakea.DatabaseContext, e *MessageLogEntry) error {
|
||||
_, err := col.InsertOne(ctx, e)
|
||||
return err
|
||||
}
|
||||
func GetMessageLogs(db *laniakea.DatabaseContext) (extypes.Slice[*MessageLogEntry], error) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
col := database.GetMongoCollection(db, "msg_logs")
|
||||
opts := options.Find()
|
||||
opts.SetLimit(5)
|
||||
opts.SetSort(bson.D{{"_id", 1}})
|
||||
cursor, err := col.Find(ctx, bson.D{}, opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer cursor.Close(ctx)
|
||||
result := make(extypes.Slice[*MessageLogEntry], 0)
|
||||
err = cursor.All(ctx, &result)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -9,7 +9,7 @@ import (
|
||||
"go.mongodb.org/mongo-driver/v2/bson"
|
||||
)
|
||||
|
||||
type RPChatMessage struct {
|
||||
type AiChatMessage struct {
|
||||
Id bson.ObjectID `bson:"_id"`
|
||||
ChatID string `bson:"chatId"`
|
||||
Role string `bson:"role"`
|
||||
@@ -17,7 +17,7 @@ type RPChatMessage struct {
|
||||
Index int `bson:"index"`
|
||||
}
|
||||
|
||||
func GetChatHistory(db *laniakea.DatabaseContext, chatId string) ([]*RPChatMessage, error) {
|
||||
func GetRPChatHistory(db *laniakea.DatabaseContext, chatId string) ([]*AiChatMessage, error) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
col := database.GetMongoCollection(db, "rp_chat_messages")
|
||||
@@ -25,15 +25,15 @@ func GetChatHistory(db *laniakea.DatabaseContext, chatId string) ([]*RPChatMessa
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make([]*RPChatMessage, 0)
|
||||
result := make([]*AiChatMessage, 0)
|
||||
err = cursor.All(ctx, &result)
|
||||
return result, err
|
||||
}
|
||||
func UpdateChatHistory(db *laniakea.DatabaseContext, chatId, role, message string, index int) error {
|
||||
func UpdateRPChatHistory(db *laniakea.DatabaseContext, chatId, role, message string, index int) error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
col := database.GetMongoCollection(db, "rp_chat_messages")
|
||||
_, err := col.InsertOne(ctx, RPChatMessage{
|
||||
_, err := col.InsertOne(ctx, AiChatMessage{
|
||||
bson.NewObjectID(),
|
||||
chatId,
|
||||
role,
|
||||
@@ -41,13 +41,13 @@ func UpdateChatHistory(db *laniakea.DatabaseContext, chatId, role, message strin
|
||||
})
|
||||
return err
|
||||
}
|
||||
func GetChatHistorySize(db *laniakea.DatabaseContext, chatId string) (int64, error) {
|
||||
func GetRPChatHistorySize(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{"chatId": chatId})
|
||||
}
|
||||
func DeleteChatEntry(db *laniakea.DatabaseContext, entry *RPChatMessage) error {
|
||||
func DeleteRPChatEntry(db *laniakea.DatabaseContext, entry *AiChatMessage) error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
col := database.GetMongoCollection(db, "rp_chat_messages")
|
||||
|
||||
35
database/psql/ai.go
Normal file
35
database/psql/ai.go
Normal file
@@ -0,0 +1,35 @@
|
||||
package psql
|
||||
|
||||
import (
|
||||
"git.nix13.pw/scuroneko/laniakea"
|
||||
"github.com/vinovest/sqlx"
|
||||
)
|
||||
|
||||
type AIModel struct {
|
||||
ID string
|
||||
Key string
|
||||
Name string
|
||||
ContextSize int `db:"context_size"`
|
||||
}
|
||||
|
||||
type AIRepository struct {
|
||||
db *sqlx.DB
|
||||
}
|
||||
|
||||
func newAiRepository(db *sqlx.DB) *AIRepository {
|
||||
return &AIRepository{db}
|
||||
}
|
||||
func NewAIRepository(db *laniakea.DatabaseContext) *AIRepository {
|
||||
return newAiRepository(db.PostgresSQL)
|
||||
}
|
||||
|
||||
func (rep *AIRepository) GetModel(id string) (*AIModel, error) {
|
||||
model := new(AIModel)
|
||||
err := rep.db.Get(model, "SELECT * FROM ai_models WHERE id=$1;", id)
|
||||
return model, err
|
||||
}
|
||||
func (rep *AIRepository) GetAllModels() ([]*AIModel, error) {
|
||||
models := make([]*AIModel, 0)
|
||||
err := rep.db.Select(&models, "SELECT * FROM ai_models ORDER BY id;")
|
||||
return models, err
|
||||
}
|
||||
@@ -21,12 +21,6 @@ type RPScenario struct {
|
||||
Description string
|
||||
Prompt string
|
||||
}
|
||||
type RPModel struct {
|
||||
ID string
|
||||
Key string
|
||||
Name string
|
||||
ContextSize int `db:"context_size"`
|
||||
}
|
||||
type RPSetting struct {
|
||||
ID int
|
||||
Name string
|
||||
@@ -41,7 +35,7 @@ type RPUser struct {
|
||||
SelectedPreset string `db:"selected_preset"`
|
||||
Preset *RPPreset
|
||||
SelectedModel string `db:"selected_model"`
|
||||
Model *RPModel
|
||||
Model *AIModel
|
||||
|
||||
CompressMethod string `db:"compress_method"`
|
||||
CompressLimit int `db:"compress_limit"`
|
||||
@@ -72,7 +66,8 @@ func (rep *RPRepository) CreateUser(id int64) (*RPUser, error) {
|
||||
if err != nil {
|
||||
return user, err
|
||||
}
|
||||
user.Model, err = rep.GetModel(user.SelectedModel)
|
||||
aiRep := newAiRepository(rep.db)
|
||||
user.Model, err = aiRep.GetModel(user.SelectedModel)
|
||||
return user, err
|
||||
}
|
||||
func (rep *RPRepository) GetUser(id int64) (*RPUser, error) {
|
||||
@@ -85,7 +80,8 @@ func (rep *RPRepository) GetUser(id int64) (*RPUser, error) {
|
||||
if err != nil {
|
||||
return user, err
|
||||
}
|
||||
user.Model, err = rep.GetModel(user.SelectedModel)
|
||||
aiRep := newAiRepository(rep.db)
|
||||
user.Model, err = aiRep.GetModel(user.SelectedModel)
|
||||
return user, err
|
||||
}
|
||||
func (rep *RPRepository) UpdateUser(user *RPUser) error {
|
||||
@@ -133,17 +129,6 @@ func (rep *RPRepository) GetScenario(id int) (*RPScenario, error) {
|
||||
return scenario, err
|
||||
}
|
||||
|
||||
func (rep *RPRepository) GetModel(id string) (*RPModel, error) {
|
||||
model := new(RPModel)
|
||||
err := rep.db.Get(model, "SELECT * FROM rp_models WHERE id=$1;", id)
|
||||
return model, err
|
||||
}
|
||||
func (rep *RPRepository) GetAllModels() ([]*RPModel, error) {
|
||||
models := make([]*RPModel, 0)
|
||||
err := rep.db.Select(&models, "SELECT * FROM rp_models ORDER BY id;")
|
||||
return models, err
|
||||
}
|
||||
|
||||
func (rep *RPRepository) GetAllSettings() ([]*RPSetting, error) {
|
||||
settings := make([]*RPSetting, 0)
|
||||
err := rep.db.Select(&settings, "SELECT * FROM rp_settings ORDER BY id;")
|
||||
|
||||
35
database/red/ai_chats.go
Normal file
35
database/red/ai_chats.go
Normal file
@@ -0,0 +1,35 @@
|
||||
package red
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"git.nix13.pw/scuroneko/laniakea"
|
||||
"github.com/google/uuid"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
type AiRepository struct {
|
||||
client *redis.Client
|
||||
}
|
||||
|
||||
func NewAiRepository(db *laniakea.DatabaseContext) *AiRepository {
|
||||
return &AiRepository{client: db.Redis}
|
||||
}
|
||||
|
||||
func (rep *AiRepository) SetChatId(userId int, chatId string) error {
|
||||
key := fmt.Sprintf("ai.chats.gpt.%d", userId)
|
||||
return rep.client.Set(ctx, key, chatId, 0).Err()
|
||||
}
|
||||
func (rep *AiRepository) GetChatId(userId int) (string, error) {
|
||||
key := fmt.Sprintf("ai.chats.gpt.%d", userId)
|
||||
return rep.client.Get(ctx, key).Result()
|
||||
}
|
||||
func (rep *AiRepository) GetOrCreateChatId(userId int) (string, error) {
|
||||
key := fmt.Sprintf("ai.chats.gpt.%d", userId)
|
||||
res := rep.client.Get(ctx, key)
|
||||
if res.Err() != nil {
|
||||
chatId := uuid.New().String()
|
||||
return chatId, rep.SetChatId(userId, chatId)
|
||||
}
|
||||
return res.Result()
|
||||
}
|
||||
Reference in New Issue
Block a user