91 lines
2.3 KiB
Go
91 lines
2.3 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sort"
|
|
"strings"
|
|
|
|
"rag-local/common"
|
|
"rag-local/kb/dao"
|
|
"rag-local/kb/model/entity"
|
|
)
|
|
|
|
var KgRelationService = &kgRelationService{}
|
|
|
|
type kgRelationService struct{}
|
|
|
|
const (
|
|
kgLinkTopN = 3 // 实体链接命中前 N 个实体
|
|
kgNeighborLimit = 20 // 一跳邻居三元组上限
|
|
)
|
|
|
|
// GraphEnhance 图增强检索：问题分词与实体名匹配（Top3）→ 一跳邻居三元组 → 格式化文本（供注入提示词）
|
|
func (s *kgRelationService) GraphEnhance(ctx context.Context, datasetId int64, question string) ([]string, error) {
|
|
names, err := dao.KgEntity.ListNames(ctx, datasetId)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
linked := linkEntities(question, names, kgLinkTopN)
|
|
if len(linked) == 0 {
|
|
return nil, nil
|
|
}
|
|
triples, err := dao.KgRelation.Neighbors(ctx, datasetId, linked, kgNeighborLimit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out := make([]string, 0, len(triples))
|
|
for _, t := range triples {
|
|
out = append(out, fmt.Sprintf("%s -%s-> %s", t.Head, t.Relation, t.Tail))
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// linkEntities 问题分词后按 token 命中实体名的个数打分，取前 topN；同分按名称长度优先
|
|
func linkEntities(question string, names []string, topN int) []string {
|
|
tokens := strings.Fields(common.Tokenize(question))
|
|
if len(tokens) == 0 {
|
|
return nil
|
|
}
|
|
type scored struct {
|
|
name string
|
|
score int
|
|
}
|
|
var hits []scored
|
|
for _, n := range names {
|
|
if n == "" {
|
|
continue
|
|
}
|
|
score := 0
|
|
if strings.Contains(question, n) {
|
|
score += 5 // 问题中出现完整实体名，强相关
|
|
}
|
|
for _, t := range tokens {
|
|
if strings.Contains(n, t) {
|
|
score += len([]rune(t)) // 命中 token 越长相关性越高
|
|
}
|
|
}
|
|
if score > 0 {
|
|
hits = append(hits, scored{name: n, score: score})
|
|
}
|
|
}
|
|
sort.Slice(hits, func(i, j int) bool {
|
|
if hits[i].score != hits[j].score {
|
|
return hits[i].score > hits[j].score
|
|
}
|
|
return len(hits[i].name) > len(hits[j].name)
|
|
})
|
|
if len(hits) > topN {
|
|
hits = hits[:topN]
|
|
}
|
|
out := make([]string, 0, len(hits))
|
|
for _, h := range hits {
|
|
out = append(out, h.name)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (s *kgRelationService) List(ctx context.Context, datasetId int64, page, pageSize int) ([]*entity.KgRelation, int, error) {
|
|
return dao.KgRelation.List(ctx, datasetId, page, pageSize)
|
|
}
|