306 lines
10 KiB
Go
306 lines
10 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"strings"
|
||
"sync"
|
||
|
||
"rag-local/common"
|
||
"rag-local/kb/consts"
|
||
"rag-local/kb/dao"
|
||
"rag-local/kb/model/entity"
|
||
|
||
emodel "github.com/cloudwego/eino/components/model"
|
||
"github.com/cloudwego/eino/schema"
|
||
"github.com/gogf/gf/v2/errors/gerror"
|
||
)
|
||
|
||
var KgEntityService = &kgEntityService{}
|
||
|
||
type kgEntityService struct{}
|
||
|
||
type kgEntityItem struct {
|
||
Name string `json:"name"`
|
||
Type string `json:"type"`
|
||
}
|
||
|
||
type kgRelationItem struct {
|
||
Head string `json:"head"`
|
||
Relation string `json:"relation"`
|
||
Tail string `json:"tail"`
|
||
}
|
||
|
||
type kgExtractResult struct {
|
||
Entities []kgEntityItem `json:"entities"`
|
||
Relations []kgRelationItem `json:"relations"`
|
||
}
|
||
|
||
// ExtractDocument 对文档全部新分块做 LLM 抽取（挂在解析流水线分块落库之后）。
|
||
// 每个分块一次调用，任何失败只记日志，不阻断解析流水线；返回未成功抽取的分块数，供调用方标记图谱未构建。
|
||
// 并行：LLM 调用（纯 IO）提交 common.KgExtractPool 并发执行，落库（SQLite 写）收敛回本 goroutine 串行。
|
||
func (s *kgEntityService) ExtractDocument(ctx context.Context, datasetId, documentId int64) (int, error) {
|
||
chunks, _, err := dao.Chunk.ListByDocument(ctx, documentId, 1, 100000)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
model, err := s.buildModel(ctx)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
// 池内并发共享 model 实例，只读字段须在此（单 goroutine）预置好
|
||
model.DisableThinking()
|
||
|
||
type extractOut struct {
|
||
chunkId int64
|
||
entities []kgEntityItem
|
||
relations []kgRelationItem
|
||
err error
|
||
}
|
||
ch := make(chan extractOut, len(chunks))
|
||
var wg sync.WaitGroup
|
||
for _, c := range chunks {
|
||
wg.Add(1)
|
||
if err := common.KgExtractPool.AddWithRecover(ctx, func(ctx context.Context) {
|
||
defer wg.Done()
|
||
out := extractOut{chunkId: c.Id}
|
||
out.entities, out.relations, out.err = s.callExtract(ctx, model, c)
|
||
ch <- out
|
||
}, func(ctx context.Context, e error) {
|
||
defer wg.Done()
|
||
ch <- extractOut{chunkId: c.Id, err: e}
|
||
}); err != nil {
|
||
wg.Done()
|
||
fmt.Printf("kg extract chunk %d submit failed: %v\n", c.Id, err)
|
||
}
|
||
}
|
||
go func() { wg.Wait(); close(ch) }()
|
||
|
||
failed := 0
|
||
for out := range ch {
|
||
if out.err != nil {
|
||
failed++
|
||
fmt.Printf("kg extract chunk %d failed: %v\n", out.chunkId, out.err)
|
||
continue
|
||
}
|
||
if err := s.saveExtract(ctx, datasetId, out.chunkId, out.entities, out.relations); err != nil {
|
||
failed++
|
||
fmt.Printf("kg extract chunk %d failed: %v\n", out.chunkId, err)
|
||
}
|
||
}
|
||
return failed, nil
|
||
}
|
||
|
||
func (s *kgEntityService) buildModel(ctx context.Context) (*OpenAIChatModel, error) {
|
||
defaultChatModel, err := dao.ModelConfig.GetDefault(ctx, consts.ModelTypeChat)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if defaultChatModel <= 0 {
|
||
return nil, gerror.New("未设置默认对话模型")
|
||
}
|
||
return BuildChatModel(ctx, defaultChatModel)
|
||
}
|
||
|
||
// callExtract 池内执行：LLM 抽取 + JSON 解析（纯 IO，不做任何写库）
|
||
func (s *kgEntityService) callExtract(ctx context.Context, model *OpenAIChatModel, chunk *entity.Chunk) ([]kgEntityItem, []kgRelationItem, error) {
|
||
msgs := []*schema.Message{
|
||
{Role: schema.System, Content: "你是知识抽取助手。从文档片段中抽取实体（人名、组织、地名、产品等专有名词）及实体间的关系（动词或介词短语）。只输出 JSON，不要 markdown 代码块或任何解释，格式：{\"entities\":[{\"name\":\"实体名\",\"type\":\"类型\"}],\"relations\":[{\"head\":\"主体\",\"relation\":\"关系\",\"tail\":\"客体\"}]}"},
|
||
{Role: schema.User, Content: "文档片段：\n" + chunk.Content},
|
||
}
|
||
resp, err := model.Generate(ctx, msgs, emodel.WithMaxTokens(consts.MaxExtractTokens))
|
||
if err != nil {
|
||
return nil, nil, err
|
||
}
|
||
data, err := parseKgJSON(resp.Content)
|
||
if err != nil {
|
||
// 模型偶发返回空内容/截断 JSON，重试一次通常正常
|
||
fmt.Printf("kg extract chunk %d parse failed (%v), retrying\n", chunk.Id, err)
|
||
resp, err = model.Generate(ctx, msgs, emodel.WithMaxTokens(consts.MaxExtractTokens))
|
||
if err != nil {
|
||
return nil, nil, err
|
||
}
|
||
data, err = parseKgJSON(resp.Content)
|
||
if err != nil {
|
||
return nil, nil, err
|
||
}
|
||
}
|
||
entities := make([]kgEntityItem, 0, len(data.Entities))
|
||
for _, e := range data.Entities {
|
||
if name := strings.TrimSpace(e.Name); name != "" {
|
||
entities = append(entities, kgEntityItem{Name: name, Type: strings.TrimSpace(e.Type)})
|
||
}
|
||
}
|
||
relations := make([]kgRelationItem, 0, len(data.Relations))
|
||
for _, r := range data.Relations {
|
||
head, relation, tail := strings.TrimSpace(r.Head), strings.TrimSpace(r.Relation), strings.TrimSpace(r.Tail)
|
||
if head == "" || relation == "" || tail == "" || head == tail {
|
||
continue
|
||
}
|
||
relations = append(relations, kgRelationItem{Head: head, Relation: relation, Tail: tail})
|
||
}
|
||
return entities, relations, nil
|
||
}
|
||
|
||
// saveExtract 主 goroutine 串行落库：实体与关系各一条批量 SQL（多行 VALUES）。
|
||
// 落库前归一化（实体名/谓词，防别名分裂）+ chunk 内去重 + 过滤空名/自环。
|
||
func (s *kgEntityService) saveExtract(ctx context.Context, datasetId, chunkId int64, entities []kgEntityItem, relations []kgRelationItem) error {
|
||
if len(entities) > 0 {
|
||
// 归一化 + chunk 内按名去重（同名取首个类型）
|
||
seen := make(map[string]string, len(entities))
|
||
for _, e := range entities {
|
||
name := common.NormalizeKgTerm(e.Name, consts.KgAliasMap)
|
||
if name == "" {
|
||
continue
|
||
}
|
||
if _, dup := seen[name]; dup {
|
||
continue
|
||
}
|
||
seen[name] = strings.TrimSpace(e.Type)
|
||
}
|
||
items := make([]dao.KgEntityItem, 0, len(seen))
|
||
for name, typ := range seen {
|
||
items = append(items, dao.KgEntityItem{Name: name, EntityType: typ})
|
||
}
|
||
if err := dao.KgEntity.UpsertBatch(ctx, datasetId, chunkId, items); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
if len(relations) > 0 {
|
||
// 归一化（head/tail 用实体别名表、谓词用谓词表）+ chunk 内去重 + 空/自环过滤
|
||
seen := make(map[string]bool, len(relations))
|
||
items := make([]dao.KgRelationItem, 0, len(relations))
|
||
for _, r := range relations {
|
||
head := common.NormalizeKgTerm(r.Head, consts.KgAliasMap)
|
||
rel := common.NormalizeKgTerm(r.Relation, consts.KgPredicateAliasMap)
|
||
tail := common.NormalizeKgTerm(r.Tail, consts.KgAliasMap)
|
||
if head == "" || rel == "" || tail == "" || head == tail {
|
||
continue
|
||
}
|
||
key := head + "|" + rel + "|" + tail
|
||
if seen[key] {
|
||
continue
|
||
}
|
||
seen[key] = true
|
||
items = append(items, dao.KgRelationItem{Head: head, Relation: rel, Tail: tail})
|
||
}
|
||
if err := dao.KgRelation.InsertBatch(ctx, datasetId, chunkId, items); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// graphNotBuiltMsg 与 parse_task_service 写入 error_msg 的图谱未构建标记保持一致
|
||
const graphNotBuiltMsg = "知识图谱未构建"
|
||
|
||
// IsGraphMissing 判定文档知识图谱是否缺失：error_msg 标记过未构建，或该文档分块未产出过关系
|
||
func (s *kgEntityService) IsGraphMissing(ctx context.Context, doc *entity.Document) (bool, error) {
|
||
if strings.Contains(doc.ErrorMsg, graphNotBuiltMsg) {
|
||
return true, nil
|
||
}
|
||
has, err := dao.KgEntity.HasGraphByDocument(ctx, doc.Id)
|
||
if err != nil {
|
||
return false, err
|
||
}
|
||
return !has, nil
|
||
}
|
||
|
||
// RebuildDocument 重建文档图谱：删旧（实体+关系按该文档分块）→ 全量重抽，返回未成功抽取的分块数
|
||
func (s *kgEntityService) RebuildDocument(ctx context.Context, datasetId, documentId int64) (int, error) {
|
||
chunks, _, err := dao.Chunk.ListByDocument(ctx, documentId, 1, 100000)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
chunkIds := make([]int64, 0, len(chunks))
|
||
for _, c := range chunks {
|
||
chunkIds = append(chunkIds, c.Id)
|
||
}
|
||
if err := dao.KgEntity.DeleteByChunkIds(ctx, chunkIds); err != nil {
|
||
return 0, err
|
||
}
|
||
if err := dao.KgRelation.DeleteByChunkIds(ctx, chunkIds); err != nil {
|
||
return 0, err
|
||
}
|
||
return s.ExtractDocument(ctx, datasetId, documentId)
|
||
}
|
||
|
||
// parseKgJSON 解析模型输出的 JSON，容忍 ```json 代码块包裹及首个 JSON 对象后的多余内容（模型偶尔续写 ",{...}"）
|
||
func parseKgJSON(content string) (*kgExtractResult, error) {
|
||
content = strings.TrimSpace(content)
|
||
content = strings.TrimPrefix(content, "```json")
|
||
content = strings.TrimPrefix(content, "```")
|
||
content = strings.TrimSuffix(content, "```")
|
||
content = strings.TrimSpace(content)
|
||
var out kgExtractResult
|
||
if err := json.Unmarshal([]byte(content), &out); err != nil {
|
||
if fixed, ok := extractFirstJSON(content); ok {
|
||
if err2 := json.Unmarshal([]byte(fixed), &out); err2 == nil {
|
||
return &out, nil
|
||
}
|
||
}
|
||
return nil, gerror.Wrapf(err, "解析抽取 JSON 失败，模型输出：%q", truncateRunes(content, 200))
|
||
}
|
||
return &out, nil
|
||
}
|
||
|
||
// extractFirstJSON 提取字符串中第一个完整 JSON 对象（平衡花括号，跳过字符串内的花括号）
|
||
func extractFirstJSON(s string) (string, bool) {
|
||
start := strings.IndexByte(s, '{')
|
||
if start < 0 {
|
||
return "", false
|
||
}
|
||
depth := 0
|
||
inStr := false
|
||
escaped := false
|
||
for i := start; i < len(s); i++ {
|
||
ch := s[i]
|
||
if inStr {
|
||
if escaped {
|
||
escaped = false
|
||
} else if ch == '\\' {
|
||
escaped = true
|
||
} else if ch == '"' {
|
||
inStr = false
|
||
}
|
||
continue
|
||
}
|
||
switch ch {
|
||
case '"':
|
||
inStr = true
|
||
case '{':
|
||
depth++
|
||
case '}':
|
||
depth--
|
||
if depth == 0 {
|
||
return s[start : i+1], true
|
||
}
|
||
}
|
||
}
|
||
return "", false
|
||
}
|
||
|
||
func (s *kgEntityService) List(ctx context.Context, datasetId int64, page, pageSize int) ([]*entity.KgEntity, int, error) {
|
||
return dao.KgEntity.List(ctx, datasetId, page, pageSize)
|
||
}
|
||
|
||
// KgEntitySource 实体 → 出现文件清单（图谱页节点标注来源）
|
||
type KgEntitySource struct {
|
||
Name string `json:"name"`
|
||
Files []string `json:"files"`
|
||
}
|
||
|
||
// Sources 数据集全部实体的来源文件（单表拆查 + 内存组装，见 dao.SourcesByDataset）
|
||
func (s *kgEntityService) Sources(ctx context.Context, datasetId int64) ([]*KgEntitySource, error) {
|
||
rows, err := dao.KgEntity.SourcesByDataset(ctx, datasetId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
out := make([]*KgEntitySource, 0, len(rows))
|
||
for _, r := range rows {
|
||
out = append(out, &KgEntitySource{Name: r.Name, Files: r.Files})
|
||
}
|
||
return out, nil
|
||
}
|