mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-12 22:55:45 +08:00
Feature: task executor migration to go (#16549)
### Summary Feature: Integrate parser
This commit is contained in:
4
.gitignore
vendored
4
.gitignore
vendored
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
/*
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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{})
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
463
internal/ingestion/pipeline/real_storage_integration_test.go
Normal file
463
internal/ingestion/pipeline/real_storage_integration_test.go
Normal file
@@ -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
|
||||
}
|
||||
@@ -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]
|
||||
|
||||
72
internal/ingestion/service/execute_task_test.go
Normal file
72
internal/ingestion/service/execute_task_test.go
Normal file
@@ -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")
|
||||
}
|
||||
}
|
||||
404
internal/ingestion/service/ingestion_service.go
Normal file
404
internal/ingestion/service/ingestion_service.go
Normal file
@@ -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")
|
||||
}
|
||||
138
internal/ingestion/service/real_consumer_dataflow_test.go
Normal file
138
internal/ingestion/service/real_consumer_dataflow_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
72
internal/ingestion/service/test_helpers.go
Normal file
72
internal/ingestion/service/test_helpers.go
Normal file
@@ -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)
|
||||
}
|
||||
168
internal/ingestion/task/chunk_builder.go
Normal file
168
internal/ingestion/task/chunk_builder.go
Normal file
@@ -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)
|
||||
}
|
||||
461
internal/ingestion/task/chunk_builder_test.go
Normal file
461
internal/ingestion/task/chunk_builder_test.go
Normal file
@@ -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")
|
||||
}
|
||||
}
|
||||
246
internal/ingestion/task/chunk_process.go
Normal file
246
internal/ingestion/task/chunk_process.go
Normal file
@@ -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")
|
||||
}
|
||||
400
internal/ingestion/task/chunk_process_test.go
Normal file
400
internal/ingestion/task/chunk_process_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
164
internal/ingestion/task/chunk_utils.go
Normal file
164
internal/ingestion/task/chunk_utils.go
Normal file
@@ -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
|
||||
}
|
||||
427
internal/ingestion/task/chunk_utils_test.go
Normal file
427
internal/ingestion/task/chunk_utils_test.go
Normal file
@@ -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": "<p>Hello</p>",
|
||||
}
|
||||
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 != "<p>Hello</p>" {
|
||||
t.Errorf("text = %q, want \"<p>Hello</p>\"", 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
|
||||
}
|
||||
29
internal/ingestion/task/constants.go
Normal file
29
internal/ingestion/task/constants.go
Normal file
@@ -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"
|
||||
)
|
||||
14
internal/ingestion/task/dataflow_chunk_assets.go
Normal file
14
internal/ingestion/task/dataflow_chunk_assets.go
Normal file
@@ -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
|
||||
}
|
||||
370
internal/ingestion/task/dataflow_e2e_test.go
Normal file
370
internal/ingestion/task/dataflow_e2e_test.go
Normal file
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
597
internal/ingestion/task/dataflow_service.go
Normal file
597
internal/ingestion/task/dataflow_service.go
Normal file
@@ -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
|
||||
}
|
||||
826
internal/ingestion/task/dataflow_service_test.go
Normal file
826
internal/ingestion/task/dataflow_service_test.go
Normal file
@@ -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)
|
||||
)
|
||||
226
internal/ingestion/task/embedding_test.go
Normal file
226
internal/ingestion/task/embedding_test.go
Normal file
@@ -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
|
||||
}
|
||||
55
internal/ingestion/task/embedding_token_test.go
Normal file
55
internal/ingestion/task/embedding_token_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
54
internal/ingestion/task/golden_compare.go
Normal file
54
internal/ingestion/task/golden_compare.go
Normal file
@@ -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,
|
||||
}
|
||||
}
|
||||
71
internal/ingestion/task/golden_compare_test.go
Normal file
71
internal/ingestion/task/golden_compare_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
74
internal/ingestion/task/has_vectors_test.go
Normal file
74
internal/ingestion/task/has_vectors_test.go
Normal file
@@ -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")
|
||||
}
|
||||
}
|
||||
269
internal/ingestion/task/naive_merge.go
Normal file
269
internal/ingestion/task/naive_merge.go
Normal file
@@ -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)
|
||||
}
|
||||
417
internal/ingestion/task/naive_merge_test.go
Normal file
417
internal/ingestion/task/naive_merge_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
178
internal/ingestion/task/pipeline_run_output_checklist.md
Normal file
178
internal/ingestion/task/pipeline_run_output_checklist.md
Normal file
@@ -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_<dim>_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_<dim>_vec`
|
||||
- `embedding_token_consumption`
|
||||
52
internal/ingestion/task/position.go
Normal file
52
internal/ingestion/task/position.go
Normal file
@@ -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
|
||||
}
|
||||
92
internal/ingestion/task/position_test.go
Normal file
92
internal/ingestion/task/position_test.go
Normal file
@@ -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])
|
||||
}
|
||||
}
|
||||
170
internal/ingestion/task/real_consumer_test.go
Normal file
170
internal/ingestion/task/real_consumer_test.go
Normal file
@@ -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)
|
||||
}
|
||||
190
internal/ingestion/task/task_context.go
Normal file
190
internal/ingestion/task/task_context.go
Normal file
@@ -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
|
||||
}
|
||||
187
internal/ingestion/task/task_context_test.go
Normal file
187
internal/ingestion/task/task_context_test.go
Normal file
@@ -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")
|
||||
}
|
||||
}
|
||||
167
internal/ingestion/task/task_handler.go
Normal file
167
internal/ingestion/task/task_handler.go
Normal file
@@ -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
|
||||
}
|
||||
232
internal/ingestion/task/task_handler_test.go
Normal file
232
internal/ingestion/task/task_handler_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
288
internal/ingestion/task/tool/compare_dataflow_golden.go
Normal file
288
internal/ingestion/task/tool/compare_dataflow_golden.go
Normal file
@@ -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/<case_id>")
|
||||
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)
|
||||
}
|
||||
246
internal/ingestion/testutil/helpers.go
Normal file
246
internal/ingestion/testutil/helpers.go
Normal file
@@ -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",
|
||||
}
|
||||
}
|
||||
174
internal/ingestion/testutil/mocks.go
Normal file
174
internal/ingestion/testutil/mocks.go
Normal file
@@ -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
|
||||
}
|
||||
86
internal/parser/parser/factory.go
Normal file
86
internal/parser/parser/factory.go
Normal file
@@ -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"} }
|
||||
81
internal/parser/parser/factory_test.go
Normal file
81
internal/parser/parser/factory_test.go
Normal file
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
116
internal/utility/metadata.go
Normal file
116
internal/utility/metadata.go
Normal file
@@ -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
|
||||
}
|
||||
115
internal/utility/metadata_test.go
Normal file
115
internal/utility/metadata_test.go
Normal file
@@ -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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user