From 7db39822db81854519004d9dc16aab1c30bd177e Mon Sep 17 00:00:00 2001 From: Jack Date: Fri, 10 Jul 2026 14:30:28 +0800 Subject: [PATCH] Feature: user select pipeline support (#16788) ### Summary Feature: user select pipeline support --- internal/admin/service.go | 2 +- internal/dao/ingestion.go | 16 ++- .../ingestion/service/execute_task_test.go | 29 ++++- .../ingestion/service/ingestion_service.go | 49 +-------- .../service/real_consumer_dataflow_test.go | 8 +- internal/ingestion/task/dataflow_e2e_test.go | 19 +--- ...dataflow_real_pipeline_integration_test.go | 9 +- internal/ingestion/task/dataflow_service.go | 40 +++---- .../ingestion/task/dataflow_service_test.go | 60 ++-------- internal/ingestion/task/real_consumer_test.go | 7 +- internal/ingestion/task/task_context.go | 98 +++-------------- internal/ingestion/task/task_context_test.go | 92 ++++++++++++++++ internal/ingestion/task/task_handler.go | 33 +----- internal/ingestion/task/task_handler_test.go | 104 ++++++++---------- internal/service/document.go | 94 +++++++--------- internal/service/document_test.go | 103 ++++++++++++++++- web/vite.config.ts | 2 +- 17 files changed, 386 insertions(+), 379 deletions(-) diff --git a/internal/admin/service.go b/internal/admin/service.go index 4b47b23e80..39638a967c 100644 --- a/internal/admin/service.go +++ b/internal/admin/service.go @@ -156,7 +156,7 @@ func (s *Service) RemoveIngestionTasks(tasks []string) ([]map[string]string, err taskRecord := map[string]string{ "task_id": taskID, } - _, err := s.ingestionTaskDAO.RemoveByAPIServerOrAdminServer(taskID, nil) + _, err := s.ingestionTaskDAO.Delete(taskID, nil) if err != nil { taskRecord["remove"] = fmt.Sprintf("fail: %s", err.Error()) } else { diff --git a/internal/dao/ingestion.go b/internal/dao/ingestion.go index 860f81f930..e8b5f632f3 100644 --- a/internal/dao/ingestion.go +++ b/internal/dao/ingestion.go @@ -41,7 +41,6 @@ func NewIngestionTaskDAO() *IngestionTaskDAO { // failed → created : Retry (back to start) // canceled → created : Retry/re-execute (back to start) func (dao *IngestionTaskDAO) CheckAndCreate(ingestionTask *entity.IngestionTask) (*entity.IngestionTask, error) { - tx := DB.Begin() if tx.Error != nil { return nil, tx.Error @@ -222,8 +221,7 @@ type TaskInfo struct { FilesToDelete []string `json:"files_to_delete"` } -func (dao *IngestionTaskDAO) RemoveByAPIServerOrAdminServer(taskID string, userID *string) (*TaskInfo, error) { - +func (dao *IngestionTaskDAO) Delete(taskID string, userID *string) (*TaskInfo, error) { tx := DB.Begin() if tx.Error != nil { return nil, tx.Error @@ -327,9 +325,15 @@ func (dao *IngestionTaskDAO) GetByID(id string) (*entity.IngestionTask, error) { } func (dao *IngestionTaskDAO) GetByDocumentID(documentId string) (*entity.IngestionTask, error) { - var task *entity.IngestionTask - err := DB.Where("document_id = ?", documentId).First(&task).Error - return task, err + var tasks []*entity.IngestionTask + err := DB.Where("document_id = ?", documentId).Limit(1).Find(&tasks).Error + if err != nil { + return nil, err + } + if len(tasks) == 0 { + return nil, nil + } + return tasks[0], nil } type IngestionTaskLogDAO struct{} diff --git a/internal/ingestion/service/execute_task_test.go b/internal/ingestion/service/execute_task_test.go index 51fb74cd68..1767d0a1bd 100644 --- a/internal/ingestion/service/execute_task_test.go +++ b/internal/ingestion/service/execute_task_test.go @@ -48,7 +48,6 @@ func TestExecuteTask_CheckpointParseFailureDoesNotKillProcess(t *testing.T) { taskCtx := taskpkg.NewTaskContextForScheduling( context.Background(), &entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING}, - nil, ) // Execute the task - this should NOT panic or fatal exit (this is our main validation!) @@ -69,6 +68,33 @@ func TestExecuteTask_CheckpointParseFailureDoesNotKillProcess(t *testing.T) { } } +func TestDefaultRunDocumentTask_RequiresConfiguredPipelineID(t *testing.T) { + db := testutil.SetupTestDB(t) + cleanup := testutil.ReplaceDBForTest(t, db) + defer cleanup() + + _, kbID, docID, taskID := testutil.SeedTestData(t, db, + testutil.WithTenantID("tenant-1"), + testutil.WithKBID("kb-1"), + testutil.WithDocID("doc-1"), + testutil.WithTaskID("task-1"), + ) + + ingestor := NewIngestor("test", 1, []string{"pdf"}) + err := ingestor.defaultRunDocumentTask(context.Background(), &entity.IngestionTask{ + ID: taskID, + DocumentID: docID, + DatasetID: kbID, + Status: common.RUNNING, + }) + if err == nil { + t.Fatal("expected error when no pipeline is configured") + } + if err.Error() != "ingestion task task-1: no pipeline_id configured for document doc-1 or dataset kb-1" { + t.Fatalf("unexpected error: %v", err) + } +} + func TestExecuteTask_DataflowRoutesToTaskHandler(t *testing.T) { db := testutil.SetupTestDB(t) cleanup := testutil.ReplaceDBForTest(t, db) @@ -99,7 +125,6 @@ func TestExecuteTask_DataflowRoutesToTaskHandler(t *testing.T) { taskCtx := taskpkg.NewTaskContextForScheduling( context.Background(), &entity.IngestionTask{ID: taskID, DocumentID: docID, DatasetID: "kb-1", Status: common.RUNNING}, - nil, ) ingestor.executeTask(taskCtx) diff --git a/internal/ingestion/service/ingestion_service.go b/internal/ingestion/service/ingestion_service.go index b3244e49a7..71a0b3c916 100644 --- a/internal/ingestion/service/ingestion_service.go +++ b/internal/ingestion/service/ingestion_service.go @@ -18,14 +18,9 @@ package service import ( "context" - "encoding/json" "errors" "fmt" - "os" - "path/filepath" "ragflow/internal/utility" - goruntime "runtime" - "strings" "sync" "ragflow/internal/common" @@ -34,7 +29,6 @@ import ( "ragflow/internal/entity" taskpkg "ragflow/internal/ingestion/task" - "github.com/google/uuid" "google.golang.org/grpc" "gorm.io/gorm" ) @@ -170,7 +164,7 @@ func (e *Ingestor) Start() error { } // Construct TaskContext with parent context - taskCtx := taskpkg.NewTaskContextForScheduling(e.ctx, task, taskHandle) + taskCtx := taskpkg.NewTaskContextForScheduling(e.ctx, task) // Push to task channel; if full, reject the task (backpressure) select { @@ -294,50 +288,15 @@ func (e *Ingestor) executeTask(taskCtx *taskpkg.TaskContext) { 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 - } - - 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 (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) + if docTaskCtx.PipelineID == "" { + return fmt.Errorf("ingestion task %s: no pipeline_id configured for document %s or dataset %s", ingestionTask.ID, docTaskCtx.Doc.ID, docTaskCtx.KB.ID) } - + docTaskCtx.Ctx = ctx return taskpkg.NewTaskHandler(docTaskCtx).Handle() } diff --git a/internal/ingestion/service/real_consumer_dataflow_test.go b/internal/ingestion/service/real_consumer_dataflow_test.go index 8777720727..cc2a88c99e 100644 --- a/internal/ingestion/service/real_consumer_dataflow_test.go +++ b/internal/ingestion/service/real_consumer_dataflow_test.go @@ -102,10 +102,10 @@ func TestRealConsumer_DataflowMessageRoutesToExecuteTask(t *testing.T) { taskCtx := taskpkg.NewTaskContextForScheduling( context.Background(), task, - taskHandle, ) + var finalProgress float64 taskCtx.ProgressFunc = func(prog float64, msg string) { - taskCtx.Progress = int32(prog * 100) + finalProgress = prog progressEvents = append(progressEvents, msg) } ingestor.runDocumentTask = func(ctx context.Context, ingestionTask *entity.IngestionTask) error { @@ -120,8 +120,8 @@ func TestRealConsumer_DataflowMessageRoutesToExecuteTask(t *testing.T) { 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 finalProgress != 1.0 { + t.Fatalf("finalProgress = %v, want 1.0", finalProgress) } if len(progressEvents) != 2 { t.Fatalf("progressEvents = %v, want 2 events", progressEvents) diff --git a/internal/ingestion/task/dataflow_e2e_test.go b/internal/ingestion/task/dataflow_e2e_test.go index 69e8b8c66e..5abe220172 100644 --- a/internal/ingestion/task/dataflow_e2e_test.go +++ b/internal/ingestion/task/dataflow_e2e_test.go @@ -28,7 +28,6 @@ import ( "ragflow/internal/engine" "ragflow/internal/engine/elasticsearch" "ragflow/internal/engine/infinity" - "ragflow/internal/entity" "ragflow/internal/ingestion/testutil" "ragflow/internal/server" ) @@ -200,20 +199,14 @@ func TestDataflowE2E_TaskHandlerToDataflowService(t *testing.T) { 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) + ingestionTask, err := dao.NewIngestionTaskDAO().GetByID(taskID) if err != nil { - t.Fatalf("LoadTaskContext failed: %v", err) + t.Fatalf("GetByID failed: %v", err) + } + taskCtx, err := LoadFromIngestionTask(ingestionTask) + if err != nil { + t.Fatalf("LoadFromIngestionTask failed: %v", err) } // Track what was called diff --git a/internal/ingestion/task/dataflow_real_pipeline_integration_test.go b/internal/ingestion/task/dataflow_real_pipeline_integration_test.go index 29f8421355..3aedaa55ad 100644 --- a/internal/ingestion/task/dataflow_real_pipeline_integration_test.go +++ b/internal/ingestion/task/dataflow_real_pipeline_integration_test.go @@ -112,7 +112,6 @@ func TestDataflowService_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) { DocumentID: docID, DatasetID: kbID, }, - TaskType: "dataflow", Doc: entity.Document{ ID: docID, KbID: kbID, @@ -136,7 +135,7 @@ func TestDataflowService_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) { }). WithChunkCounter(&stubChunkCounter{}) - if err := svc.Run(context.Background()); err != nil { + if err := svc.Execute(context.Background()); err != nil { t.Fatalf("Run: %v", err) } if len(inserted) != 1 { @@ -251,7 +250,6 @@ func TestDataflowService_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *test DocumentID: docID, DatasetID: kbID, }, - TaskType: "dataflow", Doc: entity.Document{ ID: docID, KbID: kbID, @@ -270,7 +268,7 @@ func TestDataflowService_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *test svc := mustNewDataflowService(t, taskCtx, canvasID, 0, 0). WithChunkCounter(&stubChunkCounter{}) - if err := svc.Run(context.Background()); err != nil { + if err := svc.Execute(context.Background()); err != nil { t.Fatalf("Run: %v", err) } @@ -375,7 +373,6 @@ func TestRunDataflow_RealPipelineOutput_ProducesIndexFields(t *testing.T) { DocumentID: docID, DatasetID: kbID, }, - TaskType: "dataflow", Doc: entity.Document{ ID: docID, KbID: kbID, @@ -398,7 +395,7 @@ func TestRunDataflow_RealPipelineOutput_ProducesIndexFields(t *testing.T) { }). WithChunkCounter(&stubChunkCounter{}) - if err := svc.RunDataflow(context.Background(), pipelineOut); err != nil { + if err := svc.processOutput(context.Background(), pipelineOut); err != nil { t.Fatalf("RunDataflow: %v", err) } diff --git a/internal/ingestion/task/dataflow_service.go b/internal/ingestion/task/dataflow_service.go index cef839b79d..88f3b82552 100644 --- a/internal/ingestion/task/dataflow_service.go +++ b/internal/ingestion/task/dataflow_service.go @@ -228,7 +228,7 @@ 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 { +func (s *PipelineExecutor) Execute(ctx context.Context) error { if err := ctx.Err(); err != nil { return err } @@ -251,7 +251,7 @@ func (s *PipelineExecutor) Run(ctx context.Context) error { return nil } - if err := s.RunDataflow(ctx, pipelineOutput); err != nil { + if err := s.processOutput(ctx, pipelineOutput); err != nil { return err } @@ -261,7 +261,7 @@ func (s *PipelineExecutor) Run(ctx context.Context) error { return nil } -func (s *PipelineExecutor) RunDataflow(ctx context.Context, pipelineOutput map[string]any) error { +func (s *PipelineExecutor) processOutput(ctx context.Context, pipelineOutput map[string]any) error { taskStart := time.Now() if pipelineOutput == nil { return nil @@ -425,30 +425,22 @@ func (s *PipelineExecutor) defaultLoadDSL(ctx context.Context, dataflowID string 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) + canvas, err := dao.NewUserCanvasDAO().GetByID(dataflowID) if err != nil { - return "", "", fmt.Errorf("marshal pipeline log dsl %s: %w", dataflowID, err) + return "", "", fmt.Errorf("load dataflow canvas %s: %w", dataflowID, err) } - correctedID := dataflowID - if pipelineLog.PipelineID != nil && *pipelineLog.PipelineID != "" { - correctedID = *pipelineLog.PipelineID + + canvasTitle := "" + if canvas.Title != nil { + canvasTitle = *canvas.Title } - return string(raw), correctedID, nil + common.Info(fmt.Sprintf("load dataflow canvas %s, name %s", dataflowID, canvasTitle)) + + raw, err := json.Marshal(canvas.DSL) + if err != nil { + return "", "", fmt.Errorf("marshal canvas dsl %s: %w", dataflowID, err) + } + return string(raw), dataflowID, nil } func (s *PipelineExecutor) tokenizerEmbedderResolver() componentpkg.EmbedderResolver { diff --git a/internal/ingestion/task/dataflow_service_test.go b/internal/ingestion/task/dataflow_service_test.go index 9ac6b0e3ca..1d111e5424 100644 --- a/internal/ingestion/task/dataflow_service_test.go +++ b/internal/ingestion/task/dataflow_service_test.go @@ -29,7 +29,6 @@ func makeTaskCtx() *TaskContext { ID: "task-1", DocumentID: "doc-1", }, - TaskType: "dataflow", Doc: entity.Document{ ID: "doc-1", KbID: "kb-1", @@ -77,7 +76,8 @@ func TestDataflowService_DefaultLoadDSL_UsesUserCanvas(t *testing.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 { + title := "title 1" + if err := dao.NewUserCanvasDAO().Create(&entity.UserCanvas{Title: &title, ID: "canvas-1", UserID: "u1", Permission: "me", CanvasCategory: "agent_canvas", DSL: dslMap}); err != nil { t.Fatalf("create user canvas: %v", err) } @@ -99,48 +99,6 @@ func TestDataflowService_DefaultLoadDSL_UsesUserCanvas(t *testing.T) { } } -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 // ============================================================================= @@ -429,7 +387,7 @@ func TestIncrementChunkNum_ProcessDuration(t *testing.T) { func TestRunDataflow_NilOutput(t *testing.T) { svc := mustNewDataflowService(t, makeTaskCtx(), "flow-1", 0, 0) - err := svc.RunDataflow(context.Background(), nil) + err := svc.processOutput(context.Background(), nil) if err != nil { t.Errorf("expected nil error for nil output, got %v", err) } @@ -439,7 +397,7 @@ 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{}) + err := svc.processOutput(context.Background(), map[string]any{}) if err != nil { t.Errorf("expected nil error for empty output, got %v", err) } @@ -449,7 +407,7 @@ 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": ""}) + err := svc.processOutput(context.Background(), map[string]any{"markdown": ""}) if err != nil { t.Errorf("expected nil error for empty normalized output, got %v", err) } @@ -480,7 +438,7 @@ func TestRunDataflow_FullFlow(t *testing.T) { {"text": "world"}, }, } - err := svc.RunDataflow(context.Background(), output) + err := svc.processOutput(context.Background(), output) if err != nil { t.Errorf("unexpected error: %v", err) } @@ -511,7 +469,7 @@ func TestRunDataflow_AlreadyHasVectors(t *testing.T) { {"text": "hello", "q_768_vec": []float64{0.1, 0.2}}, }, } - err := svc.RunDataflow(context.Background(), output) + err := svc.processOutput(context.Background(), output) if err != nil { t.Errorf("unexpected error: %v", err) } @@ -522,7 +480,7 @@ func TestRunDataflow_ContextCanceled(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - err := svc.RunDataflow(ctx, map[string]any{ + err := svc.processOutput(ctx, map[string]any{ "chunks": []map[string]any{{"text": "hello"}}, }) if err == nil { @@ -609,7 +567,7 @@ func TestDataflowService_Run_MainFlowWithStubs(t *testing.T) { return nil }) - err := svc.Run(context.Background()) + err := svc.Execute(context.Background()) if err != nil { t.Fatalf("unexpected error: %v", err) } diff --git a/internal/ingestion/task/real_consumer_test.go b/internal/ingestion/task/real_consumer_test.go index afa96bd77b..37ebd3ee5f 100644 --- a/internal/ingestion/task/real_consumer_test.go +++ b/internal/ingestion/task/real_consumer_test.go @@ -148,12 +148,9 @@ func TestRealProducerConsumer(t *testing.T) { } // ── 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) + tc, err := LoadFromIngestionTask(task) if err != nil { - t.Fatalf("LoadTaskContext: %v", err) + t.Fatalf("LoadFromIngestionTask: %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) diff --git a/internal/ingestion/task/task_context.go b/internal/ingestion/task/task_context.go index 76ce31068c..228ea6d950 100644 --- a/internal/ingestion/task/task_context.go +++ b/internal/ingestion/task/task_context.go @@ -19,75 +19,39 @@ package task import ( "context" "fmt" - "time" + "strings" - "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. +// TaskContext holds the execution inputs for an ingestion document task. type TaskContext struct { - // Context and execution control (from service package) - Ctx context.Context - TaskHandle common.TaskHandle + Ctx context.Context - // 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 any - // 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 { +func NewTaskContextForScheduling(ctx context.Context, task *entity.IngestionTask) *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. +// 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 { @@ -104,14 +68,10 @@ func LoadFromIngestionTask(ingestionTask *entity.IngestionTask) (*TaskContext, e return nil, fmt.Errorf("error when load tenant %s: %w", kb.TenantID, err) } - pipelineID := "" - if doc.PipelineID != nil { - pipelineID = *doc.PipelineID - } + pipelineID := resolvePipelineID(doc, kb) return &TaskContext{ IngestionTask: ingestionTask, - TaskType: "dataflow", PipelineID: pipelineID, Doc: *doc, KB: *kb, @@ -119,40 +79,16 @@ func LoadFromIngestionTask(ingestionTask *entity.IngestionTask) (*TaskContext, e }, 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) +func resolvePipelineID(doc *entity.Document, kb *entity.Knowledgebase) string { + if doc != nil && doc.PipelineID != nil { + if pipelineID := strings.TrimSpace(*doc.PipelineID); pipelineID != "" { + return pipelineID + } } - - 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) + if kb != nil && kb.PipelineID != nil { + if pipelineID := strings.TrimSpace(*kb.PipelineID); pipelineID != "" { + return pipelineID + } } - - 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 + return "" } diff --git a/internal/ingestion/task/task_context_test.go b/internal/ingestion/task/task_context_test.go index 541027e83e..23cd7bbdb3 100644 --- a/internal/ingestion/task/task_context_test.go +++ b/internal/ingestion/task/task_context_test.go @@ -59,3 +59,95 @@ func pushDB(t *testing.T, db *gorm.DB) { } func ptr(s string) *string { return &s } + +func TestLoadFromIngestionTask_FallsBackToKnowledgebasePipelineID(t *testing.T) { + db := setupTestDB(t) + pushDB(t, db) + + if err := db.Create(&entity.Document{ + ID: "doc-1", + KbID: "kb-1", + ParserID: "naive", + ParserConfig: entity.JSONMap{}, + Name: ptr("doc.pdf"), + Status: ptr("1"), + }).Error; err != nil { + t.Fatalf("create document: %v", err) + } + kbPipelineID := "kb-flow-1" + if err := db.Create(&entity.Knowledgebase{ + ID: "kb-1", + TenantID: "tenant-1", + EmbdID: "embd-1", + PipelineID: &kbPipelineID, + ParserConfig: entity.JSONMap{}, + Status: ptr("1"), + }).Error; err != nil { + t.Fatalf("create knowledgebase: %v", err) + } + if err := db.Create(&entity.Tenant{ + ID: "tenant-1", + Status: ptr("1"), + }).Error; err != nil { + t.Fatalf("create tenant: %v", err) + } + + ctx, err := LoadFromIngestionTask(&entity.IngestionTask{ + ID: "task-1", + DocumentID: "doc-1", + DatasetID: "kb-1", + }) + if err != nil { + t.Fatalf("LoadFromIngestionTask: %v", err) + } + if ctx.PipelineID != kbPipelineID { + t.Fatalf("PipelineID = %q, want %q", ctx.PipelineID, kbPipelineID) + } +} + +func TestLoadFromIngestionTask_PrefersDocumentPipelineID(t *testing.T) { + db := setupTestDB(t) + pushDB(t, db) + + docPipelineID := "doc-flow-1" + if err := db.Create(&entity.Document{ + ID: "doc-1", + KbID: "kb-1", + ParserID: "naive", + ParserConfig: entity.JSONMap{}, + PipelineID: &docPipelineID, + Name: ptr("doc.pdf"), + Status: ptr("1"), + }).Error; err != nil { + t.Fatalf("create document: %v", err) + } + kbPipelineID := "kb-flow-1" + if err := db.Create(&entity.Knowledgebase{ + ID: "kb-1", + TenantID: "tenant-1", + EmbdID: "embd-1", + PipelineID: &kbPipelineID, + ParserConfig: entity.JSONMap{}, + Status: ptr("1"), + }).Error; err != nil { + t.Fatalf("create knowledgebase: %v", err) + } + if err := db.Create(&entity.Tenant{ + ID: "tenant-1", + Status: ptr("1"), + }).Error; err != nil { + t.Fatalf("create tenant: %v", err) + } + + ctx, err := LoadFromIngestionTask(&entity.IngestionTask{ + ID: "task-1", + DocumentID: "doc-1", + DatasetID: "kb-1", + }) + if err != nil { + t.Fatalf("LoadFromIngestionTask: %v", err) + } + if ctx.PipelineID != docPipelineID { + t.Fatalf("PipelineID = %q, want %q", ctx.PipelineID, docPipelineID) + } +} diff --git a/internal/ingestion/task/task_handler.go b/internal/ingestion/task/task_handler.go index 1807fee9be..6f85bedae9 100644 --- a/internal/ingestion/task/task_handler.go +++ b/internal/ingestion/task/task_handler.go @@ -49,37 +49,8 @@ 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() - } + return h.handleDataflow() } func (h *TaskHandler) handleMemory() error { @@ -99,7 +70,7 @@ func (h *TaskHandler) handleDataflow() error { if runCtx == nil { runCtx = context.Background() } - return svc.Run(runCtx) + return svc.Execute(runCtx) } func (h *TaskHandler) handleRaptor() error { diff --git a/internal/ingestion/task/task_handler_test.go b/internal/ingestion/task/task_handler_test.go index f3ea61dc36..613c5e0913 100644 --- a/internal/ingestion/task/task_handler_test.go +++ b/internal/ingestion/task/task_handler_test.go @@ -11,18 +11,12 @@ import ( func testStrPtr(s string) *string { return &s } -func makeTaskHandlerTestContext(taskType string) *TaskContext { - var progressCalls []float64 - pipelineID := "" - if strings.HasPrefix(taskType, "dataflow") { - pipelineID = "flow-1" - } +func makeTaskHandlerTestContext(pipelineID string) *TaskContext { return &TaskContext{ IngestionTask: &entity.IngestionTask{ ID: "task-1", DocumentID: "doc-1", }, - TaskType: taskType, PipelineID: pipelineID, Doc: entity.Document{ ID: "doc-1", @@ -40,9 +34,7 @@ func makeTaskHandlerTestContext(taskType string) *TaskContext { ID: "tenant-1", LLMID: "gpt-4", }, - ProgressFunc: func(prog float64, msg string) { - progressCalls = append(progressCalls, prog) - }, + ProgressFunc: func(prog float64, msg string) {}, } } @@ -60,9 +52,10 @@ func newNoopDataflowService(ctx *TaskContext, dataflowID string) (*PipelineExecu }). 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}}, - }, + "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) { @@ -83,44 +76,22 @@ 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}, +func TestTaskHandler_HandleRejectsNilContext(t *testing.T) { + if err := NewTaskHandler(nil).Handle(); err == nil { + t.Fatal("expected error for nil context") } +} - 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_HandleRequiresPipelineID(t *testing.T) { + ctx := makeTaskHandlerTestContext("") + handler := NewTaskHandler(ctx) + if err := handler.Handle(); err == nil { + t.Fatal("expected error for empty pipeline id") } } func TestTaskHandler_DefaultDataflowServiceInjectsProgress(t *testing.T) { - ctx := makeTaskHandlerTestContext("dataflow") + ctx := makeTaskHandlerTestContext("flow-1") handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { svc, err := NewDataflowService(ctx, dataflowID, 0, 0) if err != nil { @@ -137,7 +108,7 @@ func TestTaskHandler_DefaultDataflowServiceInjectsProgress(t *testing.T) { } func TestTaskHandler_Dataflow_UsesTaskContext(t *testing.T) { - ctx := makeTaskHandlerTestContext("dataflow") + ctx := makeTaskHandlerTestContext("flow-1") type ctxKey string const key ctxKey = "trace" ctx.Ctx = context.WithValue(context.Background(), key, "task-ctx") @@ -165,8 +136,34 @@ func TestTaskHandler_Dataflow_UsesTaskContext(t *testing.T) { } } +func TestTaskHandler_Dataflow_UsesBackgroundContextWhenMissing(t *testing.T) { + ctx := makeTaskHandlerTestContext("flow-1") + + handler := NewTaskHandler(ctx).WithDataflowServiceFactory(func(ctx *TaskContext, dataflowID string) (*PipelineExecutor, error) { + return 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(runCtx context.Context, dsl string) (map[string]any, string, error) { + if runCtx == nil { + t.Fatal("runCtx is nil") + } + return map[string]any{"chunks": []map[string]any{{"text": "stub", "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{}), nil + }) + if err := handler.Handle(); err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + func TestTaskHandler_Dataflow_ShowsProgressAndPipelineLog(t *testing.T) { - ctx := makeTaskHandlerTestContext("dataflow") + ctx := makeTaskHandlerTestContext("flow-1") ctx.Doc.PipelineID = testStrPtr("flow-1") ctx.Doc.Name = testStrPtr("verify-dataflow.pdf") @@ -180,7 +177,6 @@ func TestTaskHandler_Dataflow_ShowsProgressAndPipelineLog(t *testing.T) { 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) { @@ -190,24 +186,20 @@ func TestTaskHandler_Dataflow_ShowsProgressAndPipelineLog(t *testing.T) { }). 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}}, - }, + "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{}). diff --git a/internal/service/document.go b/internal/service/document.go index d56f4183c5..e65433b2dc 100644 --- a/internal/service/document.go +++ b/internal/service/document.go @@ -70,6 +70,7 @@ type DocumentService struct { fileDAO *dao.FileDAO canvasDAO *dao.UserCanvasDAO api4ConvDAO *dao.API4ConversationDAO + publishTask func(subject string, payload []byte) error } // NewDocumentService create document service @@ -88,6 +89,13 @@ func NewDocumentService() *DocumentService { fileDAO: dao.NewFileDAO(), canvasDAO: dao.NewUserCanvasDAO(), api4ConvDAO: dao.NewAPI4ConversationDAO(), + publishTask: func(subject string, payload []byte) error { + msgQueueEngine := engine.GetMessageQueueEngine() + if msgQueueEngine == nil { + return fmt.Errorf("message queue engine is not initialized") + } + return msgQueueEngine.PublishTask(subject, payload) + }, } } @@ -687,7 +695,7 @@ func (s *DocumentService) DeleteDocuments(ids []string, deleteAll bool, datasetI deleted := 0 for _, docID := range ids { if err := s.deleteDocumentFull(docID); err != nil { - common.Logger.Warn(fmt.Sprintf("DeleteDocuments: failed to delete %s: %v", docID, err)) + common.Warn(fmt.Sprintf("DeleteDocuments: failed to delete %s: %v", docID, err)) continue } deleted++ @@ -706,9 +714,20 @@ func (s *DocumentService) deleteDocumentFull(docID string) error { } // Delete tasks from DB - if _, delErr := s.taskDAO.DeleteByDocIDs([]string{docID}); delErr != nil { - common.Logger.Warn(fmt.Sprintf("failed to delete tasks for %s: %v", docID, delErr)) + ingestionTask, err := s.ingestionTaskDAO.GetByDocumentID(docID) + if err != nil { + common.Error(fmt.Sprintf("failed to get ingestion task by doc:%s", docID), err) + return err } + if ingestionTask != nil { + taskInfo, err := s.ingestionTaskDAO.Delete(ingestionTask.ID, &ingestionTask.UserID) + if err != nil { + return err + } + // FIXME: need to add logic to delete files in taskInfo + common.Warn(fmt.Sprintf("need to delete files from taskInfo: %v", taskInfo)) + } + s.deleteDocEngineData(docID, kb.TenantID, doc.KbID) if err := s.deleteDocRecordWithCounters(doc, kb.ID); err != nil { return err @@ -1264,7 +1283,7 @@ func (s *DocumentService) RemoveIngestionTasks(tasks []string, userID string) ([ taskRecord := map[string]string{ "task_id": taskID, } - _, err := s.ingestionTaskDAO.RemoveByAPIServerOrAdminServer(taskID, &userID) + _, err := s.ingestionTaskDAO.Delete(taskID, &userID) if err != nil { taskRecord["remove"] = fmt.Sprintf("fail: %s", err.Error()) } else { @@ -1392,7 +1411,7 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com } } if doc.PipelineID != nil && strings.TrimSpace(*doc.PipelineID) != "" { - if err := s.queueDocumentDataflowTask(kb, doc, strings.TrimSpace(*doc.PipelineID), 0); err != nil { + if err := s.queueDocumentDataflowTask(kb, doc, userID); err != nil { return common.CodeExceptionError, err } continue @@ -1423,12 +1442,10 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com return common.CodeExceptionError, err } - 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) @@ -1542,28 +1559,30 @@ func (s *DocumentService) countDoneDocuments(datasetID string) (int64, error) { return count, err } -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 - } +func (s *DocumentService) queueDocumentDataflowTask(kb *entity.Knowledgebase, doc *entity.Document, userID string) error { if err := s.beginDocumentParse(doc.ID); err != nil { return err } - task := s.newDocumentParseTask(doc, 0, maximumTaskPageNumber, priority) - task.TaskType = "dataflow" - if err := s.taskDAO.CreateMany([]*entity.Task{task}); err != nil { + task := &entity.IngestionTask{ + DocumentID: doc.ID, + UserID: userID, + DatasetID: kb.ID, + Schema: nil, + Status: common.CREATED, + } + task, err := s.ingestionTaskDAO.CheckAndCreate(task) + if err != nil { return err } - message := documentTaskMessage(task) - message["task_type"] = task.TaskType - message["kb_id"] = doc.KbID - message["tenant_id"] = kb.TenantID - message["dataflow_id"] = flowID - message["file"] = nil - if redisClient := redis.Get(); redisClient == nil || !redisClient.QueueProduct(documentParseQueueName(doc, priority), message) { - return fmt.Errorf("Can't access Redis. Please check the Redis' status.") + taskMessage := common.TaskMessage{ + TaskID: task.ID, + TaskType: common.TaskTypeIngestionTask, } - return nil + payload, err := json.Marshal(taskMessage) + if err != nil { + return err + } + return s.publishTask("tasks.RAGFLOW", payload) } func (s *DocumentService) newDocumentParseTasks(doc *entity.Document, bucket, objectName string, priority int64) ([]*entity.Task, error) { @@ -1888,35 +1907,6 @@ func (s *DocumentService) beginDocumentParse(docID string) error { }).Error } -func documentParseQueueName(doc *entity.Document, priority int64) string { - suffix := "common" - if doc.ParserID == string(entity.ParserTypeResume) { - suffix = "resume" - } - return fmt.Sprintf("te.%d.%s", priority, suffix) -} - -func documentTaskMessage(task *entity.Task) map[string]interface{} { - beginAt := "" - if task.BeginAt != nil { - beginAt = task.BeginAt.Format("2006-01-02 15:04:05") - } - digest := "" - if task.Digest != nil { - digest = *task.Digest - } - return map[string]interface{}{ - "id": task.ID, - "doc_id": task.DocID, - "from_page": task.FromPage, - "to_page": task.ToPage, - "progress": task.Progress, - "priority": task.Priority, - "begin_at": beginAt, - "digest": digest, - } -} - func documentParseTaskDigest(doc *entity.Document, fromPage, toPage int64) string { hasher := xxhash.New() config := map[string]interface{}{ diff --git a/internal/service/document_test.go b/internal/service/document_test.go index f620bd4421..0ec66ed8fe 100644 --- a/internal/service/document_test.go +++ b/internal/service/document_test.go @@ -388,6 +388,8 @@ func setupServiceTestDB(t *testing.T) *gorm.DB { &entity.Document{}, &entity.Knowledgebase{}, &entity.Task{}, + &entity.IngestionTask{}, + &entity.IngestionTaskLog{}, &entity.File2Document{}, &entity.File{}, &entity.User{}, @@ -419,6 +421,7 @@ func testDocumentService(t *testing.T) *DocumentService { taskDAO: dao.NewTaskDAO(), file2DocumentDAO: dao.NewFile2DocumentDAO(), fileDAO: dao.NewFileDAO(), + ingestionTaskDAO: dao.NewIngestionTaskDAO(), docEngine: nil, metadataSvc: nil, // nil engine → metadata ops skipped } @@ -511,6 +514,21 @@ func insertTestTask(t *testing.T, id, docID string) { } } +func insertTestIngestionTask(t *testing.T, id, userID, docID, datasetID string) { + t.Helper() + task := &entity.IngestionTask{ + ID: id, + UserID: userID, + DocumentID: docID, + DatasetID: datasetID, + Schema: entity.JSONMap{}, + Status: common.CREATED, + } + if err := dao.DB.Create(task).Error; err != nil { + t.Fatalf("insert test ingestion task: %v", err) + } +} + func insertTestFile2Document(t *testing.T, id, fileID, docID string) { t.Helper() f2d := &entity.File2Document{ @@ -570,7 +588,7 @@ func TestDeleteDocumentFull_Basic(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 3, 100, 50) insertTestDoc(t, "doc-1", "kb-1", 30, 10) - insertTestTask(t, "task-1", "doc-1") + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) @@ -625,6 +643,7 @@ func TestDeleteDocumentFull_CleansUpFile2Document(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) insertTestDoc(t, "doc-1", "kb-1", 10, 5) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") loc := "path/to/blob" insertTestFile(t, "file-1", "kb-1", "test.pdf", &loc) insertTestFile2Document(t, "f2d-1", "file-1", "doc-1") @@ -657,6 +676,7 @@ func TestDeleteDocumentFull_SharedFilePreserved(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 2, 20, 10) insertTestDoc(t, "doc-1", "kb-1", 10, 5) insertTestDoc(t, "doc-2", "kb-1", 10, 5) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") loc := "shared/blob" insertTestFile(t, "file-shared", "kb-1", "shared.pdf", &loc) @@ -900,6 +920,9 @@ func TestDeleteDocuments_DeleteAll(t *testing.T) { insertTestDoc(t, "doc-1", "kb-1", 30, 10) insertTestDoc(t, "doc-2", "kb-1", 40, 20) insertTestDoc(t, "doc-3", "kb-1", 30, 20) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + insertTestIngestionTask(t, "task-2", "user-1", "doc-2", "kb-1") + insertTestIngestionTask(t, "task-3", "user-1", "doc-3", "kb-1") svc := testDocumentService(t) @@ -927,6 +950,9 @@ func TestDeleteDocuments_ByIDs(t *testing.T) { insertTestDoc(t, "doc-1", "kb-1", 30, 10) insertTestDoc(t, "doc-2", "kb-1", 40, 20) insertTestDoc(t, "doc-3", "kb-1", 30, 20) // won't be deleted + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") + insertTestIngestionTask(t, "task-2", "user-1", "doc-2", "kb-1") + insertTestIngestionTask(t, "task-3", "user-1", "doc-3", "kb-1") svc := testDocumentService(t) @@ -1003,6 +1029,7 @@ func TestDeleteDocuments_Deduplicate(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5) insertTestDoc(t, "doc-1", "kb-1", 10, 5) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) @@ -1048,6 +1075,79 @@ func insertTestTaskWithProgress(t *testing.T, id, docID string, progress float64 } } +func TestQueueDocumentDataflowTask_PublishesIngestionTaskMessage(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + + insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0) + insertTestDoc(t, "doc-1", "kb-1", 9, 4) + + var publishedSubject string + var publishedPayload []byte + svc := testDocumentService(t) + svc.publishTask = func(subject string, payload []byte) error { + publishedSubject = subject + publishedPayload = append([]byte(nil), payload...) + return nil + } + + kb, err := svc.kbDAO.GetByID("kb-1") + if err != nil { + t.Fatalf("load kb: %v", err) + } + doc, err := svc.documentDAO.GetByID("doc-1") + if err != nil { + t.Fatalf("load doc: %v", err) + } + + if err := svc.queueDocumentDataflowTask(kb, doc, "user-1"); err != nil { + t.Fatalf("queueDocumentDataflowTask: %v", err) + } + + if publishedSubject != "tasks.RAGFLOW" { + t.Fatalf("subject = %q, want %q", publishedSubject, "tasks.RAGFLOW") + } + if len(publishedPayload) == 0 { + t.Fatal("expected published payload") + } + + var msg common.TaskMessage + if err := json.Unmarshal(publishedPayload, &msg); err != nil { + t.Fatalf("unmarshal payload: %v", err) + } + if msg.TaskType != common.TaskTypeIngestionTask { + t.Fatalf("task type = %q, want %q", msg.TaskType, common.TaskTypeIngestionTask) + } + if msg.TaskID == "" { + t.Fatal("expected non-empty task id") + } + + ingestionTask, err := svc.ingestionTaskDAO.GetByID(msg.TaskID) + if err != nil { + t.Fatalf("load ingestion task: %v", err) + } + if ingestionTask.DocumentID != "doc-1" { + t.Fatalf("ingestion task document id = %q, want %q", ingestionTask.DocumentID, "doc-1") + } + if ingestionTask.DatasetID != "kb-1" { + t.Fatalf("ingestion task dataset id = %q, want %q", ingestionTask.DatasetID, "kb-1") + } + if ingestionTask.UserID != "user-1" { + t.Fatalf("ingestion task user id = %q, want %q", ingestionTask.UserID, "user-1") + } + if ingestionTask.Status != common.CREATED { + t.Fatalf("ingestion task status = %q, want %q", ingestionTask.Status, common.CREATED) + } + + tasks, err := svc.taskDAO.GetByDocID("doc-1") + if err != nil { + t.Fatalf("load legacy tasks: %v", err) + } + if len(tasks) != 0 { + t.Fatalf("expected no legacy task rows, got %d", len(tasks)) + } +} + func TestStopParseDocuments_Success(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) @@ -1222,6 +1322,7 @@ func TestDeleteDocument_DeligatesToFullCleanup(t *testing.T) { insertTestKB(t, "kb-1", "tenant-1", 1, 5, 2) insertTestDoc(t, "doc-1", "kb-1", 5, 2) + insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1") svc := testDocumentService(t) diff --git a/web/vite.config.ts b/web/vite.config.ts index 9c8ec557d9..1113a92b84 100644 --- a/web/vite.config.ts +++ b/web/vite.config.ts @@ -75,7 +75,7 @@ export default defineConfig(({ mode }) => { }, }, hybrid: { - '^(/v1/document)|^(/v1/llm/list)|^(/api/v1/datasets)|^(/api/v1/memories)|^(/v1/user)|^(/v1/user/tenant_info)|^(/v1/tenant/list)|^(/v1/system/config)|^(/v1/user/login)|^(/v1/user/logout)|^(/api/v1/files)': + '^(/v1/document)|^(/v1/llm/list)|^(/api/v1/datasets)|^(/api/v1/documents/ingest)|^(/api/v1/memories)|^(/v1/user)|^(/v1/user/tenant_info)|^(/v1/tenant/list)|^(/v1/system/config)|^(/v1/user/login)|^(/v1/user/logout)|^(/api/v1/files)': { target: 'http://127.0.0.1:9384/', changeOrigin: true,