feat: rag初始版

service/document_chunk.go

@@ -0,0 +1,176 @@
+package service
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"fmt"
+	"rag/consts/document"
+	"rag/consts/public"
+	"rag/dao"
+	"rag/model/dto"
+	"rag/model/entity"
+
+	"gitea.com/red-future/common/beans"
+	"gitea.com/red-future/common/rag/eino"
+	gmq "github.com/bjang03/gmq/core/gmq"
+	"github.com/bjang03/gmq/mq"
+	"github.com/bjang03/gmq/types"
+	"github.com/gogf/gf/v2/frame/g"
+	"github.com/gogf/gf/v2/util/gconv"
+	"github.com/pgvector/pgvector-go"
+)
+
+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)
+	return
+}
+
+// List 获取文件块列表
+func (s *documentChunkService) List(ctx context.Context, req *dto.ListDocumentChunkReq) (res *dto.ListDocumentChunkRes, err error) {
+	list, total, err := dao.DocumentChunk.List(ctx, req)
+	if err != nil {
+		return
+	}
+	res = &dto.ListDocumentChunkRes{
+		Total: total,
+	}
+	err = gconv.Struct(list, &res.List)
+	return
+}
+
+func (s *documentChunkService) DocsChunkMsg(ctx context.Context, msg any) (err error) {
+	var req = make([]*dto.VectorDocumentChunkMsg, 0)
+	msgMap := gconv.Map(msg)
+	if err = gconv.Structs(msgMap["data"], &req); err != nil {
+		g.Log().Error(ctx, "DocsChunkMsg err:", err)
+		return
+	}
+	if len(req) == 0 {
+		g.Log().Error(ctx, "DocsChunkMsg err:", "msg is empty")
+		return
+	}
+
+	ctx = context.WithValue(ctx, "user", &beans.User{
+		TenantId: req[0].TenantId,
+		UserName: req[0].Creator,
+	})
+
+	// 调用eino接口获取向量
+	var vectorDocsStr = make([]string, 0, len(req))
+	for _, t := range req {
+		vectorDocsStr = append(vectorDocsStr, t.Content)
+	}
+	embeddings, err := eino.EmbedStrings(ctx, vectorDocsStr)
+	if err != nil {
+		g.Log().Error(ctx, "DocsChunkMsg err:", err)
+		err = s.publishKnowledgeDocumentMsg(ctx, req[0].TenantId, req[0].Creator, req[0].DocumentId, document.VectorStatusFailed.Code())
+		return
+	}
+
+	// 获取向量维度
+	dimension := 0
+	if len(embeddings) > 0 {
+		dimension = len(embeddings[0])
+	}
+
+	// 创建或更新DatasetIndex
+	err = s.createOrUpdateDatasetIndex(ctx, req[0].DatasetId, dimension, int64(len(req)))
+	if err != nil {
+		g.Log().Error(ctx, "CreateOrUpdateDatasetIndex err:", err)
+		err = s.publishKnowledgeDocumentMsg(ctx, req[0].TenantId, req[0].Creator, req[0].DocumentId, document.VectorStatusFailed.Code())
+		return
+	}
+
+	// 更新向量文档
+	for i, embedding := range embeddings {
+		req[i].Vector = pgvector.NewVector(gconv.Float32s(embedding))
+		req[i].VectorStatus = document.VectorStatusCompleted.Code()
+		req[i].Status = document.StatusEnable.Code()
+	}
+	_, err = dao.DocumentChunk.BatchInsert(ctx, req)
+	if err != nil {
+		g.Log().Error(ctx, "DocsChunkMsg err:", err)
+		err = s.publishKnowledgeDocumentMsg(ctx, req[0].TenantId, req[0].Creator, req[0].DocumentId, document.VectorStatusFailed.Code())
+		return
+	}
+
+	err = s.publishKnowledgeDocumentMsg(ctx, req[0].TenantId, req[0].Creator, req[0].DocumentId, document.VectorStatusCompleted.Code())
+
+	return
+}
+
+// createOrUpdateDatasetIndex 创建或更新数据集索引
+func (s *documentChunkService) createOrUpdateDatasetIndex(ctx context.Context, datasetId int64, dimension int, vectorCount int64) (err error) {
+	// 查询数据集是否已有索引
+	existIndex, err := dao.DatasetIndex.GetByDatasetId(ctx, datasetId)
+	if err != nil && !errors.Is(err, sql.ErrNoRows) {
+		return err
+	}
+
+	// 已有索引 → 只更新数量
+	if existIndex != nil {
+		_ = dao.DatasetIndex.IncVectorCount(ctx, existIndex.Id, vectorCount)
+		return nil
+	}
+
+	// ====================== 创建新索引 ======================
+	indexName := fmt.Sprintf("idx_dataset_%d_vector", datasetId) // 真实PG索引名
+	// 1. 插入索引配置
+	index := &entity.DatasetIndex{
+		DatasetId:   datasetId,
+		Name:        indexName,
+		Dimension:   dimension,
+		FieldType:   "float",
+		MetricType:  "COSINE",
+		Status:      gconv.PtrInt8(1),
+		VectorCount: vectorCount,
+		Description: fmt.Sprintf("数据集%d向量索引", datasetId),
+	}
+	_, err = dao.DatasetIndex.Insert(ctx, index)
+	if err != nil {
+		return err
+	}
+
+	// 2. 真正创建 PGVector 索引（唯一真实索引！）
+	err = s.createRealPGVectorIndex(ctx, indexName)
+	return err
+}
+
+// createRealPGVectorIndex 真正在PostgreSQL创建向量索引（真实可用）
+func (s *documentChunkService) createRealPGVectorIndex(ctx context.Context, indexName string) error {
+	// 执行真实建索引语句
+	err := dao.DatasetIndex.InsertIndex(ctx, indexName)
+	if err != nil {
+		g.Log().Error(ctx, "创建向量索引失败：", err)
+		return err
+	}
+	g.Log().Info(ctx, "PGVector真实索引创建成功："+indexName)
+	return nil
+}
+
+// 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,
+		},
+	})
+	return
+}