198 lines
7.0 KiB
Go
198 lines
7.0 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"customer-server/dao"
|
||
"customer-server/model/entity"
|
||
"fmt"
|
||
|
||
"gitea.com/red-future/common/ragflow"
|
||
"github.com/gogf/gf/v2/errors/gerror"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
)
|
||
|
||
// EnsureTenantDataset 确保租户知识库存在
|
||
// 逻辑：
|
||
// 1. 用租户ID查找MongoDB中的知识库ID
|
||
// 2. 测试知识库ID是否在RAGFlow中可用
|
||
// 3. 如果不可用，就直接根据名字在RAGFlow中查找
|
||
// 4. 找到了就把知识库ID存到MongoDB
|
||
// 5. 如果名字也找不到，就创建新的
|
||
// 6. 把新的ID和名字存到MongoDB
|
||
func EnsureTenantDataset(ctx context.Context, tenantId string) (datasetId string, err error) {
|
||
ragflowClient := ragflow.GetGlobalClient()
|
||
datasetName := fmt.Sprintf("dataset_tenant_%s", tenantId)
|
||
|
||
// 1. 用租户ID查找MongoDB中的知识库ID
|
||
datasetId, err = dao.RAGFlowConfig.FindDatasetIdByTenant(ctx, tenantId)
|
||
if err != nil && err.Error() != "mongo: no documents in result" {
|
||
return "", gerror.Wrap(err, "查询租户知识库ID失败")
|
||
}
|
||
|
||
// 2. 测试知识库ID是否在RAGFlow中可用
|
||
if datasetId != "" {
|
||
g.Log().Infof(ctx, "租户%s已有知识库ID: %s，测试是否可用", tenantId, datasetId)
|
||
|
||
listReq := &ragflow.ListDatasetsReq{Page: 1, PageSize: 100}
|
||
listRes, listErr := ragflowClient.ListDatasets(ctx, listReq)
|
||
|
||
if listErr == nil && listRes != nil {
|
||
g.Log().Infof(ctx, "ListDatasets返回%d个知识库", len(listRes.Data))
|
||
for i, ds := range listRes.Data {
|
||
g.Log().Infof(ctx, "知识库[%d]: id=%s, name=%s", i, ds.Id, ds.Name)
|
||
}
|
||
|
||
// 测试ID是否可用
|
||
for _, ds := range listRes.Data {
|
||
if ds.Id == datasetId {
|
||
g.Log().Infof(ctx, "知识库ID可用: %s", datasetId)
|
||
return datasetId, nil
|
||
}
|
||
}
|
||
|
||
// 3. ID不可用，直接根据名字查找
|
||
g.Log().Warningf(ctx, "知识库ID不可用，根据名字查找: %s", datasetName)
|
||
for _, ds := range listRes.Data {
|
||
if ds.Name == datasetName {
|
||
datasetId = ds.Id
|
||
g.Log().Infof(ctx, "找到同名知识库: name=%s, id=%s", datasetName, datasetId)
|
||
|
||
// 4. 把知识库ID存到MongoDB
|
||
if updateErr := updateDatasetIdInMongo(ctx, tenantId, datasetId); updateErr != nil {
|
||
g.Log().Errorf(ctx, "更新MongoDB失败: %v", updateErr)
|
||
}
|
||
|
||
return datasetId, nil
|
||
}
|
||
}
|
||
|
||
g.Log().Warningf(ctx, "未找到同名知识库: %s", datasetName)
|
||
} else {
|
||
g.Log().Errorf(ctx, "ListDatasets失败: %v", listErr)
|
||
}
|
||
}
|
||
|
||
// 5. 名字也找不到，创建新的
|
||
g.Log().Infof(ctx, "创建新知识库: name=%s", datasetName)
|
||
|
||
// 从config读取embedding模型，如果未配置则使用默认值
|
||
embeddingModel := g.Cfg().MustGet(ctx, "ragflow.embedding_model", "text-embedding-v4").String()
|
||
|
||
createReq := &ragflow.CreateDatasetReq{
|
||
Name: datasetName,
|
||
EmbeddingModel: embeddingModel,
|
||
ChunkMethod: "naive",
|
||
}
|
||
|
||
dataset, err := ragflowClient.CreateDataset(ctx, createReq)
|
||
|
||
// 不依赖CreateDataset的返回值，因为RAGFlow可能返回code=0但data=null
|
||
// 立即用ListDatasets查找新创建的知识库来确认
|
||
g.Log().Infof(ctx, "创建请求已发送，立即查找确认: name=%s", datasetName)
|
||
|
||
listReq2 := &ragflow.ListDatasetsReq{Page: 1, PageSize: 100}
|
||
listRes2, listErr2 := ragflowClient.ListDatasets(ctx, listReq2)
|
||
|
||
if listErr2 == nil && listRes2 != nil {
|
||
for _, ds := range listRes2.Data {
|
||
if ds.Name == datasetName {
|
||
datasetId = ds.Id
|
||
g.Log().Infof(ctx, "确认知识库已创建: id=%s, name=%s", datasetId, datasetName)
|
||
goto saveRecord
|
||
}
|
||
}
|
||
}
|
||
|
||
// 如果ListDatasets也找不到，说明创建真的失败了
|
||
if err != nil {
|
||
return "", gerror.Wrapf(err, "创建知识库失败且无法确认")
|
||
}
|
||
if dataset == nil || dataset.Id == "" {
|
||
return "", gerror.Newf("创建知识库失败，无法通过ListDatasets确认: name=%s", datasetName)
|
||
}
|
||
|
||
// 使用CreateDataset返回的ID（兜底）
|
||
datasetId = dataset.Id
|
||
g.Log().Infof(ctx, "知识库创建成功: id=%s, name=%s", datasetId, datasetName)
|
||
|
||
saveRecord:
|
||
|
||
// 6. 把新的ID和名字存到MongoDB
|
||
if err := saveSystemDatasetRecord(ctx, tenantId, datasetId, datasetName); err != nil {
|
||
g.Log().Errorf(ctx, "保存到MongoDB失败: %v", err)
|
||
return "", gerror.Wrap(err, "保存到MongoDB失败")
|
||
}
|
||
|
||
return datasetId, nil
|
||
}
|
||
|
||
// uploadPlaceholderDocument 上传占位文档
|
||
func uploadPlaceholderDocument(ctx context.Context, datasetId, tenantId string) error {
|
||
ragflowClient := ragflow.GetGlobalClient()
|
||
placeholderContent := "欢迎使用智能客服系统。这是一个占位文档，实际产品和话术会在后续上传。"
|
||
placeholderFilename := fmt.Sprintf("placeholder_tenant_%s.txt", tenantId)
|
||
|
||
documentId, err := ragflowClient.UploadDocumentFromText(ctx, datasetId, placeholderContent, placeholderFilename)
|
||
if err != nil {
|
||
return gerror.Wrap(err, "上传占位文档失败")
|
||
}
|
||
|
||
g.Log().Infof(ctx, "占位文档上传成功: document_id=%s", documentId)
|
||
|
||
if parseErr := ragflowClient.ParseDocuments(ctx, datasetId, []string{documentId}); parseErr != nil {
|
||
g.Log().Warningf(ctx, "解析占位文档失败: %v", parseErr)
|
||
return parseErr
|
||
}
|
||
|
||
g.Log().Infof(ctx, "占位文档解析已启动（后台异步进行）")
|
||
return nil
|
||
}
|
||
|
||
// saveSystemDatasetRecord 保存系统占位记录到MongoDB
|
||
func saveSystemDatasetRecord(ctx context.Context, tenantId, datasetId, datasetName string) error {
|
||
// 先检查是否已有记录
|
||
existing, err := dao.RAGFlowConfig.FindDatasetIdByTenant(ctx, tenantId)
|
||
if err == nil && existing != "" {
|
||
g.Log().Infof(ctx, "租户%s已有知识库记录，无需创建系统占位记录", tenantId)
|
||
return nil
|
||
}
|
||
|
||
// 创建系统占位记录
|
||
systemAccountName := fmt.Sprintf("_system_tenant_%s", tenantId)
|
||
config := &entity.RAGFlowConfig{
|
||
AccountName: systemAccountName,
|
||
Platform: "system",
|
||
DatasetId: datasetId,
|
||
DatasetIds: []string{datasetId},
|
||
DatasetName: datasetName,
|
||
ChatId: "", // 系统记录不需要Chat
|
||
Prompt: "",
|
||
DocumentIds: []string{},
|
||
SimilarityThreshold: 0.2,
|
||
KeywordsSimilarityWeight: 0.7,
|
||
TopN: 8,
|
||
EmptyResponse: "抱歉，我暂时无法回答这个问题。",
|
||
SyncStatus: "synced",
|
||
}
|
||
// 统一使用string类型存储tenantId到MongoDB
|
||
config.TenantId = tenantId
|
||
config.IsDeleted = false
|
||
|
||
if err := dao.RAGFlowConfig.Insert(ctx, config); err != nil {
|
||
return gerror.Wrap(err, "插入系统占位记录失败")
|
||
}
|
||
|
||
g.Log().Infof(ctx, "系统占位记录已保存: tenant_id=%s, dataset_id=%s", tenantId, datasetId)
|
||
return nil
|
||
}
|
||
|
||
// updateDatasetIdInMongo 更新MongoDB中租户的所有datasetId记录
|
||
func updateDatasetIdInMongo(ctx context.Context, tenantId, newDatasetId string) error {
|
||
// 更新该租户所有ragflow_config记录的datasetId
|
||
if err := dao.RAGFlowConfig.UpdateDatasetIdByTenant(ctx, tenantId, newDatasetId); err != nil {
|
||
return gerror.Wrap(err, "更新租户datasetId失败")
|
||
}
|
||
g.Log().Infof(ctx, "已更新租户%s的所有记录的datasetId: %s", tenantId, newDatasetId)
|
||
return nil
|
||
}
|