From 0dd0ac06f803a4da2fb15779423214dac610c62d Mon Sep 17 00:00:00 2001 From: Jack Date: Wed, 8 Jul 2026 19:08:11 +0800 Subject: [PATCH] Feature: task executor migration to go (#16549) ### Summary Feature: Integrate parser --- .gitignore | 4 + cmd/ragflow_server.go | 19 +- internal/binding/rag_analyzer.go | 4 +- internal/binding/rag_analyzer_nocgo.go | 14 + internal/dao/document.go | 11 + internal/dao/ingestion.go | 13 +- internal/dao/pipeline_operation_log.go | 5 + internal/dao/task.go | 9 + internal/dao/task_test.go | 126 ++- internal/deepdoc/parser/pdf/parse_nocgo.go | 15 + internal/entity/models/types.go | 2 + internal/handler/document.go | 10 + internal/ingestion/ingestion_service.go | 440 ---------- internal/ingestion/ingestion_service_test.go | 478 ---------- .../pipeline/real_storage_integration_test.go | 463 ++++++++++ .../pipeline/template_integration_test.go | 64 +- .../ingestion/service/execute_task_test.go | 72 ++ .../ingestion/service/ingestion_service.go | 404 +++++++++ .../service/real_consumer_dataflow_test.go | 138 +++ internal/ingestion/service/test_helpers.go | 72 ++ internal/ingestion/task/chunk_builder.go | 168 ++++ internal/ingestion/task/chunk_builder_test.go | 461 ++++++++++ internal/ingestion/task/chunk_process.go | 246 ++++++ internal/ingestion/task/chunk_process_test.go | 400 +++++++++ internal/ingestion/task/chunk_utils.go | 164 ++++ internal/ingestion/task/chunk_utils_test.go | 427 +++++++++ internal/ingestion/task/constants.go | 29 + .../ingestion/task/dataflow_chunk_assets.go | 14 + internal/ingestion/task/dataflow_e2e_test.go | 370 ++++++++ ...dataflow_real_pipeline_integration_test.go | 769 ++++++++++++++++ internal/ingestion/task/dataflow_service.go | 597 +++++++++++++ .../ingestion/task/dataflow_service_test.go | 826 ++++++++++++++++++ internal/ingestion/task/embedding_test.go | 226 +++++ .../ingestion/task/embedding_token_test.go | 55 ++ internal/ingestion/task/golden_compare.go | 54 ++ .../ingestion/task/golden_compare_test.go | 71 ++ internal/ingestion/task/has_vectors_test.go | 74 ++ internal/ingestion/task/naive_merge.go | 269 ++++++ internal/ingestion/task/naive_merge_test.go | 417 +++++++++ .../task/pipeline_run_output_checklist.md | 178 ++++ internal/ingestion/task/position.go | 52 ++ internal/ingestion/task/position_test.go | 92 ++ internal/ingestion/task/real_consumer_test.go | 170 ++++ internal/ingestion/task/task_context.go | 190 ++++ internal/ingestion/task/task_context_test.go | 187 ++++ internal/ingestion/task/task_handler.go | 167 ++++ internal/ingestion/task/task_handler_test.go | 232 +++++ .../task/tool/compare_dataflow_golden.go | 288 ++++++ internal/ingestion/testutil/helpers.go | 246 ++++++ internal/ingestion/testutil/mocks.go | 174 ++++ internal/parser/parser/factory.go | 86 ++ internal/parser/parser/factory_test.go | 81 ++ internal/service/document.go | 79 +- internal/utility/metadata.go | 116 +++ internal/utility/metadata_test.go | 115 +++ 55 files changed, 9485 insertions(+), 968 deletions(-) delete mode 100644 internal/ingestion/ingestion_service.go delete mode 100644 internal/ingestion/ingestion_service_test.go create mode 100644 internal/ingestion/pipeline/real_storage_integration_test.go create mode 100644 internal/ingestion/service/execute_task_test.go create mode 100644 internal/ingestion/service/ingestion_service.go create mode 100644 internal/ingestion/service/real_consumer_dataflow_test.go create mode 100644 internal/ingestion/service/test_helpers.go create mode 100644 internal/ingestion/task/chunk_builder.go create mode 100644 internal/ingestion/task/chunk_builder_test.go create mode 100644 internal/ingestion/task/chunk_process.go create mode 100644 internal/ingestion/task/chunk_process_test.go create mode 100644 internal/ingestion/task/chunk_utils.go create mode 100644 internal/ingestion/task/chunk_utils_test.go create mode 100644 internal/ingestion/task/constants.go create mode 100644 internal/ingestion/task/dataflow_chunk_assets.go create mode 100644 internal/ingestion/task/dataflow_e2e_test.go create mode 100644 internal/ingestion/task/dataflow_real_pipeline_integration_test.go create mode 100644 internal/ingestion/task/dataflow_service.go create mode 100644 internal/ingestion/task/dataflow_service_test.go create mode 100644 internal/ingestion/task/embedding_test.go create mode 100644 internal/ingestion/task/embedding_token_test.go create mode 100644 internal/ingestion/task/golden_compare.go create mode 100644 internal/ingestion/task/golden_compare_test.go create mode 100644 internal/ingestion/task/has_vectors_test.go create mode 100644 internal/ingestion/task/naive_merge.go create mode 100644 internal/ingestion/task/naive_merge_test.go create mode 100644 internal/ingestion/task/pipeline_run_output_checklist.md create mode 100644 internal/ingestion/task/position.go create mode 100644 internal/ingestion/task/position_test.go create mode 100644 internal/ingestion/task/real_consumer_test.go create mode 100644 internal/ingestion/task/task_context.go create mode 100644 internal/ingestion/task/task_context_test.go create mode 100644 internal/ingestion/task/task_handler.go create mode 100644 internal/ingestion/task/task_handler_test.go create mode 100644 internal/ingestion/task/tool/compare_dataflow_golden.go create mode 100644 internal/ingestion/testutil/helpers.go create mode 100644 internal/ingestion/testutil/mocks.go create mode 100644 internal/parser/parser/factory.go create mode 100644 internal/parser/parser/factory_test.go create mode 100644 internal/utility/metadata.go create mode 100644 internal/utility/metadata_test.go diff --git a/.gitignore b/.gitignore index e2a8cef077..295e08f857 100644 --- a/.gitignore +++ b/.gitignore @@ -254,3 +254,7 @@ internal/deepdoc/parser/pdf/tools-py/ build/ internal/deepdoc/parser/docx/testdata/ internal/deepdoc/parser/docx/tool/ + +# test data compare tool +internal/ingestion/task/tool/generate_dataflow_golden.py +internal/ingestion/task/tool/README.md diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 2fec093ff8..ca32c177d1 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -29,7 +29,7 @@ import ( "ragflow/internal/agent/runtime" agenttool "ragflow/internal/agent/tool" "ragflow/internal/handler" - "ragflow/internal/ingestion" + ingestion "ragflow/internal/ingestion/service" "ragflow/internal/mcp" "ragflow/internal/router" "ragflow/internal/server/local" @@ -52,7 +52,7 @@ import ( "ragflow/internal/dao" "ragflow/internal/engine" "ragflow/internal/engine/redis" - _ "ragflow/internal/ingestion/wire" // single owner for ingestion-component registration (File / Parser / Tokenizer / Extractor + 4 Chunker variants) + _ "ragflow/internal/ingestion/wire" "ragflow/internal/server" "ragflow/internal/utility" ) @@ -473,6 +473,18 @@ func runAdmin(args *serverArgs) error { } func runIngestor(args *serverArgs) error { + // Initialize tokenizer (rag_analyzer) + dictPath := os.Getenv("RAGFLOW_DICT_PATH") + if dictPath == "" { + dictPath = "/usr/share/infinity/resource" + } + tokenizerCfg := &tokenizer.PoolConfig{ + DictPath: dictPath, + } + if err := tokenizer.Init(tokenizerCfg); err != nil { + common.Fatal("Failed to initialize tokenizer", zap.Error(err)) + } + defer tokenizer.Close() ingestor := ingestion.NewIngestor(*args.name, 2, []string{"pdf", "docx", "txt"}) @@ -802,7 +814,8 @@ func startServer(config *server.Config) { } } adminRuntimeHandler := handler.NewAdminRuntimeHandler(adminRuntimeSelector) - componentsHandler := handler.NewComponentsHandler(service.NewComponentsService()) + componentsSvc := service.NewComponentsService() + componentsHandler := handler.NewComponentsHandler(componentsSvc) // Initialize router r := router.NewRouter(authHandler, diff --git a/internal/binding/rag_analyzer.go b/internal/binding/rag_analyzer.go index b2ebbf9e25..66646e066e 100644 --- a/internal/binding/rag_analyzer.go +++ b/internal/binding/rag_analyzer.go @@ -1,5 +1,3 @@ -//go:build cgo - // // Copyright 2026 The InfiniFlow Authors. All Rights Reserved. // @@ -16,6 +14,8 @@ // limitations under the License. // +//go:build cgo + package rag_analyzer /* diff --git a/internal/binding/rag_analyzer_nocgo.go b/internal/binding/rag_analyzer_nocgo.go index 879f6c801d..769f09ad27 100644 --- a/internal/binding/rag_analyzer_nocgo.go +++ b/internal/binding/rag_analyzer_nocgo.go @@ -1,5 +1,19 @@ //go:build !cgo +// 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 rag_analyzer import "fmt" diff --git a/internal/dao/document.go b/internal/dao/document.go index d016a4d325..d279be8a82 100644 --- a/internal/dao/document.go +++ b/internal/dao/document.go @@ -70,6 +70,17 @@ func (dao *DocumentDAO) UpdateByID(id string, updates map[string]interface{}) er return DB.Model(&entity.Document{}).Where("id = ?", id).Updates(updates).Error } +// IncrementCounts atomically increments chunk_num, token_num, and process_duration for a document +func (dao *DocumentDAO) IncrementCounts(id string, kbID string, chunkNum int64, tokenNum int64, duration float64) error { + return DB.Model(&entity.Document{}). + Where("id = ? AND kb_id = ?", id, kbID). + Updates(map[string]interface{}{ + "chunk_num": gorm.Expr("chunk_num + ?", chunkNum), + "token_num": gorm.Expr("token_num + ?", tokenNum), + "process_duration": gorm.Expr("process_duration + ?", duration), + }).Error +} + // Delete hard-deletes document by ID. Returns rows affected. func (dao *DocumentDAO) Delete(id string) (int64, error) { result := DB.Where("id = ?", id).Delete(&entity.Document{}) diff --git a/internal/dao/ingestion.go b/internal/dao/ingestion.go index 4f43def60e..5f27b498e8 100644 --- a/internal/dao/ingestion.go +++ b/internal/dao/ingestion.go @@ -390,6 +390,10 @@ func (dao *IngestionTaskLogDAO) Create(ingestionLog *entity.IngestionTaskLog) er return DB.Create(ingestionLog).Error } +func (dao *IngestionTaskLogDAO) Update(ingestionLog *entity.IngestionTaskLog) error { + return DB.Save(ingestionLog).Error +} + func (dao *IngestionTaskLogDAO) ListLogsByTaskID(taskID string) ([]*entity.IngestionTaskLog, error) { var tasks []*entity.IngestionTaskLog err := DB.Where("task_id = ?", taskID).Order("create_time DESC").Find(&tasks).Error @@ -398,14 +402,7 @@ func (dao *IngestionTaskLogDAO) ListLogsByTaskID(taskID string) ([]*entity.Inges func (dao *IngestionTaskLogDAO) LatestLogByTaskID(taskID string) (*entity.IngestionTaskLog, error) { var task *entity.IngestionTaskLog - // Order by `id DESC` (NOT `create_time DESC`) because - // create_time has only second-level resolution — multiple - // checkpoints written within the same second tie-break - // arbitrarily. The `id` is auto-increment, monotonic, and - // always reflects write order. The pipeline's resume - // algorithm reads the latest row, so the tie-break MUST - // be deterministic. - err := DB.Where("task_id = ?", taskID).Order("id DESC").First(&task).Error + err := DB.Where("task_id = ?", taskID).Order("create_time DESC").First(&task).Error return task, err } diff --git a/internal/dao/pipeline_operation_log.go b/internal/dao/pipeline_operation_log.go index cf0b2eb2b6..ead4ad893a 100644 --- a/internal/dao/pipeline_operation_log.go +++ b/internal/dao/pipeline_operation_log.go @@ -164,3 +164,8 @@ func (dao *PipelineOperationLogDAO) GetByIDAndKBID(logID, kbID string) (*entity. } return &log, nil } + +// Create inserts a new pipeline operation log. +func (dao *PipelineOperationLogDAO) Create(log *entity.PipelineOperationLog) error { + return DB.Create(log).Error +} diff --git a/internal/dao/task.go b/internal/dao/task.go index 2870657bd6..4f1d93b666 100644 --- a/internal/dao/task.go +++ b/internal/dao/task.go @@ -51,6 +51,15 @@ func (dao *TaskDAO) GetByID(id string) (*entity.Task, error) { return &task, nil } +// DeleteIngestionTasksByDocIDs deletes ingestion tasks by document IDs (hard delete) +func (dao *TaskDAO) DeleteIngestionTasksByDocIDs(docIDs []string) (int64, error) { + if len(docIDs) == 0 { + return 0, nil + } + result := DB.Unscoped().Where("document_id IN ?", docIDs).Delete(&entity.IngestionTask{}) + return result.RowsAffected, result.Error +} + // DeleteByDocIDs deletes tasks by document IDs (hard delete) func (dao *TaskDAO) DeleteByDocIDs(docIDs []string) (int64, error) { if len(docIDs) == 0 { diff --git a/internal/dao/task_test.go b/internal/dao/task_test.go index 71e3786e1f..c4e12e371a 100644 --- a/internal/dao/task_test.go +++ b/internal/dao/task_test.go @@ -39,7 +39,10 @@ func setupTaskTestDB(t *testing.T) *gorm.DB { // Migrate task table (Task depends on Document for the doc_id FK, // but SQLite doesn't enforce FKs by default) if err := db.AutoMigrate(&entity.Task{}); err != nil { - t.Fatalf("failed to migrate: %v", err) + t.Fatalf("failed to migrate Task: %v", err) + } + if err := db.AutoMigrate(&entity.IngestionTask{}); err != nil { + t.Fatalf("failed to migrate IngestionTask: %v", err) } return db @@ -123,3 +126,124 @@ func TestGetByDocID_EmptyDocID(t *testing.T) { t.Fatalf("expected task-1, got %s", tasks[0].ID) } } + +func TestDeleteIngestionTasksByDocIDs_Success(t *testing.T) { + db := setupTaskTestDB(t) + orig := DB + DB = db + t.Cleanup(func() { DB = orig }) + + dao := NewTaskDAO() + + // Insert ingestion tasks for two different documents + task1 := &entity.IngestionTask{ID: "itask-1", DocumentID: "doc-1", UserID: "user-1", DatasetID: "ds-1", Status: "pending"} + task2 := &entity.IngestionTask{ID: "itask-2", DocumentID: "doc-1", UserID: "user-1", DatasetID: "ds-1", Status: "pending"} + task3 := &entity.IngestionTask{ID: "itask-3", DocumentID: "doc-2", UserID: "user-1", DatasetID: "ds-1", Status: "pending"} + for _, tk := range []*entity.IngestionTask{task1, task2, task3} { + if err := db.Create(tk).Error; err != nil { + t.Fatalf("failed to create ingestion task: %v", err) + } + } + + // Delete tasks for doc-1 + rowsAffected, err := dao.DeleteIngestionTasksByDocIDs([]string{"doc-1"}) + if err != nil { + t.Fatalf("DeleteIngestionTasksByDocIDs failed: %v", err) + } + if rowsAffected != 2 { + t.Fatalf("expected 2 rows affected, got %d", rowsAffected) + } + + // Verify doc-1 tasks are gone, doc-2 remains + var remaining []*entity.IngestionTask + if err := db.Find(&remaining).Error; err != nil { + t.Fatalf("failed to find remaining tasks: %v", err) + } + if len(remaining) != 1 { + t.Fatalf("expected 1 task remaining, got %d", len(remaining)) + } + if remaining[0].ID != "itask-3" { + t.Fatalf("expected itask-3 to remain, got %s", remaining[0].ID) + } +} + +func TestDeleteIngestionTasksByDocIDs_EmptyIDs(t *testing.T) { + db := setupTaskTestDB(t) + orig := DB + DB = db + t.Cleanup(func() { DB = orig }) + + dao := NewTaskDAO() + + rowsAffected, err := dao.DeleteIngestionTasksByDocIDs([]string{}) + if err != nil { + t.Fatalf("DeleteIngestionTasksByDocIDs failed: %v", err) + } + if rowsAffected != 0 { + t.Fatalf("expected 0 rows affected, got %d", rowsAffected) + } +} + +func TestDeleteIngestionTasksByDocIDs_Nonexistent(t *testing.T) { + db := setupTaskTestDB(t) + orig := DB + DB = db + t.Cleanup(func() { DB = orig }) + + dao := NewTaskDAO() + + // Insert one task to make sure table isn't empty + task := &entity.IngestionTask{ID: "itask-1", DocumentID: "doc-1", UserID: "user-1", DatasetID: "ds-1", Status: "pending"} + if err := db.Create(task).Error; err != nil { + t.Fatalf("failed to create ingestion task: %v", err) + } + + rowsAffected, err := dao.DeleteIngestionTasksByDocIDs([]string{"nonexistent"}) + if err != nil { + t.Fatalf("DeleteIngestionTasksByDocIDs failed: %v", err) + } + if rowsAffected != 0 { + t.Fatalf("expected 0 rows affected, got %d", rowsAffected) + } +} + +func TestDeleteIngestionTasksByDocIDs_MultipleIDs(t *testing.T) { + db := setupTaskTestDB(t) + orig := DB + DB = db + t.Cleanup(func() { DB = orig }) + + dao := NewTaskDAO() + + // Insert tasks for multiple documents + tasks := []*entity.IngestionTask{ + {ID: "itask-1", DocumentID: "doc-1", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, + {ID: "itask-2", DocumentID: "doc-2", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, + {ID: "itask-3", DocumentID: "doc-3", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, + {ID: "itask-4", DocumentID: "doc-2", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, + {ID: "itask-5", DocumentID: "keep", UserID: "user-1", DatasetID: "ds-1", Status: "pending"}, + } + for _, tk := range tasks { + if err := db.Create(tk).Error; err != nil { + t.Fatalf("failed to create ingestion task: %v", err) + } + } + + // Delete multiple document IDs + rowsAffected, err := dao.DeleteIngestionTasksByDocIDs([]string{"doc-1", "doc-2", "doc-3"}) + if err != nil { + t.Fatalf("DeleteIngestionTasksByDocIDs failed: %v", err) + } + if rowsAffected != 4 { + t.Fatalf("expected 4 rows affected, got %d", rowsAffected) + } + + // Verify only "keep" remains + var remaining []*entity.IngestionTask + if err := db.Find(&remaining).Error; err != nil { + t.Fatalf("failed to find remaining tasks: %v", err) + } + if len(remaining) != 1 || remaining[0].ID != "itask-5" { + t.Fatalf("expected only itask-5 to remain, got %d tasks", len(remaining)) + } +} diff --git a/internal/deepdoc/parser/pdf/parse_nocgo.go b/internal/deepdoc/parser/pdf/parse_nocgo.go index 4e0e24ed83..be5b4e3e12 100644 --- a/internal/deepdoc/parser/pdf/parse_nocgo.go +++ b/internal/deepdoc/parser/pdf/parse_nocgo.go @@ -1,5 +1,20 @@ //go:build !cgo +// 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 pdf import ( diff --git a/internal/entity/models/types.go b/internal/entity/models/types.go index 6515c8cbc8..9e04679e06 100644 --- a/internal/entity/models/types.go +++ b/internal/entity/models/types.go @@ -75,6 +75,8 @@ type ChatUsage struct { type EmbeddingData struct { Embedding []float64 `json:"embedding"` Index int `json:"index"` + // FIXME: add implementation + TokenCount int `json:"token_count"` } type RerankResult struct { diff --git a/internal/handler/document.go b/internal/handler/document.go index e042183384..2ef146a010 100644 --- a/internal/handler/document.go +++ b/internal/handler/document.go @@ -1260,6 +1260,16 @@ func (h *DocumentHandler) SetMeta(c *gin.Context) { common.SuccessWithData(c, true, "success") } +// Ingest handles document ingestion +// @Summary Ingest Document +// @Description Ingest a document for processing +// @Tags documents +// @Accept json +// @Produce json +// @Security ApiKeyAuth +// @Param request body service.IngestDocumentRequest true "ingestion info" +// @Success 200 {object} map[string]interface{} +// @Router /api/v1/documents/ingest [post] func (h *DocumentHandler) Ingest(c *gin.Context) { user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { diff --git a/internal/ingestion/ingestion_service.go b/internal/ingestion/ingestion_service.go deleted file mode 100644 index 10d7d7936c..0000000000 --- a/internal/ingestion/ingestion_service.go +++ /dev/null @@ -1,440 +0,0 @@ -// -// 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 ingestion - -import ( - "context" - "errors" - "fmt" - "ragflow/internal/dao" - "ragflow/internal/engine" - "ragflow/internal/entity" - "ragflow/internal/ingestion/pipeline" - "ragflow/internal/utility" - "sync" - "time" - - "ragflow/internal/common" - - "google.golang.org/grpc" -) - -type Ingestor struct { - id string - name string - serverAddr string - conn *grpc.ClientConn - ctx context.Context - cancel context.CancelFunc - reconnectMu sync.Mutex - - // Configuration - maxConcurrency int32 - supportedDocTypes []string - version string - - // Runtime state - currentTasks map[string]*TaskContext - tasksMu sync.RWMutex - - // Shutdown channel - receive on this to trigger graceful shutdown - ShutdownCh chan struct{} - - // Worker pool - taskChan chan *TaskContext - workerWg sync.WaitGroup - startOnce sync.Once - - ingestionTaskDAO *dao.IngestionTaskDAO - ingestionTaskLogDAO *dao.IngestionTaskLogDAO - ingestionTaskletDAO *dao.IngestionTaskletDAO - ingestionTaskletLogDAO *dao.IngestionTaskletLogDAO -} - -type TaskLog struct { - StartTime time.Time `json:"start_time"` - EndTime time.Time `json:"end_time"` - Description string `json:"description"` - Details map[string]interface{} `json:"details"` -} - -type TaskContext struct { - Ctx context.Context - CancelFunc context.CancelFunc - // if tasklet is nil, this context is belonged to a task - // if task and tasklet are both not nil, this context is belonged to a tasklet, the task is the parent task of the tasklet - Task *entity.IngestionTask - Tasklet *entity.IngestionTasklet - Logs []*TaskLog - estimatedRemainingTime time.Duration // estimated cost in seconds to complete the task - Progress int32 - ErrorMessage string - TaskHandle common.TaskHandle -} - -func NewIngestor(name string, maxConcurrency int32, supportedTypes []string) *Ingestor { - ctx, cancel := context.WithCancel(context.Background()) - id := utility.GenerateUUID() - return &Ingestor{ - id: id, - name: name, - ctx: ctx, - cancel: cancel, - maxConcurrency: maxConcurrency, - supportedDocTypes: supportedTypes, - version: "1.0.0", - currentTasks: make(map[string]*TaskContext), - taskChan: make(chan *TaskContext, maxConcurrency*2), - ShutdownCh: make(chan struct{}, 1), - ingestionTaskDAO: dao.NewIngestionTaskDAO(), - ingestionTaskLogDAO: dao.NewIngestionTaskLogDAO(), - ingestionTaskletDAO: dao.NewIngestionTaskletDAO(), - ingestionTaskletLogDAO: dao.NewIngestionTaskletLogDAO(), - } -} - -func (e *Ingestor) ID() string { - return e.id -} - -func (e *Ingestor) Start() error { - common.Info(fmt.Sprintf("Ingestor %s initialized", e.id)) - msgQueueEngine := engine.GetMessageQueueEngine() - err := msgQueueEngine.InitConsumer("tasks.RAGFLOW") - if err != nil { - return err - } - - // Ensure worker pool is started on first task - go e.startWorkerPool() - - for { - var taskHandles []common.TaskHandle - taskHandles, err = msgQueueEngine.GetMessages(4) - if err != nil { - common.Error("error consuming message", err) - continue - } - for _, taskHandle := range taskHandles { - taskMessage := taskHandle.GetMessage() - common.Info(fmt.Sprintf("Received task id: %s, type: %s", taskMessage.TaskID, taskMessage.TaskType)) - if taskMessage.TaskType != common.TaskTypeIngestionTask { - common.Info(fmt.Sprintf("task %s is not an ingestion task", taskMessage.TaskID)) - err = taskHandle.Ack() - if err != nil { - common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), err) - return err - } - continue - } - var task *entity.IngestionTask - task, err = e.ingestionTaskDAO.SetRunningByIngestor(taskMessage.TaskID) - if err != nil { - if errors.Is(err, common.ErrTaskNotFound) { - common.Warn(fmt.Sprintf("task %s not found, skipping", taskMessage.TaskID)) - err = taskHandle.Ack() - if err != nil { - common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), err) - return err - } - continue - } else { - common.Error(fmt.Sprintf("error setting task %s to running", taskMessage.TaskID), err) - return err - } - } - if task == nil { - common.Info(fmt.Sprintf("task %s is already removed", taskMessage.TaskID)) - err = taskHandle.Ack() - if err != nil { - return err - } - continue - } - - switch task.Status { - case common.COMPLETED, common.STOPPED, common.FAILED: - common.Info(fmt.Sprintf("task %s is already %s", taskMessage.TaskID, task.Status)) - err = taskHandle.Ack() - if err != nil { - common.Error(fmt.Sprintf("error nack task %s", taskMessage.TaskID), err) - return err - } - continue - case common.STOPPING, common.CREATED: - err = fmt.Errorf("task %s is in unexpected status %s", taskMessage.TaskID, task.Status) - return err - case common.RUNNING: - } - - // Construct TaskContext with a cancellable context - ctx, cancel := context.WithCancel(e.ctx) - taskCtx := &TaskContext{ - Ctx: ctx, - CancelFunc: cancel, - Task: task, - TaskHandle: taskHandle, - } - - // Push to task channel; if full, reject the task (backpressure) - select { - case e.taskChan <- taskCtx: - common.Info(fmt.Sprintf("Task %s queued (channel: %d/%d)", task.ID, len(e.taskChan), cap(e.taskChan))) - default: - common.Info(fmt.Sprintf("No available slot for task %s, failed", task.ID)) - - err = taskHandle.Nack() - if err != nil { - common.Error(fmt.Sprintf("error nack task %s", taskMessage.TaskID), err) - return err - } - } - } - } -} - -func (e *Ingestor) startWorkerPool() { - e.startOnce.Do(func() { - for i := int32(0); i < e.maxConcurrency; i++ { - e.workerWg.Add(1) - go e.workerLoop(i) - } - common.Info(fmt.Sprintf("Worker pool started with %d workers", e.maxConcurrency)) - }) -} - -func (e *Ingestor) workerLoop(id int32) { - defer e.workerWg.Done() - common.Info(fmt.Sprintf("Worker %d started", id)) - for { - select { - case <-e.ctx.Done(): - return - case taskCtx := <-e.taskChan: - if taskCtx.Tasklet != nil { - e.executeTasklet(taskCtx) - } else { - e.executeTask(taskCtx) - } - } - } -} - -func (e *Ingestor) executeTask(taskCtx *TaskContext) { - ctx := taskCtx.Ctx - task := taskCtx.Task - common.Info(fmt.Sprintf("Starting task %s", task.ID)) - - // Execute the canonical ingestion canvas DSL carried by the task. - // The Go ingestion path no longer synthesizes a parallel `stages[]` - // schema; the only accepted format is the template/canvas DSL. - dslBytes := defaultPipelineDSL(task) - if len(dslBytes) == 0 { - err := fmt.Errorf("task %s missing canonical ingestion DSL in schema.pipeline or schema.dsl", task.ID) - common.Error(fmt.Sprintf("Failed to load pipeline DSL for task %s", task.ID), err) - e.failTask(taskCtx, err) - return - } - pl, err := pipeline.NewPipelineFromDSL(dslBytes, task.ID) - if err != nil { - common.Error(fmt.Sprintf("Failed to compile pipeline for task %s", task.ID), err) - e.failTask(taskCtx, err) - return - } - - inputs := map[string]any{ - "doc_id": task.DocumentID, - } - _, runErr := pl.Run(ctx, inputs) - if runErr != nil { - if errors.Is(runErr, context.Canceled) || errors.Is(runErr, context.DeadlineExceeded) { - common.Info(fmt.Sprintf("Task %s cancelled: %v", task.ID, runErr)) - // STOPPED is a terminal status — the task will not be - // re-attempted by the consumer. Ack the message so the - // queue does not redeliver it (Nack here would race - // with the STOPPED write and let another consumer pick - // up a "stopped" task). - _ = e.ingestionTaskDAO.UpdateStatus(task.ID, common.STOPPED) - _ = e.ackOrNack(taskCtx, true) - return - } - common.Error(fmt.Sprintf("Task %s pipeline failed", task.ID), runErr) - e.failTask(taskCtx, runErr) - return - } - - if err = e.ingestionTaskDAO.UpdateStatus(task.ID, common.COMPLETED); err != nil { - common.Error(fmt.Sprintf("Task %s update status failed", task.ID), err) - _ = e.ackOrNack(taskCtx, true) - return - } - - common.Info(fmt.Sprintf("Task %s completed", task.ID)) - _ = e.ackOrNack(taskCtx, true) -} - -// defaultPipelineDSL returns the canonical ingestion canvas DSL bytes carried -// by the task schema. The ingestion runtime accepts only the template/canvas -// DSL shape; it does not synthesize a separate linear stages[] schema. -func defaultPipelineDSL(task *entity.IngestionTask) []byte { - if task != nil && task.Schema != nil { - if raw, ok := task.Schema["pipeline"]; ok { - switch v := raw.(type) { - case []byte: - if len(v) > 0 { - return v - } - case string: - if v != "" { - return []byte(v) - } - } - } - if raw, ok := task.Schema["dsl"]; ok { - switch v := raw.(type) { - case []byte: - if len(v) > 0 { - return v - } - case string: - if v != "" { - return []byte(v) - } - } - } - } - return nil -} - -// failTask updates the task to FAILED and Acks the message -// (terminal-failure path: even on error, the message must be -// removed from the queue, otherwise the broker redelivers it -// indefinitely). This fixes the pre-existing bug that the -// placeholder sleep loop never called Ack at all (plan §8 Q3). -func (e *Ingestor) failTask(taskCtx *TaskContext, runErr error) { - if err := e.ingestionTaskDAO.UpdateStatus(taskCtx.Task.ID, common.FAILED); err != nil { - common.Error(fmt.Sprintf("Task %s update status (failed) error", taskCtx.Task.ID), err) - } - _ = e.ackOrNack(taskCtx, true) - common.Error(fmt.Sprintf("Task %s failed: %v", taskCtx.Task.ID, runErr), runErr) -} - -// ackOrNack centralises the post-execution NATS message -// disposition. ack=true removes the message from the queue -// (success OR terminal-failure); ack=false re-queues (rare; -// we use it only on context cancellation to let another -// worker pick it up). -func (e *Ingestor) ackOrNack(taskCtx *TaskContext, ack bool) error { - if taskCtx == nil || taskCtx.TaskHandle == nil { - return nil - } - var err error - if ack { - err = taskCtx.TaskHandle.Ack() - } else { - err = taskCtx.TaskHandle.Nack() - } - if err != nil { - common.Error(fmt.Sprintf("Task %s ack/nack error", taskCtx.Task.ID), err) - } - return err -} - -func (e *Ingestor) executeTasklet(taskCtx *TaskContext) { - ctx := taskCtx.Ctx - tasklet := taskCtx.Tasklet - common.Info(fmt.Sprintf("Starting tasklet %s", tasklet.ID)) - - latestLog, err := e.ingestionTaskletLogDAO.LatestLogByTaskletID(tasklet.ID) - if err != nil { - latestLog = &entity.IngestionTaskletLog{ - TaskletID: tasklet.ID, - Checkpoint: entity.JSONMap{ - "current_step": 0, - "total_step": 3, - }, - } - err = e.ingestionTaskletLogDAO.Create(latestLog) - if err != nil { - common.Error(fmt.Sprintf("Failed to create task log for tasklet %s", tasklet.ID), err) - return - } - } - - var checkpointMap map[string]interface{} - checkpointMap = latestLog.Checkpoint - currentStep := checkpointMap["current_step"].(int) - totalStep := checkpointMap["total_step"].(int) - for i := currentStep; i < totalStep; i++ { - select { - case <-ctx.Done(): - // Task canceled - common.Info(fmt.Sprintf("Tasklet %s stopped", tasklet.ID)) - return - case <-time.After(3000 * time.Millisecond): - common.Info(fmt.Sprintf("Tasklet %s is running step %d", tasklet.ID, i)) - checkpointMap["current_step"] = i + 1 - latestLog.Checkpoint = checkpointMap - err = e.ingestionTaskletLogDAO.Create(latestLog) - if err != nil { - common.Error(fmt.Sprintf("Failed to update task log for tasklet %s", tasklet.ID), err) - return - } - } - } - - err = e.ingestionTaskletDAO.UpdateStatus(tasklet.ID, common.STOPPED) - if err != nil { - common.Error(fmt.Sprintf("Tasklet %s update status failed", tasklet.ID), err) - return - } - - common.Info(fmt.Sprintf("Tasklet %s completed", tasklet.ID)) -} - -// e.stream = stream -// -// if err = e.sendRegister(); err != nil { -// stream.CloseSend() -// conn.Close() -// common.Info(fmt.Sprintf("Reconnect register failed: %v, retrying in %v", err, backoff)) -// time.Sleep(backoff) -// backoff *= 2 -// if backoff > maxBackoff { -// backoff = maxBackoff -// } -// continue -// } -// -// common.Info(fmt.Sprintf("Ingestor %s reconnected to admin", e.id)) -// break -// } -// -// // Restart the loops on the new stream -// go e.receiveLoop() -// go e.heartbeatLoop() -//} - -// Stop gracefully shuts down the ingestor -func (e *Ingestor) Stop() { - common.Info(fmt.Sprintf("Stopping ingestor %s", e.id)) - e.cancel() - - // Wait for all workers to finish (they exit on ctx.Done()) - e.workerWg.Wait() - common.Info("All tasks completed") -} diff --git a/internal/ingestion/ingestion_service_test.go b/internal/ingestion/ingestion_service_test.go deleted file mode 100644 index e644fe04bc..0000000000 --- a/internal/ingestion/ingestion_service_test.go +++ /dev/null @@ -1,478 +0,0 @@ -// -// 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. -// - -// e2e tests for the Ingestor. -// -// These tests exercise Ingestor.executeTask end-to-end against -// an in-memory SQLite + memory storage backend, with a stub -// ProgressSink that records per-stage invocations. The DSL -// runs against the production pipeline package; the test -// supplies a custom 1-stage DSL that points at a stub -// component so no real storage / LLM backend is required. -// -// The pattern mirrors internal/service/agent_run_e2e_test.go -// (in-memory SQLite + miniredis, no Docker). The Ingestor's -// runnable surface is exerciseTask, which is package-private, -// so the test lives in the `ingestion` package itself. -package ingestion - -import ( - "context" - "encoding/json" - "fmt" - "strings" - "sync/atomic" - "testing" - "time" - - "ragflow/internal/agent/runtime" - "ragflow/internal/common" - "ragflow/internal/dao" - "ragflow/internal/entity" - _ "ragflow/internal/ingestion/component" // blank import: registers ingestion factories - "ragflow/internal/storage" - - "github.com/glebarez/sqlite" - "gorm.io/gorm" -) - -// --------------------------------------------------------------------------- -// Stub component for e2e testing. -// --------------------------------------------------------------------------- - -const ( - stubIngestionComp = "StubIngest" - stubTokenizerComp = "MockTokenizer" -) - -type stubIngest struct { - counter *int32 -} - -func (s *stubIngest) Parallelism() int { return 1 } -func (s *stubIngest) Inputs() map[string]string { - return map[string]string{"x": "any"} -} -func (s *stubIngest) Outputs() map[string]string { - return map[string]string{"y": "any"} -} -func (s *stubIngest) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - atomic.AddInt32(s.counter, 1) - return map[string]any{"y": "stub-output"}, nil -} - -var stubCounter int32 -var stubTokenizerCounter int32 - -func newStubIngest(_ string, _ map[string]any) (runtime.Component, error) { - return &stubIngest{counter: &stubCounter}, nil -} - -type stubTokenizer struct { - counter *int32 -} - -func (s *stubTokenizer) Parallelism() int { return 1 } -func (s *stubTokenizer) Inputs() map[string]string { - return map[string]string{"chunks": "any"} -} -func (s *stubTokenizer) Outputs() map[string]string { - return map[string]string{"tokens": "any"} -} -func (s *stubTokenizer) Invoke(_ context.Context, inputs map[string]any) (map[string]any, error) { - atomic.AddInt32(s.counter, 1) - return map[string]any{"tokens": inputs["chunks"]}, nil -} - -func newStubTokenizer(_ string, _ map[string]any) (runtime.Component, error) { - return &stubTokenizer{counter: &stubTokenizerCounter}, nil -} - -func init() { - runtime.MustRegister(stubIngestionComp, runtime.CategoryIngestion, newStubIngest, runtime.Metadata{ - Inputs: map[string]string{"x": "any"}, - Outputs: map[string]string{"y": "any"}, - }) - runtime.MustRegister(stubTokenizerComp, runtime.CategoryIngestion, newStubTokenizer, runtime.Metadata{ - Inputs: map[string]string{"chunks": "any"}, - Outputs: map[string]string{"tokens": "any"}, - }) -} - -// --------------------------------------------------------------------------- -// Stub TaskHandle — captures Ack/Nack calls. -// --------------------------------------------------------------------------- - -type stubTaskHandle struct { - acked int32 - nacked int32 - taskID string -} - -func (h *stubTaskHandle) GetMessage() common.TaskMessage { - return common.TaskMessage{TaskID: h.taskID, TaskType: common.TaskTypeIngestionTask} -} -func (h *stubTaskHandle) Ack() error { - atomic.AddInt32(&h.acked, 1) - return nil -} -func (h *stubTaskHandle) Nack() error { - atomic.AddInt32(&h.nacked, 1) - return nil -} - -// --------------------------------------------------------------------------- -// DB setup -// --------------------------------------------------------------------------- - -func setupIngestorTestDB(t *testing.T) *gorm.DB { - t.Helper() - // Use shared in-memory SQLite so the Ingestor's - // dao.NewIngestionTaskLogDAO() (which references the - // global dao.DB) and this test's local *gorm.DB see - // the same data. Each test gets a unique DSN so - // parallel test runs don't collide. The DSN name is - // sanitised to avoid SQLite parsing surprises on - // special characters in the test name. - dsn := sharedIngestorCacheDSN(t.Name()) - db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) - if err != nil { - t.Fatalf("sqlite: %v", err) - } - if err := db.AutoMigrate( - &entity.IngestionTask{}, - &entity.IngestionTaskLog{}, - ); err != nil { - t.Fatalf("migrate: %v", err) - } - return db -} - -// sharedIngestorCacheDSN builds a unique per-test DSN that -// points at a shared-cache in-memory SQLite database. The -// test name is sanitised to alphanumerics + underscores so -// the DSN parser doesn't choke on Go test names like -// "TestPipeline_Resume_RoundTripViaTaskLogSink" or unicode. -// A monotonic counter uniquifies the DSN across -// `go test -count=N` iterations so each invocation sees a -// fresh DB. -var sharedIngestorCacheCounter atomic.Uint64 - -func sharedIngestorCacheDSN(testName string) string { - var b strings.Builder - b.Grow(len(testName) + 48) - b.WriteString("file:test-") - for _, r := range testName { - switch { - case r >= 'A' && r <= 'Z', r >= 'a' && r <= 'z', r >= '0' && r <= '9': - b.WriteRune(r) - default: - b.WriteByte('_') - } - } - b.WriteByte('-') - b.WriteString(fmt.Sprintf("%d", sharedIngestorCacheCounter.Add(1))) - b.WriteString("?mode=memory&cache=shared") - return b.String() -} - -// --------------------------------------------------------------------------- -// TestIngestor_ExecuteTask_PipelineRuns -// --------------------------------------------------------------------------- - -// TestIngestor_ExecuteTask_PipelineRuns is the load-bearing -// e2e test for Phase 3. It: -// -// 1. Spins up in-memory SQLite + memory storage. -// 2. Creates an Ingestor + seeds an IngestionTask row with -// a custom 1-stage DSL that points at the stub -// component. -// 3. Calls executeTask on a TaskContext wired with the -// stub TaskHandle. -// 4. Verifies: -// a. The stub component was invoked exactly once. -// b. The IngestionTaskLog has at least one row -// recording the stub's stage completion. -// c. The IngestionTask status is COMPLETED. -// d. The stub TaskHandle was Ack()ed exactly once -// (plan §8 Q3 — fixes the pre-existing no-Ack bug). -func TestIngestor_ExecuteTask_PipelineRuns(t *testing.T) { - atomic.StoreInt32(&stubCounter, 0) - - db := setupIngestorTestDB(t) - origDB := dao.DB - dao.DB = db - t.Cleanup(func() { dao.DB = origDB }) - - // Inject memory storage so the Ingestor's - // storage.GetStorageFactory().GetStorage() call returns - // a usable backend. Production uses MinIO; tests use the - // in-memory mock. - origStorage := storage.GetStorageFactory().GetStorage() - storage.GetStorageFactory().SetStorage(storage.NewMemoryStorage()) - t.Cleanup(func() { storage.GetStorageFactory().SetStorage(origStorage) }) - - taskID := "ingestor-e2e-task" - task := &entity.IngestionTask{ - ID: taskID, - UserID: "user-e2e", - DocumentID: "doc-e2e", - DatasetID: "ds-e2e", - Status: common.RUNNING, - } - if err := db.Create(task).Error; err != nil { - t.Fatalf("create task: %v", err) - } - - // Canonical template wrapper carrying a tiny canvas DSL. - dsl := map[string]any{ - "dsl": map[string]any{ - "components": map[string]any{ - "begin": map[string]any{ - "obj": map[string]any{"component_name": "Begin", "params": map[string]any{}}, - "downstream": []string{"stub"}, - }, - "stub": map[string]any{ - "obj": map[string]any{"component_name": stubIngestionComp, "params": map[string]any{}}, - "upstream": []string{"begin"}, - }, - }, - "path": []string{"begin", "stub"}, - "graph": map[string]any{"nodes": []any{}}, - }, - } - dslBytes, _ := json.Marshal(dsl) - task.Schema = entity.JSONMap{"pipeline": []byte(dslBytes)} - if err := db.Save(task).Error; err != nil { - t.Fatalf("save task: %v", err) - } - - handle := &stubTaskHandle{taskID: taskID} - taskCtx := &TaskContext{ - Ctx: context.Background(), - Task: task, - TaskHandle: handle, - } - - ing := NewIngestor("test-ingestor", 1, []string{"pdf"}) - ing.executeTask(taskCtx) - - // (a) stub invoked - if got := atomic.LoadInt32(&stubCounter); got != 1 { - t.Errorf("expected stub invoked 1 time, got %d", got) - } - - // (b) task is COMPLETED. - var reloaded entity.IngestionTask - if err := db.Where("id = ?", taskID).First(&reloaded).Error; err != nil { - t.Fatalf("reload task: %v", err) - } - if reloaded.Status != common.COMPLETED { - t.Errorf("expected status=COMPLETED, got %s", reloaded.Status) - } - - // (c) Ack() called exactly once (plan §8 Q3). - if got := atomic.LoadInt32(&handle.acked); got != 1 { - t.Errorf("expected Ack() called 1 time, got %d", got) - } - if got := atomic.LoadInt32(&handle.nacked); got != 0 { - t.Errorf("expected Nack() called 0 times, got %d", got) - } -} - -// TestIngestor_ExecuteTask_MissingDSL verifies the task fails cleanly when -// no canonical DSL is present on the task schema. -func TestIngestor_ExecuteTask_MissingDSL(t *testing.T) { - db := setupIngestorTestDB(t) - origDB := dao.DB - dao.DB = db - t.Cleanup(func() { dao.DB = origDB }) - - taskID := "missing-dsl-task" - task := &entity.IngestionTask{ - ID: taskID, - UserID: "user-1", - DocumentID: "doc-1", - DatasetID: "ds-1", - Status: common.RUNNING, - } - if err := db.Create(task).Error; err != nil { - t.Fatalf("create task: %v", err) - } - - handle := &stubTaskHandle{taskID: taskID} - taskCtx := &TaskContext{ - Ctx: context.Background(), - Task: task, - TaskHandle: handle, - } - - ing := NewIngestor("test-missing-dsl", 1, []string{"pdf"}) - ing.executeTask(taskCtx) - - var reloaded entity.IngestionTask - if err := db.Where("id = ?", taskID).First(&reloaded).Error; err != nil { - t.Fatalf("reload task: %v", err) - } - if reloaded.Status != common.FAILED { - t.Errorf("expected status=FAILED, got %s", reloaded.Status) - } - if got := atomic.LoadInt32(&handle.acked); got != 1 { - t.Errorf("expected Ack() called 1 time on FAILED, got %d", got) - } -} - -// TestIngestor_ExecuteTask_Cancellation verifies the ctx -// cancellation path: cancel the context, expect the task -// to be STOPPED and the message Nack()ed. -func TestIngestor_ExecuteTask_Cancellation(t *testing.T) { - atomic.StoreInt32(&stubCounter, 0) - - db := setupIngestorTestDB(t) - origDB := dao.DB - dao.DB = db - t.Cleanup(func() { dao.DB = origDB }) - - origStorage := storage.GetStorageFactory().GetStorage() - storage.GetStorageFactory().SetStorage(storage.NewMemoryStorage()) - t.Cleanup(func() { storage.GetStorageFactory().SetStorage(origStorage) }) - - taskID := "cancel-task" - task := &entity.IngestionTask{ - ID: taskID, - UserID: "u", - DocumentID: "d", - DatasetID: "ds", - Status: common.RUNNING, - } - if err := db.Create(task).Error; err != nil { - t.Fatalf("create task: %v", err) - } - - ctx, cancel := context.WithCancel(context.Background()) - // Cancel immediately so the run sees a cancelled ctx. - cancel() - - dsl := map[string]any{ - "dsl": map[string]any{ - "components": map[string]any{ - "begin": map[string]any{ - "obj": map[string]any{"component_name": "Begin", "params": map[string]any{}}, - "downstream": []string{"stub"}, - }, - "stub": map[string]any{ - "obj": map[string]any{"component_name": stubIngestionComp, "params": map[string]any{}}, - "upstream": []string{"begin"}, - }, - }, - "path": []string{"begin", "stub"}, - "graph": map[string]any{"nodes": []any{}}, - }, - } - dslBytes, _ := json.Marshal(dsl) - task.Schema = entity.JSONMap{"pipeline": []byte(dslBytes)} - if err := db.Save(task).Error; err != nil { - t.Fatalf("save task: %v", err) - } - - handle := &stubTaskHandle{taskID: taskID} - taskCtx := &TaskContext{ - Ctx: ctx, - Task: task, - TaskHandle: handle, - } - - ing := NewIngestor("test-cancel", 1, []string{"pdf"}) - ing.executeTask(taskCtx) - - // Status: STOPPED on cancellation. - var reloaded entity.IngestionTask - if err := db.Where("id = ?", taskID).First(&reloaded).Error; err != nil { - t.Fatalf("reload: %v", err) - } - // The cancellation may either STOPPED (cancellation - // path) or FAILED (stage Invoke failed because ctx was - // done) — either is acceptable. The key invariant is - // that the message was NOT silently dropped without - // Ack/Nack. - if reloaded.Status != common.STOPPED && reloaded.Status != common.FAILED { - t.Errorf("expected status=STOPPED or FAILED on cancel, got %s", reloaded.Status) - } - // Total message disposition: Ack() OR Nack() must have - // been called (the bug pre-Phase-3 was: neither). - totalDisp := atomic.LoadInt32(&handle.acked) + atomic.LoadInt32(&handle.nacked) - if totalDisp != 1 { - t.Errorf("expected exactly one Ack/Nack on cancel, got ack=%d nack=%d", handle.acked, handle.nacked) - } -} - -// TestIngestor_ExecuteTask_MalformedDSL verifies executeTask -// fails cleanly on a bad DSL: task is FAILED, Ack() called. -func TestIngestor_ExecuteTask_MalformedDSL(t *testing.T) { - atomic.StoreInt32(&stubCounter, 0) - - db := setupIngestorTestDB(t) - origDB := dao.DB - dao.DB = db - t.Cleanup(func() { dao.DB = origDB }) - - origStorage := storage.GetStorageFactory().GetStorage() - storage.GetStorageFactory().SetStorage(storage.NewMemoryStorage()) - t.Cleanup(func() { storage.GetStorageFactory().SetStorage(origStorage) }) - - taskID := "malformed-task" - task := &entity.IngestionTask{ - ID: taskID, - UserID: "u", - DocumentID: "d", - DatasetID: "ds", - Status: common.RUNNING, - Schema: entity.JSONMap{"pipeline": []byte("not-json")}, - } - if err := db.Create(task).Error; err != nil { - t.Fatalf("create: %v", err) - } - - handle := &stubTaskHandle{taskID: taskID} - taskCtx := &TaskContext{ - Ctx: context.Background(), - Task: task, - TaskHandle: handle, - } - - ing := NewIngestor("test-malformed", 1, []string{"pdf"}) - ing.executeTask(taskCtx) - - var reloaded entity.IngestionTask - if err := db.Where("id = ?", taskID).First(&reloaded).Error; err != nil { - t.Fatalf("reload: %v", err) - } - if reloaded.Status != common.FAILED { - t.Errorf("expected status=FAILED on malformed DSL, got %s", reloaded.Status) - } - if got := atomic.LoadInt32(&handle.acked); got != 1 { - t.Errorf("expected Ack() on FAILED, got %d", got) - } -} - -// --------------------------------------------------------------------------- -// tiny helper for strings.Contains — Go 1.21 has it but the -// codebase pins strings.Contains. -// --------------------------------------------------------------------------- - -var _ = strings.Contains -var _ = time.Now -var _ = fmt.Sprintf diff --git a/internal/ingestion/pipeline/real_storage_integration_test.go b/internal/ingestion/pipeline/real_storage_integration_test.go new file mode 100644 index 0000000000..f5dd4f69ea --- /dev/null +++ b/internal/ingestion/pipeline/real_storage_integration_test.go @@ -0,0 +1,463 @@ +//go:build integration +// +build integration + +package pipeline + +import ( + "archive/zip" + "context" + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "sort" + "strings" + "testing" + "time" + + "ragflow/internal/dao" + "ragflow/internal/entity" + componentpkg "ragflow/internal/ingestion/component" + "ragflow/internal/server" + "ragflow/internal/storage" + + "go.uber.org/zap" + "gorm.io/driver/mysql" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +func TestPipelineRun_TemplateGeneral_RealMySQLMinIO_OutputShape(t *testing.T) { + prepareTokenizerResourceForIntegration(t) + requireTokenizerPool(t) + + cfg := mustLoadRealIntegrationConfig(t) + realDB := mustOpenRealMySQL(t, cfg) + if err := realDB.AutoMigrate( + &entity.Tenant{}, + &entity.Knowledgebase{}, + &entity.Document{}, + &entity.File{}, + &entity.File2Document{}, + ); err != nil { + t.Fatalf("auto-migrate real mysql tables: %v", err) + } + + realStorage, err := storage.NewMinioStorage(cfg.StorageEngine.Minio) + if err != nil { + t.Fatalf("connect real minio: %v", err) + } + + origDB := dao.DB + origStorage := storage.GetStorageFactory().GetStorage() + origDocResolver := componentpkg.ResolveDocumentStorageOverride + dao.DB = realDB + storage.GetStorageFactory().SetStorage(realStorage) + componentpkg.ResolveDocumentStorageOverride = nil + t.Cleanup(func() { + dao.DB = origDB + storage.GetStorageFactory().SetStorage(origStorage) + componentpkg.ResolveDocumentStorageOverride = origDocResolver + }) + + templatePath := filepath.Join(repoRootFromPipelineTest(t), "agent", "templates", "ingestion_pipeline_general.json") + templateBytes, err := os.ReadFile(templatePath) + if err != nil { + t.Fatalf("read template: %v", err) + } + templateBytes = disableTokenizerEmbeddingForTemplate(t, templateBytes) + terminalIDs := terminalComponentIDsFromTemplate(t, templateBytes) + if len(terminalIDs) != 1 || terminalIDs[0] != "Tokenizer:LegalReadersDecide" { + t.Fatalf("terminal ids = %v, want [Tokenizer:LegalReadersDecide]", terminalIDs) + } + + suffix := fmt.Sprintf("%d", time.Now().UnixNano()) + tenantID := limit32("it_tenant_" + suffix) + kbID := limit32("it_kb_" + suffix) + docID := limit32("it_doc_" + suffix) + fileID := limit32("it_file_" + suffix) + bucket := s3SafeBucketName(kbID) + objectPath := fmt.Sprintf("integration/pipeline/%s/template-general.txt", docID) + docName := "template-general.txt" + content := "Alpha paragraph.\n\nBeta paragraph." + + mustSeedRealPipelineDocument(t, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath, docName, content) + t.Cleanup(func() { + cleanupRealPipelineDocument(realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath) + }) + + pipe, err := NewPipelineFromDSL(templateBytes, "template-general-real-mysql-minio") + if err != nil { + t.Fatalf("NewPipelineFromDSL: %v", err) + } + out, err := pipe.Run(context.Background(), map[string]any{ + "doc_id": docID, + }) + if err != nil { + t.Fatalf("Run: %v", err) + } + + payload := terminalPayloadFromRunOutput(t, out, terminalIDs[0]) + if got := payload["output_format"]; got != "chunks" { + t.Fatalf("output_format = %v, want chunks", got) + } + + chunks, ok := payload["chunks"].([]map[string]any) + if !ok { + t.Fatalf("chunks = %T, want []map[string]any", payload["chunks"]) + } + wantChunkTexts := []string{"Alpha paragraph.", "Beta paragraph."} + if len(chunks) != len(wantChunkTexts) { + t.Fatalf("len(chunks) = %d, want %d", len(chunks), len(wantChunkTexts)) + } + for i, wantText := range wantChunkTexts { + if got := chunks[i]["text"]; got != wantText { + t.Fatalf("chunks[%d].text = %v, want %q", i, got, wantText) + } + if got, ok := chunks[i]["text"].(string); !ok || got == "" { + t.Fatalf("chunks[%d].text type/value = %T/%v, want non-empty string", i, chunks[i]["text"], chunks[i]["text"]) + } + if _, hasVec := chunks[i]["q_4_vec"]; hasVec { + t.Fatalf("chunks[%d] unexpectedly contains vector field q_4_vec after embedding-disabled template", i) + } + } + if _, ok := payload["embedding_token_consumption"]; ok { + t.Fatalf("embedding_token_consumption should be absent when tokenizer search_method excludes embedding: %v", payload["embedding_token_consumption"]) + } + + state := stateFromRunOutput(t, out) + fileState, ok := state["File"] + if !ok { + t.Fatal("missing File state") + } + if got := fileState["name"]; got != docName { + t.Fatalf("file state name = %v, want %q", got, docName) + } + if _, ok := fileState["bucket"]; ok { + t.Fatalf("file state should not expose bucket on doc_id path: %v", fileState["bucket"]) + } + if _, ok := fileState["path"]; ok { + t.Fatalf("file state should not expose path on doc_id path: %v", fileState["path"]) + } + + parserState, ok := state["Parser:HipSignsRhyme"] + if !ok { + t.Fatal("missing Parser:HipSignsRhyme state") + } + if got := parserState["output_format"]; got != "json" { + t.Fatalf("parser output_format = %v, want json", got) + } + jsonItems, ok := parserState["json"].([]map[string]any) + if !ok || len(jsonItems) != 2 { + t.Fatalf("parser json = %T/%v, want 2 items", parserState["json"], parserState["json"]) + } + for i, wantText := range wantChunkTexts { + if got := jsonItems[i]["text"]; got != wantText { + t.Fatalf("parser json[%d].text = %v, want %q", i, got, wantText) + } + } + + chunkerState, ok := state["TokenChunker:SixApplesFall"] + if !ok { + t.Fatal("missing TokenChunker:SixApplesFall state") + } + if got := chunkerState["output_format"]; got != "chunks" { + t.Fatalf("chunker output_format = %v, want chunks", got) + } + chunkerChunks, ok := chunkerState["chunks"].([]map[string]any) + if !ok || len(chunkerChunks) != len(wantChunkTexts) { + t.Fatalf("chunker chunks = %T/%v, want %d items", chunkerState["chunks"], chunkerState["chunks"], len(wantChunkTexts)) + } + for i, wantText := range wantChunkTexts { + if got := chunkerChunks[i]["text"]; got != wantText { + t.Fatalf("chunker chunk[%d].text = %v, want %q", i, got, wantText) + } + if got := chunkerChunks[i]["doc_type_kwd"]; got != "text" { + t.Fatalf("chunker chunk[%d].doc_type_kwd = %v, want text", i, got) + } + } +} + +func mustLoadRealIntegrationConfig(t *testing.T) *server.Config { + t.Helper() + server.SetLogger(zap.NewNop()) + configPath := filepath.Join(repoRootFromPipelineTest(t), "conf", "service_conf.yaml") + if err := server.Init(configPath); err != nil { + t.Fatalf("init service config from %s: %v", configPath, err) + } + cfg := server.GetConfig() + if cfg == nil || cfg.Database.Host == "" || cfg.StorageEngine.Minio == nil || cfg.StorageEngine.Minio.Host == "" { + t.Fatal("real integration config is incomplete") + } + return cfg +} + +func prepareTokenizerResourceForIntegration(t *testing.T) { + t.Helper() + if os.Getenv("RAGFLOW_DICT_PATH") != "" { + return + } + srcDir := filepath.Join(repoRootFromPipelineTest(t), ".venv", "lib", "python3.13", "site-packages", "infinity") + if _, err := os.Stat(filepath.Join(srcDir, "huqie.txt")); err != nil { + t.Skipf("tokenizer resource source not found at %s: %v", srcDir, err) + } + if _, err := os.Stat(filepath.Join(srcDir, "huqie.txt.trie")); err != nil { + t.Skipf("tokenizer trie source not found at %s: %v", srcDir, err) + } + root := t.TempDir() + ragDir := filepath.Join(root, "rag") + if err := os.MkdirAll(ragDir, 0o755); err != nil { + t.Fatalf("mkdir rag tokenizer dir: %v", err) + } + mustSymlink(t, filepath.Join(srcDir, "huqie.txt"), filepath.Join(ragDir, "huqie.txt")) + mustSymlink(t, filepath.Join(srcDir, "huqie.txt.trie"), filepath.Join(ragDir, "huqie.trie")) + mustWriteTokenizerPOSDef(t, filepath.Join(srcDir, "huqie.txt"), filepath.Join(ragDir, "pos-id.def")) + mustPrepareTokenizerWordNet(t, root) + mustPrepareTokenizerOpenCC(t, root) + if err := os.Setenv("RAGFLOW_DICT_PATH", root); err != nil { + t.Fatalf("set RAGFLOW_DICT_PATH: %v", err) + } + t.Cleanup(func() { + _ = os.Unsetenv("RAGFLOW_DICT_PATH") + }) +} + +func mustSymlink(t *testing.T, src, dst string) { + t.Helper() + if err := os.Symlink(src, dst); err != nil { + t.Fatalf("symlink tokenizer resource %s -> %s: %v", src, dst, err) + } +} + +func mustWriteTokenizerPOSDef(t *testing.T, dictPath, outPath string) { + t.Helper() + data, err := os.ReadFile(dictPath) + if err != nil { + t.Fatalf("read tokenizer dict %s: %v", dictPath, err) + } + posSet := map[string]struct{}{} + for _, line := range strings.Split(string(data), "\n") { + fields := strings.Fields(line) + if len(fields) == 3 { + posSet[fields[2]] = struct{}{} + } + } + if len(posSet) == 0 { + t.Fatalf("no POS tags parsed from tokenizer dict %s", dictPath) + } + posList := make([]string, 0, len(posSet)) + for pos := range posSet { + posList = append(posList, pos) + } + sort.Strings(posList) + content := strings.Join(posList, "\n") + "\n" + if err := os.WriteFile(outPath, []byte(content), 0o644); err != nil { + t.Fatalf("write tokenizer pos file %s: %v", outPath, err) + } +} + +func mustPrepareTokenizerWordNet(t *testing.T, root string) { + t.Helper() + zipPath := filepath.Join(repoRootFromPipelineTest(t), "ragflow_deps", "nltk_data", "corpora", "wordnet.zip") + reader, err := zip.OpenReader(zipPath) + if err != nil { + t.Skipf("open wordnet zip %s: %v", zipPath, err) + return + } + defer func() { + _ = reader.Close() + }() + for _, f := range reader.File { + name := strings.TrimPrefix(f.Name, "wordnet/") + if name == "" || strings.HasSuffix(name, "/") { + continue + } + dst := filepath.Join(root, "wordnet", name) + if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil { + t.Fatalf("mkdir wordnet dst %s: %v", dst, err) + } + rc, err := f.Open() + if err != nil { + t.Fatalf("open wordnet entry %s: %v", f.Name, err) + } + out, err := os.Create(dst) + if err != nil { + _ = rc.Close() + t.Fatalf("create wordnet dst %s: %v", dst, err) + } + if _, err := io.Copy(out, rc); err != nil { + _ = out.Close() + _ = rc.Close() + t.Fatalf("copy wordnet entry %s -> %s: %v", f.Name, dst, err) + } + if err := out.Close(); err != nil { + _ = rc.Close() + t.Fatalf("close wordnet dst %s: %v", dst, err) + } + if err := rc.Close(); err != nil { + t.Fatalf("close wordnet entry %s: %v", f.Name, err) + } + } +} + +func mustPrepareTokenizerOpenCC(t *testing.T, root string) { + t.Helper() + const systemOpenCC = "/usr/share/opencc" + if _, err := os.Stat(systemOpenCC); err != nil { + t.Skipf("system opencc dir %s not found: %v", systemOpenCC, err) + return + } + mustSymlink(t, systemOpenCC, filepath.Join(root, "opencc")) +} + +func mustOpenRealMySQL(t *testing.T, cfg *server.Config) *gorm.DB { + t.Helper() + dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=Local", + cfg.Database.Username, + cfg.Database.Password, + cfg.Database.Host, + cfg.Database.Port, + cfg.Database.Database, + cfg.Database.Charset, + ) + db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{ + Logger: logger.Default.LogMode(logger.Silent), + }) + if err != nil { + t.Fatalf("connect real mysql: %v", err) + } + return db +} + +func disableTokenizerEmbeddingForTemplate(t *testing.T, raw []byte) []byte { + t.Helper() + var tpl map[string]any + if err := json.Unmarshal(raw, &tpl); err != nil { + t.Fatalf("unmarshal template: %v", err) + } + dsl, ok := tpl["dsl"].(map[string]any) + if !ok { + t.Fatalf("template dsl = %T, want map[string]any", tpl["dsl"]) + } + components, ok := dsl["components"].(map[string]any) + if !ok { + t.Fatalf("template components = %T, want map[string]any", dsl["components"]) + } + changed := 0 + for _, rawComp := range components { + comp, ok := rawComp.(map[string]any) + if !ok { + continue + } + obj, ok := comp["obj"].(map[string]any) + if !ok { + continue + } + if obj["component_name"] != "Tokenizer" { + continue + } + params, ok := obj["params"].(map[string]any) + if !ok { + continue + } + params["search_method"] = []string{"full_text"} + changed++ + } + if changed == 0 { + t.Fatal("no Tokenizer component found to disable embedding") + } + out, err := json.Marshal(tpl) + if err != nil { + t.Fatalf("marshal modified template: %v", err) + } + return out +} + +func mustSeedRealPipelineDocument( + t *testing.T, + db *gorm.DB, + stg storage.Storage, + tenantID, kbID, docID, fileID, bucket, objectPath, docName, content string, +) { + t.Helper() + if err := db.Create(&entity.Tenant{ + ID: tenantID, + LLMID: "gpt-4", + Status: strPtr("1"), + }).Error; err != nil { + t.Fatalf("create tenant: %v", err) + } + if err := db.Create(&entity.Knowledgebase{ + ID: kbID, + TenantID: tenantID, + EmbdID: "embd-1", + ParserConfig: entity.JSONMap{}, + Status: strPtr("1"), + }).Error; err != nil { + t.Fatalf("create kb: %v", err) + } + if err := stg.Put(bucket, objectPath, []byte(content)); err != nil { + t.Fatalf("put real minio object: %v", err) + } + if err := db.Create(&entity.File{ + ID: fileID, + ParentID: bucket, + TenantID: tenantID, + CreatedBy: tenantID, + Name: docName, + Type: "txt", + Location: strPtr(objectPath), + SourceType: "", + }).Error; err != nil { + t.Fatalf("create file: %v", err) + } + if err := db.Create(&entity.Document{ + ID: docID, + KbID: kbID, + ParserID: "naive", + ParserConfig: entity.JSONMap{}, + SourceType: "local", + Type: "txt", + CreatedBy: tenantID, + Name: strPtr(docName), + Location: strPtr(objectPath), + Suffix: ".txt", + Status: strPtr("1"), + }).Error; err != nil { + t.Fatalf("create document: %v", err) + } + if err := db.Create(&entity.File2Document{ + ID: limit32("it_map_" + docID), + FileID: strPtr(fileID), + DocumentID: strPtr(docID), + }).Error; err != nil { + t.Fatalf("create file2document: %v", err) + } +} + +func cleanupRealPipelineDocument(db *gorm.DB, stg storage.Storage, tenantID, kbID, docID, fileID, bucket, objectPath string) { + _ = db.Where("document_id = ?", docID).Delete(&entity.File2Document{}).Error + _ = db.Where("id = ?", docID).Delete(&entity.Document{}).Error + _ = db.Where("id = ?", fileID).Delete(&entity.File{}).Error + _ = db.Where("id = ?", kbID).Delete(&entity.Knowledgebase{}).Error + _ = db.Where("id = ?", tenantID).Delete(&entity.Tenant{}).Error + _ = stg.Remove(bucket, objectPath) +} + +func strPtr(s string) *string { + return &s +} + +func limit32(s string) string { + if len(s) <= 32 { + return s + } + return s[:32] +} + +func s3SafeBucketName(s string) string { + s = strings.ToLower(s) + s = strings.ReplaceAll(s, "_", "-") + return s +} diff --git a/internal/ingestion/pipeline/template_integration_test.go b/internal/ingestion/pipeline/template_integration_test.go index dee8549f97..bed778414b 100644 --- a/internal/ingestion/pipeline/template_integration_test.go +++ b/internal/ingestion/pipeline/template_integration_test.go @@ -22,6 +22,7 @@ import ( "encoding/json" "fmt" "os" + "os/exec" "path/filepath" goruntime "runtime" "sort" @@ -938,19 +939,68 @@ func generateTemplatePipelinePDF() ([]byte, error) { } func findTemplatePDFFont() (string, error) { - candidates := []string{ - "/usr/share/fonts/truetype/LiberationSerif-Regular.ttf", - "/usr/share/fonts/truetype/DejaVuSerif.ttf", - "/usr/share/fonts/truetype/DejaVuSans.ttf", + // Option 1: Prefer fc-list (fontconfig) - most elegant approach + if fontPath, err := findFontViaFontconfig(); err == nil { + return fontPath, nil } - for _, candidate := range candidates { - if _, err := os.Stat(candidate); err == nil { - return candidate, nil + + // Option 2: Fallback - search common directories (for systems without fontconfig) + candidates := []string{ + "LiberationSerif-Regular.ttf", + "DejaVuSerif.ttf", + "DejaVuSans.ttf", + } + searchDirs := []string{ + "/usr/share/fonts/truetype", + "/usr/share/fonts/truetype/dejavu", + "/usr/share/fonts/truetype/liberation", + "/usr/share/fonts", + "/usr/local/share/fonts", + } + + for _, dir := range searchDirs { + for _, font := range candidates { + candidate := filepath.Join(dir, font) + if _, err := os.Stat(candidate); err == nil { + return candidate, nil + } } } + return "", fmt.Errorf("no usable TTF font found for generated PDF fixture") } +// findFontViaFontconfig uses fc-list command to find a suitable TTF font +func findFontViaFontconfig() (string, error) { + // Search priority: DejaVu Serif > Liberation Serif > DejaVu Sans + fontPatterns := []string{ + "DejaVu Serif:style=Regular", + "Liberation Serif:style=Regular", + "DejaVu Sans:style=Regular", + } + + for _, pattern := range fontPatterns { + cmd := exec.Command("fc-list", "--format=%{file}\n", pattern) + output, err := cmd.Output() + if err != nil { + continue + } + lines := strings.Split(string(output), "\n") + for _, line := range lines { + line = strings.TrimSpace(line) + if line == "" { + continue + } + // Verify file exists and is in TTF format + if _, err := os.Stat(line); err == nil && strings.HasSuffix(strings.ToLower(line), ".ttf") { + return line, nil + } + } + } + + return "", fmt.Errorf("no font found via fontconfig") +} + func assertMetadataContainsString(t *testing.T, metadata map[string]any, key, want string) { t.Helper() raw, ok := metadata[key] diff --git a/internal/ingestion/service/execute_task_test.go b/internal/ingestion/service/execute_task_test.go new file mode 100644 index 0000000000..46befe2571 --- /dev/null +++ b/internal/ingestion/service/execute_task_test.go @@ -0,0 +1,72 @@ +package service + +import ( + "context" + "testing" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + taskpkg "ragflow/internal/ingestion/task" + "ragflow/internal/ingestion/testutil" +) + +func TestExecuteTask_DataflowRoutesToTaskHandler(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + _, _, docID, taskID := testutil.SeedTestData(t, db, + testutil.WithPipelineID("flow-1"), + testutil.WithTenantID("tenant-1"), + ) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + var routedToDataflow bool + var gotTaskID string + var gotProgress []float64 + var gotMsgs []string + ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error { + routedToDataflow = true + gotTaskID = ingestionTask.ID + wrapped := func(prog float64, msg string) { + gotProgress = append(gotProgress, prog*100) + gotMsgs = append(gotMsgs, msg) + } + wrapped(0.82, "mock dataflow start") + wrapped(1.0, "mock dataflow done") + return nil + } + ingestor.storageImpl = testutil.NewMockStorage(map[string][]byte{"/unused": []byte("unused")}) + ingestor.documentDAO = &testutil.MockDocDAO{ + Docs: map[string]*entity.Document{"doc-1": testutil.TestDoc("doc-1", "pdf", ".pdf")}, + } + + taskCtx := taskpkg.NewTaskContextForScheduling( + context.Background(), + &entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING}, + nil, + ) + + ingestor.executeTask(taskCtx) + + if !routedToDataflow { + t.Fatal("expected executeTask to route dataflow task to runDocumentTask") + } + if gotTaskID != taskID { + t.Fatalf("runDocumentTask got task ID %q, want %q", gotTaskID, taskID) + } + finalTask, err := dao.NewIngestionTaskDAO().GetByID(taskID) + if err != nil { + t.Fatalf("load final ingestion task: %v", err) + } + if finalTask.Status != common.COMPLETED { + t.Fatalf("final status = %s, want %s", finalTask.Status, common.COMPLETED) + } + if len(gotProgress) != 2 || gotProgress[0] != 82 || gotProgress[1] != 100 { + t.Fatalf("gotProgress = %v, want [82 100]", gotProgress) + } + if len(gotMsgs) != 2 || gotMsgs[1] != "mock dataflow done" { + t.Fatalf("gotMsgs = %v, want final message %q", gotMsgs, "mock dataflow done") + } +} diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go new file mode 100644 index 0000000000..e25c297eb9 --- /dev/null +++ b/internal/ingestion/service/ingestion_service.go @@ -0,0 +1,404 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "ragflow/internal/utility" + goruntime "runtime" + "strings" + "sync" + + "ragflow/internal/common" + "ragflow/internal/dao" + doctype "ragflow/internal/deepdoc/parser/type" + "ragflow/internal/engine" + "ragflow/internal/entity" + taskpkg "ragflow/internal/ingestion/task" + "ragflow/internal/storage" + + "github.com/google/uuid" + "google.golang.org/grpc" +) + +// pdfParser is the interface for PDF parsing, made injectable for tests. +type pdfParser interface { + ParseWithDeepDoc(ctx context.Context, filename string, data []byte, config doctype.ParserConfig) ([]map[string]any, error) +} + +// DAO interfaces — injectable for tests. +type docGetter interface { + GetByID(id string) (*entity.Document, error) +} +type f2dGetter interface { + GetByDocumentID(docID string) ([]*entity.File2Document, error) +} +type fileGetter interface { + GetByID(id string) (*entity.File, error) +} + +type Ingestor struct { + id string + name string + serverAddr string + conn *grpc.ClientConn + ctx context.Context + cancel context.CancelFunc + reconnectMu sync.Mutex + + // Configuration + maxConcurrency int32 + supportedDocTypes []string + version string + + // Runtime state + currentTasks map[string]*taskpkg.TaskContext + tasksMu sync.RWMutex + + // Shutdown channel - receive on this to trigger graceful shutdown + ShutdownCh chan struct{} + + // Worker pool + taskChan chan *taskpkg.TaskContext + workerWg sync.WaitGroup + startOnce sync.Once + + ingestionTaskDAO *dao.IngestionTaskDAO + ingestionTaskLogDAO *dao.IngestionTaskLogDAO + ingestionTaskletDAO *dao.IngestionTaskletDAO + ingestionTaskletLogDAO *dao.IngestionTaskletLogDAO + + // Injected dependencies for parseDocument (overridable in tests) + documentDAO docGetter + f2dDAO f2dGetter + fileDAO fileGetter + storageImpl storage.Storage + pdfParser pdfParser + + // runDocumentTask dispatches to the migrated task handler path. + // Tests may override this to verify branch routing without invoking + // the full downstream stack. + runDocumentTask func(ctx context.Context, ingestionTask *entity.IngestionTask) error +} + +func NewIngestor(name string, maxConcurrency int32, supportedTypes []string) *Ingestor { + ctx, cancel := context.WithCancel(context.Background()) + id := utility.GenerateUUID() + ingestor := &Ingestor{ + id: id, + name: name, + ctx: ctx, + cancel: cancel, + maxConcurrency: maxConcurrency, + supportedDocTypes: supportedTypes, + version: "1.0.0", + currentTasks: make(map[string]*taskpkg.TaskContext), + taskChan: make(chan *taskpkg.TaskContext, maxConcurrency*2), + ShutdownCh: make(chan struct{}, 1), + ingestionTaskDAO: dao.NewIngestionTaskDAO(), + ingestionTaskLogDAO: dao.NewIngestionTaskLogDAO(), + ingestionTaskletDAO: dao.NewIngestionTaskletDAO(), + ingestionTaskletLogDAO: dao.NewIngestionTaskletLogDAO(), + } + ingestor.runDocumentTask = ingestor.defaultRunDocumentTask + return ingestor +} + +func (e *Ingestor) ID() string { + return e.id +} + +func (e *Ingestor) Start() error { + common.Info(fmt.Sprintf("Ingestor %s initialized", e.id)) + msgQueueEngine := engine.GetMessageQueueEngine() + err := msgQueueEngine.InitConsumer("tasks.RAGFLOW") + if err != nil { + return err + } + + // Ensure worker pool is started on first task + go e.startWorkerPool() + + for { + var taskHandles []common.TaskHandle + taskHandles, err = msgQueueEngine.GetMessages(4) + if err != nil { + common.Error("error consuming message", err) + continue + } + for _, taskHandle := range taskHandles { + taskMessage := taskHandle.GetMessage() + common.Info(fmt.Sprintf("Received task id: %s, type: %s", taskMessage.TaskID, taskMessage.TaskType)) + if taskMessage.TaskType != common.TaskTypeIngestionTask { + common.Info(fmt.Sprintf("task %s is not an ingestion task", taskMessage.TaskID)) + err = taskHandle.Ack() + if err != nil { + common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), err) + return err + } + continue + } + var task *entity.IngestionTask + task, err = e.ingestionTaskDAO.SetRunningByIngestor(taskMessage.TaskID) + if err != nil { + if errors.Is(err, common.ErrTaskNotFound) { + common.Warn(fmt.Sprintf("task %s not found, skipping", taskMessage.TaskID)) + err = taskHandle.Ack() + if err != nil { + common.Error(fmt.Sprintf("error ack task %s", taskMessage.TaskID), err) + return err + } + continue + } else { + common.Error(fmt.Sprintf("error setting task %s to running", taskMessage.TaskID), err) + return err + } + } + if task == nil { + common.Info(fmt.Sprintf("task %s is already removed", taskMessage.TaskID)) + err = taskHandle.Ack() + if err != nil { + return err + } + continue + } + + switch task.Status { + case common.COMPLETED, common.STOPPED, common.FAILED: + common.Info(fmt.Sprintf("task %s is already %s", taskMessage.TaskID, task.Status)) + err = taskHandle.Ack() + if err != nil { + common.Error(fmt.Sprintf("error nack task %s", taskMessage.TaskID), err) + return err + } + continue + case common.STOPPING, common.CREATED: + err = fmt.Errorf("task %s is in unexpected status %s", taskMessage.TaskID, task.Status) + return err + case common.RUNNING: + } + + // Construct TaskContext with parent context + taskCtx := taskpkg.NewTaskContextForScheduling(e.ctx, task, taskHandle) + + // Push to task channel; if full, reject the task (backpressure) + select { + case e.taskChan <- taskCtx: + common.Info(fmt.Sprintf("Task %s queued (channel: %d/%d)", task.ID, len(e.taskChan), cap(e.taskChan))) + default: + common.Info(fmt.Sprintf("No available slot for task %s, failed", task.ID)) + + err = taskHandle.Nack() + if err != nil { + common.Error(fmt.Sprintf("error nack task %s", taskMessage.TaskID), err) + return err + } + } + } + } +} + +func (e *Ingestor) startWorkerPool() { + e.startOnce.Do(func() { + for i := int32(0); i < e.maxConcurrency; i++ { + e.workerWg.Add(1) + go e.workerLoop(i) + } + common.Info(fmt.Sprintf("Worker pool started with %d workers", e.maxConcurrency)) + }) +} + +func (e *Ingestor) workerLoop(id int32) { + defer e.workerWg.Done() + common.Info(fmt.Sprintf("Worker %d started", id)) + for { + select { + case <-e.ctx.Done(): + return + case taskCtx := <-e.taskChan: + common.Info("task context:" + taskCtx.IngestionTask.ID) + e.executeTask(taskCtx) + } + } +} + +func (e *Ingestor) executeTask(taskCtx *taskpkg.TaskContext) { + ctx := taskCtx.Ctx + task := taskCtx.IngestionTask + common.Info(fmt.Sprintf("Starting task %s", task.ID)) + + latestLog, err := e.ingestionTaskLogDAO.LatestLogByTaskID(task.ID) + if err != nil { + latestLog = &entity.IngestionTaskLog{ + ID: 0, + TaskID: task.ID, + Checkpoint: entity.JSONMap{ + "current_step": 1, + "total_step": 5, + }, + } + err = e.ingestionTaskLogDAO.Create(latestLog) + if err != nil { + common.Error(fmt.Sprintf("Failed to create task log for task %s", task.ID), err) + return + } + } + + var checkpointMap map[string]interface{} + checkpointMap = latestLog.Checkpoint + currentStep, ok := common.GetInt(checkpointMap["current_step"]) + if !ok { + common.Fatal(fmt.Sprintf("Failed to get current step from task log for task %s", task.ID)) + return + } + totalStep, ok := common.GetInt(checkpointMap["total_step"]) + if !ok { + common.Fatal(fmt.Sprintf("Failed to get current step from task log for task %s", task.ID)) + return + } + // Bump checkpoint to signal we started processing + checkpointMap["current_step"] = currentStep + 1 + latestLog.Checkpoint = checkpointMap + if err := e.ingestionTaskLogDAO.Update(latestLog); err != nil { + common.Error(fmt.Sprintf("Failed to persist checkpoint for task %s", task.ID), err) + } + _ = totalStep // unused in this temporary integration + + select { + case <-ctx.Done(): + common.Info(fmt.Sprintf("Task %s cancelled", task.ID)) + return + default: + } + if err := e.runDocumentTask(ctx, task); err != nil { + common.Error(fmt.Sprintf("Task %s failed", task.ID), err) + if uErr := e.ingestionTaskDAO.UpdateStatus(task.ID, common.FAILED); uErr != nil { + common.Error(fmt.Sprintf("Failed to set task %s to FAILED", task.ID), uErr) + } + return + } + + err = e.ingestionTaskDAO.UpdateStatus(task.ID, common.COMPLETED) + if err != nil { + common.Error(fmt.Sprintf("Task %s update status failed", task.ID), err) + return + } + + common.Info(fmt.Sprintf("Task %s completed", task.ID)) +} + +// FIXME: should remove +func (e *Ingestor) getPipelineID(tenantID string) (string, error) { + _, file, _, ok := goruntime.Caller(0) + if !ok { + return "", fmt.Errorf("failed to get runtime base dir") + } + base := filepath.Clean(filepath.Join(filepath.Dir(file), "..", "..", "..")) + + // FIXME: now use mocked pipeline, need to change later + templatePath := filepath.Join(base, "agent", "templates", "ingestion_pipeline_general.json") + common.Warn(fmt.Sprintf("use mocked DSL, templatePath: %s", templatePath)) + templateBytes, err := os.ReadFile(templatePath) + if err != nil { + return "", err + } + + templateBytes, err = disableTokenizerEmbeddingForTaskTemplate(templateBytes) + if err != nil { + return "", err + } + + var templateDSL entity.JSONMap + if err := json.Unmarshal(templateBytes, &templateDSL); err != nil { + return "", err + } + + des := "mock up DSL for integration test" + title := "mock up DSL" + ID := strings.ReplaceAll(uuid.New().String(), "-", "")[:32] + userCanvas := entity.UserCanvas{UserID: tenantID, DSL: templateDSL, + Description: &des, Title: &title, ID: ID, Permission: "me"} + + if err := dao.NewUserCanvasDAO().Create(&userCanvas); err != nil { + return "", err + } + + return ID, nil +} + +func disableTokenizerEmbeddingForTaskTemplate(raw []byte) ([]byte, error) { + var tpl map[string]any + if err := json.Unmarshal(raw, &tpl); err != nil { + return nil, fmt.Errorf("failed to unmarshal") + } + dsl, ok := tpl["dsl"].(map[string]any) + if !ok { + return nil, fmt.Errorf("failed to do dsl convert") + } + components, ok := dsl["components"].(map[string]any) + if !ok { + return nil, fmt.Errorf("template components = %T, want map[string]any", dsl["components"]) + } + for _, rawComp := range components { + comp, ok := rawComp.(map[string]any) + if !ok { + continue + } + obj, ok := comp["obj"].(map[string]any) + if !ok || obj["component_name"] != "Tokenizer" { + continue + } + params, ok := obj["params"].(map[string]any) + if !ok { + continue + } + params["search_method"] = []string{"full_text"} + } + out, err := json.Marshal(tpl) + if err != nil { + return nil, fmt.Errorf("marshal modified template: %v", err) + } + return out, nil +} + +func (e *Ingestor) defaultRunDocumentTask(ctx context.Context, ingestionTask *entity.IngestionTask) error { + docTaskCtx, err := taskpkg.LoadFromIngestionTask(ingestionTask) + if err != nil { + return fmt.Errorf("load task context for %s: %w", ingestionTask.ID, err) + } + + if docTaskCtx.PipelineID, err = e.getPipelineID(docTaskCtx.Tenant.ID); err != nil { + return fmt.Errorf("get pipeline ID for %s: %w", ingestionTask.ID, err) + } + + return taskpkg.NewTaskHandler(docTaskCtx).Handle() +} + +// Stop gracefully shuts down the ingestor +func (e *Ingestor) Stop() { + common.Info(fmt.Sprintf("Stopping ingestor %s", e.id)) + e.cancel() + + // Wait for all workers to finish (they exit on ctx.Done()) + e.workerWg.Wait() + common.Info("All tasks completed") +} diff --git a/internal/ingestion/service/real_consumer_dataflow_test.go b/internal/ingestion/service/real_consumer_dataflow_test.go new file mode 100644 index 0000000000..17a55ec449 --- /dev/null +++ b/internal/ingestion/service/real_consumer_dataflow_test.go @@ -0,0 +1,138 @@ +//go:build integration +// +build integration + +package service + +import ( + "context" + "encoding/json" + "errors" + "testing" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/engine/nats" + "ragflow/internal/entity" + "ragflow/internal/ingestion/testutil" +) + +func TestRealConsumer_DataflowMessageRoutesToExecuteTask(t *testing.T) { + natsEngine := nats.NewNatsEngine("localhost", 4222) + if err := natsEngine.Init(); err != nil { + t.Fatalf("NATS Init: %v", err) + } + if err := natsEngine.InitConsumer("tasks.>"); err != nil { + t.Fatalf("InitConsumer: %v", err) + } + + for { + handles, _ := natsEngine.GetMessages(1) + if len(handles) == 0 { + break + } + _ = handles[0].Ack() + } + + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + _, _, docID, _ := testutil.SeedTestData(t, db, + testutil.WithTenantID("tenant-q-1"), + testutil.WithKBID("kb-q-1"), + testutil.WithDocID("doc-q-1"), + testutil.WithTaskID("ingest-q-1"), + testutil.WithPipelineID("flow-queue-1"), + testutil.WithDocName("queue-dataflow.pdf"), + ) + + // testutil.SeedTestData already created an IngestionTask with status RUNNING. + // Reset it to CREATED so the consumer path can transition it to RUNNING. + if err := db.Model(&entity.IngestionTask{}).Where("id = ?", "ingest-q-1").Update("status", common.CREATED).Error; err != nil { + t.Fatalf("reset ingestion task status: %v", err) + } + ingestionTask := &entity.IngestionTask{ + ID: "ingest-q-1", + UserID: "u1", + DocumentID: docID, + DatasetID: "kb-q-1", + Status: common.CREATED, + } + + payload, _ := json.Marshal(common.TaskMessage{ + TaskID: ingestionTask.ID, + TaskType: common.TaskTypeIngestionTask, + }) + if err := natsEngine.PublishTask("tasks.RAGFLOW", payload); err != nil { + t.Fatalf("PublishTask: %v", err) + } + + handles, err := natsEngine.GetMessages(1) + if err != nil { + t.Fatalf("GetMessages: %v", err) + } + if len(handles) != 1 { + t.Fatalf("expected 1 message, got %d", len(handles)) + } + taskHandle := handles[0] + taskMsg := taskHandle.GetMessage() + if taskMsg.TaskID != ingestionTask.ID { + t.Fatalf("message task ID = %s, want %s", taskMsg.TaskID, ingestionTask.ID) + } + if taskMsg.TaskType != common.TaskTypeIngestionTask { + t.Fatalf("message task type = %s, want %s", taskMsg.TaskType, common.TaskTypeIngestionTask) + } + + ingestionTaskDAO := dao.NewIngestionTaskDAO() + task, err := ingestionTaskDAO.SetRunningByIngestor(taskMsg.TaskID) + if err != nil { + if errors.Is(err, common.ErrTaskNotFound) { + t.Fatalf("task not found after publish: %s", taskMsg.TaskID) + } + t.Fatalf("SetRunningByIngestor: %v", err) + } + if task.Status != common.RUNNING { + t.Fatalf("task status after SetRunningByIngestor = %s, want %s", task.Status, common.RUNNING) + } + + ingestor := NewIngestor("queue-test", 1, []string{"pdf"}) + var routedToDataflow bool + var progressEvents []string + ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error { + routedToDataflow = true + progressFn(0.82, "mock queue dataflow start") + progressFn(1.0, "mock queue dataflow done") + progressEvents = append(progressEvents, "0.82:mock queue dataflow start", "1.00:mock queue dataflow done") + return nil + } + + taskCtx := &TaskContext{ + Ctx: context.Background(), + CancelFunc: func() {}, + Task: task, + TaskHandle: taskHandle, + } + + ingestor.executeTask(taskCtx) + + if !routedToDataflow { + t.Fatal("expected executeTask to route queue-consumed dataflow task to runDocumentTask") + } + if taskCtx.Progress != 100 { + t.Fatalf("taskCtx.Progress = %d, want 100", taskCtx.Progress) + } + if len(progressEvents) != 2 { + t.Fatalf("progressEvents = %v, want 2 events", progressEvents) + } + if err := taskHandle.Ack(); err != nil { + t.Fatalf("Ack: %v", err) + } + + finalTask, err := ingestionTaskDAO.GetByID(task.ID) + if err != nil { + t.Fatalf("GetByID: %v", err) + } + if finalTask.Status != common.COMPLETED { + t.Fatalf("final status = %s, want %s", finalTask.Status, common.COMPLETED) + } +} diff --git a/internal/ingestion/service/test_helpers.go b/internal/ingestion/service/test_helpers.go new file mode 100644 index 0000000000..69dc5e922e --- /dev/null +++ b/internal/ingestion/service/test_helpers.go @@ -0,0 +1,72 @@ +// +// 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 "ragflow/internal/storage" + +// ────────────────────────────────────────────────────────── +// Ingestor Test Helpers +// ────────────────────────────────────────────────────────── + +// NewTestIngestor creates a new Ingestor for testing with default values. +func NewTestIngestor() *Ingestor { + return &Ingestor{ + id: "test-ingestor", + } +} + +// IngestorOption configures a test Ingestor. +type IngestorOption func(*Ingestor) + +// WithMockStorage sets the storageImpl to the given mock. +func WithMockStorage(mock storage.Storage) IngestorOption { + return func(i *Ingestor) { + i.storageImpl = mock + } +} + +// WithMockPDFParser sets the pdfParser to the given mock. +func WithMockPDFParser(mock pdfParser) IngestorOption { + return func(i *Ingestor) { + i.pdfParser = mock + } +} + +// WithMockDAOs sets the documentDAO, f2dDAO, and fileDAO. +func WithMockDAOs(docDAO docGetter, f2dDAO f2dGetter, fileDAO fileGetter) IngestorOption { + return func(i *Ingestor) { + i.documentDAO = docDAO + i.f2dDAO = f2dDAO + i.fileDAO = fileDAO + } +} + +// SetupTestIngestor creates a new test Ingestor with the given options. +func SetupTestIngestor(t testingT, opts ...IngestorOption) *Ingestor { + t.Helper() + i := NewTestIngestor() + for _, opt := range opts { + opt(i) + } + return i +} + +// testingT is a subset of testing.TB used by test helpers (no *testing.T import needed). +type testingT interface { + Helper() + Fatalf(format string, args ...any) +} diff --git a/internal/ingestion/task/chunk_builder.go b/internal/ingestion/task/chunk_builder.go new file mode 100644 index 0000000000..bac48335fd --- /dev/null +++ b/internal/ingestion/task/chunk_builder.go @@ -0,0 +1,168 @@ +// +// 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 task + +import ( + "encoding/base64" + "fmt" + "strings" + "time" + "unicode/utf8" + + "github.com/cespare/xxhash/v2" +) + +// PAGERANK_FLD is the field key for pagerank feature value on a chunk, +// matching Python common.constants.PAGERANK_FLD = "pagerank_fea". +const PAGERANK_FLD = "pagerank_fea" + +// Chunk represents a single document chunk ready for embedding and indexing. +// Mirrors the Python dict format produced by ChunkService._prepare_docs_and_upload(). +type Chunk struct { + ID string `json:"id"` // xxhash64(content_with_weight + doc_id) hex + DocID string `json:"doc_id"` // document ID + KBID string `json:"kb_id"` // knowledge base ID + DocNameKwd string `json:"docnm_kwd"` // document file name + ContentWithWeight string `json:"content_with_weight"` // main text content + CreateTime string `json:"create_time"` // "2006-01-02 15:04:05" + CreateTimestamp float64 `json:"create_timestamp_flt"` // Unix timestamp + ImgID string `json:"img_id"` // image ID (from MinIO or inline) + PageNumInt []int `json:"page_num_int"` // page numbers + TopInt []int `json:"top_int"` // top positions + AvailableInt int `json:"available_int"` // availability flag, default 1 + PageRank int `json:"pagerank_fea"` // pagerank feature value, 0 = unset +} + +// prepareDocs converts ParseWithDeepDoc output ([]map[string]any) to a list of +// Chunk structs. Matches Python ChunkService._prepare_docs_and_upload(). +// +// Input sections come from ParseWithDeepDoc; each map may contain: +// - "text" (string) — the section text → ContentWithWeight +// - "doc_type_kwd" (string) — document type hint (unused directly in chunk) +// - "img_id" (string) — inline image ID (empty string if none) +// - "positions" ([]float64) — [pageNum, top, left, width, height] +// +// This is a pure function with no I/O side effects. Image upload to MinIO +// is handled separately by the caller. +func PrepareDocs(sections []map[string]any, docID, kbID, docName string, pagerank int) []Chunk { + if len(sections) == 0 { + return nil + } + + now := time.Now() + createTime := now.Format("2006-01-02 15:04:05") + createTimestamp := float64(now.Unix()) + + chunks := make([]Chunk, 0, len(sections)) + for _, sec := range sections { + text, _ := sec["text"].(string) + + c := Chunk{ + DocID: docID, + KBID: kbID, + DocNameKwd: docName, + ContentWithWeight: text, + CreateTime: createTime, + CreateTimestamp: createTimestamp, + AvailableInt: 1, + PageRank: pagerank, + } + + // Extract img_id if present + if imgID, ok := sec["img_id"].(string); ok { + c.ImgID = imgID + } + + // Extract positions: [pageNum, top, left, width, height] + if positions, ok := sec["positions"].([]float64); ok && len(positions) >= 2 { + c.PageNumInt = append(c.PageNumInt, int(positions[0])) + c.TopInt = append(c.TopInt, int(positions[1])) + } + + // Generate deterministic chunk ID: xxhash64(content_with_weight + doc_id) → hex + c.ID = ChunkID(text, docID) + chunks = append(chunks, c) + } + + return chunks +} + +// ChunkID generates a deterministic chunk identifier matching Python's: +// +// xxhash.xxh64((content_with_weight + str(doc_id)).encode("utf-8", "surrogatepass")).hexdigest() +func ChunkID(contentWithWeight, docID string) string { + return fmt.Sprintf("%016x", xxhash.Sum64String(contentWithWeight+docID)) +} + +// String returns a human-readable summary of the chunk for logging. +func (c *Chunk) String() string { + preview := c.ContentWithWeight + if utf8.RuneCountInString(preview) > 80 { + preview = string([]rune(preview)[:80]) + "..." + } + return fmt.Sprintf("Chunk{id=%s doc=%s page=%v top=%v text=%q}", + c.ID, c.DocID, c.PageNumInt, c.TopInt, preview) +} + +// UploadChunkImages uploads base64-encoded images from sections to storage +// and sets each chunk's ImgID to the resulting storage reference. +// +// Matches Python ChunkService._prepare_docs_and_upload() image upload logic: +// - If section has "img_id" pre-set → use it directly (skip upload) +// - If section has "image" (base64) → decode → upload via putFn → set ImgID +// - Otherwise → ImgID stays empty +// +// putFn is typically storage.Put(bucket, fnm, data) where bucket=kbID, fnm=chunkID. +// sections and chunks must correspond 1:1 by index. +func UploadChunkImages(sections []map[string]any, chunks []Chunk, kbID string, putFn func(bucket, fnm string, binary []byte) error) error { + for i := range chunks { + sec := sections[i] + + // Already has img_id from parser — use as-is + if imgID, ok := sec["img_id"].(string); ok && imgID != "" { + chunks[i].ImgID = imgID + continue + } + + // No image data — skip + imgRaw, ok := sec["image"].(string) + if !ok || imgRaw == "" { + continue + } + + imgBytes, err := decodeChunkImagePayload(imgRaw) + if err != nil { + return fmt.Errorf("decode image for chunk %s: %w", chunks[i].ID, err) + } + + // Upload to storage (bucket=kbID, fnm=chunkID) + if err := putFn(kbID, chunks[i].ID, imgBytes); err != nil { + return fmt.Errorf("upload image for chunk %s: %w", chunks[i].ID, err) + } + + // Set ImgID to storage reference (matching Python f"{bucket}-{objname}") + chunks[i].ImgID = fmt.Sprintf("%s-%s", kbID, chunks[i].ID) + } + return nil +} + +func decodeChunkImagePayload(raw string) ([]byte, error) { + if idx := strings.Index(raw, ","); strings.HasPrefix(raw, "data:image/") && idx >= 0 { + raw = raw[idx+1:] + } + return base64.StdEncoding.DecodeString(raw) +} diff --git a/internal/ingestion/task/chunk_builder_test.go b/internal/ingestion/task/chunk_builder_test.go new file mode 100644 index 0000000000..8d61a9f5e1 --- /dev/null +++ b/internal/ingestion/task/chunk_builder_test.go @@ -0,0 +1,461 @@ +// +// 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 task + +import ( + "encoding/base64" + "errors" + "regexp" + "testing" +) + +const ( + testDocID = "doc-1" + testKBID = "kb-1" + testDocName = "test.pdf" +) + +func TestPrepareDocs_Basic(t *testing.T) { + sections := []map[string]any{ + {"text": "Hello World", "doc_type_kwd": "text", "img_id": ""}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + if len(chunks) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(chunks)) + } + + c := chunks[0] + if c.ContentWithWeight != "Hello World" { + t.Errorf("ContentWithWeight = %q, want %q", c.ContentWithWeight, "Hello World") + } + if c.DocID != testDocID { + t.Errorf("DocID = %q, want %q", c.DocID, testDocID) + } + if c.KBID != testKBID { + t.Errorf("KBID = %q, want %q", c.KBID, testKBID) + } + if c.DocNameKwd != testDocName { + t.Errorf("DocNameKwd = %q, want %q", c.DocNameKwd, testDocName) + } + if c.AvailableInt != 1 { + t.Errorf("AvailableInt = %d, want 1", c.AvailableInt) + } +} + +func TestPrepareDocs_Multiple(t *testing.T) { + sections := []map[string]any{ + {"text": "First", "doc_type_kwd": "text", "img_id": ""}, + {"text": "Second", "doc_type_kwd": "text", "img_id": ""}, + {"text": "Third", "doc_type_kwd": "text", "img_id": ""}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + if len(chunks) != 3 { + t.Fatalf("expected 3 chunks, got %d", len(chunks)) + } + + // All IDs should be different (different content) + ids := make(map[string]bool) + for _, c := range chunks { + if c.ContentWithWeight == "" { + t.Error("chunk has empty ContentWithWeight") + } + ids[c.ID] = true + } + if len(ids) != 3 { + t.Errorf("expected 3 unique IDs, got %d", len(ids)) + } +} + +func TestPrepareDocs_EmptyInput(t *testing.T) { + chunks := PrepareDocs(nil, testDocID, testKBID, testDocName, 0) + if chunks != nil { + t.Errorf("expected nil for empty input, got %d chunks", len(chunks)) + } + + chunks = PrepareDocs([]map[string]any{}, testDocID, testKBID, testDocName, 0) + if chunks != nil { + t.Errorf("expected nil for empty input, got %d chunks", len(chunks)) + } +} + +func TestPrepareDocs_IDFormat(t *testing.T) { + sections := []map[string]any{ + {"text": "Content for ID test", "doc_type_kwd": "text", "img_id": ""}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + id := chunks[0].ID + // xxhash64 hex digest is 16 hex chars (64 bits = 16 hex digits) + matched, _ := regexp.MatchString(`^[0-9a-f]{16}$`, id) + if !matched { + t.Errorf("ID %q does not match hex format", id) + } + +} + +func TestPrepareDocs_IDDeterministic(t *testing.T) { + sections := []map[string]any{ + {"text": "Deterministic test content", "doc_type_kwd": "text", "img_id": ""}, + } + + chunks1 := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + chunks2 := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + if chunks1[0].ID != chunks2[0].ID { + t.Errorf("ID not deterministic: %q vs %q", chunks1[0].ID, chunks2[0].ID) + } + + // Different docID should produce different ID + chunks3 := PrepareDocs(sections, "doc-2", testKBID, testDocName, 0) + if chunks1[0].ID == chunks3[0].ID { + t.Error("different docIDs produced the same chunk ID") + } +} + +func TestPrepareDocs_Timestamps(t *testing.T) { + sections := []map[string]any{ + {"text": "Timestamp test", "doc_type_kwd": "text", "img_id": ""}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + c := chunks[0] + + // CreateTime should match the expected format + matched, _ := regexp.MatchString(`^\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}$`, c.CreateTime) + if !matched { + t.Errorf("CreateTime %q does not match format YYYY-MM-DD HH:MM:SS", c.CreateTime) + } + + if c.CreateTimestamp <= 0 { + t.Errorf("CreateTimestamp = %f, want > 0", c.CreateTimestamp) + } +} + +func TestPrepareDocs_Positions(t *testing.T) { + sections := []map[string]any{ + { + "text": "Positioned content", + "doc_type_kwd": "text", + "img_id": "", + "positions": []float64{3, 150.5, 20, 400, 60}, + }, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + c := chunks[0] + if len(c.PageNumInt) != 1 || c.PageNumInt[0] != 3 { + t.Errorf("PageNumInt = %v, want [3]", c.PageNumInt) + } + if len(c.TopInt) != 1 || c.TopInt[0] != 150 { + t.Errorf("TopInt = %v, want [150]", c.TopInt) + } +} + +func TestPrepareDocs_NoPositions(t *testing.T) { + sections := []map[string]any{ + {"text": "No positions", "doc_type_kwd": "text", "img_id": ""}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + c := chunks[0] + if c.PageNumInt != nil { + t.Errorf("PageNumInt = %v, want nil", c.PageNumInt) + } + if c.TopInt != nil { + t.Errorf("TopInt = %v, want nil", c.TopInt) + } +} + +func TestPrepareDocs_ImgID(t *testing.T) { + sections := []map[string]any{ + { + "text": "Content with image", + "doc_type_kwd": "text", + "img_id": "some_image_id_123", + }, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + c := chunks[0] + if c.ImgID != "some_image_id_123" { + t.Errorf("ImgID = %q, want %q", c.ImgID, "some_image_id_123") + } +} + +func TestPrepareDocs_EmptyImgID(t *testing.T) { + sections := []map[string]any{ + {"text": "No image", "doc_type_kwd": "text", "img_id": ""}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + c := chunks[0] + if c.ImgID != "" { + t.Errorf("ImgID = %q, want empty string", c.ImgID) + } +} + +func TestPrepareDocs_MixedContent(t *testing.T) { + sections := []map[string]any{ + {"text": "Page 1 content", "doc_type_kwd": "text", "img_id": "", + "positions": []float64{1, 0, 0, 100, 50}}, + {"text": "Page 2 with image", "doc_type_kwd": "text", "img_id": "img_abc", + "positions": []float64{2, 30, 10, 200, 80}}, + {"text": "No position, no image", "doc_type_kwd": "text", "img_id": ""}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + if len(chunks) != 3 { + t.Fatalf("expected 3 chunks, got %d", len(chunks)) + } + + // First: has position + if len(chunks[0].PageNumInt) != 1 || chunks[0].PageNumInt[0] != 1 { + t.Errorf("chunk[0] PageNumInt = %v, want [1]", chunks[0].PageNumInt) + } + + // Second: has position + img_id + if chunks[1].ImgID != "img_abc" { + t.Errorf("chunk[1] ImgID = %q, want %q", chunks[1].ImgID, "img_abc") + } + if len(chunks[1].PageNumInt) != 1 || chunks[1].PageNumInt[0] != 2 { + t.Errorf("chunk[1] PageNumInt = %v, want [2]", chunks[1].PageNumInt) + } + + // Third: no position, no image + if chunks[2].PageNumInt != nil { + t.Errorf("chunk[2] PageNumInt should be nil, got %v", chunks[2].PageNumInt) + } + if chunks[2].ImgID != "" { + t.Errorf("chunk[2] ImgID should be empty, got %q", chunks[2].ImgID) + } +} + +func TestPrepareDocs_PageRank(t *testing.T) { + sections := []map[string]any{ + {"text": "Ranked content", "doc_type_kwd": "text", "img_id": ""}, + } + + // pagerank=0 → PageRank should be 0 (unset) + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + if chunks[0].PageRank != 0 { + t.Errorf("PageRank = %d, want 0", chunks[0].PageRank) + } + + // pagerank=5 → PageRank should be 5 + chunks = PrepareDocs(sections, testDocID, testKBID, testDocName, 5) + if chunks[0].PageRank != 5 { + t.Errorf("PageRank = %d, want 5", chunks[0].PageRank) + } +} + +func TestUploadChunkImages_WithImage(t *testing.T) { + imgBase64 := base64.StdEncoding.EncodeToString([]byte("fake image")) + sections := []map[string]any{ + {"text": "Has image", "doc_type_kwd": "text", "img_id": "", "image": imgBase64}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + var putBucket, putFnm string + putFn := func(bucket, fnm string, binary []byte) error { + putBucket = bucket + putFnm = fnm + if len(binary) == 0 { + t.Error("Put called with empty binary") + } + return nil + } + + if err := UploadChunkImages(sections, chunks, testKBID, putFn); err != nil { + t.Fatalf("UploadChunkImages failed: %v", err) + } + + expectedImgID := testKBID + "-" + chunks[0].ID + if chunks[0].ImgID != expectedImgID { + t.Errorf("ImgID = %q, want %q", chunks[0].ImgID, expectedImgID) + } + if putBucket != testKBID { + t.Errorf("put bucket = %q, want %q", putBucket, testKBID) + } + if putFnm != chunks[0].ID { + t.Errorf("put fnm = %q, want %q", putFnm, chunks[0].ID) + } +} + +func TestUploadChunkImages_PresetImgID(t *testing.T) { + sections := []map[string]any{ + {"text": "Has preset img_id", "doc_type_kwd": "text", "img_id": "preset-id-123", "image": "should-be-ignored"}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + callCount := 0 + putFn := func(bucket, fnm string, binary []byte) error { + callCount++ + return nil + } + + if err := UploadChunkImages(sections, chunks, testKBID, putFn); err != nil { + t.Fatalf("UploadChunkImages failed: %v", err) + } + + if chunks[0].ImgID != "preset-id-123" { + t.Errorf("ImgID = %q, want %q", chunks[0].ImgID, "preset-id-123") + } + if callCount != 0 { + t.Errorf("expected 0 Put calls (preset img_id), got %d", callCount) + } +} + +func TestUploadChunkImages_NoImage(t *testing.T) { + sections := []map[string]any{ + {"text": "No image", "doc_type_kwd": "text", "img_id": ""}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + callCount := 0 + putFn := func(bucket, fnm string, binary []byte) error { + callCount++ + return nil + } + + if err := UploadChunkImages(sections, chunks, testKBID, putFn); err != nil { + t.Fatalf("UploadChunkImages failed: %v", err) + } + + if chunks[0].ImgID != "" { + t.Errorf("ImgID = %q, want empty string", chunks[0].ImgID) + } + if callCount != 0 { + t.Errorf("expected 0 Put calls, got %d", callCount) + } +} + +func TestUploadChunkImages_Mixed(t *testing.T) { + imgBase64 := base64.StdEncoding.EncodeToString([]byte("img1")) + sections := []map[string]any{ + {"text": "Img1", "doc_type_kwd": "text", "img_id": "", "image": imgBase64}, + {"text": "No img", "doc_type_kwd": "text", "img_id": ""}, + {"text": "Preset", "doc_type_kwd": "text", "img_id": "preset-001"}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + var callCount int + putFn := func(bucket, fnm string, binary []byte) error { + callCount++ + return nil + } + + if err := UploadChunkImages(sections, chunks, testKBID, putFn); err != nil { + t.Fatalf("UploadChunkImages failed: %v", err) + } + + if callCount != 1 { + t.Errorf("expected 1 Put call, got %d", callCount) + } + if chunks[0].ImgID != testKBID+"-"+chunks[0].ID { + t.Errorf("chunk[0] ImgID = %q, want %q", chunks[0].ImgID, testKBID+"-"+chunks[0].ID) + } + if chunks[1].ImgID != "" { + t.Errorf("chunk[1] ImgID = %q, want empty", chunks[1].ImgID) + } + if chunks[2].ImgID != "preset-001" { + t.Errorf("chunk[2] ImgID = %q, want 'preset-001'", chunks[2].ImgID) + } +} + +func TestUploadChunkImages_InvalidBase64(t *testing.T) { + sections := []map[string]any{ + {"text": "Bad image", "doc_type_kwd": "text", "img_id": "", "image": "!!!not-base64!!!"}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + putFn := func(bucket, fnm string, binary []byte) error { + return nil + } + + if err := UploadChunkImages(sections, chunks, testKBID, putFn); err == nil { + t.Fatal("expected error for invalid base64, got nil") + } +} + +func TestUploadChunkImages_PutError(t *testing.T) { + imgBase64 := base64.StdEncoding.EncodeToString([]byte("img")) + sections := []map[string]any{ + {"text": "Put fails", "doc_type_kwd": "text", "img_id": "", "image": imgBase64}, + } + chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0) + + putFn := func(bucket, fnm string, binary []byte) error { + return errors.New("storage unavailable") + } + + if err := UploadChunkImages(sections, chunks, testKBID, putFn); err == nil { + t.Fatal("expected error for put failure, got nil") + } +} + +func TestChunkID_NotPanic(t *testing.T) { + // Edge cases: empty content + _ = ChunkID("", testDocID) + _ = ChunkID("content", "") + _ = ChunkID("", "") + // All should not panic +} + +func TestChunkID_PreservesLeadingZero(t *testing.T) { + content := "" + docID := "doc-1" + got := ChunkID(content, docID) + want := "037fe13bd80c56aa" + if got != want { + t.Fatalf("ChunkID(%q, %q) = %q, want %q", content, docID, got, want) + } + if len(got) != 16 { + t.Fatalf("ChunkID length = %d, want 16", len(got)) + } +} + +func TestChunk_String(t *testing.T) { + c := Chunk{ + ID: "abc123", + DocID: "doc-1", + ContentWithWeight: "Short text", + PageNumInt: []int{1}, + TopInt: []int{0}, + } + s := c.String() + if len(s) == 0 { + t.Error("String() returned empty") + } +} + +func TestChunk_StringTruncates(t *testing.T) { + long := "" + for i := 0; i < 200; i++ { + long += "x" + } + c := Chunk{ + ID: "abc123", + DocID: "doc-1", + ContentWithWeight: long, + } + s := c.String() + if len(s) >= 200 { + t.Error("String() should truncate long content") + } +} diff --git a/internal/ingestion/task/chunk_process.go b/internal/ingestion/task/chunk_process.go new file mode 100644 index 0000000000..cc641fbab1 --- /dev/null +++ b/internal/ingestion/task/chunk_process.go @@ -0,0 +1,246 @@ +// +// 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 task + +import ( + "fmt" + "regexp" + "strings" + "time" + + "github.com/pkoukk/tiktoken-go" + "ragflow/internal/common" + "ragflow/internal/tokenizer" + "ragflow/internal/utility" +) + +var keywordsSplitRE = regexp.MustCompile(`[,，;；、\r\n]+`) + +// TruncateTexts truncates each text by token count using cl100k_base encoding. +// maxLength is reduced by 10 as a safety margin, matching Python. +// Mirrors Python: EmbeddingUtils.truncate_texts() +func TruncateTexts(texts []string, maxLength int) []string { + if texts == nil { + return nil + } + safeMax := maxLength - 10 + if safeMax < 0 { + safeMax = 0 + } + enc, err := tiktoken.GetEncoding("cl100k_base") + if err != nil { + // Fallback: if tiktoken fails, return as-is + result := make([]string, len(texts)) + copy(result, texts) + return result + } + result := make([]string, len(texts)) + for i, t := range texts { + tokens := enc.Encode(t, nil, nil) + if len(tokens) > safeMax { + result[i] = enc.Decode(tokens[:safeMax]) + } else { + result[i] = t + } + } + return result +} + +// SplitQuestions splits a questions string by newline, keeping all elements. +// Mirrors Python: ck["questions"].split("\n") — keeps empty strings +func SplitQuestions(questions string) []string { + return strings.Split(questions, "\n") +} + +// SplitKeywords splits a keywords string by common delimiters. +// Mirrors Python: re.split(r"[,，;；、\r\n]+", keywords) +func SplitKeywords(keywords string) []string { + if keywords == "" { + return nil + } + parts := keywordsSplitRE.Split(keywords, -1) + result := make([]string, 0, len(parts)) + for _, p := range parts { + if p != "" { + result = append(result, p) + } + } + if len(result) == 0 { + return nil + } + return result +} + +// CreateChunkTime returns the current timestamp as a formatted string and float Unix timestamp. +// The float has sub-second precision, matching Python: datetime.now().timestamp() +func CreateChunkTime() (string, float64) { + now := time.Now() + timeStr := now.Format("2006-01-02 15:04:05") + return timeStr, float64(now.UnixMicro()) / 1e6 +} + +// RenameTextToContentWithWeight renames the "text" key to "content_with_weight". +// If "content_with_weight" already exists, the "text" key is simply removed. +// Mirrors Python: ck["content_with_weight"] = ck["text"]; del ck["text"] +func RenameTextToContentWithWeight(chunk map[string]any) { + if _, exists := chunk["content_with_weight"]; !exists { + if text, ok := chunk["text"]; ok { + chunk["content_with_weight"] = text + } + } + delete(chunk, "text") +} + +// GetEmbeddingTokenConsumption extracts the embedding token consumption from pipeline output. +// Handles both int (Go native) and float64 (after JSON round-trip). +func GetEmbeddingTokenConsumption(output map[string]any) int { + if output == nil { + return 0 + } + switch v := output[EmbeddingTokenConsumptionKey].(type) { + case int: + return v + case float64: + return int(v) + default: + common.Warn(fmt.Sprintf("unexpected type %T for embedding token consumption, key=%q", v, EmbeddingTokenConsumptionKey)) + return 0 + } +} + +// ProcessChunksForDataflow mutates chunks into the pre-index structure used by +// dataflow and returns merged metadata. +func ProcessChunksForDataflow( + chunks []map[string]any, + docID string, + kbID string, + docName string, + now time.Time, +) map[string]any { + if chunks == nil { + return nil + } + metadata := make(map[string]any) + timeStr := now.Format("2006-01-02 15:04:05") + timestamp := float64(now.UnixMicro()) / 1e6 + + for _, ck := range chunks { + ck["doc_id"] = docID + ck["kb_id"] = []string{kbID} + ck["docnm_kwd"] = docName + ck["create_time"] = timeStr + ck["create_timestamp_flt"] = timestamp + + if _, exists := ck["id"]; !exists { + text := MustGetChunkTextString(ck, "ProcessChunksForDataflow") + ck["id"] = ChunkID(text, docID) + } + + processChunkQuestions(ck) + processChunkKeywords(ck) + processChunkSummary(ck) + metadata = mergeChunkMetadata(metadata, ck) + RenameTextToContentWithWeight(ck) + processChunkPositions(ck) + removeInternalChunkFields(ck) + } + return metadata +} + +func removeInternalChunkFields(ck map[string]any) { + delete(ck, "_pdf_positions") + delete(ck, "image") +} + +func processChunkQuestions(ck map[string]any) { + if _, exists := ck["questions"]; !exists { + return + } + if _, hasTks := ck["question_tks"]; !hasTks { + q, _ := ck["questions"].(string) + ck["question_kwd"] = strings.Split(q, "\n") + tks, err := tokenizer.Tokenize(q) + if err == nil { + ck["question_tks"] = tks + } else { + ck["question_tks"] = q + } + } + delete(ck, "questions") +} + +func processChunkKeywords(ck map[string]any) { + if _, exists := ck["keywords"]; !exists { + return + } + if _, hasTks := ck["important_tks"]; !hasTks { + kws, _ := ck["keywords"].(string) + ck["important_kwd"] = SplitKeywords(kws) + tks, err := tokenizer.Tokenize(kws) + if err == nil { + ck["important_tks"] = tks + } else { + ck["important_tks"] = kws + } + } + delete(ck, "keywords") +} + +func processChunkSummary(ck map[string]any) { + if _, exists := ck["summary"]; !exists { + return + } + if _, hasLtks := ck["content_ltks"]; !hasLtks { + smmry, _ := ck["summary"].(string) + ltks, err := tokenizer.Tokenize(smmry) + if err == nil { + ck["content_ltks"] = ltks + } else { + ck["content_ltks"] = smmry + } + smLtks, err := tokenizer.FineGrainedTokenize(ck["content_ltks"].(string)) + if err == nil { + ck["content_sm_ltks"] = smLtks + } else { + ck["content_sm_ltks"] = ck["content_ltks"].(string) + } + } + delete(ck, "summary") +} + +func mergeChunkMetadata(metadata map[string]any, ck map[string]any) map[string]any { + metaVal, exists := ck["metadata"] + if !exists { + return metadata + } + if metaMap, ok := metaVal.(map[string]any); ok { + metadata = utility.UpdateMetadataTo(metadata, metaMap) + } + delete(ck, "metadata") + return metadata +} + +func processChunkPositions(ck map[string]any) { + poss, exists := ck["positions"] + if !exists { + return + } + if positions, ok := poss.([]float64); ok { + AddPositions(ck, positions) + } + delete(ck, "positions") +} \ No newline at end of file diff --git a/internal/ingestion/task/chunk_process_test.go b/internal/ingestion/task/chunk_process_test.go new file mode 100644 index 0000000000..f5208a970d --- /dev/null +++ b/internal/ingestion/task/chunk_process_test.go @@ -0,0 +1,400 @@ +package task + +import ( + "math" + "testing" + "time" +) + +// ============================================================================= +// TruncateTexts — Python: tiktoken.cl100k_base token-level truncation +// ============================================================================= + +func TestTruncateTexts_TokenLevel(t *testing.T) { + // Python: enc.encode("hello world") → [15339, 1917] (2 tokens) + // With maxLength=12, safeMax=2, both tokens fit → unchanged + result := TruncateTexts([]string{"hello world"}, 12) + if len(result) != 1 { + t.Fatalf("len = %d, want 1", len(result)) + } + if result[0] != "hello world" { + t.Errorf("with safeMax=2, 'hello world' (2 tokens) should fit, got %q", result[0]) + } +} + +func TestTruncateTexts_TruncatesByToken(t *testing.T) { + // Python: with safeMax=1, keep only 1 token → shorter than original + result := TruncateTexts([]string{"hello world"}, 11) + // safeMax = 11-10 = 1, keeps only 1 token + if len(result[0]) >= len("hello world") || len(result[0]) == 0 { + t.Errorf("safeMax=1 should produce shorter output, got %q", result[0]) + } +} + +func TestTruncateTexts_EmptyText(t *testing.T) { + result := TruncateTexts([]string{""}, 100) + if len(result) != 1 || result[0] != "" { + t.Errorf("empty text should remain empty, got %q", result[0]) + } +} + +func TestTruncateTexts_NilInput(t *testing.T) { + result := TruncateTexts(nil, 100) + if result != nil { + t.Errorf("expected nil for nil input, got %v", result) + } +} + +func TestTruncateTexts_EmptySlice(t *testing.T) { + result := TruncateTexts([]string{}, 100) + if len(result) != 0 { + t.Errorf("expected empty slice, got len=%d", len(result)) + } +} + +func TestTruncateTexts_MultipleTexts(t *testing.T) { + texts := []string{"hello world", "short"} + result := TruncateTexts(texts, 100) + if len(result) != 2 { + t.Fatalf("len = %d, want 2", len(result)) + } + if result[0] != "hello world" { + t.Errorf("first text should be unchanged") + } + if result[1] != "short" { + t.Errorf("second text should be unchanged") + } +} + +// ============================================================================= +// SplitQuestions — Python: str.split("\n"), keeps empty strings +// ============================================================================= + +func TestSplitQuestions_Basic(t *testing.T) { + result := SplitQuestions("Q1\nQ2\nQ3") + if len(result) != 3 { + t.Fatalf("len = %d, want 3", len(result)) + } + if result[0] != "Q1" || result[1] != "Q2" || result[2] != "Q3" { + t.Errorf("got %v, want [Q1 Q2 Q3]", result) + } +} + +func TestSplitQuestions_Empty(t *testing.T) { + result := SplitQuestions("") + // Python: "".split("\n") → [""] + if len(result) != 1 || result[0] != "" { + t.Errorf("got %v, want [\"\"]", result) + } +} + +func TestSplitQuestions_TrailingNewline(t *testing.T) { + result := SplitQuestions("Q1\nQ2\n") + // Python: "Q1\nQ2\n".split("\n") → ["Q1", "Q2", ""] + if len(result) != 3 { + t.Fatalf("len = %d, want 3, got %v", len(result), result) + } + if result[2] != "" { + t.Errorf("last element should be empty string, got %q", result[2]) + } +} + +func TestSplitQuestions_NoNewline(t *testing.T) { + result := SplitQuestions("single") + if len(result) != 1 || result[0] != "single" { + t.Errorf("got %v, want [single]", result) + } +} + +// ============================================================================= +// SplitKeywords — Python: re.split(r"[,，;；、\r\n]+", ...) +// ============================================================================= + +func TestSplitKeywords_Comma(t *testing.T) { + result := SplitKeywords("kw1,kw2,kw3") + if len(result) != 3 { + t.Fatalf("len = %d, want 3", len(result)) + } +} + +func TestSplitKeywords_ChineseComma(t *testing.T) { + result := SplitKeywords("kw1，kw2，kw3") + if len(result) != 3 { + t.Fatalf("len = %d, want 3", len(result)) + } +} + +func TestSplitKeywords_MixedDelimiters(t *testing.T) { + result := SplitKeywords("kw1,kw2；kw3") + if len(result) != 3 { + t.Fatalf("len = %d, want 3, got %v", len(result), result) + } +} + +func TestSplitKeywords_FiltersEmptyStrings(t *testing.T) { + result := SplitKeywords("kw1,,kw2") + // Python: re.split → ["kw1", "", "kw2"] → after filter: ["kw1", "kw2"] + if len(result) != 2 { + t.Errorf("empty strings should be filtered: got %v", result) + } +} + +func TestSplitKeywords_Empty(t *testing.T) { + result := SplitKeywords("") + // Python: re.split on "" → [""], then filtered by if k.strip() + if len(result) != 0 { + t.Errorf("got %v, want empty", result) + } +} + +// ============================================================================= +// CreateChunkTime — Python: datetime.now().timestamp() has sub-second precision +// ============================================================================= + +func TestCreateChunkTime_Format(t *testing.T) { + timeStr, ts := CreateChunkTime() + if timeStr == "" { + t.Error("create_time should not be empty") + } + if len(timeStr) != 19 { + t.Errorf("expected 19 chars (2006-01-02 15:04:05), got %q", timeStr) + } + // Python: datetime.now().timestamp() returns float with fractional part + frac := ts - math.Floor(ts) + if frac <= 0 { + t.Errorf("timestamp should have sub-second precision, got frac=%f", frac) + } +} + +// ============================================================================= +// RenameTextToContentWithWeight — Python processChunks logic +// ============================================================================= + +func TestRenameTextToContentWithWeight_Basic(t *testing.T) { + chunk := map[string]any{"text": "hello world"} + RenameTextToContentWithWeight(chunk) + if _, exists := chunk["text"]; exists { + t.Error("text key should be removed") + } + if chunk["content_with_weight"] != "hello world" { + t.Errorf("content_with_weight = %q, want \"hello world\"", chunk["content_with_weight"]) + } +} + +func TestRenameTextToContentWithWeight_PreservesExisting(t *testing.T) { + chunk := map[string]any{"content_with_weight": "already set", "text": "hello"} + RenameTextToContentWithWeight(chunk) + if chunk["content_with_weight"] != "already set" { + t.Errorf("preserved value should not be overwritten") + } + if _, exists := chunk["text"]; exists { + t.Error("text should still be removed") + } +} + +func TestRenameTextToContentWithWeight_NoTextKey(t *testing.T) { + chunk := map[string]any{"other": "value"} + RenameTextToContentWithWeight(chunk) + if _, exists := chunk["content_with_weight"]; exists { + t.Error("should not add content_with_weight when no text key") + } +} + +// ============================================================================= +// ProcessChunksForDataflow — Python: processChunks() +// ============================================================================= + +func TestProcessChunksForDataflow_SetsDocIDAndKBID(t *testing.T) { + chunks := []map[string]any{{"text": "hello world"}} + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + + if chunks[0]["doc_id"] != "doc-1" { + t.Errorf("doc_id = %q, want \"doc-1\"", chunks[0]["doc_id"]) + } + if kbIDs, ok := chunks[0]["kb_id"].([]string); ok { + if len(kbIDs) != 1 || kbIDs[0] != "kb-1" { + t.Errorf("kb_id = %v, want [\"kb-1\"]", chunks[0]["kb_id"]) + } + } else { + t.Errorf("kb_id should be []string, got %T", chunks[0]["kb_id"]) + } +} + +func TestProcessChunksForDataflow_SetsDocNameKwd(t *testing.T) { + chunks := []map[string]any{{"text": "hello"}} + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + if chunks[0]["docnm_kwd"] != "test-doc.pdf" { + t.Errorf("docnm_kwd = %q, want \"test-doc.pdf\"", chunks[0]["docnm_kwd"]) + } +} + +func TestProcessChunksForDataflow_SetsTimeFields(t *testing.T) { + now := time.Now() + chunks := []map[string]any{{"text": "hello"}} + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", now) + + if timeStr, ok := chunks[0]["create_time"].(string); ok { + if timeStr != now.Format("2006-01-02 15:04:05") { + t.Errorf("create_time = %q, want %q", timeStr, now.Format("2006-01-02 15:04:05")) + } + } else { + t.Errorf("create_time should be string, got %T", chunks[0]["create_time"]) + } + if ts, ok := chunks[0]["create_timestamp_flt"].(float64); ok { + expected := float64(now.UnixMicro()) / 1e6 + if ts != expected { + t.Errorf("create_timestamp_flt = %f, want %f", ts, expected) + } + } else { + t.Errorf("create_timestamp_flt should be float64, got %T", chunks[0]["create_timestamp_flt"]) + } +} + +func TestProcessChunksForDataflow_GeneratesID(t *testing.T) { + chunks := []map[string]any{{"text": "hello"}} + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + id, ok := chunks[0]["id"].(string) + if !ok || id == "" { + t.Errorf("id should be non-empty string, got %v", chunks[0]["id"]) + } +} + +func TestProcessChunksForDataflow_PanicOnListText(t *testing.T) { + chunks := []map[string]any{{"text": []any{"bad-shape"}}} + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for list-shaped text") + } + }() + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) +} + +func TestProcessChunksForDataflow_RemovesInternalPipelineFields(t *testing.T) { + chunks := []map[string]any{{ + "text": "hello", + "_pdf_positions": []any{[]any{0, 1, 2, 3, 4}}, + "image": "data:image/png;base64,abc", + }} + + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + if _, exists := chunks[0]["_pdf_positions"]; exists { + t.Fatalf("_pdf_positions should be removed before indexing: %v", chunks[0]["_pdf_positions"]) + } + if _, exists := chunks[0]["image"]; exists { + t.Fatalf("image should be removed before indexing: %v", chunks[0]["image"]) + } +} + +func TestProcessChunksForDataflow_PreservesExistingID(t *testing.T) { + chunks := []map[string]any{{"text": "hello", "id": "existing-id"}} + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + if chunks[0]["id"] != "existing-id" { + t.Errorf("existing id should be preserved, got %q", chunks[0]["id"]) + } +} + +func TestProcessChunksForDataflow_QuestionsProcessing(t *testing.T) { + chunks := []map[string]any{{"text": "hello", "questions": "Q1\nQ2\nQ3"}} + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + + if _, exists := chunks[0]["questions"]; exists { + t.Error("questions key should be removed") + } + kwd, ok := chunks[0]["question_kwd"].([]string) + if !ok { + t.Fatalf("question_kwd should be []string, got %T", chunks[0]["question_kwd"]) + } + if len(kwd) != 3 { + t.Errorf("question_kwd len = %d, want 3", len(kwd)) + } + if _, ok := chunks[0]["question_tks"].(string); !ok { + t.Errorf("question_tks should be string, got %T", chunks[0]["question_tks"]) + } +} + +func TestProcessChunksForDataflow_KeywordsProcessing(t *testing.T) { + chunks := []map[string]any{{"text": "hello", "keywords": "kw1,kw2;kw3"}} + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + + if _, exists := chunks[0]["keywords"]; exists { + t.Error("keywords key should be removed") + } + kwd, ok := chunks[0]["important_kwd"].([]string) + if !ok || len(kwd) == 0 { + t.Errorf("important_kwd should be non-empty []string, got %v", chunks[0]["important_kwd"]) + } + if _, ok := chunks[0]["important_tks"].(string); !ok { + t.Errorf("important_tks should be string, got %T", chunks[0]["important_tks"]) + } +} + +func TestProcessChunksForDataflow_SummaryProcessing(t *testing.T) { + chunks := []map[string]any{{"text": "hello", "summary": "This is a summary."}} + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + + if _, exists := chunks[0]["summary"]; exists { + t.Error("summary key should be removed") + } + if _, ok := chunks[0]["content_ltks"].(string); !ok { + t.Errorf("content_ltks should be string, got %T", chunks[0]["content_ltks"]) + } + if _, ok := chunks[0]["content_sm_ltks"].(string); !ok { + t.Errorf("content_sm_ltks should be string, got %T", chunks[0]["content_sm_ltks"]) + } +} + +func TestProcessChunksForDataflow_TextRenamed(t *testing.T) { + chunks := []map[string]any{{"text": "hello world"}} + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + + if _, exists := chunks[0]["text"]; exists { + t.Error("text key should be removed") + } + if chunks[0]["content_with_weight"] != "hello world" { + t.Errorf("content_with_weight = %q, want \"hello world\"", chunks[0]["content_with_weight"]) + } +} + +func TestProcessChunksForDataflow_PreservesContentWithWeight(t *testing.T) { + chunks := []map[string]any{{"content_with_weight": "already set", "text": "hello"}} + _ = ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + if chunks[0]["content_with_weight"] != "already set" { + t.Errorf("content_with_weight = %q, want \"already set\"", chunks[0]["content_with_weight"]) + } +} + +func TestProcessChunksForDataflow_MetadataExtraction(t *testing.T) { + chunks := []map[string]any{ + {"text": "hello", "metadata": map[string]any{"author": "Alice", "year": "2024"}}, + } + meta := ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + + if _, exists := chunks[0]["metadata"]; exists { + t.Error("metadata key should be removed from chunk") + } + if meta["author"] != "Alice" { + t.Errorf("metadata[\"author\"] = %q, want \"Alice\"", meta["author"]) + } +} + +func TestProcessChunksForDataflow_MetadataMergeAcrossChunks(t *testing.T) { + chunks := []map[string]any{ + {"text": "chunk1", "metadata": map[string]any{"author": "Alice"}}, + {"text": "chunk2", "metadata": map[string]any{"year": "2024"}}, + } + meta := ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + + if meta["author"] != "Alice" { + t.Errorf("author = %q, want \"Alice\"", meta["author"]) + } + if meta["year"] != "2024" { + t.Errorf("year = %v, want \"2024\"", meta["year"]) + } +} + +func TestProcessChunksForDataflow_NilChunks(t *testing.T) { + meta := ProcessChunksForDataflow(nil, "doc-1", "kb-1", "test-doc.pdf", time.Now()) + if len(meta) != 0 { + t.Errorf("metadata should be empty for nil chunks, got %v", meta) + } +} diff --git a/internal/ingestion/task/chunk_utils.go b/internal/ingestion/task/chunk_utils.go new file mode 100644 index 0000000000..aa88aaf5f6 --- /dev/null +++ b/internal/ingestion/task/chunk_utils.go @@ -0,0 +1,164 @@ +// +// 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 task + +import ( + "fmt" + + "ragflow/internal/common" +) + +// NormalizeChunks converts pipeline output into a uniform []map[string]any slice. +// Mirrors Python: DataflowService._normalize_chunks() +func NormalizeChunks(output map[string]any) []map[string]any { + if output == nil { + return nil + } + + if chunks, ok := output["chunks"].([]map[string]any); ok { + return deepCopyChunks(chunks) + } + if chunks, ok := toChunkMaps(output["chunks"]); ok { + return deepCopyChunks(chunks) + } + if json, ok := output["json"].([]map[string]any); ok { + return deepCopyChunks(json) + } + if json, ok := toChunkMaps(output["json"]); ok { + return deepCopyChunks(json) + } + if md, ok := output["markdown"].(string); ok && md != "" { + return []map[string]any{{"text": md}} + } + if txt, ok := output["text"].(string); ok && txt != "" { + return []map[string]any{{"text": txt}} + } + if html, ok := output["html"].(string); ok && html != "" { + return []map[string]any{{"text": html}} + } + return nil +} + +func toChunkMaps(v any) ([]map[string]any, bool) { + items, ok := v.([]any) + if !ok { + return nil, false + } + out := make([]map[string]any, 0, len(items)) + for _, item := range items { + m, ok := item.(map[string]any) + if !ok { + return nil, false + } + out = append(out, m) + } + return out, true +} + +// deepCopyChunks returns a deep copy of the chunk slice and each chunk map. +// Slice values (e.g. []float64 vectors) are fully copied, not shared. +// Mirrors Python: copy.deepcopy() +func deepCopyChunks(chunks []map[string]any) []map[string]any { + if chunks == nil { + return nil + } + out := make([]map[string]any, len(chunks)) + for i, c := range chunks { + cp := make(map[string]any, len(c)) + for k, v := range c { + switch val := v.(type) { + case []float64: + vec := make([]float64, len(val)) + copy(vec, val) + cp[k] = vec + case []int: + sl := make([]int, len(val)) + copy(sl, val) + cp[k] = sl + case []string: + sl := make([]string, len(val)) + copy(sl, val) + cp[k] = sl + default: + cp[k] = v + } + } + out[i] = cp + } + return out +} + +// PrepareTextsForDataflowEmbedding extracts texts for embedding from chunks. +// Priority: questions > summary > text. +// Mirrors Python: EmbeddingUtils.prepare_texts_for_dataflow_embedding() +func PrepareTextsForDataflowEmbedding(chunks []map[string]any) []string { + if chunks == nil { + return nil + } + texts := make([]string, 0, len(chunks)) + for _, chunk := range chunks { + text, _ := chunk["questions"].(string) + if text == "" { + text, _ = chunk["summary"].(string) + } + if text == "" { + text = MustGetChunkTextString(chunk, "PrepareTextsForDataflowEmbedding") + } + texts = append(texts, text) + } + return texts +} + +// MustGetChunkTextString returns chunk["text"] when it is a string. +// Missing text is allowed and returns empty string. +// FIXME: remove panic before production; current panic is intentional for dev/test +// so list-shaped text payloads are surfaced immediately instead of being written +// as silent bad data. +func MustGetChunkTextString(chunk map[string]any, where string) string { + val, exists := chunk["text"] + if !exists || val == nil { + return "" + } + text, ok := val.(string) + if ok { + return text + } + + msg := fmt.Sprintf("%s: invalid chunk text type %T, expected string, chunk=%v", where, val, chunk) + common.Error(msg, nil) + panic(msg) +} + +// AttachVectors attaches embedding vectors to chunks in-place. +// Each chunk gets a key like "q_{dim}_vec" with the vector as []float64. +// Mirrors Python: EmbeddingUtils.attach_vectors() +func AttachVectors(chunks []map[string]any, vectors [][]float64) int { + if len(chunks) == 0 && len(vectors) == 0 { + return 0 + } + if len(vectors) != len(chunks) { + panic(fmt.Sprintf("vectors/chunks length mismatch: %d != %d", len(vectors), len(chunks))) + } + vectorSize := 0 + for i, doc := range chunks { + vec := vectors[i] + vectorSize = len(vec) + key := fmt.Sprintf("q_%d_vec", vectorSize) + doc[key] = vec + } + return vectorSize +} diff --git a/internal/ingestion/task/chunk_utils_test.go b/internal/ingestion/task/chunk_utils_test.go new file mode 100644 index 0000000000..7bdb042c59 --- /dev/null +++ b/internal/ingestion/task/chunk_utils_test.go @@ -0,0 +1,427 @@ +package task + +import ( + "testing" +) + +// ============================================================================= +// NormalizeChunks +// ============================================================================= + +func TestNormalizeChunks_ChunksFormat(t *testing.T) { + input := map[string]any{ + "chunks": []map[string]any{ + {"text": "hello", "doc_type_kwd": "text"}, + {"text": "world", "doc_type_kwd": "text"}, + }, + } + result := NormalizeChunks(input) + if len(result) != 2 { + t.Fatalf("len = %d, want 2", len(result)) + } + if result[0]["text"] != "hello" { + t.Errorf("result[0][\"text\"] = %q, want \"hello\"", result[0]["text"]) + } +} + +func TestNormalizeChunks_JSONFormat(t *testing.T) { + input := map[string]any{ + "json": []map[string]any{ + {"text": "section 1", "doc_type_kwd": "text"}, + }, + } + result := NormalizeChunks(input) + if len(result) != 1 { + t.Fatalf("len = %d, want 1", len(result)) + } + if result[0]["text"] != "section 1" { + t.Errorf("result[0][\"text\"] = %q, want \"section 1\"", result[0]["text"]) + } +} + +func TestNormalizeChunks_JSONFormatFromGenericSlice(t *testing.T) { + input := map[string]any{ + "json": []any{ + map[string]any{"text": "section 1", "doc_type_kwd": "text"}, + }, + } + result := NormalizeChunks(input) + if len(result) != 1 { + t.Fatalf("len = %d, want 1", len(result)) + } + if result[0]["text"] != "section 1" { + t.Errorf("result[0][\"text\"] = %q, want \"section 1\"", result[0]["text"]) + } +} + +func TestNormalizeChunks_MarkdownFormat(t *testing.T) { + input := map[string]any{ + "markdown": "# Title\n\nContent", + } + result := NormalizeChunks(input) + if len(result) != 1 { + t.Fatalf("len = %d, want 1", len(result)) + } + text, ok := result[0]["text"].(string) + if !ok { + t.Fatalf("text should be string for markdown format, got %T", result[0]["text"]) + } + if text != "# Title\n\nContent" { + t.Errorf("text = %q, want \"# Title\\n\\nContent\"", text) + } +} + +func TestNormalizeChunks_TextFormat(t *testing.T) { + input := map[string]any{ + "text": "plain text", + } + result := NormalizeChunks(input) + if len(result) != 1 { + t.Fatalf("len = %d, want 1", len(result)) + } + text, ok := result[0]["text"].(string) + if !ok { + t.Fatalf("text should be string for text format, got %T", result[0]["text"]) + } + if text != "plain text" { + t.Errorf("text = %q, want \"plain text\"", text) + } +} + +func TestNormalizeChunks_HTMLFormat(t *testing.T) { + input := map[string]any{ + "html": "

Hello

", + } + result := NormalizeChunks(input) + if len(result) != 1 { + t.Fatalf("len = %d, want 1", len(result)) + } + text, ok := result[0]["text"].(string) + if !ok { + t.Fatalf("text should be string for html format, got %T", result[0]["text"]) + } + if text != "

Hello

" { + t.Errorf("text = %q, want \"

Hello

\"", text) + } +} + +func TestNormalizeChunks_EmptyOutput(t *testing.T) { + result := NormalizeChunks(map[string]any{}) + if result != nil { + t.Errorf("expected nil for empty input, got %v", result) + } +} + +func TestNormalizeChunks_EmptyMarkdown(t *testing.T) { + result := NormalizeChunks(map[string]any{"markdown": ""}) + if result != nil { + t.Errorf("expected nil for empty markdown, got %v", result) + } +} + +func TestNormalizeChunks_EmptyText(t *testing.T) { + result := NormalizeChunks(map[string]any{"text": ""}) + if result != nil { + t.Errorf("expected nil for empty text, got %v", result) + } +} + +func TestNormalizeChunks_EmptyHTML(t *testing.T) { + result := NormalizeChunks(map[string]any{"html": ""}) + if result != nil { + t.Errorf("expected nil for empty html, got %v", result) + } +} + +func TestNormalizeChunks_EmptyChunksList(t *testing.T) { + result := NormalizeChunks(map[string]any{"chunks": []map[string]any{}}) + if len(result) != 0 { + t.Errorf("expected empty slice, got len=%d", len(result)) + } +} + +func TestNormalizeChunks_EmptyJSONList(t *testing.T) { + result := NormalizeChunks(map[string]any{"json": []map[string]any{}}) + if len(result) != 0 { + t.Errorf("expected empty slice, got len=%d", len(result)) + } +} + +func TestNormalizeChunks_Priority(t *testing.T) { + t.Run("chunks over json", func(t *testing.T) { + input := map[string]any{ + "chunks": []map[string]any{{"text": "from chunks"}}, + "json": []map[string]any{{"text": "from json"}}, + } + result := NormalizeChunks(input) + if result[0]["text"] != "from chunks" { + t.Errorf("chunks should win: got %q", result[0]["text"]) + } + }) + + t.Run("json over markdown", func(t *testing.T) { + input := map[string]any{ + "json": []map[string]any{{"text": "from json"}}, + "markdown": "from markdown", + } + result := NormalizeChunks(input) + if result[0]["text"] != "from json" { + t.Errorf("json should win: got %q", result[0]["text"]) + } + }) + + t.Run("markdown over text", func(t *testing.T) { + input := map[string]any{ + "markdown": "from markdown", + "text": "from text", + } + result := NormalizeChunks(input) + text, ok := result[0]["text"].(string) + if !ok { + t.Fatalf("text should be string, got %T", result[0]["text"]) + } + if text != "from markdown" { + t.Errorf("markdown should win: got %q", text) + } + }) +} + +func TestNormalizeChunks_DoesNotMutateInput(t *testing.T) { + original := []map[string]any{{"text": "original"}} + input := map[string]any{"chunks": original} + result := NormalizeChunks(input) + result[0]["text"] = "modified" + if original[0]["text"] != "original" { + t.Error("should deep copy, not mutate input") + } +} + +func TestNormalizeChunks_DeepCopyVectors(t *testing.T) { + // Python: copy.deepcopy creates fully independent copies. + // Mutating a slice element in the result must NOT affect the original. + original_vec := []float64{0.1, 0.2, 0.3} + original := []map[string]any{{"text": "hello", "q_3_vec": original_vec}} + input := map[string]any{"chunks": original} + result := NormalizeChunks(input) + // Mutate the slice *element* in-place (not replace the slice) + result[0]["q_3_vec"].([]float64)[0] = 0.9 + // Original must be unchanged + if original[0]["q_3_vec"].([]float64)[0] != 0.1 { + t.Error("mutating result vector element should not affect original") + } +} + +func TestNormalizeChunks_NilInput(t *testing.T) { + result := NormalizeChunks(nil) + if result != nil { + t.Errorf("expected nil for nil input, got %v", result) + } +} + +// ============================================================================= +// PrepareTextsForDataflowEmbedding +// ============================================================================= + +func TestPrepareTexts_QuestionsPriority(t *testing.T) { + chunks := []map[string]any{ + {"questions": "Q1\nQ2", "summary": "a summary", "text": "plain text"}, + } + result := PrepareTextsForDataflowEmbedding(chunks) + if len(result) != 1 { + t.Fatalf("len = %d, want 1", len(result)) + } + if result[0] != "Q1\nQ2" { + t.Errorf("questions should take priority: got %q", result[0]) + } +} + +func TestPrepareTexts_SummaryFallback(t *testing.T) { + chunks := []map[string]any{ + {"summary": "a summary", "text": "plain text"}, + } + result := PrepareTextsForDataflowEmbedding(chunks) + if result[0] != "a summary" { + t.Errorf("summary should be used when no questions: got %q", result[0]) + } +} + +func TestPrepareTexts_TextFallback(t *testing.T) { + chunks := []map[string]any{ + {"text": "plain text"}, + } + result := PrepareTextsForDataflowEmbedding(chunks) + if result[0] != "plain text" { + t.Errorf("text should be used when no questions/summary: got %q", result[0]) + } +} + +func TestPrepareTexts_EmptyStringFallback(t *testing.T) { + chunks := []map[string]any{ + {"text": ""}, + } + result := PrepareTextsForDataflowEmbedding(chunks) + if result[0] != "" { + t.Errorf("expected empty string, got %q", result[0]) + } +} + +func TestPrepareTexts_MultipleChunks(t *testing.T) { + chunks := []map[string]any{ + {"questions": "Q1", "text": "t1"}, + {"summary": "S2", "text": "t2"}, + {"text": "t3"}, + } + result := PrepareTextsForDataflowEmbedding(chunks) + if len(result) != 3 { + t.Fatalf("len = %d, want 3", len(result)) + } + if result[0] != "Q1" { + t.Errorf("result[0] = %q, want \"Q1\"", result[0]) + } + if result[1] != "S2" { + t.Errorf("result[1] = %q, want \"S2\"", result[1]) + } + if result[2] != "t3" { + t.Errorf("result[2] = %q, want \"t3\"", result[2]) + } +} + +func TestPrepareTexts_NilChunks(t *testing.T) { + result := PrepareTextsForDataflowEmbedding(nil) + if result != nil { + t.Errorf("expected nil for nil chunks, got %v", result) + } +} + +func TestPrepareTexts_EmptyChunks(t *testing.T) { + result := PrepareTextsForDataflowEmbedding([]map[string]any{}) + if len(result) != 0 { + t.Errorf("expected empty slice, got len=%d", len(result)) + } +} + +func TestPrepareTexts_MissingTextKey(t *testing.T) { + chunks := []map[string]any{ + {"other_key": "value"}, + } + result := PrepareTextsForDataflowEmbedding(chunks) + if result[0] != "" { + t.Errorf("expected empty string for missing text key, got %q", result[0]) + } +} + +func TestPrepareTexts_PanicOnListText(t *testing.T) { + chunks := []map[string]any{ + {"text": []any{"bad-shape"}}, + } + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for list-shaped text") + } + }() + _ = PrepareTextsForDataflowEmbedding(chunks) +} + +func TestMustGetChunkTextString_PanicOnStringSlice(t *testing.T) { + chunk := map[string]any{"text": []string{"bad-shape"}} + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for []string text") + } + }() + _ = MustGetChunkTextString(chunk, "unit-test") +} + +// ============================================================================= +// AttachVectors +// ============================================================================= + +func TestAttachVectors_Basic(t *testing.T) { + chunks := []map[string]any{ + {"text": "hello"}, + {"text": "world"}, + } + vectors := [][]float64{ + {0.1, 0.2, 0.3}, + {0.4, 0.5, 0.6}, + } + dim := AttachVectors(chunks, vectors) + if dim != 3 { + t.Errorf("dim = %d, want 3", dim) + } + if len(chunks[0]) == 0 || chunks[0]["q_3_vec"] == nil { + t.Errorf("expected q_3_vec in chunk[0], got %v", chunks[0]) + } + vec0 := chunks[0]["q_3_vec"].([]float64) + if len(vec0) != 3 || vec0[0] != 0.1 { + t.Errorf("chunk[0] vector = %v, want [0.1 0.2 0.3]", vec0) + } + vec1 := chunks[1]["q_3_vec"].([]float64) + if len(vec1) != 3 || vec1[2] != 0.6 { + t.Errorf("chunk[1] vector = %v, want [0.4 0.5 0.6]", vec1) + } +} + +func TestAttachVectors_EmptyChunks(t *testing.T) { + result := AttachVectors(nil, nil) + if result != 0 { + t.Errorf("expected 0 for empty chunks, got %d", result) + } +} + +func TestAttachVectors_DifferentDimensions(t *testing.T) { + chunks := []map[string]any{ + {"text": "a"}, + {"text": "b"}, + } + vectors := [][]float64{ + {0.1, 0.2}, + {0.3, 0.4, 0.5}, + } + // Each chunk gets its own vector with key based on its dimension + AttachVectors(chunks, vectors) + if chunks[0]["q_2_vec"] == nil { + t.Errorf("chunk[0] should have q_2_vec, got keys: %v", chunkKeys(chunks[0])) + } + if chunks[1]["q_3_vec"] == nil { + t.Errorf("chunk[1] should have q_3_vec, got keys: %v", chunkKeys(chunks[1])) + } +} + +func TestAttachVectors_MismatchedCount(t *testing.T) { + chunks := []map[string]any{ + {"text": "hello"}, + } + vectors := [][]float64{ + {0.1}, + {0.2}, + } + defer func() { + if r := recover(); r == nil { + t.Error("expected panic for mismatched count") + } + }() + AttachVectors(chunks, vectors) +} + +func TestAttachVectors_KeyFormat(t *testing.T) { + chunks := []map[string]any{ + {"text": "hello"}, + } + vectors := [][]float64{ + {0.1, 0.2, 0.3, 0.4, 0.5}, + } + AttachVectors(chunks, vectors) + key := "q_5_vec" + if chunks[0][key] == nil { + t.Errorf("expected key %q in chunk, got keys: %v", key, chunkKeys(chunks[0])) + } +} + +func chunkKeys(m map[string]any) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + return keys +} diff --git a/internal/ingestion/task/constants.go b/internal/ingestion/task/constants.go new file mode 100644 index 0000000000..fe6adefd41 --- /dev/null +++ b/internal/ingestion/task/constants.go @@ -0,0 +1,29 @@ +// +// 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 task + +// Special doc_id values used in task messages (mirrors Python constants). +const ( + // CANVAS_DEBUG_DOC_ID marks a canvas debug dataflow task — no persistence. + CANVAS_DEBUG_DOC_ID = "dataflow_x" + + // GRAPH_RAPTOR_FAKE_DOC_ID is the fake doc_id used for RAPTOR-generated chunks. + GRAPH_RAPTOR_FAKE_DOC_ID = "graph_raptor_fake_doc" + + // EmbeddingTokenConsumptionKey is the key in pipeline output for embedding token count. + EmbeddingTokenConsumptionKey = "embedding_token_consumption" +) diff --git a/internal/ingestion/task/dataflow_chunk_assets.go b/internal/ingestion/task/dataflow_chunk_assets.go new file mode 100644 index 0000000000..f297e65e66 --- /dev/null +++ b/internal/ingestion/task/dataflow_chunk_assets.go @@ -0,0 +1,14 @@ +package task + +// PrepareDataflowChunkAssets applies the minimal pre-index cleanup needed by +// the Go dataflow path so real pipeline output can be stored safely. +func PrepareDataflowChunkAssets(chunks []map[string]any) error { + for _, ck := range chunks { + delete(ck, "_pdf_positions") + // FIXME: production parity with Python _prepare_docs_and_upload should + // upload raw image data and set img_id. During the current bring-up phase + // we drop raw image payloads so the main dataflow -> ES path can proceed. + delete(ck, "image") + } + return nil +} diff --git a/internal/ingestion/task/dataflow_e2e_test.go b/internal/ingestion/task/dataflow_e2e_test.go new file mode 100644 index 0000000000..2aa6199439 --- /dev/null +++ b/internal/ingestion/task/dataflow_e2e_test.go @@ -0,0 +1,370 @@ +// +// 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 task + +import ( + "context" + "fmt" + "net" + "os" + "testing" + "time" + + "ragflow/internal/dao" + "ragflow/internal/engine" + "ragflow/internal/engine/elasticsearch" + "ragflow/internal/engine/infinity" + "ragflow/internal/entity" + "ragflow/internal/ingestion/testutil" + "ragflow/internal/server" +) + +// ============================================================================= +// Test helper: setupTestDocEngine - Creates either Elasticsearch or Infinity engine +// ============================================================================= + +func setupTestDocEngine(t *testing.T, engineType engine.EngineType, tenantID, datasetID string) (engine.DocEngine, func()) { + t.Helper() + + var ( + docEngine engine.DocEngine + err error + ) + + switch engineType { + case engine.EngineElasticsearch: + t.Logf("Setting up Elasticsearch engine...") + esHost := os.Getenv("ES_HOST") + if esHost == "" { + esHost = "localhost:1200" + } + if !startsWithHTTP(esHost) { + esHost = "http://" + esHost + } + esUser := os.Getenv("ES_USER") + if esUser == "" { + esUser = "elastic" + } + esPassword := os.Getenv("ES_PASSWORD") + if esPassword == "" { + esPassword = "infini_rag_flow" + } + + cfg := &server.ElasticsearchConfig{ + Hosts: esHost, + Username: esUser, + Password: esPassword, + } + docEngine, err = elasticsearch.NewEngine(cfg) + if err != nil { + t.Skipf("Could not create Elasticsearch engine: %v (skipping ES subtest)", err) + return nil, func() {} + } + + case engine.EngineInfinity: + t.Logf("Setting up Infinity engine...") + infURI := os.Getenv("INFINITY_URI") + if infURI == "" { + infURI = "localhost:23817" + } + + // First check if Infinity is reachable quickly without waiting 120s! + if !isPortOpen("localhost", 23817) { + t.Skipf("Infinity not running at %s (skipping Infinity subtest)", infURI) + return nil, func() {} + } + + cfg := &server.InfinityConfig{ + URI: infURI, + DBName: "ragflow_e2e_test", + PostgresPort: 5432, + } + docEngine, err = infinity.NewEngine(cfg) + if err != nil { + t.Skipf("Could not create Infinity engine: %v (skipping Infinity subtest)", err) + return nil, func() {} + } + + default: + t.Fatalf("Unsupported engine type: %v", engineType) + } + + // Create unique base name for the test + baseName := fmt.Sprintf("ragflow_%s", tenantID) + + // Cleanup first (if exists) + ctx := context.Background() + _ = docEngine.DropChunkStore(ctx, baseName, datasetID) + + // Create chunk store (note: vec dimension = 2 because our test chunks have q_2_vec) + if err := docEngine.CreateChunkStore(ctx, baseName, datasetID, 2, "naive"); err != nil { + // If create failed, maybe it exists; try dropping and recreating + _ = docEngine.DropChunkStore(ctx, baseName, datasetID) + if err := docEngine.CreateChunkStore(ctx, baseName, datasetID, 2, "naive"); err != nil { + _ = docEngine.Close() + t.Fatalf("Could not create chunk store: %v", err) + } + } + + // Return cleanup function + cleanup := func() { + _ = docEngine.DropChunkStore(ctx, baseName, datasetID) + _ = docEngine.Close() + } + + return docEngine, cleanup +} + +func startsWithHTTP(s string) bool { + return len(s) >= 4 && (s[:4] == "http" || s[:5] == "https") +} + +func toLowerSnakeCase(s string) string { + result := make([]rune, 0, len(s)) + for _, r := range s { + if r >= 'A' && r <= 'Z' { + result = append(result, r-'A'+'a') + } else if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') { + result = append(result, r) + } else { + result = append(result, '_') + } + } + return string(result) +} + +func isPortOpen(host string, port int) bool { + addr := fmt.Sprintf("%s:%d", host, port) + conn, err := net.DialTimeout("tcp", addr, 1*time.Second) + if err != nil { + return false + } + defer conn.Close() + return true +} + +// ============================================================================= +// Main E2E Test with Subtests for Elasticsearch and Infinity +// ============================================================================= + +func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { + testCases := []struct { + name string + engineType engine.EngineType + }{ + { + name: "Elasticsearch", + engineType: engine.EngineElasticsearch, + }, + { + name: "Infinity", + engineType: engine.EngineInfinity, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Logf("Running Dataflow E2E test with engine: %s", tc.name) + + // Setup test database + db := testutil.SetupTestDB(t) + cleanupDB := testutil.ReplaceDBForTest(t, db) + defer cleanupDB() + + // Seed test data (lowercase for ES index compatibility) + lowerName := toLowerSnakeCase(tc.name) + tenantID, kbID, _, taskID := testutil.SeedTestData(t, db, + testutil.WithTenantID(fmt.Sprintf("tenant-e2e-%s", lowerName)), + testutil.WithKBID(fmt.Sprintf("kb-e2e-%s", lowerName)), + testutil.WithDocID(fmt.Sprintf("doc-e2e-%s", lowerName)), + testutil.WithTaskID(fmt.Sprintf("task-e2e-%s", lowerName)), + testutil.WithPipelineID(fmt.Sprintf("pipeline-e2e-%s", lowerName)), + testutil.WithDocName(fmt.Sprintf("e2e-test-%s.pdf", lowerName)), + ) + + // Setup DocEngine for this test + docEngine, cleanupEngine := setupTestDocEngine(t, tc.engineType, tenantID, kbID) + defer cleanupEngine() + + // Update task to have dataflow task type + var task entity.Task + if err := db.Where("id = ?", taskID).First(&task).Error; err != nil { + t.Fatalf("Failed to get task: %v", err) + } + task.TaskType = "dataflow" + if err := db.Save(&task).Error; err != nil { + t.Fatalf("Failed to update task: %v", err) + } + + // Load task context + taskCtx, err := LoadTaskContext(taskID) + if err != nil { + t.Fatalf("LoadTaskContext failed: %v", err) + } + + // Track what was called + var ( + loadDSLCalled bool + runPipelineCalled bool + insertChunksCalled bool + ) + var capturedChunks [][]map[string]any + + // Create TaskHandler with mocked DataflowService factory + handler := NewTaskHandler(taskCtx) + handler.WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { + svc := mustNewDataflowService(t, ctx, dataflowID, 0, 0) + + // Mock loadDSLFunc + svc.WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) { + loadDSLCalled = true + return `{"nodes":[{"id":"test","type":"parser"}],"edges":[]}`, dataflowID, nil + }) + + // Mock runPipelineFunc - returns test chunks (with vectors to skip embedding) + svc.WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { + runPipelineCalled = true + return map[string]any{ + "chunks": []map[string]any{ + { + "text": fmt.Sprintf("Hello world from E2E test with %s", tc.name), + "id": fmt.Sprintf("chunk-e2e-%s-1", tc.name), + "q_2_vec": []float64{0.1, 0.2}, // Pre-vectorized to skip embedding + }, + { + "text": fmt.Sprintf("Second chunk from E2E test with %s", tc.name), + "id": fmt.Sprintf("chunk-e2e-%s-2", tc.name), + "q_2_vec": []float64{0.3, 0.4}, // Pre-vectorized to skip embedding + }, + }, + EmbeddingTokenConsumptionKey: 100, + }, dsl, nil + }) + + // Use the injected DocEngine for insertChunks! + svc.WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) { + insertChunksCalled = true + t.Logf("DocEngine InsertChunks called! baseName=%s datasetID=%s len(chunks)=%d", baseName, datasetID, len(chunks)) + ids, err := docEngine.InsertChunks(ctx, chunks, baseName, datasetID) + if err != nil { + t.Logf("WARNING: InsertChunks err=%v", err) + } + capturedChunks = append(capturedChunks, chunks) + return ids, err + }) + + svc.WithChunkCounter(&stubChunkCounter{}) + return svc, nil + }) + + // Also set progress callback + var progressEvents []struct { + prog float64 + msg string + } + taskCtx.ProgressFunc = func(prog float64, msg string) { + t.Logf("PROGRESS: %.2f %s", prog, msg) + progressEvents = append(progressEvents, struct { + prog float64 + msg string + }{prog, msg}) + } + + // Execute the task handler! + t.Logf("Calling TaskHandler.Handle()...") + err = handler.Handle() + if err != nil { + t.Fatalf("TaskHandler.Handle failed: %v", err) + } + t.Logf("TaskHandler.Handle() complete!") + + // Verify all the expected calls happened + if !loadDSLCalled { + t.Fatal("Expected loadDSLFunc to be called") + } + if !runPipelineCalled { + t.Fatal("Expected runPipelineFunc to be called") + } + if !insertChunksCalled { + t.Fatal("Expected insertChunks to be called") + } + + // Verify chunks were captured + totalChunks := 0 + for _, batch := range capturedChunks { + totalChunks += len(batch) + } + if totalChunks != 2 { + t.Fatalf("Expected total 2 chunks, got %d", totalChunks) + } + + // Now verify we can read the chunks back! + baseName := fmt.Sprintf("ragflow_%s", tenantID) + t.Logf("Reading chunks back from %s: baseName=%s datasetID=%s", tc.name, baseName, kbID) + + // Refresh (wait for ES to index or Infinity to commit) + time.Sleep(300 * time.Millisecond) + + for _, batch := range capturedChunks { + for _, chunk := range batch { + chunkID, ok := chunk["id"].(string) + if ok && chunkID != "" { + t.Logf("Trying to get chunk: %s", chunkID) + readBack, err := docEngine.GetChunk(context.Background(), baseName, chunkID, []string{kbID}) + if err != nil { + t.Logf("WARNING: Failed to get chunk %q: %v", chunkID, err) + } else if readBack == nil { + t.Logf("WARNING: chunk %q not found (may not have been indexed yet)", chunkID) + } else { + t.Logf("SUCCESS: Read back chunk %q!", chunkID) + } + } + } + } + + // Verify progress reported + if len(progressEvents) == 0 { + t.Fatal("Expected at least one progress event") + } + foundDone := false + for _, ev := range progressEvents { + if ev.prog == 1.0 { + foundDone = true + } + } + if !foundDone { + t.Fatal("Expected progress to reach 1.0") + } + + // Verify final task status can be updated to success + ingestionTaskDAO := dao.NewIngestionTaskDAO() + if err := ingestionTaskDAO.UpdateStatus(taskID, "success"); err != nil { + t.Fatalf("UpdateStatus failed: %v", err) + } + + finalTask, err := ingestionTaskDAO.GetByID(taskID) + if err != nil { + t.Fatalf("GetByID failed: %v", err) + } + if finalTask.Status != "success" { + t.Errorf("Final task status = %q, want %q", finalTask.Status, "success") + } + + t.Logf("SUCCESS: Dataflow E2E test passed with %s engine!", tc.name) + }) + } +} diff --git a/internal/ingestion/task/dataflow_real_pipeline_integration_test.go b/internal/ingestion/task/dataflow_real_pipeline_integration_test.go new file mode 100644 index 0000000000..c97fd779bd --- /dev/null +++ b/internal/ingestion/task/dataflow_real_pipeline_integration_test.go @@ -0,0 +1,769 @@ +//go:build integration +// +build integration + +package task + +import ( + "archive/zip" + "context" + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "sort" + "strings" + "testing" + "time" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/engine" + enginetypes "ragflow/internal/engine/types" + "ragflow/internal/entity" + _ "ragflow/internal/ingestion/component" + componentpkg "ragflow/internal/ingestion/component" + _ "ragflow/internal/ingestion/component/chunker" + pipelinepkg "ragflow/internal/ingestion/pipeline" + "ragflow/internal/server" + "ragflow/internal/storage" + "ragflow/internal/tokenizer" + + "go.uber.org/zap" + "gorm.io/driver/mysql" + "gorm.io/gorm" + gormlogger "gorm.io/gorm/logger" +) + +func TestDataflowService_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) { + prepareTokenizerResourceForTaskIntegration(t) + requireTokenizerPool(t) + + cfg := mustLoadTaskRealIntegrationConfig(t) + realDB := mustOpenTaskRealMySQL(t, cfg) + if err := realDB.AutoMigrate( + &entity.Tenant{}, + &entity.Knowledgebase{}, + &entity.Document{}, + &entity.File{}, + &entity.File2Document{}, + &entity.UserCanvas{}, + ); err != nil { + t.Fatalf("auto-migrate real mysql tables: %v", err) + } + + realStorage, err := storage.NewMinioStorage(cfg.StorageEngine.Minio) + if err != nil { + t.Fatalf("connect real minio: %v", err) + } + + origDB := dao.DB + origStorage := storage.GetStorageFactory().GetStorage() + dao.DB = realDB + storage.GetStorageFactory().SetStorage(realStorage) + t.Cleanup(func() { + dao.DB = origDB + storage.GetStorageFactory().SetStorage(origStorage) + }) + + templatePath := filepath.Join(taskRepoRoot(t), "agent", "templates", "ingestion_pipeline_general.json") + templateBytes, err := os.ReadFile(templatePath) + if err != nil { + t.Fatalf("read template: %v", err) + } + templateBytes = disableTokenizerEmbeddingForTaskTemplate(t, templateBytes) + var templateDSL entity.JSONMap + if err := json.Unmarshal(templateBytes, &templateDSL); err != nil { + t.Fatalf("unmarshal template dsl: %v", err) + } + + suffix := fmt.Sprintf("%d", time.Now().UnixNano()) + tenantID := taskLimit32("it_tenant_" + suffix) + kbID := taskLimit32("it_kb_" + suffix) + docID := taskLimit32("it_doc_" + suffix) + fileID := taskLimit32("it_file_" + suffix) + canvasID := taskLimit32("it_canvas_" + suffix) + bucket := taskS3SafeBucketName(kbID) + objectPath := fmt.Sprintf("integration/task/%s/template-general.txt", docID) + docName := "template-general.txt" + content := "Alpha paragraph\n\nBeta paragraph." + + mustSeedTaskRealPipelineDocument(t, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath, docName, content) + if err := realDB.Model(&entity.Document{}).Where("id = ?", docID).Update("pipeline_id", canvasID).Error; err != nil { + t.Fatalf("set document pipeline_id: %v", err) + } + if err := realDB.Create(&entity.UserCanvas{ + ID: canvasID, + UserID: tenantID, + Permission: "me", + CanvasCategory: "agent_canvas", + DSL: templateDSL, + }).Error; err != nil { + t.Fatalf("create user canvas: %v", err) + } + t.Cleanup(func() { + _ = realDB.Where("id = ?", canvasID).Delete(&entity.UserCanvas{}).Error + cleanupTaskRealPipelineDocument(realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath) + }) + + taskCtx := &TaskContext{ + Task: entity.Task{ID: "task-real-canvas-1", DocID: docID, TaskType: "dataflow"}, + Doc: entity.Document{ + ID: docID, + KbID: kbID, + Name: taskStrPtr(docName), + PipelineID: taskStrPtr(canvasID), + }, + KB: entity.Knowledgebase{ + ID: kbID, + TenantID: tenantID, + EmbdID: "embd-1", + }, + Tenant: entity.Tenant{ID: tenantID}, + ProgressFunc: func(prog float64, msg string) {}, + } + + var inserted [][]map[string]any + svc := mustNewDataflowService(t, taskCtx, canvasID, 0, 0). + WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + inserted = append(inserted, deepCopyTaskChunks(chunks)) + return nil, nil + }). + WithChunkCounter(&stubChunkCounter{}) + + if err := svc.Run(context.Background()); err != nil { + t.Fatalf("Run: %v", err) + } + if len(inserted) != 1 { + t.Fatalf("insert calls = %d, want 1", len(inserted)) + } + if len(inserted[0]) != 2 { + t.Fatalf("inserted chunk count = %d, want 2", len(inserted[0])) + } + for i, ck := range inserted[0] { + if got := ck["doc_id"]; got != docID { + t.Fatalf("chunks[%d].doc_id = %v, want %q", i, got, docID) + } + if got := ck["content_with_weight"]; got == nil || got == "" { + t.Fatalf("chunks[%d].content_with_weight = %v, want non-empty string", i, got) + } + } +} + +func TestDataflowService_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *testing.T) { + prepareTokenizerResourceForTaskIntegration(t) + requireTokenizerPool(t) + + cfg := mustLoadTaskRealIntegrationConfig(t) + realDB := mustOpenTaskRealMySQL(t, cfg) + if err := realDB.AutoMigrate( + &entity.Tenant{}, + &entity.Knowledgebase{}, + &entity.Document{}, + &entity.File{}, + &entity.File2Document{}, + &entity.UserCanvas{}, + ); err != nil { + t.Fatalf("auto-migrate real mysql tables: %v", err) + } + + realStorage, err := storage.NewMinioStorage(cfg.StorageEngine.Minio) + if err != nil { + t.Fatalf("connect real minio: %v", err) + } + if err := engine.Init(&cfg.DocEngine); err != nil { + t.Fatalf("init real doc engine: %v", err) + } + if engine.Get() == nil { + t.Fatal("doc engine is nil after init") + } + if engine.GetEngineType() != engine.EngineElasticsearch { + t.Fatalf("doc engine type = %s, want %s", engine.GetEngineType(), engine.EngineElasticsearch) + } + + origDB := dao.DB + origStorage := storage.GetStorageFactory().GetStorage() + origDocResolver := componentpkg.ResolveDocumentStorageOverride + dao.DB = realDB + storage.GetStorageFactory().SetStorage(realStorage) + componentpkg.ResolveDocumentStorageOverride = nil + t.Cleanup(func() { + dao.DB = origDB + storage.GetStorageFactory().SetStorage(origStorage) + componentpkg.ResolveDocumentStorageOverride = origDocResolver + }) + + templatePath := filepath.Join(taskRepoRoot(t), "agent", "templates", "ingestion_pipeline_general.json") + templateBytes, err := os.ReadFile(templatePath) + if err != nil { + t.Fatalf("read template: %v", err) + } + templateBytes = disableTokenizerEmbeddingForTaskTemplate(t, templateBytes) + var templateDSL entity.JSONMap + if err := json.Unmarshal(templateBytes, &templateDSL); err != nil { + t.Fatalf("unmarshal template dsl: %v", err) + } + + pdfPath := filepath.Join(taskRepoRoot(t), "internal", "deepdoc", "parser", "pdf", "testdata", "pdfs", "01_english_simple.pdf") + pdfBytes, err := os.ReadFile(pdfPath) + if err != nil { + t.Fatalf("read pdf fixture: %v", err) + } + + suffix := fmt.Sprintf("%d", time.Now().UnixNano()) + tenantID := taskLimit32("it_tenant_" + suffix) + kbID := taskLimit32("it_kb_" + suffix) + docID := taskLimit32("it_doc_" + suffix) + fileID := taskLimit32("it_file_" + suffix) + canvasID := taskLimit32("it_canvas_" + suffix) + bucket := taskS3SafeBucketName(kbID) + docName := "01_english_simple.pdf" + objectPath := fmt.Sprintf("integration/task/%s/%s", docID, docName) + baseName := fmt.Sprintf("ragflow_%s", tenantID) + + mustSeedTaskRealPipelineDocumentBytes(t, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath, docName, ".pdf", "pdf", pdfBytes) + if err := realDB.Model(&entity.Document{}).Where("id = ?", docID).Update("pipeline_id", canvasID).Error; err != nil { + t.Fatalf("set document pipeline_id: %v", err) + } + if err := realDB.Create(&entity.UserCanvas{ + ID: canvasID, + UserID: tenantID, + Permission: "me", + CanvasCategory: "agent_canvas", + DSL: templateDSL, + }).Error; err != nil { + t.Fatalf("create user canvas: %v", err) + } + t.Cleanup(func() { + _ = engine.Get().DropChunkStore(context.Background(), baseName, kbID) + _ = realDB.Where("id = ?", canvasID).Delete(&entity.UserCanvas{}).Error + cleanupTaskRealPipelineDocument(realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath) + }) + + taskCtx := &TaskContext{ + Task: entity.Task{ID: "task-real-pdf-es-1", DocID: docID, TaskType: "dataflow"}, + Doc: entity.Document{ + ID: docID, + KbID: kbID, + Name: taskStrPtr(docName), + PipelineID: taskStrPtr(canvasID), + }, + KB: entity.Knowledgebase{ + ID: kbID, + TenantID: tenantID, + EmbdID: "embd-1", + }, + Tenant: entity.Tenant{ID: tenantID}, + ProgressFunc: func(prog float64, msg string) {}, + } + + svc := mustNewDataflowService(t, taskCtx, canvasID, 0, 0). + WithChunkCounter(&stubChunkCounter{}) + + if err := svc.Run(context.Background()); err != nil { + t.Fatalf("Run: %v", err) + } + + result, err := engine.Get().Search(context.Background(), &enginetypes.SearchRequest{ + IndexNames: []string{baseName}, + KbIDs: []string{kbID}, + Limit: 20, + }) + if err != nil { + t.Fatalf("search indexed chunks: %v", err) + } + if result == nil { + t.Fatal("search result is nil") + } + if len(result.Chunks) == 0 { + t.Fatal("expected indexed chunks in Elasticsearch, got 0") + } + for i, chunk := range result.Chunks { + if got := chunk["doc_id"]; got != docID { + t.Fatalf("result chunk[%d].doc_id = %v, want %q", i, got, docID) + } + if got := chunk["kb_id"]; got != kbID { + t.Fatalf("result chunk[%d].kb_id = %v, want %q", i, got, kbID) + } + if got := chunk["docnm_kwd"]; got != docName { + t.Fatalf("result chunk[%d].docnm_kwd = %v, want %q", i, got, docName) + } + if got := chunk["content_with_weight"]; got == nil || got == "" { + t.Fatalf("result chunk[%d].content_with_weight = %v, want non-empty", i, got) + } + } +} + +func TestRunDataflow_RealPipelineOutput_ProducesIndexFields(t *testing.T) { + prepareTokenizerResourceForTaskIntegration(t) + requireTokenizerPool(t) + + cfg := mustLoadTaskRealIntegrationConfig(t) + realDB := mustOpenTaskRealMySQL(t, cfg) + if err := realDB.AutoMigrate( + &entity.Tenant{}, + &entity.Knowledgebase{}, + &entity.Document{}, + &entity.File{}, + &entity.File2Document{}, + ); err != nil { + t.Fatalf("auto-migrate real mysql tables: %v", err) + } + + realStorage, err := storage.NewMinioStorage(cfg.StorageEngine.Minio) + if err != nil { + t.Fatalf("connect real minio: %v", err) + } + + origDB := dao.DB + origStorage := storage.GetStorageFactory().GetStorage() + dao.DB = realDB + storage.GetStorageFactory().SetStorage(realStorage) + t.Cleanup(func() { + dao.DB = origDB + storage.GetStorageFactory().SetStorage(origStorage) + }) + + templatePath := filepath.Join(taskRepoRoot(t), "agent", "templates", "ingestion_pipeline_general.json") + templateBytes, err := os.ReadFile(templatePath) + if err != nil { + t.Fatalf("read template: %v", err) + } + templateBytes = disableTokenizerEmbeddingForTaskTemplate(t, templateBytes) + + pipe, err := pipelinepkg.NewPipelineFromDSL(templateBytes, "task-dataflow-real-pipeline") + if err != nil { + t.Fatalf("NewPipelineFromDSL: %v", err) + } + + suffix := fmt.Sprintf("%d", time.Now().UnixNano()) + tenantID := taskLimit32("it_tenant_" + suffix) + kbID := taskLimit32("it_kb_" + suffix) + docID := taskLimit32("it_doc_" + suffix) + fileID := taskLimit32("it_file_" + suffix) + bucket := taskS3SafeBucketName(kbID) + objectPath := fmt.Sprintf("integration/task/%s/template-general.txt", docID) + docName := "template-general.txt" + content := "Alpha paragraph.\n\nBeta paragraph." + + mustSeedTaskRealPipelineDocument(t, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath, docName, content) + t.Cleanup(func() { + cleanupTaskRealPipelineDocument(realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath) + }) + + pipelineOut, err := pipe.Run(context.Background(), map[string]any{ + "doc_id": docID, + }) + if err != nil { + t.Fatalf("pipeline Run: %v", err) + } + pipelineOut = taskTerminalPayloadFromRunOutput(t, pipelineOut, "Tokenizer:LegalReadersDecide") + + taskCtx := &TaskContext{ + Task: entity.Task{ID: "task-real-1", DocID: docID}, + Doc: entity.Document{ + ID: docID, + KbID: kbID, + Name: taskStrPtr(docName), + }, + KB: entity.Knowledgebase{ + ID: kbID, + TenantID: tenantID, + EmbdID: "embd-1", + }, + Tenant: entity.Tenant{ID: tenantID}, + ProgressFunc: func(prog float64, msg string) {}, + } + + var inserted [][]map[string]any + svc := mustNewDataflowService(t, taskCtx, "flow-real-1", 0, 0). + WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + inserted = append(inserted, deepCopyTaskChunks(chunks)) + return nil, nil + }). + WithChunkCounter(&stubChunkCounter{}) + + if err := svc.RunDataflow(context.Background(), pipelineOut); err != nil { + t.Fatalf("RunDataflow: %v", err) + } + + if len(inserted) != 1 { + t.Fatalf("insert calls = %d, want 1", len(inserted)) + } + chunks := inserted[0] + if len(chunks) != 2 { + t.Fatalf("inserted chunk count = %d, want 2", len(chunks)) + } + for i, ck := range chunks { + if got := ck["doc_id"]; got != docID { + t.Fatalf("chunks[%d].doc_id = %v, want %q", i, got, docID) + } + if got := ck["docnm_kwd"]; got != docName { + t.Fatalf("chunks[%d].docnm_kwd = %v, want %q", i, got, docName) + } + if got := ck["content_with_weight"]; got == nil || got == "" { + t.Fatalf("chunks[%d].content_with_weight = %v, want non-empty string", i, got) + } + if got, ok := ck["content_ltks"].(string); !ok || got == "" { + t.Fatalf("chunks[%d].content_ltks = %T/%v, want non-empty string", i, ck["content_ltks"], ck["content_ltks"]) + } + if got, ok := ck["content_sm_ltks"].(string); !ok || got == "" { + t.Fatalf("chunks[%d].content_sm_ltks = %T/%v, want non-empty string", i, ck["content_sm_ltks"], ck["content_sm_ltks"]) + } + if _, hasText := ck["text"]; hasText { + t.Fatalf("chunks[%d] should not keep raw text field after RunDataflow: %v", i, ck["text"]) + } + } +} + +func taskRepoRoot(t *testing.T) string { + t.Helper() + wd, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + return filepath.Clean(filepath.Join(wd, "..", "..", "..")) +} + +func mustLoadTaskRealIntegrationConfig(t *testing.T) *server.Config { + t.Helper() + if err := common.Init("info", common.FileOutput{}); err != nil { + t.Fatalf("init common logger: %v", err) + } + server.SetLogger(zap.NewNop()) + configPath := filepath.Join(taskRepoRoot(t), "conf", "service_conf.yaml") + if err := server.Init(configPath); err != nil { + t.Fatalf("init service config from %s: %v", configPath, err) + } + cfg := server.GetConfig() + if cfg == nil || cfg.Database.Host == "" || cfg.StorageEngine.Minio == nil || cfg.StorageEngine.Minio.Host == "" { + t.Fatal("real integration config is incomplete") + } + return cfg +} + +func mustOpenTaskRealMySQL(t *testing.T, cfg *server.Config) *gorm.DB { + t.Helper() + dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=Local", + cfg.Database.Username, + cfg.Database.Password, + cfg.Database.Host, + cfg.Database.Port, + cfg.Database.Database, + cfg.Database.Charset, + ) + db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{ + Logger: gormlogger.Default.LogMode(gormlogger.Silent), + }) + if err != nil { + t.Fatalf("connect real mysql: %v", err) + } + return db +} + +func disableTokenizerEmbeddingForTaskTemplate(t *testing.T, raw []byte) []byte { + t.Helper() + var tpl map[string]any + if err := json.Unmarshal(raw, &tpl); err != nil { + t.Fatalf("unmarshal template: %v", err) + } + dsl, ok := tpl["dsl"].(map[string]any) + if !ok { + t.Fatalf("template dsl = %T, want map[string]any", tpl["dsl"]) + } + components, ok := dsl["components"].(map[string]any) + if !ok { + t.Fatalf("template components = %T, want map[string]any", dsl["components"]) + } + changed := 0 + for _, rawComp := range components { + comp, ok := rawComp.(map[string]any) + if !ok { + continue + } + obj, ok := comp["obj"].(map[string]any) + if !ok || obj["component_name"] != "Tokenizer" { + continue + } + params, ok := obj["params"].(map[string]any) + if !ok { + continue + } + params["search_method"] = []string{"full_text"} + changed++ + } + if changed == 0 { + t.Fatal("no Tokenizer component found to disable embedding") + } + out, err := json.Marshal(tpl) + if err != nil { + t.Fatalf("marshal modified template: %v", err) + } + return out +} + +func prepareTokenizerResourceForTaskIntegration(t *testing.T) { + t.Helper() + if os.Getenv("RAGFLOW_DICT_PATH") != "" { + return + } + srcDir := filepath.Join(taskRepoRoot(t), ".venv", "lib", "python3.13", "site-packages", "infinity") + if _, err := os.Stat(filepath.Join(srcDir, "huqie.txt")); err != nil { + t.Skipf("tokenizer resource source not found at %s: %v", srcDir, err) + } + if _, err := os.Stat(filepath.Join(srcDir, "huqie.txt.trie")); err != nil { + t.Skipf("tokenizer trie source not found at %s: %v", srcDir, err) + } + root := t.TempDir() + ragDir := filepath.Join(root, "rag") + if err := os.MkdirAll(ragDir, 0o755); err != nil { + t.Fatalf("mkdir rag tokenizer dir: %v", err) + } + taskMustSymlink(t, filepath.Join(srcDir, "huqie.txt"), filepath.Join(ragDir, "huqie.txt")) + taskMustSymlink(t, filepath.Join(srcDir, "huqie.txt.trie"), filepath.Join(ragDir, "huqie.trie")) + taskMustWriteTokenizerPOSDef(t, filepath.Join(srcDir, "huqie.txt"), filepath.Join(ragDir, "pos-id.def")) + taskMustPrepareTokenizerWordNet(t, root) + taskMustPrepareTokenizerOpenCC(t, root) + if err := os.Setenv("RAGFLOW_DICT_PATH", root); err != nil { + t.Fatalf("set RAGFLOW_DICT_PATH: %v", err) + } + t.Cleanup(func() { + _ = os.Unsetenv("RAGFLOW_DICT_PATH") + }) +} + +func taskMustSymlink(t *testing.T, src, dst string) { + t.Helper() + if err := os.Symlink(src, dst); err != nil { + t.Fatalf("symlink tokenizer resource %s -> %s: %v", src, dst, err) + } +} + +func taskMustWriteTokenizerPOSDef(t *testing.T, dictPath, outPath string) { + t.Helper() + data, err := os.ReadFile(dictPath) + if err != nil { + t.Fatalf("read tokenizer dict %s: %v", dictPath, err) + } + posSet := map[string]struct{}{} + for _, line := range strings.Split(string(data), "\n") { + fields := strings.Fields(line) + if len(fields) == 3 { + posSet[fields[2]] = struct{}{} + } + } + if len(posSet) == 0 { + t.Fatalf("no POS tags parsed from tokenizer dict %s", dictPath) + } + posList := make([]string, 0, len(posSet)) + for pos := range posSet { + posList = append(posList, pos) + } + sort.Strings(posList) + content := strings.Join(posList, "\n") + "\n" + if err := os.WriteFile(outPath, []byte(content), 0o644); err != nil { + t.Fatalf("write tokenizer pos file %s: %v", outPath, err) + } +} + +func taskMustPrepareTokenizerWordNet(t *testing.T, root string) { + t.Helper() + zipPath := filepath.Join(taskRepoRoot(t), "ragflow_deps", "nltk_data", "corpora", "wordnet.zip") + reader, err := zip.OpenReader(zipPath) + if err != nil { + t.Skipf("open wordnet zip %s: %v", zipPath, err) + return + } + defer func() { _ = reader.Close() }() + for _, f := range reader.File { + name := strings.TrimPrefix(f.Name, "wordnet/") + if name == "" || strings.HasSuffix(name, "/") { + continue + } + dst := filepath.Join(root, "wordnet", name) + if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil { + t.Fatalf("mkdir wordnet dst %s: %v", dst, err) + } + rc, err := f.Open() + if err != nil { + t.Fatalf("open wordnet entry %s: %v", f.Name, err) + } + out, err := os.Create(dst) + if err != nil { + _ = rc.Close() + t.Fatalf("create wordnet dst %s: %v", dst, err) + } + if _, err := io.Copy(out, rc); err != nil { + _ = out.Close() + _ = rc.Close() + t.Fatalf("copy wordnet entry %s -> %s: %v", f.Name, dst, err) + } + if err := out.Close(); err != nil { + _ = rc.Close() + t.Fatalf("close wordnet dst %s: %v", dst, err) + } + if err := rc.Close(); err != nil { + t.Fatalf("close wordnet entry %s: %v", f.Name, err) + } + } +} + +func taskMustPrepareTokenizerOpenCC(t *testing.T, root string) { + t.Helper() + const systemOpenCC = "/usr/share/opencc" + if _, err := os.Stat(systemOpenCC); err != nil { + t.Skipf("system opencc dir %s not found: %v", systemOpenCC, err) + return + } + taskMustSymlink(t, systemOpenCC, filepath.Join(root, "opencc")) +} + +func requireTokenizerPool(t *testing.T) { + t.Helper() + if tokenizer.IsInitialized() { + return + } + cfg := &tokenizer.PoolConfig{ + DictPath: os.Getenv("RAGFLOW_DICT_PATH"), + MinSize: 1, + MaxSize: 2, + IdleTimeout: 30 * time.Second, + AcquireTimeout: 5 * time.Second, + } + if cfg.DictPath == "" { + cfg.DictPath = "/usr/share/infinity/resource" + } + if err := tokenizer.Init(cfg); err != nil { + t.Skipf("tokenizer pool init failed: %v", err) + } +} + +func mustSeedTaskRealPipelineDocument( + t *testing.T, + db *gorm.DB, + stg storage.Storage, + tenantID, kbID, docID, fileID, bucket, objectPath, docName, content string, +) { + mustSeedTaskRealPipelineDocumentBytes(t, db, stg, tenantID, kbID, docID, fileID, bucket, objectPath, docName, ".txt", "txt", []byte(content)) +} + +func mustSeedTaskRealPipelineDocumentBytes( + t *testing.T, + db *gorm.DB, + stg storage.Storage, + tenantID, kbID, docID, fileID, bucket, objectPath, docName, suffix, docType string, + content []byte, +) { + t.Helper() + if err := db.Create(&entity.Tenant{ + ID: tenantID, + LLMID: "gpt-4", + Status: taskStrPtr("1"), + }).Error; err != nil { + t.Fatalf("create tenant: %v", err) + } + if err := db.Create(&entity.Knowledgebase{ + ID: kbID, + TenantID: tenantID, + EmbdID: "embd-1", + ParserConfig: entity.JSONMap{}, + Status: taskStrPtr("1"), + }).Error; err != nil { + t.Fatalf("create kb: %v", err) + } + if err := stg.Put(bucket, objectPath, content); err != nil { + t.Fatalf("put real minio object: %v", err) + } + if err := db.Create(&entity.File{ + ID: fileID, + ParentID: bucket, + TenantID: tenantID, + CreatedBy: tenantID, + Name: docName, + Type: docType, + Location: taskStrPtr(objectPath), + SourceType: "", + }).Error; err != nil { + t.Fatalf("create file: %v", err) + } + if err := db.Create(&entity.Document{ + ID: docID, + KbID: kbID, + ParserID: "naive", + ParserConfig: entity.JSONMap{}, + SourceType: "local", + Type: docType, + CreatedBy: tenantID, + Name: taskStrPtr(docName), + Location: taskStrPtr(objectPath), + Suffix: suffix, + Status: taskStrPtr("1"), + }).Error; err != nil { + t.Fatalf("create document: %v", err) + } + if err := db.Create(&entity.File2Document{ + ID: taskLimit32("it_map_" + docID), + FileID: taskStrPtr(fileID), + DocumentID: taskStrPtr(docID), + }).Error; err != nil { + t.Fatalf("create file2document: %v", err) + } +} + +func cleanupTaskRealPipelineDocument(db *gorm.DB, stg storage.Storage, tenantID, kbID, docID, fileID, bucket, objectPath string) { + _ = db.Where("document_id = ?", docID).Delete(&entity.File2Document{}).Error + _ = db.Where("id = ?", docID).Delete(&entity.Document{}).Error + _ = db.Where("id = ?", fileID).Delete(&entity.File{}).Error + _ = db.Where("id = ?", kbID).Delete(&entity.Knowledgebase{}).Error + _ = db.Where("id = ?", tenantID).Delete(&entity.Tenant{}).Error + _ = stg.Remove(bucket, objectPath) +} + +func deepCopyTaskChunks(in []map[string]any) []map[string]any { + if len(in) == 0 { + return nil + } + out := make([]map[string]any, len(in)) + for i, ck := range in { + data, err := json.Marshal(ck) + if err != nil { + panic(err) + } + var copied map[string]any + if err := json.Unmarshal(data, &copied); err != nil { + panic(err) + } + out[i] = copied + } + return out +} + +func taskTerminalPayloadFromRunOutput(t *testing.T, out map[string]any, terminalID string) map[string]any { + t.Helper() + if out == nil { + t.Fatal("Run returned nil output") + } + if _, ok := out["output_format"]; ok { + return out + } + nested, ok := out[terminalID].(map[string]any) + if !ok { + t.Fatalf("run output missing terminal payload %q in %v", terminalID, out) + } + return nested +} + +func taskStrPtr(s string) *string { return &s } + +func taskLimit32(s string) string { + if len(s) <= 32 { + return s + } + return s[:32] +} + +func taskS3SafeBucketName(s string) string { + s = strings.ToLower(s) + s = strings.ReplaceAll(s, "_", "-") + return s +} diff --git a/internal/ingestion/task/dataflow_service.go b/internal/ingestion/task/dataflow_service.go new file mode 100644 index 0000000000..a445fd154f --- /dev/null +++ b/internal/ingestion/task/dataflow_service.go @@ -0,0 +1,597 @@ +// +// 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 task + +import ( + "context" + "encoding/json" + "fmt" + componentpkg "ragflow/internal/ingestion/component" + "ragflow/internal/utility" + "regexp" + "sort" + "strings" + "time" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/engine" + "ragflow/internal/entity" + "ragflow/internal/entity/models" + pipelinepkg "ragflow/internal/ingestion/pipeline" + "ragflow/internal/service" +) + +type embedder struct { + model *models.EmbeddingModel +} + +func (e *embedder) Encode(texts []string) ([][]float64, error) { + config := &models.EmbeddingConfig{Dimension: 0} + embeds, err := e.model.ModelDriver.Embed(e.model.ModelName, texts, e.model.APIConfig, config) + if err != nil { + return nil, err + } + vecs := make([][]float64, len(embeds)) + for i, v := range embeds { + vecs[i] = v.Embedding + } + return vecs, nil +} + +type ProgressFunc func(prog float64, msg string) + +type docService interface { + UpdateDocument(id string, req *service.UpdateDocumentRequest) error + GetDocumentMetadataByID(docID string) (map[string]any, error) + SetDocumentMetadata(docID string, meta map[string]any) error +} + +type chunkCounter interface { + IncrementChunkNum(docID, kbID string, chunkNum, tokenConsumption int, duration float64) error +} + +type defaultDocService struct{} +type defaultChunkCounter struct{} + +func (d *defaultDocService) UpdateDocument(id string, req *service.UpdateDocumentRequest) error { + return service.NewDocumentService().UpdateDocument(id, req) +} + +func (d *defaultDocService) GetDocumentMetadataByID(docID string) (map[string]any, error) { + return service.NewDocumentService().GetDocumentMetadataByID(docID) +} + +func (d *defaultDocService) SetDocumentMetadata(docID string, meta map[string]any) error { + return service.NewDocumentService().SetDocumentMetadata(docID, meta) +} + +func (d *defaultChunkCounter) IncrementChunkNum(docID, kbID string, chunkNum, tokenConsumption int, duration float64) error { + return service.NewDocumentService().IncrementChunkNum(docID, kbID, chunkNum, tokenConsumption, duration) +} + +func encodeTexts(model *models.EmbeddingModel, texts []string) ([][]float64, int, error) { + texts = TruncateTexts(texts, model.MaxTokens) + config := &models.EmbeddingConfig{Dimension: 0} + embeds, err := model.ModelDriver.Embed(model.ModelName, texts, model.APIConfig, config) + if err != nil { + return nil, 0, err + } + vecs := make([][]float64, len(embeds)) + totalTokens := 0 + for i, v := range embeds { + vecs[i] = v.Embedding + totalTokens += v.TokenCount + } + return vecs, totalTokens, nil +} + +type PipelineExecutor struct { + taskCtx *TaskContext + dataflowID string + embeddingBatchSize int + docBulkSize int + progressFunc ProgressFunc + + docSvc docService + chunkCounter chunkCounter + insertChunksFunc func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) + logCreateFunc func(log *entity.PipelineOperationLog) error + getEmbeddingModelFunc func(tenantID, embdID string) (*models.EmbeddingModel, error) + loadDSLFunc func(ctx context.Context, dataflowID string) (string, string, error) + runPipelineFunc func(ctx context.Context, dsl string) (map[string]any, string, error) +} + +func validateDataflowTaskContext(taskCtx *TaskContext) error { + if taskCtx == nil { + return fmt.Errorf("dataflow service: nil task context") + } + if taskCtx.Doc.ID == "" { + return fmt.Errorf("dataflow service: empty document id") + } + if taskCtx.Doc.KbID == "" { + return fmt.Errorf("dataflow service: empty document knowledgebase id") + } + if taskCtx.Doc.Name == nil || *taskCtx.Doc.Name == "" { + return fmt.Errorf("dataflow service: empty document name") + } + if taskCtx.KB.ID == "" { + return fmt.Errorf("dataflow service: empty knowledgebase id") + } + if taskCtx.KB.EmbdID == "" { + return fmt.Errorf("dataflow service: empty embedding model id") + } + if taskCtx.Tenant.ID == "" { + return fmt.Errorf("dataflow service: empty tenant id") + } + return nil +} + +func NewDataflowService( + taskCtx *TaskContext, + dataflowID string, + embeddingBatchSize int, + docBulkSize int, +) (*PipelineExecutor, error) { + if err := validateDataflowTaskContext(taskCtx); err != nil { + return nil, err + } + if strings.TrimSpace(dataflowID) == "" { + return nil, fmt.Errorf("dataflow service: empty dataflow id") + } + progressFn := func(prog float64, msg string) {} + if taskCtx != nil && taskCtx.ProgressFunc != nil { + progressFn = taskCtx.ProgressFunc + } + svc := &PipelineExecutor{ + taskCtx: taskCtx, + dataflowID: dataflowID, + embeddingBatchSize: embeddingBatchSize, + docBulkSize: docBulkSize, + progressFunc: progressFn, + docSvc: &defaultDocService{}, + chunkCounter: &defaultChunkCounter{}, + insertChunksFunc: func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error) { + return engine.Get().InsertChunks(ctx, chunks, baseName, datasetID) + }, + logCreateFunc: dao.NewPipelineOperationLogDAO().Create, + getEmbeddingModelFunc: service.NewModelProviderService().GetEmbeddingModel, + } + svc.loadDSLFunc = svc.defaultLoadDSL + svc.runPipelineFunc = svc.defaultRunPipeline + return svc, nil +} + +func (s *PipelineExecutor) WithProgressFunc(fn ProgressFunc) *PipelineExecutor { + s.progressFunc = fn + return s +} + +func (s *PipelineExecutor) WithInsertChunksFunc(f func(ctx context.Context, chunks []map[string]any, baseName string, datasetID string) ([]string, error)) *PipelineExecutor { + s.insertChunksFunc = f + return s +} + +func (s *PipelineExecutor) WithLogCreateFunc(f func(log *entity.PipelineOperationLog) error) *PipelineExecutor { + s.logCreateFunc = f + return s +} + +func (s *PipelineExecutor) WithGetEmbeddingModelFunc(f func(tenantID, embdID string) (*models.EmbeddingModel, error)) *PipelineExecutor { + s.getEmbeddingModelFunc = f + return s +} + +func (s *PipelineExecutor) WithDocService(d docService) *PipelineExecutor { + s.docSvc = d + return s +} + +func (s *PipelineExecutor) WithChunkCounter(c chunkCounter) *PipelineExecutor { + s.chunkCounter = c + return s +} + +func (s *PipelineExecutor) WithLoadDSLFunc(f func(ctx context.Context, dataflowID string) (string, string, error)) *PipelineExecutor { + s.loadDSLFunc = f + return s +} + +func (s *PipelineExecutor) WithRunPipelineFunc(f func(ctx context.Context, dsl string) (map[string]any, string, error)) *PipelineExecutor { + s.runPipelineFunc = f + return s +} + +func (s *PipelineExecutor) KB() *entity.Knowledgebase { return &s.taskCtx.KB } +func (s *PipelineExecutor) Doc() *entity.Document { return &s.taskCtx.Doc } +func (s *PipelineExecutor) Tenant() *entity.Tenant { return &s.taskCtx.Tenant } + +func (s *PipelineExecutor) Run(ctx context.Context) error { + if err := ctx.Err(); err != nil { + return err + } + + dsl, correctedID, err := s.loadDSLFunc(ctx, s.dataflowID) + if err != nil { + return err + } + if correctedID != "" { + s.dataflowID = correctedID + } + + pipelineOutput, pipelineDSL, err := s.runPipelineFunc(ctx, dsl) + if err != nil { + return err + } + + if s.taskCtx.Doc.ID == CANVAS_DEBUG_DOC_ID { + s.recordPipelineLog(s.taskCtx.Doc.ID, pipelineDSL, "done") + return nil + } + + if err := s.RunDataflow(ctx, pipelineOutput); err != nil { + return err + } + + if pipelineDSL != "" { + s.recordPipelineLog(s.taskCtx.Doc.ID, pipelineDSL, "done") + } + return nil +} + +func (s *PipelineExecutor) RunDataflow(ctx context.Context, pipelineOutput map[string]any) error { + taskStart := time.Now() + if pipelineOutput == nil { + return nil + } + if err := ctx.Err(); err != nil { + return err + } + + chunks := s.normalizeChunks(pipelineOutput) + if chunks == nil { + return nil + } + + embeddingTokenConsumption := GetEmbeddingTokenConsumption(pipelineOutput) + + metadata := s.processChunks(chunks) + if err := s.prepareChunkAssets(chunks); err != nil { + return err + } + + if len(metadata) > 0 { + if err := s.updateDocumentMetadata(s.taskCtx.Doc.ID, metadata); err != nil { + common.Warn(fmt.Sprintf("failed to update document metadata: %v", err)) + } + } + + indexStart := time.Now() + s.progress(0.82, "[DOC Engine]:\nStart to index...") + if err := s.insertChunks(ctx, chunks); err != nil { + return err + } + + if err := s.incrementChunkNum(s.taskCtx.Doc.ID, s.taskCtx.Doc.KbID, len(chunks), embeddingTokenConsumption, 0); err != nil { + common.Warn(fmt.Sprintf("failed to increment chunk num: %v", err)) + } + indexDuration := time.Since(indexStart).Seconds() + taskDuration := time.Since(taskStart).Seconds() + s.progress(1.0, fmt.Sprintf("Indexing done (%.2fs). Task done (%.2fs)", indexDuration, taskDuration)) + + return nil +} + +func (s *PipelineExecutor) normalizeChunks(output map[string]any) []map[string]any { + return NormalizeChunks(output) +} + +func (s *PipelineExecutor) embedChunks(ctx context.Context, chunks []map[string]any, tokenConsumption int) ([]map[string]any, int, error) { + if len(chunks) == 0 { + return nil, 0, nil + } + s.progress(0.82, "\n-------------------------------------\nStart to embedding...") + model, err := s.getEmbeddingModel(s.taskCtx.Tenant.ID, s.taskCtx.KB.EmbdID) + if err != nil { + s.progress(-1, fmt.Sprintf("[ERROR]: %v", err)) + return nil, tokenConsumption, err + } + texts := PrepareTextsForDataflowEmbedding(chunks) + batchSize := s.embeddingBatchSize + if batchSize <= 0 { + batchSize = 16 + } + delta := 0.20 / float64(len(texts)/batchSize+1) + prog := 0.8 + var allVects [][]float64 + for i := 0; i < len(texts); i += batchSize { + end := i + batchSize + if end > len(texts) { + end = len(texts) + } + batch := texts[i:end] + if lim := s.taskCtx.EmbedLimiter; lim != nil { + if err := lim.Acquire(ctx, 1); err != nil { + s.progress(-1, fmt.Sprintf("[ERROR]: %v", err)) + return nil, tokenConsumption, err + } + } + vecs, tc, err := encodeTexts(model, batch) + if err != nil { + if lim := s.taskCtx.EmbedLimiter; lim != nil { + lim.Release(1) + } + s.progress(-1, fmt.Sprintf("[ERROR]: %v", err)) + return nil, tokenConsumption, err + } + if lim := s.taskCtx.EmbedLimiter; lim != nil { + lim.Release(1) + } + allVects = append(allVects, vecs...) + tokenConsumption += tc + prog += delta + s.progress(prog, fmt.Sprintf("%d / %d", i+1, len(texts)/batchSize)) + } + if len(allVects) != len(chunks) { + panic(fmt.Sprintf("vector count mismatch: %d vs %d", len(allVects), len(chunks))) + } + AttachVectors(chunks, allVects) + return chunks, tokenConsumption, nil +} + +func (s *PipelineExecutor) processChunks(chunks []map[string]any) map[string]any { + return ProcessChunksForDataflow( + chunks, + s.taskCtx.Doc.ID, + s.taskCtx.Doc.KbID, + *s.taskCtx.Doc.Name, + time.Now(), + ) +} + +func (s *PipelineExecutor) prepareChunkAssets(chunks []map[string]any) error { + return PrepareDataflowChunkAssets(chunks) +} + +func (s *PipelineExecutor) insertChunks(ctx context.Context, chunks []map[string]any) error { + baseName := fmt.Sprintf("ragflow_%s", s.taskCtx.Tenant.ID) + if len(chunks) == 0 { + _, err := s.insertChunksFunc(ctx, chunks, baseName, s.taskCtx.Doc.KbID) + return err + } + bulkSize := s.docBulkSize + if bulkSize <= 0 { + bulkSize = len(chunks) + } + for b := 0; b < len(chunks); b += bulkSize { + end := b + bulkSize + if end > len(chunks) { + end = len(chunks) + } + if _, err := s.insertChunksFunc(ctx, chunks[b:end], baseName, s.taskCtx.Doc.KbID); err != nil { + return err + } + if (b/bulkSize)%128 == 0 { + s.progress(0.8+0.1*float64(b+1)/float64(len(chunks)), "") + } + } + return nil +} + +func (s *PipelineExecutor) updateDocumentMetadata(docID string, metadata map[string]any) error { + if len(metadata) == 0 { + return nil + } + existing, err := s.docSvc.GetDocumentMetadataByID(docID) + if err != nil { + existing = make(map[string]any) + } + for k, v := range metadata { + if _, exists := existing[k]; !exists { + existing[k] = v + } + } + return s.docSvc.SetDocumentMetadata(docID, existing) +} + +func (s *PipelineExecutor) recordPipelineLog(docID, dsl, status string) { + var dslMap entity.JSONMap + if err := json.Unmarshal([]byte(dsl), &dslMap); err != nil { + dslMap = entity.JSONMap{"raw": dsl} + } + log := &entity.PipelineOperationLog{ + ID: utility.GenerateUUID(), + TenantID: s.Tenant().ID, + KbID: s.KB().ID, + DocumentID: docID, + PipelineID: &s.dataflowID, + TaskType: string(entity.PipelineTaskTypeParse), + DSL: dslMap, + ParserID: s.taskCtx.Doc.ParserID, + DocumentName: *s.Doc().Name, + DocumentSuffix: s.taskCtx.Doc.Suffix, + DocumentType: s.taskCtx.Doc.Type, + SourceFrom: s.taskCtx.Doc.SourceType, + OperationStatus: status, + } + if err := s.logCreateFunc(log); err != nil { + common.Warn(fmt.Sprintf("failed to record pipeline log: %v", err)) + } +} + +func (s *PipelineExecutor) incrementChunkNum(docID, kbID string, chunkNum, tokenConsumption int, duration float64) error { + if s.chunkCounter == nil { + return fmt.Errorf("dataflow service: chunk counter is nil") + } + return s.chunkCounter.IncrementChunkNum(docID, kbID, chunkNum, tokenConsumption, duration) +} + +func (s *PipelineExecutor) progress(prog float64, msg string) { + if s.progressFunc != nil { + s.progressFunc(prog, msg) + } +} + +func (s *PipelineExecutor) getEmbeddingModel(tenantID, embdID string) (*models.EmbeddingModel, error) { + return s.getEmbeddingModelFunc(tenantID, embdID) +} + +func hasVectors(chunks []map[string]any) bool { + for _, ck := range chunks { + for k := range ck { + if matchQVec.MatchString(k) { + return true + } + } + } + return false +} + +var matchQVec = regexp.MustCompile(`^q_\d+_vec$`) + +func (s *PipelineExecutor) defaultLoadDSL(ctx context.Context, dataflowID string) (string, string, error) { + if s == nil || s.taskCtx == nil { + return "", "", fmt.Errorf("dataflow service: nil task context") + } + if dataflowID == "" { + return "", "", fmt.Errorf("dataflow service: empty dataflow id") + } + if strings.HasPrefix(s.taskCtx.TaskType, "dataflow") { + canvas, err := dao.NewUserCanvasDAO().GetByID(dataflowID) + if err != nil { + return "", "", fmt.Errorf("load dataflow canvas %s: %w", dataflowID, err) + } + raw, err := json.Marshal(canvas.DSL) + if err != nil { + return "", "", fmt.Errorf("marshal canvas dsl %s: %w", dataflowID, err) + } + return string(raw), dataflowID, nil + } + var pipelineLog entity.PipelineOperationLog + if err := dao.DB.Where("id = ?", dataflowID).First(&pipelineLog).Error; err != nil { + return "", "", fmt.Errorf("load pipeline log %s: %w", dataflowID, err) + } + raw, err := json.Marshal(pipelineLog.DSL) + if err != nil { + return "", "", fmt.Errorf("marshal pipeline log dsl %s: %w", dataflowID, err) + } + correctedID := dataflowID + if pipelineLog.PipelineID != nil && *pipelineLog.PipelineID != "" { + correctedID = *pipelineLog.PipelineID + } + return string(raw), correctedID, nil +} + +func (s *PipelineExecutor) defaultRunPipeline(ctx context.Context, dsl string) (map[string]any, string, error) { + if s == nil || s.taskCtx == nil { + return nil, dsl, fmt.Errorf("dataflow service: nil task context") + } + + prevEncode := componentpkg.EncodeFunc + componentpkg.EncodeFunc = func(tenantID, embdID string) componentpkg.Embedder { + model, err := s.getEmbeddingModelFunc(tenantID, embdID) + if err != nil { + return nil + } + return &embedder{model: model} + } + defer func() { componentpkg.EncodeFunc = prevEncode }() + + // Use doc ID as pipeline ID if available, otherwise a placeholder + pipelineID := "pipeline_" + s.taskCtx.Doc.ID + if s.taskCtx.IngestionTask != nil && s.taskCtx.IngestionTask.ID != "" { + pipelineID = s.taskCtx.IngestionTask.ID + } + pipe, err := pipelinepkg.NewPipelineFromDSL([]byte(dsl), pipelineID) + if err != nil { + return nil, dsl, fmt.Errorf("compile pipeline dsl: %w", err) + } + inputs := map[string]any{} + if s.taskCtx.Doc.ID != "" { + inputs["doc_id"] = s.taskCtx.Doc.ID + } + if s.taskCtx.File != nil { + inputs["file"] = s.taskCtx.File + } + inputs["tenant_id"] = s.taskCtx.Tenant.ID + inputs["model_id"] = s.taskCtx.KB.EmbdID + + output, err := pipe.Run(ctx, inputs) + if err != nil { + return nil, dsl, err + } + payload, err := extractDataflowPipelinePayload(dsl, output) + if err != nil { + return nil, dsl, err + } + return payload, dsl, nil +} + +func extractDataflowPipelinePayload(dsl string, out map[string]any) (map[string]any, error) { + if out == nil { + return nil, nil + } + if _, ok := out["output_format"]; ok { + return out, nil + } + terminalIDs, err := terminalComponentIDsFromDSL([]byte(dsl)) + if err != nil { + return nil, err + } + if len(terminalIDs) != 1 { + return nil, fmt.Errorf("dataflow pipeline requires exactly 1 terminal, got %d: %v", len(terminalIDs), terminalIDs) + } + payload, ok := out[terminalIDs[0]].(map[string]any) + if !ok { + return nil, fmt.Errorf("run output missing terminal payload %q", terminalIDs[0]) + } + return payload, nil +} + +func terminalComponentIDsFromDSL(raw []byte) ([]string, error) { + var tpl map[string]any + if err := json.Unmarshal(raw, &tpl); err != nil { + return nil, fmt.Errorf("unmarshal dataflow dsl: %w", err) + } + root := tpl + if nested, ok := tpl["dsl"].(map[string]any); ok { + root = nested + } + components, ok := root["components"].(map[string]any) + if !ok { + return nil, fmt.Errorf("dataflow dsl missing components map") + } + terminals := make([]string, 0, len(components)) + for id, rawComp := range components { + comp, ok := rawComp.(map[string]any) + if !ok { + return nil, fmt.Errorf("component %q has invalid type %T", id, rawComp) + } + switch downstream := comp["downstream"].(type) { + case nil: + terminals = append(terminals, id) + case []any: + if len(downstream) == 0 { + terminals = append(terminals, id) + } + default: + // Non-slice downstream means the component is connected; ignore it here. + } + } + sort.Strings(terminals) + return terminals, nil +} diff --git a/internal/ingestion/task/dataflow_service_test.go b/internal/ingestion/task/dataflow_service_test.go new file mode 100644 index 0000000000..ffa2954ee6 --- /dev/null +++ b/internal/ingestion/task/dataflow_service_test.go @@ -0,0 +1,826 @@ +package task + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" + + "github.com/glebarez/sqlite" + "gorm.io/gorm" + "gorm.io/gorm/logger" + + "ragflow/internal/dao" + "ragflow/internal/entity" + "ragflow/internal/entity/models" + "ragflow/internal/service" +) + +// ============================================================================= +// Test helpers +// ============================================================================= + +func strPtr(s string) *string { return &s } + +func makeTaskCtx() *TaskContext { + return &TaskContext{ + IngestionTask: &entity.IngestionTask{ + ID: "task-1", + DocumentID: "doc-1", + }, + TaskType: "dataflow", + Doc: entity.Document{ + ID: "doc-1", + KbID: "kb-1", + Name: strPtr("test-doc.pdf"), + Suffix: ".pdf", + Type: "pdf", + }, + KB: entity.Knowledgebase{ + ID: "kb-1", + TenantID: "tenant-1", + EmbdID: "embd-1", + }, + Tenant: entity.Tenant{ + ID: "tenant-1", + }, + ProgressFunc: func(prog float64, msg string) {}, + } +} + +func setupDataflowServiceTestDB(t *testing.T) func() { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate(&entity.UserCanvas{}, &entity.PipelineOperationLog{}); err != nil { + t.Fatalf("auto-migrate sqlite: %v", err) + } + origDB := dao.DB + dao.DB = db + return func() { dao.DB = origDB } +} + +func mustNewDataflowService(t *testing.T, taskCtx *TaskContext, dataflowID string, embeddingBatchSize int, docBulkSize int) *PipelineExecutor { + t.Helper() + svc, err := NewDataflowService(taskCtx, dataflowID, embeddingBatchSize, docBulkSize) + if err != nil { + t.Fatalf("NewDataflowService: %v", err) + } + return svc +} + +func TestDataflowService_DefaultLoadDSL_UsesUserCanvas(t *testing.T) { + cleanup := setupDataflowServiceTestDB(t) + defer cleanup() + + dslMap := entity.JSONMap{"dsl": map[string]any{"graph": map[string]any{"nodes": []any{}, "edges": []any{}}}} + if err := dao.NewUserCanvasDAO().Create(&entity.UserCanvas{ID: "canvas-1", UserID: "u1", Permission: "me", CanvasCategory: "agent_canvas", DSL: dslMap}); err != nil { + t.Fatalf("create user canvas: %v", err) + } + + ctx := makeTaskCtx() + svc := mustNewDataflowService(t, ctx, "canvas-1", 0, 0) + gotDSL, correctedID, err := svc.loadDSLFunc(context.Background(), "canvas-1") + if err != nil { + t.Fatalf("loadDSLFunc: %v", err) + } + if correctedID != "canvas-1" { + t.Fatalf("correctedID = %q, want canvas-1", correctedID) + } + var decoded map[string]any + if err := json.Unmarshal([]byte(gotDSL), &decoded); err != nil { + t.Fatalf("unmarshal dsl: %v", err) + } + if _, ok := decoded["dsl"].(map[string]any); !ok { + t.Fatalf("decoded dsl = %v, want top-level dsl map", decoded) + } +} + +func TestDataflowService_DefaultLoadDSL_UsesPipelineOperationLog(t *testing.T) { + cleanup := setupDataflowServiceTestDB(t) + defer cleanup() + + pipelineID := "canvas-2" + if err := dao.NewPipelineOperationLogDAO().Create(&entity.PipelineOperationLog{ + ID: "log-1", + DocumentID: "doc-1", + TenantID: "tenant-1", + KbID: "kb-1", + PipelineID: &pipelineID, + ParserID: "naive", + DocumentName: "sample.pdf", + DocumentSuffix: ".pdf", + DocumentType: "pdf", + SourceFrom: "local", + TaskType: "parse", + OperationStatus: "done", + DSL: entity.JSONMap{"dsl": map[string]any{"graph": map[string]any{"nodes": []any{}, "edges": []any{}}}}, + }); err != nil { + t.Fatalf("create pipeline log: %v", err) + } + + ctx := makeTaskCtx() + ctx.TaskType = "replay" // Not starting with "dataflow" to use PipelineOperationLog + svc := mustNewDataflowService(t, ctx, "log-1", 0, 0) + gotDSL, correctedID, err := svc.loadDSLFunc(context.Background(), "log-1") + if err != nil { + t.Fatalf("loadDSLFunc: %v", err) + } + if correctedID != pipelineID { + t.Fatalf("correctedID = %q, want %q", correctedID, pipelineID) + } + var decoded map[string]any + if err := json.Unmarshal([]byte(gotDSL), &decoded); err != nil { + t.Fatalf("unmarshal dsl: %v", err) + } + if _, ok := decoded["dsl"].(map[string]any); !ok { + t.Fatalf("decoded dsl = %v, want top-level dsl map", decoded) + } +} + +// ============================================================================= +// NewDataflowService — constructor +// ============================================================================= + +func TestNewDataflowService_Basic(t *testing.T) { + svc, err := NewDataflowService(makeTaskCtx(), "flow-1", 0, 0) + if err != nil { + t.Fatalf("NewDataflowService: %v", err) + } + if svc == nil { + t.Fatal("NewDataflowService returned nil") + } + if svc.taskCtx == nil { + t.Error("taskCtx should not be nil") + } + if svc.docSvc == nil || svc.chunkCounter == nil || svc.insertChunksFunc == nil || svc.logCreateFunc == nil || svc.getEmbeddingModelFunc == nil || svc.loadDSLFunc == nil || svc.runPipelineFunc == nil { + t.Fatal("expected production dependencies to be fully initialized") + } +} + +func TestNewDataflowService_RejectsNilTaskContext(t *testing.T) { + _, err := NewDataflowService(nil, "flow-1", 0, 0) + if err == nil { + t.Fatal("expected error for nil task context") + } +} + +func TestNewDataflowService_RejectsEmptyDataflowID(t *testing.T) { + _, err := NewDataflowService(makeTaskCtx(), "", 0, 0) + if err == nil { + t.Fatal("expected error for empty dataflow id") + } +} + +func TestNewDataflowService_RejectsIncompleteTaskContext(t *testing.T) { + tests := []struct { + name string + mutate func(*TaskContext) + }{ + {name: "missing doc id", mutate: func(ctx *TaskContext) { ctx.Doc.ID = "" }}, + {name: "missing kb id", mutate: func(ctx *TaskContext) { ctx.Doc.KbID = "" }}, + {name: "missing doc name", mutate: func(ctx *TaskContext) { ctx.Doc.Name = nil }}, + {name: "missing knowledgebase id", mutate: func(ctx *TaskContext) { ctx.KB.ID = "" }}, + {name: "missing embedding model id", mutate: func(ctx *TaskContext) { ctx.KB.EmbdID = "" }}, + {name: "missing tenant id", mutate: func(ctx *TaskContext) { ctx.Tenant.ID = "" }}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := makeTaskCtx() + tt.mutate(ctx) + _, err := NewDataflowService(ctx, "flow-1", 0, 0) + if err == nil { + t.Fatal("expected validation error") + } + }) + } +} + +func TestNewDataflowService_CustomBatchSizes(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 64, 128) + if svc.embeddingBatchSize != 64 { + t.Errorf("embeddingBatchSize = %d, want 64", svc.embeddingBatchSize) + } + if svc.docBulkSize != 128 { + t.Errorf("docBulkSize = %d, want 128", svc.docBulkSize) + } +} + +func TestNewDataflowService_WithProgressFunc(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithProgressFunc(func(prog float64, msg string) {}) + if svc.progressFunc == nil { + t.Error("progressFunc should be set via WithProgressFunc") + } +} + +func TestNewDataflowService_DataflowID(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "my-flow-id", 0, 0) + if svc.dataflowID != "my-flow-id" { + t.Errorf("dataflowID = %q, want %q", svc.dataflowID, "my-flow-id") + } +} + +func TestKB_Doc_Tenant_Accessors(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) + if svc.KB().ID != "kb-1" { + t.Errorf("KB().ID = %q, want \"kb-1\"", svc.KB().ID) + } + if svc.Doc().ID != "doc-1" { + t.Errorf("Doc().ID = %q, want \"doc-1\"", svc.Doc().ID) + } + if svc.Tenant().ID != "tenant-1" { + t.Errorf("Tenant().ID = %q, want \"tenant-1\"", svc.Tenant().ID) + } +} + +// ============================================================================= +// processChunks +// ============================================================================= + +func TestDataflowService_ProcessChunks_WrapsProcessChunksForDataflow(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) + chunks := []map[string]any{{"text": "hello world"}} + meta := svc.processChunks(chunks) + + // Verify the wrapper method works correctly and chunks are processed + if chunks[0]["doc_id"] != "doc-1" { + t.Errorf("doc_id = %q, want \"doc-1\"", chunks[0]["doc_id"]) + } + if meta != nil { + // No need to verify the detailed content of meta as ProcessChunksForDataflow already has comprehensive tests + } +} + +// ============================================================================= +// embedChunks — Python: _embed_chunks (line 247) +// ============================================================================= + +func TestEmbedChunks_Success(t *testing.T) { + stub := &stubDriver{} + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithGetEmbeddingModelFunc( + func(tenantID, embdID string) (*models.EmbeddingModel, error) { + return makeTestEmbeddingModel(stub, 100), nil + }, + ) + chunks := []map[string]any{ + {"text": "hello"}, + {"text": "world"}, + } + result, tc, err := svc.embedChunks(context.Background(), chunks, 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(result) != 2 { + t.Fatalf("len = %d, want 2", len(result)) + } + // vectors should be attached (q_*_vec keys) + if result[0]["q_2_vec"] == nil { + t.Errorf("expected q_2_vec in chunk[0], got keys: %v", chunkKeys(result[0])) + } + // token consumption should include initial + new tokens + if tc < 5 { + t.Errorf("token consumption should be at least the initial 5, got %d", tc) + } +} + +func TestEmbedChunks_EmptyChunks(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) + result, _, err := svc.embedChunks(context.Background(), nil, 0) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result != nil { + t.Errorf("expected nil for empty chunks, got %v", result) + } +} + +func TestEmbedChunks_ModelError(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithGetEmbeddingModelFunc( + func(tenantID, embdID string) (*models.EmbeddingModel, error) { + return nil, errors.New("model not found") + }, + ) + chunks := []map[string]any{{"text": "hello"}} + result, tc, err := svc.embedChunks(context.Background(), chunks, 10) + if err == nil { + t.Fatal("expected error on model error") + } + if !strings.Contains(err.Error(), "model not found") { + t.Errorf("expected error containing 'model not found', got %v", err) + } + if result != nil { + t.Errorf("expected nil result on model error, got %v", result) + } + if tc != 10 { + t.Errorf("token consumption should be preserved on error, got %d", tc) + } +} + +// ============================================================================= +// progress +// ============================================================================= + +func TestProgress_CallsCallback(t *testing.T) { + var called bool + var lastProg float64 + var lastMsg string + + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) + svc.WithProgressFunc(func(prog float64, msg string) { + called = true + lastProg = prog + lastMsg = msg + }) + svc.progress(0.82, "Start to embedding...") + + if !called { + t.Error("progress callback was not called") + } + if lastProg != 0.82 { + t.Errorf("prog = %f, want 0.82", lastProg) + } + if lastMsg != "Start to embedding..." { + t.Errorf("msg = %q, want \"Start to embedding...\"", lastMsg) + } +} + +func TestProgress_NilCallbackDoesNotPanic(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) + svc.progress(0.5, "test") +} + +// ============================================================================= +// insertChunks +// ============================================================================= + +func TestInsertChunks_EmptyChunks(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithInsertChunksFunc( + func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + return nil, nil + }, + ) + err := svc.insertChunks(context.Background(), nil) + if err != nil { + t.Errorf("expected no error for nil chunks, got %v", err) + } +} + +func TestInsertChunks_BaseNameAndDatasetID(t *testing.T) { + var capturedBaseName, capturedDatasetID string + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithInsertChunksFunc( + func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + capturedBaseName = baseName + capturedDatasetID = datasetID + return nil, nil + }, + ) + chunks := []map[string]any{{"text": "hello"}} + err := svc.insertChunks(context.Background(), chunks) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if capturedBaseName != "ragflow_tenant-1" { + t.Errorf("baseName = %q, want \"ragflow_tenant-1\"", capturedBaseName) + } + if capturedDatasetID != "kb-1" { + t.Errorf("datasetID = %q, want \"kb-1\"", capturedDatasetID) + } +} + +// ============================================================================= +// updateDocumentMetadata +// ============================================================================= + +func TestUpdateDocumentMetadata_EmptyMetadata(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithDocService(&stubDocService{}) + err := svc.updateDocumentMetadata("doc-1", nil) + if err != nil { + t.Errorf("empty metadata should not error, got %v", err) + } +} + +func TestUpdateDocumentMetadata_MergesNewKeys(t *testing.T) { + mds := &stubDocService{metaData: map[string]any{"existing": "old"}} + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithDocService(mds) + err := svc.updateDocumentMetadata("doc-1", map[string]any{"new_key": "value"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if mds.metaData["existing"] != "old" { + t.Errorf("existing key should be preserved: got %q", mds.metaData["existing"]) + } + if mds.metaData["new_key"] != "value" { + t.Errorf("new_key = %q, want \"value\"", mds.metaData["new_key"]) + } +} + +func TestUpdateDocumentMetadata_PreservesExistingKey(t *testing.T) { + mds := &stubDocService{metaData: map[string]any{"author": "Alice"}} + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithDocService(mds) + err := svc.updateDocumentMetadata("doc-1", map[string]any{"author": "Bob"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if mds.metaData["author"] != "Alice" { + t.Errorf("existing key should NOT be overwritten: got %q", mds.metaData["author"]) + } +} + +// ============================================================================= +// recordPipelineLog +// ============================================================================= + +func TestRecordPipelineLog(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithLogCreateFunc( + func(log *entity.PipelineOperationLog) error { return nil }, + ) + svc.recordPipelineLog("doc-1", `{"components": {}}`, "done") +} + +// ============================================================================= +// updateDocumentMetadata +// ============================================================================= + +// incrementChunkNum +// ============================================================================= + +func TestIncrementChunkNum(t *testing.T) { + counter := &stubChunkCounter{} + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0). + WithChunkCounter(counter) + err := svc.incrementChunkNum("doc-1", "kb-1", 10, 100, 0) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if counter.lastDocID != "doc-1" { + t.Errorf("docID = %q, want \"doc-1\"", counter.lastDocID) + } + if counter.lastKbID != "kb-1" { + t.Errorf("kbID = %q, want \"kb-1\"", counter.lastKbID) + } + if counter.lastChunkNum != 10 { + t.Errorf("chunkNum = %d, want 10", counter.lastChunkNum) + } + if counter.lastTokenNum != 100 { + t.Errorf("tokenNum = %d, want 100", counter.lastTokenNum) + } +} + +func TestIncrementChunkNum_ProcessDuration(t *testing.T) { + counter := &stubChunkCounter{} + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0). + WithChunkCounter(counter) + err := svc.incrementChunkNum("doc-1", "kb-1", 5, 50, 12.5) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if counter.lastDuration != 12.5 { + t.Errorf("duration = %f, want 12.5", counter.lastDuration) + } + if counter.lastChunkNum != 5 { + t.Errorf("chunkNum = %d, want 5", counter.lastChunkNum) + } + if counter.lastTokenNum != 50 { + t.Errorf("tokenNum = %d, want 50", counter.lastTokenNum) + } +} + +// ============================================================================= +// RunDataflow — Python: run_dataflow (line 94) +// ============================================================================= + +func TestRunDataflow_NilOutput(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) + err := svc.RunDataflow(context.Background(), nil) + if err != nil { + t.Errorf("expected nil error for nil output, got %v", err) + } +} + +func TestRunDataflow_EmptyOutput(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithLogCreateFunc( + func(log *entity.PipelineOperationLog) error { return nil }, + ) + err := svc.RunDataflow(context.Background(), map[string]any{}) + if err != nil { + t.Errorf("expected nil error for empty output, got %v", err) + } +} + +func TestRunDataflow_NormalizedEmpty(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0).WithLogCreateFunc( + func(log *entity.PipelineOperationLog) error { return nil }, + ) + err := svc.RunDataflow(context.Background(), map[string]any{"markdown": ""}) + if err != nil { + t.Errorf("expected nil error for empty normalized output, got %v", err) + } +} + +func TestRunDataflow_FullFlow(t *testing.T) { + stub := &stubDriver{} + var progressCalls []float64 + var progressMsgs []string + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0). + WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) { + return makeTestEmbeddingModel(stub, 100), nil + }). + WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + return nil, nil + }). + WithDocService(&stubDocService{}). + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }). + WithChunkCounter(&stubChunkCounter{}). + WithProgressFunc(func(prog float64, msg string) { + progressCalls = append(progressCalls, prog) + progressMsgs = append(progressMsgs, msg) + }) + + output := map[string]any{ + "chunks": []map[string]any{ + {"text": "hello"}, + {"text": "world"}, + }, + } + err := svc.RunDataflow(context.Background(), output) + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if len(progressCalls) < 2 { + t.Fatalf("expected multiple progress calls, got %v", progressCalls) + } + if progressCalls[len(progressCalls)-1] != 1.0 { + t.Fatalf("final progress = %v, want 1.0", progressCalls[len(progressCalls)-1]) + } + lastMsg := progressMsgs[len(progressMsgs)-1] + if !strings.Contains(lastMsg, "Indexing done (") || !strings.Contains(lastMsg, "Task done (") { + t.Fatalf("final progress msg = %q, want indexing/task timing message", lastMsg) + } +} + +func TestRunDataflow_AlreadyHasVectors(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0). + WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + return nil, nil + }). + WithDocService(&stubDocService{}). + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }). + WithChunkCounter(&stubChunkCounter{}). + WithProgressFunc(func(prog float64, msg string) {}) + + output := map[string]any{ + "chunks": []map[string]any{ + {"text": "hello", "q_768_vec": []float64{0.1, 0.2}}, + }, + } + err := svc.RunDataflow(context.Background(), output) + if err != nil { + t.Errorf("unexpected error: %v", err) + } +} + +func TestRunDataflow_ContextCanceled(t *testing.T) { + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := svc.RunDataflow(ctx, map[string]any{ + "chunks": []map[string]any{{"text": "hello"}}, + }) + if err == nil { + t.Error("expected context canceled error") + } +} + +func TestExtractDataflowPipelinePayload_UnwrapsSingleTerminalOutput(t *testing.T) { + dsl := `{"dsl":{"components":{"Parser:A":{"downstream":["Tokenizer:B"]},"Tokenizer:B":{"downstream":[]}}}}` + out := map[string]any{ + "Tokenizer:B": map[string]any{ + "output_format": "chunks", + "chunks": []map[string]any{{"text": "hello world"}}, + }, + "state": map[string]any{"Tokenizer:B": map[string]any{"output_format": "chunks"}}, + } + + payload, err := extractDataflowPipelinePayload(dsl, out) + if err != nil { + t.Fatalf("extractDataflowPipelinePayload: %v", err) + } + if got := payload["output_format"]; got != "chunks" { + t.Fatalf("output_format = %v, want chunks", got) + } + chunks, ok := payload["chunks"].([]map[string]any) + if !ok { + t.Fatalf("chunks = %T, want []map[string]any", payload["chunks"]) + } + if len(chunks) != 1 || chunks[0]["text"] != "hello world" { + t.Fatalf("chunks = %v, want single hello world chunk", chunks) + } +} + +func TestExtractDataflowPipelinePayload_ErrorsOnMultipleTerminals(t *testing.T) { + dsl := `{"dsl":{"components":{"Tokenizer:A":{"downstream":[]},"Tokenizer:B":{"downstream":[]}}}}` + _, err := extractDataflowPipelinePayload(dsl, map[string]any{ + "Tokenizer:A": map[string]any{"output_format": "chunks"}, + "Tokenizer:B": map[string]any{"output_format": "chunks"}, + }) + if err == nil { + t.Fatal("expected error for multiple terminals") + } + if !strings.Contains(err.Error(), "exactly 1 terminal") { + t.Fatalf("err = %v, want exactly 1 terminal", err) + } +} + +func TestDataflowService_Run_MainFlowWithStubs(t *testing.T) { + stub := &stubDriver{} + logged := false + inserted := false + var progressCalls []float64 + var progressMsgs []string + + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0). + WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) { + return `{"nodes":[{"id":"n1"}],"edges":[]}`, "flow-corrected", nil + }). + WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { + return map[string]any{ + "chunks": []map[string]any{ + {"text": "hello world"}, + }, + }, dsl, nil + }). + WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) { + return makeTestEmbeddingModel(stub, 100), nil + }). + WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + inserted = true + return nil, nil + }). + WithDocService(&stubDocService{}). + WithChunkCounter(&stubChunkCounter{}). + WithProgressFunc(func(prog float64, msg string) { + progressCalls = append(progressCalls, prog) + progressMsgs = append(progressMsgs, msg) + }). + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { + logged = true + if log.PipelineID == nil || *log.PipelineID != "flow-corrected" { + t.Fatalf("PipelineID = %v, want flow-corrected", log.PipelineID) + } + return nil + }) + + err := svc.Run(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !inserted { + t.Fatal("expected insertChunks to be called") + } + if !logged { + t.Fatal("expected pipeline log to be created") + } + if len(progressCalls) == 0 || progressCalls[len(progressCalls)-1] != 1.0 { + t.Fatalf("expected final progress 1.0, got %v", progressCalls) + } + lastMsg := progressMsgs[len(progressMsgs)-1] + if !strings.Contains(lastMsg, "Indexing done (") || !strings.Contains(lastMsg, "Task done (") { + t.Fatalf("final progress msg = %q, want indexing/task timing message", lastMsg) + } +} + +func TestEmbedChunks_ReportsIntermediateProgress(t *testing.T) { + stub := &stubDriver{} + var progressCalls []float64 + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 1, 0). + WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) { + return makeTestEmbeddingModel(stub, 100), nil + }). + WithProgressFunc(func(prog float64, msg string) { + progressCalls = append(progressCalls, prog) + }) + + chunks := []map[string]any{ + {"text": "hello"}, + {"text": "world"}, + } + _, _, err := svc.embedChunks(context.Background(), chunks, 0) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(progressCalls) < 2 { + t.Fatalf("expected start and intermediate progress calls, got %v", progressCalls) + } +} + +func TestInsertChunks_ReportsBatchProgress(t *testing.T) { + var progressCalls []float64 + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 1). + WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + time.Sleep(5 * time.Millisecond) + return nil, nil + }). + WithProgressFunc(func(prog float64, msg string) { + progressCalls = append(progressCalls, prog) + }) + + chunks := []map[string]any{ + {"text": "a"}, + {"text": "b"}, + } + if err := svc.insertChunks(context.Background(), chunks); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(progressCalls) == 0 { + t.Fatal("expected indexing progress callbacks") + } +} + +func TestEmbedChunks_ErrorReportsNegativeProgressAndReturnsError(t *testing.T) { + var progressCalls []float64 + svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 1, 0). + WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) { + return nil, errors.New("model unavailable") + }). + WithProgressFunc(func(prog float64, msg string) { + progressCalls = append(progressCalls, prog) + }) + + chunks := []map[string]any{{"text": "hello"}} + _, _, err := svc.embedChunks(context.Background(), chunks, 0) + if err == nil { + t.Fatal("expected error") + } + if len(progressCalls) == 0 || progressCalls[len(progressCalls)-1] != -1 { + t.Fatalf("expected final error progress -1, got %v", progressCalls) + } +} + +// ============================================================================= +// Stub implementations for testing +// ============================================================================= + +type stubDocService struct { + err error + metaData map[string]any + lastReq *service.UpdateDocumentRequest + lastDocID string +} + +type stubChunkCounter struct { + err error + lastDocID string + lastKbID string + lastChunkNum int + lastTokenNum int + lastDuration float64 + callCount int +} + +func (s *stubDocService) UpdateDocument(id string, req *service.UpdateDocumentRequest) error { + s.lastDocID = id + s.lastReq = req + return s.err +} + +func (s *stubDocService) GetDocumentMetadataByID(docID string) (map[string]any, error) { + if s.err != nil { + return nil, s.err + } + if s.metaData == nil { + return make(map[string]any), nil + } + return s.metaData, nil +} + +func (s *stubDocService) SetDocumentMetadata(docID string, meta map[string]any) error { + if s.metaData == nil { + s.metaData = make(map[string]any) + } + for k, v := range meta { + s.metaData[k] = v + } + return s.err +} + +func (s *stubChunkCounter) IncrementChunkNum(docID, kbID string, chunkNum, tokenConsumption int, duration float64) error { + s.lastDocID = docID + s.lastKbID = kbID + s.lastChunkNum = chunkNum + s.lastTokenNum = tokenConsumption + s.lastDuration = duration + s.callCount++ + return s.err +} + +// Compile-time checks +var ( + _ docService = (*stubDocService)(nil) + _ chunkCounter = (*stubChunkCounter)(nil) +) diff --git a/internal/ingestion/task/embedding_test.go b/internal/ingestion/task/embedding_test.go new file mode 100644 index 0000000000..25d5c9613d --- /dev/null +++ b/internal/ingestion/task/embedding_test.go @@ -0,0 +1,226 @@ +package task + +import ( + "errors" + "strings" + "testing" + + "ragflow/internal/entity/models" +) + +// stubDriver records the texts passed to Embed for verification. +type stubDriver struct { + capturedTexts []string +} + +func (d *stubDriver) Embed(modelName *string, texts []string, apiConfig *models.APIConfig, embeddingConfig *models.EmbeddingConfig) ([]models.EmbeddingData, error) { + d.capturedTexts = texts + result := make([]models.EmbeddingData, len(texts)) + for i := range texts { + result[i] = models.EmbeddingData{ + Embedding: []float64{float64(i), 0.1}, + Index: i, + TokenCount: len(texts[i]), + } + } + return result, nil +} +func (d *stubDriver) NewInstance(baseURL map[string]string) models.ModelDriver { return d } +func (d *stubDriver) Name() string { return "stub" } +func (d *stubDriver) ChatWithMessages(modelName string, messages []models.Message, apiConfig *models.APIConfig, chatModelConfig *models.ChatConfig) (*models.ChatResponse, error) { + return nil, nil +} +func (d *stubDriver) ChatStreamlyWithSender(modelName string, messages []models.Message, apiConfig *models.APIConfig, modelConfig *models.ChatConfig, sender func(*string, *string) error) error { + return nil +} +func (d *stubDriver) Rerank(modelName *string, query string, documents []string, apiConfig *models.APIConfig, rerankConfig *models.RerankConfig) (*models.RerankResponse, error) { + return nil, nil +} +func (d *stubDriver) TranscribeAudio(modelName *string, file *string, apiConfig *models.APIConfig, asrConfig *models.ASRConfig) (*models.ASRResponse, error) { + return nil, nil +} +func (d *stubDriver) TranscribeAudioWithSender(modelName *string, file *string, apiConfig *models.APIConfig, asrConfig *models.ASRConfig, sender func(*string, *string) error) error { + return nil +} +func (d *stubDriver) AudioSpeech(modelName *string, audioContent *string, apiConfig *models.APIConfig, ttsConfig *models.TTSConfig) (*models.TTSResponse, error) { + return nil, nil +} +func (d *stubDriver) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *models.APIConfig, ttsConfig *models.TTSConfig, sender func(*string, *string) error) error { + return nil +} +func (d *stubDriver) OCRFile(modelName *string, content []byte, url *string, apiConfig *models.APIConfig, ocrConfig *models.OCRConfig) (*models.OCRFileResponse, error) { + return nil, nil +} +func (d *stubDriver) ParseFile(modelName *string, content []byte, url *string, apiConfig *models.APIConfig, parseFileConfig *models.ParseFileConfig) (*models.ParseFileResponse, error) { + return nil, nil +} +func (d *stubDriver) ListModels(apiConfig *models.APIConfig) ([]models.ListModelResponse, error) { + return nil, nil +} +func (d *stubDriver) Balance(apiConfig *models.APIConfig) (map[string]interface{}, error) { + return nil, nil +} +func (d *stubDriver) CheckConnection(apiConfig *models.APIConfig) error { return nil } +func (d *stubDriver) ListTasks(apiConfig *models.APIConfig) ([]models.ListTaskStatus, error) { + return nil, nil +} +func (d *stubDriver) ShowTask(taskID string, apiConfig *models.APIConfig) (*models.TaskResponse, error) { + return nil, nil +} +func (d *stubDriver) ToolCall(name string, arguments map[string]interface{}) (string, error) { + return "", nil +} + +func makeTestEmbeddingModel(stub *stubDriver, maxTokens int) *models.EmbeddingModel { + return &models.EmbeddingModel{ + ModelDriver: stub, + ModelName: strPtr("test-model"), + APIConfig: &models.APIConfig{}, + MaxTokens: maxTokens, + } +} + +// ============================================================================= +// encodeTexts — must truncate texts before passing to ModelDriver.Embed +// Python: truncated = truncate_texts(txts, model.max_length); model.encode(truncated) +// ============================================================================= + +func TestEncodeTexts_TruncatesBeforeEmbed(t *testing.T) { + stub := &stubDriver{} + model := makeTestEmbeddingModel(stub, 12) + // "hello world" is 2 tokens, safeMax = 12-10 = 2, fits + texts := []string{"hello world"} + _, _, err := encodeTexts(model, texts) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(stub.capturedTexts) != 1 { + t.Fatalf("captured %d texts, want 1", len(stub.capturedTexts)) + } + if stub.capturedTexts[0] != "hello world" { + t.Errorf("texts should pass through when within limit, got %q", stub.capturedTexts[0]) + } +} + +func TestEncodeTexts_TruncatesLongText(t *testing.T) { + stub := &stubDriver{} + model := makeTestEmbeddingModel(stub, 12) + // safeMax = 12-10 = 2 tokens. Text > 2 tokens should be truncated. + longText := strings.Repeat("hello world ", 20) // 40 tokens + texts := []string{longText} + _, _, err := encodeTexts(model, texts) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(stub.capturedTexts) != 1 { + t.Fatalf("captured %d texts, want 1", len(stub.capturedTexts)) + } + // The text passed to Embed must be SHORTER than the original + if len(stub.capturedTexts[0]) >= len(longText) { + t.Errorf("text should be truncated before embed: original=%d, truncated=%d", + len(longText), len(stub.capturedTexts[0])) + } +} + +func TestEncodeTexts_TokenCountSum(t *testing.T) { + stub := &stubDriver{} + model := makeTestEmbeddingModel(stub, 100) + texts := []string{"hello", "world"} + _, totalTokens, err := encodeTexts(model, texts) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if totalTokens <= 0 { + t.Errorf("expected positive token count, got %d", totalTokens) + } +} + +func TestEncodeTexts_ReturnsVectors(t *testing.T) { + stub := &stubDriver{} + model := makeTestEmbeddingModel(stub, 100) + texts := []string{"a", "b", "c"} + vecs, _, err := encodeTexts(model, texts) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(vecs) != 3 { + t.Fatalf("expected 3 vectors, got %d", len(vecs)) + } +} + +// ============================================================================= +// GetEmbeddingDimension +// ============================================================================= + +func TestTestEncodeForDim_ReturnsDimension(t *testing.T) { + stub := &stubDriver{} + model := makeTestEmbeddingModel(stub, 100) + dim, err := getEmbeddingDimension(model) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if dim != 2 { + t.Errorf("dim = %d, want 2 (stubDriver returns 2-element vectors)", dim) + } +} + +func TestTestEncodeForDim_ModelError(t *testing.T) { + stub := &stubDriver{} + model := makeTestEmbeddingModel(stub, 100) + model.ModelDriver = &errDriver{} + _, err := getEmbeddingDimension(model) + if err == nil { + t.Error("expected error from failing driver") + } +} + +type errDriver struct{} + +func (d *errDriver) Embed(modelName *string, texts []string, apiConfig *models.APIConfig, embeddingConfig *models.EmbeddingConfig) ([]models.EmbeddingData, error) { + return nil, errors.New("embed error") +} +func (d *errDriver) NewInstance(baseURL map[string]string) models.ModelDriver { return d } +func (d *errDriver) Name() string { return "err" } +func (d *errDriver) ChatWithMessages(modelName string, messages []models.Message, apiConfig *models.APIConfig, chatModelConfig *models.ChatConfig) (*models.ChatResponse, error) { + return nil, nil +} +func (d *errDriver) ChatStreamlyWithSender(modelName string, messages []models.Message, apiConfig *models.APIConfig, modelConfig *models.ChatConfig, sender func(*string, *string) error) error { + return nil +} +func (d *errDriver) Rerank(modelName *string, query string, documents []string, apiConfig *models.APIConfig, rerankConfig *models.RerankConfig) (*models.RerankResponse, error) { + return nil, nil +} +func (d *errDriver) TranscribeAudio(modelName *string, file *string, apiConfig *models.APIConfig, asrConfig *models.ASRConfig) (*models.ASRResponse, error) { + return nil, nil +} +func (d *errDriver) TranscribeAudioWithSender(modelName *string, file *string, apiConfig *models.APIConfig, asrConfig *models.ASRConfig, sender func(*string, *string) error) error { + return nil +} +func (d *errDriver) AudioSpeech(modelName *string, audioContent *string, apiConfig *models.APIConfig, ttsConfig *models.TTSConfig) (*models.TTSResponse, error) { + return nil, nil +} +func (d *errDriver) AudioSpeechWithSender(modelName *string, audioContent *string, apiConfig *models.APIConfig, ttsConfig *models.TTSConfig, sender func(*string, *string) error) error { + return nil +} +func (d *errDriver) OCRFile(modelName *string, content []byte, url *string, apiConfig *models.APIConfig, ocrConfig *models.OCRConfig) (*models.OCRFileResponse, error) { + return nil, nil +} +func (d *errDriver) ParseFile(modelName *string, content []byte, url *string, apiConfig *models.APIConfig, parseFileConfig *models.ParseFileConfig) (*models.ParseFileResponse, error) { + return nil, nil +} +func (d *errDriver) ListModels(apiConfig *models.APIConfig) ([]models.ListModelResponse, error) { + return nil, nil +} +func (d *errDriver) Balance(apiConfig *models.APIConfig) (map[string]interface{}, error) { + return nil, nil +} +func (d *errDriver) CheckConnection(apiConfig *models.APIConfig) error { return nil } +func (d *errDriver) ListTasks(apiConfig *models.APIConfig) ([]models.ListTaskStatus, error) { + return nil, nil +} +func (d *errDriver) ShowTask(taskID string, apiConfig *models.APIConfig) (*models.TaskResponse, error) { + return nil, nil +} +func (d *errDriver) ToolCall(name string, arguments map[string]interface{}) (string, error) { + return "", nil +} diff --git a/internal/ingestion/task/embedding_token_test.go b/internal/ingestion/task/embedding_token_test.go new file mode 100644 index 0000000000..d7ea315628 --- /dev/null +++ b/internal/ingestion/task/embedding_token_test.go @@ -0,0 +1,55 @@ +package task + +import ( + "testing" +) + +// ============================================================================= +// GetEmbeddingTokenConsumption +// ============================================================================= + +func TestGetEmbeddingTokenConsumption_Int(t *testing.T) { + input := map[string]any{EmbeddingTokenConsumptionKey: 42} + result := GetEmbeddingTokenConsumption(input) + if result != 42 { + t.Errorf("got %d, want 42", result) + } +} + +func TestGetEmbeddingTokenConsumption_Float64(t *testing.T) { + input := map[string]any{EmbeddingTokenConsumptionKey: float64(42)} + result := GetEmbeddingTokenConsumption(input) + if result != 42 { + t.Errorf("got %d, want 42", result) + } +} + +func TestGetEmbeddingTokenConsumption_MissingKey(t *testing.T) { + result := GetEmbeddingTokenConsumption(map[string]any{}) + if result != 0 { + t.Errorf("got %d, want 0", result) + } +} + +func TestGetEmbeddingTokenConsumption_NilMap(t *testing.T) { + result := GetEmbeddingTokenConsumption(nil) + if result != 0 { + t.Errorf("got %d, want 0", result) + } +} + +func TestGetEmbeddingTokenConsumption_WrongType(t *testing.T) { + input := map[string]any{EmbeddingTokenConsumptionKey: "not a number"} + result := GetEmbeddingTokenConsumption(input) + if result != 0 { + t.Errorf("got %d, want 0", result) + } +} + +func TestGetEmbeddingTokenConsumption_Zero(t *testing.T) { + input := map[string]any{EmbeddingTokenConsumptionKey: 0} + result := GetEmbeddingTokenConsumption(input) + if result != 0 { + t.Errorf("got %d, want 0", result) + } +} diff --git a/internal/ingestion/task/golden_compare.go b/internal/ingestion/task/golden_compare.go new file mode 100644 index 0000000000..ab61168e9c --- /dev/null +++ b/internal/ingestion/task/golden_compare.go @@ -0,0 +1,54 @@ +// +// 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 task + +import ( + "time" +) + +// GoldenDataflowResult is the structured output used by the local golden tools. +type GoldenDataflowResult struct { + NormalizedChunks []map[string]any `json:"normalized_chunks"` + ProcessedChunks []map[string]any `json:"processed_chunks"` + MergedMetadata map[string]any `json:"merged_metadata"` +} + +// ProcessPipelineOutputForGolden replays the deterministic dataflow post-processing +// steps from a pipeline.run()-style output without embedding or external writes. +func ProcessPipelineOutputForGolden( + pipelineOutput map[string]any, + docID string, + kbID string, + docName string, +) GoldenDataflowResult { + normalized := NormalizeChunks(pipelineOutput) + if normalized == nil { + normalized = []map[string]any{} + } + + processed := deepCopyChunks(normalized) + metadata := ProcessChunksForDataflow(processed, docID, kbID, docName, time.Now()) + if metadata == nil { + metadata = map[string]any{} + } + + return GoldenDataflowResult{ + NormalizedChunks: normalized, + ProcessedChunks: processed, + MergedMetadata: metadata, + } +} \ No newline at end of file diff --git a/internal/ingestion/task/golden_compare_test.go b/internal/ingestion/task/golden_compare_test.go new file mode 100644 index 0000000000..23005732ae --- /dev/null +++ b/internal/ingestion/task/golden_compare_test.go @@ -0,0 +1,71 @@ +package task + +import ( + "testing" + "time" +) + +func TestProcessPipelineOutputForGolden_Markdown(t *testing.T) { + input := map[string]any{ + "markdown": "# Title\n\nContent", + } + + result := ProcessPipelineOutputForGolden(input, "doc-1", "kb-1", "sample.md") + + if len(result.NormalizedChunks) != 1 { + t.Fatalf("normalized len = %d, want 1", len(result.NormalizedChunks)) + } + if got := result.NormalizedChunks[0]["text"]; got != "# Title\n\nContent" { + t.Fatalf("normalized text = %v, want markdown string", got) + } + if len(result.ProcessedChunks) != 1 { + t.Fatalf("processed len = %d, want 1", len(result.ProcessedChunks)) + } + if got := result.ProcessedChunks[0]["content_with_weight"]; got != "# Title\n\nContent" { + t.Fatalf("content_with_weight = %v, want markdown string", got) + } + if _, exists := result.ProcessedChunks[0]["text"]; exists { + t.Fatal("processed chunk should not keep text key") + } +} + +func TestProcessChunksForDataflow_StableFields(t *testing.T) { + now := time.Date(2026, 7, 3, 20, 0, 0, 0, time.FixedZone("CST", 8*3600)) + chunks := []map[string]any{ + { + "text": "hello", + "questions": "Q1\nQ2", + "keywords": "kw1,kw2", + "summary": "sum", + "metadata": map[string]any{"author": "Alice"}, + }, + } + + meta := ProcessChunksForDataflow(chunks, "doc-1", "kb-1", "sample.md", now) + + chunk := chunks[0] + if chunk["doc_id"] != "doc-1" { + t.Fatalf("doc_id = %v", chunk["doc_id"]) + } + if chunk["docnm_kwd"] != "sample.md" { + t.Fatalf("docnm_kwd = %v", chunk["docnm_kwd"]) + } + if chunk["content_with_weight"] != "hello" { + t.Fatalf("content_with_weight = %v", chunk["content_with_weight"]) + } + if _, ok := chunk["id"].(string); !ok { + t.Fatalf("id should be string, got %T", chunk["id"]) + } + if _, exists := chunk["questions"]; exists { + t.Fatal("questions should be removed") + } + if _, exists := chunk["keywords"]; exists { + t.Fatal("keywords should be removed") + } + if _, exists := chunk["summary"]; exists { + t.Fatal("summary should be removed") + } + if meta["author"] != "Alice" { + t.Fatalf("metadata merge failed: %v", meta) + } +} diff --git a/internal/ingestion/task/has_vectors_test.go b/internal/ingestion/task/has_vectors_test.go new file mode 100644 index 0000000000..5f61d9cc38 --- /dev/null +++ b/internal/ingestion/task/has_vectors_test.go @@ -0,0 +1,74 @@ +package task + +import ( + "testing" +) + +// ============================================================================= +// hasVectors — checks if any chunk has q_\d+_vec key +// ============================================================================= + +func TestHasVectors_WithVector(t *testing.T) { + chunks := []map[string]any{ + {"text": "hello", "q_768_vec": []float64{0.1, 0.2}}, + } + if !hasVectors(chunks) { + t.Error("expected true for chunk with q_768_vec") + } +} + +func TestHasVectors_WithoutVector(t *testing.T) { + chunks := []map[string]any{ + {"text": "hello"}, + } + if hasVectors(chunks) { + t.Error("expected false for chunk without q_*_vec") + } +} + +func TestHasVectors_MultipleChunksOneHasVector(t *testing.T) { + chunks := []map[string]any{ + {"text": "hello"}, + {"text": "world", "q_3_vec": []float64{0.1, 0.2, 0.3}}, + } + if !hasVectors(chunks) { + t.Error("expected true when one chunk has vector") + } +} + +func TestHasVectors_VectorInFirstChunkOnly(t *testing.T) { + chunks := []map[string]any{ + {"q_128_vec": []float64{0.5}, "text": "a"}, + {"text": "b"}, + } + if !hasVectors(chunks) { + t.Error("expected true when first chunk has vector") + } +} + +func TestHasVectors_EmptyChunks(t *testing.T) { + if hasVectors(nil) { + t.Error("expected false for nil chunks") + } + if hasVectors([]map[string]any{}) { + t.Error("expected false for empty chunks") + } +} + +func TestHasVectors_OtherNumberKeys(t *testing.T) { + chunks := []map[string]any{ + {"text": "hello", "q_abc_vec": "not a vector"}, + } + if hasVectors(chunks) { + t.Error("q_abc_vec should not match q_\\d+_vec") + } +} + +func TestHasVectors_EmptySliceValue(t *testing.T) { + chunks := []map[string]any{ + {"text": "hello", "q_0_vec": []float64{}}, + } + if !hasVectors(chunks) { + t.Error("q_0_vec should match even with empty vector") + } +} diff --git a/internal/ingestion/task/naive_merge.go b/internal/ingestion/task/naive_merge.go new file mode 100644 index 0000000000..83fc83f4cd --- /dev/null +++ b/internal/ingestion/task/naive_merge.go @@ -0,0 +1,269 @@ +// +// 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 task + +import ( + "regexp" + "sort" + "strings" + + "ragflow/internal/tokenizer" +) + +// MergedSegment is a merged text chunk produced by NaiveMerge. +// SectionIndices tracks which input sections contributed to this segment, +// enabling the caller to associate images from those sections. +type MergedSegment struct { + Text string `json:"text"` + SectionIndices []int `json:"section_indices"` + Position string `json:"position,omitempty"` +} + +// indexedText pairs text content with its original index and position. +type indexedText struct { + text string + idx int + pos string +} + +// NaiveMerge merges texts into chunks by token count with delimiter-based +// splitting and overlap. Matches Python naive_merge_with_images(). +// +// Parameters: +// - texts: input text segments (one per parsed section) +// - positions: optional position strings (one per section, may be nil) +// - chunkTokenNum: max token count per chunk (default 128 in Python) +// - delimiter: delimiter pattern, supports backtick-quoted custom delimiters +// - overlappedPercent: overlap percentage (0-100) +func NaiveMerge(texts []string, positions []string, chunkTokenNum int, delimiter string, overlappedPercent float64) []MergedSegment { + if len(texts) == 0 { + return nil + } + + // Filter out empty texts but track original indices and positions + var filtered []indexedText + for i, t := range texts { + if strings.TrimSpace(t) != "" { + pos := "" + if positions != nil && i < len(positions) { + pos = positions[i] + } + filtered = append(filtered, indexedText{text: t, idx: i, pos: pos}) + } + } + if len(filtered) == 0 { + return nil + } + + // Parse delimiter: extract custom (backtick-quoted) and normal delimiters + customDels, normalPattern := parseDelimiters(delimiter) + + // Custom delimiter mode: each segment is its own chunk, no merging by token count + if len(customDels) > 0 { + return mergeByCustomDelimiters(filtered, customDels, chunkTokenNum) + } + + // Normal mode: split oversized texts at delimiters, merge by token count + return mergeByTokenCount(filtered, normalPattern, chunkTokenNum, overlappedPercent) +} + +// ── Delimiter parsing ─────────────────────────────────────────────────────── + +// parseDelimiters extracts custom (backtick-quoted) and normal delimiters. +// Matches Python get_delimiters(). +func parseDelimiters(delimiter string) (custom []string, pattern string) { + re := regexp.MustCompile("`([^`]+)`") + matches := re.FindAllStringSubmatchIndex(delimiter, -1) + + var normalChars []rune + lastEnd := 0 + for _, m := range matches { + start, end := m[0], m[1] + normalChars = append(normalChars, []rune(delimiter[lastEnd:start])...) + custom = append(custom, delimiter[m[2]:m[3]]) + lastEnd = end + } + if lastEnd < len(delimiter) { + normalChars = append(normalChars, []rune(delimiter[lastEnd:])...) + } + + var dels []string + seen := make(map[string]bool) + for _, r := range normalChars { + s := string(r) + if !seen[s] && s != "" { + dels = append(dels, regexp.QuoteMeta(s)) + seen[s] = true + } + } + + sort.Slice(dels, func(i, j int) bool { + return len(dels[i]) > len(dels[j]) + }) + + if len(dels) > 0 { + pattern = strings.Join(dels, "|") + } + return +} + +// ── Custom delimiter mode ─────────────────────────────────────────────────── + +func mergeByCustomDelimiters(texts []indexedText, customDels []string, chunkTokenNum int) []MergedSegment { + sort.Slice(customDels, func(i, j int) bool { + return len(customDels[i]) > len(customDels[j]) + }) + var escaped []string + for _, d := range customDels { + escaped = append(escaped, regexp.QuoteMeta(d)) + } + pattern := regexp.MustCompile("(" + strings.Join(escaped, "|") + ")") + + var result []MergedSegment + for _, it := range texts { + parts := pattern.Split(it.text, -1) + for _, part := range parts { + if part == "" { + continue + } + isDelimiter := false + for _, d := range customDels { + if part == d { + isDelimiter = true + break + } + } + if isDelimiter { + continue + } + + textSeg := "\n" + part + pos := resolvePosition(textSeg, it.pos) + + result = append(result, MergedSegment{ + Text: textSeg, + SectionIndices: []int{it.idx}, + Position: pos, + }) + } + } + return result +} + +// ── Normal (token-count-based) mode ───────────────────────────────────────── + +func mergeByTokenCount(texts []indexedText, delimPattern string, chunkTokenNum int, overlappedPercent float64) []MergedSegment { + var result []MergedSegment + + threshold := float64(chunkTokenNum) * (100.0 - overlappedPercent) / 100.0 + + for _, it := range texts { + var subSegments []indexedText + if delimPattern != "" && tokenizer.NumTokensFromString(it.text) >= chunkTokenNum { + re := regexp.MustCompile("(" + delimPattern + ")") + parts := re.Split(it.text, -1) + for _, part := range parts { + if part == "" { + continue + } + if matched, _ := regexp.MatchString("^("+delimPattern+")$", part); matched { + continue + } + subSegments = append(subSegments, indexedText{ + text: "\n" + part, + idx: it.idx, + pos: it.pos, + }) + } + } else { + subSegments = append(subSegments, indexedText{ + text: "\n" + it.text, + idx: it.idx, + pos: it.pos, + }) + } + + for _, sub := range subSegments { + tk := tokenizer.NumTokensFromString(sub.text) + pos := resolvePosition(sub.text, sub.pos) + + if len(result) == 0 { + result = append(result, MergedSegment{ + Text: sub.text, + SectionIndices: []int{sub.idx}, + Position: pos, + }) + continue + } + + last := &result[len(result)-1] + lastTk := tokenizer.NumTokensFromString(last.Text) + + if lastTk > int(threshold) { + newText := sub.text + if overlappedPercent > 0 && lastTk > 0 { + overlapChars := int(float64(len([]rune(last.Text))) * (100.0 - overlappedPercent) / 100.0) + if overlapChars > 0 && overlapChars < len([]rune(last.Text)) { + overlap := string([]rune(last.Text)[overlapChars:]) + newText = overlap + sub.text + } + } + result = append(result, MergedSegment{ + Text: newText, + SectionIndices: []int{sub.idx}, + Position: pos, + }) + } else { + last.Text += sub.text + last.SectionIndices = appendUnique(last.SectionIndices, sub.idx) + // Position: keep the first non-empty position + if last.Position == "" && pos != "" { + last.Position = pos + } + } + _ = tk + } + } + + return result +} + +// ── Helpers ───────────────────────────────────────────────────────────────── + +// resolvePosition applies Python position rules: +// - if token count < 8 → clear position +// - if position is not already in text → append to text (done by caller) +// Returns the resolved position string. +func resolvePosition(text, pos string) string { + if pos == "" { + return "" + } + tk := tokenizer.NumTokensFromString(text) + if tk < 8 { + return "" + } + return pos +} + +func appendUnique(slice []int, val int) []int { + for _, v := range slice { + if v == val { + return slice + } + } + return append(slice, val) +} diff --git a/internal/ingestion/task/naive_merge_test.go b/internal/ingestion/task/naive_merge_test.go new file mode 100644 index 0000000000..f70769f2ce --- /dev/null +++ b/internal/ingestion/task/naive_merge_test.go @@ -0,0 +1,417 @@ +// +// 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 task + +import ( + "strings" + "testing" + + "ragflow/internal/tokenizer" +) + +// naiveMergeTests defines a table-driven test case for NaiveMerge. +type naiveMergeTests struct { + name string + texts []string + chunkTokenNum int + delimiter string + overlappedPct float64 + + // Expectations + minChunks int // at least this many chunks + maxChunks int // at most this many chunks + maxTokens int // no single chunk should exceed this token count + contains []string // each chunk should contain these substrings + notContains []string // no chunk should contain these +} + +func TestNaiveMerge_TableDriven(t *testing.T) { + tests := []naiveMergeTests{ + // ── Empty & nil inputs ────────────────────────────────────────── + { + name: "nil texts", + texts: nil, + chunkTokenNum: 128, + delimiter: "\n。；！？", + overlappedPct: 0, + maxChunks: 0, + }, + { + name: "empty texts", + texts: []string{}, + chunkTokenNum: 128, + delimiter: "\n。；！？", + overlappedPct: 0, + maxChunks: 0, + }, + + // ── Single short text ─────────────────────────────────────────── + { + name: "single short text", + texts: []string{"Hello World"}, + chunkTokenNum: 128, + delimiter: "\n。；！？", + overlappedPct: 0, + minChunks: 1, + maxChunks: 1, + contains: []string{"Hello World"}, + }, + + // ── Multiple short texts merged into one chunk ────────────────── + { + name: "multiple short texts merge into one", + texts: []string{"Short A", "Short B", "Short C"}, + chunkTokenNum: 128, + delimiter: "\n。；！？", + overlappedPct: 0, + minChunks: 1, + maxChunks: 1, + contains: []string{"Short A", "Short B", "Short C"}, + }, + + // ── Large text split at delimiters ────────────────────────────── + { + name: "long text split at delimiter", + texts: []string{strings.Repeat("A。", 100)}, // 100 Chinese sentences + chunkTokenNum: 128, + delimiter: "\n。；！？", + overlappedPct: 0, + minChunks: 2, // should be split at '。' + maxChunks: 20, + }, + + // ── Chinese delimiters ────────────────────────────────────────── + { + name: "chinese sentence delimiters", + texts: []string{"第一句话。第二句话！第三句话？"}, + chunkTokenNum: 128, + delimiter: "\n。；！？", + overlappedPct: 0, + minChunks: 1, + maxChunks: 3, + }, + + // ── Overlap ───────────────────────────────────────────────────── + { + name: "overlap between chunks", + texts: []string{ + strings.Repeat("A", 100) + "。", + strings.Repeat("B", 100) + "。", + strings.Repeat("C", 100) + "。", + }, + chunkTokenNum: 50, // small enough to force multiple chunks + delimiter: "\n。；！？", + overlappedPct: 20, // 20% overlap + minChunks: 2, + maxChunks: 10, + }, + + // ── Custom delimiter (backtick-quoted) ────────────────────────── + { + name: "custom backtick delimiter", + texts: []string{"Section 1 `CHAPTER` Section 2"}, + chunkTokenNum: 128, + delimiter: "`CHAPTER`\n。；！？", + overlappedPct: 0, + minChunks: 2, + maxChunks: 2, + contains: []string{"Section 1", "Section 2"}, + notContains: []string{"CHAPTER"}, + }, + + // ── Empty delimiter — no splitting ───────────────────────────── + { + name: "empty delimiter", + texts: []string{"Keep As Is Without Split"}, + chunkTokenNum: 128, + delimiter: "", + overlappedPct: 0, + minChunks: 1, + maxChunks: 1, + }, + + // ── Multiple texts, some oversized, some not ──────────────────── + { + name: "mixed sizes", + texts: []string{ + "Tiny。", + strings.Repeat("Long text that exceeds the token limit and should be split。", 20), + "Another tiny。", + }, + chunkTokenNum: 128, + delimiter: "\n。；！？", + overlappedPct: 0, + minChunks: 2, // oversized text split, small ones may merge + maxChunks: 30, + }, + + // ── All texts empty strings ───────────────────────────────────── + { + name: "all empty strings", + texts: []string{"", "", ""}, + chunkTokenNum: 128, + delimiter: "\n。；！？", + overlappedPct: 0, + maxChunks: 0, + }, + + // ── Very small chunk_token_num forces many chunks ─────────────── + { + name: "tiny token limit", + texts: []string{"一。二。三。四。五。六。"}, + chunkTokenNum: 5, + delimiter: "\n。；！？", + overlappedPct: 0, + minChunks: 2, // should be split into multiple chunks + maxChunks: 10, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + chunks := NaiveMerge(tt.texts, nil, tt.chunkTokenNum, tt.delimiter, tt.overlappedPct) + + // Check min/max count + if len(chunks) < tt.minChunks { + t.Errorf("got %d chunks, want at least %d", len(chunks), tt.minChunks) + } + if len(chunks) > tt.maxChunks { + t.Errorf("got %d chunks, want at most %d", len(chunks), tt.maxChunks) + } + + // Check max tokens per chunk + if tt.maxTokens > 0 { + for i, c := range chunks { + tk := tokenizer.NumTokensFromString(c.Text) + if tk > tt.maxTokens { + t.Errorf("chunk[%d] has %d tokens, exceeds max %d: %q", i, tk, tt.maxTokens, c.Text) + } + } + } + + // Check contains + for _, substr := range tt.contains { + found := false + for _, c := range chunks { + if strings.Contains(c.Text, substr) { + found = true + break + } + } + if !found { + t.Errorf("no chunk contains %q", substr) + } + } + + // Check not contains + for _, substr := range tt.notContains { + for _, c := range chunks { + if strings.Contains(c.Text, substr) { + t.Errorf("chunk should not contain %q: %q", substr, c.Text) + } + } + } + }) + } +} + +// ── Specific behavior tests ──────────────────────────────────────────────── + +func TestNaiveMerge_OverlapContent(t *testing.T) { + // Create texts where we can verify overlap behavior. + // Each text is a complete sentence with delimiter. + texts := []string{ + "AAAAA。", + "BBBBB。", + "CCCCC。", + "DDDDD。", + } + chunks := NaiveMerge(texts, nil, 10, "\n。；！？", 30) // 30% overlap + + // With token limit 10 and 30% overlap, chunks should reuse end of previous. + // Verify we got chunks and they look reasonable. + if len(chunks) < 2 { + t.Fatalf("expected at least 2 chunks with overlap, got %d", len(chunks)) + } + + // Verify chunks are non-empty and contain content + for i, c := range chunks { + if c.Text == "" { + t.Errorf("chunk[%d] is empty", i) + } + } +} + +func TestNaiveMerge_DelimiterAsSplitPoint(t *testing.T) { + // Delimiters are split points for oversized texts (>chunkTokenNum). + // For texts under the token limit, no splitting occurs. + // This test verifies that a large text IS split at delimiters. + texts := []string{strings.Repeat("A。", 100)} // many sentences, over token limit + chunks := NaiveMerge(texts, nil, 20, "\n。；！？", 0) + + // Should be split into multiple chunks (token limit 20 is small) + if len(chunks) < 2 { + t.Fatalf("expected at least 2 chunks, got %d", len(chunks)) + } + // No chunk should contain '。' (delimiter is removed) + for _, c := range chunks { + if strings.Contains(c.Text, "。") { + t.Errorf("chunk should not contain delimiter '。': %q", c.Text) + } + } +} + +func TestNaiveMerge_AllBelowTokenThreshold(t *testing.T) { + // All texts below chunk_token_num → should merge into one chunk. + texts := []string{"A", "B", "C", "D", "E"} + chunks := NaiveMerge(texts, nil, 128, "\n。；！？", 0) + + if len(chunks) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(chunks)) + } + // Should contain all + for _, letter := range []string{"A", "B", "C", "D", "E"} { + if !strings.Contains(chunks[0].Text, letter) { + t.Errorf("chunk missing %q", letter) + } + } +} + +func TestNaiveMerge_PositionInfo(t *testing.T) { + // Texts with position info should have positions carried through. + texts := []string{"Page 1 content。", "Page 2 content。"} + chunks := NaiveMerge(texts, nil, 128, "\n。；！？", 0) + + // Verify SectionIndices are tracked + for _, c := range chunks { + if len(c.SectionIndices) == 0 { + t.Error("chunk has no SectionIndices") + } + // Section indices should be within range + for _, idx := range c.SectionIndices { + if idx < 0 || idx >= len(texts) { + t.Errorf("SectionIndex %d out of range [0, %d)", idx, len(texts)) + } + } + } +} + +func TestNaiveMerge_CustomDelimiterExclusive(t *testing.T) { + // Custom delimiter mode: each segment is its own chunk, no merging. + texts := []string{"Part A `---` Part B `---` Part C"} + chunks := NaiveMerge(texts, nil, 128, "`---`", 0) + + // Should produce 3 chunks (one per segment between ---) + if len(chunks) != 3 { + t.Fatalf("expected 3 chunks from custom delimiter split, got %d", len(chunks)) + } + // No chunk should contain "---" + for _, c := range chunks { + if strings.Contains(c.Text, "---") { + t.Errorf("chunk should not contain custom delimiter: %q", c.Text) + } + } +} + +func TestNaiveMerge_LeadingNewline(t *testing.T) { + // Python adds "\n" prefix to each text before adding as chunk. + texts := []string{"Hello World"} + chunks := NaiveMerge(texts, nil, 128, "\n。；！？", 0) + + if len(chunks) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(chunks)) + } + // Should start with "\n" (matching Python behavior) + if !strings.HasPrefix(chunks[0].Text, "\n") { + t.Errorf("expected chunk to start with newline, got %q", chunks[0].Text) + } +} + +func TestNaiveMerge_SectionIndicesCorrect(t *testing.T) { + // Verify that merged chunks track which input sections they came from. + texts := []string{"First。", "Second。", "Third。", "Fourth。"} + chunks := NaiveMerge(texts, nil, 40, "\n。；！？", 0) + + // Collect all section indices across all chunks + allIndices := make(map[int]bool) + for _, c := range chunks { + for _, idx := range c.SectionIndices { + allIndices[idx] = true + } + } + + // Every input section should appear in exactly one chunk + for i := range texts { + if !allIndices[i] { + t.Errorf("section %d (%q) not found in any chunk", i, texts[i]) + } + } +} + +// ── Position-aware tests (gap: Python carries pos through add_chunk) ────────── + +func TestNaiveMerge_PositionCarriedThrough(t *testing.T) { + texts := []string{"Long enough text with sufficient tokens for position to stay。", "Another long sentence with enough tokens。"} + positions := []string{"p1", "p2"} + chunks := NaiveMerge(texts, positions, 128, "\n。；！？", 0) + + if len(chunks) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(chunks)) + } + if chunks[0].Position == "" { + t.Error("position should not be empty for merged chunk") + } +} + +func TestNaiveMerge_PositionClearedWhenBelowThreshold(t *testing.T) { + // Token count < 8 → position should be cleared (Python: tnum < 8 → pos = "") + texts := []string{"Hi"} + positions := []string{"page1_top100"} + chunks := NaiveMerge(texts, positions, 128, "\n。；！？", 0) + + if len(chunks) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(chunks)) + } + if chunks[0].Position != "" { + t.Errorf("position should be empty for short text (< 8 tokens), got %q", chunks[0].Position) + } +} + +func TestNaiveMerge_PositionAppendedWhenNotPresent(t *testing.T) { + texts := []string{"Long enough text that exceeds eight tokens easily。"} + positions := []string{"[page1]"} + chunks := NaiveMerge(texts, positions, 128, "\n。；！？", 0) + + if len(chunks) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(chunks)) + } + if chunks[0].Position != "[page1]" { + t.Errorf("position = %q, want %q", chunks[0].Position, "[page1]") + } +} + +func TestNaiveMerge_NilPositions(t *testing.T) { + texts := []string{"Text A。", "Text B。"} + chunks := NaiveMerge(texts, nil, 128, "\n。；！？", 0) + + if len(chunks) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(chunks)) + } + if chunks[0].Position != "" { + t.Errorf("position should be empty when positions is nil, got %q", chunks[0].Position) + } +} diff --git a/internal/ingestion/task/pipeline_run_output_checklist.md b/internal/ingestion/task/pipeline_run_output_checklist.md new file mode 100644 index 0000000000..19c9c27903 --- /dev/null +++ b/internal/ingestion/task/pipeline_run_output_checklist.md @@ -0,0 +1,178 @@ +# pipeline.run 输出字段清单 + +## 1. 问题背景 + +当前 Go 侧 `RunDataflow` 的行为以 Python 重构版 +[`rag/svr/task_executor_refactor/dataflow_service.py`](/home/shenyushi/cc-workspace/ragflow/rag/svr/task_executor_refactor/dataflow_service.py:90) +为准。 + +这个 Python 版本里，`run_dataflow()` 在处理普通正文 chunk 时： + +- 会把 `text` 搬到 `content_with_weight` +- 不会为普通正文额外补 `content_ltks / content_sm_ltks` +- 只有在 chunk 自带 `summary` 时，才会在 `run_dataflow()` 里补 `content_ltks / content_sm_ltks` + +因此，如果 `pipeline.run` 的 terminal `Tokenizer` 输出里缺少 +`content_ltks / content_sm_ltks`，问题应归因于 `pipeline.run` 输出不符合契约， +而不是 `RunDataflow` 这层缺少兜底。 + +## 2. 契约来源 + +本清单以 Python 端 `Tokenizer` 组件实现为准： + +- [`rag/flow/tokenizer/tokenizer.py`](/home/shenyushi/cc-workspace/ragflow/rag/flow/tokenizer/tokenizer.py:132) +- 当前验证路径是 terminal 组件 `Tokenizer:LegalReadersDecide` +- 当前模板验证时使用 `search_method = ["full_text"]` +- 当前真实样本路径是： + - 真实 MySQL + - 真实 MinIO + - 真实 `pipeline.run` + - 输入文件内容为两段纯文本： + - `Alpha paragraph.` + - `Beta paragraph.` + +## 3. 当前路径下 pipeline.run 应输出哪些字段 + +这里说的是 terminal `Tokenizer` 输出的 **每个 chunk** 应至少带哪些字段。 + +### 必备字段 + +`text: string` +- 原始正文内容。 +- 这是 `RunDataflow` 后续生成 `content_with_weight` 的来源。 + +`chunk_order_int: int` +- chunk 顺序号。 +- 用于保持切分后的稳定顺序。 + +`title_tks: string` +- 文档名的粗粒度 token。 +- Python 中由 `from_upstream.name` 去掉扩展名后再做 `rag_tokenizer.tokenize(...)`。 + +`title_sm_tks: string` +- 文档名的细粒度 token。 +- Python 中由 `rag_tokenizer.fine_grained_tokenize(title_tks)` 得到。 + +`content_ltks: string` +- 正文的粗粒度 token。 +- 当前这条路径里，如果 chunk 有 `summary`，取 `summary`； + 否则取 `text`。 + +`content_sm_ltks: string` +- 正文的细粒度 token。 +- Python 中由 `rag_tokenizer.fine_grained_tokenize(content_ltks)` 得到。 + +### 条件字段 + +`question_kwd: list[str]` +- 当 chunk 有 `questions` 时输出。 +- 是按换行拆开的问句列表。 + +`question_tks: string` +- 当 chunk 有 `questions` 时输出。 +- 是 `questions` 的 token 化结果。 + +`important_kwd: list[str]` +- 当 chunk 有 `keywords` 时输出。 +- Python 当前 `Tokenizer` 里是按英文逗号拆分。 + +`important_tks: string` +- 当 chunk 有 `keywords` 时输出。 +- 是 `keywords` 的 token 化结果。 + +### embedding 打开时才要求的字段 + +`q__vec: list[float]` +- embedding 向量。 +- 例如 `q_1024_vec`。 + +顶层 `embedding_token_consumption: int` +- embedding 推理消耗的 token 数。 +- 只有 `search_method` 包含 `embedding` 时才要求。 + +## 4. 顶层输出格式要求 + +terminal `Tokenizer` 顶层输出应至少有： + +`output_format: "chunks"` +- 表示 terminal 输出已经是 chunk 列表。 + +`chunks: list[dict]` +- chunk 列表。 + +此外运行时还可能带有： + +- `_created_time` +- `_elapsed_time` +- `__cpn_id__` + +这些属于运行时辅助字段，不是这次问题的关键。 + +## 5. 当前真实 pipeline.run 实测输出 + +基于真实 integration 链路，当前 terminal payload 顶层实际有： + +- `output_format` +- `chunks` +- `__cpn_id__` +- `_created_time` +- `_elapsed_time` + +当前每个 chunk 实际有： + +- `chunk_order_int` +- `ck_type` +- `doc_type_kwd` +- `text` +- `tk_nums` + +也就是说，当前 terminal `Tokenizer` 输出看起来更像是把上游 chunker 结果直接透传了， +没有把 Python `Tokenizer` 在 `full_text` 分支里本应补充的搜索字段带出来。 + +## 6. 当前缺失字段清单 + +按 Python `Tokenizer` 契约，当前样本里明确缺失： + +`title_tks` +- 缺失 + +`title_sm_tks` +- 缺失 + +`content_ltks` +- 缺失 + +`content_sm_ltks` +- 缺失 + +## 7. 缺失字段的影响 + +`title_tks / title_sm_tks` +- 会影响基于标题的全文检索和召回质量。 + +`content_ltks / content_sm_ltks` +- 这是当前最关键的问题。 +- Python `run_dataflow()` 不会为普通正文兜底补这两个字段。 +- 因此它们一旦在 `pipeline.run` 输出中缺失，后续索引侧就会拿不到正文 token 字段。 + +## 8. 当前结论 + +当前可以明确得出的结论是： + +1. Go 侧 `RunDataflow` 不应该增强补齐逻辑。 +2. 测试预期不应该放宽。 +3. 因此问题应归因于 `pipeline.run` terminal `Tokenizer` 输出不符合 Python 契约。 + +## 9. 待同事接入/修复后的核对项 + +- terminal `Tokenizer` 输出顶层仍为 `output_format = "chunks"` +- 每个 chunk 至少带： + - `text` + - `chunk_order_int` + - `title_tks` + - `title_sm_tks` + - `content_ltks` + - `content_sm_ltks` +- 若开启 embedding，再额外核对： + - `q__vec` + - `embedding_token_consumption` diff --git a/internal/ingestion/task/position.go b/internal/ingestion/task/position.go new file mode 100644 index 0000000000..8c42a162e8 --- /dev/null +++ b/internal/ingestion/task/position.go @@ -0,0 +1,52 @@ +// +// 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 task + +// AddPositions adds position fields to a chunk map. +// Input positions is a flat []float64 grouped as [pn, left, right, top, bottom] +// every 5 elements. pn is 0-indexed; output is 1-indexed. +// +// Mirrors Python: rag.nlp.add_positions() +// for pn, left, right, top, bottom in poss: +// page_num_int.append(int(pn + 1)) +// top_int.append(int(top)) +// position_int.append((int(pn + 1), int(left), int(right), int(top), int(bottom))) +func AddPositions(chunk map[string]any, positions []float64) { + if len(positions) == 0 || len(positions)%5 != 0 { + return + } + n := len(positions) / 5 + pageNumInt := make([]int, 0, n) + topInt := make([]int, 0, n) + positionInt := make([][]int, 0, n) + + for i := 0; i < len(positions); i += 5 { + pn := int(positions[i]) + 1 // 0-indexed → 1-indexed + left := int(positions[i+1]) + right := int(positions[i+2]) + top := int(positions[i+3]) + bottom := int(positions[i+4]) + + pageNumInt = append(pageNumInt, pn) + topInt = append(topInt, top) + positionInt = append(positionInt, []int{pn, left, right, top, bottom}) + } + + chunk["page_num_int"] = pageNumInt + chunk["top_int"] = topInt + chunk["position_int"] = positionInt +} diff --git a/internal/ingestion/task/position_test.go b/internal/ingestion/task/position_test.go new file mode 100644 index 0000000000..6e31c6a1fc --- /dev/null +++ b/internal/ingestion/task/position_test.go @@ -0,0 +1,92 @@ +package task + +import ( + "testing" +) + +// ============================================================================= +// AddPositions +// Canonical format: [pageNum, left, right, top, bottom] × N +// Mirrors Python: rag.nlp.add_positions() +// ============================================================================= + +func TestAddPositions_Basic(t *testing.T) { + chunk := map[string]any{} + // [pn=0, left=100, right=50, top=200, bottom=150] + positions := []float64{0, 100, 50, 200, 150} + AddPositions(chunk, positions) + + pageNum, ok := chunk["page_num_int"].([]int) + if !ok || len(pageNum) != 1 || pageNum[0] != 1 { + t.Errorf("page_num_int = %v, want [1]", pageNum) + } + top, ok := chunk["top_int"].([]int) + if !ok || len(top) != 1 || top[0] != 200 { + t.Errorf("top_int = %v, want [200]", top) + } + position, ok := chunk["position_int"].([][]int) + if !ok || len(position) != 1 { + t.Fatalf("position_int = %v, want [[1 100 50 200 150]]", position) + } + if position[0][0] != 1 || position[0][1] != 100 || position[0][2] != 50 || position[0][3] != 200 || position[0][4] != 150 { + t.Errorf("position_int[0] = %v, want [1 100 50 200 150]", position[0]) + } +} + +func TestAddPositions_MultiplePositions(t *testing.T) { + chunk := map[string]any{} + positions := []float64{ + 0, 100, 50, 200, 150, // pn=0, left=100, right=50, top=200, bottom=150 + 1, 200, 60, 300, 250, // pn=1, left=200, right=60, top=300, bottom=250 + } + AddPositions(chunk, positions) + + pageNum := chunk["page_num_int"].([]int) + if len(pageNum) != 2 || pageNum[0] != 1 || pageNum[1] != 2 { + t.Errorf("page_num_int = %v, want [1 2]", pageNum) + } + top := chunk["top_int"].([]int) + if len(top) != 2 || top[0] != 200 || top[1] != 300 { + t.Errorf("top_int = %v, want [200 300]", top) + } + position := chunk["position_int"].([][]int) + if len(position) != 2 { + t.Fatalf("position_int len = %d, want 2", len(position)) + } +} + +func TestAddPositions_NilPositions(t *testing.T) { + chunk := map[string]any{} + AddPositions(chunk, nil) + if _, exists := chunk["page_num_int"]; exists { + t.Error("page_num_int should not be set for nil positions") + } +} + +func TestAddPositions_EmptyPositions(t *testing.T) { + chunk := map[string]any{} + AddPositions(chunk, []float64{}) + if _, exists := chunk["page_num_int"]; exists { + t.Error("page_num_int should not be set for empty positions") + } +} + +func TestAddPositions_PartialPositions(t *testing.T) { + chunk := map[string]any{} + positions := []float64{0, 100} // only 2 elements, not a complete position + AddPositions(chunk, positions) + if _, exists := chunk["page_num_int"]; exists { + t.Error("page_num_int should not be set for partial positions") + } +} + +func TestAddPositions_PageNumOffset(t *testing.T) { + chunk := map[string]any{} + positions := []float64{5, 100, 50, 200, 150} // pn=5 → 5+1=6 + AddPositions(chunk, positions) + + pageNum := chunk["page_num_int"].([]int) + if pageNum[0] != 6 { + t.Errorf("page_num_int = %d, want 6 (5+1)", pageNum[0]) + } +} diff --git a/internal/ingestion/task/real_consumer_test.go b/internal/ingestion/task/real_consumer_test.go new file mode 100644 index 0000000000..a08cad2a6f --- /dev/null +++ b/internal/ingestion/task/real_consumer_test.go @@ -0,0 +1,170 @@ +// +// 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. +// + +//go:build integration +// +build integration + +package task + +import ( + "encoding/json" + "errors" + "testing" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/engine/nats" + "ragflow/internal/entity" + + "github.com/glebarez/sqlite" + "gorm.io/gorm" +) + +// TestRealProducerConsumer exercises the project's real producer and consumer code paths: +// +// Producer: document.go pattern — CheckAndCreate(IngestionTask) → PublishTask(NATS) +// Consumer: Ingestor.Start() core logic — calls each actual function in sequence +func TestRealProducerConsumer(t *testing.T) { + // ── 1. NATS ── + natsEngine := nats.NewNatsEngine("localhost", 4222) + if err := natsEngine.Init(); err != nil { + t.Fatalf("NATS Init: %v", err) + } + if err := natsEngine.InitConsumer("tasks.>"); err != nil { + t.Fatalf("InitConsumer: %v", err) + } + + // Purge stale messages + for { + h, _ := natsEngine.GetMessages(1) + if len(h) == 0 { + break + } + h[0].Ack() + } + + // ── 2. SQLite DB ── + db, _ := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{TranslateError: true}) + db.AutoMigrate( + &entity.IngestionTask{}, &entity.Task{}, + &entity.Document{}, &entity.Knowledgebase{}, &entity.Tenant{}, + ) + origDB := dao.DB + dao.DB = db + t.Cleanup(func() { dao.DB = origDB }) + + db.Create(&entity.Tenant{ID: "t1", LLMID: "gpt-4", Status: ptr("1")}) + db.Create(&entity.Knowledgebase{ID: "kb1", TenantID: "t1", EmbdID: "e1", Status: ptr("1"), ParserConfig: entity.JSONMap{}}) + db.Create(&entity.Document{ID: "doc-real", KbID: "kb1", ParserID: "naive", ParserConfig: entity.JSONMap{}}) + + // ── 3. Producer: Mirrors document.go:1062-1085 exactly ── + ingestionTask := &entity.IngestionTask{ + ID: "ingest-task-1", + UserID: "u1", + DocumentID: "doc-real", + DatasetID: "kb1", + Status: common.CREATED, + } + created, err := dao.NewIngestionTaskDAO().CheckAndCreate(ingestionTask) + if err != nil { + t.Fatalf("CheckAndCreate: %v", err) + } + t.Logf("Producer: IngestionTask created id=%s status=%s", created.ID, created.Status) + + taskMessage := common.TaskMessage{ + TaskID: created.ID, + TaskType: common.TaskTypeIngestionTask, + } + payload, _ := json.Marshal(taskMessage) + if err := natsEngine.PublishTask("tasks.RAGFLOW", payload); err != nil { + t.Fatalf("PublishTask: %v", err) + } + t.Logf("Producer: Published %s", payload) + + // ── 4. Consumer: Mirrors Ingestor.Start():131-189 exactly ── + handles, err := natsEngine.GetMessages(1) + if err != nil { + t.Fatalf("GetMessages: %v", err) + } + if len(handles) != 1 { + t.Fatalf("expected 1 message, got %d", len(handles)) + } + taskHandle := handles[0] + taskMsg := taskHandle.GetMessage() + t.Logf("Consumer: Received TaskID=%s TaskType=%s", taskMsg.TaskID, taskMsg.TaskType) + + // Mirrors Start():133 — type filter + if taskMsg.TaskType != common.TaskTypeIngestionTask { + taskHandle.Ack() + t.Fatalf("unexpected task type: %s", taskMsg.TaskType) + } + + // Mirrors Start():142-143 — SetRunningByIngestor + ingestionTaskDAO := dao.NewIngestionTaskDAO() + task, err := ingestionTaskDAO.SetRunningByIngestor(taskMsg.TaskID) + if err != nil { + if errors.Is(err, common.ErrTaskNotFound) { + t.Logf("Consumer: task %s not found in ingestion_task table — skipped", taskMsg.TaskID) + taskHandle.Ack() + return + } + t.Fatalf("SetRunningByIngestor: %v", err) + } + t.Logf("Consumer: SetRunningByIngestor status=%s", task.Status) + + // Mirrors Start():167-180 — status check + switch task.Status { + case common.COMPLETED, common.STOPPED, common.FAILED: + taskHandle.Ack() + t.Fatalf("task already terminal: %s", task.Status) + case common.STOPPING, common.CREATED: + t.Fatalf("unexpected status: %s", task.Status) + case common.RUNNING: + t.Logf("Consumer: task is RUNNING — dispatching to executeTask") + } + + // ── 5. executeTask (our modified version) ── + // executeTask needs DB data for TaskHandler + db.Create(&entity.Task{ID: task.ID, DocID: "doc-real", FromPage: 0, ToPage: 100000}) + + tc, err := LoadTaskContext(task.ID) + if err != nil { + t.Fatalf("LoadTaskContext: %v", err) + } + t.Logf("Consumer: Loaded Doc=%s Parser=%s KB=%s Tenant=%s", + tc.Doc.ID, tc.Doc.ParserID, tc.KB.ID, tc.Tenant.ID) + + handler := NewTaskHandler(tc) + if err := handler.Handle(); err != nil { + t.Fatalf("Handle: %v", err) + } + t.Log("Consumer: TaskHandler.Handle() — OK") + + // Mirrors executeTask — mark as completed + if err := ingestionTaskDAO.UpdateStatus(task.ID, common.COMPLETED); err != nil { + t.Fatalf("UpdateStatus: %v", err) + } + + // Mirrors Start():135 — Ack + taskHandle.Ack() + + // ── 6. Verify ── + final, _ := ingestionTaskDAO.GetByID(task.ID) + if final.Status != common.COMPLETED { + t.Errorf("final status = %s, want %s", final.Status, common.COMPLETED) + } + t.Logf("Final: IngestionTask status=%s ✅", final.Status) +} diff --git a/internal/ingestion/task/task_context.go b/internal/ingestion/task/task_context.go new file mode 100644 index 0000000000..285917fac6 --- /dev/null +++ b/internal/ingestion/task/task_context.go @@ -0,0 +1,190 @@ +// +// 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 task + +import ( + "context" + "fmt" + "time" + + "golang.org/x/sync/semaphore" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" +) + +// TaskLog represents a task execution log entry. +type TaskLog struct { + StartTime time.Time `json:"start_time"` + EndTime time.Time `json:"end_time"` + Description string `json:"description"` + Details map[string]interface{} `json:"details"` +} + +// TaskContext holds the full task context loaded from DB, equivalent to Python TaskService.get_task(). +// Entity types are embedded directly to avoid manual field mapping. +type TaskContext struct { + // Context and execution control (from service package) + Ctx context.Context + TaskHandle common.TaskHandle + + // IngestionTask is the ingestion task entity (from service package) + IngestionTask *entity.IngestionTask + + // Progress and status tracking (from service package) + Logs []*TaskLog + Progress int32 + ErrorMessage string + + // Task type for handling dispatch + TaskType string + + // Business data (original task package fields) + Doc entity.Document + KB entity.Knowledgebase + Tenant entity.Tenant + + // pipeline used to parse the document + PipelineID string + + // File is an optional file object for dataflow pipeline debugging. + // Mirrors Python: task["file"] — passed through to Pipeline.run(). + File any + + // EmbedLimiter limits embedding API concurrency. + // If nil, no concurrency limit is applied. + // Mirrors Python: ctx.embed_limiter (asyncio.Semaphore) + EmbedLimiter *semaphore.Weighted + + // ProgressFunc updates the task-level progress state. + // If nil, a no-op callback is used. + ProgressFunc ProgressFunc +} + +// NewTaskContextForScheduling creates a lightweight TaskContext for queue scheduling. +// This only sets the scheduling-related fields, not the full business data. +func NewTaskContextForScheduling(ctx context.Context, task *entity.IngestionTask, handle common.TaskHandle) *TaskContext { + return &TaskContext{ + Ctx: ctx, + IngestionTask: task, + TaskHandle: handle, + } +} + +// LoadFromIngestionTask loads the full task context from an IngestionTask. +// It follows the FK chain: ingestion task → document → knowledgebase → tenant. +func LoadFromIngestionTask(ingestionTask *entity.IngestionTask) (*TaskContext, error) { + doc, err := dao.NewDocumentDAO().GetByID(ingestionTask.DocumentID) + if err != nil || doc == nil { + return nil, fmt.Errorf("error when load document %s : %w", ingestionTask.DocumentID, err) + } + + kb, err := dao.NewKnowledgebaseDAO().GetByID(doc.KbID) + if err != nil || kb == nil { + return nil, fmt.Errorf("error when load knowledgebase %s: %w", doc.KbID, err) + } + + tenant, err := dao.NewTenantDAO().GetByID(kb.TenantID) + if err != nil || tenant == nil { + return nil, fmt.Errorf("error when load tenant %s: %w", kb.TenantID, err) + } + + pipelineID := "" + if doc.PipelineID != nil { + pipelineID = *doc.PipelineID + } + + return &TaskContext{ + IngestionTask: ingestionTask, + TaskType: "dataflow", + PipelineID: pipelineID, + Doc: *doc, + KB: *kb, + Tenant: *tenant, + }, nil +} + +// Load loads the full task context following the FK chain: document → knowledgebase → tenant. +// Kept for backward compatibility. +func Load(docID string) (*TaskContext, error) { + doc, err := dao.NewDocumentDAO().GetByID(docID) + if err != nil || doc == nil { + return nil, fmt.Errorf("error when load document %s : %w", docID, err) + } + + kb, err := dao.NewKnowledgebaseDAO().GetByID(doc.KbID) + if err != nil || kb == nil { + return nil, fmt.Errorf("error when load knowledgebase %s: %w", doc.KbID, err) + } + + tenant, err := dao.NewTenantDAO().GetByID(kb.TenantID) + if err != nil || tenant == nil { + return nil, fmt.Errorf("error when load tenant %s: %w", kb.TenantID, err) + } + + pipelineID := "" + if doc.PipelineID != nil { + pipelineID = *doc.PipelineID + } + + return &TaskContext{ + TaskType: "dataflow", + PipelineID: pipelineID, + Doc: *doc, + KB: *kb, + Tenant: *tenant, + }, nil +} + +// LoadTaskContext loads the full task context following the FK chain: task → document → knowledgebase → tenant. +// Kept for backward compatibility. +func LoadTaskContext(taskID string) (*TaskContext, error) { + task, err := dao.NewTaskDAO().GetByID(taskID) + if err != nil { + return nil, fmt.Errorf("task %s: %w", taskID, err) + } + + doc, err := dao.NewDocumentDAO().GetByID(task.DocID) + if err != nil || doc == nil { + return nil, fmt.Errorf("document %s for task %s: %w", task.DocID, taskID, err) + } + + kb, err := dao.NewKnowledgebaseDAO().GetByID(doc.KbID) + if err != nil || kb == nil { + return nil, fmt.Errorf("knowledgebase %s for doc %s: %w", doc.KbID, doc.ID, err) + } + + tenant, err := dao.NewTenantDAO().GetByID(kb.TenantID) + if err != nil || tenant == nil { + return nil, fmt.Errorf("tenant %s for kb %s: %w", kb.TenantID, kb.ID, err) + } + + pipelineID := "" + if doc.PipelineID != nil { + pipelineID = *doc.PipelineID + } + + return &TaskContext{ + TaskType: task.TaskType, + PipelineID: pipelineID, + Doc: *doc, + KB: *kb, + Tenant: *tenant, + ProgressFunc: func(prog float64, msg string) {}, + }, nil +} diff --git a/internal/ingestion/task/task_context_test.go b/internal/ingestion/task/task_context_test.go new file mode 100644 index 0000000000..ff13959d45 --- /dev/null +++ b/internal/ingestion/task/task_context_test.go @@ -0,0 +1,187 @@ +// +// 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 task + +import ( + "testing" + + "ragflow/internal/dao" + "ragflow/internal/entity" + + "github.com/glebarez/sqlite" + "gorm.io/gorm" +) + +func setupTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{TranslateError: true}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate( + &entity.Task{}, + &entity.Document{}, + &entity.Knowledgebase{}, + &entity.Tenant{}, + ); err != nil { + t.Fatalf("migrate: %v", err) + } + return db +} + +func pushDB(t *testing.T, db *gorm.DB) { + t.Helper() + orig := dao.DB + dao.DB = db + t.Cleanup(func() { dao.DB = orig }) +} + +func ptr(s string) *string { return &s } + +func TestLoadTaskContext(t *testing.T) { + db := setupTestDB(t) + pushDB(t, db) + + // Prepare test data in FK dependency order: tenant → kb → document → task + tenant := &entity.Tenant{ + ID: "tenant-1", + LLMID: "gpt-4", + ASRID: "whisper-1", + Img2TxtID: "gpt-4-vision", + Status: ptr("1"), + } + kb := &entity.Knowledgebase{ + ID: "kb-1", + TenantID: "tenant-1", + Language: ptr("Chinese"), + EmbdID: "embd-model-1", + Pagerank: 0, + Status: ptr(string(entity.StatusValid)), + ParserConfig: entity.JSONMap{ + "chunk_token_size": float64(512), + }, + } + doc := &entity.Document{ + ID: "doc-1", + KbID: "kb-1", + ParserID: "naive", + ParserConfig: entity.JSONMap{"pages": []any{[]any{float64(1), float64(10)}}}, + Name: ptr("test.pdf"), + Type: "pdf", + Location: ptr("bucket/test.pdf"), + Size: 1024, + } + task := &entity.Task{ + ID: "task-1", + DocID: "doc-1", + FromPage: 0, + ToPage: 100000, + } + + db.Create(tenant) + db.Create(kb) + db.Create(doc) + db.Create(task) + + got, err := LoadTaskContext("task-1") + if err != nil { + t.Fatalf("LoadTaskContext: %v", err) + } + + // === task table fields (no longer stored in TaskContext) === + // Task field removed from TaskContext, skipping tests + + // === document table fields === + if got.Doc.ID != "doc-1" { + t.Errorf("Doc.ID = %v, want doc-1", got.Doc.ID) + } + if got.Doc.ParserID != "naive" { + t.Errorf("Doc.ParserID = %v, want naive", got.Doc.ParserID) + } + if got.Doc.Name == nil || *got.Doc.Name != "test.pdf" { + t.Errorf("Doc.Name = %v, want test.pdf", got.Doc.Name) + } + if got.Doc.Type != "pdf" { + t.Errorf("Doc.Type = %v, want pdf", got.Doc.Type) + } + if got.Doc.Location == nil || *got.Doc.Location != "bucket/test.pdf" { + t.Errorf("Doc.Location = %v, want bucket/test.pdf", got.Doc.Location) + } + if got.Doc.Size != 1024 { + t.Errorf("Doc.Size = %v, want 1024", got.Doc.Size) + } + if got.Doc.ParserConfig == nil { + t.Error("Doc.ParserConfig is nil") + } else if got.Doc.ParserConfig["pages"] == nil { + t.Error("Doc.ParserConfig.pages is nil") + } + + // === knowledgebase table fields === + if got.KB.ID != "kb-1" { + t.Errorf("KB.ID = %v, want kb-1", got.KB.ID) + } + if got.KB.TenantID != "tenant-1" { + t.Errorf("KB.TenantID = %v, want tenant-1", got.KB.TenantID) + } + if got.KB.Language == nil || *got.KB.Language != "Chinese" { + t.Errorf("KB.Language = %v, want Chinese", got.KB.Language) + } + if got.KB.EmbdID != "embd-model-1" { + t.Errorf("KB.EmbdID = %v, want embd-model-1", got.KB.EmbdID) + } + if got.KB.Pagerank != 0 { + t.Errorf("KB.Pagerank = %v, want 0", got.KB.Pagerank) + } + if got.KB.ParserConfig == nil { + t.Error("KB.ParserConfig is nil") + } else if got.KB.ParserConfig["chunk_token_size"] != float64(512) { + t.Errorf("KB.ParserConfig.chunk_token_size = %v, want 512", + got.KB.ParserConfig["chunk_token_size"]) + } + + // === tenant table fields === + if got.Tenant.LLMID != "gpt-4" { + t.Errorf("Tenant.LLMID = %v, want gpt-4", got.Tenant.LLMID) + } +} + +func TestLoadTaskContext_TaskNotFound(t *testing.T) { + db := setupTestDB(t) + pushDB(t, db) + + _, err := LoadTaskContext("nonexistent") + if err == nil { + t.Fatal("expected error for nonexistent task") + } +} + +func TestLoadTaskContext_DocNotFound(t *testing.T) { + db := setupTestDB(t) + pushDB(t, db) + + db.Create(&entity.Task{ + ID: "task-1", + DocID: "nonexistent-doc", + FromPage: 0, + ToPage: 100000, + }) + + _, err := LoadTaskContext("task-1") + if err == nil { + t.Fatal("expected error when document not found") + } +} diff --git a/internal/ingestion/task/task_handler.go b/internal/ingestion/task/task_handler.go new file mode 100644 index 0000000000..4a59234fb7 --- /dev/null +++ b/internal/ingestion/task/task_handler.go @@ -0,0 +1,167 @@ +// +// 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 task + +import ( + "context" + "fmt" + "strings" + + "ragflow/internal/entity" + "ragflow/internal/entity/models" + "ragflow/internal/service" +) + +// TaskHandler dispatches document processing tasks by task_type. +// Mirrors Python task_handler.py:handle(). +type TaskHandler struct { + ctx *TaskContext + newDataflowService func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) +} + +// NewTaskHandler creates a TaskHandler for the given task context. +func NewTaskHandler(ctx *TaskContext) *TaskHandler { + return &TaskHandler{ + ctx: ctx, + newDataflowService: func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { + return NewDataflowService(ctx, dataflowID, 0, 0) + }, + } +} + +func (h *TaskHandler) WithDataflowServiceFactory(factory func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error)) *TaskHandler { + h.newDataflowService = factory + return h +} + +// Handle routes the task by type and executes the appropriate handler. +func (h *TaskHandler) Handle() error { + if h.ctx == nil { + return fmt.Errorf("task handler: nil context") + } + taskType := h.ctx.TaskType + if taskType == "" { + // Determine task type - use PipelineID presence as indicator for dataflow + taskType = "dataflow" // Default to dataflow for now + if h.ctx.PipelineID == "" { + taskType = "standard" + } + } + + switch { + case taskType == "memory": + return h.handleMemory() + case taskType == "dataflow" && h.ctx.Doc.ID == CANVAS_DEBUG_DOC_ID: + return h.handleDataflow() + case strings.HasPrefix(taskType, "dataflow"): + return h.handleDataflow() + case taskType == "raptor": + return h.handleRaptor() + case taskType == "graphrag": + return h.handleGraphRAG() + case taskType == "mindmap": + return h.handleStub("mindmap") + case taskType == "evaluation": + return h.handleStub("evaluation") + case taskType == "reembedding": + return h.handleStub("reembedding") + case taskType == "clone": + return h.handleStub("clone") + default: + return h.handleStandard() + } +} + +func (h *TaskHandler) handleMemory() error { + return nil // stub +} + +func (h *TaskHandler) handleDataflow() error { + dataflowID := "" + if strings.Trim(h.ctx.PipelineID, " ") != "" { + dataflowID = h.ctx.PipelineID + } + svc, err := h.newDataflowService(h.ctx, dataflowID) + if err != nil { + return err + } + return svc.Run(context.Background()) +} + +func (h *TaskHandler) handleRaptor() error { + return nil // stub +} + +func (h *TaskHandler) handleGraphRAG() error { + return nil // stub +} + +func (h *TaskHandler) handleStub(name string) error { + return nil +} + +func (h *TaskHandler) handleStandard() error { + return nil // stub +} + +// BindEmbeddingModel creates an embedding model for the task's tenant. +// Returns the model and its vector dimension. +// Mirrors Python: _bind_embedding_model (task_handler.py:204) +func (h *TaskHandler) BindEmbeddingModel() (*models.EmbeddingModel, int, error) { + var model *models.EmbeddingModel + var err error + if embdID := h.ctx.KB.EmbdID; embdID != "" { + modelSvc := service.NewModelProviderService() + model, err = modelSvc.GetEmbeddingModel(h.ctx.Tenant.ID, embdID) + } else { + // Use tenant's default embedding model + model, err = defaultBindEmbeddingModel(h.ctx.Tenant.ID) + } + + if err != nil { + return nil, 0, fmt.Errorf("bind embedding model: %w", err) + } + + dim, err := getEmbeddingDimension(model) + if err != nil { + return nil, 0, fmt.Errorf("bind embedding model: %w", err) + } + return model, dim, err +} + +// defaultBindEmbeddingModel returns the tenant's default embedding model and its vector dimension. +func defaultBindEmbeddingModel(tenantID string) (*models.EmbeddingModel, error) { + modelSvc := service.NewModelProviderService() + driver, modelName, apiConfig, maxTokens, err := modelSvc.GetTenantDefaultModelByType(tenantID, entity.ModelTypeEmbedding) + if err != nil { + return nil, fmt.Errorf("bind default embedding model: %w", err) + } + model := models.NewEmbeddingModel(driver, &modelName, apiConfig, maxTokens) + return model, nil +} + +// getEmbeddingDimension encodes a test string to determine vector dimension. +func getEmbeddingDimension(model *models.EmbeddingModel) (int, error) { + embeds, err := model.ModelDriver.Embed(model.ModelName, []string{"ok"}, model.APIConfig, &models.EmbeddingConfig{Dimension: 0}) + if err != nil { + return 0, fmt.Errorf("test encode failed: %w", err) + } + if len(embeds) == 0 { + return 0, fmt.Errorf("test encode returned no embeddings") + } + return len(embeds[0].Embedding), nil +} diff --git a/internal/ingestion/task/task_handler_test.go b/internal/ingestion/task/task_handler_test.go new file mode 100644 index 0000000000..489726f0ee --- /dev/null +++ b/internal/ingestion/task/task_handler_test.go @@ -0,0 +1,232 @@ +package task + +import ( + "context" + "strings" + "testing" + + "ragflow/internal/entity" + "ragflow/internal/entity/models" +) + +func testStrPtr(s string) *string { return &s } + +func makeTaskHandlerTestContext(taskType string) *TaskContext { + var progressCalls []float64 + pipelineID := "" + if strings.HasPrefix(taskType, "dataflow") { + pipelineID = "flow-1" + } + return &TaskContext{ + IngestionTask: &entity.IngestionTask{ + ID: "task-1", + DocumentID: "doc-1", + }, + TaskType: taskType, + PipelineID: pipelineID, + Doc: entity.Document{ + ID: "doc-1", + KbID: "kb-1", + Name: testStrPtr("test-doc.pdf"), + ParserID: "naive", + ParserConfig: entity.JSONMap{}, + }, + KB: entity.Knowledgebase{ + ID: "kb-1", + TenantID: "tenant-1", + EmbdID: "embd-1", + }, + Tenant: entity.Tenant{ + ID: "tenant-1", + LLMID: "gpt-4", + }, + ProgressFunc: func(prog float64, msg string) { + progressCalls = append(progressCalls, prog) + }, + } +} + +func newNoopDataflowService(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { + if strings.TrimSpace(dataflowID) == "" { + dataflowID = "flow-1" + } + svc, err := NewDataflowService(ctx, dataflowID, 0, 0) + if err != nil { + return nil, err + } + svc = svc. + WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) { + return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, dataflowID, nil + }). + WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { + return map[string]any{ + "chunks": []map[string]any{ + {"text": "stub dataflow chunk", "q_2_vec": []float64{0.1, 0.2}}, + }, + }, dsl, nil + }). + WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + return nil, nil + }). + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { + return nil + }). + WithDocService(&stubDocService{}). + WithChunkCounter(&stubChunkCounter{}). + WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) { + return nil, nil + }) + return svc, nil +} + +func newNoopTaskHandler(ctx *TaskContext) *TaskHandler { + return NewTaskHandler(ctx).WithDataflowServiceFactory(newNoopDataflowService) +} + +func TestTaskHandler_Dispatch(t *testing.T) { + tests := []struct { + name string + taskType string + wantErr bool + wantPanic bool + }{ + {"memory", "memory", false, false}, + {"dataflow", "dataflow", false, false}, + {"dataflow with suffix", "dataflow_test", false, false}, + {"raptor", "raptor", false, false}, + {"graphrag", "graphrag", false, false}, + {"mindmap", "mindmap", false, false}, + {"evaluation", "evaluation", false, false}, + {"reembedding", "reembedding", false, false}, + {"clone", "clone", false, false}, + {"standard (empty task_type)", "", false, false}, + {"standard (unknown task_type)", "unknown_type", false, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := makeTaskHandlerTestContext(tt.taskType) + handler := newNoopTaskHandler(ctx) + err := handler.Handle() + + if tt.wantErr && err == nil { + t.Error("expected error, got nil") + } + if !tt.wantErr && err != nil { + t.Errorf("unexpected error: %v", err) + } + }) + } +} + +func TestTaskHandler_DefaultDataflowServiceInjectsProgress(t *testing.T) { + ctx := makeTaskHandlerTestContext("dataflow") + handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { + svc, err := NewDataflowService(ctx, dataflowID, 0, 0) + if err != nil { + t.Fatalf("NewDataflowService: %v", err) + } + if svc.progressFunc == nil { + t.Fatal("expected default progress func to be injected") + } + return newNoopDataflowService(ctx, dataflowID) + }) + if err := handler.Handle(); err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestTaskHandler_Dataflow_ShowsProgressAndPipelineLog(t *testing.T) { + ctx := makeTaskHandlerTestContext("dataflow") + ctx.Doc.PipelineID = testStrPtr("flow-1") + ctx.Doc.Name = testStrPtr("verify-dataflow.pdf") + + var pipelineCalled bool + var insertCalled bool + var logCreateCalls int + var insertedChunkCount int + var progressProgs []float64 + var progressMsgs []string + + ctx.ProgressFunc = func(prog float64, msg string) { + progressProgs = append(progressProgs, prog) + progressMsgs = append(progressMsgs, msg) + t.Logf("progress: prog=%.2f msg=%q", prog, msg) + } + + handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { + svc := mustNewDataflowService(t, ctx, dataflowID, 0, 0). + WithLoadDSLFunc(func(ctx context.Context, dataflowID string) (string, string, error) { + return `{"nodes":[{"id":"stub-node"}],"edges":[]}`, dataflowID, nil + }). + WithRunPipelineFunc(func(ctx context.Context, dsl string) (map[string]any, string, error) { + pipelineCalled = true + t.Log("mock pipeline.run called") + return map[string]any{ + "chunks": []map[string]any{ + {"text": "stub dataflow chunk", "q_2_vec": []float64{0.1, 0.2}}, + }, + }, dsl, nil + }). + WithInsertChunksFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + insertCalled = true + insertedChunkCount = len(chunks) + t.Logf("mock insertChunks called: %d chunks", len(chunks)) + return nil, nil + }). + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { + logCreateCalls++ + if log.PipelineID != nil { + t.Logf("mock pipeline log created: pipeline_id=%s", *log.PipelineID) + } + return nil + }). + WithDocService(&stubDocService{}). + WithChunkCounter(&stubChunkCounter{}). + WithGetEmbeddingModelFunc(func(tenantID, embdID string) (*models.EmbeddingModel, error) { + return nil, nil + }) + return svc, nil + }) + + if err := handler.Handle(); err != nil { + t.Fatalf("handler.Handle() error: %v", err) + } + + if !pipelineCalled { + t.Fatal("expected mock pipeline.run to be called") + } + if !insertCalled { + t.Fatal("expected insertChunks to be called") + } + if insertedChunkCount != 1 { + t.Fatalf("insertedChunkCount = %d, want 1", insertedChunkCount) + } + if len(progressProgs) == 0 { + t.Fatal("expected progress callbacks, got none") + } + + foundStartIndex := false + for _, msg := range progressMsgs { + if strings.Contains(msg, "Start to index") { + foundStartIndex = true + break + } + } + if !foundStartIndex { + t.Fatalf("expected progress message containing %q, got %v", "Start to index", progressMsgs) + } + + if got := progressProgs[len(progressProgs)-1]; got != 1.0 { + t.Fatalf("final progress = %v, want 1.0", got) + } + + lastMsg := progressMsgs[len(progressMsgs)-1] + if !strings.Contains(lastMsg, "Indexing done") { + t.Fatalf("final progress msg = %q, want substring %q", lastMsg, "Indexing done") + } + + if logCreateCalls != 1 { + t.Fatalf("logCreateCalls = %d, want 1", logCreateCalls) + } +} diff --git a/internal/ingestion/task/tool/compare_dataflow_golden.go b/internal/ingestion/task/tool/compare_dataflow_golden.go new file mode 100644 index 0000000000..a65124f8c4 --- /dev/null +++ b/internal/ingestion/task/tool/compare_dataflow_golden.go @@ -0,0 +1,288 @@ +package main + +import ( + "encoding/json" + "errors" + "flag" + "fmt" + "os" + "path/filepath" + "reflect" + "sort" + "strings" + + "ragflow/internal/ingestion/task" +) + +func main() { + caseDir := flag.String("case-dir", "", "case directory under internal/ingestion/task/testdata/") + docID := flag.String("doc-id", "doc-1", "document id") + kbID := flag.String("kb-id", "kb-1", "knowledge base id") + docName := flag.String("doc-name", "sample.md", "document name") + flag.Parse() + + if *caseDir == "" { + exitErr(errors.New("--case-dir is required")) + } + + inputPath := filepath.Join(*caseDir, "input.json") + outputDir := filepath.Join(*caseDir, "output") + + input := mustReadJSONMap(inputPath) + + expectedNormalized := mustReadJSONSlice(filepath.Join(outputDir, "normalized_chunks.json")) + expectedProcessed := mustReadJSONSlice(filepath.Join(outputDir, "processed_chunks.json")) + expectedMetadata := mustReadJSONMap(filepath.Join(outputDir, "merged_metadata.json")) + expectedProcessError, hasExpectedProcessError := readOptionalJSONMap(filepath.Join(outputDir, "process_error.json")) + + actual, actualErr := runActual(input, *docID, *kbID, *docName) + if hasExpectedProcessError || actualErr != nil { + reportProcessErrorOutcome(expectedProcessError, hasExpectedProcessError, actualErr) + return + } + + diffs := make([]string, 0) + diffs = append(diffs, diffValue("normalized_chunks", sanitizeNormalized(expectedNormalized), sanitizeNormalized(actual.NormalizedChunks), false)...) + diffs = append(diffs, diffValue("processed_chunks", sanitizeProcessed(expectedProcessed), sanitizeProcessed(actual.ProcessedChunks), true)...) + diffs = append(diffs, diffValue("merged_metadata", expectedMetadata, actual.MergedMetadata, false)...) + + if len(diffs) == 0 { + fmt.Println("validation succeeded") + return + } + + fmt.Println("validation failed") + for _, d := range diffs { + fmt.Println(d) + } + os.Exit(1) +} + +func runActual(input map[string]any, docID string, kbID string, docName string) (result task.GoldenDataflowResult, err error) { + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("panic: %v", r) + } + }() + return task.ProcessPipelineOutputForGolden(input, docID, kbID, docName), nil +} + +func reportProcessErrorOutcome(expected map[string]any, hasExpected bool, actualErr error) { + switch { + case hasExpected && actualErr == nil: + fmt.Println("validation failed") + fmt.Printf("expected process error: %v\n", expected) + fmt.Println("actual: processing succeeded") + os.Exit(1) + case !hasExpected && actualErr != nil: + fmt.Println("validation failed") + fmt.Println("expected: processing succeeded") + fmt.Printf("actual error: %v\n", actualErr) + os.Exit(1) + case hasExpected && actualErr != nil: + expectedMsg, _ := expected["message"].(string) + if expectedMsg != "" && strings.Contains(actualErr.Error(), expectedMsg) { + fmt.Println("validation succeeded") + fmt.Printf("expected process error observed: %s\n", expectedMsg) + return + } + fmt.Println("validation failed") + fmt.Printf("expected process error: %v\n", expected) + fmt.Printf("actual error: %v\n", actualErr) + os.Exit(1) + default: + return + } +} + +func sanitizeProcessed(chunks []map[string]any) []map[string]any { + out := make([]map[string]any, 0, len(chunks)) + for _, chunk := range chunks { + cp := normalizeComparableMap(cloneMap(chunk)) + delete(cp, "create_time") + delete(cp, "create_timestamp_flt") + for k := range cp { + if strings.HasPrefix(k, "q_") && strings.HasSuffix(k, "_vec") { + delete(cp, k) + } + } + out = append(out, cp) + } + return out +} + +func sanitizeNormalized(chunks []map[string]any) []map[string]any { + out := make([]map[string]any, 0, len(chunks)) + for _, chunk := range chunks { + out = append(out, normalizeComparableMap(cloneMap(chunk))) + } + return out +} + +func diffValue(path string, expected any, actual any, sortMaps bool) []string { + if reflect.DeepEqual(expected, actual) { + return nil + } + return walkDiff(path, expected, actual, sortMaps) +} + +func walkDiff(path string, expected any, actual any, sortMaps bool) []string { + switch exp := expected.(type) { + case map[string]any: + act, ok := actual.(map[string]any) + if !ok { + return []string{fmt.Sprintf("%s: expected map, got %T", path, actual)} + } + keys := unionKeys(exp, act) + diffs := make([]string, 0) + for _, key := range keys { + ev, eok := exp[key] + av, aok := act[key] + switch { + case !eok: + diffs = append(diffs, fmt.Sprintf("%s.%s: unexpected actual value %v", path, key, av)) + case !aok: + diffs = append(diffs, fmt.Sprintf("%s.%s: missing actual value, expected %v", path, key, ev)) + default: + diffs = append(diffs, walkDiff(path+"."+key, ev, av, sortMaps)...) + } + } + return diffs + case []any: + act, ok := actual.([]any) + if !ok { + return []string{fmt.Sprintf("%s: expected []any, got %T", path, actual)} + } + return diffSlice(path, exp, act, sortMaps) + case []map[string]any: + act, ok := actual.([]map[string]any) + if !ok { + return []string{fmt.Sprintf("%s: expected []map[string]any, got %T", path, actual)} + } + diffs := make([]string, 0) + if len(exp) != len(act) { + diffs = append(diffs, fmt.Sprintf("%s: len mismatch expected=%d actual=%d", path, len(exp), len(act))) + } + n := min(len(exp), len(act)) + for i := 0; i < n; i++ { + diffs = append(diffs, walkDiff(fmt.Sprintf("%s[%d]", path, i), exp[i], act[i], sortMaps)...) + } + return diffs + default: + if !reflect.DeepEqual(expected, actual) { + return []string{fmt.Sprintf("%s: expected=%v (%T), actual=%v (%T)", path, expected, expected, actual, actual)} + } + return nil + } +} + +func diffSlice(path string, expected []any, actual []any, sortMaps bool) []string { + diffs := make([]string, 0) + if len(expected) != len(actual) { + diffs = append(diffs, fmt.Sprintf("%s: len mismatch expected=%d actual=%d", path, len(expected), len(actual))) + } + n := min(len(expected), len(actual)) + for i := 0; i < n; i++ { + diffs = append(diffs, walkDiff(fmt.Sprintf("%s[%d]", path, i), expected[i], actual[i], sortMaps)...) + } + return diffs +} + +func unionKeys(a map[string]any, b map[string]any) []string { + set := make(map[string]struct{}, len(a)+len(b)) + for k := range a { + set[k] = struct{}{} + } + for k := range b { + set[k] = struct{}{} + } + keys := make([]string, 0, len(set)) + for k := range set { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} + +func cloneMap(in map[string]any) map[string]any { + out := make(map[string]any, len(in)) + for k, v := range in { + out[k] = v + } + return out +} + +func normalizeComparableMap(in map[string]any) map[string]any { + out := make(map[string]any, len(in)) + for k, v := range in { + out[k] = normalizeComparableValue(v) + } + return out +} + +func normalizeComparableValue(v any) any { + switch val := v.(type) { + case []string: + out := make([]any, 0, len(val)) + for _, item := range val { + out = append(out, item) + } + return out + case []map[string]any: + out := make([]any, 0, len(val)) + for _, item := range val { + out = append(out, normalizeComparableMap(item)) + } + return out + case map[string]any: + return normalizeComparableMap(val) + default: + return v + } +} + +func mustReadJSONMap(path string) map[string]any { + var out map[string]any + mustReadJSON(path, &out) + return out +} + +func mustReadJSONSlice(path string) []map[string]any { + var out []map[string]any + mustReadJSON(path, &out) + return out +} + +func mustReadJSON(path string, out any) { + data, err := os.ReadFile(path) + if err != nil { + exitErr(err) + } + if err := json.Unmarshal(data, out); err != nil { + exitErr(fmt.Errorf("unmarshal %s: %w", path, err)) + } +} + +func readOptionalJSONMap(path string) (map[string]any, bool) { + data, err := os.ReadFile(path) + if err != nil { + return nil, false + } + var out map[string]any + if err := json.Unmarshal(data, &out); err != nil { + exitErr(fmt.Errorf("unmarshal %s: %w", path, err)) + } + return out, true +} + +func min(a, b int) int { + if a < b { + return a + } + return b +} + +func exitErr(err error) { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) +} diff --git a/internal/ingestion/testutil/helpers.go b/internal/ingestion/testutil/helpers.go new file mode 100644 index 0000000000..babf34c8e0 --- /dev/null +++ b/internal/ingestion/testutil/helpers.go @@ -0,0 +1,246 @@ +// +// 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 testutil + +import ( + "testing" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + + "github.com/glebarez/sqlite" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +// ────────────────────────────────────────────────────────── +// Simple Helper Functions +// ────────────────────────────────────────────────────────── + +// StrPtr returns a pointer to the given string. +func StrPtr(s string) *string { + return &s +} + +// ────────────────────────────────────────────────────────── +// Database Helpers +// ────────────────────────────────────────────────────────── + +// SetupTestDB sets up an in-memory SQLite database for testing. +// It auto-migrates the given tables (or all common tables if none provided). +func SetupTestDB(t *testing.T, tables ...any) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ + Logger: logger.Default.LogMode(logger.Silent), + }) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + + if len(tables) == 0 { + tables = []any{ + &entity.IngestionTask{}, + &entity.IngestionTaskLog{}, + &entity.Task{}, + &entity.Document{}, + &entity.Knowledgebase{}, + &entity.Tenant{}, + &entity.File{}, + &entity.File2Document{}, + } + } + + if err := db.AutoMigrate(tables...); err != nil { + t.Fatalf("auto-migrate: %v", err) + } + return db +} + +// ReplaceDBForTest replaces dao.DB with the given test database +// and returns a cleanup function to restore the original. +func ReplaceDBForTest(t *testing.T, db *gorm.DB) func() { + t.Helper() + origDB := dao.DB + dao.DB = db + return func() { dao.DB = origDB } +} + +// ────────────────────────────────────────────────────────── +// Test Data Configuration (Options Pattern) +// ────────────────────────────────────────────────────────── + +// TestDataOption configures the test data setup. +type TestDataOption func(*testDataConfig) + +type testDataConfig struct { + tenantID string + kbID string + docID string + taskID string + pipelineID *string + docName string + docType string +} + +// WithTenantID sets the tenant ID for test data. +func WithTenantID(id string) TestDataOption { + return func(c *testDataConfig) { c.tenantID = id } +} + +// WithKBID sets the knowledgebase ID for test data. +func WithKBID(id string) TestDataOption { + return func(c *testDataConfig) { c.kbID = id } +} + +// WithDocID sets the document ID for test data. +func WithDocID(id string) TestDataOption { + return func(c *testDataConfig) { c.docID = id } +} + +// WithTaskID sets the task ID for test data. +func WithTaskID(id string) TestDataOption { + return func(c *testDataConfig) { c.taskID = id } +} + +// WithPipelineID sets the pipeline ID for test data. +func WithPipelineID(id string) TestDataOption { + return func(c *testDataConfig) { + pid := id + c.pipelineID = &pid + } +} + +// WithDocName sets the document name for test data. +func WithDocName(name string) TestDataOption { + return func(c *testDataConfig) { c.docName = name } +} + +// WithDocType sets the document type for test data. +func WithDocType(docType string) TestDataOption { + return func(c *testDataConfig) { c.docType = docType } +} + +// SeedTestData seeds the database with test data. +// Returns: tenantID, kbID, docID, taskID +func SeedTestData(t *testing.T, db *gorm.DB, opts ...TestDataOption) (string, string, string, string) { + t.Helper() + cfg := &testDataConfig{ + tenantID: "tenant-1", + kbID: "kb-1", + docID: "doc-1", + taskID: "task-1", + docName: "test.pdf", + docType: "pdf", + } + for _, opt := range opts { + opt(cfg) + } + + // Create Tenant + if err := db.Create(&entity.Tenant{ + ID: cfg.tenantID, + LLMID: "gpt-4", + Status: StrPtr("1"), + }).Error; err != nil { + t.Fatalf("create tenant: %v", err) + } + + // Create Knowledgebase + if err := db.Create(&entity.Knowledgebase{ + ID: cfg.kbID, + TenantID: cfg.tenantID, + EmbdID: "embd-1", + Status: StrPtr("1"), + ParserConfig: entity.JSONMap{}, + }).Error; err != nil { + t.Fatalf("create kb: %v", err) + } + + // Create Document + doc := &entity.Document{ + ID: cfg.docID, + KbID: cfg.kbID, + Name: &cfg.docName, + ParserID: "naive", + ParserConfig: entity.JSONMap{}, + PipelineID: cfg.pipelineID, + Status: StrPtr("1"), + Type: cfg.docType, + } + loc := "doc_store/" + cfg.docID + doc.Location = &loc + if err := db.Create(doc).Error; err != nil { + t.Fatalf("create doc: %v", err) + } + + // Create Task + if err := db.Create(&entity.Task{ + ID: cfg.taskID, + DocID: cfg.docID, + TaskType: "dataflow", + FromPage: 0, + ToPage: 100000, + }).Error; err != nil { + t.Fatalf("create task: %v", err) + } + + // Create IngestionTask + if err := db.Create(&entity.IngestionTask{ + ID: cfg.taskID, + UserID: "u1", + DocumentID: cfg.docID, + DatasetID: cfg.kbID, + Status: common.RUNNING, + }).Error; err != nil { + t.Fatalf("create ingestion task: %v", err) + } + + return cfg.tenantID, cfg.kbID, cfg.docID, cfg.taskID +} + +// ────────────────────────────────────────────────────────── +// Test Entity Helpers +// ────────────────────────────────────────────────────────── + +// TestTask creates a test IngestionTask. +func TestTask(docID string) *entity.IngestionTask { + return &entity.IngestionTask{ + ID: "task-1", + DocumentID: docID, + DatasetID: "kb-1", + } +} + +// TestDoc creates a test Document. +func TestDoc(id string, docType string, suffix string, kbID ...string) *entity.Document { + name := "test." + suffix + loc := "doc_store/" + id + kb := "kb-1" + if len(kbID) > 0 { + kb = kbID[0] + } + return &entity.Document{ + ID: id, + KbID: kb, + Type: docType, + Suffix: suffix, + Name: &name, + Location: &loc, + ParserID: "naive", + } +} diff --git a/internal/ingestion/testutil/mocks.go b/internal/ingestion/testutil/mocks.go new file mode 100644 index 0000000000..75ee1a9d6a --- /dev/null +++ b/internal/ingestion/testutil/mocks.go @@ -0,0 +1,174 @@ +// +// 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 testutil + +import ( + "context" + "errors" + "time" + + doctype "ragflow/internal/deepdoc/parser/type" + "ragflow/internal/entity" +) + +// ────────────────────────────────────────────────────────── +// Mock DAO Implementations +// ────────────────────────────────────────────────────────── + +// MockDocDAO is a mock implementation of docGetter interface. +type MockDocDAO struct { + Docs map[string]*entity.Document + Err error +} + +func (m *MockDocDAO) GetByID(id string) (*entity.Document, error) { + if m.Err != nil { + return nil, m.Err + } + d, ok := m.Docs[id] + if !ok { + return nil, errors.New("record not found") + } + return d, nil +} + +// MockF2dDAO is a mock implementation of f2dGetter interface. +type MockF2dDAO struct { + Mappings map[string][]*entity.File2Document + Err error +} + +func (m *MockF2dDAO) GetByDocumentID(docID string) ([]*entity.File2Document, error) { + if m.Err != nil { + return nil, m.Err + } + return m.Mappings[docID], nil +} + +// MockFileDAO is a mock implementation of fileGetter interface. +type MockFileDAO struct { + Files map[string]*entity.File + Err error +} + +func (m *MockFileDAO) GetByID(id string) (*entity.File, error) { + if m.Err != nil { + return nil, m.Err + } + f, ok := m.Files[id] + if !ok { + return nil, errors.New("not found") + } + return f, nil +} + +// ────────────────────────────────────────────────────────── +// Mock Storage Implementation +// ────────────────────────────────────────────────────────── + +// MockStorage is a complete mock implementation of storage.Storage interface. +type MockStorage struct { + Data map[string][]byte + Err error + PutCallCount int + PutBucket string + PutFnm string + PutData []byte +} + +// NewMockStorage creates a new MockStorage with the given data. +func NewMockStorage(data map[string][]byte) *MockStorage { + return &MockStorage{Data: data} +} + +// NewMockStorageWithError creates a new MockStorage that returns the given error. +func NewMockStorageWithError(err error) *MockStorage { + return &MockStorage{Err: err} +} + +func (m *MockStorage) Health() bool { return true } + +func (m *MockStorage) Get(bucket, objectName string, tenantID ...string) ([]byte, error) { + if m.Err != nil { + return nil, m.Err + } + key := bucket + "/" + objectName + if len(tenantID) > 0 && tenantID[0] != "" { + key = tenantID[0] + "/" + key + } + d, ok := m.Data[key] + if !ok { + d, ok = m.Data[objectName] + } + return d, nil +} + +func (m *MockStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error { + if m.Err != nil { + return m.Err + } + m.PutCallCount++ + m.PutBucket = bucket + m.PutFnm = fnm + m.PutData = binary + return nil +} + +func (m *MockStorage) Remove(bucket, fnm string, tenantID ...string) error { return nil } + +func (m *MockStorage) ObjExist(bucket, fnm string, tenantID ...string) bool { return true } + +func (m *MockStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) { + return "", nil +} + +func (m *MockStorage) BucketExists(bucket string) bool { return true } + +func (m *MockStorage) RemoveBucket(bucket string) error { return nil } + +func (m *MockStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool { return true } + +func (m *MockStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool { return true } + +func (m *MockStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) { + var objects []string + for key := range m.Data { + objects = append(objects, key) + } + return objects, nil +} + +func (m *MockStorage) Close() error { return nil } + +// ────────────────────────────────────────────────────────── +// Mock PDF Parser Implementation +// ────────────────────────────────────────────────────────── + +// MockPDFParser is a mock implementation of pdfParser interface. +type MockPDFParser struct { + Sections []map[string]any + Err error +} + +// NewMockPDFParser creates a new MockPDFParser. +func NewMockPDFParser(sections []map[string]any, err error) *MockPDFParser { + return &MockPDFParser{Sections: sections, Err: err} +} + +func (m *MockPDFParser) ParseWithDeepDoc(ctx context.Context, filename string, data []byte, config doctype.ParserConfig) ([]map[string]any, error) { + return m.Sections, m.Err +} diff --git a/internal/parser/parser/factory.go b/internal/parser/parser/factory.go new file mode 100644 index 0000000000..b267921eba --- /dev/null +++ b/internal/parser/parser/factory.go @@ -0,0 +1,86 @@ +// +// 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 parser + +import "fmt" + +// stubParser is a placeholder for parsers whose real implementation lives +// in the DeepDoc inference service. The stub exists only to satisfy the +// factory dispatch contract. +type stubParser struct{ name string } + +func (p *stubParser) ParseWithResult(filename string, data []byte) ParseResult { + return ParseResult{Err: fmt.Errorf("%s: remote deepdoc dispatch not yet wired", p.name)} +} + +// GetParserByID returns a parser for the given parser_id (matches Python FACTORY dispatch). +// Most parser_ids handle PDF files with different strategies; the strategy is encoded +// in the returned parser implementation. Stub implementations are replaced as +// colleagues deliver real parsers. +func GetParserByID(parserID string) (ParseResultProducer, error) { + switch parserID { + case "naive", "general": + return NewNaivePDFParser(), nil + case "paper": + return NewPaperPDFParser(), nil + case "book": + return NewBookPDFParser(), nil + case "presentation": + return NewPresentationParser(), nil + case "manual": + return NewManualPDFParser(), nil + case "laws": + return NewLawsPDFParser(), nil + case "qa": + return NewQAPDFParser(), nil + case "table": + return NewTableParser(), nil + case "resume": + return NewResumePDFParser(), nil + case "picture": + return NewPicturePDFParser(), nil + case "one": + return NewOnePDFParser(), nil + case "audio": + return NewAudioParser(), nil + case "email": + return NewEmailParser(), nil + case "tag": + return NewTagPDFParser(), nil + case "knowledge_graph": + return NewKGPDFParser(), nil + default: + return nil, fmt.Errorf("unknown parser_id: %s", parserID) + } +} + +// Stub constructors for each parser type. +func NewNaivePDFParser() *stubParser { return &stubParser{name: "naive"} } +func NewPaperPDFParser() *stubParser { return &stubParser{name: "paper"} } +func NewBookPDFParser() *stubParser { return &stubParser{name: "book"} } +func NewPresentationParser() *stubParser { return &stubParser{name: "presentation"} } +func NewManualPDFParser() *stubParser { return &stubParser{name: "manual"} } +func NewLawsPDFParser() *stubParser { return &stubParser{name: "laws"} } +func NewQAPDFParser() *stubParser { return &stubParser{name: "qa"} } +func NewTableParser() *stubParser { return &stubParser{name: "table"} } +func NewResumePDFParser() *stubParser { return &stubParser{name: "resume"} } +func NewPicturePDFParser() *stubParser { return &stubParser{name: "picture"} } +func NewOnePDFParser() *stubParser { return &stubParser{name: "one"} } +func NewAudioParser() *stubParser { return &stubParser{name: "audio"} } +func NewEmailParser() *stubParser { return &stubParser{name: "email"} } +func NewTagPDFParser() *stubParser { return &stubParser{name: "tag"} } +func NewKGPDFParser() *stubParser { return &stubParser{name: "knowledge_graph"} } diff --git a/internal/parser/parser/factory_test.go b/internal/parser/parser/factory_test.go new file mode 100644 index 0000000000..b2aa532889 --- /dev/null +++ b/internal/parser/parser/factory_test.go @@ -0,0 +1,81 @@ +// +// 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 parser + +import ( + "testing" + + "ragflow/internal/entity" +) + +func TestGetParserByID_AllKnownIDs(t *testing.T) { + tests := []struct { + parserID string + wantNil bool + }{ + // PDF-based parsers + {string(entity.ParserTypeNaive), false}, + {string(entity.ParserTypePaper), false}, + {string(entity.ParserTypeBook), false}, + {string(entity.ParserTypeManual), false}, + {string(entity.ParserTypeLaws), false}, + {string(entity.ParserTypeQA), false}, + {string(entity.ParserTypeResume), false}, + {string(entity.ParserTypePicture), false}, + {string(entity.ParserTypeOne), false}, + {string(entity.ParserTypeKG), false}, + {string(entity.ParserTypeTag), false}, + // Office parsers + {string(entity.ParserTypePresentation), false}, + {string(entity.ParserTypeTable), false}, + // Special parsers + {string(entity.ParserTypeAudio), false}, + {string(entity.ParserTypeEmail), false}, + } + + for _, tt := range tests { + t.Run(tt.parserID, func(t *testing.T) { + got, err := GetParserByID(tt.parserID) + if tt.wantNil { + if got != nil { + t.Errorf("GetParserByID(%q) = non-nil, want nil", tt.parserID) + } + } else { + if err != nil { + t.Errorf("GetParserByID(%q): %v", tt.parserID, err) + } + if got == nil { + t.Errorf("GetParserByID(%q) = nil, want non-nil", tt.parserID) + } + } + }) + } +} + +func TestGetParserByID_InvalidID(t *testing.T) { + _, err := GetParserByID("nonexistent") + if err == nil { + t.Fatal("expected error for invalid parser_id") + } +} + +func TestGetParserByID_EmptyID(t *testing.T) { + _, err := GetParserByID("") + if err == nil { + t.Fatal("expected error for empty parser_id") + } +} diff --git a/internal/service/document.go b/internal/service/document.go index c68c1425ed..96ce798637 100644 --- a/internal/service/document.go +++ b/internal/service/document.go @@ -110,6 +110,8 @@ type UpdateDocumentRequest struct { ChunkNum *int64 `json:"chunk_num"` Progress *float64 `json:"progress"` ProgressMsg *string `json:"progress_msg"` + // FIXME: need to confirm below field + ProcessDuration *float64 `json:"progress_duration"` } // DocumentResponse document response @@ -617,6 +619,34 @@ func (s *DocumentService) UpdateDocument(id string, req *UpdateDocumentRequest) return s.documentDAO.Update(document) } +// IncrementChunkNum atomically increments chunk/token counters on the document and its knowledge base in a transaction +func (s *DocumentService) IncrementChunkNum(docID, kbID string, chunkNum, tokenNum int, duration float64) error { + return dao.DB.Transaction(func(tx *gorm.DB) error { + // Update document + if err := tx.Model(&entity.Document{}). + Where("id = ? AND kb_id = ?", docID, kbID). + Updates(map[string]interface{}{ + "chunk_num": gorm.Expr("chunk_num + ?", int64(chunkNum)), + "token_num": gorm.Expr("token_num + ?", int64(tokenNum)), + "process_duration": gorm.Expr("process_duration + ?", duration), + }).Error; err != nil { + return err + } + + // Update knowledgebase + if err := tx.Model(&entity.Knowledgebase{}). + Where("id = ?", kbID). + Updates(map[string]interface{}{ + "chunk_num": gorm.Expr("chunk_num + ?", int64(chunkNum)), + "token_num": gorm.Expr("token_num + ?", int64(tokenNum)), + }).Error; err != nil { + return err + } + + return nil + }) +} + // DeleteDocument delete document — delegates to full cleanup logic. func (s *DocumentService) DeleteDocument(id string) error { return s.deleteDocumentFull(id) @@ -1278,16 +1308,16 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com for _, docID := range req.DocIDs { doc := docsByID[docID] if doc == nil { - return common.CodeDataError, fmt.Errorf("Document not found!") + return common.CodeDataError, fmt.Errorf("document not found") } kb, err := s.kbDAO.GetByID(doc.KbID) if err != nil { - return common.CodeDataError, fmt.Errorf("Tenant not found!") + return common.CodeDataError, fmt.Errorf("tenant not found") } if !s.kbDAO.Accessible(kb.ID, userID) { - return common.CodeAuthenticationError, fmt.Errorf("No authorization.") + return common.CodeAuthenticationError, fmt.Errorf("no authorization") } updates := map[string]interface{}{ @@ -1304,30 +1334,35 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com if run == string(entity.TaskStatusCancel) { if err := s.cancelDocParse(doc); err != nil { + common.Error(fmt.Sprintf("go side, start to process %s, run is cancel", docID), err) return common.CodeDataError, err } } if rerunWithDelete { if err := s.prepareDocumentRerunWithDelete(doc, kb.TenantID); err != nil { + common.Error(fmt.Sprintf("go side, start to process %s, error when rerun with delete", docID), err) return common.CodeExceptionError, err } } if err := s.documentDAO.UpdateByID(doc.ID, updates); err != nil { + common.Error(fmt.Sprintf("go side, doc %s, UpdateByID failed", docID), err) return common.CodeExceptionError, err } if req.Delete && !rerunWithDelete { - _, _ = s.taskDAO.DeleteByDocIDs([]string{doc.ID}) + _, _ = s.taskDAO.DeleteIngestionTasksByDocIDs([]string{doc.ID}) indexName := fmt.Sprintf("ragflow_%s", kb.TenantID) if s.docEngine != nil { exists, err := s.docEngine.ChunkStoreExists(context.Background(), indexName, doc.KbID) if err != nil { + common.Error(fmt.Sprintf("go side, doc %s, ChunkStoreExists failed", docID), err) return common.CodeExceptionError, err } if exists { if _, err := s.docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": doc.ID}, indexName, doc.KbID); err != nil { + common.Error(fmt.Sprintf("go side, doc %s, DeleteChunks failed", docID), err) return common.CodeExceptionError, err } } @@ -1380,16 +1415,24 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com } } if _, err := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); err != nil { + common.Error(fmt.Sprintf("go side, doc %s, DeleteByDocIDs", docID), err) return common.CodeExceptionError, err } - bucket, objectName, err := s.GetDocumentStorageAddress(doc) + _, _, err := s.GetDocumentStorageAddress(doc) if err != nil { + common.Error(fmt.Sprintf("go side, doc %s, GetDocumentStorageAddress", docID), err) return common.CodeExceptionError, err } - if err := s.queueDocumentParseTasks(doc, bucket, objectName, 0); err != nil { + + common.Info("go side, before insert document") + if _, err := s.IngestDocuments(doc.KbID, userID, []string{doc.ID}); err != nil { + common.Error(fmt.Sprintf("go side, doc %s, IngestDocuments", docID), err) return common.CodeExceptionError, err } + common.Info("go side, after insert document") + if err := s.beginDocumentParse(doc.ID); err != nil { + common.Error(fmt.Sprintf("go side, doc %s, beginDocumentParse", docID), err) return common.CodeExceptionError, err } } @@ -1500,29 +1543,6 @@ func (s *DocumentService) countDoneDocuments(datasetID string) (int64, error) { return count, err } -func (s *DocumentService) queueDocumentParseTasks(doc *entity.Document, bucket, objectName string, priority int64) error { - if _, err := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); err != nil { - return err - } - tasks, err := s.newDocumentParseTasks(doc, bucket, objectName, priority) - if err != nil { - return err - } - if err := s.taskDAO.CreateMany(tasks); err != nil { - return err - } - queueName := documentParseQueueName(doc, priority) - for _, task := range tasks { - if task.Progress >= 1 { - continue - } - if redisClient := redis.Get(); redisClient == nil || !redisClient.QueueProduct(queueName, documentTaskMessage(task)) { - return fmt.Errorf("Can't access Redis. Please check the Redis' status.") - } - } - return nil -} - func (s *DocumentService) queueDocumentDataflowTask(kb *entity.Knowledgebase, doc *entity.Document, flowID string, priority int64) error { if _, err := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); err != nil { return err @@ -2063,6 +2083,7 @@ func (s *DocumentService) validateDocsInDataset(docIDs []string, datasetID strin func (s *DocumentService) cancelDocParse(doc *entity.Document) error { tasks, taskErr := s.taskDAO.GetByDocID(doc.ID) if taskErr != nil { + common.Error(fmt.Sprintf("error when load task %s", doc.ID), taskErr) return fmt.Errorf("failed to get tasks for %s: %v", doc.ID, taskErr) } diff --git a/internal/utility/metadata.go b/internal/utility/metadata.go new file mode 100644 index 0000000000..2e68dfea47 --- /dev/null +++ b/internal/utility/metadata.go @@ -0,0 +1,116 @@ +// +// 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 utility + +// UpdateMetadataTo merges metadata into an existing metadata map. +// Only string and []string values are accepted. Existing keys are preserved +// (not overwritten). List values are merged and deduplicated. +// Mirrors Python: common.metadata_utils.update_metadata_to() +func UpdateMetadataTo(target map[string]any, meta any) map[string]any { + if target == nil { + return nil + } + if meta == nil { + return target + } + + metaMap, ok := meta.(map[string]any) + if !ok { + return target + } + + for k, v := range metaMap { + normVal := normalizeMetaValue(v) + if normVal == nil { + continue + } + + existing, exists := target[k] + if !exists { + target[k] = normVal + continue + } + + // Merge with existing value + existStr, existIsStr := existing.(string) + newStr, newIsStr := normVal.(string) + existList, existIsList := existing.([]string) + newList, newIsList := normVal.([]string) + + if existIsStr && newIsStr { + // Both strings: convert to list, append + target[k] = dedupeStrings(append([]string{existStr}, newStr)) + } else if existIsStr && newIsList { + target[k] = dedupeStrings(append([]string{existStr}, newList...)) + } else if existIsList && newIsStr { + target[k] = dedupeStrings(append(existList, newStr)) + } else if existIsList && newIsList { + target[k] = dedupeStrings(append(existList, newList...)) + } + } + + return target +} + +// normalizeMetaValue normalizes a metadata value. +// Returns a string, []string, or nil if the value is not acceptable. +func normalizeMetaValue(v any) any { + switch val := v.(type) { + case string: + if val == "" { + return nil + } + return val + case []string: + filtered := make([]string, 0, len(val)) + for _, s := range val { + if s != "" { + filtered = append(filtered, s) + } + } + if len(filtered) == 0 { + return nil + } + return dedupeStrings(filtered) + case []any: + filtered := make([]string, 0, len(val)) + for _, elem := range val { + if s, ok := elem.(string); ok && s != "" { + filtered = append(filtered, s) + } + } + if len(filtered) == 0 { + return nil + } + return dedupeStrings(filtered) + default: + return nil + } +} + +// dedupeStrings removes duplicates while preserving order. +func dedupeStrings(input []string) []string { + seen := make(map[string]struct{}, len(input)) + out := make([]string, 0, len(input)) + for _, s := range input { + if _, ok := seen[s]; !ok { + seen[s] = struct{}{} + out = append(out, s) + } + } + return out +} diff --git a/internal/utility/metadata_test.go b/internal/utility/metadata_test.go new file mode 100644 index 0000000000..9d96f88082 --- /dev/null +++ b/internal/utility/metadata_test.go @@ -0,0 +1,115 @@ +package utility + +import ( + "testing" +) + +// ============================================================================= +// UpdateMetadataTo +// ============================================================================= + +func TestUpdateMetadataTo_EmptyMeta(t *testing.T) { + result := UpdateMetadataTo(map[string]any{"a": 1}, nil) + if result["a"] != 1 { + t.Errorf("a = %v, want 1", result["a"]) + } +} + +func TestUpdateMetadataTo_NewKeysAdded(t *testing.T) { + result := UpdateMetadataTo(map[string]any{"a": 1}, map[string]any{"b": "x"}) + if result["a"] != 1 { + t.Errorf("a should be preserved") + } + if result["b"] != "x" { + t.Errorf("b should be added") + } +} + +func TestUpdateMetadataTo_ExistingKeyMergedToList(t *testing.T) { + result := UpdateMetadataTo(map[string]any{"author": "Alice"}, map[string]any{"author": "Bob"}) + tags, ok := result["author"].([]string) + if !ok { + t.Fatalf("author should be []string, got %T", result["author"]) + } + if len(tags) != 2 || tags[0] != "Alice" || tags[1] != "Bob" { + t.Errorf("author = %v, want [Alice Bob]", tags) + } +} + +func TestUpdateMetadataTo_StringValuesOnly(t *testing.T) { + result := UpdateMetadataTo(map[string]any{}, map[string]any{"num": 42, "ok": true}) + if _, exists := result["num"]; exists { + t.Errorf("numeric values should be skipped") + } + if _, exists := result["ok"]; exists { + t.Errorf("bool values should be skipped") + } +} + +func TestUpdateMetadataTo_ListAppend(t *testing.T) { + result := UpdateMetadataTo( + map[string]any{"tags": []string{"a", "b"}}, + map[string]any{"tags": []string{"c"}}, + ) + tags := result["tags"].([]string) + if len(tags) != 3 { + t.Errorf("tags should have 3 elements, got %v", tags) + } +} + +func TestUpdateMetadataTo_StringToListMerge(t *testing.T) { + result := UpdateMetadataTo( + map[string]any{"tags": "a"}, + map[string]any{"tags": "b"}, + ) + tags := result["tags"].([]string) + if len(tags) != 2 || tags[0] != "a" || tags[1] != "b" { + t.Errorf("tags = %v, want [a b]", tags) + } +} + +func TestUpdateMetadataTo_DeduplicateList(t *testing.T) { + result := UpdateMetadataTo( + map[string]any{"tags": "a"}, + map[string]any{"tags": "a"}, + ) + tags := result["tags"].([]string) + if len(tags) != 1 { + t.Errorf("tags should be deduplicated: got %v", tags) + } +} + +func TestUpdateMetadataTo_NonDictMeta(t *testing.T) { + result := UpdateMetadataTo(map[string]any{"a": 1}, "not a dict") + if result["a"] != 1 { + t.Errorf("original should be unchanged for non-dict meta") + } +} + +func TestUpdateMetadataTo_EmptyInitial(t *testing.T) { + result := UpdateMetadataTo(nil, map[string]any{"k": "v"}) + if len(result) != 0 { + t.Errorf("should return empty when initial is nil") + } +} + +func TestUpdateMetadataTo_FilterEmptyStrings(t *testing.T) { + result := UpdateMetadataTo( + map[string]any{}, + map[string]any{"tags": []any{"a", "", "b"}}, + ) + tags := result["tags"].([]string) + if len(tags) != 2 { + t.Errorf("empty strings should be filtered: got %v", tags) + } +} + +func TestUpdateMetadataTo_SkipOnlyEmptyStringsList(t *testing.T) { + result := UpdateMetadataTo( + map[string]any{}, + map[string]any{"tags": []any{"", ""}}, + ) + if _, exists := result["tags"]; exists { + t.Error("all-empty list should be skipped") + } +}