92 lines
2.6 KiB
Go
92 lines
2.6 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"rag/common/eino"
|
|
"rag/consts/document"
|
|
"rag/consts/public"
|
|
"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"
|
|
"github.com/gogf/gf/v2/util/gconv"
|
|
)
|
|
|
|
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 docs = make([]*schema.Document, 0)
|
|
msgMap := gconv.Map(msg)
|
|
if err = gconv.Structs(msgMap["data"], &docs); err != nil {
|
|
g.Log().Error(ctx, "DocsChunkMsg err:", err)
|
|
return
|
|
}
|
|
if len(docs) == 0 {
|
|
g.Log().Error(ctx, "DocsChunkMsg err:", "msg is empty")
|
|
return
|
|
}
|
|
|
|
idx := eino.NewPGVectorIndexer(&eino.PGVectorIndexerOptions{
|
|
BatchSize: 10,
|
|
})
|
|
rows, err := idx.Store(ctx, docs, indexer.WithEmbedding(eino.EmbedderDashscope))
|
|
if err != nil || rows == 0 {
|
|
g.Log().Error(ctx, "DocsChunkMsg rows: , err:", rows, err)
|
|
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,
|
|
},
|
|
})
|
|
return
|
|
}
|