55 lines
1.7 KiB
Go
55 lines
1.7 KiB
Go
package dao
|
|
|
|
import (
|
|
"context"
|
|
|
|
"rag-local/kb/consts"
|
|
"rag-local/kb/model/entity"
|
|
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
"github.com/gogf/gf/v2/os/gtime"
|
|
)
|
|
|
|
var Message = &messageDao{}
|
|
|
|
type messageDao struct{}
|
|
|
|
func init() {
|
|
ctx := context.Background()
|
|
_, err := g.DB(consts.DbGroupChat).Exec(ctx, `CREATE TABLE IF NOT EXISTS `+consts.TableNameMessage+` (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
conversation_id INTEGER NOT NULL DEFAULT 0,
|
|
role TEXT NOT NULL DEFAULT 'user',
|
|
content TEXT NOT NULL DEFAULT '',
|
|
citations TEXT NOT NULL DEFAULT '',
|
|
created_at DATETIME DEFAULT (datetime('now','localtime'))
|
|
)`)
|
|
if err != nil {
|
|
g.Log().Warningf(ctx, "create chat_message table failed: %v", err)
|
|
}
|
|
if _, err := g.DB(consts.DbGroupChat).Exec(ctx, "CREATE INDEX IF NOT EXISTS idx_chat_message_conversation ON "+consts.TableNameMessage+"(conversation_id)"); err != nil {
|
|
g.Log().Warningf(ctx, "create index idx_chat_message_conversation failed: %v", err)
|
|
}
|
|
}
|
|
|
|
func (d *messageDao) List(ctx context.Context, conversationId int64) ([]*entity.Message, error) {
|
|
var list []*entity.Message
|
|
err := g.DB(consts.DbGroupChat).Model(consts.TableNameMessage).Ctx(ctx).
|
|
Where("conversation_id", conversationId).OrderAsc("id").Scan(&list)
|
|
return list, err
|
|
}
|
|
|
|
func (d *messageDao) Insert(ctx context.Context, conversationId int64, role, content, citations string) (int64, error) {
|
|
r, err := g.DB(consts.DbGroupChat).Model(consts.TableNameMessage).Ctx(ctx).Data(g.Map{
|
|
"conversation_id": conversationId,
|
|
"role": role,
|
|
"content": content,
|
|
"citations": citations,
|
|
"created_at": gtime.Now().Format("Y-m-d H:i:s"),
|
|
}).Insert()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return r.LastInsertId()
|
|
}
|