mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-17 21:27:23 +08:00
Feat: ingestion cleanup (#16953)
### Summary 1. Remove dead code (replaced by builtin ingestion pipeline) 2. Refactor (move document parsing progress from http api into ingestion executor)
This commit is contained in:
@@ -194,6 +194,21 @@ func (dao *IngestionTaskDAO) GetByDocumentID(documentId string) (*entity.Ingesti
|
||||
return tasks[0], nil
|
||||
}
|
||||
|
||||
// DeleteIfTerminal deletes ingestion tasks for a document that are in a
|
||||
// terminal state (COMPLETED, STOPPED, FAILED) or still queued (CREATED).
|
||||
// RUNNING and STOPPING tasks are NOT deleted because an in-flight worker
|
||||
// would keep writing chunks and corrupt a new run's results.
|
||||
// Returns the number of rows deleted.
|
||||
func (dao *IngestionTaskDAO) DeleteIfTerminal(documentID string) (int64, error) {
|
||||
result := DB.Where("document_id = ? AND status NOT IN (?, ?)",
|
||||
documentID, common.RUNNING, common.STOPPING).
|
||||
Delete(&entity.IngestionTask{})
|
||||
if result.Error != nil {
|
||||
return 0, result.Error
|
||||
}
|
||||
return result.RowsAffected, nil
|
||||
}
|
||||
|
||||
type IngestionTaskLogDAO struct{}
|
||||
|
||||
func NewIngestionTaskLogDAO() *IngestionTaskLogDAO {
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"fmt"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/entity"
|
||||
|
||||
@@ -133,3 +134,61 @@ func TestIngestionTaskDAOUpdateStatusIfCurrentRejectsMismatchedStatus(t *testing
|
||||
t.Fatalf("status = %q, want %q", reloaded.Status, common.STOPPING)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIngestionTaskDAODeleteIfTerminal_RemovesOnlyTerminal(t *testing.T) {
|
||||
db := setupTaskTestDB(t)
|
||||
orig := DB
|
||||
DB = db
|
||||
t.Cleanup(func() { DB = orig })
|
||||
|
||||
// Create tasks in different statuses, each with a unique docID.
|
||||
statuses := []string{common.CREATED, common.RUNNING, common.STOPPING, common.COMPLETED, common.STOPPED, common.FAILED}
|
||||
for i, status := range statuses {
|
||||
docID := fmt.Sprintf("doc-%d", i)
|
||||
task := &entity.IngestionTask{
|
||||
ID: fmt.Sprintf("task-%d", i),
|
||||
UserID: "user-1",
|
||||
DocumentID: docID,
|
||||
DatasetID: "kb-1",
|
||||
Status: status,
|
||||
}
|
||||
if err := db.Create(task).Error; err != nil {
|
||||
t.Fatalf("create task %s: %v", status, err)
|
||||
}
|
||||
}
|
||||
|
||||
// DeleteIfTerminal deletes everything except RUNNING and STOPPING.
|
||||
// CREATED is safe to delete (no worker has claimed it yet);
|
||||
// COMPLETED/STOPPED/FAILED are terminal.
|
||||
// Call it for every doc and verify the negative cases survived.
|
||||
for i := 0; i < len(statuses); i++ {
|
||||
docID := fmt.Sprintf("doc-%d", i)
|
||||
_, err := NewIngestionTaskDAO().DeleteIfTerminal(docID)
|
||||
if err != nil {
|
||||
t.Fatalf("DeleteIfTerminal(doc-%d): %v", i, err)
|
||||
}
|
||||
}
|
||||
|
||||
// RUNNING and STOPPING must survive.
|
||||
for _, i := range []int{1, 2} {
|
||||
docID := fmt.Sprintf("doc-%d", i)
|
||||
task, err := NewIngestionTaskDAO().GetByDocumentID(docID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetByDocumentID %s: %v", docID, err)
|
||||
}
|
||||
if task == nil {
|
||||
t.Fatalf("%s task (doc=%d) must not be deleted", statuses[i], i)
|
||||
}
|
||||
}
|
||||
// CREATED, COMPLETED, STOPPED, FAILED must be gone.
|
||||
for _, i := range []int{0, 3, 4, 5} {
|
||||
docID := fmt.Sprintf("doc-%d", i)
|
||||
task, err := NewIngestionTaskDAO().GetByDocumentID(docID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetByDocumentID %s: %v", docID, err)
|
||||
}
|
||||
if task != nil {
|
||||
t.Fatalf("%s task (doc=%d) should be deleted, still present", statuses[i], i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -49,6 +49,25 @@ func (m *mockCanvasStage) Invoke(_ context.Context, inputs map[string]any) (map[
|
||||
func (m *mockCanvasStage) Inputs() map[string]string { return map[string]string{"name": "string"} }
|
||||
func (m *mockCanvasStage) Outputs() map[string]string { return map[string]string{"output": "any"} }
|
||||
|
||||
// oneShotErrStage errors on the first Invoke (simulating a component crash
|
||||
// mid-run), then delegates to the embedded mock on subsequent calls. Used to
|
||||
// test that a second pipeline Run on the same taskID resumes past non-terminal
|
||||
// checkpoints instead of re-executing completed components.
|
||||
type oneShotErrStage struct {
|
||||
mockCanvasStage
|
||||
n int
|
||||
}
|
||||
|
||||
func (s *oneShotErrStage) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
|
||||
s.n++
|
||||
if s.n == 1 {
|
||||
return nil, errors.New("simulated crash")
|
||||
}
|
||||
return s.mockCanvasStage.Invoke(ctx, inputs)
|
||||
}
|
||||
func (s *oneShotErrStage) Inputs() map[string]string { return s.mockCanvasStage.Inputs() }
|
||||
func (s *oneShotErrStage) Outputs() map[string]string { return s.mockCanvasStage.Outputs() }
|
||||
|
||||
func TestPipelineRunHappyPath(t *testing.T) {
|
||||
stageA := &mockCanvasStage{output: map[string]any{"a": 1}}
|
||||
stageB := &mockCanvasStage{output: map[string]any{"b": 2}}
|
||||
@@ -361,8 +380,81 @@ func TestPipelineRunResumableAutoResumes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineRun_RequireResumeRejectsWithoutStore verifies M4 (plan §6.a
|
||||
// 方案 A): with WithRequireResume set and no checkpoint store resolvable (no
|
||||
// TestPipelineRunResumableCrossRunResume validates crash-recovery resume
|
||||
// across two Run calls: when the terminal component errors mid-run (simulated
|
||||
// crash), non-terminal checkpoints + interrupt state persist. The second Run
|
||||
// on the same taskID resumes past completed non-terminal components instead of
|
||||
// re-executing them. A non-terminal stage (A) runs exactly once across both runs; the
|
||||
// terminal stage (B, a oneShotErrStage) errors on run 1 and succeeds on run 2.
|
||||
func TestPipelineRunResumableCrossRunResume(t *testing.T) {
|
||||
mockA := &mockCanvasStage{output: map[string]any{"a": 1}}
|
||||
mockB := &mockCanvasStage{output: map[string]any{"b": 2}}
|
||||
termStage := &oneShotErrStage{mockCanvasStage: *mockB}
|
||||
|
||||
const (
|
||||
nameA = "p.XRunStageA"
|
||||
nameB = "p.XRunStageB"
|
||||
)
|
||||
runtime.MustRegister(nameA, runtime.CategoryIngestion,
|
||||
func(_ string, _ map[string]any) (runtime.Component, error) { return mockA, nil },
|
||||
runtime.Metadata{Version: "1.0.0"})
|
||||
runtime.MustRegister(nameB, runtime.CategoryIngestion,
|
||||
func(_ string, _ map[string]any) (runtime.Component, error) { return termStage, nil },
|
||||
runtime.Metadata{Version: "1.0.0"})
|
||||
|
||||
store := newMemCheckpointStore()
|
||||
mr := miniredis.RunT(t)
|
||||
client := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() { client.Close() })
|
||||
tracker := canvas.NewRunTrackerWithClient(client, time.Hour)
|
||||
|
||||
const taskID = "task-cross-run-resume"
|
||||
pipe, err := NewPipelineFromDSL([]byte(`{
|
||||
"dsl": {
|
||||
"components": {
|
||||
"begin": {"obj": {"component_name": "Begin", "params": {}}, "downstream": ["a"]},
|
||||
"a": {"obj": {"component_name": "`+nameA+`", "params": {}}, "upstream": ["begin"], "downstream": ["b"]},
|
||||
"b": {"obj": {"component_name": "`+nameB+`", "params": {}}, "upstream": ["a"]}
|
||||
},
|
||||
"path": ["begin", "a", "b"],
|
||||
"graph": {"nodes": []}
|
||||
}
|
||||
}`), taskID, WithCheckPointStore(store), WithRunTracker(tracker))
|
||||
if err != nil {
|
||||
t.Fatalf("NewPipelineFromDSL: %v", err)
|
||||
}
|
||||
|
||||
// Run 1: terminal B errors (oneShotErrStage n=1), simulating a crash.
|
||||
// Non-terminal A's checkpoint + interrupt persist because the error path
|
||||
// does not call ClearInterruptID or store.Delete.
|
||||
_, err = pipe.Run(context.Background(), map[string]any{"name": "doc-cross-run"})
|
||||
if err == nil {
|
||||
t.Fatal("Run 1: expected error from simulated crash, got nil")
|
||||
}
|
||||
if mockA.calls != 1 {
|
||||
t.Fatalf("Run 1: expected A to run once, got %d", mockA.calls)
|
||||
}
|
||||
// oneShotErrStage did not delegate to its embedded mock on the first call.
|
||||
if termStage.calls != 0 {
|
||||
t.Fatalf("Run 1: expected B (embedded mock) calls=0 (error before delegate), got %d", termStage.calls)
|
||||
}
|
||||
|
||||
// Run 2: resume from after A via tracker.GetInterruptID. A is skipped;
|
||||
// B's oneShotErrStage (n=2) delegates to its embedded mock successfully.
|
||||
_, err = pipe.Run(context.Background(), map[string]any{"name": "doc-cross-run"})
|
||||
if err != nil {
|
||||
t.Fatalf("Run 2: expected recovery success, got error: %v", err)
|
||||
}
|
||||
if mockA.calls != 1 {
|
||||
t.Fatalf("Run 2: expected A to still have calls=1 (was skipped by resume), got %d", mockA.calls)
|
||||
}
|
||||
if termStage.calls != 1 {
|
||||
t.Fatalf("Run 2: expected B (embedded mock) calls=1 (delegated once), got %d", termStage.calls)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPipelineRun_RequireResumeRejectsWithoutStore verifies that with
|
||||
// WithRequireResume set and no checkpoint store resolvable (no
|
||||
// injected store, no global Redis in unit scope), Run must refuse to start
|
||||
// and return ErrResumeUnavailable — a clear, distinguishable signal — rather
|
||||
// than silently degrading to a non-resumable runPlain. The reject fires
|
||||
|
||||
@@ -50,12 +50,26 @@ func TestProcessMessage_TaskNotFoundAcks(t *testing.T) {
|
||||
}
|
||||
|
||||
// TestProcessMessage_AlreadyCompletedAcks: a task already in a terminal state
|
||||
// (COMPLETED) is acked and skipped — no enqueue, no status change.
|
||||
// (COMPLETED) is acked and skipped — no enqueue, and the document is NOT
|
||||
// resurrected to RUNNING. A redelivered terminal task must not reset a
|
||||
// finished document's run status or counters.
|
||||
func TestProcessMessage_AlreadyCompletedAcks(t *testing.T) {
|
||||
db := testutil.SetupTestDB(t)
|
||||
cleanup := testutil.ReplaceDBForTest(t, db)
|
||||
defer cleanup()
|
||||
_, _, _, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1"))
|
||||
_, _, docID, taskID := testutil.SeedTestData(t, db, testutil.WithPipelineID("flow-1"))
|
||||
|
||||
// Mark the document as a finished parse: a non-RUNNING run status and
|
||||
// non-zero counters that StartRunning would clobber to "1"/0.
|
||||
finishedRun := "3"
|
||||
if err := db.Model(&entity.Document{}).Where("id = ?", docID).
|
||||
Updates(map[string]interface{}{
|
||||
"run": finishedRun,
|
||||
"chunk_num": 42,
|
||||
"token_num": 99,
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("set document finished sentinel: %v", err)
|
||||
}
|
||||
|
||||
// Set the task COMPLETED so the status switch skips it.
|
||||
if err := db.Model(&entity.IngestionTask{}).Where("id = ?", taskID).
|
||||
@@ -73,6 +87,18 @@ func TestProcessMessage_AlreadyCompletedAcks(t *testing.T) {
|
||||
if len(ingestor.taskChan) != 0 {
|
||||
t.Fatal("expected no task enqueued for completed task")
|
||||
}
|
||||
|
||||
// The finished document must be untouched - no resurrection to RUNNING.
|
||||
var doc entity.Document
|
||||
if err := db.Where("id = ?", docID).First(&doc).Error; err != nil {
|
||||
t.Fatalf("reload document: %v", err)
|
||||
}
|
||||
if doc.Run == nil || *doc.Run != finishedRun {
|
||||
t.Fatalf("document.run = %v, want %q (must not be resurrected to RUNNING)", doc.Run, finishedRun)
|
||||
}
|
||||
if doc.ChunkNum != 42 || doc.TokenNum != 99 {
|
||||
t.Fatalf("document counters changed: chunk_num=%d token_num=%d, want 42/99", doc.ChunkNum, doc.TokenNum)
|
||||
}
|
||||
}
|
||||
|
||||
// TestProcessMessage_ClaimFailsAcks: a RUNNING task whose claim fails
|
||||
|
||||
121
internal/parser/parser/xlsx_parse_method_test.go
Normal file
121
internal/parser/parser/xlsx_parse_method_test.go
Normal file
@@ -0,0 +1,121 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/xuri/excelize/v2"
|
||||
)
|
||||
|
||||
// TestNormalizeXLSXParseMethod verifies the parse_method normalization
|
||||
// shared by the XLSX/XLS/CSV parsers. "deepdoc" is the default
|
||||
// spreadsheet parse_method (see schema.ParserParam.Defaults and the
|
||||
// matching Python ParserParam), and the DSL templates ship "DeepDOC".
|
||||
// Both must normalize to "" so the default Excelize/CSV path is taken,
|
||||
// matching rag/flow/parser/parser.py:_spreadsheet which only special-cases
|
||||
// "tcadp parser" and routes everything else to the default parser.
|
||||
func TestNormalizeXLSXParseMethod(t *testing.T) {
|
||||
cases := []struct {
|
||||
in, want string
|
||||
}{
|
||||
{"", ""},
|
||||
{"deepdoc", ""},
|
||||
{"DeepDOC", ""}, // canonical casing used by DSL templates
|
||||
{"DEEPDOC", ""},
|
||||
{" deepdoc ", ""},
|
||||
{"deepdoc parser", ""},
|
||||
{"DeepDOC Parser", ""},
|
||||
{"tcadp parser", "tcadp"},
|
||||
{"TCADP Parser", "tcadp"},
|
||||
{"excelize", "excelize"},
|
||||
{"csv", "csv"},
|
||||
{"unknown", "unknown"}, // preserved so the switch rejects it
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := normalizeXLSXParseMethod(c.in); got != c.want {
|
||||
t.Errorf("normalizeXLSXParseMethod(%q) = %q, want %q", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestXLSXParser_DeepDocParseMethod verifies that both the lowercase "deepdoc"
|
||||
// and the uppercase "DeepDOC" (as shipped by the ingestion pipeline DSL templates)
|
||||
// parse_method values produce the default HTML table output.
|
||||
func TestXLSXParser_DeepDocParseMethod(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
method string
|
||||
cellValue string
|
||||
}{
|
||||
{"Lowercase", "deepdoc", "hello"},
|
||||
{"Uppercase", "DeepDOC", "world"},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
f := excelize.NewFile()
|
||||
defer f.Close()
|
||||
if err := f.SetCellValue("Sheet1", "A1", tc.cellValue); err != nil {
|
||||
t.Fatalf("SetCellValue: %v", err)
|
||||
}
|
||||
buf, err := f.WriteToBuffer()
|
||||
if err != nil {
|
||||
t.Fatalf("WriteToBuffer: %v", err)
|
||||
}
|
||||
|
||||
p, err := NewXLSXParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewXLSXParser: %v", err)
|
||||
}
|
||||
p.ConfigureFromSetup(map[string]any{"parse_method": tc.method})
|
||||
|
||||
res := p.ParseWithResult("test.xlsx", buf.Bytes())
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult(%s): %v", tc.method, res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "html"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if !strings.Contains(res.HTML, tc.cellValue) {
|
||||
t.Fatalf("HTML = %q, want it to contain cell content %q", res.HTML, tc.cellValue)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestCSVParser_DeepDocParseMethod asserts the CSV parser accepts the
|
||||
// default "deepdoc" parse_method and renders the default HTML table.
|
||||
func TestCSVParser_DeepDocParseMethod(t *testing.T) {
|
||||
p := NewCSVParser()
|
||||
p.ConfigureFromSetup(map[string]any{"parse_method": "deepdoc"})
|
||||
|
||||
res := p.ParseWithResult("test.csv", []byte("a,b\n1,2"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult(deepdoc): %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "html"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if !strings.Contains(res.HTML, "<table>") {
|
||||
t.Fatalf("HTML = %q, want a rendered <table>", res.HTML)
|
||||
}
|
||||
}
|
||||
|
||||
// TestXLSParser_DeepDocParseMethod_NoUnsupportedError asserts the XLS
|
||||
// parser no longer rejects "deepdoc". A real .xls blob is hard to
|
||||
// synthesize in a unit test, so we only assert that the error is not the
|
||||
// "unsupported XLS parse method" rejection; an open/parse failure from
|
||||
// the fake blob is acceptable. The shared normalization is already
|
||||
// covered by TestNormalizeXLSXParseMethod.
|
||||
func TestXLSParser_DeepDocParseMethod_NoUnsupportedError(t *testing.T) {
|
||||
p, err := NewXLSParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewXLSParser: %v", err)
|
||||
}
|
||||
p.ConfigureFromSetup(map[string]any{"parse_method": "deepdoc"})
|
||||
|
||||
res := p.ParseWithResult("test.xls", []byte("not a real xls"))
|
||||
if res.Err != nil && strings.Contains(res.Err.Error(), "unsupported XLS parse method") {
|
||||
t.Fatalf("deepdoc must not be rejected as unsupported: %v", res.Err)
|
||||
}
|
||||
}
|
||||
@@ -78,6 +78,15 @@ func normalizeXLSXParseMethod(raw string) string {
|
||||
if method == "tcadp parser" {
|
||||
return "tcadp"
|
||||
}
|
||||
// "deepdoc" / "deepdoc parser" is the default spreadsheet parse_method
|
||||
// (see schema.ParserParam.Defaults and the matching Python ParserParam),
|
||||
// and the DSL templates ship "DeepDOC". Normalize to "" so the default
|
||||
// Excelize/CSV path is taken by all three spreadsheet parsers, mirroring
|
||||
// rag/flow/parser/parser.py:_spreadsheet which only special-cases
|
||||
// "tcadp parser" and routes every other value to the default parser.
|
||||
if method == "deepdoc" || method == "deepdoc parser" {
|
||||
return ""
|
||||
}
|
||||
return method
|
||||
}
|
||||
|
||||
|
||||
@@ -17,28 +17,18 @@
|
||||
package chunk
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/csv"
|
||||
"encoding/json"
|
||||
"encoding/xml"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"image/draw"
|
||||
"image/jpeg"
|
||||
"io"
|
||||
"math"
|
||||
"math/rand"
|
||||
"path/filepath"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/engine/redis"
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/entity/models"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -61,11 +51,6 @@ import (
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
|
||||
const (
|
||||
maximumPageNumber = 100000
|
||||
maximumTaskPageNumber = maximumPageNumber * 1000
|
||||
)
|
||||
|
||||
var chunkImageMergeLocks = struct {
|
||||
sync.Mutex
|
||||
locks map[string]*chunkImageMergeLock
|
||||
@@ -97,20 +82,24 @@ type ChunkService struct {
|
||||
taskDAO *dao.TaskDAO
|
||||
searchService *service.SearchService
|
||||
|
||||
accessibleFunc func(string, string) bool
|
||||
getKnowledgebaseByIDFunc func(string) (*entity.Knowledgebase, error)
|
||||
getDocumentsByIDsFunc func([]string) ([]*entity.Document, error)
|
||||
getDocumentStorageAddressFunc func(*entity.Document) (string, string, error)
|
||||
queueParseTasksFunc func(*entity.Document, string, string, int64) error
|
||||
beginParseDocumentFunc func(string) error
|
||||
deleteTasksByDocIDsFunc func([]string) (int64, error)
|
||||
getEmbeddingModelFunc func(string, string) (*models.EmbeddingModel, error)
|
||||
incrementChunkStatsFunc func(string, string, int64, int64, float64) error
|
||||
decrementChunkStatsFunc func(string, string, int64, int64, float64) error
|
||||
storeChunkImageFunc func(string, string, []byte) error
|
||||
tokenizeFunc func(string) (string, error)
|
||||
fineGrainedTokenizeFunc func(string) (string, error)
|
||||
numTokensFunc func(string) int
|
||||
accessibleFunc func(string, string) bool
|
||||
getKnowledgebaseByIDFunc func(string) (*entity.Knowledgebase, error)
|
||||
getDocumentsByIDsFunc func([]string) ([]*entity.Document, error)
|
||||
// startParseDocumentsFunc overrides the DSL start-parse flow. Production
|
||||
// uses service.DocumentService.StartParseDocuments; tests inject a fake
|
||||
// to avoid the MQ publisher.
|
||||
startParseDocumentsFunc func(doc *entity.Document, kb *entity.Knowledgebase, userID string, opts service.StartParseOptions) error
|
||||
// cancelIngestionTaskFunc overrides the document-parsing cancellation.
|
||||
// Production uses service.DocumentService.CancelDocParse; tests inject
|
||||
// a fake to avoid the MQ publisher.
|
||||
cancelIngestionTaskFunc func(doc *entity.Document) error
|
||||
getEmbeddingModelFunc func(string, string) (*models.EmbeddingModel, error)
|
||||
incrementChunkStatsFunc func(string, string, int64, int64, float64) error
|
||||
decrementChunkStatsFunc func(string, string, int64, int64, float64) error
|
||||
storeChunkImageFunc func(string, string, []byte) error
|
||||
tokenizeFunc func(string) (string, error)
|
||||
fineGrainedTokenizeFunc func(string) (string, error)
|
||||
numTokensFunc func(string) int
|
||||
}
|
||||
|
||||
// NewChunkService creates chunk service
|
||||
@@ -611,26 +600,12 @@ const (
|
||||
docStopParsingInvalidStateErrorCode = "DOC_STOP_PARSING_INVALID_STATE"
|
||||
)
|
||||
|
||||
func (s *ChunkService) cancelAllTasksOfDoc(docID string) error {
|
||||
tasks, err := s.taskDAO.GetByDocID(docID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get tasks for document %s: %w", docID, err)
|
||||
func (s *ChunkService) cancelAllTasksOfDoc(doc *entity.Document) error {
|
||||
cancel := s.cancelIngestionTaskFunc
|
||||
if cancel == nil {
|
||||
cancel = service.NewDocumentService().CancelDocParse
|
||||
}
|
||||
|
||||
redisClient := redis.Get()
|
||||
if redisClient == nil {
|
||||
common.Warn(fmt.Sprintf("Redis unavailable; cannot cancel tasks for document %s", docID))
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, task := range tasks {
|
||||
if task == nil {
|
||||
continue
|
||||
}
|
||||
redisClient.Set(fmt.Sprintf("%s-cancel", task.ID), "x", 0)
|
||||
}
|
||||
|
||||
return nil
|
||||
return cancel(doc)
|
||||
}
|
||||
|
||||
func (s *ChunkService) StopParsing(userID, datasetID string, req service.StopParsingRequest) (*service.StopParsingResponse, common.ErrorCode, error) {
|
||||
@@ -642,15 +617,13 @@ func (s *ChunkService) StopParsing(userID, datasetID string, req service.StopPar
|
||||
return nil, common.CodeDataError, fmt.Errorf("`document_ids` is required")
|
||||
}
|
||||
|
||||
kb, err := s.kbDAO.GetByID(datasetID)
|
||||
_, err := s.kbDAO.GetByID(datasetID)
|
||||
if err != nil {
|
||||
return nil, common.CodeDataError, fmt.Errorf("You don't own the dataset %s", datasetID)
|
||||
}
|
||||
|
||||
docIDs, duplicateMessages := service.CheckDuplicateIDs(req.DocumentIDs, "document")
|
||||
successCount := 0
|
||||
ctx := context.Background()
|
||||
indexName := service.IndexName(kb.TenantID)
|
||||
|
||||
for _, docID := range docIDs {
|
||||
doc, err := s.documentDAO.GetByDocumentIDAndDatasetID(docID, datasetID)
|
||||
@@ -658,40 +631,21 @@ func (s *ChunkService) StopParsing(userID, datasetID string, req service.StopPar
|
||||
return nil, common.CodeDataError, fmt.Errorf("You don't own the document %s", docID)
|
||||
}
|
||||
|
||||
if doc.Run == nil || *doc.Run != string(entity.TaskStatusRunning) {
|
||||
task, _ := dao.NewIngestionTaskDAO().GetByDocumentID(docID)
|
||||
if task == nil || task.Status == common.COMPLETED ||
|
||||
task.Status == common.STOPPED || task.Status == common.FAILED {
|
||||
return &service.StopParsingResponse{
|
||||
Data: map[string]interface{}{"error_code": docStopParsingInvalidStateErrorCode},
|
||||
}, common.CodeDataError, fmt.Errorf("%s", docStopParsingInvalidStateMessage)
|
||||
}
|
||||
|
||||
if err := s.cancelAllTasksOfDoc(docID); err != nil {
|
||||
if err := s.cancelAllTasksOfDoc(doc); err != nil {
|
||||
return nil, common.CodeServerError, err
|
||||
}
|
||||
|
||||
updates := map[string]interface{}{
|
||||
"run": string(entity.TaskStatusCancel),
|
||||
"progress": 0,
|
||||
"chunk_num": 0,
|
||||
}
|
||||
if err := s.documentDAO.UpdateByID(doc.ID, updates); err != nil {
|
||||
return nil, common.CodeServerError, fmt.Errorf("failed to update document %s: %w", doc.ID, err)
|
||||
}
|
||||
|
||||
if s.docEngine != nil {
|
||||
exists, err := s.docEngine.ChunkStoreExists(ctx, indexName, datasetID)
|
||||
if err != nil {
|
||||
return nil, common.CodeServerError, fmt.Errorf("failed to check chunk store %s/%s: %w", indexName, datasetID, err)
|
||||
}
|
||||
if exists {
|
||||
if _, err := s.docEngine.DeleteChunks(ctx, map[string]interface{}{"doc_id": doc.ID}, indexName, datasetID); err != nil {
|
||||
return nil, common.CodeServerError, fmt.Errorf("failed to delete chunks for document %s: %w", doc.ID, err)
|
||||
}
|
||||
} else {
|
||||
common.Info(fmt.Sprintf("Skipping chunk delete during stop_parsing for doc %s: index %s/%s does not exist", doc.ID, indexName, datasetID))
|
||||
}
|
||||
} else {
|
||||
common.Info(fmt.Sprintf("Skipping chunk delete during stop_parsing for doc %s: index %s/%s does not exist", doc.ID, indexName, datasetID))
|
||||
}
|
||||
// CancelDocParse (inside cancelAllTasksOfDoc) already issues
|
||||
// RequestStop (STOPPING) and updates doc.run=CANCEL. Defer
|
||||
// destruction (chunk deletion, counter reset) until the worker
|
||||
// detects STOPPING and reaches a terminal state.
|
||||
|
||||
successCount++
|
||||
}
|
||||
@@ -729,184 +683,6 @@ func checkDuplicateIDs(documentIDs []string, idTypes string) ([]string, []string
|
||||
return uniqueDocIDs, duplicateMessages
|
||||
}
|
||||
|
||||
func (s *ChunkService) queueParseTasks(doc *entity.Document, bucket, objectName string, priority int64) error {
|
||||
if s.queueParseTasksFunc != nil {
|
||||
return s.queueParseTasksFunc(doc, bucket, objectName, priority)
|
||||
}
|
||||
tasks, err := s.buildParseTasks(doc, bucket, objectName, priority)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(tasks) == 0 {
|
||||
return nil
|
||||
}
|
||||
if err := dao.NewTaskDAO().CreateMany(tasks); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
queueName := s.parseQueueName(doc, priority)
|
||||
for _, task := range tasks {
|
||||
if task.Progress >= 1 {
|
||||
continue
|
||||
}
|
||||
message := parseTaskMessage(task)
|
||||
if ok := redis.Get().QueueProduct(queueName, message); !ok {
|
||||
if _, err := dao.NewTaskDAO().DeleteByDocIDs([]string{doc.ID}); err != nil {
|
||||
common.Warn("Failed to clean parse tasks after Redis enqueue failure",
|
||||
zap.String("docID", doc.ID),
|
||||
zap.Error(err))
|
||||
}
|
||||
return fmt.Errorf("Can't access Redis. Please check the Redis' status.")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *ChunkService) buildParseTasks(doc *entity.Document, bucket, objectName string, priority int64) ([]*entity.Task, error) {
|
||||
now := time.Now()
|
||||
ranges, err := s.parseTaskRanges(doc, bucket, objectName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tasks := make([]*entity.Task, 0, len(ranges))
|
||||
for _, pageRange := range ranges {
|
||||
taskID := utility.GenerateUUID()
|
||||
progressMsg := ""
|
||||
digest := s.parseTaskDigest(doc, pageRange.from, pageRange.to)
|
||||
chunkIDs := ""
|
||||
tasks = append(tasks, &entity.Task{
|
||||
ID: taskID,
|
||||
DocID: doc.ID,
|
||||
FromPage: pageRange.from,
|
||||
ToPage: pageRange.to,
|
||||
TaskType: "",
|
||||
Priority: priority,
|
||||
BeginAt: &now,
|
||||
Progress: 0,
|
||||
ProgressMsg: &progressMsg,
|
||||
Digest: &digest,
|
||||
ChunkIDs: &chunkIDs,
|
||||
})
|
||||
}
|
||||
return tasks, nil
|
||||
}
|
||||
|
||||
type parsePageRange struct {
|
||||
from int64
|
||||
to int64
|
||||
}
|
||||
|
||||
func (s *ChunkService) parseTaskRanges(doc *entity.Document, bucket, objectName string) ([]parsePageRange, error) {
|
||||
if doc.Type == "pdf" {
|
||||
return s.pdfParseTaskRanges(doc, bucket, objectName)
|
||||
}
|
||||
if doc.ParserID == string(entity.ParserTypeTable) {
|
||||
return s.tableParseTaskRanges(doc, bucket, objectName)
|
||||
}
|
||||
return []parsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil
|
||||
}
|
||||
|
||||
func (s *ChunkService) pdfParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]parsePageRange, error) {
|
||||
binary, err := s.getStorageBinary(bucket, objectName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pages := estimatePDFPageCount(binary)
|
||||
pageSize := int64(parserConfigInt(doc.ParserConfig, "task_page_size", 12))
|
||||
if doc.ParserID == string(entity.ParserTypePaper) {
|
||||
pageSize = int64(parserConfigInt(doc.ParserConfig, "task_page_size", 22))
|
||||
}
|
||||
if doc.ParserID == string(entity.ParserTypeOne) ||
|
||||
parserConfigString(doc.ParserConfig, "layout_recognize", "DeepDOC") != "DeepDOC" ||
|
||||
parserConfigBool(doc.ParserConfig, "toc_extraction", false) {
|
||||
pageSize = maximumTaskPageNumber
|
||||
}
|
||||
if pageSize <= 0 {
|
||||
pageSize = 12
|
||||
}
|
||||
|
||||
pageRanges := parserConfigPageRanges(doc.ParserConfig)
|
||||
ranges := make([]parsePageRange, 0)
|
||||
for _, configuredRange := range pageRanges {
|
||||
start := configuredRange.from - 1
|
||||
if start < 0 {
|
||||
start = 0
|
||||
}
|
||||
end := configuredRange.to - 1
|
||||
if pages >= 0 && end > pages {
|
||||
end = pages
|
||||
}
|
||||
for page := start; page < end; page += pageSize {
|
||||
to := page + pageSize
|
||||
if to > end {
|
||||
to = end
|
||||
}
|
||||
ranges = append(ranges, parsePageRange{from: page, to: to})
|
||||
}
|
||||
}
|
||||
if len(ranges) == 0 {
|
||||
ranges = append(ranges, parsePageRange{from: 0, to: maximumTaskPageNumber})
|
||||
}
|
||||
return ranges, nil
|
||||
}
|
||||
|
||||
func (s *ChunkService) tableParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]parsePageRange, error) {
|
||||
binary, err := s.getStorageBinary(bucket, objectName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows := estimateTableRowCount(docName(doc), binary)
|
||||
if rows <= 0 {
|
||||
return []parsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil
|
||||
}
|
||||
ranges := make([]parsePageRange, 0, (rows+2999)/3000)
|
||||
for row := int64(0); row < int64(rows); row += 3000 {
|
||||
to := row + 3000
|
||||
if to > int64(rows) {
|
||||
to = int64(rows)
|
||||
}
|
||||
ranges = append(ranges, parsePageRange{from: row, to: to})
|
||||
}
|
||||
return ranges, nil
|
||||
}
|
||||
|
||||
func (s *ChunkService) getStorageBinary(bucket, objectName string) ([]byte, error) {
|
||||
storageImpl := storage.GetStorageFactory().GetStorage()
|
||||
if storageImpl == nil {
|
||||
return nil, fmt.Errorf("storage not initialized")
|
||||
}
|
||||
return storageImpl.Get(bucket, objectName)
|
||||
}
|
||||
|
||||
func (s *ChunkService) beginParseDocument(docID string) error {
|
||||
if s.beginParseDocumentFunc != nil {
|
||||
return s.beginParseDocumentFunc(docID)
|
||||
}
|
||||
now := time.Now()
|
||||
return dao.GetDB().Model(&entity.Document{}).Where("id = ?", docID).Updates(map[string]interface{}{
|
||||
"progress_msg": "Task is queued...",
|
||||
"process_begin_at": now,
|
||||
"progress": rand.Float64() * 0.01,
|
||||
"run": string(entity.TaskStatusRunning),
|
||||
"chunk_num": 0,
|
||||
"token_num": 0,
|
||||
}).Error
|
||||
}
|
||||
|
||||
func (s *ChunkService) getDocumentStorageAddress(doc *entity.Document) (string, string, error) {
|
||||
if s.getDocumentStorageAddressFunc != nil {
|
||||
return s.getDocumentStorageAddressFunc(doc)
|
||||
}
|
||||
return service.NewDocumentService().GetDocumentStorageAddress(doc)
|
||||
}
|
||||
|
||||
func (s *ChunkService) deleteTasksByDocIDs(docIDs []string) (int64, error) {
|
||||
if s.deleteTasksByDocIDsFunc != nil {
|
||||
return s.deleteTasksByDocIDsFunc(docIDs)
|
||||
}
|
||||
return dao.NewTaskDAO().DeleteByDocIDs(docIDs)
|
||||
}
|
||||
|
||||
func (s *ChunkService) accessible(datasetID, userID string) bool {
|
||||
if s.accessibleFunc != nil {
|
||||
return s.accessibleFunc(datasetID, userID)
|
||||
@@ -928,282 +704,6 @@ func (s *ChunkService) getDocumentsByIDs(docIDs []string) ([]*entity.Document, e
|
||||
return s.documentDAO.GetByIDs(docIDs)
|
||||
}
|
||||
|
||||
func (s *ChunkService) parseQueueName(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 (s *ChunkService) parseTaskDigest(doc *entity.Document, fromPage, toPage int64) string {
|
||||
hasher := xxhash.New()
|
||||
config := chunkingConfigForDigest(doc)
|
||||
keys := make([]string, 0, len(config))
|
||||
for key := range config {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
for _, key := range keys {
|
||||
hasher.WriteString(stableString(config[key]))
|
||||
}
|
||||
hasher.WriteString(doc.ID)
|
||||
hasher.WriteString(strconv.FormatInt(fromPage, 10))
|
||||
hasher.WriteString(strconv.FormatInt(toPage, 10))
|
||||
return fmt.Sprintf("%x", hasher.Sum64())
|
||||
}
|
||||
|
||||
func parseTaskMessage(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 chunkingConfigForDigest(doc *entity.Document) map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"doc_id": doc.ID,
|
||||
"kb_id": doc.KbID,
|
||||
"parser_id": doc.ParserID,
|
||||
"parser_config": copyParserConfigForDigest(doc.ParserConfig),
|
||||
}
|
||||
}
|
||||
|
||||
func copyParserConfigForDigest(config map[string]interface{}) map[string]interface{} {
|
||||
copied := make(map[string]interface{}, len(config))
|
||||
for key, value := range config {
|
||||
if key == "raptor" || key == "graphrag" {
|
||||
continue
|
||||
}
|
||||
copied[key] = value
|
||||
}
|
||||
return copied
|
||||
}
|
||||
|
||||
func stableString(value interface{}) string {
|
||||
binary, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return fmt.Sprint(value)
|
||||
}
|
||||
return string(binary)
|
||||
}
|
||||
|
||||
func parserConfigInt(config map[string]interface{}, key string, fallback int) int {
|
||||
value, ok := config[key]
|
||||
if !ok || value == nil {
|
||||
return fallback
|
||||
}
|
||||
switch typedValue := value.(type) {
|
||||
case int:
|
||||
return typedValue
|
||||
case int64:
|
||||
return int(typedValue)
|
||||
case float64:
|
||||
return int(typedValue)
|
||||
case json.Number:
|
||||
if intValue, err := typedValue.Int64(); err == nil {
|
||||
return int(intValue)
|
||||
}
|
||||
case string:
|
||||
if intValue, err := strconv.Atoi(strings.TrimSpace(typedValue)); err == nil {
|
||||
return intValue
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func parserConfigString(config map[string]interface{}, key, fallback string) string {
|
||||
value, ok := config[key]
|
||||
if !ok || value == nil {
|
||||
return fallback
|
||||
}
|
||||
if stringValue, ok := value.(string); ok {
|
||||
return stringValue
|
||||
}
|
||||
return fmt.Sprint(value)
|
||||
}
|
||||
|
||||
func parserConfigBool(config map[string]interface{}, key string, fallback bool) bool {
|
||||
value, ok := config[key]
|
||||
if !ok || value == nil {
|
||||
return fallback
|
||||
}
|
||||
switch typedValue := value.(type) {
|
||||
case bool:
|
||||
return typedValue
|
||||
case string:
|
||||
switch strings.ToLower(strings.TrimSpace(typedValue)) {
|
||||
case "true", "1", "yes", "on":
|
||||
return true
|
||||
case "false", "0", "no", "off":
|
||||
return false
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func parserConfigPageRanges(config map[string]interface{}) []parsePageRange {
|
||||
defaultRanges := []parsePageRange{{from: 1, to: maximumPageNumber}}
|
||||
raw, ok := config["pages"]
|
||||
if !ok || raw == nil {
|
||||
return defaultRanges
|
||||
}
|
||||
rawRanges, ok := raw.([]interface{})
|
||||
if !ok || len(rawRanges) == 0 {
|
||||
return defaultRanges
|
||||
}
|
||||
|
||||
ranges := make([]parsePageRange, 0, len(rawRanges))
|
||||
for _, rawRange := range rawRanges {
|
||||
rangeValues, ok := rawRange.([]interface{})
|
||||
if !ok || len(rangeValues) < 2 {
|
||||
continue
|
||||
}
|
||||
from, okFrom := toInt64(rangeValues[0])
|
||||
to, okTo := toInt64(rangeValues[1])
|
||||
if okFrom && okTo && to > from {
|
||||
ranges = append(ranges, parsePageRange{from: from, to: to})
|
||||
}
|
||||
}
|
||||
if len(ranges) == 0 {
|
||||
return defaultRanges
|
||||
}
|
||||
return ranges
|
||||
}
|
||||
|
||||
func toInt64(value interface{}) (int64, bool) {
|
||||
switch typedValue := value.(type) {
|
||||
case int:
|
||||
return int64(typedValue), true
|
||||
case int64:
|
||||
return typedValue, true
|
||||
case float64:
|
||||
return int64(typedValue), true
|
||||
case json.Number:
|
||||
intValue, err := typedValue.Int64()
|
||||
return intValue, err == nil
|
||||
case string:
|
||||
intValue, err := strconv.ParseInt(strings.TrimSpace(typedValue), 10, 64)
|
||||
return intValue, err == nil
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
var pdfPagePattern = regexp.MustCompile(`/Type\s*/Page\b`)
|
||||
|
||||
func estimatePDFPageCount(binary []byte) int64 {
|
||||
if len(binary) == 0 {
|
||||
return 0
|
||||
}
|
||||
return int64(len(pdfPagePattern.FindAll(binary, -1)))
|
||||
}
|
||||
|
||||
func estimateTableRowCount(name string, binary []byte) int {
|
||||
switch strings.ToLower(filepath.Ext(name)) {
|
||||
case ".xlsx":
|
||||
if rows, err := countXLSXRows(binary); err == nil {
|
||||
return rows
|
||||
}
|
||||
case ".csv", ".tsv", ".txt":
|
||||
return countDelimitedRows(name, binary)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func countDelimitedRows(name string, binary []byte) int {
|
||||
reader := csv.NewReader(bytes.NewReader(binary))
|
||||
reader.FieldsPerRecord = -1
|
||||
reader.ReuseRecord = true
|
||||
if strings.EqualFold(filepath.Ext(name), ".tsv") {
|
||||
reader.Comma = '\t'
|
||||
}
|
||||
rows := 0
|
||||
for {
|
||||
_, err := reader.Read()
|
||||
if err == nil {
|
||||
rows++
|
||||
continue
|
||||
}
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
rows += bytes.Count(binary, []byte{'\n'})
|
||||
if len(binary) > 0 && binary[len(binary)-1] != '\n' {
|
||||
rows++
|
||||
}
|
||||
break
|
||||
}
|
||||
return rows
|
||||
}
|
||||
|
||||
func countXLSXRows(binary []byte) (int, error) {
|
||||
zipReader, err := zip.NewReader(bytes.NewReader(binary), int64(len(binary)))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
maxRows := 0
|
||||
for _, file := range zipReader.File {
|
||||
if !strings.HasPrefix(file.Name, "xl/worksheets/") || !strings.HasSuffix(file.Name, ".xml") {
|
||||
continue
|
||||
}
|
||||
rows, err := countWorksheetRows(file)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if rows > maxRows {
|
||||
maxRows = rows
|
||||
}
|
||||
}
|
||||
return maxRows, nil
|
||||
}
|
||||
|
||||
func countWorksheetRows(file *zip.File) (int, error) {
|
||||
reader, err := file.Open()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer reader.Close()
|
||||
|
||||
decoder := xml.NewDecoder(reader)
|
||||
rows := 0
|
||||
for {
|
||||
token, err := decoder.Token()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
start, ok := token.(xml.StartElement)
|
||||
if ok && start.Name.Local == "row" {
|
||||
rows++
|
||||
}
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
func docName(doc *entity.Document) string {
|
||||
if doc.Name == nil {
|
||||
return ""
|
||||
}
|
||||
return *doc.Name
|
||||
}
|
||||
|
||||
func (s *ChunkService) Parse(userID, datasetID string, req *service.ParseFileRequest) (map[string]interface{}, common.ErrorCode, error) {
|
||||
if !s.accessible(datasetID, userID) {
|
||||
return nil, common.CodeOperatingError, fmt.Errorf("You don't own the dataset %s.", datasetID)
|
||||
@@ -1244,34 +744,23 @@ func (s *ChunkService) Parse(userID, datasetID string, req *service.ParseFileReq
|
||||
}
|
||||
}
|
||||
|
||||
// Batch pre-check: refuse the whole request if any document's ingestion
|
||||
// task is non-terminal (RUNNING/STOPPING), so we never partially clean.
|
||||
if err := (service.NewDocumentService().AssertIngestionTasksTerminal(docIDs)); err != nil {
|
||||
return nil, common.CodeDataError, err
|
||||
}
|
||||
|
||||
startParse := s.startParseDocumentsFunc
|
||||
if startParse == nil {
|
||||
docSvc := service.NewDocumentService()
|
||||
startParse = docSvc.StartParseDocuments
|
||||
}
|
||||
|
||||
successCount := 0
|
||||
|
||||
for _, docID := range docIDs {
|
||||
doc := docByID[docID]
|
||||
|
||||
if s.docEngine != nil {
|
||||
indexName := fmt.Sprintf("ragflow_%s", kb.TenantID)
|
||||
if _, err := s.docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": docID}, indexName, datasetID); err != nil {
|
||||
return nil, common.CodeServerError, err
|
||||
}
|
||||
}
|
||||
if _, err := s.deleteTasksByDocIDs([]string{docID}); err != nil {
|
||||
return nil, common.CodeServerError, err
|
||||
}
|
||||
|
||||
bucket, objectName, err := s.getDocumentStorageAddress(doc)
|
||||
if err != nil {
|
||||
return nil, common.CodeServerError, err
|
||||
}
|
||||
if err := s.queueParseTasks(doc, bucket, objectName, 0); err != nil {
|
||||
return nil, common.CodeServerError, err
|
||||
}
|
||||
if err := s.beginParseDocument(doc.ID); err != nil {
|
||||
if _, delErr := s.deleteTasksByDocIDs([]string{doc.ID}); delErr != nil {
|
||||
common.Warn("Failed to clean parse tasks after document state update failure",
|
||||
zap.String("docID", doc.ID),
|
||||
zap.Error(delErr))
|
||||
}
|
||||
if err := startParse(doc, kb, userID, service.StartParseOptions{RerunWithDelete: true}); err != nil {
|
||||
return nil, common.CodeServerError, err
|
||||
}
|
||||
successCount++
|
||||
|
||||
@@ -274,200 +274,6 @@ func TestParseRejectsRunningDocument(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseReturnsServerErrorWhenDeleteChunksFails(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
userID, datasetID := "user-1", "kb-1"
|
||||
insertChunkTestKB(t, datasetID, userID)
|
||||
insertChunkTestDoc(t, "doc-1", datasetID)
|
||||
|
||||
deleteErr := errors.New("delete chunks failed")
|
||||
engine := &parseTestDocEngine{deleteChunksErr: deleteErr}
|
||||
svc := newParseTestService(t)
|
||||
svc.docEngine = engine
|
||||
|
||||
_, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
|
||||
if !errors.Is(err, deleteErr) {
|
||||
t.Fatalf("expected delete chunks error, got %v", err)
|
||||
}
|
||||
if code != common.CodeServerError {
|
||||
t.Fatalf("expected CodeServerError, got %v", code)
|
||||
}
|
||||
if engine.deleteChunksCalls != 1 {
|
||||
t.Fatalf("expected DeleteChunks once, got %d", engine.deleteChunksCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseReturnsServerErrorWhenDeleteTasksFails(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
userID, datasetID := "user-1", "kb-1"
|
||||
insertChunkTestKB(t, datasetID, userID)
|
||||
insertChunkTestDoc(t, "doc-1", datasetID)
|
||||
|
||||
deleteErr := errors.New("delete tasks failed")
|
||||
svc := newParseTestService(t)
|
||||
svc.deleteTasksByDocIDsFunc = func([]string) (int64, error) { return 0, deleteErr }
|
||||
|
||||
_, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
|
||||
if !errors.Is(err, deleteErr) {
|
||||
t.Fatalf("expected delete tasks error, got %v", err)
|
||||
}
|
||||
if code != common.CodeServerError {
|
||||
t.Fatalf("expected CodeServerError, got %v", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseReturnsServerErrorWhenStorageAddressFails(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
userID, datasetID := "user-1", "kb-1"
|
||||
insertChunkTestKB(t, datasetID, userID)
|
||||
insertChunkTestDoc(t, "doc-1", datasetID)
|
||||
|
||||
storageErr := errors.New("storage address failed")
|
||||
svc := newParseTestService(t)
|
||||
svc.getDocumentStorageAddressFunc = func(*entity.Document) (string, string, error) {
|
||||
return "", "", storageErr
|
||||
}
|
||||
|
||||
_, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
|
||||
if !errors.Is(err, storageErr) {
|
||||
t.Fatalf("expected storage error, got %v", err)
|
||||
}
|
||||
if code != common.CodeServerError {
|
||||
t.Fatalf("expected CodeServerError, got %v", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseReturnsServerErrorWhenQueueFails(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
userID, datasetID := "user-1", "kb-1"
|
||||
insertChunkTestKB(t, datasetID, userID)
|
||||
insertChunkTestDoc(t, "doc-1", datasetID)
|
||||
|
||||
queueErr := errors.New("queue failed")
|
||||
svc := newParseTestService(t)
|
||||
svc.queueParseTasksFunc = func(*entity.Document, string, string, int64) error {
|
||||
return queueErr
|
||||
}
|
||||
|
||||
_, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
|
||||
if !errors.Is(err, queueErr) {
|
||||
t.Fatalf("expected queue error, got %v", err)
|
||||
}
|
||||
if code != common.CodeServerError {
|
||||
t.Fatalf("expected CodeServerError, got %v", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseCleansTasksWhenBeginParseFails(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
userID, datasetID := "user-1", "kb-1"
|
||||
insertChunkTestKB(t, datasetID, userID)
|
||||
insertChunkTestDoc(t, "doc-1", datasetID)
|
||||
|
||||
beginErr := errors.New("begin parse failed")
|
||||
deleteCalls := 0
|
||||
svc := newParseTestService(t)
|
||||
svc.beginParseDocumentFunc = func(string) error { return beginErr }
|
||||
svc.deleteTasksByDocIDsFunc = func([]string) (int64, error) {
|
||||
deleteCalls++
|
||||
return 1, nil
|
||||
}
|
||||
|
||||
_, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
|
||||
if !errors.Is(err, beginErr) {
|
||||
t.Fatalf("expected begin error, got %v", err)
|
||||
}
|
||||
if code != common.CodeServerError {
|
||||
t.Fatalf("expected CodeServerError, got %v", code)
|
||||
}
|
||||
if deleteCalls != 2 {
|
||||
t.Fatalf("expected initial delete and cleanup delete, got %d calls", deleteCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseReturnsPartialSuccessForDuplicateDocumentIDs(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
userID, datasetID := "user-1", "kb-1"
|
||||
insertChunkTestKB(t, datasetID, userID)
|
||||
insertChunkTestDoc(t, "doc-1", datasetID)
|
||||
|
||||
svc := newParseTestService(t)
|
||||
result, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{
|
||||
DocumentIDs: []string{"doc-1", "doc-1"},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected duplicate warning error")
|
||||
}
|
||||
if code != common.CodeSuccess {
|
||||
t.Fatalf("expected CodeSuccess, got %v", code)
|
||||
}
|
||||
if result["success_count"] != 1 {
|
||||
t.Fatalf("expected one successful parse, got %v", result["success_count"])
|
||||
}
|
||||
errorsValue, ok := result["errors"].([]string)
|
||||
if !ok || len(errorsValue) != 1 || !strings.Contains(errorsValue[0], "Duplicate document ids: doc-1") {
|
||||
t.Fatalf("unexpected duplicate errors: %#v", result["errors"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseQueuesAndBeginsDocument(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
userID, datasetID := "user-1", "kb-1"
|
||||
insertChunkTestKB(t, datasetID, userID)
|
||||
insertChunkTestDoc(t, "doc-1", datasetID)
|
||||
insertChunkTestTask(t, "task-1", "doc-1")
|
||||
|
||||
queueCalls := 0
|
||||
svc := newParseTestService(t)
|
||||
svc.queueParseTasksFunc = func(doc *entity.Document, bucket, objectName string, priority int64) error {
|
||||
queueCalls++
|
||||
if doc.ID != "doc-1" || bucket != datasetID || objectName != "doc-1.txt" || priority != 0 {
|
||||
t.Fatalf("unexpected queue args: doc=%s bucket=%s object=%s priority=%d", doc.ID, bucket, objectName, priority)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
result, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
|
||||
if err != nil {
|
||||
t.Fatalf("expected parse success, got %v", err)
|
||||
}
|
||||
if code != common.CodeSuccess {
|
||||
t.Fatalf("expected CodeSuccess, got %v", code)
|
||||
}
|
||||
if result != nil {
|
||||
t.Fatalf("expected nil result, got %#v", result)
|
||||
}
|
||||
if queueCalls != 1 {
|
||||
t.Fatalf("expected queue once, got %d", queueCalls)
|
||||
}
|
||||
|
||||
var taskCount int64
|
||||
if err := dao.DB.Model(&entity.Task{}).Where("doc_id = ?", "doc-1").Count(&taskCount).Error; err != nil {
|
||||
t.Fatalf("count tasks: %v", err)
|
||||
}
|
||||
if taskCount != 0 {
|
||||
t.Fatalf("expected old tasks to be deleted, got %d", taskCount)
|
||||
}
|
||||
|
||||
doc, err := dao.NewDocumentDAO().GetByID("doc-1")
|
||||
if err != nil {
|
||||
t.Fatalf("get doc: %v", err)
|
||||
}
|
||||
if doc.Run == nil || *doc.Run != string(entity.TaskStatusRunning) {
|
||||
t.Fatalf("expected document to be running, got %v", doc.Run)
|
||||
}
|
||||
if doc.ChunkNum != 0 {
|
||||
t.Fatalf("expected chunk_num reset to 0, got %d", doc.ChunkNum)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAddChunkSuccess(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
@@ -961,6 +767,7 @@ func setupChunkTestDB(t *testing.T) *gorm.DB {
|
||||
&entity.Knowledgebase{},
|
||||
&entity.Document{},
|
||||
&entity.Task{},
|
||||
&entity.IngestionTask{},
|
||||
); err != nil {
|
||||
t.Fatalf("failed to migrate: %v", err)
|
||||
}
|
||||
@@ -984,15 +791,6 @@ func newParseTestService(t *testing.T) *ChunkService {
|
||||
docEngine: nil,
|
||||
kbDAO: dao.NewKnowledgebaseDAO(),
|
||||
documentDAO: dao.NewDocumentDAO(),
|
||||
getDocumentStorageAddressFunc: func(doc *entity.Document) (string, string, error) {
|
||||
if doc.Location == nil {
|
||||
return doc.KbID, "", nil
|
||||
}
|
||||
return doc.KbID, *doc.Location, nil
|
||||
},
|
||||
queueParseTasksFunc: func(*entity.Document, string, string, int64) error {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1057,6 +855,21 @@ func insertChunkTestDoc(t *testing.T, id, kbID string) {
|
||||
}
|
||||
}
|
||||
|
||||
func insertChunkTestIngestionTask(t *testing.T, id, userID, docID, datasetID, status string) {
|
||||
t.Helper()
|
||||
task := &entity.IngestionTask{
|
||||
ID: id,
|
||||
UserID: userID,
|
||||
DocumentID: docID,
|
||||
DatasetID: datasetID,
|
||||
Schema: entity.JSONMap{},
|
||||
Status: status,
|
||||
}
|
||||
if err := dao.DB.Create(task).Error; err != nil {
|
||||
t.Fatalf("insert ingestion task: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func setChunkTestStats(t *testing.T, docID, kbID string, docTokenNum, docChunkNum, kbTokenNum, kbChunkNum int64) {
|
||||
t.Helper()
|
||||
|
||||
@@ -1093,6 +906,7 @@ type parseTestDocEngine struct {
|
||||
deleteChunksCondition map[string]interface{}
|
||||
deleteIndexName string
|
||||
deleteDatasetID string
|
||||
chunkStoreExists bool // if true, ChunkStoreExists returns true
|
||||
}
|
||||
|
||||
func (e *parseTestDocEngine) CreateChunkStore(context.Context, string, string, int, string) error {
|
||||
@@ -1128,7 +942,7 @@ func (e *parseTestDocEngine) DropChunkStore(context.Context, string, string) err
|
||||
}
|
||||
|
||||
func (e *parseTestDocEngine) ChunkStoreExists(context.Context, string, string) (bool, error) {
|
||||
return false, nil
|
||||
return e.chunkStoreExists, nil
|
||||
}
|
||||
|
||||
func (e *parseTestDocEngine) CreateMetadataStore(context.Context, string) error {
|
||||
@@ -1517,3 +1331,175 @@ func copyMap(in map[string]interface{}) map[string]interface{} {
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestChunkServiceParse_RejectsBatchWithRunningIngestionTask(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
insertChunkTestUserTenant(t, "user-1", "tenant-1")
|
||||
insertChunkTestKB(t, "kb-1", "tenant-1")
|
||||
insertChunkTestDoc(t, "doc-1", "kb-1")
|
||||
insertChunkTestDoc(t, "doc-2", "kb-1")
|
||||
insertChunkTestIngestionTask(t, "task-2", "user-1", "doc-2", "kb-1", common.RUNNING)
|
||||
|
||||
svc := newParseTestService(t)
|
||||
svc.accessibleFunc = func(string, string) bool { return true }
|
||||
_, _, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1", "doc-2"}})
|
||||
if err == nil {
|
||||
t.Fatal("expected error for RUNNING ingestion task, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkServiceParse_CallsStartParseDocumentsWithRerunWithDelete(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
insertChunkTestUserTenant(t, "user-1", "tenant-1")
|
||||
insertChunkTestKB(t, "kb-1", "tenant-1")
|
||||
insertChunkTestDoc(t, "doc-1", "kb-1")
|
||||
|
||||
svc := newParseTestService(t)
|
||||
svc.accessibleFunc = func(string, string) bool { return true }
|
||||
var calledOpts service.StartParseOptions
|
||||
var calledDocID string
|
||||
svc.startParseDocumentsFunc = func(doc *entity.Document, kb *entity.Knowledgebase, userID string, opts service.StartParseOptions) error {
|
||||
calledOpts = opts
|
||||
calledDocID = doc.ID
|
||||
return nil
|
||||
}
|
||||
|
||||
_, code, err := svc.Parse("user-1", "kb-1", &service.ParseFileRequest{DocumentIDs: []string{"doc-1"}})
|
||||
if err != nil {
|
||||
t.Fatalf("Parse: %v", err)
|
||||
}
|
||||
if code != common.CodeSuccess {
|
||||
t.Fatalf("code = %d, want success", code)
|
||||
}
|
||||
if calledDocID != "doc-1" {
|
||||
t.Fatalf("startParseDocuments called with doc %s, want doc-1", calledDocID)
|
||||
}
|
||||
if !calledOpts.RerunWithDelete {
|
||||
t.Fatal("expected RerunWithDelete=true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkServiceParse_ReturnsPartialSuccessForDuplicateDocumentIDs(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
userID, datasetID := "user-1", "kb-1"
|
||||
insertChunkTestKB(t, datasetID, userID)
|
||||
insertChunkTestDoc(t, "doc-1", datasetID)
|
||||
|
||||
svc := newParseTestService(t)
|
||||
svc.startParseDocumentsFunc = func(*entity.Document, *entity.Knowledgebase, string, service.StartParseOptions) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
result, code, err := svc.Parse(userID, datasetID, &service.ParseFileRequest{
|
||||
DocumentIDs: []string{"doc-1", "doc-1"},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected duplicate warning error")
|
||||
}
|
||||
if code != common.CodeSuccess {
|
||||
t.Fatalf("expected CodeSuccess, got %v", code)
|
||||
}
|
||||
if result["success_count"] != 1 {
|
||||
t.Fatalf("expected one successful parse, got %v", result["success_count"])
|
||||
}
|
||||
errorsValue, ok := result["errors"].([]string)
|
||||
if !ok || len(errorsValue) != 1 || !strings.Contains(errorsValue[0], "Duplicate document ids: doc-1") {
|
||||
t.Fatalf("unexpected duplicate errors: %#v", result["errors"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopParsing_CallsCancelIngestionTask(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
insertChunkTestKB(t, "kb-1", "user-1") // owner = user-1, Accessible passes
|
||||
insertChunkTestDoc(t, "doc-1", "kb-1")
|
||||
// StopParsing now checks the IngestionTask status (not doc.Run).
|
||||
insertChunkTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1", common.RUNNING)
|
||||
|
||||
svc := newParseTestService(t)
|
||||
var calledDocID string
|
||||
svc.cancelIngestionTaskFunc = func(doc *entity.Document) error {
|
||||
calledDocID = doc.ID
|
||||
return nil
|
||||
}
|
||||
|
||||
_, _, err := svc.StopParsing("user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}})
|
||||
if err != nil {
|
||||
t.Fatalf("StopParsing: %v", err)
|
||||
}
|
||||
if calledDocID != "doc-1" {
|
||||
t.Fatalf("cancelIngestionTaskFunc called with doc %s, want doc-1", calledDocID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopParsing_RejectsDocumentWithoutIngestionTask(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
insertChunkTestKB(t, "kb-1", "user-1")
|
||||
insertChunkTestDoc(t, "doc-1", "kb-1")
|
||||
// No IngestionTask — should be rejected (nothing to stop).
|
||||
|
||||
svc := newParseTestService(t)
|
||||
var called bool
|
||||
svc.cancelIngestionTaskFunc = func(doc *entity.Document) error { called = true; return nil }
|
||||
|
||||
_, code, err := svc.StopParsing("user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}})
|
||||
if err == nil {
|
||||
t.Fatal("expected error for document without ingestion task")
|
||||
}
|
||||
if code != common.CodeDataError {
|
||||
t.Fatalf("expected CodeDataError, got %v", code)
|
||||
}
|
||||
if called {
|
||||
t.Fatal("cancel should not be called for doc without ingestion task")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopParsing_DoesNotDeleteChunksOrResetCountersAfterCancel(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
pushChunkTestDB(t, db)
|
||||
insertChunkTestKB(t, "kb-1", "user-1")
|
||||
insertChunkTestDoc(t, "doc-1", "kb-1")
|
||||
// Pre-set counters on the doc so we can verify they are NOT reset.
|
||||
if err := dao.DB.Model(&entity.Document{}).Where("id = ?", "doc-1").Updates(map[string]interface{}{
|
||||
"token_num": int64(10),
|
||||
"chunk_num": int64(5),
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("set counters: %v", err)
|
||||
}
|
||||
insertChunkTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1", common.RUNNING)
|
||||
|
||||
svc := newParseTestService(t)
|
||||
svc.cancelIngestionTaskFunc = func(doc *entity.Document) error {
|
||||
// Simulate CancelDocParse: set doc.run=CANCEL.
|
||||
return dao.DB.Model(&entity.Document{}).Where("id = ?", doc.ID).
|
||||
Update("run", string(entity.TaskStatusCancel)).Error
|
||||
}
|
||||
engine := &parseTestDocEngine{chunkStoreExists: true}
|
||||
svc.docEngine = engine
|
||||
|
||||
_, _, err := svc.StopParsing("user-1", "kb-1", service.StopParsingRequest{DocumentIDs: []string{"doc-1"}})
|
||||
if err != nil {
|
||||
t.Fatalf("StopParsing: %v", err)
|
||||
}
|
||||
|
||||
// Counters must NOT be reset.
|
||||
doc, err := dao.NewDocumentDAO().GetByID("doc-1")
|
||||
if err != nil {
|
||||
t.Fatalf("reload doc: %v", err)
|
||||
}
|
||||
if doc.TokenNum != 10 {
|
||||
t.Fatalf("token_num = %d, want 10 (not reset)", doc.TokenNum)
|
||||
}
|
||||
if doc.ChunkNum != 5 {
|
||||
t.Fatalf("chunk_num = %d, want 5 (not reset)", doc.ChunkNum)
|
||||
}
|
||||
|
||||
// Chunks must NOT be deleted during cancel.
|
||||
if engine.deleteChunksCalls != 0 {
|
||||
t.Fatalf("DeleteChunks called %d times, want 0", engine.deleteChunksCalls)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -17,17 +17,12 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"encoding/json"
|
||||
"encoding/xml"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"math/rand"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
@@ -40,16 +35,13 @@ import (
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/deepdoc/parser/pdf/pdfoxide"
|
||||
"ragflow/internal/engine"
|
||||
"ragflow/internal/engine/redis"
|
||||
enginetypes "ragflow/internal/engine/types"
|
||||
"ragflow/internal/entity"
|
||||
"ragflow/internal/storage"
|
||||
"ragflow/internal/tokenizer"
|
||||
"ragflow/internal/utility"
|
||||
|
||||
"github.com/cespare/xxhash/v2"
|
||||
"go.uber.org/zap"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
@@ -721,7 +713,7 @@ func (s *DocumentService) deleteDocumentFull(docID string) error {
|
||||
// Delete tasks from DB
|
||||
ingestionTask, err := s.ingestionTaskDAO.GetByDocumentID(docID)
|
||||
if err != nil {
|
||||
common.Error(fmt.Sprintf("failed to get ingestion task by doc:%s", docID), err)
|
||||
common.Error(fmt.Sprintf("failed to get ingestion task by doc:%s", doc.ID), err)
|
||||
return err
|
||||
}
|
||||
if ingestionTask != nil {
|
||||
@@ -1200,9 +1192,63 @@ type IngestDocumentRequest struct {
|
||||
ApplyKB bool `json:"apply_kb"`
|
||||
}
|
||||
|
||||
type documentParsePageRange struct {
|
||||
from int64
|
||||
to int64
|
||||
// StartParseOptions controls StartParseDocuments behavior.
|
||||
type StartParseOptions struct {
|
||||
// ApplyKB merges the knowledgebase's parser_config (llm_id, metadata)
|
||||
// into the document before parsing.
|
||||
ApplyKB bool
|
||||
// RerunWithDelete clears prior chunks/tasks/counters before re-parsing.
|
||||
RerunWithDelete bool
|
||||
}
|
||||
|
||||
// StartParseDocuments starts parsing a document via the DSL ingestion
|
||||
// pipeline. It optionally clears prior results (RerunWithDelete), applies
|
||||
// KB config (ApplyKB), validates storage, and enqueues an ingestion task.
|
||||
// The document run status is NOT set here; IngestionTaskService.StartRunning
|
||||
// sets it to RUNNING when the worker picks up the task and transitions it from
|
||||
// CREATED. Extracted from Ingest so other entry points (e.g. ChunkService.Parse)
|
||||
// can reuse the same start-parse flow.
|
||||
func (s *DocumentService) StartParseDocuments(doc *entity.Document, kb *entity.Knowledgebase, userID string, opts StartParseOptions) error {
|
||||
// Validate storage first so we don't clear prior results and then fail
|
||||
// because the document can't be read, leaving the document with neither
|
||||
// old nor new parse results.
|
||||
if _, _, err := s.GetDocumentStorageAddress(doc); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if opts.RerunWithDelete {
|
||||
if err := s.clearDocumentParseResults(doc, kb.TenantID); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
if opts.ApplyKB {
|
||||
if doc.ParserConfig == nil {
|
||||
doc.ParserConfig = entity.JSONMap{}
|
||||
}
|
||||
config := map[string]interface{}{
|
||||
"llm_id": kb.ParserConfig["llm_id"],
|
||||
"enable_metadata": false,
|
||||
"metadata": map[string]interface{}{},
|
||||
}
|
||||
if value, ok := kb.ParserConfig["enable_metadata"]; ok {
|
||||
config["enable_metadata"] = value
|
||||
}
|
||||
if value, ok := kb.ParserConfig["metadata"]; ok {
|
||||
config["metadata"] = value
|
||||
}
|
||||
if err := s.updateDocumentParserConfig(doc.ID, config); err != nil {
|
||||
return err
|
||||
}
|
||||
for key, value := range config {
|
||||
doc.ParserConfig[key] = value
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := s.IngestDocuments(doc.KbID, userID, []string{doc.ID}); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (common.ErrorCode, error) {
|
||||
@@ -1220,150 +1266,148 @@ func (s *DocumentService) Ingest(userID string, req *IngestDocumentRequest) (com
|
||||
}
|
||||
}
|
||||
|
||||
tableDoneCountByKB := make(map[string]int64)
|
||||
|
||||
// First pass: validate every document exists and is accessible before
|
||||
// mutating any state, so a single invalid doc rejects the whole request.
|
||||
type validatedDoc struct {
|
||||
doc *entity.Document
|
||||
kb *entity.Knowledgebase
|
||||
}
|
||||
validated := make([]validatedDoc, 0, len(req.DocIDs))
|
||||
validatedIDs := make([]string, 0, len(req.DocIDs))
|
||||
for _, docID := range req.DocIDs {
|
||||
doc := docsByID[docID]
|
||||
if doc == nil {
|
||||
return common.CodeDataError, fmt.Errorf("document not found")
|
||||
}
|
||||
|
||||
kb, err := s.kbDAO.GetByID(doc.KbID)
|
||||
if err != nil {
|
||||
return common.CodeDataError, fmt.Errorf("dataset not found")
|
||||
}
|
||||
|
||||
if !s.kbDAO.Accessible(kb.ID, userID) {
|
||||
return common.CodeAuthenticationError, fmt.Errorf("no authorization")
|
||||
}
|
||||
validated = append(validated, validatedDoc{doc, kb})
|
||||
validatedIDs = append(validatedIDs, docID)
|
||||
}
|
||||
|
||||
updates := map[string]interface{}{
|
||||
"run": run,
|
||||
"progress": 0,
|
||||
// Batch pre-check for re-parse with delete: use the validated doc IDs
|
||||
// so we don't silently skip non-existent or unauthorized documents.
|
||||
if run == string(entity.TaskStatusRunning) && req.Delete {
|
||||
if err := s.AssertIngestionTasksTerminal(validatedIDs); err != nil {
|
||||
return common.CodeDataError, err
|
||||
}
|
||||
}
|
||||
|
||||
rerunWithDelete := run == string(entity.TaskStatusRunning) && req.Delete
|
||||
if rerunWithDelete {
|
||||
updates["progress_msg"] = ""
|
||||
updates["chunk_num"] = 0
|
||||
updates["token_num"] = 0
|
||||
}
|
||||
for _, vd := range validated {
|
||||
doc := vd.doc
|
||||
kb := vd.kb
|
||||
|
||||
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)
|
||||
// Start parsing: delegates to the shared start-parse flow. The
|
||||
// document run status is set by IngestionTaskService.StartRunning
|
||||
// when the task transitions from CREATED, not here.
|
||||
if run == string(entity.TaskStatusRunning) {
|
||||
if err := s.StartParseDocuments(doc, kb, userID, StartParseOptions{
|
||||
ApplyKB: req.ApplyKB,
|
||||
RerunWithDelete: req.Delete,
|
||||
}); err != nil {
|
||||
common.Error(fmt.Sprintf("go side, doc %s, start parse", doc.ID), err)
|
||||
return common.CodeExceptionError, err
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if err := s.documentDAO.UpdateByID(doc.ID, updates); err != nil {
|
||||
common.Error(fmt.Sprintf("go side, doc %s, UpdateByID failed", docID), err)
|
||||
// Cancel: RequestStop (STOPPING) and update doc state. Do NOT
|
||||
// delete the ingestion task or chunks here — deletion races with
|
||||
// the worker's async markStopped/settleToTerminal flow. Once the
|
||||
// worker detects STOPPING and transitions to STOPPED, the task
|
||||
// is terminal and can be safely cleaned up.
|
||||
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", doc.ID), err)
|
||||
return common.CodeDataError, err
|
||||
}
|
||||
if err := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{
|
||||
"run": string(entity.TaskStatusCancel),
|
||||
"progress": 0,
|
||||
}); err != nil {
|
||||
common.Error(fmt.Sprintf("go side, doc %s, UpdateByID failed", doc.ID), err)
|
||||
return common.CodeExceptionError, err
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Delete-only: user asked to remove prior parse results without
|
||||
// starting a new parse. RUNNING already continued above.
|
||||
if err := s.documentDAO.UpdateByID(doc.ID, map[string]interface{}{
|
||||
"run": run,
|
||||
"progress": 0,
|
||||
}); err != nil {
|
||||
common.Error(fmt.Sprintf("go side, doc %s, UpdateByID failed", doc.ID), err)
|
||||
return common.CodeExceptionError, err
|
||||
}
|
||||
|
||||
if req.Delete && !rerunWithDelete {
|
||||
if req.Delete {
|
||||
_, _ = 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)
|
||||
common.Error(fmt.Sprintf("go side, doc %s, ChunkStoreExists failed", doc.ID), 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)
|
||||
common.Error(fmt.Sprintf("go side, doc %s, DeleteChunks failed", doc.ID), err)
|
||||
return common.CodeExceptionError, err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if run == string(entity.TaskStatusRunning) {
|
||||
if req.ApplyKB {
|
||||
if doc.ParserConfig == nil {
|
||||
doc.ParserConfig = entity.JSONMap{}
|
||||
}
|
||||
config := map[string]interface{}{
|
||||
"llm_id": kb.ParserConfig["llm_id"],
|
||||
"enable_metadata": false,
|
||||
"metadata": map[string]interface{}{},
|
||||
}
|
||||
if value, ok := kb.ParserConfig["enable_metadata"]; ok {
|
||||
config["enable_metadata"] = value
|
||||
}
|
||||
if value, ok := kb.ParserConfig["metadata"]; ok {
|
||||
config["metadata"] = value
|
||||
}
|
||||
if err := s.updateDocumentParserConfig(doc.ID, config); err != nil {
|
||||
return common.CodeExceptionError, err
|
||||
}
|
||||
for key, value := range config {
|
||||
doc.ParserConfig[key] = value
|
||||
}
|
||||
}
|
||||
if doc.PipelineID != nil && strings.TrimSpace(*doc.PipelineID) != "" {
|
||||
if err := s.queueDocumentDataflowTask(kb, doc, userID); err != nil {
|
||||
return common.CodeExceptionError, err
|
||||
}
|
||||
continue
|
||||
}
|
||||
if doc.ParserID == string(entity.ParserTypeTable) {
|
||||
doneCount, ok := tableDoneCountByKB[doc.KbID]
|
||||
if !ok {
|
||||
count, err := s.countDoneDocuments(doc.KbID)
|
||||
if err != nil {
|
||||
return common.CodeExceptionError, err
|
||||
}
|
||||
doneCount = count
|
||||
tableDoneCountByKB[doc.KbID] = doneCount
|
||||
if doneCount <= 0 {
|
||||
if err := s.kbDAO.DeleteFieldMap(doc.KbID); err != nil && !dao.IsNotFoundErr(err) {
|
||||
return common.CodeExceptionError, err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
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
|
||||
}
|
||||
_, _, 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.IngestDocuments(doc.KbID, userID, []string{doc.ID}); err != nil {
|
||||
common.Error(fmt.Sprintf("go side, doc %s, IngestDocuments", docID), err)
|
||||
return common.CodeExceptionError, err
|
||||
}
|
||||
|
||||
if err := s.beginDocumentParse(doc.ID); err != nil {
|
||||
common.Error(fmt.Sprintf("go side, doc %s, beginDocumentParse", docID), err)
|
||||
return common.CodeExceptionError, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return common.CodeSuccess, nil
|
||||
}
|
||||
|
||||
func (s *DocumentService) prepareDocumentRerunWithDelete(doc *entity.Document, tenantID string) error {
|
||||
// AssertIngestionTasksTerminal verifies none of the documents has an
|
||||
// in-flight (RUNNING/STOPPING) ingestion task. Used as a batch pre-check
|
||||
// before re-parsing so a single non-terminal doc rejects the whole request
|
||||
// up front instead of partially cleaning some docs then failing.
|
||||
func (s *DocumentService) AssertIngestionTasksTerminal(docIDs []string) error {
|
||||
for _, docID := range docIDs {
|
||||
task, err := s.ingestionTaskDAO.GetByDocumentID(docID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("check ingestion task for %s: %w", docID, err)
|
||||
}
|
||||
if task == nil {
|
||||
continue
|
||||
}
|
||||
if task.Status == common.RUNNING || task.Status == common.STOPPING {
|
||||
return fmt.Errorf("document %s ingestion task is %s; stop it and wait for a terminal state before re-parsing", docID, task.Status)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *DocumentService) clearDocumentParseResults(doc *entity.Document, tenantID string) error {
|
||||
if doc == nil {
|
||||
return fmt.Errorf("document is nil")
|
||||
}
|
||||
|
||||
s.cancelExistingParseTasksBestEffort(doc.ID)
|
||||
// Refuse to clear a non-terminal ingestion task. An in-flight worker
|
||||
// (RUNNING) or one mid-stop (STOPPING) would keep writing chunks and
|
||||
// corrupt the new run's results. The caller must stop the task first
|
||||
// and wait for a terminal state (COMPLETED/STOPPED/FAILED) or CREATED.
|
||||
if task, _ := s.ingestionTaskDAO.GetByDocumentID(doc.ID); task != nil {
|
||||
if task.Status == common.RUNNING || task.Status == common.STOPPING {
|
||||
return fmt.Errorf("document %s ingestion task is %s; stop it and wait for a terminal state before re-parsing", doc.ID, task.Status)
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := s.taskDAO.DeleteByDocIDs([]string{doc.ID}); err != nil {
|
||||
// Delete terminal and CREATED ingestion tasks atomically, leaving
|
||||
// RUNNING/STOPPING tasks untouched so the check-then-delete window
|
||||
// between GetByDocumentID and the delete above cannot delete a task
|
||||
// that just transitioned to RUNNING.
|
||||
if _, err := s.ingestionTaskDAO.DeleteIfTerminal(doc.ID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -1389,24 +1433,6 @@ func (s *DocumentService) prepareDocumentRerunWithDelete(doc *entity.Document, t
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *DocumentService) cancelExistingParseTasksBestEffort(docID string) {
|
||||
tasks, err := s.taskDAO.GetByDocID(docID)
|
||||
if err != nil {
|
||||
common.Logger.Warn(fmt.Sprintf("cancelExistingParseTasksBestEffort: failed to get tasks for %s: %v", docID, err))
|
||||
return
|
||||
}
|
||||
redisClient := redis.Get()
|
||||
if redisClient == nil {
|
||||
return
|
||||
}
|
||||
for _, task := range tasks {
|
||||
if task == nil {
|
||||
continue
|
||||
}
|
||||
redisClient.Set(fmt.Sprintf("%s-cancel", task.ID), "x", 24*time.Hour)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *DocumentService) clearDocumentAndKBCountersForRerun(docID, kbID string) error {
|
||||
return dao.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var current entity.Document
|
||||
@@ -1458,368 +1484,6 @@ func (s *DocumentService) countDoneDocuments(datasetID string) (int64, error) {
|
||||
return count, err
|
||||
}
|
||||
|
||||
func (s *DocumentService) queueDocumentDataflowTask(kb *entity.Knowledgebase, doc *entity.Document, userID string) error {
|
||||
if err := s.beginDocumentParse(doc.ID); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := s.ingestionTaskSvc.CreateAndEnqueue(&entity.IngestionTask{
|
||||
DocumentID: doc.ID,
|
||||
UserID: userID,
|
||||
DatasetID: kb.ID,
|
||||
Schema: nil,
|
||||
Status: common.CREATED,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *DocumentService) newDocumentParseTasks(doc *entity.Document, bucket, objectName string, priority int64) ([]*entity.Task, error) {
|
||||
ranges, err := documentParseTaskRanges(doc, bucket, objectName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tasks := make([]*entity.Task, 0, len(ranges))
|
||||
for _, pageRange := range ranges {
|
||||
tasks = append(tasks, s.newDocumentParseTask(doc, pageRange.from, pageRange.to, priority))
|
||||
}
|
||||
return tasks, nil
|
||||
}
|
||||
|
||||
func (s *DocumentService) newDocumentParseTask(doc *entity.Document, fromPage, toPage, priority int64) *entity.Task {
|
||||
now := time.Now()
|
||||
progressMsg := ""
|
||||
digest := documentParseTaskDigest(doc, fromPage, toPage)
|
||||
chunkIDs := ""
|
||||
return &entity.Task{
|
||||
ID: utility.GenerateUUID(),
|
||||
DocID: doc.ID,
|
||||
FromPage: fromPage,
|
||||
ToPage: toPage,
|
||||
TaskType: "",
|
||||
Priority: priority,
|
||||
BeginAt: &now,
|
||||
Progress: 0,
|
||||
ProgressMsg: &progressMsg,
|
||||
Digest: &digest,
|
||||
ChunkIDs: &chunkIDs,
|
||||
}
|
||||
}
|
||||
|
||||
func documentParseTaskRanges(doc *entity.Document, bucket, objectName string) ([]documentParsePageRange, error) {
|
||||
if doc.Type == "pdf" {
|
||||
binary, err := documentStorageBinary(bucket, objectName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pages := documentEstimatePDFPageCount(binary)
|
||||
pageSize := int64(documentParserConfigInt(doc.ParserConfig, "task_page_size", 12))
|
||||
if doc.ParserID == string(entity.ParserTypePaper) {
|
||||
pageSize = int64(documentParserConfigInt(doc.ParserConfig, "task_page_size", 22))
|
||||
}
|
||||
if doc.ParserID == string(entity.ParserTypeOne) ||
|
||||
documentParserConfigBool(doc.ParserConfig, "toc_extraction", false) {
|
||||
pageSize = maximumTaskPageNumber
|
||||
}
|
||||
if pageSize <= 0 {
|
||||
pageSize = 12
|
||||
}
|
||||
ranges := make([]documentParsePageRange, 0)
|
||||
for _, configuredRange := range documentParserConfigPageRanges(doc.ParserConfig) {
|
||||
start := configuredRange.from - 1
|
||||
if start < 0 {
|
||||
start = 0
|
||||
}
|
||||
end := configuredRange.to - 1
|
||||
if pages >= 0 && end > pages {
|
||||
end = pages
|
||||
}
|
||||
for page := start; page < end; page += pageSize {
|
||||
to := page + pageSize
|
||||
if to > end {
|
||||
to = end
|
||||
}
|
||||
ranges = append(ranges, documentParsePageRange{from: page, to: to})
|
||||
}
|
||||
}
|
||||
if len(ranges) == 0 {
|
||||
// pages == 0 means page count detection failed (e.g. compressed
|
||||
// PDF where both regex and pdfoxide fallbacks failed). Fall back
|
||||
// to maximumTaskPageNumber so the Python parser processes all
|
||||
// pages via slicing (Python gracefully caps at actual page count).
|
||||
ranges = append(ranges, documentParsePageRange{from: 0, to: maximumTaskPageNumber})
|
||||
}
|
||||
return ranges, nil
|
||||
}
|
||||
if doc.ParserID == string(entity.ParserTypeTable) {
|
||||
binary, err := documentStorageBinary(bucket, objectName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows := documentEstimateTableRowCount(documentName(doc), binary)
|
||||
if rows <= 0 {
|
||||
return []documentParsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil
|
||||
}
|
||||
ranges := make([]documentParsePageRange, 0, (rows+2999)/3000)
|
||||
for row := int64(0); row < int64(rows); row += 3000 {
|
||||
to := row + 3000
|
||||
if to > int64(rows) {
|
||||
to = int64(rows)
|
||||
}
|
||||
ranges = append(ranges, documentParsePageRange{from: row, to: to})
|
||||
}
|
||||
return ranges, nil
|
||||
}
|
||||
return []documentParsePageRange{{from: 0, to: maximumTaskPageNumber}}, nil
|
||||
}
|
||||
|
||||
func documentStorageBinary(bucket, objectName string) ([]byte, error) {
|
||||
storageImpl := storage.GetStorageFactory().GetStorage()
|
||||
if storageImpl == nil {
|
||||
return nil, fmt.Errorf("storage not initialized")
|
||||
}
|
||||
return storageImpl.Get(bucket, objectName)
|
||||
}
|
||||
|
||||
func documentName(doc *entity.Document) string {
|
||||
if doc == nil || doc.Name == nil {
|
||||
return ""
|
||||
}
|
||||
return *doc.Name
|
||||
}
|
||||
|
||||
func documentParserConfigInt(config map[string]interface{}, key string, fallback int) int {
|
||||
value, ok := config[key]
|
||||
if !ok || value == nil {
|
||||
return fallback
|
||||
}
|
||||
switch typedValue := value.(type) {
|
||||
case int:
|
||||
return typedValue
|
||||
case int64:
|
||||
return int(typedValue)
|
||||
case float64:
|
||||
return int(typedValue)
|
||||
case json.Number:
|
||||
if intValue, err := typedValue.Int64(); err == nil {
|
||||
return int(intValue)
|
||||
}
|
||||
case string:
|
||||
if intValue, err := strconv.Atoi(strings.TrimSpace(typedValue)); err == nil {
|
||||
return intValue
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func documentParserConfigBool(config map[string]interface{}, key string, fallback bool) bool {
|
||||
value, ok := config[key]
|
||||
if !ok || value == nil {
|
||||
return fallback
|
||||
}
|
||||
switch typedValue := value.(type) {
|
||||
case bool:
|
||||
return typedValue
|
||||
case string:
|
||||
switch strings.ToLower(strings.TrimSpace(typedValue)) {
|
||||
case "true", "1", "yes", "on":
|
||||
return true
|
||||
case "false", "0", "no", "off":
|
||||
return false
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func documentParserConfigPageRanges(config map[string]interface{}) []documentParsePageRange {
|
||||
defaultRanges := []documentParsePageRange{{from: 1, to: 100000}}
|
||||
raw, ok := config["pages"]
|
||||
if !ok || raw == nil {
|
||||
return defaultRanges
|
||||
}
|
||||
rawRanges, ok := raw.([]interface{})
|
||||
if !ok || len(rawRanges) == 0 {
|
||||
return defaultRanges
|
||||
}
|
||||
ranges := make([]documentParsePageRange, 0, len(rawRanges))
|
||||
for _, rawRange := range rawRanges {
|
||||
rangeValues, ok := rawRange.([]interface{})
|
||||
if !ok || len(rangeValues) < 2 {
|
||||
continue
|
||||
}
|
||||
from, okFrom := documentToInt64(rangeValues[0])
|
||||
to, okTo := documentToInt64(rangeValues[1])
|
||||
if okFrom && okTo && to > from {
|
||||
ranges = append(ranges, documentParsePageRange{from: from, to: to})
|
||||
}
|
||||
}
|
||||
if len(ranges) == 0 {
|
||||
return defaultRanges
|
||||
}
|
||||
return ranges
|
||||
}
|
||||
|
||||
func documentToInt64(value interface{}) (int64, bool) {
|
||||
switch typedValue := value.(type) {
|
||||
case int:
|
||||
return int64(typedValue), true
|
||||
case int64:
|
||||
return typedValue, true
|
||||
case float64:
|
||||
return int64(typedValue), true
|
||||
case json.Number:
|
||||
intValue, err := typedValue.Int64()
|
||||
return intValue, err == nil
|
||||
case string:
|
||||
intValue, err := strconv.ParseInt(strings.TrimSpace(typedValue), 10, 64)
|
||||
return intValue, err == nil
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
var documentPDFPagePattern = regexp.MustCompile(`/Type\s*/Page\b`)
|
||||
|
||||
func documentEstimatePDFPageCount(binary []byte) int64 {
|
||||
if len(binary) == 0 {
|
||||
return 0
|
||||
}
|
||||
// Fast path: regex works for uncompressed PDFs.
|
||||
count := int64(len(documentPDFPagePattern.FindAll(binary, -1)))
|
||||
if count > 0 {
|
||||
return count
|
||||
}
|
||||
// Fallback for compressed PDFs where /Type /Page is inside a
|
||||
// compressed object stream: use pdf_oxide to get the real page count.
|
||||
if doc, err := pdfoxide.OpenBytes(binary); err == nil {
|
||||
defer doc.Close()
|
||||
if pages, err := doc.PageCount(); err == nil {
|
||||
return int64(pages)
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func documentEstimateTableRowCount(name string, binary []byte) int {
|
||||
switch strings.ToLower(filepath.Ext(name)) {
|
||||
case ".xlsx":
|
||||
if rows, err := documentCountXLSXRows(binary); err == nil {
|
||||
return rows
|
||||
}
|
||||
case ".csv", ".tsv", ".txt":
|
||||
return documentCountDelimitedRows(name, binary)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func documentCountDelimitedRows(name string, binary []byte) int {
|
||||
reader := csv.NewReader(bytes.NewReader(binary))
|
||||
reader.FieldsPerRecord = -1
|
||||
reader.ReuseRecord = true
|
||||
if strings.EqualFold(filepath.Ext(name), ".tsv") {
|
||||
reader.Comma = '\t'
|
||||
}
|
||||
rows := 0
|
||||
for {
|
||||
_, err := reader.Read()
|
||||
if err == nil {
|
||||
rows++
|
||||
continue
|
||||
}
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
rows += bytes.Count(binary, []byte{'\n'})
|
||||
if len(binary) > 0 && binary[len(binary)-1] != '\n' {
|
||||
rows++
|
||||
}
|
||||
break
|
||||
}
|
||||
return rows
|
||||
}
|
||||
|
||||
func documentCountXLSXRows(binary []byte) (int, error) {
|
||||
zipReader, err := zip.NewReader(bytes.NewReader(binary), int64(len(binary)))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
maxRows := 0
|
||||
for _, file := range zipReader.File {
|
||||
if !strings.HasPrefix(file.Name, "xl/worksheets/") || !strings.HasSuffix(file.Name, ".xml") {
|
||||
continue
|
||||
}
|
||||
rows, err := documentCountWorksheetRows(file)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if rows > maxRows {
|
||||
maxRows = rows
|
||||
}
|
||||
}
|
||||
return maxRows, nil
|
||||
}
|
||||
|
||||
func documentCountWorksheetRows(file *zip.File) (int, error) {
|
||||
reader, err := file.Open()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer reader.Close()
|
||||
decoder := xml.NewDecoder(reader)
|
||||
rows := 0
|
||||
for {
|
||||
token, err := decoder.Token()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
start, ok := token.(xml.StartElement)
|
||||
if ok && start.Name.Local == "row" {
|
||||
rows++
|
||||
}
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
func (s *DocumentService) beginDocumentParse(docID string) error {
|
||||
now := time.Now()
|
||||
return dao.GetDB().Model(&entity.Document{}).Where("id = ?", docID).Updates(map[string]interface{}{
|
||||
"progress_msg": "Task is queued...",
|
||||
"process_begin_at": now,
|
||||
"progress": rand.Float64() * 0.01,
|
||||
"run": string(entity.TaskStatusRunning),
|
||||
"chunk_num": 0,
|
||||
"token_num": 0,
|
||||
}).Error
|
||||
}
|
||||
|
||||
func documentParseTaskDigest(doc *entity.Document, fromPage, toPage int64) string {
|
||||
hasher := xxhash.New()
|
||||
config := map[string]interface{}{
|
||||
"doc_id": doc.ID,
|
||||
"kb_id": doc.KbID,
|
||||
"parser_id": doc.ParserID,
|
||||
"parser_config": doc.ParserConfig,
|
||||
}
|
||||
keys := make([]string, 0, len(config))
|
||||
for key := range config {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
for _, key := range keys {
|
||||
b, err := json.Marshal(config[key])
|
||||
if err != nil {
|
||||
hasher.WriteString(fmt.Sprint(config[key]))
|
||||
} else {
|
||||
hasher.Write(b)
|
||||
}
|
||||
}
|
||||
hasher.WriteString(doc.ID)
|
||||
hasher.WriteString(strconv.FormatInt(fromPage, 10))
|
||||
hasher.WriteString(strconv.FormatInt(toPage, 10))
|
||||
return fmt.Sprintf("%x", hasher.Sum64())
|
||||
}
|
||||
|
||||
func (s *DocumentService) clearKBChunkNumWhenRerun(doc *entity.Document) error {
|
||||
if doc == nil {
|
||||
return fmt.Errorf("document is nil")
|
||||
@@ -1915,7 +1579,7 @@ func (s *DocumentService) StopParseDocuments(datasetID string, docIDs []string)
|
||||
var errors []string
|
||||
successCount := 0
|
||||
for _, doc := range docs {
|
||||
if cancelErr := s.cancelDocParse(doc); cancelErr != nil {
|
||||
if cancelErr := s.CancelDocParse(doc); cancelErr != nil {
|
||||
errors = append(errors, cancelErr.Error())
|
||||
continue
|
||||
}
|
||||
@@ -1952,10 +1616,9 @@ func (s *DocumentService) validateDocsInDataset(docIDs []string, datasetID strin
|
||||
return docs, nil
|
||||
}
|
||||
|
||||
// cancelDocParse stops the ingestion task for the document by calling
|
||||
// RequestStop (which sets Redis {taskID}-cancel + DB STOPPING), then marks
|
||||
// the document run status as CANCEL.
|
||||
func (s *DocumentService) cancelDocParse(doc *entity.Document) error {
|
||||
// CancelDocParse stops the ingestion task for the document by calling
|
||||
// RequestStop (STOPPING), then marks the document run status as CANCEL.
|
||||
func (s *DocumentService) CancelDocParse(doc *entity.Document) error {
|
||||
task, err := s.ingestionTaskDAO.GetByDocumentID(doc.ID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get ingestion task for %s: %v", doc.ID, err)
|
||||
|
||||
@@ -517,6 +517,11 @@ func insertTestTask(t *testing.T, id, docID string) {
|
||||
}
|
||||
|
||||
func insertTestIngestionTask(t *testing.T, id, userID, docID, datasetID string) {
|
||||
t.Helper()
|
||||
insertTestIngestionTaskWithStatus(t, id, userID, docID, datasetID, common.CREATED)
|
||||
}
|
||||
|
||||
func insertTestIngestionTaskWithStatus(t *testing.T, id, userID, docID, datasetID, status string) {
|
||||
t.Helper()
|
||||
task := &entity.IngestionTask{
|
||||
ID: id,
|
||||
@@ -524,7 +529,7 @@ func insertTestIngestionTask(t *testing.T, id, userID, docID, datasetID string)
|
||||
DocumentID: docID,
|
||||
DatasetID: datasetID,
|
||||
Schema: entity.JSONMap{},
|
||||
Status: common.CREATED,
|
||||
Status: status,
|
||||
}
|
||||
if err := dao.DB.Create(task).Error; err != nil {
|
||||
t.Fatalf("insert test ingestion task: %v", err)
|
||||
@@ -1077,12 +1082,16 @@ func insertTestTaskWithProgress(t *testing.T, id, docID string, progress float64
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueueDocumentDataflowTask_PublishesIngestionTaskMessage(t *testing.T) {
|
||||
func TestStartParseDocuments_EnqueuesIngestionTask(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)
|
||||
// GetDocumentStorageAddress requires a non-empty document location.
|
||||
if err := dao.DB.Model(&entity.Document{}).Where("id = ?", "doc-1").Update("location", "loc-1").Error; err != nil {
|
||||
t.Fatalf("set location: %v", err)
|
||||
}
|
||||
|
||||
publisher := &recordingTaskPublisher{}
|
||||
svc := testDocumentService(t)
|
||||
@@ -1097,8 +1106,8 @@ func TestQueueDocumentDataflowTask_PublishesIngestionTaskMessage(t *testing.T) {
|
||||
t.Fatalf("load doc: %v", err)
|
||||
}
|
||||
|
||||
if err := svc.queueDocumentDataflowTask(kb, doc, "user-1"); err != nil {
|
||||
t.Fatalf("queueDocumentDataflowTask: %v", err)
|
||||
if err := svc.StartParseDocuments(doc, kb, "user-1", StartParseOptions{}); err != nil {
|
||||
t.Fatalf("StartParseDocuments: %v", err)
|
||||
}
|
||||
|
||||
if publisher.subject != "tasks.RAGFLOW" {
|
||||
@@ -1805,12 +1814,12 @@ func TestResetDocumentForReparseSkipsSecondCounterDecrement(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareDocumentRerunWithDeleteClearsCountersTasksAndChunks(t *testing.T) {
|
||||
func TestClearDocumentParseResultsClearsCountersTasksAndChunks(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5)
|
||||
insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5)
|
||||
insertTestTask(t, "task-1", "doc-1")
|
||||
insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED)
|
||||
|
||||
doc, err := dao.NewDocumentDAO().GetByID("doc-1")
|
||||
if err != nil {
|
||||
@@ -1821,8 +1830,8 @@ func TestPrepareDocumentRerunWithDeleteClearsCountersTasksAndChunks(t *testing.T
|
||||
svc := testDocumentService(t)
|
||||
svc.docEngine = engine
|
||||
|
||||
if err := svc.prepareDocumentRerunWithDelete(doc, "tenant-1"); err != nil {
|
||||
t.Fatalf("prepareDocumentRerunWithDelete failed: %v", err)
|
||||
if err := svc.clearDocumentParseResults(doc, "tenant-1"); err != nil {
|
||||
t.Fatalf("clearDocumentParseResults failed: %v", err)
|
||||
}
|
||||
|
||||
updatedDoc, _ := dao.NewDocumentDAO().GetByID("doc-1")
|
||||
@@ -1833,12 +1842,10 @@ func TestPrepareDocumentRerunWithDeleteClearsCountersTasksAndChunks(t *testing.T
|
||||
if kb.TokenNum != 0 || kb.ChunkNum != 0 {
|
||||
t.Fatalf("kb counters = token:%d chunk:%d, want zero", kb.TokenNum, kb.ChunkNum)
|
||||
}
|
||||
var taskCount int64
|
||||
if err := dao.DB.Model(&entity.Task{}).Where("doc_id = ?", "doc-1").Count(&taskCount).Error; err != nil {
|
||||
t.Fatalf("count tasks: %v", err)
|
||||
}
|
||||
if taskCount != 0 {
|
||||
t.Fatalf("task count = %d, want zero", taskCount)
|
||||
// The completed ingestion task must be deleted so the new run can proceed.
|
||||
remainingTask, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1")
|
||||
if remainingTask != nil {
|
||||
t.Fatalf("ingestion task should be deleted, status was %q", remainingTask.Status)
|
||||
}
|
||||
if engine.deleteCalls != 1 {
|
||||
t.Fatalf("deleteCalls = %d, want 1", engine.deleteCalls)
|
||||
@@ -1848,7 +1855,7 @@ func TestPrepareDocumentRerunWithDeleteClearsCountersTasksAndChunks(t *testing.T
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareDocumentRerunWithDeleteIsIdempotentForStaleDocSnapshot(t *testing.T) {
|
||||
func TestClearDocumentParseResultsIsIdempotentForStaleDocSnapshot(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5)
|
||||
@@ -1860,11 +1867,11 @@ func TestPrepareDocumentRerunWithDeleteIsIdempotentForStaleDocSnapshot(t *testin
|
||||
}
|
||||
|
||||
svc := testDocumentService(t)
|
||||
if err := svc.prepareDocumentRerunWithDelete(staleDoc, "tenant-1"); err != nil {
|
||||
t.Fatalf("first prepareDocumentRerunWithDelete failed: %v", err)
|
||||
if err := svc.clearDocumentParseResults(staleDoc, "tenant-1"); err != nil {
|
||||
t.Fatalf("first clearDocumentParseResults failed: %v", err)
|
||||
}
|
||||
if err := svc.prepareDocumentRerunWithDelete(staleDoc, "tenant-1"); err != nil {
|
||||
t.Fatalf("second prepareDocumentRerunWithDelete failed: %v", err)
|
||||
if err := svc.clearDocumentParseResults(staleDoc, "tenant-1"); err != nil {
|
||||
t.Fatalf("second clearDocumentParseResults failed: %v", err)
|
||||
}
|
||||
|
||||
doc, _ := dao.NewDocumentDAO().GetByID("doc-1")
|
||||
@@ -1877,6 +1884,126 @@ func TestPrepareDocumentRerunWithDeleteIsIdempotentForStaleDocSnapshot(t *testin
|
||||
}
|
||||
}
|
||||
|
||||
// TestClearDocumentParseResults_RejectsNonTerminalIngestionTask verifies
|
||||
// that re-parsing is refused while a document's ingestion task is still
|
||||
// RUNNING or STOPPING. Deleting a non-terminal task would let the in-flight
|
||||
// worker keep writing chunks and corrupt the new run's results.
|
||||
func TestClearDocumentParseResults_RejectsNonTerminalIngestionTask(t *testing.T) {
|
||||
for _, status := range []string{common.RUNNING, common.STOPPING} {
|
||||
t.Run(status, func(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5)
|
||||
insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusRunning), 10, 5)
|
||||
insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", status)
|
||||
|
||||
doc, err := dao.NewDocumentDAO().GetByID("doc-1")
|
||||
if err != nil {
|
||||
t.Fatalf("get doc: %v", err)
|
||||
}
|
||||
svc := testDocumentService(t)
|
||||
|
||||
if err := svc.clearDocumentParseResults(doc, "tenant-1"); err == nil {
|
||||
t.Fatalf("expected error for %s ingestion task, got nil", status)
|
||||
}
|
||||
// The non-terminal task must NOT be deleted.
|
||||
task, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1")
|
||||
if task == nil {
|
||||
t.Fatalf("%s ingestion task must not be deleted", status)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestClearDocumentParseResults_DeletesTerminalIngestionTask verifies
|
||||
// that a terminal ingestion task (STOPPED/COMPLETED/FAILED) and a CREATED
|
||||
// (queued) task are deleted so the new run can proceed.
|
||||
func TestClearDocumentParseResults_DeletesTerminalIngestionTask(t *testing.T) {
|
||||
for _, status := range []string{common.CREATED, common.COMPLETED, common.STOPPED, common.FAILED} {
|
||||
t.Run(status, func(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertTestKB(t, "kb-1", "tenant-1", 1, 10, 5)
|
||||
insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5)
|
||||
insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", status)
|
||||
|
||||
doc, err := dao.NewDocumentDAO().GetByID("doc-1")
|
||||
if err != nil {
|
||||
t.Fatalf("get doc: %v", err)
|
||||
}
|
||||
svc := testDocumentService(t)
|
||||
|
||||
if err := svc.clearDocumentParseResults(doc, "tenant-1"); err != nil {
|
||||
t.Fatalf("clearDocumentParseResults for %s task: %v", status, err)
|
||||
}
|
||||
task, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1")
|
||||
if task != nil {
|
||||
t.Fatalf("%s ingestion task should be deleted, still present", status)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssertIngestionTasksTerminal_RejectsNonTerminal(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0)
|
||||
insertTestDoc(t, "doc-1", "kb-1", 0, 0)
|
||||
insertTestDoc(t, "doc-2", "kb-1", 0, 0)
|
||||
insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED)
|
||||
insertTestIngestionTaskWithStatus(t, "task-2", "user-1", "doc-2", "kb-1", common.RUNNING)
|
||||
|
||||
svc := testDocumentService(t)
|
||||
if err := svc.AssertIngestionTasksTerminal([]string{"doc-1", "doc-2"}); err == nil {
|
||||
t.Fatal("expected error for RUNNING task, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssertIngestionTasksTerminal_AcceptsAllTerminal(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0)
|
||||
insertTestDoc(t, "doc-1", "kb-1", 0, 0)
|
||||
insertTestDoc(t, "doc-2", "kb-1", 0, 0)
|
||||
insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED)
|
||||
insertTestIngestionTaskWithStatus(t, "task-2", "user-1", "doc-2", "kb-1", common.STOPPED)
|
||||
|
||||
svc := testDocumentService(t)
|
||||
if err := svc.AssertIngestionTasksTerminal([]string{"doc-1", "doc-2"}); err != nil {
|
||||
t.Fatalf("expected nil for all-terminal batch, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestIngest_RerunWithDelete_RejectsBatchWithRunningTask verifies the batch
|
||||
// pre-check: when re-parsing with delete, if any document's ingestion task is
|
||||
// non-terminal, the whole request is rejected up front so no document is
|
||||
// partially cleaned.
|
||||
func TestIngest_RerunWithDelete_RejectsBatchWithRunningTask(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertUserTenantForAccessCheck(t, "user-1", "tenant-1")
|
||||
insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0)
|
||||
insertTestDoc(t, "doc-1", "kb-1", 0, 0)
|
||||
insertTestDoc(t, "doc-2", "kb-1", 0, 0)
|
||||
insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED)
|
||||
insertTestIngestionTaskWithStatus(t, "task-2", "user-1", "doc-2", "kb-1", common.RUNNING)
|
||||
|
||||
svc := testDocumentService(t)
|
||||
_, err := svc.Ingest("user-1", &IngestDocumentRequest{
|
||||
DocIDs: []string{"doc-1", "doc-2"},
|
||||
Run: string(entity.TaskStatusRunning),
|
||||
Delete: true,
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected error for batch with RUNNING task, got nil")
|
||||
}
|
||||
// doc-1 (terminal) task must NOT be deleted - the whole batch was rejected.
|
||||
task1, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1")
|
||||
if task1 == nil {
|
||||
t.Fatal("doc-1 terminal task should not be deleted when batch is rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateDatasetDocumentPropagatesMetadataDeleteFailure(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
@@ -2506,3 +2633,69 @@ func TestGetThumbnails_AlignsWithPythonFormatting(t *testing.T) {
|
||||
t.Fatalf("did not expect other tenant doc in result: %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestStartParseDocuments_FailsBeforeClearing verifies that GetDocumentStorageAddress
|
||||
// runs before clearDocumentParseResults: when storage is invalid, we fail without
|
||||
// deleting existing parse results.
|
||||
func TestStartParseDocuments_FailsBeforeClearing(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertTestKB(t, "kb-1", "tenant-1", 0, 10, 5)
|
||||
insertTestDocWithRun(t, "doc-1", "kb-1", string(entity.TaskStatusDone), 10, 5)
|
||||
insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.COMPLETED)
|
||||
|
||||
doc, err := dao.NewDocumentDAO().GetByID("doc-1")
|
||||
if err != nil {
|
||||
t.Fatalf("get doc: %v", err)
|
||||
}
|
||||
kb, err := dao.NewKnowledgebaseDAO().GetByID("kb-1")
|
||||
if err != nil {
|
||||
t.Fatalf("get kb: %v", err)
|
||||
}
|
||||
|
||||
// Force GetDocumentStorageAddress to fail by clearing the document location.
|
||||
if err := dao.DB.Model(&entity.Document{}).Where("id = ?", "doc-1").Update("location", "").Error; err != nil {
|
||||
t.Fatalf("clear location: %v", err)
|
||||
}
|
||||
|
||||
svc := testDocumentService(t)
|
||||
err = svc.StartParseDocuments(doc, kb, "user-1", StartParseOptions{RerunWithDelete: true})
|
||||
if err == nil {
|
||||
t.Fatal("expected error from GetDocumentStorageAddress, got nil")
|
||||
}
|
||||
|
||||
// The old ingestion task must still exist — we failed before clearing.
|
||||
remaining, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1")
|
||||
if remaining == nil {
|
||||
t.Fatal("ingestion task should NOT be deleted when storage validation fails")
|
||||
}
|
||||
}
|
||||
|
||||
// TestIngest_CancelDoesNotDeleteIngestionTask verifies that cancel+delete
|
||||
// no longer races with the worker by deleting the ingestion task while the
|
||||
// worker is still stopping. Cancel only issues RequestStop (STOPPING);
|
||||
// deletion is deferred until the task reaches a terminal state.
|
||||
func TestIngest_CancelDoesNotDeleteIngestionTask(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertUserTenantForAccessCheck(t, "user-1", "tenant-1")
|
||||
insertTestKB(t, "kb-1", "tenant-1", 0, 0, 0)
|
||||
insertTestDoc(t, "doc-1", "kb-1", 10, 5)
|
||||
insertTestIngestionTaskWithStatus(t, "task-1", "user-1", "doc-1", "kb-1", common.RUNNING)
|
||||
|
||||
svc := testDocumentService(t)
|
||||
_, err := svc.Ingest("user-1", &IngestDocumentRequest{
|
||||
DocIDs: []string{"doc-1"},
|
||||
Run: string(entity.TaskStatusCancel),
|
||||
Delete: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Ingest(cancel+delete): %v", err)
|
||||
}
|
||||
|
||||
// The ingestion task must NOT be deleted — cancel only stops it.
|
||||
remaining, _ := svc.ingestionTaskDAO.GetByDocumentID("doc-1")
|
||||
if remaining == nil {
|
||||
t.Fatal("ingestion task must NOT be deleted by cancel")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -205,7 +205,27 @@ func (s *IngestionTaskService) StartRunning(taskID string) (*entity.IngestionTas
|
||||
}
|
||||
switch task.Status {
|
||||
case common.CREATED:
|
||||
return s.transition(taskID, common.RUNNING)
|
||||
task, err := s.transition(taskID, common.RUNNING)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// The task just started running: mirror it to the document so its
|
||||
// run status and progress counters reflect real processing, not
|
||||
// just API acceptance. Best-effort - a DB blip here must not fail
|
||||
// the task transition and trigger a redelivery loop. run uses the
|
||||
// document's numeric TaskStatus enum ("1"), not the task's string
|
||||
// status label.
|
||||
if err := s.documentDAO.UpdateByID(task.DocumentID, map[string]interface{}{
|
||||
"run": string(entity.TaskStatusRunning),
|
||||
"progress": float64(0),
|
||||
"chunk_num": int64(0),
|
||||
"token_num": int64(0),
|
||||
"process_begin_at": time.Now(),
|
||||
"progress_msg": "",
|
||||
}); err != nil {
|
||||
common.Warn(fmt.Sprintf("StartRunning: mark document %s running for task %s: %v", task.DocumentID, taskID, err))
|
||||
}
|
||||
return task, nil
|
||||
case common.STOPPING:
|
||||
return s.transition(taskID, common.STOPPED)
|
||||
case common.RUNNING, common.COMPLETED, common.STOPPED, common.FAILED:
|
||||
|
||||
@@ -213,6 +213,102 @@ func TestIngestionTaskServiceStartRunningTransitionsCreatedTask(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestStartRunningMarksDocumentRunning locks in that starting a CREATED task
|
||||
// mirrors the transition to its document: run=RUNNING and progress counters
|
||||
// reset, with a fresh process_begin_at. The document bookkeeping is owned by
|
||||
// the task-lifecycle transition, not the ingestion worker's execution path.
|
||||
func TestStartRunningMarksDocumentRunning(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertTestKB(t, "kb-1", "tenant-1", 1, 0, 0)
|
||||
insertTestDoc(t, "doc-1", "kb-1", 100, 10)
|
||||
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
||||
|
||||
// Seed the document as a partially-processed, non-RUNNING state that the
|
||||
// start transition must clobber.
|
||||
if err := db.Model(&entity.Document{}).Where("id = ?", "doc-1").
|
||||
Updates(map[string]interface{}{
|
||||
"run": string(entity.TaskStatusDone),
|
||||
"progress": float64(0.5),
|
||||
"progress_msg": "partial",
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("seed document: %v", err)
|
||||
}
|
||||
|
||||
svc := NewIngestionTaskService()
|
||||
if _, err := svc.StartRunning("task-1"); err != nil {
|
||||
t.Fatalf("StartRunning failed: %v", err)
|
||||
}
|
||||
|
||||
var doc entity.Document
|
||||
if err := db.Where("id = ?", "doc-1").First(&doc).Error; err != nil {
|
||||
t.Fatalf("reload document: %v", err)
|
||||
}
|
||||
if doc.Run == nil || *doc.Run != string(entity.TaskStatusRunning) {
|
||||
t.Fatalf("run = %v, want RUNNING(%q)", doc.Run, string(entity.TaskStatusRunning))
|
||||
}
|
||||
if doc.Progress != 0 {
|
||||
t.Fatalf("progress = %f, want 0", doc.Progress)
|
||||
}
|
||||
if doc.ChunkNum != 0 {
|
||||
t.Fatalf("chunk_num = %d, want 0", doc.ChunkNum)
|
||||
}
|
||||
if doc.TokenNum != 0 {
|
||||
t.Fatalf("token_num = %d, want 0", doc.TokenNum)
|
||||
}
|
||||
if doc.ProgressMsg != nil && *doc.ProgressMsg != "" {
|
||||
t.Fatalf("progress_msg = %q, want empty", *doc.ProgressMsg)
|
||||
}
|
||||
if doc.ProcessBeginAt == nil || doc.ProcessBeginAt.IsZero() {
|
||||
t.Fatal("process_begin_at not set")
|
||||
}
|
||||
}
|
||||
|
||||
// TestStartRunningLeavesTerminalDocumentUntouched locks in the no-resurrection
|
||||
// invariant: a task already in a terminal status is returned as-is by
|
||||
// StartRunning, and its document's finished run status/counters are not reset.
|
||||
func TestStartRunningLeavesTerminalDocumentUntouched(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
insertTestKB(t, "kb-1", "tenant-1", 1, 0, 0)
|
||||
insertTestDoc(t, "doc-1", "kb-1", 100, 10)
|
||||
insertTestIngestionTask(t, "task-1", "user-1", "doc-1", "kb-1")
|
||||
|
||||
finishedRun := string(entity.TaskStatusDone)
|
||||
if err := db.Model(&entity.Document{}).Where("id = ?", "doc-1").
|
||||
Updates(map[string]interface{}{
|
||||
"run": finishedRun,
|
||||
"progress": float64(1.0),
|
||||
"progress_msg": "done",
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("seed document: %v", err)
|
||||
}
|
||||
if err := db.Model(&entity.IngestionTask{}).Where("id = ?", "task-1").
|
||||
Update("status", common.COMPLETED).Error; err != nil {
|
||||
t.Fatalf("set COMPLETED: %v", err)
|
||||
}
|
||||
|
||||
svc := NewIngestionTaskService()
|
||||
task, err := svc.StartRunning("task-1")
|
||||
if err != nil {
|
||||
t.Fatalf("StartRunning failed: %v", err)
|
||||
}
|
||||
if task.Status != common.COMPLETED {
|
||||
t.Fatalf("status = %q, want %q (terminal must be preserved)", task.Status, common.COMPLETED)
|
||||
}
|
||||
|
||||
var doc entity.Document
|
||||
if err := db.Where("id = ?", "doc-1").First(&doc).Error; err != nil {
|
||||
t.Fatalf("reload document: %v", err)
|
||||
}
|
||||
if doc.Run == nil || *doc.Run != finishedRun {
|
||||
t.Fatalf("run = %v, want %q (terminal document must not be resurrected)", doc.Run, finishedRun)
|
||||
}
|
||||
if doc.ChunkNum != 10 || doc.TokenNum != 100 {
|
||||
t.Fatalf("counters changed: chunk_num=%d token_num=%d, want 10/100", doc.ChunkNum, doc.TokenNum)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIngestionTaskServiceRequestStopTransitionsCreatedTaskToStopped(t *testing.T) {
|
||||
db := setupServiceTestDB(t)
|
||||
pushServiceDB(t, db)
|
||||
|
||||
Reference in New Issue
Block a user