136 lines
4.8 KiB
Go
136 lines
4.8 KiB
Go
package dao
|
||
|
||
import (
|
||
"context"
|
||
"customer-server/model/entity"
|
||
|
||
"gitea.com/red-future/common/db/mongo"
|
||
"github.com/gogf/gf/v2/errors/gerror"
|
||
"github.com/gogf/gf/v2/util/gconv"
|
||
"go.mongodb.org/mongo-driver/v2/bson"
|
||
)
|
||
|
||
var RAGFlowConfig = new(ragflowConfig)
|
||
|
||
type ragflowConfig struct{}
|
||
|
||
// FindByAccountName 根据客服账号名称查询配置（带租户隔离，兼容tenantId类型不一致）
|
||
func (d *ragflowConfig) FindByAccountName(ctx context.Context, accountName string) (*entity.RAGFlowConfig, error) {
|
||
// 先查询客服账号获取tenantId
|
||
account, err := CustomerServiceAccount.FindByAccountName(ctx, accountName)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if account == nil {
|
||
return nil, nil
|
||
}
|
||
|
||
// 使用accountName + tenantId查询RAGFlow配置（租户隔离）
|
||
// 先尝试原始类型查询
|
||
filter := bson.M{"accountName": accountName, "tenantId": account.TenantId, "isDeleted": false}
|
||
var config entity.RAGFlowConfig
|
||
err = MongoDAO.FindOne(ctx, filter, &config, entity.RAGFlowConfigCollection)
|
||
|
||
// 如果未找到且tenantId可以转为string，尝试用string查询（兼容性处理）
|
||
if (err != nil || config.Id == nil || config.Id.IsZero()) && account.TenantId != nil {
|
||
tenantIdStr := gconv.String(account.TenantId)
|
||
if tenantIdStr != "" {
|
||
filter = bson.M{"accountName": accountName, "tenantId": tenantIdStr, "isDeleted": false}
|
||
err = MongoDAO.FindOne(ctx, filter, &config, entity.RAGFlowConfigCollection)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
}
|
||
}
|
||
|
||
if config.Id.IsZero() {
|
||
return nil, nil
|
||
}
|
||
return &config, nil
|
||
}
|
||
|
||
// FindDatasetIdByTenant 根据租户ID查询知识库ID（从任意一条RAGFlowConfig记录中获取）
|
||
func (d *ragflowConfig) FindDatasetIdByTenant(ctx context.Context, tenantId string) (datasetId string, err error) {
|
||
// 先尝试字符串查询
|
||
filter := bson.M{"tenantId": tenantId, "isDeleted": false}
|
||
var config entity.RAGFlowConfig
|
||
err = MongoDAO.FindOne(ctx, filter, &config, entity.RAGFlowConfigCollection)
|
||
|
||
// 如果未找到且tenantId可以转为数字，尝试用数字查询（兼容MongoDB中可能存储为int的情况）
|
||
if err != nil || config.Id.IsZero() {
|
||
tenantIdInt := gconv.Int(tenantId)
|
||
if tenantIdInt > 0 {
|
||
filter = bson.M{"tenantId": tenantIdInt, "isDeleted": false}
|
||
err = MongoDAO.FindOne(ctx, filter, &config, entity.RAGFlowConfigCollection)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
}
|
||
}
|
||
|
||
if config.Id.IsZero() {
|
||
return "", nil // 未找到记录
|
||
}
|
||
return config.DatasetId, nil
|
||
}
|
||
|
||
// Insert 插入配置
|
||
func (d *ragflowConfig) Insert(ctx context.Context, config *entity.RAGFlowConfig) error {
|
||
_, err := mongo.DB().Insert(ctx, []interface{}{config}, entity.RAGFlowConfigCollection)
|
||
return err
|
||
}
|
||
|
||
// UpdateEntity 更新配置（避免双重token验证冲突）
|
||
func (d *ragflowConfig) UpdateEntity(ctx context.Context, config *entity.RAGFlowConfig) error {
|
||
filter := bson.M{"_id": config.Id, "isDeleted": false}
|
||
|
||
// 将实体转换为bson.M
|
||
updateDoc := bson.M{}
|
||
data, _ := bson.Marshal(config)
|
||
bson.Unmarshal(data, &updateDoc)
|
||
|
||
update := bson.M{"$set": updateDoc}
|
||
// 使用MongoDAO.UpdateOne（不需要token验证）
|
||
_, _, err := MongoDAO.UpdateOne(ctx, filter, update, entity.RAGFlowConfigCollection)
|
||
if err != nil {
|
||
return gerror.Wrap(err, "更新RAGFlow配置失败")
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// UpdateDocumentIds 更新文档ID列表（避免双重token验证冲突）
|
||
func (d *ragflowConfig) UpdateDocumentIds(ctx context.Context, accountName string, documentIds []string) error {
|
||
filter := bson.M{"accountName": accountName, "isDeleted": false}
|
||
update := bson.M{"$set": bson.M{"documentIds": documentIds}}
|
||
// 使用MongoDAO.UpdateOne（不需要token验证）
|
||
_, _, err := MongoDAO.UpdateOne(ctx, filter, update, entity.RAGFlowConfigCollection)
|
||
return err
|
||
}
|
||
|
||
// UpdateDatasetIdByTenant 批量更新租户的所有datasetId记录（兼容tenantId类型不一致）
|
||
func (d *ragflowConfig) UpdateDatasetIdByTenant(ctx context.Context, tenantId, newDatasetId string) error {
|
||
// 先尝试字符串查询
|
||
filter := bson.M{"tenantId": tenantId, "isDeleted": false}
|
||
update := bson.M{"$set": bson.M{"datasetId": newDatasetId}}
|
||
matchedCount, _, err := MongoDAO.UpdateMany(ctx, filter, update, entity.RAGFlowConfigCollection)
|
||
|
||
// 如果未匹配到且tenantId可以转为数字，尝试用数字查询（兼容MongoDB中可能存储为int的情况）
|
||
if (err != nil || matchedCount == 0) && gconv.Int(tenantId) > 0 {
|
||
filter = bson.M{"tenantId": gconv.Int(tenantId), "isDeleted": false}
|
||
matchedCount, _, err = MongoDAO.UpdateMany(ctx, filter, update, entity.RAGFlowConfigCollection)
|
||
if err != nil {
|
||
return gerror.Wrap(err, "批量更新datasetId失败（数字类型尝试）")
|
||
}
|
||
}
|
||
|
||
if err != nil {
|
||
return gerror.Wrap(err, "批量更新datasetId失败")
|
||
}
|
||
|
||
if matchedCount == 0 {
|
||
return gerror.Newf("未找到租户%s的记录", tenantId)
|
||
}
|
||
|
||
return nil
|
||
}
|