119 lines
3.7 KiB
Go
119 lines
3.7 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"strings"
|
|
|
|
"rag-local/kb/consts"
|
|
"rag-local/kb/dao"
|
|
"rag-local/kb/model/entity"
|
|
|
|
"github.com/cloudwego/eino/schema"
|
|
"github.com/gogf/gf/v2/errors/gerror"
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
)
|
|
|
|
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 抽取（挂在解析流水线分块落库之后）。
|
|
// 每个分块一次调用，任何失败只记日志，不阻断解析流水线；未配置默认对话模型时直接跳过。
|
|
func (s *kgEntityService) ExtractDocument(ctx context.Context, datasetId, documentId int64) error {
|
|
chunks, _, err := dao.Chunk.ListByDocument(ctx, documentId, 1, 100000)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
model, err := s.buildModel(ctx)
|
|
if err != nil {
|
|
g.Log().Infof(ctx, "kg extract skipped: %v", err)
|
|
return nil
|
|
}
|
|
for _, c := range chunks {
|
|
if err := s.extractChunk(ctx, model, datasetId, c); err != nil {
|
|
g.Log().Warningf(ctx, "kg extract chunk %d failed: %v", c.Id, err)
|
|
}
|
|
}
|
|
return 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)
|
|
}
|
|
|
|
func (s *kgEntityService) extractChunk(ctx context.Context, model *OpenAIChatModel, datasetId int64, chunk *entity.Chunk) 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)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
data, err := parseKgJSON(resp.Content)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, e := range data.Entities {
|
|
name := strings.TrimSpace(e.Name)
|
|
if name == "" {
|
|
continue
|
|
}
|
|
if err := dao.KgEntity.Upsert(ctx, datasetId, chunk.Id, name, strings.TrimSpace(e.Type)); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
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
|
|
}
|
|
if err := dao.KgRelation.Insert(ctx, datasetId, chunk.Id, head, relation, tail); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// parseKgJSON 解析模型输出的 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 {
|
|
return nil, gerror.Wrap(err, "解析抽取 JSON 失败")
|
|
}
|
|
return &out, nil
|
|
}
|
|
|
|
func (s *kgEntityService) List(ctx context.Context, datasetId int64, page, pageSize int) ([]*entity.KgEntity, int, error) {
|
|
return dao.KgEntity.List(ctx, datasetId, page, pageSize)
|
|
}
|