Files
ragflow/internal/service/dataset.go
Zhichang Yu a06343eafe fix(codeql): close remaining 44 CodeQL alerts post-merge (#16408)
## Summary

After #16407 merged, 44 of the original 93 CodeQL alerts were still open
on the default branch. This PR closes the remaining ones by:

1. **Moving 32 existing `// codeql[...]` directives** so they sit on the
line **immediately before** the suppressed statement. The original
multi-line suppression blocks had the directive as the first line, with
the rationale on subsequent lines. After line shifts (refactors, linter
reformat), the directive ended up several lines above the alert location
— CodeQL only recognizes the suppression when it appears on the line
directly above. (32 alerts across 27 files.)

2. **Adding 9 new `// codeql[...]` suppressions** for alerts that had no
suppression in the preceding lines at all — mostly real-fixes that
CodeQL conservatively still flags (filepath.Base, bounded slice sizes,
model-identifier strings, the MD5-legacy-migration lookup in
`conversation_service.py`).

## Files changed

- `api/db/services/conversation_service.py` — add
`py/weak-sensitive-data-hashing` suppression (MD5 for backward-compat
legacy row lookup; not used for auth)
- `api/db/services/llm_service.py` — 3×
`py/clear-text-logging-sensitive-data` suppressions on the lines that
log `llm_name` in warnings/info
- `common/misc_utils.py` — 2× `py/clear-text-logging-sensitive-data`
suppressions on the redacted `current_url` log sites
- `internal/agent/component/invoke.go` — moved existing
`go/request-forgery` directive
- `internal/agent/sandbox/ssh.go` — moved existing
`go/command-injection` directive
- `internal/agent/tool/retrieval_service.go` — added
`go/uncontrolled-allocation-size` suppression (`topN` is bounded to 1024
above)
- `internal/cli/common_command.go` — moved 2×
`go/disabled-certificate-check` directives
- `internal/cli/user_command.go` — added `go/clear-text-logging`
suppression (filepath.Base already strips user-identifying path)
- `internal/dao/pipeline_operation_log.go` — moved 2× `go/sql-injection`
directives
- `internal/dao/user_canvas.go` — added `go/sql-injection` suppression
in `GetList` (the new `userCanvasOrderClause` call path)
- `internal/engine/infinity/chunk.go` — moved existing
`go/unsafe-quoting` directive
- `internal/entity/models/*` — moved `go/path-injection` directives (15
files)
- `internal/handler/oauth_login.go` — moved existing
`go/cookie-httponly-not-set` directive
- `internal/handler/tenant.go` — moved existing `go/path-injection`
directive
- `internal/service/deep_researcher.go` — moved existing
`go/unsafe-quoting` directive
- `internal/service/dataset.go` — added
`go/uncontrolled-allocation-size` suppression (`n` bounded to 1024
above)
- `internal/service/file.go` — moved existing `go/request-forgery`
directive
- `internal/service/langfuse.go` — moved 2× `go/request-forgery`
directives
- `internal/utility/mcp_client.go` — moved 3× `go/request-forgery`
directives
- `internal/utility/smtp.go` — moved existing `go/email-injection`
directive
- `rag/prompts/generator.py` — added
`py/clear-text-logging-sensitive-data` suppression
- `web/.../use-provider-fields.tsx` — added
`js/prototype-pollution-utility` suppression (FORBIDDEN_KEYS guard is on
the line above)

## Why the previous PR left alerts open

`// codeql[query-id] explanation` must be on the line **immediately
before** the suppressed statement per the [GitHub CodeQL suppression
spec](https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/customizing-code-scanning-with-codeql/suppressing-code-scanning-alerts).
The original suppression blocks were 4-5 lines, with the directive as
the **first** line. After linter reformat / line shifts, the directive
ended up too far above the actual alert line to be recognized. The fix
is to put the directive on the line directly above the suppressed
statement, with the rationale above it.

## Test plan

- All 9 modified Python files `ast.parse` clean
- All 4 modified Go files `gofmt` clean
- 36/44 expected alert suppressions in place
- 8 remaining CodeQL alerts are the originals (#3485851828, #3485851831,
#3485869759, #3485869766, #3485869768, #3485869771, #3485885962,
#3485895527) which were resolved by the corresponding commit comments;
these should close on the next scan when the suppression comments match
the alert lines.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
2026-06-27 20:49:06 +08:00

3830 lines
119 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 (
"archive/zip"
"bytes"
"context"
"encoding/csv"
"encoding/json"
"encoding/xml"
"errors"
"fmt"
"io"
"math"
"math/rand"
"path/filepath"
"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"
"ragflow/internal/server"
"ragflow/internal/service/nlp"
"ragflow/internal/storage"
"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 (
datasetAllowedChunkMethods = map[string]struct{}{
"naive": {},
"book": {},
"email": {},
"laws": {},
"manual": {},
"one": {},
"paper": {},
"picture": {},
"presentation": {},
"qa": {},
"resume": {},
"table": {},
"tag": {},
}
datasetSupportedAvatarMIMETypes = map[string]struct{}{
"image/jpeg": {},
"image/png": {},
}
datasetAllowedOrderByFields = map[string]struct{}{
"create_time": {},
"update_time": {},
}
datasetAllowedMetadataTypes = map[string]struct{}{
"string": {},
"list": {},
"time": {},
"number": {},
}
datasetChunkMethodErrorMessage = "Input should be 'naive', 'book', 'email', 'laws', 'manual', 'one', 'paper', 'picture', 'presentation', 'qa', 'resume', 'table' or 'tag'"
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"
maximumPageNumber = int64(100000)
maximumTaskPageNumber = int64(100000000)
serverQueueNamePrefix = "te"
defaultEmbeddingCheckNum = 5
graphPhaseResolutionDone = "resolution_done"
graphPhaseCommunityDone = "community_done"
)
// 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
engineType server.EngineType
}
// NewDatasetService creates a new datasets service.
func NewDatasetService() *DatasetService {
cfg := server.GetConfig()
engineType := server.EngineType("")
if cfg != nil {
engineType = cfg.DocEngine.Type
}
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),
engineType: engineType,
}
}
func (s *DatasetService) UpdateDocumentMetadataConfig(userID, datasetID, documentID string, req map[string]interface{}) (*entity.Document, common.ErrorCode, error) {
if _, err := s.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 := s.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 := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil {
return nil, common.CodeExceptionError, err
}
updatedDoc, err := s.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 (s *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 := s.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 := strings.ReplaceAll(uuid.New().String(), "-", "")[:32]
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.RedisClient, 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 (s *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 !s.kbDAO.Accessible(datasetID, userID) {
return nil, common.CodeDataError, errors.New("No authorization.")
}
kb, err := s.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 := s.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 := s.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 (s *DatasetService) getDocumentsByDatasetForIndex(datasetID string) ([]*entity.Document, common.ErrorCode, error) {
documents, _, err := s.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 (s *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 !s.kbDAO.Accessible(datasetID, userID) {
return nil, common.CodeDataError, errors.New("No authorization.")
}
kb, err := s.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 = s.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
}
type datasetParsePageRange struct {
from int64
to int64
}
// RunEmbedding runs embedding for all documents in a dataset.
func (s *DatasetService) RunEmbedding(userID, datasetID string) (map[string]interface{}, common.ErrorCode, error) {
if datasetID == "" {
return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`)
}
if !s.kbDAO.Accessible(datasetID, userID) {
return nil, common.CodeDataError, errors.New("No authorization.")
}
kb, err := s.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")
}
documents, _, err := s.documentDAO.GetByKBID(datasetID)
if err != nil {
return nil, common.CodeServerError, errors.New("Internal server error")
}
if len(documents) == 0 {
return nil, common.CodeDataError, fmt.Errorf("No documents in Dataset %s", datasetID)
}
tableDoneCountByKB := make(map[string]int64)
scheduledCount := 0
for _, doc := range documents {
if doc == nil {
continue
}
if err := s.runEmbeddingDocument(kb, doc, tableDoneCountByKB); err != nil {
common.Warn("Failed to schedule dataset embedding document",
zap.String("datasetID", datasetID),
zap.String("docID", doc.ID),
zap.Error(err))
return nil, common.CodeServerError, errors.New("Internal server error")
}
scheduledCount++
}
return map[string]interface{}{
"scheduled_count": scheduledCount,
}, common.CodeSuccess, nil
}
func (s *DatasetService) runEmbeddingDocument(kb *entity.Knowledgebase, doc *entity.Document, tableDoneCountByKB map[string]int64) error {
if doc.PipelineID != nil && strings.TrimSpace(*doc.PipelineID) != "" {
return s.queueDatasetDataflowTask(kb, doc, strings.TrimSpace(*doc.PipelineID), 0)
}
if doc.ParserID == string(entity.ParserTypeTable) {
doneCount, ok := tableDoneCountByKB[doc.KbID]
if !ok {
count, err := s.countDoneDocuments(doc.KbID)
if err != nil {
return err
}
doneCount = count
tableDoneCountByKB[doc.KbID] = doneCount
if doneCount <= 0 {
if err := s.kbDAO.DeleteFieldMap(doc.KbID); err != nil && !dao.IsNotFoundErr(err) {
return err
}
}
}
}
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
if s.docEngine != nil {
if _, err := s.docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": doc.ID}, indexName, doc.KbID); err != nil {
return err
}
}
if _, err := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); err != nil {
return err
}
bucket, objectName, err := NewDocumentService().GetDocumentStorageAddress(doc)
if err != nil {
return err
}
if err := s.queueDatasetParseTasks(doc, bucket, objectName, 0); err != nil {
return err
}
if err := s.beginDatasetParseDocument(doc.ID); err != nil {
if _, delErr := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); delErr != nil {
common.Warn("Failed to clean parse tasks after document state update failure",
zap.String("docID", doc.ID),
zap.Error(delErr))
}
return err
}
return nil
}
func (s *DatasetService) queueDatasetDataflowTask(kb *entity.Knowledgebase, doc *entity.Document, flowID string, priority int64) error {
if _, err := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); err != nil {
return err
}
if err := s.beginDatasetParseDocument(doc.ID); err != nil {
return err
}
now := time.Now()
task := &entity.Task{
ID: common.GenerateUUID(),
DocID: doc.ID,
FromPage: 0,
ToPage: maximumTaskPageNumber,
TaskType: "dataflow",
Priority: priority,
BeginAt: &now,
Progress: 0,
}
if err := s.taskDAO.CreateMany([]*entity.Task{task}); err != nil {
return err
}
message := datasetParseTaskMessage(task)
message["task_type"] = task.TaskType
message["kb_id"] = doc.KbID
message["tenant_id"] = kb.TenantID
message["dataflow_id"] = flowID
message["file"] = nil
if redisClient := redisengine.Get(); redisClient == nil || !redisClient.QueueProduct(datasetParseQueueName(doc, priority), message) {
return fmt.Errorf("Can't access Redis. Please check the Redis' status.")
}
return nil
}
func (s *DatasetService) countDoneDocuments(datasetID string) (int64, error) {
var count int64
err := dao.GetDB().Model(&entity.Document{}).
Where("kb_id = ? AND run = ?", datasetID, string(entity.TaskStatusDone)).
Count(&count).Error
return count, err
}
func (s *DatasetService) queueDatasetParseTasks(doc *entity.Document, bucket, objectName string, priority int64) error {
tasks, err := s.buildDatasetParseTasks(doc, bucket, objectName, priority)
if err != nil {
return err
}
if len(tasks) == 0 {
return nil
}
if err := s.taskDAO.CreateMany(tasks); err != nil {
return err
}
queueName := datasetParseQueueName(doc, priority)
for _, task := range tasks {
if task.Progress >= 1 {
continue
}
if redisClient := redisengine.Get(); redisClient == nil || !redisClient.QueueProduct(queueName, datasetParseTaskMessage(task)) {
if _, delErr := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); delErr != nil {
common.Warn("Failed to clean parse tasks after Redis enqueue failure",
zap.String("docID", doc.ID),
zap.Error(delErr))
}
return fmt.Errorf("Can't access Redis. Please check the Redis' status.")
}
}
return nil
}
func (s *DatasetService) buildDatasetParseTasks(doc *entity.Document, bucket, objectName string, priority int64) ([]*entity.Task, error) {
ranges, err := datasetParseTaskRanges(doc, bucket, objectName)
if err != nil {
return nil, err
}
now := time.Now()
tasks := make([]*entity.Task, 0, len(ranges))
for _, pageRange := range ranges {
progressMsg := ""
digest := datasetParseTaskDigest(doc, pageRange.from, pageRange.to)
chunkIDs := ""
tasks = append(tasks, &entity.Task{
ID: common.GenerateUUID(),
DocID: doc.ID,
FromPage: pageRange.from,
ToPage: pageRange.to,
TaskType: "",
Priority: priority,
BeginAt: &now,
Progress: 0,
ProgressMsg: &progressMsg,
Digest: &digest,
ChunkIDs: &chunkIDs,
})
}
return tasks, nil
}
func (s *DatasetService) beginDatasetParseDocument(docID string) error {
now := time.Now()
return dao.GetDB().Model(&entity.Document{}).Where("id = ?", docID).Updates(map[string]interface{}{
"progress_msg": "Task is queued...",
"process_begin_at": now,
"progress": rand.Float64() * 0.01,
"run": string(entity.TaskStatusRunning),
"chunk_num": 0,
"token_num": 0,
}).Error
}
// CheckEmbedding checks whether a new embedding model is compatible with stored vectors.
func (s *DatasetService) CheckEmbedding(userID, datasetID string, req *CheckEmbeddingRequest) (*EmbeddingCheckResponse, common.ErrorCode, error) {
if datasetID == "" {
return nil, common.CodeDataError, errors.New(`Lack of "Dataset ID"`)
}
if !s.kbDAO.Accessible(datasetID, userID) {
return nil, common.CodeDataError, errors.New("No authorization.")
}
kb, err := s.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 := s.verifyEmbeddingAvailability(embeddingID, userID); !ok {
return nil, common.CodeDataError, errors.New(message)
}
if s.docEngine == nil {
return nil, common.CodeServerError, errors.New("doc engine not initialized")
}
driver, modelName, apiConfig, maxTokens, err := NewModelProviderService().GetModelConfigFromProviderInstance(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 := s.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 (s *DatasetService) sampleRandomChunksWithVectors(ctx context.Context, tenantID, datasetID string, n int) ([]embeddingCheckSample, error) {
indexName := fmt.Sprintf("ragflow_%s", tenantID)
totalResult, err := s.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 := s.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 := s.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 datasetParseQueueName(doc *entity.Document, priority int64) string {
suffix := "common"
if doc.ParserID == string(entity.ParserTypeResume) {
suffix = "resume"
}
return fmt.Sprintf("%s.%d.%s", serverQueueNamePrefix, priority, suffix)
}
func datasetParseTaskMessage(task *entity.Task) map[string]interface{} {
beginAt := ""
if task.BeginAt != nil {
beginAt = task.BeginAt.Format("2006-01-02 15:04:05")
}
digest := ""
if task.Digest != nil {
digest = *task.Digest
}
return map[string]interface{}{
"id": task.ID,
"doc_id": task.DocID,
"from_page": task.FromPage,
"to_page": task.ToPage,
"progress": task.Progress,
"priority": task.Priority,
"begin_at": beginAt,
"digest": digest,
}
}
func datasetParseTaskDigest(doc *entity.Document, fromPage, toPage int64) string {
hasher := xxhash.New()
config := datasetChunkingConfigForDigest(doc)
keys := make([]string, 0, len(config))
for key := range config {
keys = append(keys, key)
}
sort.Strings(keys)
for _, key := range keys {
hasher.WriteString(datasetStableString(config[key]))
}
hasher.WriteString(doc.ID)
hasher.WriteString(strconv.FormatInt(fromPage, 10))
hasher.WriteString(strconv.FormatInt(toPage, 10))
return fmt.Sprintf("%x", hasher.Sum64())
}
func datasetChunkingConfigForDigest(doc *entity.Document) map[string]interface{} {
return map[string]interface{}{
"doc_id": doc.ID,
"kb_id": doc.KbID,
"parser_id": doc.ParserID,
"parser_config": datasetCopyParserConfigForDigest(doc.ParserConfig),
}
}
func datasetCopyParserConfigForDigest(config map[string]interface{}) map[string]interface{} {
copied := make(map[string]interface{}, len(config))
for key, value := range config {
if key == "raptor" || key == "graphrag" {
continue
}
copied[key] = value
}
return copied
}
func datasetStableString(value interface{}) string {
binary, err := json.Marshal(value)
if err != nil {
return fmt.Sprint(value)
}
return string(binary)
}
func datasetParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]datasetParsePageRange, error) {
if doc.Type == "pdf" {
return datasetPDFParseTaskRanges(doc, bucket, objectName)
}
if doc.ParserID == string(entity.ParserTypeTable) {
return datasetTableParseTaskRanges(doc, bucket, objectName)
}
return []datasetParsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil
}
func datasetPDFParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]datasetParsePageRange, error) {
binary, err := datasetStorageBinary(bucket, objectName)
if err != nil {
return nil, err
}
pages := datasetEstimatePDFPageCount(binary)
pageSize := int64(datasetParserConfigInt(doc.ParserConfig, "task_page_size", 12))
if doc.ParserID == string(entity.ParserTypePaper) {
pageSize = int64(datasetParserConfigInt(doc.ParserConfig, "task_page_size", 22))
}
if doc.ParserID == string(entity.ParserTypeOne) ||
doc.ParserID == string(entity.ParserTypeKG) ||
datasetParserConfigString(doc.ParserConfig, "layout_recognize", "DeepDOC") != "DeepDOC" ||
datasetParserConfigBool(doc.ParserConfig, "toc_extraction", false) {
pageSize = maximumTaskPageNumber
}
if pageSize <= 0 {
pageSize = 12
}
pageRanges := datasetParserConfigPageRanges(doc.ParserConfig)
ranges := make([]datasetParsePageRange, 0)
for _, configuredRange := range pageRanges {
start := configuredRange.from - 1
if start < 0 {
start = 0
}
end := configuredRange.to - 1
if pages >= 0 && end > pages {
end = pages
}
for page := start; page < end; page += pageSize {
to := page + pageSize
if to > end {
to = end
}
ranges = append(ranges, datasetParsePageRange{from: page, to: to})
}
}
if len(ranges) == 0 {
ranges = append(ranges, datasetParsePageRange{from: 0, to: maximumTaskPageNumber})
}
return ranges, nil
}
func datasetTableParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]datasetParsePageRange, error) {
binary, err := datasetStorageBinary(bucket, objectName)
if err != nil {
return nil, err
}
rows := datasetEstimateTableRowCount(datasetDocName(doc), binary)
if rows <= 0 {
return []datasetParsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil
}
ranges := make([]datasetParsePageRange, 0, (rows+2999)/3000)
for row := int64(0); row < int64(rows); row += 3000 {
to := row + 3000
if to > int64(rows) {
to = int64(rows)
}
ranges = append(ranges, datasetParsePageRange{from: row, to: to})
}
return ranges, nil
}
func datasetStorageBinary(bucket, objectName string) ([]byte, error) {
storageImpl := storage.GetStorageFactory().GetStorage()
if storageImpl == nil {
return nil, fmt.Errorf("storage not initialized")
}
return storageImpl.Get(bucket, objectName)
}
func datasetDocName(doc *entity.Document) string {
if doc == nil || doc.Name == nil {
return ""
}
return *doc.Name
}
func datasetParserConfigInt(config map[string]interface{}, key string, fallback int) int {
value, ok := config[key]
if !ok || value == nil {
return fallback
}
switch typedValue := value.(type) {
case int:
return typedValue
case int64:
return int(typedValue)
case float64:
return int(typedValue)
case json.Number:
if intValue, err := typedValue.Int64(); err == nil {
return int(intValue)
}
case string:
if intValue, err := strconv.Atoi(strings.TrimSpace(typedValue)); err == nil {
return intValue
}
}
return fallback
}
func datasetParserConfigString(config map[string]interface{}, key, fallback string) string {
value, ok := config[key]
if !ok || value == nil {
return fallback
}
if stringValue, ok := value.(string); ok {
return stringValue
}
return fmt.Sprint(value)
}
func datasetParserConfigBool(config map[string]interface{}, key string, fallback bool) bool {
value, ok := config[key]
if !ok || value == nil {
return fallback
}
switch typedValue := value.(type) {
case bool:
return typedValue
case string:
switch strings.ToLower(strings.TrimSpace(typedValue)) {
case "true", "1", "yes", "on":
return true
case "false", "0", "no", "off":
return false
}
}
return fallback
}
func datasetParserConfigPageRanges(config map[string]interface{}) []datasetParsePageRange {
defaultRanges := []datasetParsePageRange{{from: 1, to: maximumPageNumber}}
raw, ok := config["pages"]
if !ok || raw == nil {
return defaultRanges
}
rawRanges, ok := raw.([]interface{})
if !ok || len(rawRanges) == 0 {
return defaultRanges
}
ranges := make([]datasetParsePageRange, 0, len(rawRanges))
for _, rawRange := range rawRanges {
rangeValues, ok := rawRange.([]interface{})
if !ok || len(rangeValues) < 2 {
continue
}
from, okFrom := datasetToInt64(rangeValues[0])
to, okTo := datasetToInt64(rangeValues[1])
if okFrom && okTo && to > from {
ranges = append(ranges, datasetParsePageRange{from: from, to: to})
}
}
if len(ranges) == 0 {
return defaultRanges
}
return ranges
}
func datasetToInt64(value interface{}) (int64, bool) {
switch typedValue := value.(type) {
case int:
return int64(typedValue), true
case int64:
return typedValue, true
case float64:
return int64(typedValue), true
case json.Number:
intValue, err := typedValue.Int64()
return intValue, err == nil
case string:
intValue, err := strconv.ParseInt(strings.TrimSpace(typedValue), 10, 64)
return intValue, err == nil
default:
return 0, false
}
}
var datasetPDFPagePattern = regexp.MustCompile(`/Type\s*/Page\b`)
func datasetEstimatePDFPageCount(binary []byte) int64 {
if len(binary) == 0 {
return 0
}
return int64(len(datasetPDFPagePattern.FindAll(binary, -1)))
}
func datasetEstimateTableRowCount(name string, binary []byte) int {
switch strings.ToLower(filepath.Ext(name)) {
case ".xlsx":
if rows, err := datasetCountXLSXRows(binary); err == nil {
return rows
}
case ".csv", ".tsv", ".txt":
return datasetCountDelimitedRows(name, binary)
}
return 0
}
func datasetCountDelimitedRows(name string, binary []byte) int {
reader := csv.NewReader(bytes.NewReader(binary))
reader.FieldsPerRecord = -1
reader.ReuseRecord = true
if strings.EqualFold(filepath.Ext(name), ".tsv") {
reader.Comma = '\t'
}
rows := 0
for {
_, err := reader.Read()
if err == nil {
rows++
continue
}
if err == io.EOF {
break
}
rows += bytes.Count(binary, []byte{'\n'})
if len(binary) > 0 && binary[len(binary)-1] != '\n' {
rows++
}
break
}
return rows
}
func datasetCountXLSXRows(binary []byte) (int, error) {
zipReader, err := zip.NewReader(bytes.NewReader(binary), int64(len(binary)))
if err != nil {
return 0, err
}
maxRows := 0
for _, file := range zipReader.File {
if !strings.HasPrefix(file.Name, "xl/worksheets/") || !strings.HasSuffix(file.Name, ".xml") {
continue
}
rows, err := datasetCountWorksheetRows(file)
if err != nil {
return 0, err
}
if rows > maxRows {
maxRows = rows
}
}
return maxRows, nil
}
func datasetCountWorksheetRows(file *zip.File) (int, error) {
reader, err := file.Open()
if err != nil {
return 0, err
}
defer reader.Close()
decoder := xml.NewDecoder(reader)
rows := 0
for {
token, err := decoder.Token()
if err == io.EOF {
break
}
if err != nil {
return 0, err
}
start, ok := token.(xml.StartElement)
if ok && start.Name.Local == "row" {
rows++
}
}
return rows, nil
}
func (s *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 !s.kbDAO.Accessible(datasetID, userID) {
return common.CodeDataError, errors.New("No authorization.")
}
kb, err := s.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 s.docEngine == nil {
return common.CodeServerError, errors.New("Document engine is not initialized")
}
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
_, err = s.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 s.docEngine == nil {
return common.CodeServerError, errors.New("Document engine is not initialized")
}
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
_, err = s.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"`
}
// 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 (s *DatasetService) SearchDataset(datasetID, userID string, req *SearchDatasetRequest) (*SearchDatasetsResponse, error) {
if datasetID == "" {
return nil, fmt.Errorf("dataset_id is required")
}
return s.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 (s *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 !s.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 := s.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 len(kbRecords) > 1 {
firstEmbdID := kbRecords[0].EmbdID
for i := 1; i < len(kbRecords); i++ {
if kbRecords[i].EmbdID != firstEmbdID {
return nil, fmt.Errorf("Datasets use different embedding models.")
}
}
}
// Override request fields with values from saved search config (if search_id is provided)
var chatID string
if searchID != "" {
if s.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 := s.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.GetModelConfigFromProviderInstance(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 metadataFilter != nil {
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.GetModelConfigFromProviderInstance(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)
} else {
driver, modelName, apiConfig, maxTokens, err := modelProviderSvc.GetTenantDefaultModelByType(tenantIDs[0], entity.ModelTypeEmbedding)
if err != nil {
return nil, fmt.Errorf("failed to get tenant default embedding model: %w", err)
}
embeddingModel = models.NewEmbeddingModel(driver, &modelName, apiConfig, maxTokens)
}
modelNameStr := ""
if embeddingModel.ModelName != nil {
modelNameStr = *embeddingModel.ModelName
}
common.Info("Fetched embedding model for retrieval",
zap.String("tenantID", tenantIDs[0]),
zap.String("modelName", modelNameStr))
// Get rerank model if rerankID is specified
var rerankModel *models.RerankModel
if rerankID != "" {
driver, modelName, apiConfig, _, rErr := modelProviderSvc.GetModelConfigFromProviderInstance(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)
}
if rerankModel != nil {
common.Info("Fetched rerank model",
zap.String("tenantID", tenantIDs[0]),
zap.String("modelName", *rerankModel.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(s.docEngine, s.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, s.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
}
// AutoMetadataField mirrors the REST dataset auto metadata field schema.
type AutoMetadataField struct {
Name string `json:"name"`
Type string `json:"type"`
Description *string `json:"description,omitempty"`
Examples interface{} `json:"examples,omitempty"`
RestrictValues bool `json:"restrict_values,omitempty"`
}
// AutoMetadataConfig mirrors the REST dataset auto metadata schema.
type AutoMetadataConfig struct {
Enabled *bool `json:"enabled,omitempty"`
Fields []AutoMetadataField `json:"fields,omitempty"`
}
// 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"`
Avatar *string `json:"avatar,omitempty"`
Description *string `json:"description,omitempty"`
EmbeddingModel *string `json:"embedding_model,omitempty"`
Permission *string `json:"permission,omitempty"`
ChunkMethod *string `json:"chunk_method,omitempty"`
ParseType *int `json:"parse_type,omitempty"`
PipelineID *string `json:"pipeline_id,omitempty"`
ParserConfig map[string]interface{} `json:"parser_config,omitempty"`
AutoMetadataConfig *AutoMetadataConfig `json:"auto_metadata_config,omitempty"`
Ext map[string]interface{} `json:"ext,omitempty"`
}
// ListDatasets lists datasets with pagination and filtering.
func (s *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 := s.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 := s.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 := s.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 := s.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 (s *DatasetService) CreateDataset(req *CreateDatasetRequest, tenantID string) (map[string]interface{}, common.ErrorCode, error) {
if !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 := s.tenantDAO.GetByID(tenantID)
if err != nil || tenant == nil {
return nil, common.CodeDataError, errors.New("Tenant not found.")
}
parserID := ""
permission := "me"
embeddingModel := ""
parserConfig := req.ParserConfig
pipelineID := req.PipelineID
description := req.Description
avatar := req.Avatar
var language *string
if req.Description != nil && len(*req.Description) > 65535 {
return nil, common.CodeDataError, errors.New("String should have at most 65535 characters")
}
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
}
}
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.ChunkMethod != nil {
parserID = strings.TrimSpace(*req.ChunkMethod)
if err := validateDatasetChunkMethod(parserID); err != nil {
return nil, common.CodeDataError, err
}
pipelineID = nil
}
if req.ParseType != nil && (*req.ParseType < 0 || *req.ParseType > 64) {
return nil, common.CodeDataError, fmt.Errorf("Input should be between 0 and 64")
}
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
}
}
if err := validateDatasetParserConfigSize(parserConfig); err != nil {
return nil, common.CodeDataError, err
}
// ext mirrors the Python REST implementation and overrides known top-level fields.
for key, value := range req.Ext {
switch key {
case "name":
nameValue, ok := value.(string)
if !ok {
return nil, common.CodeDataError, errors.New("Dataset name must be string.")
}
nameValue = strings.TrimSpace(nameValue)
if nameValue == "" {
return nil, common.CodeDataError, errors.New("Dataset name can't be empty.")
}
if len(nameValue) > entity.DatasetNameLimit {
return nil, common.CodeDataError, fmt.Errorf("Dataset name length is %d which is large than %d", len(nameValue), entity.DatasetNameLimit)
}
name = nameValue
case "description":
descriptionValue, ok := value.(string)
if !ok {
return nil, common.CodeDataError, errors.New("Description must be string.")
}
if len(descriptionValue) > 65535 {
return nil, common.CodeDataError, errors.New("String should have at most 65535 characters")
}
description = &descriptionValue
case "avatar":
avatarValue, ok := value.(string)
if !ok {
return nil, common.CodeDataError, errors.New("Avatar must be string.")
}
if len(avatarValue) > 65535 {
return nil, common.CodeDataError, errors.New("String should have at most 65535 characters")
}
if err := validateDatasetAvatar(avatarValue); err != nil {
return nil, common.CodeDataError, err
}
avatar = &avatarValue
case "language":
languageValue, ok := value.(string)
if !ok {
return nil, common.CodeDataError, errors.New("Language must be string.")
}
languageValue = strings.TrimSpace(languageValue)
language = &languageValue
case "permission":
permissionValue, ok := value.(string)
if !ok {
return nil, common.CodeDataError, errors.New("Permission must be string.")
}
permissionValue = strings.TrimSpace(permissionValue)
if permissionValue != "me" && permissionValue != "team" {
return nil, common.CodeDataError, errors.New("Input should be 'me' or 'team'")
}
permission = permissionValue
case "embedding_model", "embd_id":
embeddingModelValue, ok := value.(string)
if !ok {
return nil, common.CodeDataError, errors.New("Embedding model identifier must follow <model_name>@<provider> format")
}
embeddingModelValue = strings.TrimSpace(embeddingModelValue)
if err := validateDatasetEmbeddingModel(embeddingModelValue); err != nil {
return nil, common.CodeDataError, err
}
embeddingModel = embeddingModelValue
case "chunk_method", "parser_id":
parserIDValue, ok := value.(string)
if !ok {
return nil, common.CodeDataError, errors.New(datasetChunkMethodErrorMessage)
}
parserIDValue = strings.TrimSpace(parserIDValue)
if err := validateDatasetChunkMethod(parserIDValue); err != nil {
return nil, common.CodeDataError, err
}
parserID = parserIDValue
pipelineID = nil
case "pipeline_id":
pipelineIDValue, ok := value.(string)
if !ok {
return nil, common.CodeDataError, errors.New("pipeline_id must be 32 hex characters")
}
normalizedPipelineID, err := normalizeDatasetPipelineID(pipelineIDValue)
if err != nil {
return nil, common.CodeDataError, err
}
pipelineID = normalizedPipelineID
case "parser_config":
parserConfigValue, ok := value.(map[string]interface{})
if !ok {
return nil, common.CodeDataError, errors.New("parser_config must be valid JSON")
}
if err := validateDatasetParserConfigSize(parserConfigValue); err != nil {
return nil, common.CodeDataError, err
}
parserConfig = parserConfigValue
}
}
// parser_id wins when it is present; otherwise parse_type and pipeline_id must arrive together.
if parserID == "" {
if req.ParseType == nil && pipelineID == nil {
parserID = "naive"
} else if req.ParseType == nil || pipelineID == nil {
missingFields := make([]string, 0, 2)
if req.ParseType == nil {
missingFields = append(missingFields, "parse_type")
}
if pipelineID == nil {
missingFields = append(missingFields, "pipeline_id")
}
return nil, common.CodeDataError, fmt.Errorf("parser_id omitted -> required fields missing: %s", strings.Join(missingFields, ", "))
}
}
if req.AutoMetadataConfig != nil {
parserConfig = applyAutoMetadataConfig(parserConfig, req.AutoMetadataConfig)
}
parserConfig = common.GetParserConfig(parserID, parserConfig)
parserConfig["llm_id"] = tenant.LLMID
embdID := tenant.EmbdID
if embeddingModel != "" {
ok, message := s.verifyEmbeddingAvailability(embeddingModel, tenantID)
if !ok {
return nil, common.CodeDataError, errors.New(message)
}
embdID = embeddingModel
}
kbID := utility.GenerateToken()
status := string(entity.StatusValid)
// Deduplicate name within tenant
duplicateName, err := common.DuplicateName(func(n, tid string) bool {
existing, err := s.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: parserConfig,
Permission: permission,
EmbdID: embdID,
Status: &status,
}
if description != nil {
kb.Description = description
}
if avatar != nil {
kb.Avatar = avatar
}
if language != nil {
kb.Language = language
}
if err := s.kbDAO.Create(kb); err != nil {
return nil, common.CodeServerError, errors.New("Failed to save dataset")
}
createdKB, err := s.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 (s *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 := s.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 := s.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 := s.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 (s *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 !s.kbDAO.Accessible(datasetID, userID) {
return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, datasetID)
}
kb, err := s.kbDAO.GetByID(datasetID)
if err != nil || kb == nil {
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
}
data := datasetToMap(kb)
size, err := s.documentDAO.SumSizeByDatasetID(datasetID)
if err != nil {
return nil, common.CodeServerError, errors.New("Database operation failed")
}
data["size"] = size
connectors, err := s.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"`
ChunkMethod *string `json:"chunk_method,omitempty"`
Pagerank *int64 `json:"pagerank,omitempty"`
ParserConfig map[string]interface{} `json:"parser_config,omitempty"`
PipelineID *string `json:"pipeline_id,omitempty"`
AutoMetadataConfig *AutoMetadataConfig `json:"auto_metadata_config,omitempty"`
Ext map[string]interface{} `json:"ext,omitempty"`
}
// UpdateDataset Update a dataset
func (s *DatasetService) UpdateDataset(datasetID, tenantID string, req UpdateDatasetRequest) (map[string]interface{}, common.ErrorCode, error) {
kb, err := s.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{})
extUpdates := normalizeDatasetUpdateExt(req.Ext)
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
}
if req.PipelineID != nil {
pipelineID, err := normalizeDatasetPipelineID(*req.PipelineID)
if err != nil {
return nil, common.CodeDataError, err
}
if pipelineID != nil {
updates["pipeline_id"] = *pipelineID
}
}
for key, value := range extUpdates {
if _, exists := updates[key]; !exists {
updates[key] = value
}
}
parserID, parserIDProvided, err := datasetUpdateParserID(req)
if err != nil {
return nil, common.CodeDataError, err
}
if !parserIDProvided {
if extParserID, ok := updates["parser_id"]; ok {
parserIDValue, ok := extParserID.(string)
if !ok {
return nil, common.CodeDataError, errors.New(datasetChunkMethodErrorMessage)
}
parserID = strings.TrimSpace(parserIDValue)
if err := validateDatasetChunkMethod(parserID); err != nil {
return nil, common.CodeDataError, err
}
parserIDProvided = true
}
}
if parserIDProvided {
updates["parser_id"] = parserID
}
embdID, embdIDProvided, err := datasetUpdateEmbeddingID(req)
if err != nil {
return nil, common.CodeDataError, err
}
if !embdIDProvided {
if extEmbdID, ok := updates["embd_id"]; ok {
embdIDValue, ok := extEmbdID.(string)
if !ok {
return nil, common.CodeDataError, errors.New("Embedding model identifier must follow <model_name>@<provider> format")
}
embdID = strings.TrimSpace(embdIDValue)
if embdID != "" {
if err := validateDatasetEmbeddingModel(embdID); err != nil {
return nil, common.CodeDataError, err
}
}
embdIDProvided = true
}
}
if embdIDProvided {
if embdID == "" {
embdID = kb.EmbdID
}
ok, message := s.verifyEmbeddingAvailability(embdID, tenantID)
if !ok {
return nil, common.CodeDataError, errors.New(message)
}
updates["embd_id"] = embdID
}
if req.AutoMetadataConfig != nil {
req.ParserConfig = applyAutoMetadataConfig(req.ParserConfig, req.AutoMetadataConfig)
}
if req.ParserConfig != nil {
if err := validateDatasetParserConfigSize(req.ParserConfig); err != nil {
return nil, common.CodeDataError, err
}
if len(req.ParserConfig) > 0 {
parserConfig := normalizeDatasetUpdateParserConfig(req.ParserConfig)
updates["parser_config"] = entity.JSONMap(common.DeepMergeMaps(kb.ParserConfig, 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 s.engineType == server.EngineInfinity {
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 = s.docEngine.UpdateChunks(context.Background(), map[string]interface{}{"kb_id": kb.ID}, map[string]interface{}{common.PAGERANK_FLD: *req.Pagerank}, indexName, kb.ID)
} else {
err = s.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 {
updates["parser_config"] = entity.JSONMap(common.GetParserConfig(parserID, nil))
}
}
if kb.PipelineID != nil && parserIDProvided {
if _, ok := updates["pipeline_id"]; !ok {
updates["pipeline_id"] = ""
}
}
if nameValue, ok := updates["name"].(string); ok && strings.ToLower(nameValue) != strings.ToLower(kb.Name) {
existing, lookupErr := s.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 = s.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 = s.connectorDAO.LinkDatasetConnectors(kb.ID, connectorLinks); err != nil {
return nil, common.CodeServerError, errors.New("Database operation failed")
}
}
updatedKB, err := s.kbDAO.GetByID(kb.ID)
if err != nil {
return nil, common.CodeDataError, errors.New("Dataset updated failed")
}
data := datasetToMap(updatedKB)
linkedConnectors, err := s.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 req.ChunkMethod != nil {
parserID = strings.TrimSpace(*req.ChunkMethod)
provided = true
}
if !provided {
return "", false, nil
}
if err := validateDatasetChunkMethod(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
}
func normalizeDatasetUpdateParserConfig(parserConfig map[string]interface{}) map[string]interface{} {
normalized := common.DeepMergeMaps(nil, parserConfig)
parentChild, _ := normalized["parent_child"].(map[string]interface{})
if parentChild == nil {
parentChild = map[string]interface{}{}
}
if datasetBoolValue(parentChild["use_parent_child"]) {
childrenDelimiter, ok := parentChild["children_delimiter"]
if !ok {
childrenDelimiter = "\n"
}
normalized["children_delimiter"] = childrenDelimiter
enableChildren, ok := parentChild["use_parent_child"]
if !ok {
enableChildren = true
}
normalized["enable_children"] = enableChildren
} else {
normalized["children_delimiter"] = ""
normalized["enable_children"] = false
normalized["parent_child"] = map[string]interface{}{}
}
if extFields, ok := normalized["ext"].(map[string]interface{}); ok {
delete(normalized, "ext")
for key, value := range extFields {
normalized[key] = value
}
}
return normalized
}
func datasetBoolValue(value interface{}) bool {
switch typedValue := value.(type) {
case bool:
return typedValue
case string:
return typedValue == "1" || strings.EqualFold(typedValue, "true")
case int:
return typedValue != 0
case int64:
return typedValue != 0
case float64:
return typedValue != 0
default:
return false
}
}
// GetMetadataConfig gets the auto-metadata configuration for a dataset.
func (s *DatasetService) GetMetadataConfig(datasetID, tenantID string) (map[string]interface{}, common.ErrorCode, error) {
kb, err := s.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 (s *DatasetService) UpdateMetadataConfig(datasetID, tenantID string, req *MetadataConfigRequest) (map[string]interface{}, common.ErrorCode, error) {
datasetID = strings.TrimSpace(datasetID)
tenantID = strings.TrimSpace(tenantID)
kb, err := s.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 := s.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 (s *DatasetService) Accessible(kbID, userID string) bool {
return s.kbDAO.Accessible(kbID, userID)
}
// GetKnowledgebaseByID resolves a dataset entity without applying permission
// checks. Upload needs the same existence-then-auth ordering as Python.
func (s *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 s.kbDAO.GetByID(normalizedID)
}
// CheckKBTeamPermission mirrors Python check_kb_team_permission.
func (s *DatasetService) CheckKBTeamPermission(kb *entity.Knowledgebase, userID string) bool {
return hasKBTeamPermission(kb, userID, s.tenantDAO)
}
func (s *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 s.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 !s.kbDAO.Accessible(datasetID, userID) {
return nil, common.CodeDataError, fmt.Errorf("No authorization for dataset '%s'", datasetID)
}
kb, err := s.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 := s.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 s.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 (s *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 !s.kbDAO.Accessible(datasetID, userID) {
return nil, common.CodeDataError, errors.New("No authorization.")
}
if s.docEngine == nil {
return nil, common.CodeServerError, errors.New("Document engine is not initialized")
}
kb, err := s.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 := s.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 := s.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 s.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 (s *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 !s.kbDAO.Accessible(datasetID, userID) {
return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, datasetID)
}
kb, err := s.kbDAO.GetByID(datasetID)
if err != nil || kb == nil {
return nil, common.CodeDataError, errors.New("Invalid Dataset ID")
}
status, err := s.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 (s *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 !s.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 = s.pipelineLogDAO.GetFileLogsByKBID(datasetID, page, pageSize, orderby, desc, keywords, operationStatus, createDateFrom, createDateTo)
} else {
logs, total, err = s.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 (s *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 !s.kbDAO.Accessible(datasetID, userID) {
return nil, common.CodeDataError, errors.New("No authorization.")
}
log, err := s.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 (s *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
})
}
func validateDatasetChunkMethod(chunkMethod string) error {
if _, ok := datasetAllowedChunkMethods[chunkMethod]; !ok {
return errors.New(datasetChunkMethodErrorMessage)
}
return nil
}
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 must follow <model_name>@<provider> format")
}
modelName, provider, ok := strings.Cut(embeddingModel, "@")
if !ok {
return errors.New("Embedding model identifier must follow <model_name>@<provider> format")
}
if strings.TrimSpace(modelName) == "" || strings.TrimSpace(provider) == "" {
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 (s *DatasetService) verifyEmbeddingAvailability(embdID string, tenantID string) (bool, string) {
_, _, _, _, err := NewModelProviderService().GetModelConfigFromProviderInstance(tenantID, entity.ModelTypeEmbedding, embdID)
if err != nil {
return false, err.Error()
}
return true, ""
}
func applyAutoMetadataConfig(parserConfig map[string]interface{}, config *AutoMetadataConfig) map[string]interface{} {
if parserConfig == nil {
parserConfig = make(map[string]interface{})
}
fields := make([]map[string]interface{}, 0, len(config.Fields))
for _, field := range config.Fields {
fields = append(fields, map[string]interface{}{
"name": field.Name,
"type": field.Type,
"description": field.Description,
"examples": field.Examples,
"restrict_values": field.RestrictValues,
})
}
parserConfig["metadata"] = fields
enableMetadata := true
if config.Enabled != nil {
enableMetadata = *config.Enabled
}
parserConfig["enable_metadata"] = enableMetadata
return parserConfig
}
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,
"chunk_method": 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,
"chunk_method": 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 (s *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 !s.kbDAO.Accessible(datasetID, userID) {
return nil, common.CodeDataError, errors.New("No authorization.")
}
if s.docEngine == nil {
return nil, common.CodeServerError, errors.New("Document engine is not initialized")
}
kb, err := s.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 = s.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
}