refactor: 重构文档处理流程和任务管理

service/document_chunk.go

@@ -3,15 +3,11 @@ package service
 import (
 	"context"
 	"rag/common/eino"
-	"rag/consts/document"
-	"rag/consts/public"
+	"rag/common/task"
 	"rag/dao"
 	"rag/model/dto"
 	"rag/model/entity"
 
-	gmq "github.com/bjang03/gmq/core/gmq"
-	"github.com/bjang03/gmq/mq"
-	"github.com/bjang03/gmq/types"
 	"github.com/cloudwego/eino/components/indexer"
 	"github.com/cloudwego/eino/schema"
 	"github.com/gogf/gf/v2/frame/g"
@@ -22,10 +18,6 @@ var DocumentChunk = new(documentChunkService)
 
 type documentChunkService struct{}
 
-const (
-	DatasetIndexStatusReady = "ready"
-)
-
 // Update 更新文件块
 func (s *documentChunkService) Update(ctx context.Context, req *dto.UpdateDocumentChunkReq) (err error) {
 	_, err = dao.DocumentChunk.Update(ctx, req)
@@ -60,32 +52,29 @@ func (s *documentChunkService) DocsChunkMsg(ctx context.Context, msg any) (err e
 	idx := eino.NewPGVectorIndexer(&eino.PGVectorIndexerOptions{
 		BatchSize: 10,
 	})
+	documentId := gconv.Int64(docs[0].MetaData[entity.DocumentChunkCol.DocumentId])
 	rows, err := idx.Store(ctx, docs, indexer.WithEmbedding(eino.EmbedderDashscope))
 	if err != nil || rows == 0 {
 		g.Log().Error(ctx, "DocsChunkMsg rows: , err:", rows, err)
+		// 写入任务进度失败 任务类型为sql存储
+		remark := " 向量存储数量: " + gconv.String(rows)
+		if err != nil {
+			remark = "向量存储失败: " + err.Error()
+		}
+		err = Task.WriteTaskProgress(ctx, &dto.WriteTaskProgressReq{
+			TaskId:   documentId,
+			TaskType: task.TaskTypeGenerateVector,
+			Status:   task.TaskStatusFailed,
+			Remark:   remark,
+		})
 		return
 	}
-	tenantId := gconv.Uint64(docs[0].MetaData[entity.DocumentChunkCol.TenantId])
-	creator := gconv.String(docs[0].MetaData[entity.DocumentChunkCol.Creator])
-	documentId := gconv.Int64(docs[0].MetaData[entity.DocumentChunkCol.DocumentId])
-	err = s.publishKnowledgeDocumentMsg(ctx, tenantId, creator, documentId, document.VectorStatusCompleted.Code())
-
-	return
-}
-
-// publishKnowledgeDocumentMsg 发布消息
-func (s *documentChunkService) publishKnowledgeDocumentMsg(ctx context.Context, tenantId uint64, creator string, documentId int64, vectorStatus document.VectorStatus) (err error) {
-	knowledgeDocumentMsg := dto.KnowledgeDocumentMsg{
-		TenantId:     tenantId,
-		Creator:      creator,
-		Id:           documentId,
-		VectorStatus: vectorStatus,
-	}
-	err = gmq.GetGmq("primary").GmqPublish(ctx, &mq.RedisPubMessage{
-		PubMessage: types.PubMessage{
-			Topic: public.KnowledgeDocumentVectorStatusTopic,
-			Data:  knowledgeDocumentMsg,
-		},
+	// 写入任务进度成功 任务类型为sql存储
+	err = Task.WriteTaskProgress(ctx, &dto.WriteTaskProgressReq{
+		TaskId:   documentId,
+		TaskType: task.TaskTypeGenerateVector,
+		Status:   task.TaskStatusCompleted,
+		Remark:   "向量生成完成",
 	})
 	return
 }