118 lines
3.4 KiB
Go
118 lines
3.4 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sort"
|
|
"strings"
|
|
|
|
"rag-local/common"
|
|
"rag-local/kb/consts"
|
|
"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；同分按名称长度优先。
|
|
// 别名展开：归一化后的问题包含某别名（按 key 长度降序，最长优先）且标准名在库中，
|
|
// 标准名以高分直接加入命中（用户用简称提问时能链接到全称实体）。
|
|
func linkEntities(question string, names []string, topN int) []string {
|
|
normQ := common.NormalizeKgTerm(question, consts.KgAliasMap)
|
|
tokens := strings.Fields(common.Tokenize(normQ))
|
|
if len(tokens) == 0 {
|
|
return nil
|
|
}
|
|
type scored struct {
|
|
name string
|
|
score int
|
|
}
|
|
byName := make(map[string]bool, len(names))
|
|
var hits []scored
|
|
for _, n := range names {
|
|
if n == "" {
|
|
continue
|
|
}
|
|
byName[n] = true
|
|
score := 0
|
|
if strings.Contains(normQ, 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})
|
|
}
|
|
}
|
|
// 别名展开：最长 key 优先（防"民诉法"吃掉"民诉法解释"），标准名须存在于库中
|
|
aliasKeys := make([]string, 0, len(consts.KgAliasMap))
|
|
for k := range consts.KgAliasMap {
|
|
aliasKeys = append(aliasKeys, k)
|
|
}
|
|
sort.Slice(aliasKeys, func(i, j int) bool { return len(aliasKeys[i]) > len(aliasKeys[j]) })
|
|
hitByName := make(map[string]bool, len(hits))
|
|
for _, h := range hits {
|
|
hitByName[h.name] = true
|
|
}
|
|
for _, k := range aliasKeys {
|
|
if !strings.Contains(normQ, k) {
|
|
continue
|
|
}
|
|
std := consts.KgAliasMap[k]
|
|
if std == "" || !byName[std] || hitByName[std] {
|
|
continue
|
|
}
|
|
hitByName[std] = true
|
|
hits = append(hits, scored{name: std, score: 10}) // 明确以别名提及，强于 +5 的完整名命中
|
|
}
|
|
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)
|
|
}
|