250 lines
8.3 KiB
Go
250 lines
8.3 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 KgEntity = &kgEntityDao{}
|
|
|
|
type kgEntityDao struct{}
|
|
|
|
func init() {
|
|
ctx := context.Background()
|
|
_, err := g.DB(consts.DbGroupDefault).Exec(ctx, `CREATE TABLE IF NOT EXISTS `+consts.TableNameKgEntity+` (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
dataset_id INTEGER NOT NULL DEFAULT 0,
|
|
name TEXT NOT NULL DEFAULT '',
|
|
entity_type TEXT NOT NULL DEFAULT '',
|
|
chunk_id INTEGER NOT NULL DEFAULT 0,
|
|
created_at DATETIME DEFAULT (datetime('now','localtime')),
|
|
updated_at DATETIME DEFAULT (datetime('now','localtime')),
|
|
UNIQUE(dataset_id, name)
|
|
)`)
|
|
if err != nil {
|
|
g.Log().Warningf(ctx, "create kg_entity table failed: %v", err)
|
|
}
|
|
if _, err := g.DB(consts.DbGroupDefault).Exec(ctx, "CREATE INDEX IF NOT EXISTS idx_kg_entity_dataset ON "+consts.TableNameKgEntity+"(dataset_id)"); err != nil {
|
|
g.Log().Warningf(ctx, "create index idx_kg_entity_dataset failed: %v", err)
|
|
}
|
|
}
|
|
|
|
func (d *kgEntityDao) GetOne(ctx context.Context, id int64) (*entity.KgEntity, error) {
|
|
var m entity.KgEntity
|
|
err := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgEntity).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 *kgEntityDao) List(ctx context.Context, datasetId int64, page, pageSize int) ([]*entity.KgEntity, int, error) {
|
|
if page < 1 {
|
|
page = 1
|
|
}
|
|
if pageSize < 1 {
|
|
pageSize = 20
|
|
}
|
|
m := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgEntity).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.KgEntity
|
|
err = m.Page(page, pageSize).OrderDesc("id").Scan(&list)
|
|
if list == nil {
|
|
list = make([]*entity.KgEntity, 0)
|
|
}
|
|
return list, total, err
|
|
}
|
|
|
|
// ListNames 数据集全部实体名（实体链接在内存中打分，本地库规模可控）
|
|
func (d *kgEntityDao) ListNames(ctx context.Context, datasetId int64) ([]string, error) {
|
|
var list []*entity.KgEntity
|
|
err := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgEntity).Ctx(ctx).
|
|
Fields("name").Where("dataset_id", datasetId).Scan(&list)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
names := make([]string, 0, len(list))
|
|
for _, e := range list {
|
|
names = append(names, e.Name)
|
|
}
|
|
return names, nil
|
|
}
|
|
|
|
// KgEntityItem 批量写入条目
|
|
type KgEntityItem struct {
|
|
Name string
|
|
EntityType string
|
|
}
|
|
|
|
// UpsertBatch 按 (dataset_id, name) 去重批量写入，已存在则更新类型与来源（单条 SQL 多行 VALUES）
|
|
func (d *kgEntityDao) UpsertBatch(ctx context.Context, datasetId, chunkId int64, items []KgEntityItem) 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.TableNameKgEntity + `
|
|
(dataset_id, name, entity_type, chunk_id, created_at, updated_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].Name, items[i].EntityType, chunkId, now, now)
|
|
}
|
|
sb.WriteString(`
|
|
ON CONFLICT(dataset_id, name) DO UPDATE SET
|
|
entity_type=excluded.entity_type, chunk_id=excluded.chunk_id, updated_at=excluded.updated_at`)
|
|
if _, err := g.DB(consts.DbGroupDefault).Exec(ctx, sb.String(), args...); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// HasGraphByDocument 判定某文档是否产出过图谱（单表约束：先取分块 id，再数关系）。
|
|
// 不用实体行计数：实体按 (dataset_id,name) 去重、chunk_id 会被跨文件 upsert 覆盖，按实体计数会漏计。
|
|
// 关系每 chunk 至少一条、不受覆盖影响。
|
|
func (d *kgEntityDao) HasGraphByDocument(ctx context.Context, documentId int64) (bool, error) {
|
|
chunkIds, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameChunk).Ctx(ctx).
|
|
Fields("id").Where("document_id", documentId).Array()
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
// 分块数可能超 SQLite 变量上限，按 100 分批
|
|
for start := 0; start < len(chunkIds); start += 100 {
|
|
end := min(start+100, len(chunkIds))
|
|
ids := make([]int64, 0, end-start)
|
|
for _, v := range chunkIds[start:end] {
|
|
ids = append(ids, v.Int64())
|
|
}
|
|
n, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgRelation).Ctx(ctx).
|
|
WhereIn("chunk_id", ids).Count()
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if n > 0 {
|
|
return true, nil
|
|
}
|
|
}
|
|
return false, nil
|
|
}
|
|
|
|
func (d *kgEntityDao) 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.TableNameKgEntity+" WHERE chunk_id IN ("+placeholders+")", args...)
|
|
return err
|
|
}
|
|
|
|
func (d *kgEntityDao) DeleteByDataset(ctx context.Context, datasetId int64) error {
|
|
_, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgEntity).Ctx(ctx).
|
|
Where("dataset_id", datasetId).Delete()
|
|
return err
|
|
}
|
|
|
|
// KgEntitySource 实体 → 出现文件清单（dao 内已按名聚合）
|
|
type KgEntitySource struct {
|
|
Name string
|
|
Files []string
|
|
}
|
|
|
|
// SourcesByDataset 实体名 → 出现文件清单（单表约束：拆 4 条单表查询 + 内存组装）。
|
|
// 有关系的实体经关系表溯源（关系不去重、每条带来源 chunk，覆盖实体全部出现文件）；
|
|
// 孤立实体（无任何关系）由实体行 chunk_id 弱引用兜底（该弱引用是 upsert 的最近来源，不保证全部出现文件）。
|
|
// 不再建实体-分块关联表：该组合已覆盖"图谱页标注来源文件"的展示需求。
|
|
func (d *kgEntityDao) SourcesByDataset(ctx context.Context, datasetId int64) ([]KgEntitySource, error) {
|
|
// 1. 关系表（单表）：数据集全部三元组的 head/tail 及其来源分块
|
|
relRows, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgRelation).Ctx(ctx).
|
|
Fields("head", "tail", "chunk_id").Where("dataset_id", datasetId).All()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// 2. 实体表（单表）：孤立实体弱引用兜底
|
|
entRows, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameKgEntity).Ctx(ctx).
|
|
Fields("name", "chunk_id").Where("dataset_id", datasetId).All()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// 3. 分块表（单表）：数据集全部分块 chunk_id → document_id
|
|
chunkRows, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameChunk).Ctx(ctx).
|
|
Fields("id", "document_id").Where("dataset_id", datasetId).All()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// 4. 文档表（单表）：数据集全部文档 id → 文件名
|
|
docRows, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameDocument).Ctx(ctx).
|
|
Fields("id", "filename").Where("dataset_id", datasetId).All()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
docByID := make(map[int64]string, len(docRows))
|
|
for _, row := range docRows {
|
|
docByID[row["id"].Int64()] = row["filename"].String()
|
|
}
|
|
docByChunk := make(map[int64]int64, len(chunkRows))
|
|
for _, row := range chunkRows {
|
|
docByChunk[row["id"].Int64()] = row["document_id"].Int64()
|
|
}
|
|
filesByName := make(map[string]map[string]bool)
|
|
add := func(name string, chunkId int64) {
|
|
docId, ok := docByChunk[chunkId]
|
|
if !ok {
|
|
return
|
|
}
|
|
fname := docByID[docId]
|
|
if name == "" || fname == "" {
|
|
return
|
|
}
|
|
if filesByName[name] == nil {
|
|
filesByName[name] = make(map[string]bool)
|
|
}
|
|
filesByName[name][fname] = true
|
|
}
|
|
for _, row := range relRows {
|
|
add(row["head"].String(), row["chunk_id"].Int64())
|
|
add(row["tail"].String(), row["chunk_id"].Int64())
|
|
}
|
|
for _, row := range entRows {
|
|
add(row["name"].String(), row["chunk_id"].Int64())
|
|
}
|
|
out := make([]KgEntitySource, 0, len(filesByName))
|
|
for name, set := range filesByName {
|
|
files := make([]string, 0, len(set))
|
|
for f := range set {
|
|
files = append(files, f)
|
|
}
|
|
out = append(out, KgEntitySource{Name: name, Files: files})
|
|
}
|
|
return out, nil
|
|
}
|