mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-21 15:11:08 +08:00
10c00a9614
### Summary 1. list and get by id API for builtin DSL 2. add DSL default component param values validation 3. remove all hard code keys for parser config
3191 lines
102 KiB
Go
3191 lines
102 KiB
Go
//
|
|
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
|
|
package service
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"math"
|
|
"math/rand"
|
|
"ragflow/internal/common"
|
|
"ragflow/internal/dao"
|
|
"ragflow/internal/engine"
|
|
redisengine "ragflow/internal/engine/redis"
|
|
"ragflow/internal/engine/types"
|
|
enginetypes "ragflow/internal/engine/types"
|
|
"ragflow/internal/entity"
|
|
"ragflow/internal/entity/models"
|
|
pipelinepkg "ragflow/internal/ingestion/pipeline"
|
|
"ragflow/internal/service/nlp"
|
|
"ragflow/internal/utility"
|
|
"regexp"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/cespare/xxhash/v2"
|
|
"github.com/google/uuid"
|
|
"go.uber.org/zap"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
)
|
|
|
|
var (
|
|
datasetSupportedAvatarMIMETypes = map[string]struct{}{
|
|
"image/jpeg": {},
|
|
"image/png": {},
|
|
}
|
|
datasetAllowedOrderByFields = map[string]struct{}{
|
|
"create_time": {},
|
|
"update_time": {},
|
|
}
|
|
datasetAllowedMetadataTypes = map[string]struct{}{
|
|
"string": {},
|
|
"list": {},
|
|
"time": {},
|
|
"number": {},
|
|
}
|
|
validIndexTypes = []string{"graph", "raptor", "mindmap"}
|
|
indexTypeToTaskType = map[string]string{"graph": "graphrag", "raptor": "raptor", "mindmap": "mindmap"}
|
|
indexTypeToDisplayName = map[string]string{"graph": "Graph", "raptor": "RAPTOR", "mindmap": "Mindmap"}
|
|
)
|
|
|
|
const (
|
|
// Keep the legacy worker marker in queue payloads; persisted tasks use a real document ID.
|
|
graphRaptorQueueDocID = "graph_raptor_x"
|
|
maximumTaskPageNumber = int64(100000000)
|
|
serverQueueNamePrefix = "te"
|
|
defaultEmbeddingCheckNum = 5
|
|
|
|
graphPhaseResolutionDone = "resolution_done"
|
|
graphPhaseCommunityDone = "community_done"
|
|
)
|
|
|
|
// validateDatasetEmbeddingModels checks that all given datasets use the same
|
|
// embedding model (or all use none). Returns an error on mismatch.
|
|
func validateDatasetEmbeddingModels(kbs []*entity.Knowledgebase) error {
|
|
embdIDs := make(map[string]struct{})
|
|
hasEmbd := false
|
|
noEmbd := false
|
|
for _, kb := range kbs {
|
|
if kb.EmbdID != "" {
|
|
hasEmbd = true
|
|
baseName := kb.EmbdID
|
|
if idx := strings.LastIndex(kb.EmbdID, "@"); idx > 0 {
|
|
baseName = kb.EmbdID[:idx]
|
|
// Strip the second-to-last @-segment too (instance name),
|
|
// matching Python's _base_model_name which uses rsplit("@", 2).
|
|
if idx2 := strings.LastIndex(baseName, "@"); idx2 > 0 {
|
|
baseName = baseName[:idx2]
|
|
}
|
|
}
|
|
embdIDs[baseName] = struct{}{}
|
|
} else {
|
|
noEmbd = true
|
|
}
|
|
}
|
|
if hasEmbd && noEmbd {
|
|
return fmt.Errorf("Cannot search across datasets where some have embedding models and others do not.")
|
|
}
|
|
if len(embdIDs) > 1 {
|
|
return fmt.Errorf("Datasets use different embedding models: %v", getEmbdIDs(kbs))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// getEmbdIDs extracts embedding IDs from knowledge bases.
|
|
func getEmbdIDs(kbs []*entity.Knowledgebase) []string {
|
|
ids := make([]string, len(kbs))
|
|
for i, kb := range kbs {
|
|
ids[i] = kb.EmbdID
|
|
}
|
|
return ids
|
|
}
|
|
|
|
// DatasetService implements the RESTful dataset APIs from dataset_api.py.
|
|
type DatasetService struct {
|
|
kbDAO *dao.KnowledgebaseDAO
|
|
documentDAO *dao.DocumentDAO
|
|
connectorDAO *dao.ConnectorDAO
|
|
tenantDAO *dao.TenantDAO
|
|
tenantLLMDAO *dao.TenantLLMDAO
|
|
pipelineLogDAO *dao.PipelineOperationLogDAO
|
|
userTenantDAO *dao.UserTenantDAO
|
|
taskDAO *dao.TaskDAO
|
|
searchService *SearchService
|
|
docEngine engine.DocEngine
|
|
embeddingCache *utility.EmbeddingLRU
|
|
}
|
|
|
|
// NewDatasetService creates a new datasets service.
|
|
func NewDatasetService() *DatasetService {
|
|
return &DatasetService{
|
|
kbDAO: dao.NewKnowledgebaseDAO(),
|
|
documentDAO: dao.NewDocumentDAO(),
|
|
connectorDAO: dao.NewConnectorDAO(),
|
|
tenantDAO: dao.NewTenantDAO(),
|
|
tenantLLMDAO: dao.NewTenantLLMDAO(),
|
|
pipelineLogDAO: dao.NewPipelineOperationLogDAO(),
|
|
userTenantDAO: dao.NewUserTenantDAO(),
|
|
taskDAO: dao.NewTaskDAO(),
|
|
searchService: NewSearchService(),
|
|
docEngine: engine.Get(),
|
|
embeddingCache: utility.NewEmbeddingLRU(1000),
|
|
}
|
|
}
|
|
|
|
func (d *DatasetService) UpdateDocumentMetadataConfig(userID, datasetID, documentID string, req map[string]interface{}) (*entity.Document, common.ErrorCode, error) {
|
|
if _, err := d.kbDAO.GetByIDAndTenantID(datasetID, userID); err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeDataError, errors.New("You don't own the dataset.")
|
|
}
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
|
|
doc, err := d.documentDAO.GetByDocumentIDAndDatasetID(documentID, datasetID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeDataError, fmt.Errorf("Document %s not found in dataset %s", documentID, datasetID)
|
|
}
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
metadata, ok := req["metadata"]
|
|
if !ok {
|
|
return nil, common.CodeArgumentError, errors.New("metadata is required")
|
|
}
|
|
|
|
parserConfig := doc.ParserConfig
|
|
if parserConfig == nil {
|
|
parserConfig = entity.JSONMap{}
|
|
}
|
|
parserConfig["metadata"] = metadata
|
|
|
|
if err := d.documentDAO.UpdateByID(doc.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil {
|
|
return nil, common.CodeExceptionError, err
|
|
}
|
|
|
|
updatedDoc, err := d.documentDAO.GetByID(doc.ID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeDataError, errors.New("Document not found!")
|
|
}
|
|
return nil, common.CodeExceptionError, err
|
|
}
|
|
|
|
return updatedDoc, common.CodeSuccess, nil
|
|
}
|
|
|
|
// checkType reports whether indexType is supported by dataset index APIs.
|
|
func checkType(indexType string) bool {
|
|
haveType := false
|
|
for _, t := range validIndexTypes {
|
|
if indexType == t {
|
|
haveType = true
|
|
}
|
|
}
|
|
return haveType
|
|
}
|
|
|
|
func (d *DatasetService) newRaptorOrGraphRagTask(sampleDoc *entity.Document, taskType string, taskDocID string, queueDocID string, docIDs []string) (*entity.Task, map[string]interface{}, error) {
|
|
if docIDs == nil || len(docIDs) == 0 {
|
|
docIDs = make([]string, 0)
|
|
}
|
|
if !checkIndexTaskType(taskType) {
|
|
return nil, nil, errors.New("type should be graphrag, raptor or mindmap")
|
|
}
|
|
|
|
chunkingConfig, err := d.documentDAO.GetChunkingConfig(sampleDoc.ID)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
hasher := xxhash.New()
|
|
keys := make([]string, 0, len(chunkingConfig))
|
|
for key := range chunkingConfig {
|
|
keys = append(keys, key)
|
|
}
|
|
sort.Strings(keys)
|
|
for _, key := range keys {
|
|
_, _ = hasher.Write([]byte(key))
|
|
_, _ = hasher.Write([]byte{0})
|
|
v, mErr := json.Marshal(chunkingConfig[key])
|
|
if mErr != nil {
|
|
return nil, nil, mErr
|
|
}
|
|
_, _ = hasher.Write(v)
|
|
_, _ = hasher.Write([]byte{0})
|
|
}
|
|
|
|
taskID := utility.GenerateUUID()
|
|
beginAt := time.Now().Truncate(time.Second)
|
|
progressMsg := beginAt.Format("15:04:05") + " created task " + taskType
|
|
|
|
for _, field := range []interface{}{taskDocID, maximumTaskPageNumber, maximumTaskPageNumber, taskType} {
|
|
_, _ = hasher.Write([]byte(fmt.Sprint(field)))
|
|
}
|
|
digest := fmt.Sprintf("%016x", hasher.Sum64())
|
|
task := &entity.Task{
|
|
ID: taskID,
|
|
DocID: taskDocID,
|
|
FromPage: maximumTaskPageNumber,
|
|
ToPage: maximumTaskPageNumber,
|
|
TaskType: taskType,
|
|
ProgressMsg: &progressMsg,
|
|
BeginAt: &beginAt,
|
|
Digest: &digest,
|
|
}
|
|
|
|
queueMessage := map[string]interface{}{
|
|
"id": taskID,
|
|
"doc_id": queueDocID,
|
|
"from_page": maximumTaskPageNumber,
|
|
"to_page": maximumTaskPageNumber,
|
|
"task_type": taskType,
|
|
"progress_msg": progressMsg,
|
|
"begin_at": beginAt.Format("2006-01-02 15:04:05"),
|
|
"digest": digest,
|
|
"doc_ids": docIDs,
|
|
}
|
|
|
|
return task, queueMessage, nil
|
|
}
|
|
|
|
func createDatasetIndexTaskInTx(tx *gorm.DB, task *entity.Task, queueDocID string) (*entity.Document, error) {
|
|
if task == nil {
|
|
return nil, errors.New("task is required")
|
|
}
|
|
if err := tx.Create(task).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if queueDocID == "" {
|
|
return nil, nil
|
|
}
|
|
|
|
var document entity.Document
|
|
err := tx.Select("id", "progress_msg", "process_begin_at").Where("id = ?", queueDocID).First(&document).Error
|
|
if err != nil {
|
|
if errors.Is(err, gorm.ErrRecordNotFound) {
|
|
return nil, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
|
|
beginAt := time.Now().Truncate(time.Second)
|
|
if task.BeginAt != nil {
|
|
beginAt = *task.BeginAt
|
|
}
|
|
if err := tx.Model(&entity.Document{}).Where("id = ?", queueDocID).Updates(map[string]interface{}{
|
|
"progress_msg": "Task is queued...",
|
|
"process_begin_at": beginAt,
|
|
}).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &document, nil
|
|
}
|
|
|
|
func enqueueDatasetIndexTask(priority int, queueMessage map[string]interface{}) error {
|
|
redisClient := redisengine.Get()
|
|
if redisClient == nil || !redisClient.QueueProduct(datasetIndexQueueName(priority), queueMessage) {
|
|
return errors.New("Can't access Redis. Please check the Redis' status")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func cleanupFailedDatasetIndexTask(taskID string, updatedDocument *entity.Document, kbID string, indexType string) error {
|
|
return dao.DB.Transaction(func(tx *gorm.DB) error {
|
|
if err := tx.Unscoped().Where("id = ?", taskID).Delete(&entity.Task{}).Error; err != nil {
|
|
return fmt.Errorf("delete task %s: %w", taskID, err)
|
|
}
|
|
|
|
if column := datasetIndexTaskIDColumn(indexType); kbID != "" && column != "" {
|
|
if err := tx.Model(&entity.Knowledgebase{}).Where("id = ? AND "+column+" = ?", kbID, taskID).Update(column, nil).Error; err != nil {
|
|
return fmt.Errorf("clear dataset task id %s: %w", taskID, err)
|
|
}
|
|
}
|
|
|
|
if updatedDocument == nil {
|
|
return nil
|
|
}
|
|
|
|
return tx.Model(&entity.Document{}).Where("id = ?", updatedDocument.ID).Updates(map[string]interface{}{
|
|
"progress_msg": updatedDocument.ProgressMsg,
|
|
"process_begin_at": updatedDocument.ProcessBeginAt,
|
|
}).Error
|
|
})
|
|
}
|
|
|
|
func datasetIndexTaskIDColumn(indexType string) string {
|
|
switch indexType {
|
|
case "graph":
|
|
return "graphrag_task_id"
|
|
case "raptor":
|
|
return "raptor_task_id"
|
|
case "mindmap":
|
|
return "mindmap_task_id"
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func datasetIndexTaskFinishAtColumn(indexType string) string {
|
|
switch indexType {
|
|
case "graph":
|
|
return "graphrag_task_finish_at"
|
|
case "raptor":
|
|
return "raptor_task_finish_at"
|
|
case "mindmap":
|
|
return "mindmap_task_finish_at"
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func checkIndexTaskType(taskType string) bool {
|
|
switch taskType {
|
|
case "graphrag", "raptor", "mindmap":
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func datasetIndexTaskID(kb *entity.Knowledgebase, indexType string) string {
|
|
if kb == nil {
|
|
return ""
|
|
}
|
|
switch indexType {
|
|
case "graph":
|
|
if kb.GraphragTaskID != nil {
|
|
return *kb.GraphragTaskID
|
|
}
|
|
case "raptor":
|
|
if kb.RaptorTaskID != nil {
|
|
return *kb.RaptorTaskID
|
|
}
|
|
case "mindmap":
|
|
if kb.MindmapTaskID != nil {
|
|
return *kb.MindmapTaskID
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func datasetIndexTaskIDUpdate(indexType, taskID string) map[string]interface{} {
|
|
switch indexType {
|
|
case "graph":
|
|
return map[string]interface{}{"graphrag_task_id": taskID}
|
|
case "raptor":
|
|
return map[string]interface{}{"raptor_task_id": taskID}
|
|
case "mindmap":
|
|
return map[string]interface{}{"mindmap_task_id": taskID}
|
|
default:
|
|
return map[string]interface{}{}
|
|
}
|
|
}
|
|
|
|
func datasetIndexTaskIDs(kb *entity.Knowledgebase) []string {
|
|
if kb == nil {
|
|
return nil
|
|
}
|
|
|
|
taskIDs := make([]string, 0, 3)
|
|
for _, taskID := range []*string{kb.GraphragTaskID, kb.RaptorTaskID, kb.MindmapTaskID} {
|
|
if taskID != nil && *taskID != "" {
|
|
taskIDs = append(taskIDs, *taskID)
|
|
}
|
|
}
|
|
return common.Deduplicate(taskIDs)
|
|
}
|
|
|
|
func datasetIndexQueueName(priority int) string {
|
|
return fmt.Sprintf("%s.%d.common", serverQueueNamePrefix, priority)
|
|
}
|
|
|
|
func interfaceSlice(items ...string) []interface{} {
|
|
result := make([]interface{}, len(items))
|
|
for i, item := range items {
|
|
result[i] = item
|
|
}
|
|
return result
|
|
}
|
|
|
|
func clearGraphPhaseMarkers(redisClient *redisengine.Client, datasetID string) {
|
|
if redisClient == nil || datasetID == "" {
|
|
return
|
|
}
|
|
for _, phase := range []string{graphPhaseResolutionDone, graphPhaseCommunityDone} {
|
|
if !redisClient.Delete(fmt.Sprintf("graphrag:phase:%s:%s", datasetID, phase)) {
|
|
common.Warn("Failed to clear GraphRAG phase marker", zap.String("dataset_id", datasetID), zap.String("phase", phase))
|
|
}
|
|
}
|
|
}
|
|
|
|
// RunIndex Run an indexing task (graph/raptor/mindmap) for a dataset.
|
|
func (d *DatasetService) RunIndex(userID, datasetID, indexType string) (map[string]interface{}, common.ErrorCode, error) {
|
|
if !checkType(indexType) {
|
|
return nil, common.CodeDataError, fmt.Errorf("Invalid index type '%s'. Must be one of %v", indexType, validIndexTypes)
|
|
}
|
|
|
|
if datasetID == "" {
|
|
return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`)
|
|
}
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return nil, common.CodeDataError, errors.New("No authorization.")
|
|
}
|
|
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
|
|
}
|
|
return nil, common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
|
|
taskType := indexTypeToTaskType[indexType]
|
|
displayName := indexTypeToDisplayName[indexType]
|
|
|
|
documents, code, err := d.getDocumentsByDatasetForIndex(datasetID)
|
|
if err != nil {
|
|
return nil, code, err
|
|
}
|
|
_ = documents
|
|
|
|
sampleDocument := documents[0]
|
|
documentIDs := make([]string, len(documents))
|
|
|
|
for i, doc := range documents {
|
|
documentIDs[i] = doc.ID
|
|
}
|
|
|
|
task, queueMessage, err := d.newRaptorOrGraphRagTask(sampleDocument, taskType, sampleDocument.ID, graphRaptorQueueDocID, documentIDs)
|
|
if err != nil {
|
|
common.Warn("Failed to build dataset index task", zap.String("dataset_id", datasetID), zap.String("task_type", taskType), zap.Error(err))
|
|
return nil, common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
|
|
var updatedDocument *entity.Document
|
|
var dataErr error
|
|
err = dao.DB.Transaction(func(tx *gorm.DB) error {
|
|
var lockedKB entity.Knowledgebase
|
|
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
|
Where("id = ? AND status = ?", kb.ID, string(entity.StatusValid)).
|
|
First(&lockedKB).Error; err != nil {
|
|
return err
|
|
}
|
|
|
|
existingTaskID := datasetIndexTaskID(&lockedKB, indexType)
|
|
if existingTaskID != "" {
|
|
var existingTask entity.Task
|
|
taskErr := tx.Where("id = ?", existingTaskID).First(&existingTask).Error
|
|
if taskErr != nil {
|
|
if errors.Is(taskErr, gorm.ErrRecordNotFound) {
|
|
} else {
|
|
return taskErr
|
|
}
|
|
} else if existingTask.Progress != 1 && existingTask.Progress != -1 {
|
|
dataErr = fmt.Errorf("Task %s in progress with status %v. A %s Task is already running.", existingTaskID, existingTask.Progress, displayName)
|
|
return dataErr
|
|
}
|
|
}
|
|
|
|
updatedDocument, err = createDatasetIndexTaskInTx(tx, task, graphRaptorQueueDocID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return tx.Model(&entity.Knowledgebase{}).Where("id = ?", lockedKB.ID).Updates(datasetIndexTaskIDUpdate(indexType, task.ID)).Error
|
|
})
|
|
if err != nil {
|
|
if dataErr != nil {
|
|
return nil, common.CodeDataError, dataErr
|
|
}
|
|
common.Warn("Failed to create dataset index task", zap.String("dataset_id", datasetID), zap.String("task_type", taskType), zap.Error(err))
|
|
return nil, common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
|
|
if err := enqueueDatasetIndexTask(0, queueMessage); err != nil {
|
|
if cleanupErr := cleanupFailedDatasetIndexTask(task.ID, updatedDocument, kb.ID, indexType); cleanupErr != nil {
|
|
err = errors.Join(err, cleanupErr)
|
|
}
|
|
common.Warn("Failed to queue dataset index task", zap.String("dataset_id", datasetID), zap.String("task_type", taskType), zap.Error(err))
|
|
return nil, common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
|
|
return map[string]interface{}{"task_id": task.ID}, common.CodeSuccess, nil
|
|
}
|
|
|
|
func (d *DatasetService) getDocumentsByDatasetForIndex(datasetID string) ([]*entity.Document, common.ErrorCode, error) {
|
|
documents, _, err := d.documentDAO.GetByKBID(datasetID)
|
|
if err != nil {
|
|
common.Warn("Failed to load dataset documents for index", zap.String("dataset_id", datasetID), zap.Error(err))
|
|
return nil, common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
if len(documents) == 0 {
|
|
return nil, common.CodeDataError, fmt.Errorf("No documents in Dataset %s", datasetID)
|
|
}
|
|
return documents, common.CodeSuccess, nil
|
|
}
|
|
|
|
type TraceIndexRequest struct {
|
|
Type string `json:"type" binding:"required"`
|
|
}
|
|
|
|
// TraceIndex Trace an indexing task (graph/raptor/mindmap) for a dataset.
|
|
func (d *DatasetService) TraceIndex(datasetID, userID, indexType string) (*entity.Task, common.ErrorCode, error) {
|
|
if !checkType(indexType) {
|
|
return nil, common.CodeDataError, fmt.Errorf("Invalid index type '%s'. Must be one of %v", indexType, validIndexTypes)
|
|
}
|
|
|
|
if datasetID == "" {
|
|
return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`)
|
|
}
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return nil, common.CodeDataError, errors.New("No authorization.")
|
|
}
|
|
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
|
|
}
|
|
return nil, common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
|
|
taskID := datasetIndexTaskID(kb, indexType)
|
|
|
|
var task *entity.Task
|
|
if taskID != "" {
|
|
task, err = d.taskDAO.GetByID(taskID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeSuccess, nil
|
|
}
|
|
return nil, common.CodeServerError, errors.New("Internal server error")
|
|
}
|
|
if task == nil {
|
|
return nil, common.CodeSuccess, nil
|
|
}
|
|
}
|
|
|
|
return task, common.CodeSuccess, nil
|
|
}
|
|
|
|
type CheckEmbeddingRequest struct {
|
|
EmbeddingID string `json:"embd_id" binding:"required"`
|
|
CheckNum *int `json:"check_num,omitempty"`
|
|
}
|
|
|
|
type EmbeddingCheckSummary struct {
|
|
KbID string `json:"kb_id"`
|
|
Model string `json:"model"`
|
|
Sampled int `json:"sampled"`
|
|
Valid int `json:"valid"`
|
|
AvgCosSim float64 `json:"avg_cos_sim"`
|
|
MinCosSim float64 `json:"min_cos_sim"`
|
|
MaxCosSim float64 `json:"max_cos_sim"`
|
|
MatchMode string `json:"match_mode"`
|
|
}
|
|
|
|
type EmbeddingCheckResult struct {
|
|
ChunkID string `json:"chunk_id"`
|
|
DocID string `json:"doc_id,omitempty"`
|
|
DocName string `json:"doc_name,omitempty"`
|
|
VectorField string `json:"vector_field,omitempty"`
|
|
VectorDim int `json:"vector_dim,omitempty"`
|
|
CosSim float64 `json:"cos_sim,omitempty"`
|
|
Reason string `json:"reason,omitempty"`
|
|
}
|
|
|
|
type EmbeddingCheckResponse struct {
|
|
Summary EmbeddingCheckSummary `json:"summary"`
|
|
Results []EmbeddingCheckResult `json:"results"`
|
|
}
|
|
|
|
type embeddingCheckSample struct {
|
|
ChunkID string
|
|
KbID string
|
|
DocID string
|
|
DocName string
|
|
VectorField string
|
|
Vector []float64
|
|
PageNum interface{}
|
|
Position interface{}
|
|
Top interface{}
|
|
ContentWithWeight string
|
|
QuestionKeywords []string
|
|
}
|
|
|
|
// CheckEmbedding checks whether a new embedding model is compatible with stored vectors.
|
|
func (d *DatasetService) CheckEmbedding(userID, datasetID string, req *CheckEmbeddingRequest) (*EmbeddingCheckResponse, common.ErrorCode, error) {
|
|
if datasetID == "" {
|
|
return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`)
|
|
}
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return nil, common.CodeDataError, errors.New("No authorization.")
|
|
}
|
|
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
|
|
}
|
|
return nil, common.CodeServerError, errors.New("Internal server error")
|
|
}
|
|
|
|
if req == nil || strings.TrimSpace(req.EmbeddingID) == "" {
|
|
return nil, common.CodeDataError, errors.New("`embd_id` is required.")
|
|
}
|
|
embeddingID := strings.TrimSpace(req.EmbeddingID)
|
|
if ok, message := d.verifyEmbeddingAvailability(embeddingID, userID); !ok {
|
|
return nil, common.CodeDataError, errors.New(message)
|
|
}
|
|
if d.docEngine == nil {
|
|
return nil, common.CodeServerError, errors.New("doc engine not initialized")
|
|
}
|
|
|
|
driver, modelName, apiConfig, maxTokens, err := NewModelProviderService().ResolveModelConfig(kb.TenantID, entity.ModelTypeEmbedding, embeddingID)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
embeddingModel := models.NewEmbeddingModel(driver, &modelName, apiConfig, maxTokens)
|
|
|
|
checkNum := defaultEmbeddingCheckNum
|
|
if req.CheckNum != nil {
|
|
checkNum = *req.CheckNum
|
|
}
|
|
if checkNum <= 0 {
|
|
checkNum = defaultEmbeddingCheckNum
|
|
}
|
|
|
|
samples, err := d.sampleRandomChunksWithVectors(context.Background(), kb.TenantID, datasetID, checkNum)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
|
|
results := make([]EmbeddingCheckResult, 0, len(samples))
|
|
effectiveSimilarities := make([]float64, 0, len(samples))
|
|
matchMode := "content_only"
|
|
for _, sample := range samples {
|
|
title := sample.DocName
|
|
if strings.TrimSpace(title) == "" {
|
|
title = "Title"
|
|
}
|
|
|
|
textInput := strings.Join(sample.QuestionKeywords, "\n")
|
|
if strings.TrimSpace(textInput) == "" {
|
|
textInput = sample.ContentWithWeight
|
|
}
|
|
textInput = datasetCleanEmbeddingText(textInput)
|
|
if textInput == "" {
|
|
results = append(results, EmbeddingCheckResult{ChunkID: sample.ChunkID, Reason: "no_text"})
|
|
continue
|
|
}
|
|
if len(sample.Vector) == 0 {
|
|
results = append(results, EmbeddingCheckResult{ChunkID: sample.ChunkID, Reason: "no_stored_vector"})
|
|
continue
|
|
}
|
|
|
|
vectors, err := datasetEncodeEmbedding(embeddingModel, []string{title, textInput})
|
|
if err != nil {
|
|
return nil, common.CodeDataError, fmt.Errorf("Embedding failure. %w", err)
|
|
}
|
|
if len(vectors) < 2 {
|
|
return nil, common.CodeDataError, errors.New("Embedding failure. embedding response is incomplete")
|
|
}
|
|
if len(vectors[1]) != len(sample.Vector) {
|
|
return nil, common.CodeDataError, fmt.Errorf("Embedding failure. The dimension (%d) of given embedding model is different from the original (%d)", len(vectors[1]), len(sample.Vector))
|
|
}
|
|
|
|
simContent := datasetCosSim(vectors[1], sample.Vector)
|
|
simMix := datasetCosSim(datasetMixVectors(vectors[0], vectors[1], 0.1), sample.Vector)
|
|
sim := simContent
|
|
matchMode = "content_only"
|
|
if simMix > sim {
|
|
sim = simMix
|
|
matchMode = "title+content"
|
|
}
|
|
sim = datasetRoundFloat(sim, 6)
|
|
|
|
effectiveSimilarities = append(effectiveSimilarities, sim)
|
|
results = append(results, EmbeddingCheckResult{
|
|
ChunkID: sample.ChunkID,
|
|
DocID: sample.DocID,
|
|
DocName: sample.DocName,
|
|
VectorField: sample.VectorField,
|
|
VectorDim: len(sample.Vector),
|
|
CosSim: sim,
|
|
})
|
|
}
|
|
|
|
summary := datasetEmbeddingCheckSummary(datasetID, embeddingID, len(samples), effectiveSimilarities, matchMode)
|
|
response := &EmbeddingCheckResponse{Summary: summary, Results: results}
|
|
if len(effectiveSimilarities) == 0 {
|
|
return nil, common.CodeDataError, errors.New("No embedded chunks are available to compare.")
|
|
}
|
|
if summary.AvgCosSim >= 0.9 {
|
|
return response, common.CodeSuccess, nil
|
|
}
|
|
return response, common.CodeNotEffective, errors.New("Embedding model switch failed: the average similarity between old and new vectors is below 0.9, indicating incompatible vector spaces.")
|
|
}
|
|
|
|
func (d *DatasetService) sampleRandomChunksWithVectors(ctx context.Context, tenantID, datasetID string, n int) ([]embeddingCheckSample, error) {
|
|
indexName := fmt.Sprintf("ragflow_%s", tenantID)
|
|
totalResult, err := d.docEngine.Search(ctx, &enginetypes.SearchRequest{
|
|
IndexNames: []string{indexName},
|
|
KbIDs: []string{datasetID},
|
|
Offset: 0,
|
|
Limit: 1,
|
|
Filter: map[string]interface{}{
|
|
"kb_id": datasetID,
|
|
"available_int": 1,
|
|
},
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if totalResult == nil || totalResult.Total <= 0 {
|
|
return []embeddingCheckSample{}, nil
|
|
}
|
|
|
|
total := int(totalResult.Total)
|
|
// Cap n to a sane upper bound so a hostile caller can't force a
|
|
// huge preallocation. The downstream `samples` slice is sized
|
|
// directly from n.
|
|
const maxEmbeddingSamples = 1024
|
|
if n < 0 {
|
|
return nil, fmt.Errorf("invalid sample size: %d", n)
|
|
}
|
|
if n > maxEmbeddingSamples {
|
|
n = maxEmbeddingSamples
|
|
}
|
|
if n > total {
|
|
n = total
|
|
}
|
|
limit := total
|
|
if limit > 1000 {
|
|
limit = 1000
|
|
}
|
|
if n > limit {
|
|
n = limit
|
|
}
|
|
offsets := rand.Perm(limit)
|
|
offsets = offsets[:n]
|
|
sort.Ints(offsets)
|
|
|
|
baseFields := []string{"docnm_kwd", "doc_id", "content_with_weight", "page_num_int", "position_int", "top_int"}
|
|
// codeql[go/uncontrolled-allocation-size] False positive: n is
|
|
// bounded to maxEmbeddingSamples (1024) at the top of this
|
|
// function, so the samples slice cannot exceed ~1 MiB
|
|
// (embeddingCheckSample is a small struct).
|
|
samples := make([]embeddingCheckSample, 0, n)
|
|
for _, offset := range offsets {
|
|
searchResult, err := d.docEngine.Search(ctx, &enginetypes.SearchRequest{
|
|
IndexNames: []string{indexName},
|
|
KbIDs: []string{datasetID},
|
|
Offset: offset,
|
|
Limit: 1,
|
|
SelectFields: baseFields,
|
|
Filter: map[string]interface{}{
|
|
"kb_id": datasetID,
|
|
"available_int": 1,
|
|
},
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if searchResult == nil || len(searchResult.Chunks) == 0 {
|
|
continue
|
|
}
|
|
chunkID := datasetChunkID(searchResult.Chunks[0])
|
|
if chunkID == "" {
|
|
continue
|
|
}
|
|
fullChunk, err := d.docEngine.GetChunk(ctx, indexName, chunkID, []string{datasetID})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
chunkMap := datasetMap(fullChunk)
|
|
if len(chunkMap) == 0 {
|
|
continue
|
|
}
|
|
vectorField := datasetGuessVecField(chunkMap)
|
|
vector := datasetAsFloatVec(chunkMap[vectorField])
|
|
samples = append(samples, embeddingCheckSample{
|
|
ChunkID: chunkID,
|
|
KbID: datasetID,
|
|
DocID: datasetString(chunkMap["doc_id"]),
|
|
DocName: datasetString(chunkMap["docnm_kwd"]),
|
|
VectorField: vectorField,
|
|
Vector: vector,
|
|
PageNum: chunkMap["page_num_int"],
|
|
Position: chunkMap["position_int"],
|
|
Top: chunkMap["top_int"],
|
|
ContentWithWeight: datasetString(chunkMap["content_with_weight"]),
|
|
QuestionKeywords: datasetStringSlice(chunkMap["question_kwd"]),
|
|
})
|
|
}
|
|
return samples, nil
|
|
}
|
|
|
|
func datasetGuessVecField(src map[string]interface{}) string {
|
|
for k := range src {
|
|
if strings.HasSuffix(k, "_vec") {
|
|
return k
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func datasetAsFloatVec(v interface{}) []float64 {
|
|
if v == nil {
|
|
return []float64{}
|
|
}
|
|
switch val := v.(type) {
|
|
case string:
|
|
parts := strings.Split(val, "\t")
|
|
res := make([]float64, 0, len(parts))
|
|
for _, p := range parts {
|
|
if p == "" {
|
|
continue
|
|
}
|
|
f, err := strconv.ParseFloat(p, 64)
|
|
if err != nil {
|
|
continue
|
|
}
|
|
res = append(res, f)
|
|
}
|
|
return res
|
|
case []float64:
|
|
return val
|
|
case []float32:
|
|
res := make([]float64, len(val))
|
|
for i, x := range val {
|
|
res[i] = float64(x)
|
|
}
|
|
return res
|
|
case []int:
|
|
res := make([]float64, len(val))
|
|
for i, x := range val {
|
|
res[i] = float64(x)
|
|
}
|
|
return res
|
|
case []interface{}:
|
|
res := make([]float64, 0, len(val))
|
|
for _, x := range val {
|
|
switch n := x.(type) {
|
|
case float64:
|
|
res = append(res, n)
|
|
case float32:
|
|
res = append(res, float64(n))
|
|
case int:
|
|
res = append(res, float64(n))
|
|
case string:
|
|
f, err := strconv.ParseFloat(n, 64)
|
|
if err == nil {
|
|
res = append(res, f)
|
|
}
|
|
}
|
|
}
|
|
return res
|
|
}
|
|
return []float64{}
|
|
}
|
|
|
|
func datasetCosSim(a, b []float64) float64 {
|
|
if len(a) == 0 || len(b) == 0 {
|
|
return 0
|
|
}
|
|
var dot, na, nb float64
|
|
n := len(a)
|
|
if len(b) < n {
|
|
n = len(b)
|
|
}
|
|
for i := 0; i < n; i++ {
|
|
dot += a[i] * b[i]
|
|
}
|
|
for _, x := range a {
|
|
na += x * x
|
|
}
|
|
for _, x := range b {
|
|
nb += x * x
|
|
}
|
|
|
|
if na == 0 || nb == 0 {
|
|
return 0
|
|
}
|
|
return dot / (math.Sqrt(na) * math.Sqrt(nb))
|
|
}
|
|
|
|
func datasetCleanEmbeddingText(s string) string {
|
|
re := regexp.MustCompile(`</?(table|td|caption|tr|th)( [^<>]{0,12})?>`)
|
|
return strings.TrimSpace(re.ReplaceAllString(s, " "))
|
|
}
|
|
|
|
func datasetEncodeEmbedding(embeddingModel *models.EmbeddingModel, texts []string) ([][]float64, error) {
|
|
embeddingConfig := &models.EmbeddingConfig{Dimension: 0}
|
|
embeddings, err := embeddingModel.ModelDriver.Embed(embeddingModel.ModelName, texts, embeddingModel.APIConfig, embeddingConfig)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
vectors := make([][]float64, len(embeddings))
|
|
for i, embedding := range embeddings {
|
|
vectors[i] = embedding.Embedding
|
|
}
|
|
return vectors, nil
|
|
}
|
|
|
|
func datasetMixVectors(titleVector, contentVector []float64, titleWeight float64) []float64 {
|
|
if len(titleVector) != len(contentVector) {
|
|
return contentVector
|
|
}
|
|
mixed := make([]float64, len(contentVector))
|
|
contentWeight := 1 - titleWeight
|
|
for i := range contentVector {
|
|
mixed[i] = titleWeight*titleVector[i] + contentWeight*contentVector[i]
|
|
}
|
|
return mixed
|
|
}
|
|
|
|
func datasetEmbeddingCheckSummary(datasetID, embeddingID string, sampled int, similarities []float64, matchMode string) EmbeddingCheckSummary {
|
|
summary := EmbeddingCheckSummary{
|
|
KbID: datasetID,
|
|
Model: embeddingID,
|
|
Sampled: sampled,
|
|
Valid: len(similarities),
|
|
MatchMode: matchMode,
|
|
}
|
|
if len(similarities) == 0 {
|
|
return summary
|
|
}
|
|
minValue := similarities[0]
|
|
maxValue := similarities[0]
|
|
total := 0.0
|
|
for _, value := range similarities {
|
|
total += value
|
|
if value < minValue {
|
|
minValue = value
|
|
}
|
|
if value > maxValue {
|
|
maxValue = value
|
|
}
|
|
}
|
|
summary.AvgCosSim = datasetRoundFloat(total/float64(len(similarities)), 6)
|
|
summary.MinCosSim = datasetRoundFloat(minValue, 6)
|
|
summary.MaxCosSim = datasetRoundFloat(maxValue, 6)
|
|
return summary
|
|
}
|
|
|
|
func datasetRoundFloat(value float64, places int) float64 {
|
|
factor := math.Pow10(places)
|
|
return math.Round(value*factor) / factor
|
|
}
|
|
|
|
func datasetChunkID(chunk map[string]interface{}) string {
|
|
for _, key := range []string{"id", "_id"} {
|
|
if value := datasetString(chunk[key]); value != "" {
|
|
return value
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func datasetMap(value interface{}) map[string]interface{} {
|
|
switch typedValue := value.(type) {
|
|
case map[string]interface{}:
|
|
return typedValue
|
|
default:
|
|
return map[string]interface{}{}
|
|
}
|
|
}
|
|
|
|
func datasetString(value interface{}) string {
|
|
switch typedValue := value.(type) {
|
|
case string:
|
|
return typedValue
|
|
case fmt.Stringer:
|
|
return typedValue.String()
|
|
case nil:
|
|
return ""
|
|
default:
|
|
return fmt.Sprint(typedValue)
|
|
}
|
|
}
|
|
|
|
func datasetStringSlice(value interface{}) []string {
|
|
switch typedValue := value.(type) {
|
|
case []string:
|
|
return typedValue
|
|
case []interface{}:
|
|
values := make([]string, 0, len(typedValue))
|
|
for _, item := range typedValue {
|
|
if s := strings.TrimSpace(datasetString(item)); s != "" {
|
|
values = append(values, s)
|
|
}
|
|
}
|
|
return values
|
|
case string:
|
|
if typedValue == "" {
|
|
return nil
|
|
}
|
|
return []string{typedValue}
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (d *DatasetService) DeleteIndex(userID, datasetID, indexType string, wipe bool) (common.ErrorCode, error) {
|
|
if !checkType(indexType) {
|
|
return common.CodeArgumentError, fmt.Errorf("Invalid index type '%s'", indexType)
|
|
}
|
|
|
|
if datasetID == "" {
|
|
return common.CodeDataError, errors.New(`Lack of "Dataset ID"`)
|
|
}
|
|
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return common.CodeDataError, errors.New("No authorization.")
|
|
}
|
|
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return common.CodeDataError, errors.New("Invalid Dataset ID")
|
|
}
|
|
return common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
|
|
taskIDField := datasetIndexTaskIDColumn(indexType)
|
|
taskFinishAtField := datasetIndexTaskFinishAtColumn(indexType)
|
|
taskID := datasetIndexTaskID(kb, indexType)
|
|
|
|
common.Info("delete_index", zap.String("dataset_id", datasetID), zap.String("index_type", indexType), zap.Bool("wipe", wipe))
|
|
|
|
if taskID != "" {
|
|
redisClient := redisengine.Get()
|
|
if redisClient == nil || !redisClient.Set(fmt.Sprintf("%s-cancel", taskID), "x", 0) {
|
|
common.Warn("Failed to set dataset index cancellation marker", zap.String("dataset_id", datasetID), zap.String("task_id", taskID))
|
|
}
|
|
if err := dao.DB.Unscoped().Where("id = ?", taskID).Delete(&entity.Task{}).Error; err != nil {
|
|
common.Warn("Failed to delete dataset index task", zap.String("dataset_id", datasetID), zap.String("task_id", taskID), zap.Error(err))
|
|
return common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
}
|
|
|
|
if wipe && indexType == "graph" {
|
|
if d.docEngine == nil {
|
|
return common.CodeServerError, errors.New("Document engine is not initialized")
|
|
}
|
|
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
|
|
_, err = d.docEngine.DeleteChunks(context.Background(), map[string]interface{}{
|
|
"knowledge_graph_kwd": interfaceSlice("graph", "subgraph", "entity", "relation", "community_report"),
|
|
"kb_id": datasetID,
|
|
}, indexName, datasetID)
|
|
if err != nil {
|
|
common.Warn("Failed to delete GraphRAG artefacts", zap.String("dataset_id", datasetID), zap.Error(err))
|
|
return common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
clearGraphPhaseMarkers(redisengine.Get(), datasetID)
|
|
common.Info("delete_index: cleared GraphRAG artefacts and phase markers", zap.String("dataset_id", datasetID))
|
|
} else if wipe && indexType == "raptor" {
|
|
if d.docEngine == nil {
|
|
return common.CodeServerError, errors.New("Document engine is not initialized")
|
|
}
|
|
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
|
|
_, err = d.docEngine.DeleteChunks(context.Background(), map[string]interface{}{
|
|
"raptor_kwd": interfaceSlice("raptor"),
|
|
"kb_id": datasetID,
|
|
}, indexName, datasetID)
|
|
if err != nil {
|
|
common.Warn("Failed to delete RAPTOR artefacts", zap.String("dataset_id", datasetID), zap.Error(err))
|
|
return common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
}
|
|
|
|
if err := dao.DB.Model(&entity.Knowledgebase{}).Where("id = ?", kb.ID).Updates(map[string]interface{}{
|
|
taskIDField: "",
|
|
taskFinishAtField: nil,
|
|
}).Error; err != nil {
|
|
common.Warn("Failed to clear dataset index task fields", zap.String("dataset_id", datasetID), zap.String("index_type", indexType), zap.Error(err))
|
|
return common.CodeDataError, errors.New("Internal server error")
|
|
}
|
|
|
|
return common.CodeSuccess, nil
|
|
}
|
|
|
|
// SearchDatasetsRequest is the request structure for searching chunks across datasets.
|
|
type SearchDatasetsRequest struct {
|
|
DatasetIDs []string `json:"dataset_ids" binding:"required"`
|
|
Question string `json:"question" binding:"required"`
|
|
Page *int `json:"page,omitempty"`
|
|
Size *int `json:"size,omitempty"`
|
|
DocIDs []string `json:"doc_ids,omitempty"`
|
|
UseKG *bool `json:"use_kg,omitempty"`
|
|
TopK *int `json:"top_k,omitempty"`
|
|
CrossLanguages []string `json:"cross_languages,omitempty"`
|
|
SearchID *string `json:"search_id,omitempty"`
|
|
MetadataFilter map[string]interface{} `json:"meta_data_filter,omitempty"`
|
|
RerankID *string `json:"rerank_id,omitempty"`
|
|
Keyword *bool `json:"keyword,omitempty"`
|
|
SimilarityThreshold *float64 `json:"similarity_threshold,omitempty"`
|
|
VectorSimilarityWeight *float64 `json:"vector_similarity_weight,omitempty"`
|
|
ForceRefresh bool `json:"force_refresh"`
|
|
}
|
|
|
|
// SearchDatasetsResponse is the response structure for dataset search results.
|
|
type SearchDatasetsResponse struct {
|
|
Chunks []map[string]interface{} `json:"chunks"`
|
|
DocAggs []map[string]interface{} `json:"doc_aggs"`
|
|
Labels *map[string]float64 `json:"labels"`
|
|
Total int64 `json:"total"`
|
|
}
|
|
|
|
// SearchDatasetRequest is the request structure for searching chunks within one dataset.
|
|
type SearchDatasetRequest struct {
|
|
Question string `json:"question"`
|
|
Page *int `json:"page,omitempty"`
|
|
Size *int `json:"size,omitempty"`
|
|
DocIDs []string `json:"doc_ids,omitempty"`
|
|
UseKG *bool `json:"use_kg,omitempty"`
|
|
TopK *int `json:"top_k,omitempty"`
|
|
CrossLanguages []string `json:"cross_languages,omitempty"`
|
|
SearchID *string `json:"search_id,omitempty"`
|
|
MetadataFilter map[string]interface{} `json:"meta_data_filter,omitempty"`
|
|
RerankID *string `json:"rerank_id,omitempty"`
|
|
Keyword *bool `json:"keyword,omitempty"`
|
|
SimilarityThreshold *float64 `json:"similarity_threshold,omitempty"`
|
|
VectorSimilarityWeight *float64 `json:"vector_similarity_weight,omitempty"`
|
|
}
|
|
|
|
// ToSearchDatasetsRequest converts a single-dataset search request into the multi-dataset form.
|
|
func (req *SearchDatasetRequest) ToSearchDatasetsRequest(datasetID string) *SearchDatasetsRequest {
|
|
if req == nil {
|
|
return &SearchDatasetsRequest{DatasetIDs: []string{datasetID}}
|
|
}
|
|
return &SearchDatasetsRequest{
|
|
DatasetIDs: []string{datasetID},
|
|
Question: req.Question,
|
|
Page: req.Page,
|
|
Size: req.Size,
|
|
DocIDs: req.DocIDs,
|
|
UseKG: req.UseKG,
|
|
TopK: req.TopK,
|
|
CrossLanguages: req.CrossLanguages,
|
|
SearchID: req.SearchID,
|
|
MetadataFilter: req.MetadataFilter,
|
|
RerankID: req.RerankID,
|
|
Keyword: req.Keyword,
|
|
SimilarityThreshold: req.SimilarityThreshold,
|
|
VectorSimilarityWeight: req.VectorSimilarityWeight,
|
|
}
|
|
}
|
|
|
|
// SearchDataset searches chunks within one knowledge base based on a question.
|
|
func (d *DatasetService) SearchDataset(datasetID, userID string, req *SearchDatasetRequest) (*SearchDatasetsResponse, error) {
|
|
if datasetID == "" {
|
|
return nil, fmt.Errorf("dataset_id is required")
|
|
}
|
|
return d.SearchDatasets(req.ToSearchDatasetsRequest(datasetID), userID)
|
|
}
|
|
|
|
// SearchDatasets searches chunks across one or more knowledge bases based on a question.
|
|
// It retrieves relevant chunks using embedding and optional reranking, applying filters,
|
|
// cross-language translation, and keyword extraction as configured.
|
|
func (d *DatasetService) SearchDatasets(req *SearchDatasetsRequest, userID string) (*SearchDatasetsResponse, error) {
|
|
if req.Question == "" {
|
|
return nil, fmt.Errorf("question is required")
|
|
}
|
|
if len(req.DatasetIDs) == 0 {
|
|
return nil, fmt.Errorf("dataset_ids is required")
|
|
}
|
|
common.Info("SearchDatasets started", zap.String("userID", userID), zap.Any("datasets", req.DatasetIDs), zap.String("question", req.Question))
|
|
|
|
page := 1
|
|
if req.Page != nil {
|
|
page = *req.Page
|
|
}
|
|
pageSize := 30
|
|
if req.Size != nil {
|
|
pageSize = *req.Size
|
|
}
|
|
useKG := false
|
|
if req.UseKG != nil {
|
|
useKG = *req.UseKG
|
|
}
|
|
similarityThreshold := 0.0
|
|
if req.SimilarityThreshold != nil {
|
|
similarityThreshold = *req.SimilarityThreshold
|
|
}
|
|
vectorSimilarityWeight := 0.3
|
|
if req.VectorSimilarityWeight != nil {
|
|
vectorSimilarityWeight = *req.VectorSimilarityWeight
|
|
}
|
|
topK := 1024
|
|
if req.TopK != nil {
|
|
topK = *req.TopK
|
|
}
|
|
if topK < 1 {
|
|
topK = 1
|
|
} else if topK > 2048 {
|
|
topK = 2048
|
|
}
|
|
keyword := false
|
|
if req.Keyword != nil {
|
|
keyword = *req.Keyword
|
|
}
|
|
searchID := ""
|
|
if req.SearchID != nil {
|
|
searchID = *req.SearchID
|
|
}
|
|
rerankID := ""
|
|
if req.RerankID != nil {
|
|
rerankID = *req.RerankID
|
|
}
|
|
|
|
question := req.Question
|
|
datasetIDs := req.DatasetIDs
|
|
metadataFilter := req.MetadataFilter
|
|
crossLanguages := req.CrossLanguages
|
|
|
|
common.Debug(fmt.Sprintf("SearchDatasets request:\n"+
|
|
" datasetIDs=%v\n"+
|
|
" question=%s\n"+
|
|
" page=%v, pageSize=%v\n"+
|
|
" docIDs=%v\n"+
|
|
" useKG=%v, topK=%v\n"+
|
|
" crossLanguages=%v\n"+
|
|
" searchID=%v\n"+
|
|
" metadataFilter=%v\n"+
|
|
" rerankID=%v\n"+
|
|
" keyword=%v\n"+
|
|
" similarityThreshold=%v, vectorSimilarityWeight=%v",
|
|
datasetIDs, req.Question,
|
|
common.PtrString(req.Page), common.PtrString(req.Size), req.DocIDs,
|
|
useKG, topK, crossLanguages, searchID,
|
|
metadataFilter,
|
|
rerankID,
|
|
keyword,
|
|
similarityThreshold, vectorSimilarityWeight))
|
|
|
|
ctx := context.Background()
|
|
modelProviderSvc := NewModelProviderService()
|
|
|
|
// Access check for all datasets
|
|
var tenantIDs []string
|
|
var kbRecords []*entity.Knowledgebase
|
|
seenTenants := make(map[string]bool)
|
|
for _, datasetID := range datasetIDs {
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
common.Warn("SearchDatasets access denied", zap.String("datasetID", datasetID), zap.String("userID", userID))
|
|
return nil, fmt.Errorf("only owner of dataset %s is authorized for this operation", datasetID)
|
|
}
|
|
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil || kb == nil {
|
|
common.Warn("SearchDatasets dataset not found", zap.String("datasetID", datasetID))
|
|
return nil, fmt.Errorf("dataset %s not found", datasetID)
|
|
}
|
|
if !seenTenants[kb.TenantID] {
|
|
seenTenants[kb.TenantID] = true
|
|
tenantIDs = append(tenantIDs, kb.TenantID)
|
|
}
|
|
kbRecords = append(kbRecords, kb)
|
|
}
|
|
|
|
// Check if all kbs have the same embedding model
|
|
if err := validateDatasetEmbeddingModels(kbRecords); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Override request fields with values from saved search config (if search_id is provided)
|
|
var chatID string
|
|
if searchID != "" {
|
|
if d.searchService == nil {
|
|
common.Warn("Search service is not initialized for search_id", zap.String("searchID", searchID))
|
|
return nil, fmt.Errorf("Invalid search_id")
|
|
}
|
|
searchDetail, err := d.searchService.GetDetail(searchID)
|
|
if err != nil || searchDetail == nil || len(searchDetail) == 0 {
|
|
common.Warn("Invalid search_id", zap.String("searchID", searchID), zap.Error(err))
|
|
return nil, fmt.Errorf("Invalid search_id")
|
|
} else if searchConfig, ok := searchDetail["search_config"].(map[string]interface{}); ok && searchConfig != nil {
|
|
if scMetadataFilter, ok := searchConfig["meta_data_filter"].(map[string]interface{}); ok {
|
|
metadataFilter = scMetadataFilter
|
|
}
|
|
if scST, ok := searchConfig["similarity_threshold"].(float64); ok {
|
|
similarityThreshold = scST
|
|
}
|
|
if scVSW, ok := searchConfig["vector_similarity_weight"].(float64); ok {
|
|
vectorSimilarityWeight = scVSW
|
|
}
|
|
if scTopK, ok := searchConfig["top_k"].(float64); ok {
|
|
topK = int(scTopK)
|
|
if topK < 1 {
|
|
topK = 1
|
|
} else if topK > 2048 {
|
|
topK = 2048
|
|
}
|
|
}
|
|
if scUseKG, ok := searchConfig["use_kg"].(bool); ok {
|
|
useKG = scUseKG
|
|
}
|
|
if scLangs, ok := searchConfig["cross_languages"].([]interface{}); ok {
|
|
crossLanguages = make([]string, len(scLangs))
|
|
for i, l := range scLangs {
|
|
if s, ok := l.(string); ok {
|
|
crossLanguages[i] = s
|
|
}
|
|
}
|
|
}
|
|
if scKeyword, ok := searchConfig["keyword"].(bool); ok {
|
|
keyword = scKeyword
|
|
}
|
|
if scRerankID, ok := searchConfig["rerank_id"].(string); ok {
|
|
rerankID = scRerankID
|
|
}
|
|
chatID, _ = searchConfig["chat_id"].(string)
|
|
|
|
common.Debug("SearchDatasets loaded Search config",
|
|
zap.String("searchID", searchID),
|
|
zap.Strings("datasetIDs", datasetIDs),
|
|
zap.Float64("vectorSimilarityWeight", vectorSimilarityWeight),
|
|
zap.Float64("fullTextWeight", 1-vectorSimilarityWeight),
|
|
zap.Float64("similarityThreshold", similarityThreshold),
|
|
zap.Int("topK", topK),
|
|
zap.Strings("crossLanguages", crossLanguages),
|
|
zap.Bool("keyword", keyword),
|
|
zap.String("rerankID", rerankID),
|
|
zap.String("chatID", chatID),
|
|
zap.Bool("useKG", useKG))
|
|
} else {
|
|
common.Warn("Invalid search_id: search_config missing or invalid", zap.String("searchID", searchID))
|
|
return nil, fmt.Errorf("Invalid search_id")
|
|
}
|
|
}
|
|
|
|
// If meta_data_filter method is auto/semi_auto, get chat model
|
|
var err error
|
|
var chatModelForFilter *models.ChatModel
|
|
if metadataFilter != nil {
|
|
method, _ := metadataFilter["method"].(string)
|
|
if method == "auto" || method == "semi_auto" {
|
|
if chatID != "" {
|
|
driver, modelName, apiConfig, _, err := modelProviderSvc.ResolveModelConfig(tenantIDs[0], entity.ModelTypeChat, chatID)
|
|
if err != nil {
|
|
common.Warn("Failed to get chat model config from search_config chat_id, using tenant default", zap.String("chatID", chatID), zap.Error(err))
|
|
} else {
|
|
chatModelForFilter = models.NewChatModel(driver, &modelName, apiConfig)
|
|
common.Info("Fetched chat model (from search_config) for metadata filter",
|
|
zap.String("chatID", chatID),
|
|
zap.String("tenantID", tenantIDs[0]))
|
|
}
|
|
}
|
|
|
|
if chatModelForFilter == nil {
|
|
driver, modelName, apiConfig, _, err := modelProviderSvc.GetTenantDefaultModelByType(tenantIDs[0], entity.ModelTypeChat)
|
|
if err != nil {
|
|
common.Warn("Failed to get tenant default chat model for meta_data_filter", zap.Error(err))
|
|
} else {
|
|
chatModelForFilter = models.NewChatModel(driver, &modelName, apiConfig)
|
|
common.Info("Fetched chat model (tenant default) for metadata filter",
|
|
zap.String("tenantID", tenantIDs[0]))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Apply meta_data_filter to get filtered doc_ids
|
|
docIDs := make([]string, len(req.DocIDs))
|
|
copy(docIDs, req.DocIDs)
|
|
if len(metadataFilter) > 0 {
|
|
metadataSvc := NewMetadataService()
|
|
flattedMeta, err := metadataSvc.GetFlattedMetaByKBs(datasetIDs)
|
|
if err != nil {
|
|
common.Warn("Failed to get flatted metadata, using empty metadata for filter", zap.Error(err))
|
|
flattedMeta = make(common.MetaData)
|
|
}
|
|
common.Info("Metadata filter conditions", zap.Any("filter", metadataFilter))
|
|
filteredDocIDs, _ := ApplyMetaDataFilter(ctx, metadataFilter, flattedMeta, question, chatModelForFilter, req.DocIDs, datasetIDs)
|
|
docIDs = filteredDocIDs
|
|
common.Info("ApplyMetaDataFilter result", zap.Strings("docIDs", docIDs))
|
|
}
|
|
|
|
// Apply cross_languages and keyword extraction
|
|
modifiedQuestion := question
|
|
if len(crossLanguages) > 0 {
|
|
// Pass tenantID and empty llmID so CrossLanguages can fetch default if needed
|
|
// This matches Python's cross_languages(tenant_id, llm_id, query, languages)
|
|
common.Info("CrossLanguages: dispatching translation",
|
|
zap.String("tenantID", tenantIDs[0]),
|
|
zap.String("llmID", ""),
|
|
zap.Strings("crossLanguages", crossLanguages))
|
|
translated, err := CrossLanguages(ctx, tenantIDs[0], "", question, crossLanguages)
|
|
if err != nil {
|
|
common.Warn("Failed to translate question", zap.String("llmID", ""), zap.Error(err))
|
|
} else {
|
|
modifiedQuestion = translated
|
|
}
|
|
}
|
|
if keyword {
|
|
driver, modelName, apiConfig, _, err := modelProviderSvc.GetTenantDefaultModelByType(tenantIDs[0], entity.ModelTypeChat)
|
|
if err != nil {
|
|
common.Warn("Failed to get default chat model for LLM transformations", zap.Error(err))
|
|
} else {
|
|
chatModel := models.NewChatModel(driver, &modelName, apiConfig)
|
|
common.Info("Fetched chat model (tenant default) for keyword_extraction",
|
|
zap.String("tenantID", tenantIDs[0]))
|
|
|
|
extractedKeywords, err := KeywordExtraction(ctx, chatModel, modifiedQuestion, 3)
|
|
if err != nil {
|
|
common.Warn("Failed to extract keywords from question", zap.Error(err))
|
|
} else if extractedKeywords != "" {
|
|
modifiedQuestion = modifiedQuestion + extractedKeywords
|
|
}
|
|
}
|
|
}
|
|
if modifiedQuestion != question {
|
|
common.Info("Modified question after transformations",
|
|
zap.String("originalQuestion", question),
|
|
zap.String("modifiedQuestion", modifiedQuestion),
|
|
zap.Strings("crossLanguages", crossLanguages),
|
|
zap.Bool("keywordExtraction", keyword))
|
|
}
|
|
|
|
// Get tag-based rank features via LabelQuestion
|
|
metadataSvc := NewMetadataService()
|
|
labels := metadataSvc.LabelQuestion(modifiedQuestion, kbRecords)
|
|
if len(labels) > 0 {
|
|
common.Debug("LabelQuestion result", zap.Any("labels", labels))
|
|
}
|
|
|
|
// Determine embedding model
|
|
var embeddingModel *models.EmbeddingModel
|
|
if kbRecords[0].EmbdID != "" {
|
|
driver, modelName, apiConfig, maxTokens, embErr := modelProviderSvc.ResolveModelConfig(tenantIDs[0], entity.ModelTypeEmbedding, kbRecords[0].EmbdID)
|
|
if embErr != nil {
|
|
return nil, fmt.Errorf("failed to get embedding model by embd_id: %w", embErr)
|
|
}
|
|
embeddingModel = models.NewEmbeddingModel(driver, &modelName, apiConfig, maxTokens)
|
|
common.Info("Fetched embedding model for retrieval",
|
|
zap.String("tenantID", tenantIDs[0]),
|
|
zap.String("modelName", modelName))
|
|
|
|
}
|
|
|
|
// Get rerank model if rerankID is specified
|
|
var rerankModel *models.RerankModel
|
|
if rerankID != "" {
|
|
driver, modelName, apiConfig, _, rErr := modelProviderSvc.ResolveModelConfig(tenantIDs[0], entity.ModelTypeRerank, rerankID)
|
|
if rErr != nil {
|
|
return nil, fmt.Errorf("failed to get rerank model by rerank_id: %w", rErr)
|
|
}
|
|
rerankModel = models.NewRerankModel(driver, &modelName, apiConfig)
|
|
common.Info("Fetched rerank model",
|
|
zap.String("tenantID", tenantIDs[0]),
|
|
zap.String("modelName", modelName))
|
|
}
|
|
|
|
retrievalReq := &nlp.RetrievalRequest{
|
|
TenantIDs: tenantIDs,
|
|
Question: modifiedQuestion,
|
|
KbIDs: datasetIDs,
|
|
DocIDs: docIDs,
|
|
Page: page,
|
|
PageSize: pageSize,
|
|
Top: &topK,
|
|
SimilarityThreshold: &similarityThreshold,
|
|
VectorSimilarityWeight: &vectorSimilarityWeight,
|
|
RerankModel: rerankModel,
|
|
RankFeature: &labels,
|
|
EmbeddingModel: embeddingModel,
|
|
}
|
|
|
|
retrievalResult, err := nlp.NewRetrievalService(d.docEngine, d.documentDAO).Retrieval(ctx, retrievalReq)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("retrieval search failed: %w", err)
|
|
}
|
|
|
|
filteredChunks := retrievalResult.Chunks
|
|
|
|
if useKG {
|
|
common.Warn("use_kg is not yet implemented in Go - skipping KG retrieval")
|
|
}
|
|
|
|
filteredChunks = nlp.RetrievalByChildren(filteredChunks, tenantIDs, d.docEngine, ctx)
|
|
|
|
for i := range filteredChunks {
|
|
delete(filteredChunks[i], "vector")
|
|
}
|
|
|
|
common.Info("SearchDatasets completed", zap.String("userID", userID), zap.Any("kbID", datasetIDs), zap.String("question", question), zap.Int64("chunkCount", int64(len(filteredChunks))))
|
|
|
|
// Convert all float64 values to PyFloat64 for Python-compatible JSON serialization
|
|
pyChunks := common.ConvertFloatsToPyFormat(filteredChunks).([]map[string]interface{})
|
|
|
|
return &SearchDatasetsResponse{
|
|
Chunks: pyChunks,
|
|
DocAggs: retrievalResult.DocAggs,
|
|
Labels: &labels,
|
|
Total: retrievalResult.Total,
|
|
}, nil
|
|
}
|
|
|
|
// MetadataConfigField mirrors one field in the dataset metadata config API.
|
|
type MetadataConfigField struct {
|
|
Key string `json:"key"`
|
|
Type string `json:"type"`
|
|
Description *string `json:"description"`
|
|
Enum []string `json:"enum"`
|
|
}
|
|
|
|
// MetadataConfigRequest mirrors PUT /datasets/:dataset_id/metadata/config.
|
|
type MetadataConfigRequest struct {
|
|
Metadata []MetadataConfigField `json:"metadata"`
|
|
BuiltInMetadata []MetadataConfigField `json:"built_in_metadata"`
|
|
}
|
|
|
|
// CreateDatasetRequest represents the request for creating a dataset.
|
|
type CreateDatasetRequest struct {
|
|
Name string `json:"name" binding:"required"`
|
|
EmbeddingModel *string `json:"embedding_model,omitempty"`
|
|
Permission *string `json:"permission,omitempty"`
|
|
ParserID *string `json:"parser_id,omitempty"`
|
|
PipelineID *string `json:"pipeline_id,omitempty"`
|
|
// ParseType indicates pipeline selection mode: 1 = BuiltIn (parser_id),
|
|
// 2 = Pipeline (pipeline_id). nil means unspecified (backward compat).
|
|
ParseType *int `json:"parse_type,omitempty"`
|
|
}
|
|
|
|
// ListDatasets lists datasets with pagination and filtering.
|
|
func (d *DatasetService) ListDatasets(id, name string, page, pageSize int, orderby string, desc bool, keywords string, ownerIDs []string, parserID, userID string) ([]map[string]interface{}, int64, common.ErrorCode, error) {
|
|
id = strings.TrimSpace(id)
|
|
if id != "" {
|
|
normalizedID, err := normalizeDatasetID(id)
|
|
if err != nil {
|
|
return nil, 0, common.CodeDataError, err
|
|
}
|
|
id = normalizedID
|
|
|
|
kbs, err := d.kbDAO.GetKBByIDAndUserID(id, userID)
|
|
if err != nil {
|
|
return nil, 0, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
if len(kbs) == 0 {
|
|
return nil, 0, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, id)
|
|
}
|
|
}
|
|
|
|
name = strings.TrimSpace(name)
|
|
if name != "" {
|
|
kbs, err := d.kbDAO.GetKBByNameAndUserID(name, userID)
|
|
if err != nil {
|
|
return nil, 0, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
if len(kbs) == 0 {
|
|
return nil, 0, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, name)
|
|
}
|
|
}
|
|
|
|
if page <= 0 {
|
|
page = 1
|
|
}
|
|
if pageSize <= 0 {
|
|
pageSize = 30
|
|
}
|
|
|
|
orderby = strings.TrimSpace(orderby)
|
|
if _, ok := datasetAllowedOrderByFields[orderby]; !ok {
|
|
orderby = "create_time"
|
|
}
|
|
|
|
keywords = strings.TrimSpace(keywords)
|
|
parserID = strings.TrimSpace(parserID)
|
|
|
|
// Empty owner ids do not change the query, so only keep the meaningful ones.
|
|
tenantIDs := make([]string, 0, len(ownerIDs))
|
|
for _, ownerID := range ownerIDs {
|
|
ownerID = strings.TrimSpace(ownerID)
|
|
if ownerID != "" {
|
|
tenantIDs = append(tenantIDs, ownerID)
|
|
}
|
|
}
|
|
if len(tenantIDs) == 0 {
|
|
joinedTenants, err := d.tenantDAO.GetJoinedTenantsByUserID(userID)
|
|
if err != nil {
|
|
return nil, 0, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
for _, joinedTenant := range joinedTenants {
|
|
if joinedTenant == nil || joinedTenant.TenantID == "" {
|
|
continue
|
|
}
|
|
tenantIDs = append(tenantIDs, joinedTenant.TenantID)
|
|
}
|
|
}
|
|
|
|
kbs, total, err := d.kbDAO.GetByTenantIDs(tenantIDs, userID, page, pageSize, orderby, desc, keywords, parserID)
|
|
if err != nil {
|
|
return nil, 0, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
|
|
data := make([]map[string]interface{}, 0, len(kbs))
|
|
for _, kb := range kbs {
|
|
if kb == nil {
|
|
continue
|
|
}
|
|
data = append(data, datasetListItemToMap(kb))
|
|
}
|
|
|
|
return data, total, common.CodeSuccess, nil
|
|
}
|
|
|
|
// CreateDataset creates a new dataset.
|
|
func (d *DatasetService) CreateDataset(req *CreateDatasetRequest, tenantID string) (map[string]interface{}, common.ErrorCode, error) {
|
|
if !common.IsValidString(req.Name) {
|
|
return nil, common.CodeDataError, errors.New("Dataset name must be string.")
|
|
}
|
|
|
|
name := strings.TrimSpace(req.Name)
|
|
if name == "" {
|
|
return nil, common.CodeDataError, errors.New("Dataset name can't be empty.")
|
|
}
|
|
if len(name) > entity.DatasetNameLimit {
|
|
return nil, common.CodeDataError, fmt.Errorf("Dataset name length is %d which is large than %d", len(name), entity.DatasetNameLimit)
|
|
}
|
|
|
|
tenant, err := d.tenantDAO.GetByID(tenantID)
|
|
if err != nil || tenant == nil {
|
|
return nil, common.CodeDataError, errors.New("Tenant not found.")
|
|
}
|
|
|
|
// parse_type explicitly signals the pipeline mode (1 = BuiltIn, 2 = Pipeline).
|
|
// When absent (nil), infer intent from which field is present for backward compat.
|
|
isPipelineMode := req.ParseType != nil && *req.ParseType == 2
|
|
isBuiltinMode := req.ParseType != nil && *req.ParseType == 1
|
|
|
|
if isBuiltinMode && req.PipelineID != nil {
|
|
// BuiltIn mode: discard pipeline_id so only parser_id matters.
|
|
req.PipelineID = nil
|
|
}
|
|
if isPipelineMode && req.ParserID != nil {
|
|
// Pipeline mode: ignore parser_id.
|
|
req.ParserID = nil
|
|
}
|
|
|
|
if req.ParseType == nil && req.ParserID != nil && req.PipelineID != nil {
|
|
return nil, common.CodeDataError, errors.New("parser_id and pipeline_id are mutually exclusive")
|
|
}
|
|
|
|
parserID := ""
|
|
permission := "me"
|
|
embeddingModel := ""
|
|
pipelineID := req.PipelineID
|
|
|
|
if req.Permission != nil {
|
|
permission = strings.TrimSpace(*req.Permission)
|
|
if permission != "me" && permission != "team" {
|
|
return nil, common.CodeDataError, errors.New("Input should be 'me' or 'team'")
|
|
}
|
|
}
|
|
if req.ParserID != nil {
|
|
parserID = strings.TrimSpace(*req.ParserID)
|
|
if err := validateParserID(parserID); err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
pipelineID = nil
|
|
}
|
|
if req.PipelineID != nil {
|
|
normalizedPipelineID, err := normalizeDatasetPipelineID(*req.PipelineID)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
pipelineID = normalizedPipelineID
|
|
}
|
|
if req.EmbeddingModel != nil {
|
|
embeddingModel = strings.TrimSpace(*req.EmbeddingModel)
|
|
if err := validateDatasetEmbeddingModel(embeddingModel); err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
}
|
|
|
|
// Reject references to a custom canvas the caller does not own or share.
|
|
// The handler passes the caller's user id in tenantID for CreateDataset.
|
|
if pipelineID != nil && strings.TrimSpace(*pipelineID) != "" {
|
|
if ok, err := canvasAccessibleForUser(tenantID, strings.TrimSpace(*pipelineID)); err != nil {
|
|
return nil, common.CodeServerError, err
|
|
} else if !ok {
|
|
return nil, common.CodeDataError, errors.New("canvas is not accessible")
|
|
}
|
|
}
|
|
|
|
// Resolve component params defaults from the DSL template. parser_config
|
|
// stores component params directly: {cpnID: {param: value}}.
|
|
parserConfig, cpErr := resolveComponentParamsDefaults(parserID, pipelineID)
|
|
if cpErr != nil {
|
|
common.Warn("failed to resolve component params defaults for dataset",
|
|
zap.String("parserID", parserID), zap.Error(cpErr))
|
|
parserConfig = entity.JSONMap{}
|
|
}
|
|
|
|
var parserConfigMap map[string]interface{} = parserConfig
|
|
|
|
embdID := tenant.EmbdID
|
|
tenantEmbdID := ptrStringValue(tenant.TenantEmbdID)
|
|
if embeddingModel != "" {
|
|
ok, message := d.verifyEmbeddingAvailability(embeddingModel, tenantID)
|
|
if !ok {
|
|
return nil, common.CodeDataError, errors.New(message)
|
|
}
|
|
embdID = embeddingModel
|
|
tenantEmbdID = ""
|
|
}
|
|
if embdID != "" && tenantEmbdID == "" {
|
|
resolvedID, err := NewModelProviderService().ResolveModelID(tenantID, entity.ModelTypeEmbedding, embdID)
|
|
if err == nil {
|
|
tenantEmbdID = resolvedID
|
|
} else {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
}
|
|
|
|
kbID := utility.GenerateToken()
|
|
status := string(entity.StatusValid)
|
|
// Deduplicate name within tenant
|
|
duplicateName, err := common.DuplicateName(func(n, tid string) bool {
|
|
existing, err := d.kbDAO.GetByName(n, tid)
|
|
return err == nil && existing != nil
|
|
}, name, tenantID)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
|
|
kb := &entity.Knowledgebase{
|
|
ID: kbID,
|
|
Name: duplicateName,
|
|
TenantID: tenantID,
|
|
CreatedBy: tenantID,
|
|
ParserID: parserID,
|
|
PipelineID: pipelineID,
|
|
ParserConfig: entity.JSONMap(parserConfigMap),
|
|
Permission: permission,
|
|
EmbdID: embdID,
|
|
TenantEmbdID: stringPtrIfNotEmpty(tenantEmbdID),
|
|
Status: &status,
|
|
}
|
|
|
|
if err = d.kbDAO.Create(kb); err != nil {
|
|
return nil, common.CodeServerError, errors.New("Failed to save dataset")
|
|
}
|
|
|
|
createdKB, err := d.kbDAO.GetByID(kbID)
|
|
if err != nil || createdKB == nil {
|
|
return nil, common.CodeServerError, errors.New("Dataset created failed")
|
|
}
|
|
|
|
return datasetToMap(createdKB), common.CodeSuccess, nil
|
|
}
|
|
|
|
// DeleteDatasets deletes multiple datasets.
|
|
func (d *DatasetService) DeleteDatasets(ids []string, deleteAll bool, tenantID string) (map[string]interface{}, common.ErrorCode, error) {
|
|
normalizedIDs := make([]string, 0, len(ids))
|
|
seenIDs := make(map[string]struct{}, len(ids))
|
|
|
|
// Canonicalize ids once so every downstream DAO call sees the same 32-char hex format.
|
|
for _, id := range ids {
|
|
normalizedID, err := normalizeDatasetID(strings.TrimSpace(id))
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
if _, exists := seenIDs[normalizedID]; exists {
|
|
return nil, common.CodeDataError, fmt.Errorf("Duplicate ids: '%s'", normalizedID)
|
|
}
|
|
seenIDs[normalizedID] = struct{}{}
|
|
normalizedIDs = append(normalizedIDs, normalizedID)
|
|
}
|
|
|
|
if len(normalizedIDs) == 0 {
|
|
if !deleteAll {
|
|
return map[string]interface{}{"success_count": 0}, common.CodeSuccess, nil
|
|
}
|
|
|
|
kbs, err := d.kbDAO.Query(map[string]interface{}{"tenant_id": tenantID})
|
|
if err != nil {
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
for _, kb := range kbs {
|
|
normalizedIDs = append(normalizedIDs, kb.ID)
|
|
}
|
|
}
|
|
|
|
kbs := make([]*entity.Knowledgebase, 0, len(normalizedIDs))
|
|
unauthorizedIDs := make([]string, 0)
|
|
for _, id := range normalizedIDs {
|
|
kb, err := d.kbDAO.GetByIDAndTenantID(id, tenantID)
|
|
if err != nil || kb == nil {
|
|
unauthorizedIDs = append(unauthorizedIDs, id)
|
|
continue
|
|
}
|
|
kbs = append(kbs, kb)
|
|
}
|
|
if len(unauthorizedIDs) > 0 {
|
|
return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for datasets: '%s'", tenantID, strings.Join(unauthorizedIDs, ", "))
|
|
}
|
|
|
|
errorsList := make([]string, 0)
|
|
successCount := 0
|
|
for _, kb := range kbs {
|
|
if err := d.deleteDataset(tenantID, kb); err != nil {
|
|
errorsList = append(errorsList, err.Error())
|
|
continue
|
|
}
|
|
successCount++
|
|
}
|
|
|
|
if len(errorsList) == 0 {
|
|
return map[string]interface{}{"success_count": successCount}, common.CodeSuccess, nil
|
|
}
|
|
|
|
details := strings.Join(errorsList, "; ")
|
|
if len(details) > 128 {
|
|
details = details[:128]
|
|
}
|
|
errorMessage := fmt.Sprintf(
|
|
"Successfully deleted %d datasets, %d failed. Details: %s...",
|
|
successCount,
|
|
len(errorsList),
|
|
details,
|
|
)
|
|
if successCount == 0 {
|
|
return nil, common.CodeDataError, errors.New(errorMessage)
|
|
}
|
|
|
|
return map[string]interface{}{
|
|
"success_count": successCount,
|
|
"errors": limitStrings(errorsList, 5),
|
|
}, common.CodeSuccess, nil
|
|
}
|
|
|
|
// GetDataset gets a single dataset with its size and linked connectors.
|
|
func (d *DatasetService) GetDataset(datasetID, userID string) (map[string]interface{}, common.ErrorCode, error) {
|
|
datasetID = strings.TrimSpace(datasetID)
|
|
if datasetID == "" {
|
|
return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"")
|
|
}
|
|
|
|
normalizedID, err := normalizeDatasetID(datasetID)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
datasetID = normalizedID
|
|
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, datasetID)
|
|
}
|
|
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil || kb == nil {
|
|
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
|
|
}
|
|
|
|
data := datasetToMap(kb)
|
|
|
|
size, err := d.documentDAO.SumSizeByDatasetID(datasetID)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
data["size"] = size
|
|
|
|
connectors, err := d.connectorDAO.ListByDatasetID(datasetID)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
data["connectors"] = datasetConnectorsOrEmpty(connectors)
|
|
|
|
return data, common.CodeSuccess, nil
|
|
}
|
|
|
|
type DatasetConnectorRequest struct {
|
|
ID string `json:"id"`
|
|
AutoParse string `json:"auto_parse,omitempty"`
|
|
}
|
|
|
|
type UpdateDatasetRequest struct {
|
|
Name *string `json:"name,omitempty"`
|
|
Avatar *string `json:"avatar,omitempty"`
|
|
Description *string `json:"description,omitempty"`
|
|
Language *string `json:"language,omitempty"`
|
|
Connectors *[]DatasetConnectorRequest `json:"connectors,omitempty"`
|
|
EmbdID *string `json:"embd_id,omitempty"`
|
|
EmbeddingModel *string `json:"embedding_model,omitempty"`
|
|
Permission *string `json:"permission,omitempty"`
|
|
ParserID *string `json:"parser_id,omitempty"`
|
|
Pagerank *int64 `json:"pagerank,omitempty"`
|
|
ParserConfig map[string]interface{} `json:"parser_config,omitempty"`
|
|
PipelineID *string `json:"pipeline_id,omitempty"`
|
|
// ParseType indicates pipeline selection mode: 1 = BuiltIn (parser_id),
|
|
// 2 = Pipeline (pipeline_id). nil means unspecified (backward compat).
|
|
ParseType *int `json:"parse_type,omitempty"`
|
|
}
|
|
|
|
// UpdateDataset Update a dataset
|
|
func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req UpdateDatasetRequest) (map[string]interface{}, common.ErrorCode, error) {
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeDataError, errors.New("Dataset not found")
|
|
}
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
|
|
if kb == nil || kb.TenantID != tenantID {
|
|
return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID)
|
|
}
|
|
|
|
connectorsProvided := req.Connectors != nil
|
|
connectors := make([]DatasetConnectorRequest, 0)
|
|
if req.Connectors != nil {
|
|
connectors = *req.Connectors
|
|
}
|
|
|
|
updates := make(map[string]interface{})
|
|
|
|
if req.Name != nil {
|
|
name := strings.TrimSpace(*req.Name)
|
|
if name == "" {
|
|
return nil, common.CodeDataError, errors.New("String should have at least 1 character")
|
|
}
|
|
if len(name) > 128 {
|
|
return nil, common.CodeDataError, errors.New("String should have at most 128 characters")
|
|
}
|
|
updates["name"] = name
|
|
}
|
|
if req.Avatar != nil {
|
|
if len(*req.Avatar) > 65535 {
|
|
return nil, common.CodeDataError, errors.New("String should have at most 65535 characters")
|
|
}
|
|
if err := validateDatasetAvatar(*req.Avatar); err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
updates["avatar"] = *req.Avatar
|
|
}
|
|
if req.Description != nil {
|
|
if len(*req.Description) > 65535 {
|
|
return nil, common.CodeDataError, errors.New("String should have at most 65535 characters")
|
|
}
|
|
updates["description"] = *req.Description
|
|
}
|
|
if req.Language != nil {
|
|
language := strings.TrimSpace(*req.Language)
|
|
if len(language) > 32 {
|
|
return nil, common.CodeDataError, errors.New("String should have at most 32 characters")
|
|
}
|
|
updates["language"] = language
|
|
}
|
|
if req.Permission != nil {
|
|
permission := strings.TrimSpace(*req.Permission)
|
|
if permission != "me" && permission != "team" {
|
|
return nil, common.CodeDataError, errors.New("Input should be 'me' or 'team'")
|
|
}
|
|
updates["permission"] = permission
|
|
}
|
|
|
|
// parse_type explicitly signals the pipeline mode (1 = BuiltIn, 2 = Pipeline).
|
|
// When absent (nil), existing mutual-exclusivity behavior is preserved.
|
|
isPipelineMode := req.ParseType != nil && *req.ParseType == 2
|
|
isBuiltinMode := req.ParseType != nil && *req.ParseType == 1
|
|
|
|
if isBuiltinMode && req.PipelineID != nil {
|
|
req.PipelineID = nil
|
|
}
|
|
if isPipelineMode && req.ParserID != nil {
|
|
req.ParserID = nil
|
|
}
|
|
|
|
if req.PipelineID != nil {
|
|
pipelineID, err := normalizeDatasetPipelineID(*req.PipelineID)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
if pipelineID != nil {
|
|
updates["pipeline_id"] = *pipelineID
|
|
}
|
|
}
|
|
|
|
parserID, parserIDProvided, err := datasetUpdateParserID(req)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
if parserIDProvided {
|
|
updates["parser_id"] = parserID
|
|
}
|
|
|
|
// When parse_type is absent, parser_id and pipeline_id are mutually exclusive.
|
|
if req.ParseType == nil && parserIDProvided && req.PipelineID != nil {
|
|
return nil, common.CodeDataError, errors.New("parser_id and pipeline_id are mutually exclusive")
|
|
}
|
|
|
|
embdID, embdIDProvided, err := datasetUpdateEmbeddingID(req)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
if embdIDProvided {
|
|
tenantEmbdID := ptrStringValue(kb.TenantEmbdID)
|
|
if embdID == "" {
|
|
embdID = kb.EmbdID
|
|
} else {
|
|
tenantEmbdID = ""
|
|
}
|
|
ok, message := d.verifyEmbeddingAvailability(embdID, tenantID)
|
|
if !ok {
|
|
return nil, common.CodeDataError, errors.New(message)
|
|
}
|
|
if embdID != "" && tenantEmbdID == "" {
|
|
resolvedID, err := NewModelProviderService().ResolveModelID(tenantID, entity.ModelTypeEmbedding, embdID)
|
|
if err == nil {
|
|
tenantEmbdID = resolvedID
|
|
}
|
|
}
|
|
updates["embd_id"] = embdID
|
|
updates["tenant_embd_id"] = stringPtrIfNotEmpty(tenantEmbdID)
|
|
}
|
|
|
|
if req.ParserConfig != nil {
|
|
if err := validateDatasetParserConfigSize(req.ParserConfig); err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
if len(req.ParserConfig) > 0 {
|
|
// Resolve the effective pipeline and load its DSL schema.
|
|
effectiveParserID := kb.ParserID
|
|
if parserIDProvided {
|
|
effectiveParserID = parserID
|
|
}
|
|
effectivePipelineID := kb.PipelineID
|
|
if req.PipelineID != nil {
|
|
if normalized, err := normalizeDatasetPipelineID(*req.PipelineID); err == nil {
|
|
effectivePipelineID = normalized
|
|
}
|
|
} else if parserIDProvided && kb.PipelineID != nil {
|
|
effectivePipelineID = nil
|
|
}
|
|
|
|
isCanvas := effectivePipelineID != nil && strings.TrimSpace(*effectivePipelineID) != ""
|
|
dslJSON, dslErr := loadPipelineDSL(isCanvas, effectiveParserID, effectivePipelineID)
|
|
if dslErr != nil {
|
|
common.Warn("failed to load pipeline DSL for building parser_config",
|
|
zap.String("parserID", effectiveParserID), zap.Error(dslErr))
|
|
}
|
|
if dslJSON != nil {
|
|
// Start from DSL defaults, overlay the cleaned incoming overrides.
|
|
// Unknown cpnIDs, unknown params, and legacy flat fields are
|
|
// dropped. Full replace — no merge with the stored config.
|
|
updates["parser_config"] = buildParserConfig(dslJSON, map[string]interface{}(req.ParserConfig))
|
|
}
|
|
}
|
|
}
|
|
|
|
if req.Pagerank != nil && *req.Pagerank != kb.Pagerank {
|
|
if *req.Pagerank < 0 || *req.Pagerank > 100 {
|
|
return nil, common.CodeDataError, errors.New("Input should be less than or equal to 100")
|
|
}
|
|
if !d.docEngine.SupportsPageRank() {
|
|
return nil, common.CodeDataError, errors.New("'pagerank' can only be set when doc_engine is elasticsearch")
|
|
}
|
|
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
|
|
if *req.Pagerank > 0 {
|
|
err = d.docEngine.UpdateChunks(context.Background(), map[string]interface{}{"kb_id": kb.ID}, map[string]interface{}{common.PAGERANK_FLD: *req.Pagerank}, indexName, kb.ID)
|
|
} else {
|
|
err = d.docEngine.UpdateChunks(context.Background(), map[string]interface{}{"exists": common.PAGERANK_FLD}, map[string]interface{}{"remove": common.PAGERANK_FLD}, indexName, kb.ID)
|
|
}
|
|
if err != nil {
|
|
return nil, common.CodeServerError, err
|
|
}
|
|
updates["pagerank"] = *req.Pagerank
|
|
}
|
|
|
|
if parserIDProvided && parserID != kb.ParserID {
|
|
if _, ok := updates["parser_config"]; !ok {
|
|
if resolved, cpErr := resolveComponentParamsDefaults(parserID, nil); cpErr != nil {
|
|
common.Warn("failed to resolve component params defaults on parser_id switch",
|
|
zap.String("parserID", parserID), zap.Error(cpErr))
|
|
} else if resolved != nil {
|
|
updates["parser_config"] = resolved
|
|
}
|
|
}
|
|
}
|
|
if kb.PipelineID != nil && parserIDProvided {
|
|
if _, ok := updates["pipeline_id"]; !ok {
|
|
updates["pipeline_id"] = nil // clear to NULL, not empty string
|
|
}
|
|
}
|
|
|
|
// Regenerate parser_config from the new pipeline's DSL defaults when
|
|
// the pipeline changes, so that component-level parameters stay in sync.
|
|
pipelineChanged := req.PipelineID != nil && (kb.PipelineID == nil || *req.PipelineID != *kb.PipelineID)
|
|
if pipelineChanged {
|
|
cfgParserID := kb.ParserID
|
|
if parserIDProvided {
|
|
cfgParserID = parserID
|
|
}
|
|
cfgPipelineID, _ := updates["pipeline_id"].(string)
|
|
var cpPipelineID *string
|
|
if cfgPipelineID != "" {
|
|
cpPipelineID = &cfgPipelineID
|
|
}
|
|
if cpDefaults, cpErr := resolveComponentParamsDefaults(cfgParserID, cpPipelineID); cpErr != nil {
|
|
common.Warn("failed to resolve component params defaults on pipeline change",
|
|
zap.String("parserID", cfgParserID), zap.Error(cpErr))
|
|
} else if cpDefaults != nil {
|
|
updates["parser_config"] = cpDefaults
|
|
}
|
|
}
|
|
|
|
if nameValue, ok := updates["name"].(string); ok && strings.ToLower(nameValue) != strings.ToLower(kb.Name) {
|
|
existing, lookupErr := d.kbDAO.GetByName(nameValue, tenantID)
|
|
if lookupErr != nil && !dao.IsNotFoundErr(lookupErr) {
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
if existing != nil {
|
|
return nil, common.CodeDataError, fmt.Errorf("Dataset name '%s' already exists", nameValue)
|
|
}
|
|
}
|
|
|
|
if len(updates) == 0 && !connectorsProvided {
|
|
return nil, common.CodeDataError, errors.New("No properties were modified")
|
|
}
|
|
|
|
if len(updates) > 0 {
|
|
if err = d.kbDAO.UpdateByID(kb.ID, updates); err != nil {
|
|
return nil, common.CodeServerError, errors.New("Update dataset error.(Database error)")
|
|
}
|
|
}
|
|
|
|
if connectorsProvided {
|
|
connectorLinks := make([]dao.DatasetConnectorLink, 0, len(connectors))
|
|
for _, connector := range connectors {
|
|
connectorID := strings.TrimSpace(connector.ID)
|
|
if connectorID == "" {
|
|
return nil, common.CodeDataError, errors.New("connector id is required")
|
|
}
|
|
connectorLinks = append(connectorLinks, dao.DatasetConnectorLink{
|
|
ID: connectorID,
|
|
AutoParse: connector.AutoParse,
|
|
})
|
|
}
|
|
if err = d.connectorDAO.LinkDatasetConnectors(kb.ID, connectorLinks); err != nil {
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
}
|
|
|
|
updatedKB, err := d.kbDAO.GetByID(kb.ID)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, errors.New("Dataset updated failed")
|
|
}
|
|
|
|
data := datasetToMap(updatedKB)
|
|
linkedConnectors, err := d.connectorDAO.ListByDatasetID(kb.ID)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
data["connectors"] = datasetConnectorsOrEmpty(linkedConnectors)
|
|
return data, common.CodeSuccess, nil
|
|
}
|
|
|
|
func datasetConnectorsOrEmpty(connectors []*dao.ConnectorDatasetListItem) []*dao.ConnectorDatasetListItem {
|
|
if connectors == nil {
|
|
return make([]*dao.ConnectorDatasetListItem, 0)
|
|
}
|
|
return connectors
|
|
}
|
|
|
|
func datasetUpdateParserID(req UpdateDatasetRequest) (string, bool, error) {
|
|
parserID := ""
|
|
provided := false
|
|
if req.ParserID != nil {
|
|
parserID = strings.TrimSpace(*req.ParserID)
|
|
provided = true
|
|
}
|
|
if !provided {
|
|
return "", false, nil
|
|
}
|
|
if err := validateParserID(parserID); err != nil {
|
|
return "", true, err
|
|
}
|
|
return parserID, true, nil
|
|
}
|
|
|
|
func datasetUpdateEmbeddingID(req UpdateDatasetRequest) (string, bool, error) {
|
|
embdID := ""
|
|
provided := false
|
|
if req.EmbdID != nil {
|
|
embdID = strings.TrimSpace(*req.EmbdID)
|
|
provided = true
|
|
}
|
|
if req.EmbeddingModel != nil {
|
|
embdID = strings.TrimSpace(*req.EmbeddingModel)
|
|
provided = true
|
|
}
|
|
if !provided {
|
|
return "", false, nil
|
|
}
|
|
if embdID != "" {
|
|
if err := validateDatasetEmbeddingModel(embdID); err != nil {
|
|
return "", true, err
|
|
}
|
|
}
|
|
return embdID, true, nil
|
|
}
|
|
|
|
func normalizeDatasetUpdateExt(ext map[string]interface{}) map[string]interface{} {
|
|
if ext == nil {
|
|
return nil
|
|
}
|
|
|
|
updates := make(map[string]interface{}, len(ext))
|
|
for key, value := range ext {
|
|
switch key {
|
|
case "embedding_model":
|
|
updates["embd_id"] = value
|
|
case "chunk_method":
|
|
updates["parser_id"] = value
|
|
case "connectors", "auto_metadata_config", "ext", "parse_type":
|
|
continue
|
|
default:
|
|
updates[key] = value
|
|
}
|
|
}
|
|
return updates
|
|
}
|
|
|
|
// GetMetadataConfig gets the auto-metadata configuration for a dataset.
|
|
func (d *DatasetService) GetMetadataConfig(datasetID, tenantID string) (map[string]interface{}, common.ErrorCode, error) {
|
|
kb, err := d.kbDAO.GetByIDAndTenantID(datasetID, tenantID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID)
|
|
}
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
if kb == nil {
|
|
return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID)
|
|
}
|
|
|
|
return map[string]interface{}{
|
|
"metadata": parserConfigValueOrEmptyList(kb.ParserConfig, "metadata"),
|
|
"built_in_metadata": parserConfigValueOrEmptyList(kb.ParserConfig, "built_in_metadata"),
|
|
}, common.CodeSuccess, nil
|
|
}
|
|
|
|
// UpdateMetadataConfig updates the auto-metadata configuration for a dataset.
|
|
func (d *DatasetService) UpdateMetadataConfig(datasetID, tenantID string, req *MetadataConfigRequest) (map[string]interface{}, common.ErrorCode, error) {
|
|
datasetID = strings.TrimSpace(datasetID)
|
|
tenantID = strings.TrimSpace(tenantID)
|
|
|
|
kb, err := d.kbDAO.GetByIDAndTenantID(datasetID, tenantID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID)
|
|
}
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
if kb == nil {
|
|
return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", tenantID, datasetID)
|
|
}
|
|
|
|
if req == nil {
|
|
req = &MetadataConfigRequest{}
|
|
}
|
|
|
|
metadata, err := normalizeMetadataConfigFields(req.Metadata, "metadata")
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
builtInMetadata, err := normalizeMetadataConfigFields(req.BuiltInMetadata, "built_in_metadata")
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
|
|
parserConfig := kb.ParserConfig
|
|
if parserConfig == nil {
|
|
parserConfig = entity.JSONMap{}
|
|
}
|
|
parserConfig["metadata"] = metadata
|
|
parserConfig["built_in_metadata"] = builtInMetadata
|
|
|
|
if err = d.kbDAO.UpdateByID(kb.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil {
|
|
return nil, common.CodeServerError, errors.New("Update auto-metadata error.(Database error)")
|
|
}
|
|
|
|
return map[string]interface{}{
|
|
"metadata": metadata,
|
|
"built_in_metadata": builtInMetadata,
|
|
}, common.CodeSuccess, nil
|
|
}
|
|
|
|
// Accessible checks if a knowledge base is accessible by a user
|
|
func (d *DatasetService) Accessible(kbID, userID string) bool {
|
|
return d.kbDAO.Accessible(kbID, userID)
|
|
}
|
|
|
|
func (d *DatasetService) GetByID(kbID string) (*entity.Knowledgebase, error) {
|
|
return d.kbDAO.GetByID(kbID)
|
|
}
|
|
|
|
// GetKnowledgebaseByID resolves a dataset entity without applying permission
|
|
// checks. Upload needs the same existence-then-auth ordering as Python.
|
|
func (d *DatasetService) GetKnowledgebaseByID(datasetID string) (*entity.Knowledgebase, error) {
|
|
datasetID = strings.TrimSpace(datasetID)
|
|
if datasetID == "" {
|
|
return nil, errors.New("Lack of \"Dataset ID\"")
|
|
}
|
|
normalizedID, err := normalizeDatasetID(datasetID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return d.kbDAO.GetByID(normalizedID)
|
|
}
|
|
|
|
// CheckKBTeamPermission mirrors Python check_kb_team_permission.
|
|
func (d *DatasetService) CheckKBTeamPermission(kb *entity.Knowledgebase, userID string) bool {
|
|
return hasKBTeamPermission(kb, userID, d.tenantDAO)
|
|
}
|
|
|
|
func (d *DatasetService) AggregateTags(datasetIDs []string, userID string) ([]map[string]interface{}, common.ErrorCode, error) {
|
|
if len(datasetIDs) == 0 {
|
|
return nil, common.CodeDataError, errors.New("Lack of dataset_ids in query parameters")
|
|
}
|
|
if d.docEngine == nil {
|
|
return nil, common.CodeServerError, errors.New("Document engine is not initialized")
|
|
}
|
|
|
|
datasetIDsByTenant := make(map[string][]string)
|
|
for _, rawID := range datasetIDs {
|
|
rawID = strings.TrimSpace(rawID)
|
|
if rawID == "" {
|
|
continue
|
|
}
|
|
datasetID, err := normalizeDatasetID(rawID)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return nil, common.CodeDataError, fmt.Errorf("No authorization for dataset '%s'", datasetID)
|
|
}
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil {
|
|
if dao.IsNotFoundErr(err) {
|
|
return nil, common.CodeDataError, fmt.Errorf("Invalid Dataset ID '%s'", datasetID)
|
|
}
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
if kb.DocNum <= 0 {
|
|
continue
|
|
}
|
|
datasetIDsByTenant[kb.TenantID] = append(datasetIDsByTenant[kb.TenantID], datasetID)
|
|
}
|
|
|
|
const pageSize = 10000
|
|
merged := make(map[string]int)
|
|
for tenantID, kbIDs := range datasetIDsByTenant {
|
|
for offset := 0; ; offset += pageSize {
|
|
searchResp, err := d.docEngine.Search(context.Background(), &types.SearchRequest{
|
|
IndexNames: []string{fmt.Sprintf("ragflow_%s", tenantID)},
|
|
KbIDs: kbIDs,
|
|
Offset: offset,
|
|
Limit: pageSize,
|
|
SelectFields: []string{"tag_kwd"},
|
|
})
|
|
if err != nil {
|
|
return nil, common.CodeServerError, fmt.Errorf("failed to aggregate tags: %w", err)
|
|
}
|
|
for _, agg := range d.docEngine.GetAggregation(searchResp.Chunks, "tag_kwd") {
|
|
tag, _ := agg["key"].(string)
|
|
if tag == "" {
|
|
continue
|
|
}
|
|
switch count := agg["count"].(type) {
|
|
case int:
|
|
merged[tag] += count
|
|
case int32:
|
|
merged[tag] += int(count)
|
|
case int64:
|
|
merged[tag] += int(count)
|
|
case float64:
|
|
merged[tag] += int(count)
|
|
}
|
|
}
|
|
|
|
chunkCount := len(searchResp.Chunks)
|
|
if chunkCount == 0 || chunkCount < pageSize {
|
|
break
|
|
}
|
|
if searchResp.Total > 0 && int64(offset+chunkCount) >= searchResp.Total {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
result := make([]map[string]interface{}, 0, len(merged))
|
|
for tag, count := range merged {
|
|
result = append(result, map[string]interface{}{
|
|
"value": tag,
|
|
"count": count,
|
|
})
|
|
}
|
|
return result, common.CodeSuccess, nil
|
|
}
|
|
|
|
func (d *DatasetService) ListTags(datasetID, userID string) ([]map[string]interface{}, common.ErrorCode, error) {
|
|
datasetID = strings.TrimSpace(datasetID)
|
|
if datasetID == "" {
|
|
return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"")
|
|
}
|
|
|
|
normalizedID, err := normalizeDatasetID(datasetID)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
datasetID = normalizedID
|
|
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return nil, common.CodeDataError, errors.New("No authorization.")
|
|
}
|
|
if d.docEngine == nil {
|
|
return nil, common.CodeServerError, errors.New("Document engine is not initialized")
|
|
}
|
|
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil || kb == nil {
|
|
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
|
|
}
|
|
|
|
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
exists, err := d.docEngine.ChunkStoreExists(ctx, indexName, datasetID)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, fmt.Errorf("failed to inspect chunk store: %w", err)
|
|
}
|
|
if !exists {
|
|
return []map[string]interface{}{}, common.CodeSuccess, nil
|
|
}
|
|
|
|
const pageSize = 10000
|
|
counts := make(map[string]int)
|
|
for offset := 0; ; offset += pageSize {
|
|
if err = ctx.Err(); err != nil {
|
|
return nil, common.CodeServerError, fmt.Errorf("list tags timeout or canceled: %w", err)
|
|
}
|
|
|
|
searchResp, err := d.docEngine.Search(ctx, &types.SearchRequest{
|
|
IndexNames: []string{indexName},
|
|
KbIDs: []string{datasetID},
|
|
Offset: offset,
|
|
Limit: pageSize,
|
|
SelectFields: []string{"tag_kwd"},
|
|
})
|
|
if err != nil {
|
|
return nil, common.CodeServerError, fmt.Errorf("failed to list tags: %w", err)
|
|
}
|
|
|
|
for _, agg := range d.docEngine.GetAggregation(searchResp.Chunks, "tag_kwd") {
|
|
tag, _ := agg["key"].(string)
|
|
if tag == "" {
|
|
continue
|
|
}
|
|
switch count := agg["count"].(type) {
|
|
case int:
|
|
counts[tag] += count
|
|
case int32:
|
|
counts[tag] += int(count)
|
|
case int64:
|
|
counts[tag] += int(count)
|
|
case float64:
|
|
counts[tag] += int(count)
|
|
}
|
|
}
|
|
|
|
chunkCount := len(searchResp.Chunks)
|
|
if chunkCount == 0 || chunkCount < pageSize {
|
|
break
|
|
}
|
|
if searchResp.Total > 0 && int64(offset+chunkCount) >= searchResp.Total {
|
|
break
|
|
}
|
|
}
|
|
|
|
if len(counts) == 0 {
|
|
return []map[string]interface{}{}, common.CodeSuccess, nil
|
|
}
|
|
|
|
tags := make([]string, 0, len(counts))
|
|
for tag := range counts {
|
|
tags = append(tags, tag)
|
|
}
|
|
sort.Slice(tags, func(i, j int) bool {
|
|
if counts[tags[i]] != counts[tags[j]] {
|
|
return counts[tags[i]] > counts[tags[j]]
|
|
}
|
|
return tags[i] < tags[j]
|
|
})
|
|
|
|
result := make([]map[string]interface{}, 0, len(tags))
|
|
for _, tag := range tags {
|
|
result = append(result, map[string]interface{}{
|
|
"key": tag,
|
|
"count": counts[tag],
|
|
})
|
|
}
|
|
|
|
return result, common.CodeSuccess, nil
|
|
}
|
|
|
|
// GetIngestionSummary returns dataset-level ingestion counters together with
|
|
// the aggregated document parsing status, mirroring
|
|
// dataset_api_service.get_ingestion_summary.
|
|
func (d *DatasetService) GetIngestionSummary(datasetID, userID string) (map[string]interface{}, common.ErrorCode, error) {
|
|
datasetID = strings.TrimSpace(datasetID)
|
|
if datasetID == "" {
|
|
return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"")
|
|
}
|
|
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, datasetID)
|
|
}
|
|
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil || kb == nil {
|
|
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
|
|
}
|
|
|
|
status, err := d.documentDAO.GetParsingStatusByKBID(datasetID)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
|
|
return map[string]interface{}{
|
|
"doc_num": kb.DocNum,
|
|
"chunk_num": kb.ChunkNum,
|
|
"token_num": kb.TokenNum,
|
|
"status": status,
|
|
}, common.CodeSuccess, nil
|
|
}
|
|
|
|
// ListIngestionLogs lists ingestion logs for a dataset, mirroring
|
|
// dataset_api_service.list_ingestion_logs. log_type selects between
|
|
// dataset-level logs ("dataset") and per-file logs ("file").
|
|
func (d *DatasetService) ListIngestionLogs(datasetID, userID string, page, pageSize int, orderby string, desc bool, operationStatus []string, createDateFrom, createDateTo, logType, keywords string) (map[string]interface{}, common.ErrorCode, error) {
|
|
datasetID = strings.TrimSpace(datasetID)
|
|
if datasetID == "" {
|
|
return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"")
|
|
}
|
|
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return nil, common.CodeDataError, errors.New("No authorization.")
|
|
}
|
|
|
|
if logType != "dataset" && logType != "file" {
|
|
return nil, common.CodeDataError, errors.New("Invalid \"log_type\", expected \"dataset\" or \"file\"")
|
|
}
|
|
|
|
var (
|
|
logs []*entity.PipelineOperationLog
|
|
total int64
|
|
err error
|
|
)
|
|
if logType == "file" {
|
|
logs, total, err = d.pipelineLogDAO.GetFileLogsByKBID(datasetID, page, pageSize, orderby, desc, keywords, operationStatus, createDateFrom, createDateTo)
|
|
} else {
|
|
logs, total, err = d.pipelineLogDAO.GetDatasetLogsByKBID(datasetID, page, pageSize, orderby, desc, operationStatus, createDateFrom, createDateTo, keywords)
|
|
}
|
|
if err != nil {
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
|
|
items := make([]map[string]interface{}, 0, len(logs))
|
|
for _, log := range logs {
|
|
if log == nil {
|
|
continue
|
|
}
|
|
if logType == "file" {
|
|
items = append(items, fileIngestionLogToMap(log))
|
|
} else {
|
|
items = append(items, datasetIngestionLogToMap(log))
|
|
}
|
|
}
|
|
|
|
return map[string]interface{}{
|
|
"total": total,
|
|
"logs": items,
|
|
}, common.CodeSuccess, nil
|
|
}
|
|
|
|
// GetIngestionLog returns a single ingestion log, mirroring
|
|
// dataset_api_service.get_ingestion_log. It returns the full record (including
|
|
// the `dsl`, `document_id`, `parser_id`, etc.) so that the front-end
|
|
// dataflow-result page can render the pipeline timeline and chunks. The
|
|
// file-level converter is a superset of the dataset-level fields, so it is
|
|
// correct for both dataset-level (graph/raptor/mindmap) and per-file logs.
|
|
func (d *DatasetService) GetIngestionLog(datasetID, userID, logID string) (map[string]interface{}, common.ErrorCode, error) {
|
|
datasetID = strings.TrimSpace(datasetID)
|
|
if datasetID == "" {
|
|
return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"")
|
|
}
|
|
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return nil, common.CodeDataError, errors.New("No authorization.")
|
|
}
|
|
|
|
log, err := d.pipelineLogDAO.GetByIDAndKBID(logID, datasetID)
|
|
if err != nil {
|
|
if errors.Is(err, gorm.ErrRecordNotFound) {
|
|
return nil, common.CodeDataError, errors.New("Log not found")
|
|
}
|
|
return nil, common.CodeServerError, errors.New("Database operation failed")
|
|
}
|
|
|
|
return fileIngestionLogToMap(log), common.CodeSuccess, nil
|
|
}
|
|
|
|
func datasetIngestionLogToMap(log *entity.PipelineOperationLog) map[string]interface{} {
|
|
return map[string]interface{}{
|
|
"id": log.ID,
|
|
"tenant_id": log.TenantID,
|
|
"kb_id": log.KbID,
|
|
"progress": log.Progress,
|
|
"progress_msg": stringPointerValue(log.ProgressMsg),
|
|
"process_begin_at": timePointerValue(log.ProcessBeginAt),
|
|
"process_duration": log.ProcessDuration,
|
|
"task_type": log.TaskType,
|
|
"operation_status": log.OperationStatus,
|
|
"avatar": stringPointerValue(log.Avatar),
|
|
"status": stringPointerValue(log.Status),
|
|
"create_time": int64PointerValue(log.CreateTime),
|
|
"create_date": timePointerValue(log.CreateDate),
|
|
"update_time": int64PointerValue(log.UpdateTime),
|
|
"update_date": timePointerValue(log.UpdateDate),
|
|
}
|
|
}
|
|
|
|
func fileIngestionLogToMap(log *entity.PipelineOperationLog) map[string]interface{} {
|
|
return map[string]interface{}{
|
|
"id": log.ID,
|
|
"document_id": log.DocumentID,
|
|
"tenant_id": log.TenantID,
|
|
"kb_id": log.KbID,
|
|
"pipeline_id": stringPointerValue(log.PipelineID),
|
|
"pipeline_title": stringPointerValue(log.PipelineTitle),
|
|
"parser_id": log.ParserID,
|
|
"document_name": log.DocumentName,
|
|
"document_suffix": log.DocumentSuffix,
|
|
"document_type": log.DocumentType,
|
|
"source_from": log.SourceFrom,
|
|
"progress": log.Progress,
|
|
"progress_msg": stringPointerValue(log.ProgressMsg),
|
|
"process_begin_at": timePointerValue(log.ProcessBeginAt),
|
|
"process_duration": log.ProcessDuration,
|
|
"dsl": jsonMapValue(log.DSL),
|
|
"task_type": log.TaskType,
|
|
"operation_status": log.OperationStatus,
|
|
"avatar": stringPointerValue(log.Avatar),
|
|
"status": stringPointerValue(log.Status),
|
|
"create_time": int64PointerValue(log.CreateTime),
|
|
"create_date": timePointerValue(log.CreateDate),
|
|
"update_time": int64PointerValue(log.UpdateTime),
|
|
"update_date": timePointerValue(log.UpdateDate),
|
|
}
|
|
}
|
|
|
|
func stringPointerValue(s *string) interface{} {
|
|
if s == nil {
|
|
return nil
|
|
}
|
|
return *s
|
|
}
|
|
|
|
func int64PointerValue(i *int64) interface{} {
|
|
if i == nil {
|
|
return nil
|
|
}
|
|
return *i
|
|
}
|
|
|
|
func timePointerValue(t *time.Time) interface{} {
|
|
if t == nil {
|
|
return nil
|
|
}
|
|
return t.Format("2006-01-02 15:04:05")
|
|
}
|
|
|
|
func jsonMapValue(m entity.JSONMap) interface{} {
|
|
if m == nil {
|
|
return nil
|
|
}
|
|
return m
|
|
}
|
|
|
|
func (d *DatasetService) deleteDataset(tenantID string, kb *entity.Knowledgebase) error {
|
|
return dao.DB.Transaction(func(tx *gorm.DB) error {
|
|
if taskIDs := datasetIndexTaskIDs(kb); len(taskIDs) > 0 {
|
|
if err := tx.Where("id IN ?", taskIDs).Delete(&entity.Task{}).Error; err != nil {
|
|
return fmt.Errorf("Delete dataset error for %s", kb.ID)
|
|
}
|
|
}
|
|
|
|
var documents []entity.Document
|
|
if err := tx.Where("kb_id = ?", kb.ID).Find(&documents).Error; err != nil {
|
|
return fmt.Errorf("Delete dataset error for %s", kb.ID)
|
|
}
|
|
|
|
docIDs := make([]string, 0, len(documents))
|
|
for _, document := range documents {
|
|
docIDs = append(docIDs, document.ID)
|
|
}
|
|
|
|
if len(docIDs) > 0 {
|
|
var mappings []entity.File2Document
|
|
if err := tx.Where("document_id IN ?", docIDs).Find(&mappings).Error; err != nil {
|
|
return fmt.Errorf("Delete dataset error for %s", kb.ID)
|
|
}
|
|
|
|
fileIDs := make([]string, 0, len(mappings))
|
|
seenFileIDs := make(map[string]struct{}, len(mappings))
|
|
for _, mapping := range mappings {
|
|
if mapping.FileID == nil || *mapping.FileID == "" {
|
|
continue
|
|
}
|
|
if _, exists := seenFileIDs[*mapping.FileID]; exists {
|
|
continue
|
|
}
|
|
seenFileIDs[*mapping.FileID] = struct{}{}
|
|
fileIDs = append(fileIDs, *mapping.FileID)
|
|
}
|
|
|
|
if err := tx.Where("doc_id IN ?", docIDs).Delete(&entity.Task{}).Error; err != nil {
|
|
return fmt.Errorf("Delete dataset error for %s", kb.ID)
|
|
}
|
|
if err := tx.Where("document_id IN ?", docIDs).Delete(&entity.File2Document{}).Error; err != nil {
|
|
return fmt.Errorf("Delete dataset error for %s", kb.ID)
|
|
}
|
|
if len(fileIDs) > 0 {
|
|
if err := tx.Unscoped().Where("id IN ?", fileIDs).Delete(&entity.File{}).Error; err != nil {
|
|
return fmt.Errorf("Delete dataset error for %s", kb.ID)
|
|
}
|
|
}
|
|
if err := tx.Where("id IN ?", docIDs).Delete(&entity.Document{}).Error; err != nil {
|
|
return fmt.Errorf("Delete dataset error for %s", kb.ID)
|
|
}
|
|
}
|
|
|
|
if err := tx.Unscoped().
|
|
Where("source_type = ? AND type = ? AND name = ? AND tenant_id = ?", string(entity.FileSourceKnowledgebase), "folder", kb.Name, tenantID).
|
|
Delete(&entity.File{}).Error; err != nil {
|
|
return fmt.Errorf("Delete dataset error for %s", kb.ID)
|
|
}
|
|
|
|
if err := tx.Where("id = ?", kb.ID).Delete(&entity.Knowledgebase{}).Error; err != nil {
|
|
return fmt.Errorf("Delete dataset error for %s", kb.ID)
|
|
}
|
|
|
|
return nil
|
|
})
|
|
}
|
|
|
|
// validateParserID validates parser_id against the built-in
|
|
// pipeline registry. The registry is the single source of truth, so the
|
|
// legacy hardcoded allow-list is gone. Legacy values (e.g. "naive") are
|
|
// accepted via registry aliases.
|
|
func validateParserID(chunkMethod string) error {
|
|
registry, err := pipelinepkg.DefaultRegistry()
|
|
if err != nil || registry == nil {
|
|
return errors.New("parser_id validation unavailable: builtin pipeline registry not loaded")
|
|
}
|
|
if registry.IsValid(chunkMethod) {
|
|
return nil
|
|
}
|
|
return parserIDError()
|
|
}
|
|
|
|
// parserIDError builds a validation error that lists the valid
|
|
// canonical parser_ids from the registry, mirroring the shape of the old
|
|
// hardcoded message but driven by the embedded templates.
|
|
func parserIDError() error {
|
|
registry, err := pipelinepkg.DefaultRegistry()
|
|
if err != nil || registry == nil {
|
|
return errors.New("invalid parser_id")
|
|
}
|
|
refs := registry.Refs()
|
|
switch len(refs) {
|
|
case 0:
|
|
return errors.New("invalid parser_id")
|
|
case 1:
|
|
return fmt.Errorf("Input should be '%s'", refs[0])
|
|
default:
|
|
return fmt.Errorf("Input should be %s or '%s'", quoteList(refs[:len(refs)-1]), refs[len(refs)-1])
|
|
}
|
|
}
|
|
|
|
// quoteList renders ["a", "b"] as "'a', 'b'".
|
|
func quoteList(items []string) string {
|
|
quoted := make([]string, len(items))
|
|
for i, v := range items {
|
|
quoted[i] = "'" + v + "'"
|
|
}
|
|
return strings.Join(quoted, ", ")
|
|
}
|
|
|
|
func validateDatasetAvatar(avatar string) error {
|
|
if !strings.Contains(avatar, ",") {
|
|
return errors.New("Missing MIME prefix. Expected format: data:<mime>;base64,<data>")
|
|
}
|
|
|
|
prefix, _, _ := strings.Cut(avatar, ",")
|
|
if !strings.HasPrefix(prefix, "data:") {
|
|
return errors.New("Invalid MIME prefix format. Must start with 'data:'")
|
|
}
|
|
|
|
mimeType, _, _ := strings.Cut(strings.TrimPrefix(prefix, "data:"), ";")
|
|
if _, ok := datasetSupportedAvatarMIMETypes[mimeType]; !ok {
|
|
return errors.New("Unsupported MIME type. Allowed: [image/jpeg image/png]")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func validateDatasetEmbeddingModel(embeddingModel string) error {
|
|
if embeddingModel == "" {
|
|
return errors.New("Embedding model identifier is required")
|
|
}
|
|
|
|
if !strings.Contains(embeddingModel, "@") {
|
|
return nil
|
|
}
|
|
|
|
parts := strings.Split(embeddingModel, "@")
|
|
for _, part := range parts {
|
|
if strings.TrimSpace(part) == "" {
|
|
return errors.New("Both model_name and provider must be non-empty strings")
|
|
}
|
|
}
|
|
if len(parts) < 2 {
|
|
return errors.New("Embedding model identifier must follow <model_name>@<provider> format")
|
|
}
|
|
if strings.TrimSpace(parts[0]) == "" || strings.TrimSpace(parts[len(parts)-1]) == "" {
|
|
return errors.New("Both model_name and provider must be non-empty strings")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func normalizeDatasetPipelineID(pipelineID string) (*string, error) {
|
|
pipelineID = strings.TrimSpace(pipelineID)
|
|
if pipelineID == "" {
|
|
return nil, nil
|
|
}
|
|
if len(pipelineID) != 32 {
|
|
return nil, errors.New("pipeline_id must be 32 hex characters")
|
|
}
|
|
for _, char := range pipelineID {
|
|
if !strings.ContainsRune("0123456789abcdefABCDEF", char) {
|
|
return nil, errors.New("pipeline_id must be hexadecimal")
|
|
}
|
|
}
|
|
|
|
normalized := strings.ToLower(pipelineID)
|
|
return &normalized, nil
|
|
}
|
|
|
|
func validateDatasetParserConfigSize(parserConfig map[string]interface{}) error {
|
|
if len(parserConfig) == 0 {
|
|
return nil
|
|
}
|
|
|
|
data, err := json.Marshal(parserConfig)
|
|
if err != nil {
|
|
return errors.New("parser_config must be valid JSON")
|
|
}
|
|
if len(data) > 65535 {
|
|
return fmt.Errorf("Parser config exceeds size limit (max 65,535 characters). Current size: %d", len(data))
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// normalizeDatasetID canonicalizes an id into the 32-char hex form used by
|
|
// the storage layer. The "UUID1" name was a legacy term from when the
|
|
// Python service generated ids with `uuid.uuid1().hex`; the Go port uses
|
|
// `uuid.New()` (v4), so we accept any RFC 4122 version. We only reject the
|
|
// Nil UUID, which is the reserved "no id" sentinel.
|
|
func normalizeDatasetID(id string) (string, error) {
|
|
parsedUUID, err := uuid.Parse(id)
|
|
if err != nil {
|
|
return "", errors.New("Invalid UUID format")
|
|
}
|
|
if parsedUUID == (uuid.UUID{}) {
|
|
return "", errors.New("Invalid UUID format")
|
|
}
|
|
return strings.ReplaceAll(parsedUUID.String(), "-", ""), nil
|
|
}
|
|
|
|
func (d *DatasetService) verifyEmbeddingAvailability(embdID string, tenantID string) (bool, string) {
|
|
_, _, _, _, err := NewModelProviderService().ResolveModelConfig(tenantID, entity.ModelTypeEmbedding, embdID)
|
|
if err != nil {
|
|
return false, err.Error()
|
|
}
|
|
return true, ""
|
|
}
|
|
|
|
func parserConfigValueOrEmptyList(parserConfig map[string]interface{}, key string) interface{} {
|
|
if parserConfig == nil {
|
|
return []interface{}{}
|
|
}
|
|
|
|
value, ok := parserConfig[key]
|
|
if !ok || value == nil {
|
|
return []interface{}{}
|
|
}
|
|
|
|
return value
|
|
}
|
|
|
|
func normalizeMetadataConfigFields(fields []MetadataConfigField, fieldName string) ([]map[string]interface{}, error) {
|
|
normalizedFields := make([]map[string]interface{}, 0, len(fields))
|
|
for i, field := range fields {
|
|
key := strings.TrimSpace(field.Key)
|
|
if key == "" {
|
|
return nil, fmt.Errorf("%s[%d].key is required", fieldName, i)
|
|
}
|
|
if len(key) > 255 {
|
|
return nil, fmt.Errorf("%s[%d].key should have at most 255 characters", fieldName, i)
|
|
}
|
|
|
|
fieldType := strings.TrimSpace(field.Type)
|
|
if _, ok := datasetAllowedMetadataTypes[fieldType]; !ok {
|
|
return nil, fmt.Errorf("%s[%d].type should be one of 'string', 'list', 'time' or 'number'", fieldName, i)
|
|
}
|
|
|
|
if field.Description != nil && len(*field.Description) > 65535 {
|
|
return nil, fmt.Errorf("%s[%d].description should have at most 65535 characters", fieldName, i)
|
|
}
|
|
|
|
normalizedFields = append(normalizedFields, map[string]interface{}{
|
|
"key": key,
|
|
"type": fieldType,
|
|
"description": field.Description,
|
|
"enum": field.Enum,
|
|
})
|
|
}
|
|
|
|
return normalizedFields, nil
|
|
}
|
|
|
|
func datasetListItemToMap(kb *entity.KnowledgebaseListItem) map[string]interface{} {
|
|
item := map[string]interface{}{
|
|
"id": kb.ID,
|
|
"name": kb.Name,
|
|
"tenant_id": kb.TenantID,
|
|
"permission": kb.Permission,
|
|
"document_count": kb.DocNum,
|
|
"token_num": kb.TokenNum,
|
|
"chunk_count": kb.ChunkNum,
|
|
"parser_id": kb.ParserID,
|
|
"embedding_model": kb.EmbdID,
|
|
"nickname": kb.Nickname,
|
|
}
|
|
|
|
if kb.Avatar != nil {
|
|
item["avatar"] = *kb.Avatar
|
|
}
|
|
if kb.Language != nil {
|
|
item["language"] = *kb.Language
|
|
}
|
|
if kb.Description != nil {
|
|
item["description"] = *kb.Description
|
|
}
|
|
if kb.TenantAvatar != nil {
|
|
item["tenant_avatar"] = *kb.TenantAvatar
|
|
}
|
|
if kb.UpdateTime != nil {
|
|
item["update_time"] = *kb.UpdateTime
|
|
}
|
|
|
|
return item
|
|
}
|
|
|
|
func datasetToMap(kb *entity.Knowledgebase) map[string]interface{} {
|
|
item := map[string]interface{}{
|
|
"id": kb.ID,
|
|
"tenant_id": kb.TenantID,
|
|
"name": kb.Name,
|
|
"embedding_model": kb.EmbdID,
|
|
"permission": kb.Permission,
|
|
"created_by": kb.CreatedBy,
|
|
"document_count": kb.DocNum,
|
|
"token_num": kb.TokenNum,
|
|
"chunk_count": kb.ChunkNum,
|
|
"similarity_threshold": kb.SimilarityThreshold,
|
|
"vector_similarity_weight": kb.VectorSimilarityWeight,
|
|
"parser_id": kb.ParserID,
|
|
"parser_config": kb.ParserConfig,
|
|
"pagerank": kb.Pagerank,
|
|
"create_time": kb.CreateTime,
|
|
}
|
|
|
|
if kb.Avatar != nil {
|
|
item["avatar"] = *kb.Avatar
|
|
}
|
|
if kb.Language != nil {
|
|
item["language"] = *kb.Language
|
|
}
|
|
if kb.Description != nil {
|
|
item["description"] = *kb.Description
|
|
}
|
|
if kb.PipelineID != nil {
|
|
item["pipeline_id"] = *kb.PipelineID
|
|
}
|
|
if kb.GraphragTaskID != nil {
|
|
item["graphrag_task_id"] = *kb.GraphragTaskID
|
|
}
|
|
if kb.GraphragTaskFinishAt != nil {
|
|
item["graphrag_task_finish_at"] = kb.GraphragTaskFinishAt.Format("2006-01-02 15:04:05")
|
|
}
|
|
if kb.RaptorTaskID != nil {
|
|
item["raptor_task_id"] = *kb.RaptorTaskID
|
|
}
|
|
if kb.RaptorTaskFinishAt != nil {
|
|
item["raptor_task_finish_at"] = kb.RaptorTaskFinishAt.Format("2006-01-02 15:04:05")
|
|
}
|
|
if kb.MindmapTaskID != nil {
|
|
item["mindmap_task_id"] = *kb.MindmapTaskID
|
|
}
|
|
if kb.MindmapTaskFinishAt != nil {
|
|
item["mindmap_task_finish_at"] = kb.MindmapTaskFinishAt.Format("2006-01-02 15:04:05")
|
|
}
|
|
if kb.UpdateTime != nil {
|
|
item["update_time"] = *kb.UpdateTime
|
|
}
|
|
|
|
return item
|
|
}
|
|
|
|
func limitStrings(values []string, limit int) []string {
|
|
if len(values) <= limit {
|
|
return values
|
|
}
|
|
return values[:limit]
|
|
}
|
|
|
|
func (d *DatasetService) RenameTag(datasetID, userID, fromTag, toTag string) (map[string]interface{}, common.ErrorCode, error) {
|
|
fromTag = strings.TrimSpace(fromTag)
|
|
toTag = strings.TrimSpace(toTag)
|
|
|
|
datasetID, err := normalizeDatasetID(datasetID)
|
|
if err != nil {
|
|
return nil, common.CodeDataError, err
|
|
}
|
|
if strings.TrimSpace(datasetID) == "" {
|
|
return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"")
|
|
}
|
|
if !d.kbDAO.Accessible(datasetID, userID) {
|
|
return nil, common.CodeDataError, errors.New("No authorization.")
|
|
}
|
|
if d.docEngine == nil {
|
|
return nil, common.CodeServerError, errors.New("Document engine is not initialized")
|
|
}
|
|
|
|
kb, err := d.kbDAO.GetByID(datasetID)
|
|
if err != nil || kb == nil {
|
|
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
|
|
}
|
|
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
|
|
|
|
condition := map[string]interface{}{
|
|
"tag_kwd": fromTag,
|
|
"kb_id": datasetID,
|
|
}
|
|
newValue := map[string]interface{}{
|
|
"remove": map[string]interface{}{
|
|
"tag_kwd": fromTag,
|
|
},
|
|
"add": map[string]interface{}{
|
|
"tag_kwd": toTag,
|
|
},
|
|
}
|
|
|
|
err = d.docEngine.UpdateChunks(context.Background(), condition, newValue, indexName, datasetID)
|
|
if err != nil {
|
|
return nil, common.CodeServerError, fmt.Errorf("failed to rename tag: %w", err)
|
|
}
|
|
|
|
return map[string]interface{}{
|
|
"from": fromTag,
|
|
"to": toTag,
|
|
}, common.CodeSuccess, nil
|
|
}
|
|
|
|
func (d *DatasetService) GetFieldMap(ids []string) (map[string]interface{}, error) {
|
|
return d.kbDAO.GetFieldMap(ids)
|
|
}
|
|
|
|
// resolveComponentParamsDefaults loads the DSL for the target pipeline and
|
|
// returns the component params defaults as an entity.JSONMap {cpnID: {param: value}}.
|
|
// For builtin templates the DSL is loaded from the embedded registry; for custom
|
|
// canvas pipelines it is loaded from the canvas row in the database.
|
|
func resolveComponentParamsDefaults(parserID string, pipelineID *string) (entity.JSONMap, error) {
|
|
isCanvas := pipelineID != nil && strings.TrimSpace(*pipelineID) != ""
|
|
var cp map[string]map[string]any
|
|
var err error
|
|
if isCanvas {
|
|
dslJSON, lerr := loadCanvasDSLJSON(strings.TrimSpace(*pipelineID))
|
|
if lerr != nil {
|
|
return nil, fmt.Errorf("load canvas DSL: %w", lerr)
|
|
}
|
|
cp, err = pipelinepkg.ComponentParamsDefaults(dslJSON)
|
|
} else {
|
|
registry, regErr := pipelinepkg.DefaultRegistry()
|
|
if regErr != nil {
|
|
return nil, fmt.Errorf("builtin registry: %w", regErr)
|
|
}
|
|
if !registry.IsValid(parserID) {
|
|
return nil, fmt.Errorf("unknown builtin parser_id: %q", parserID)
|
|
}
|
|
dslStr, dslErr := pipelinepkg.LoadBuiltinDSL(parserID)
|
|
if dslErr != nil {
|
|
return nil, fmt.Errorf("load builtin DSL: %w", dslErr)
|
|
}
|
|
cp, err = pipelinepkg.ComponentParamsDefaults([]byte(dslStr))
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out := make(entity.JSONMap, len(cp))
|
|
for k, v := range cp {
|
|
out[k] = v
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// canvasAccessibleForUser reports whether the given canvas is owned by or
|
|
// team-shared with the caller.
|
|
func canvasAccessibleForUser(userID, canvasID string) (bool, error) {
|
|
tenantIDs, _ := dao.NewUserTenantDAO().GetTenantIDsByUserID(userID)
|
|
return dao.NewUserCanvasDAO().Accessible(canvasID, userID, tenantIDs), nil
|
|
}
|