253 lines
9.7 KiB
Go
253 lines
9.7 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"path/filepath"
|
||
"time"
|
||
|
||
"rag-local/common"
|
||
"rag-local/kb/consts"
|
||
"rag-local/kb/dao"
|
||
"rag-local/kb/model/entity"
|
||
|
||
"github.com/gogf/gf/v2/errors/gerror"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gogf/gf/v2/os/gtimer"
|
||
)
|
||
|
||
var ParseTaskService = &parseTaskService{}
|
||
|
||
type parseTaskService struct{}
|
||
|
||
// StartParsePoller 启动任务轮询：gtimer 单例定时器串行消费待处理任务（job 未结束不重入，与视频工厂同模式）。
|
||
// 启动时先恢复上次运行遗留的 running 任务（进程异常退出会孤儿化，不恢复则永远不被消费）。
|
||
func (s *parseTaskService) StartParsePoller(ctx context.Context) {
|
||
if n, err := dao.ParseTask.ResetRunning(ctx); err != nil {
|
||
g.Log().Warningf(ctx, "reset running parse tasks failed: %v", err)
|
||
} else if n > 0 {
|
||
g.Log().Infof(ctx, "recovered %d running parse task(s) to pending", n)
|
||
}
|
||
g.Log().Info(ctx, "parse task poller started")
|
||
gtimer.AddSingleton(ctx, consts.ParsePollIntervalSeconds*time.Second, func(ctx context.Context) {
|
||
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 task.TaskType == consts.TaskTypeReembed {
|
||
s.processReembed(ctx, task, doc)
|
||
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
|
||
}
|
||
// 数据集分块配置（大小/重叠），未设置用默认值
|
||
chunkSize, chunkOverlap := consts.DefaultChunkSize, consts.DefaultChunkOverlap
|
||
if ds, err := dao.Dataset.GetOne(ctx, task.DatasetId); err == nil && ds != nil {
|
||
chunkSize, chunkOverlap = ds.ChunkSize, ds.ChunkOverlap
|
||
}
|
||
// 数据集必须绑定向量模型，构建失败直接失败任务（不降级全文索引）
|
||
cfgId, err := dao.Dataset.GetEmbeddingCfgId(ctx, task.DatasetId)
|
||
if err != nil {
|
||
s.fail(ctx, task, "读取数据集向量模型配置失败: "+err.Error())
|
||
return
|
||
}
|
||
if cfgId <= 0 {
|
||
s.fail(ctx, task, "数据集未绑定向量模型")
|
||
return
|
||
}
|
||
embedder, err := BuildEmbedder(ctx, cfgId)
|
||
if err != nil {
|
||
s.fail(ctx, task, "构建向量模型失败: "+err.Error())
|
||
return
|
||
}
|
||
if dim := embedder.Dim(); dim != g.Cfg().MustGet(ctx, "vector.dim", consts.DefaultEmbeddingDim).Int() {
|
||
g.Log().Warningf(ctx, "embedding 维度 %d 与 vec0 表维度不一致，请确认 vector.dim 配置", dim)
|
||
}
|
||
// 重新解析时先清空旧分块，避免新旧分块并存
|
||
if err := ChunkService.DeleteByDocument(ctx, doc.Id); err != nil {
|
||
s.fail(ctx, task, "清理旧分块失败: "+err.Error())
|
||
return
|
||
}
|
||
// 原始全文落库（分块抽屉展示用），失败不阻断流水线
|
||
if err := dao.Document.UpdateFields(ctx, doc.Id, g.Map{"content": text}); err != nil {
|
||
g.Log().Warningf(ctx, "save document content failed: %v", err)
|
||
}
|
||
chunks, unitPattern, ctxPattern, err := ChunkService.SplitAuto(ctx, text, chunkSize, chunkOverlap, embedder)
|
||
if err != nil {
|
||
s.fail(ctx, task, "分块失败: "+err.Error())
|
||
return
|
||
}
|
||
// 结构识别出的模式落库（列表展示用）；未识别出结构时保留旧值，避免清空已展示的模式
|
||
if unitPattern != "" {
|
||
if err := dao.Dataset.UpdateFields(ctx, task.DatasetId, g.Map{
|
||
"unit_pattern": unitPattern,
|
||
"context_pattern": ctxPattern,
|
||
}); err != nil {
|
||
g.Log().Warningf(ctx, "save detected structure patterns failed: %v", err)
|
||
}
|
||
}
|
||
if err := dao.Document.UpdateFields(ctx, doc.Id, g.Map{"status": consts.DocumentStatusEmbedding}); err != nil {
|
||
s.fail(ctx, task, "更新文档状态失败: "+err.Error())
|
||
return
|
||
}
|
||
if err := ChunkService.InsertAll(ctx, task.DatasetId, doc.Id, chunks, embedder); err != nil {
|
||
s.fail(ctx, task, "写入分块失败: "+err.Error())
|
||
return
|
||
}
|
||
// 第 3.5 步：知识图谱抽取（失败不阻断流水线，但文档标记图谱未构建，避免误报已完成）
|
||
if err := dao.Document.UpdateFields(ctx, doc.Id, g.Map{"status": consts.DocumentStatusKgBuilding}); err != nil {
|
||
g.Log().Warningf(ctx, "mark kg building failed: %v", err)
|
||
}
|
||
kgFailed, kgErr := KgEntityService.ExtractDocument(ctx, task.DatasetId, doc.Id)
|
||
kgMsg := ""
|
||
switch {
|
||
case kgErr != nil:
|
||
kgMsg = graphNotBuiltMsg + ": " + kgErr.Error()
|
||
case kgFailed > 0:
|
||
kgMsg = fmt.Sprintf(graphNotBuiltMsg+"：%d 个分块抽取失败", kgFailed)
|
||
}
|
||
if kgMsg != "" {
|
||
g.Log().Warningf(ctx, "kg extract incomplete for doc %d: %s", doc.Id, kgMsg)
|
||
}
|
||
if err := dao.Document.UpdateFields(ctx, doc.Id, g.Map{"status": consts.DocumentStatusDone, "error_msg": kgMsg}); err != nil {
|
||
g.Log().Warningf(ctx, "mark document done failed: %v", err)
|
||
}
|
||
if err := dao.ParseTask.UpdateStatus(ctx, task.Id, consts.TaskStatusDone, kgMsg); err != nil {
|
||
g.Log().Errorf(ctx, "mark task done failed: %v", err)
|
||
}
|
||
}
|
||
|
||
// processReembed 重新向量化任务：分块文本不变，用数据集当前绑定模型重算全部向量。
|
||
// 过程状态置为向量生成中，失败恢复原状态（旧向量仍可用）；重算后若文档知识图谱缺失
|
||
// （error_msg 标记未构建或实体为空），删旧重抽补建，成功后清空 error_msg 并置为已完成。
|
||
func (s *parseTaskService) processReembed(ctx context.Context, task *entity.ParseTask, doc *entity.Document) {
|
||
prevStatus := doc.Status
|
||
if err := dao.Document.UpdateFields(ctx, doc.Id, g.Map{"status": consts.DocumentStatusEmbedding}); err != nil {
|
||
g.Log().Warningf(ctx, "mark doc embedding failed: %v", err)
|
||
}
|
||
if _, err := DocumentService.Reembed(ctx, task.DocumentId); err != nil {
|
||
_ = dao.ParseTask.UpdateStatus(ctx, task.Id, consts.TaskStatusFailed, "重新向量化失败: "+err.Error())
|
||
_ = dao.Document.UpdateFields(ctx, doc.Id, g.Map{"status": prevStatus})
|
||
g.Log().Errorf(ctx, "reembed task %d failed: %s", task.Id, err.Error())
|
||
return
|
||
}
|
||
if missing, err := KgEntityService.IsGraphMissing(ctx, doc); err != nil {
|
||
g.Log().Warningf(ctx, "check kg missing failed for doc %d: %v", doc.Id, err)
|
||
} else if missing {
|
||
kgMsg := ""
|
||
if kgFailed, kgErr := KgEntityService.RebuildDocument(ctx, doc.DatasetId, doc.Id); kgErr != nil {
|
||
kgMsg = graphNotBuiltMsg + ": " + kgErr.Error()
|
||
} else if kgFailed > 0 {
|
||
kgMsg = fmt.Sprintf(graphNotBuiltMsg+"：%d 个分块抽取失败", kgFailed)
|
||
}
|
||
if kgMsg != "" {
|
||
g.Log().Warningf(ctx, "kg rebuild incomplete for doc %d: %s", doc.Id, kgMsg)
|
||
}
|
||
if err := dao.Document.UpdateFields(ctx, doc.Id, g.Map{"error_msg": kgMsg}); err != nil {
|
||
g.Log().Warningf(ctx, "update doc error_msg failed: %v", err)
|
||
}
|
||
}
|
||
if err := dao.Document.UpdateFields(ctx, doc.Id, g.Map{"status": consts.DocumentStatusDone}); err != nil {
|
||
g.Log().Warningf(ctx, "mark doc done failed: %v", err)
|
||
}
|
||
if err := dao.ParseTask.UpdateStatus(ctx, task.Id, consts.TaskStatusDone, ""); err != nil {
|
||
g.Log().Errorf(ctx, "mark task done failed: %v", err)
|
||
}
|
||
}
|
||
|
||
// Enqueue 单文档入队（解析/重新向量化）。
|
||
// 已有待处理/运行中任务时返回错误，防重复排队；重新向量化入队即刻把文档置为向量生成中，
|
||
// 避免等待轮询器捡起（最长一个轮询间隔）期间列表状态无变化、按钮仍可重复点击。
|
||
func (s *parseTaskService) Enqueue(ctx context.Context, documentId int64, taskType string) error {
|
||
doc, err := dao.Document.GetOne(ctx, documentId)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if doc == nil {
|
||
return gerror.New("文档不存在")
|
||
}
|
||
ok, err := dao.ParseTask.HasPending(ctx, documentId)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if ok {
|
||
return gerror.New("该文档已有待处理或处理中的任务")
|
||
}
|
||
if _, err := dao.ParseTask.Insert(ctx, documentId, doc.DatasetId, taskType); err != nil {
|
||
return err
|
||
}
|
||
if taskType == consts.TaskTypeReembed {
|
||
if err := dao.Document.UpdateFields(ctx, documentId, g.Map{"status": consts.DocumentStatusEmbedding}); err != nil {
|
||
g.Log().Warningf(ctx, "mark doc %d embedding failed: %v", documentId, err)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
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": "",
|
||
})
|
||
}
|