red-future/rag
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 支持多租户多模型对话及文档去重优化

Browse Source
This commit is contained in:
曲皓东
2026-04-16 15:47:37 +08:00
parent 4ead3f82cf
commit 27b1dd3c27
34 changed files with 2188 additions and 315 deletions
Download Patch File Download Diff File Expand all files Collapse all files

144
service/document_vector.go
View File

@@ -4,17 +4,18 @@ import (
"context"
"fmt"
"rag/common/eino"
"rag/consts/model"
"rag/consts/task"
"rag/dao"
"rag/model/dto"
"rag/model/entity"
"gitea.com/red-future/common/beans"
"github.com/cloudwego/eino/components/indexer"
"github.com/cloudwego/eino/components/retriever"
"github.com/cloudwego/eino/schema"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/util/gconv"
"github.com/pgvector/pgvector-go"
)
var DocumentVector = new(documentVectorService)
@@ -23,23 +24,32 @@ type documentVectorService struct{}
// Query 执行RAG查询
func (s *documentVectorService) Query(ctx context.Context, req *dto.RAGQueryReq) (*dto.RAGQueryRes, error) {
if req.TopK <= 0 {
req.TopK = 5
modelInfo, err := dao.Model.Get(ctx, &dto.GetModelReq{
ModelType: model.ModelTypeChat.Code(),
})
if err != nil {
g.Log().Errorf(ctx, "获取模型失败: %v", err)
return nil, fmt.Errorf("获取模型失败: %w", err)
}
if modelInfo == nil {
g.Log().Errorf(ctx, "模型不存在: %v", model.ModelTypeChat.Code())
return nil, fmt.Errorf("模型不存在: %w", err)
}
// 4. 使用向量检索器进行查询
r, err := eino.NewPGVectorRetriever(&eino.PGVectorRetrieverConfig{
Embedder: eino.EmbedderDashscope,
r, err := eino.NewPGVectorRetriever(ctx, &eino.PGVectorRetrieverConfig{
DefaultTopK: req.TopK,
})
}, model.ModelConfigTypeVectorDashScope.Code()) //TODO 后续替换成本地模型
if err != nil {
g.Log().Errorf(ctx, "初始化向量检索器失败: %v", err)
return nil, fmt.Errorf("初始化向量检索器失败: %w", err)
}
// 5. 执行向量检索
docs, err := r.Retrieve(ctx, req.Content, retriever.WithEmbedding(eino.EmbedderDashscope), retriever.WithDSLInfo(map[string]any{
"dataset_ids": req.DatasetIds,
docs, err := r.Retrieve(ctx, req.Content, retriever.WithDSLInfo(map[string]any{
"dataset_ids": req.DatasetIds,
"document_ids": req.DocumentIds,
}))
if err != nil {
g.Log().Errorf(ctx, "向量检索失败: %v", err)
@@ -53,7 +63,7 @@ func (s *documentVectorService) Query(ctx context.Context, req *dto.RAGQueryReq)
return nil, fmt.Errorf("转换历史消息失败: %w", err)
}
replyMsg, err := eino.NewChatModel(ctx, req.Content, docs, messages)
replyMsg, err := eino.NewChatModel(ctx, req.Content, docs, messages, modelInfo.ConfigType)
if err != nil {
g.Log().Errorf(ctx, "向量检索失败: %v", err)
return nil, fmt.Errorf("向量检索失败: %w", err)
@@ -98,26 +108,108 @@ func (s *documentVectorService) DocsChunkMsg(ctx context.Context, msg any) (err
TenantId: gconv.Uint64(docs[0].MetaData[entity.DocumentVectorCol.TenantId]),
UserName: gconv.String(docs[0].MetaData[entity.DocumentVectorCol.Creator]),
})
idx := eino.NewPGVectorIndexer(&eino.PGVectorIndexerOptions{
BatchSize: 10,
})
documentId := gconv.Int64(docs[0].MetaData[entity.DocumentVectorCol.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()
var docsStore = make([]*schema.Document, 0)
var docsInsert = make([]*dto.VectorDocumentVectorMsg, 0)
for _, doc := range docs {
if gconv.Bool(doc.MetaData["isNew"]) {
docsStore = append(docsStore, doc)
} else {
ck := new(dto.VectorDocumentVectorMsg)
err = gconv.Struct(doc.MetaData, ck)
ck.Content = doc.Content
ck.VectorStatus = gconv.PtrInt8(1)
ck.Status = gconv.PtrInt8(1)
docsInsert = append(docsInsert, ck)
}
err = Task.WriteTaskProgress(ctx, &dto.WriteTaskProgressReq{
TaskId: documentId,
TaskType: task.TaskTypeGenerateVector,
Status: task.TaskStatusFailed,
Remark: remark,
})
return
}
if !g.IsEmpty(docsStore) {
idx := eino.NewPGVectorIndexer(&eino.PGVectorIndexerOptions{
BatchSize: 10,
})
var rows int64
rows, err = idx.Store(ctx, docsStore, model.ModelConfigTypeVectorDashScope.Code()) //TODO 后续替换成本地模型
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
}
}
if !g.IsEmpty(docsInsert) {
// 1. 提取所有 contentHash
contentHashs := make([]string, 0, len(docsInsert))
for _, d := range docsInsert {
contentHashs = append(contentHashs, d.ContentHash)
}
// 2. 分页查询已存在的向量（一页1000，避免大查询）
var existVectors []*entity.DocumentVector
for page := 1; ; page++ {
res, total, err := dao.DocumentVector.List(ctx, &dto.ListDocumentVectorReq{
Page: &beans.Page{PageSize: 1000, PageNum: int64(page)},
ContentHashs: contentHashs,
})
if err != nil {
return err
}
if len(res) == 0 {
break
}
existVectors = append(existVectors, res...)
if len(existVectors) >= total {
break
}
}
// 3. 构建哈希 -> 向量 的映射表（O(1) 查找，性能提升巨大）
vectorMap := make(map[string]pgvector.Vector, len(existVectors))
for _, v := range existVectors {
vectorMap[v.ContentHash] = v.Vector
}
// 4. 回填向量 + 过滤掉数据库已存在的数据（避免重复插入）
for _, d := range docsInsert {
// 回填已有向量
if vec, ok := vectorMap[d.ContentHash]; ok {
d.Vector = vec
}
}
var rows int64
rows, err = dao.DocumentVector.BatchInsert(ctx, docsInsert)
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
}
}
// 写入任务进度成功 任务类型为sql存储
err = Task.WriteTaskProgress(ctx, &dto.WriteTaskProgressReq{
TaskId: documentId,