138 lines
4.1 KiB
Go
138 lines
4.1 KiB
Go
package dao
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"strings"
|
|
|
|
"rag-local/kb/consts"
|
|
"rag-local/kb/model/entity"
|
|
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
"github.com/gogf/gf/v2/os/gtime"
|
|
)
|
|
|
|
var KgRelation = &kgRelationDao{}
|
|
|
|
type kgRelationDao struct{}
|
|
|
|
func init() {
|
|
ctx := context.Background()
|
|
_, err := g.DB(consts.DbGroupDefault).Exec(ctx, `CREATE TABLE IF NOT EXISTS `+consts.TableNameKgRelation+` (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
dataset_id INTEGER NOT NULL DEFAULT 0,
|
|
head TEXT NOT NULL DEFAULT '',
|
|
relation TEXT NOT NULL DEFAULT '',
|
|
tail TEXT NOT NULL DEFAULT '',
|
|
chunk_id INTEGER NOT NULL DEFAULT 0,
|
|
created_at DATETIME DEFAULT (datetime('now','localtime'))
|
|
)`)
|
|
if err != nil {
|
|
g.Log().Warningf(ctx, "create kg_relation table failed: %v", err)
|
|
}
|
|
if _, err := g.DB(consts.DbGroupDefault).Exec(ctx, "CREATE INDEX IF NOT EXISTS idx_kg_relation_dataset ON "+consts.TableNameKgRelation+"(dataset_id)"); err != nil {
|
|
g.Log().Warningf(ctx, "create index idx_kg_relation_dataset failed: %v", err)
|
|
}
|
|
}
|
|
|
|
func (d *kgRelationDao) GetOne(ctx context.Context, id int64) (*entity.KgRelation, error) {
|
|
var m entity.KgRelation
|
|
err := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgRelation).Ctx(ctx).Where("id", id).Scan(&m)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &m, nil
|
|
}
|
|
|
|
func (d *kgRelationDao) List(ctx context.Context, datasetId int64, page, pageSize int) ([]*entity.KgRelation, int, error) {
|
|
if page < 1 {
|
|
page = 1
|
|
}
|
|
if pageSize < 1 {
|
|
pageSize = 20
|
|
}
|
|
m := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgRelation).Ctx(ctx)
|
|
if datasetId > 0 {
|
|
m = m.Where("dataset_id", datasetId)
|
|
}
|
|
total, err := m.Count()
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
var list []*entity.KgRelation
|
|
err = m.Page(page, pageSize).OrderDesc("id").Scan(&list)
|
|
return list, total, err
|
|
}
|
|
|
|
func (d *kgRelationDao) Insert(ctx context.Context, datasetId, chunkId int64, head, relation, tail string) error {
|
|
_, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgRelation).Ctx(ctx).Data(g.Map{
|
|
"dataset_id": datasetId,
|
|
"head": head,
|
|
"relation": relation,
|
|
"tail": tail,
|
|
"chunk_id": chunkId,
|
|
"created_at": gtime.Now().Format("Y-m-d H:i:s"),
|
|
}).Insert()
|
|
return err
|
|
}
|
|
|
|
// Neighbors 一跳邻居：head 或 tail 命中实体名的三元组（实体链接用）
|
|
func (d *kgRelationDao) Neighbors(ctx context.Context, datasetId int64, names []string, limit int) ([]*entity.KgRelation, error) {
|
|
if len(names) == 0 {
|
|
return nil, nil
|
|
}
|
|
if limit < 1 {
|
|
limit = 20
|
|
}
|
|
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(names)), ",")
|
|
args := make([]any, 0, len(names)*2+2)
|
|
for _, n := range names {
|
|
args = append(args, n, n)
|
|
}
|
|
args = append(args, datasetId, limit)
|
|
r, err := g.DB(consts.DbGroupDefault).Ctx(ctx).Raw(`
|
|
SELECT id, dataset_id, head, relation, tail, chunk_id, created_at
|
|
FROM `+consts.TableNameKgRelation+`
|
|
WHERE (head IN (`+placeholders+`) OR tail IN (`+placeholders+`)) AND dataset_id = ?
|
|
ORDER BY id LIMIT ?`, args...).All()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
list := make([]*entity.KgRelation, 0, len(r))
|
|
for _, row := range r {
|
|
list = append(list, &entity.KgRelation{
|
|
Id: row["id"].Int64(),
|
|
DatasetId: row["dataset_id"].Int64(),
|
|
Head: row["head"].String(),
|
|
Relation: row["relation"].String(),
|
|
Tail: row["tail"].String(),
|
|
ChunkId: row["chunk_id"].Int64(),
|
|
})
|
|
}
|
|
return list, nil
|
|
}
|
|
|
|
func (d *kgRelationDao) DeleteByChunkIds(ctx context.Context, chunkIds []int64) error {
|
|
if len(chunkIds) == 0 {
|
|
return nil
|
|
}
|
|
placeholders := strings.TrimSuffix(strings.Repeat("?,", len(chunkIds)), ",")
|
|
args := make([]any, 0, len(chunkIds))
|
|
for _, id := range chunkIds {
|
|
args = append(args, id)
|
|
}
|
|
_, err := g.DB(consts.DbGroupDefault).Exec(ctx,
|
|
"DELETE FROM "+consts.TableNameKgRelation+" WHERE chunk_id IN ("+placeholders+")", args...)
|
|
return err
|
|
}
|
|
|
|
func (d *kgRelationDao) DeleteByDataset(ctx context.Context, datasetId int64) error {
|
|
_, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgRelation).Ctx(ctx).
|
|
Where("dataset_id", datasetId).Delete()
|
|
return err
|
|
}
|