ScuroNeko/YaeMikoBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

small fix

Browse Source
This commit is contained in:
ScuroNeko
2026-02-02 17:33:03 +03:00
parent 7ab5ad91de
commit 3741017b02
2 changed files with 13 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
database/mdb/rp_chats.go
View File

@@ -11,16 +11,17 @@ import (
type RPChatMessage struct { type RPChatMessage struct {
Id bson.ObjectID `bson:"_id"` Id bson.ObjectID `bson:"_id"`
ChatID string `bson:"chat_id"` ChatID string `bson:"chatId"`
Role string `bson:"role"` Role string `bson:"role"`
Message string `bson:"message"` Message string `bson:"message"`
Index int `bson:"index"`
} }
func GetChatHistory(db *laniakea.DatabaseContext, chatId string) ([]*RPChatMessage, error) { func GetChatHistory(db *laniakea.DatabaseContext, chatId string) ([]*RPChatMessage, error) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel() defer cancel()
col := database.GetMongoCollection(db, "rp_chat_messages") col := database.GetMongoCollection(db, "rp_chat_messages")
cursor, err := col.Find(ctx, bson.M{"chat_id": chatId}) cursor, err := col.Find(ctx, bson.M{"chatId": chatId})
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -28,7 +29,7 @@ func GetChatHistory(db *laniakea.DatabaseContext, chatId string) ([]*RPChatMessa
err = cursor.All(ctx, &result) err = cursor.All(ctx, &result)
return result, err return result, err
} }
func UpdateChatHistory(db *laniakea.DatabaseContext, chatId, role, message string) error { func UpdateChatHistory(db *laniakea.DatabaseContext, chatId, role, message string, index int) error {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel() defer cancel()
col := database.GetMongoCollection(db, "rp_chat_messages") col := database.GetMongoCollection(db, "rp_chat_messages")
@@ -36,7 +37,7 @@ func UpdateChatHistory(db *laniakea.DatabaseContext, chatId, role, message strin
bson.NewObjectID(), bson.NewObjectID(),
chatId, chatId,
role, role,
message, message, index,
}) })
return err return err
} }
@@ -44,12 +45,12 @@ func GetChatHistorySize(db *laniakea.DatabaseContext, chatId string) (int64, err
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel() defer cancel()
col := database.GetMongoCollection(db, "rp_chat_messages") col := database.GetMongoCollection(db, "rp_chat_messages")
return col.CountDocuments(ctx, bson.M{"chat_id": chatId}) return col.CountDocuments(ctx, bson.M{"chatId": chatId})
} }
func DeleteChatEntry(db *laniakea.DatabaseContext, entry *RPChatMessage) error { func DeleteChatEntry(db *laniakea.DatabaseContext, entry *RPChatMessage) error {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel() defer cancel()
col := database.GetMongoCollection(db, "rp_chat_messages") col := database.GetMongoCollection(db, "rp_chat_messages")
_, err := col.DeleteOne(ctx, bson.M{"chat_id": entry.ChatID}) _, err := col.DeleteOne(ctx, bson.M{"chatId": entry.ChatID})
return err return err
} }

12
plugins/rp.go
View File

@@ -605,14 +605,15 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
return return
} }
err = mdb.UpdateChatHistory(db, chatId, "user", userMessage) counter := redisRpRep.GetCounter(ctx.FromID, waifuId)
err = mdb.UpdateChatHistory(db, chatId, "user", userMessage, counter+1)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
return return
} }
agentAnswer := res.Choices[0].Message agentAnswer := res.Choices[0].Message
answerContent := strings.TrimSpace(agentAnswer.Content) answerContent := strings.TrimSpace(agentAnswer.Content)
err = mdb.UpdateChatHistory(db, chatId, agentAnswer.Role, answerContent) err = mdb.UpdateChatHistory(db, chatId, agentAnswer.Role, answerContent, counter+2)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
} }
@@ -629,7 +630,6 @@ func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
} }
counter := redisRpRep.GetCounter(ctx.FromID, waifuId)
err = redisRpRep.SetCounter(ctx.FromID, waifuId, counter+2) err = redisRpRep.SetCounter(ctx.FromID, waifuId, counter+2)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
@@ -786,14 +786,14 @@ func _compress(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
return return
} }
err = mdb.UpdateChatHistory(db, chatId, "assistant", compressedHistory) err = mdb.UpdateChatHistory(db, chatId, "assistant", compressedHistory, 0)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
} }
offset := utils.Min(len(history), 20) offset := utils.Min(len(history), 20)
for _, m := range history[len(history)-offset:] { for i, m := range history[len(history)-offset:] {
tokens += len(m.Message) tokens += len(m.Message)
err = mdb.UpdateChatHistory(db, chatId, m.Role, m.Message) err = mdb.UpdateChatHistory(db, chatId, m.Role, m.Message, i+1)
if err != nil { if err != nil {
ctx.Error(err) ctx.Error(err)
} }