131 lines
4.1 KiB
Go
131 lines
4.1 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"rag-local/common"
|
|
"rag-local/kb/consts"
|
|
"rag-local/kb/dao"
|
|
"rag-local/kb/model/entity"
|
|
|
|
eembedding "github.com/cloudwego/eino/components/embedding"
|
|
"github.com/gogf/gf/v2/errors/gerror"
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
)
|
|
|
|
var ParseTaskService = &parseTaskService{}
|
|
|
|
type parseTaskService struct{}
|
|
|
|
// StartParsePoller 启动任务轮询：单 goroutine 串行消费待处理任务（与 video-factory StartVideoPoller 同模式）
|
|
func (s *parseTaskService) StartParsePoller(ctx context.Context) {
|
|
go func() {
|
|
g.Log().Info(ctx, "parse task poller started")
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-time.After(consts.ParsePollIntervalSeconds * time.Second):
|
|
s.processOne(ctx)
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
// processOne 处理一个待处理任务：解析 → 分块 → 落库（向量化在 M3 接入）
|
|
func (s *parseTaskService) processOne(ctx context.Context) {
|
|
task, err := dao.ParseTask.NextPending(ctx)
|
|
if err != nil {
|
|
g.Log().Errorf(ctx, "next parse task failed: %v", err)
|
|
return
|
|
}
|
|
if task == nil {
|
|
return
|
|
}
|
|
if err := dao.ParseTask.UpdateStatus(ctx, task.Id, consts.TaskStatusRunning, ""); err != nil {
|
|
g.Log().Errorf(ctx, "mark task running failed: %v", err)
|
|
return
|
|
}
|
|
|
|
doc, err := dao.Document.GetOne(ctx, task.DocumentId)
|
|
if err != nil {
|
|
s.fail(ctx, task, "读取文档失败: "+err.Error())
|
|
return
|
|
}
|
|
if doc == nil {
|
|
s.fail(ctx, task, "文档不存在")
|
|
return
|
|
}
|
|
if err := dao.Document.UpdateFields(ctx, doc.Id, g.Map{"status": consts.DocumentStatusParsing}); err != nil {
|
|
s.fail(ctx, task, "更新文档状态失败: "+err.Error())
|
|
return
|
|
}
|
|
|
|
text, err := common.ParseFile(filepath.Join("workspace", doc.FilePath))
|
|
if err != nil {
|
|
s.fail(ctx, task, "解析文件失败: "+err.Error())
|
|
return
|
|
}
|
|
chunks := ChunkService.SplitText(text)
|
|
// 数据集绑定 embedding 配置时构建向量模型，无配置降级为仅全文索引
|
|
var embedder eembedding.Embedder
|
|
if cfgId, err := dao.Dataset.GetEmbeddingCfgId(ctx, task.DatasetId); err == nil && cfgId > 0 {
|
|
if em, err := BuildEmbedder(ctx, cfgId); err == nil {
|
|
embedder = em
|
|
if dim := em.Dim(); dim != g.Cfg().MustGet(ctx, "vector.dim", consts.DefaultEmbeddingDim).Int() {
|
|
g.Log().Warningf(ctx, "embedding 维度 %d 与 vec0 表维度不一致，请确认 vector.dim 配置", dim)
|
|
}
|
|
} else {
|
|
g.Log().Warningf(ctx, "build embedder failed, fallback to fts only: %v", err)
|
|
}
|
|
}
|
|
if err := ChunkService.InsertAll(ctx, task.DatasetId, doc.Id, chunks, embedder); err != nil {
|
|
s.fail(ctx, task, "写入分块失败: "+err.Error())
|
|
return
|
|
}
|
|
// 第 3.5 步：知识图谱抽取（失败不阻断流水线）
|
|
if err := KgEntityService.ExtractDocument(ctx, task.DatasetId, doc.Id); err != nil {
|
|
g.Log().Warningf(ctx, "kg extract skipped for doc %d: %v", doc.Id, err)
|
|
}
|
|
if err := dao.ParseTask.UpdateStatus(ctx, task.Id, consts.TaskStatusDone, ""); err != nil {
|
|
g.Log().Errorf(ctx, "mark task done failed: %v", err)
|
|
}
|
|
}
|
|
|
|
func (s *parseTaskService) fail(ctx context.Context, task *entity.ParseTask, msg string) {
|
|
_ = dao.ParseTask.UpdateStatus(ctx, task.Id, consts.TaskStatusFailed, msg)
|
|
_ = dao.Document.UpdateFields(ctx, task.DocumentId, g.Map{
|
|
"status": consts.DocumentStatusFailed,
|
|
"error_msg": msg,
|
|
})
|
|
g.Log().Errorf(ctx, "parse task %d failed: %s", task.Id, msg)
|
|
}
|
|
|
|
// List 任务列表
|
|
func (s *parseTaskService) List(ctx context.Context, page, pageSize int) ([]*entity.ParseTask, int, error) {
|
|
return dao.ParseTask.List(ctx, page, pageSize)
|
|
}
|
|
|
|
// Retry 失败任务重置为待处理（文档状态同步重置）
|
|
func (s *parseTaskService) Retry(ctx context.Context, id int64) error {
|
|
task, err := dao.ParseTask.GetOne(ctx, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if task == nil {
|
|
return gerror.New("任务不存在")
|
|
}
|
|
if task.Status != consts.TaskStatusFailed {
|
|
return gerror.New("仅失败任务可重试")
|
|
}
|
|
if err := dao.ParseTask.UpdateStatus(ctx, id, consts.TaskStatusPending, ""); err != nil {
|
|
return err
|
|
}
|
|
return dao.Document.UpdateFields(ctx, task.DocumentId, g.Map{
|
|
"status": consts.DocumentStatusPending,
|
|
"error_msg": "",
|
|
})
|
|
}
|