rag-local/kb/dao/kg_relation_dao.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)
	if list == nil {
		list = make([]*entity.KgRelation, 0)
	}
	return list, total, err
}

// KgRelationItem 批量写入条目（三元组）
type KgRelationItem struct {
	Head     string
	Relation string
	Tail     string
}

// InsertBatch 批量写入三元组（单条 SQL 多行 VALUES）
func (d *kgRelationDao) InsertBatch(ctx context.Context, datasetId, chunkId int64, items []KgRelationItem) error {
	if len(items) == 0 {
		return nil
	}
	now := gtime.Now().Format("Y-m-d H:i:s")
	// SQLite 变量数上限 999，按 100 行/条分片
	for start := 0; start < len(items); start += 100 {
		end := min(start+100, len(items))
		var sb strings.Builder
		sb.WriteString("INSERT INTO " + consts.TableNameKgRelation + `
			(dataset_id, head, relation, tail, chunk_id, created_at)
			VALUES `)
		args := make([]any, 0, (end-start)*6)
		for i := start; i < end; i++ {
			if i > start {
				sb.WriteString(",")
			}
			sb.WriteString("(?,?,?,?,?,?)")
			args = append(args, datasetId, items[i].Head, items[i].Relation, items[i].Tail, chunkId, now)
		}
		if _, err := g.DB(consts.DbGroupDefault).Exec(ctx, sb.String(), args...); err != nil {
			return err
		}
	}
	return nil
}

// 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
}