From 35de240c396fcdd76ef4acd565d0d34d9ef7cec1 Mon Sep 17 00:00:00 2001 From: jay77721 <164177721+jay77721@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:47:54 +0800 Subject: [PATCH] feat(ingestion): add progress fraction reporting to VisionEnhancement and cropImageChunks (#20800) --- .../component/chunker/pdfcrop_cgo.go | 69 +++++-- .../chunker/pdfcrop_progress_cgo_test.go | 133 +++++++++++++ internal/ingestion/component/extractor.go | 44 ----- .../ingestion/component/media_dispatch.go | 2 +- internal/ingestion/component/parser.go | 13 +- .../ingestion/component/parser_dispatch.go | 32 +++- .../parser_dispatch_progress_test.go | 71 ++++++- .../ingestion/component/progress_fraction.go | 51 +++++ .../ingestion/component/vision_enhancement.go | 17 ++ .../vision_enhancement_progress_test.go | 180 ++++++++++++++++++ .../component/vision_enhancement_test.go | 47 ++--- 11 files changed, 559 insertions(+), 100 deletions(-) create mode 100644 internal/ingestion/component/chunker/pdfcrop_progress_cgo_test.go create mode 100644 internal/ingestion/component/progress_fraction.go create mode 100644 internal/ingestion/component/vision_enhancement_progress_test.go diff --git a/internal/ingestion/component/chunker/pdfcrop_cgo.go b/internal/ingestion/component/chunker/pdfcrop_cgo.go index 0bd5105cfd..704ecc6b4c 100644 --- a/internal/ingestion/component/chunker/pdfcrop_cgo.go +++ b/internal/ingestion/component/chunker/pdfcrop_cgo.go @@ -19,12 +19,13 @@ import ( "runtime" "strconv" "sync" + "sync/atomic" "go.uber.org/zap" + agentruntime "ragflow/internal/agent/runtime" "ragflow/internal/common" deepdocpdf "ragflow/internal/deepdoc/parser/pdf" - pdfpos "ragflow/internal/deepdoc/parser/pdf/type" "ragflow/internal/deepdoc/parser/pdf/util" deepdoctype "ragflow/internal/deepdoc/parser/type" "ragflow/internal/ingestion/component" @@ -118,8 +119,7 @@ func cropImageChunks(ctx context.Context, engine deepdoctype.PDFEngine, chunks [ } out := make([]schema.ChunkDoc, len(chunks)) - sem := make(chan struct{}, cropConcurrency()) - var wg sync.WaitGroup + var candidateIndices []int for i := range chunks { ck := chunks[i] out[i] = ck @@ -145,22 +145,52 @@ func cropImageChunks(ctx context.Context, engine deepdoctype.PDFEngine, chunks [ if len(raw) == 0 { raw = ck.Positions } - var matrix [][]any - if err := json.Unmarshal(raw, &matrix); err != nil || len(matrix) == 0 { + if len(raw) == 0 { continue } - positions := util.PositionsFromMatrix(matrix) - if len(positions) == 0 { - continue + candidateIndices = append(candidateIndices, i) + } + if len(candidateIndices) == 0 { + return out + } + + total := int64(len(candidateIndices)) + var settled atomic.Int64 + sem := make(chan struct{}, cropConcurrency()) + var wg sync.WaitGroup + for _, idx := range candidateIndices { + if ctx.Err() != nil { + break } wg.Add(1) sem <- struct{}{} - go func(i int, ck schema.ChunkDoc, positions []pdfpos.Position) { + go func(chunkIdx int) { defer wg.Done() defer func() { <-sem }() - if err := ctx.Err(); err != nil { + if ctx.Err() != nil { return } + defer func() { + if ctx.Err() == nil { + done := settled.Add(1) + agentruntime.ReportComponentFraction(ctx, float64(done)/float64(total)) + } + }() + + ck := out[chunkIdx] + raw := ck.PDFPositions + if len(raw) == 0 { + raw = ck.Positions + } + var matrix [][]any + if err := json.Unmarshal(raw, &matrix); err != nil || len(matrix) == 0 { + return + } + positions := util.PositionsFromMatrix(matrix) + if len(positions) == 0 { + return + } + single := make(map[int]image.Image, len(positions)) acquired := make([]int, 0, len(positions)) for _, pos := range positions { @@ -189,7 +219,7 @@ func cropImageChunks(ctx context.Context, engine deepdoctype.PDFEngine, chunks [ if img == "" { return } - out[i].Image = "data:image/png;base64," + img + out[chunkIdx].Image = "data:image/png;base64," + img // Stream the freshly cropped preview to object storage and drop // the in-memory base64 immediately, instead of carrying every @@ -200,7 +230,7 @@ func cropImageChunks(ctx context.Context, engine deepdoctype.PDFEngine, chunks [ // only in canvas debug (dry-run) mode, where no persist stage // runs and the decorator's debug branch drops the raw bytes. if kbID, docID := resolveImageUploadContext(ctx, nil); kbID != "" { - if raw, derr := base64.StdEncoding.DecodeString(img); derr == nil { + if rawBytes, derr := base64.StdEncoding.DecodeString(img); derr == nil { // Key the streamed upload under the chunk's canonical id, // so the MinIO object is stored under exactly the key the // decorator (imageUploadDecorator) later exposes as @@ -214,11 +244,11 @@ func cropImageChunks(ctx context.Context, engine deepdoctype.PDFEngine, chunks [ // positions, so reading the canonical text here does not // alter the cropped image or the chunker's later output // text. - chunkID := canonicalChunkID(docID, out[i]) - if imgID, uerr := uploadOneImage(ctx, ChunkImageUploader, kbID, chunkID, raw); uerr == nil { - out[i].ImgID = imgID - out[i].Image = "" - out[i].ID = chunkID + chunkID := canonicalChunkID(docID, out[chunkIdx]) + if imgID, uerr := uploadOneImage(ctx, ChunkImageUploader, kbID, chunkID, rawBytes); uerr == nil { + out[chunkIdx].ImgID = imgID + out[chunkIdx].Image = "" + out[chunkIdx].ID = chunkID } else { // Note: the document text is intentionally NOT logged // here (CWE-532). The upload is retried at the persist @@ -228,9 +258,12 @@ func cropImageChunks(ctx context.Context, engine deepdoctype.PDFEngine, chunks [ } } } - }(i, ck, positions) + }(idx) } wg.Wait() + if ctx.Err() == nil { + agentruntime.ReportComponentFraction(ctx, 1.0) + } return out } diff --git a/internal/ingestion/component/chunker/pdfcrop_progress_cgo_test.go b/internal/ingestion/component/chunker/pdfcrop_progress_cgo_test.go new file mode 100644 index 0000000000..93d2bd2b58 --- /dev/null +++ b/internal/ingestion/component/chunker/pdfcrop_progress_cgo_test.go @@ -0,0 +1,133 @@ +//go:build cgo + +package chunker + +import ( + "context" + "math" + "sort" + "sync" + "testing" + + agentruntime "ragflow/internal/agent/runtime" + "ragflow/internal/ingestion/component/schema" +) + +type chunkFractionRecorder struct { + mu sync.Mutex + fractions []float64 +} + +func (r *chunkFractionRecorder) context(t *testing.T, component string) context.Context { + t.Helper() + ctx := agentruntime.WithProgressFractionCallback(t.Context(), func(c string, f float64) { + r.mu.Lock() + defer r.mu.Unlock() + r.fractions = append(r.fractions, f) + }) + return agentruntime.BindComponentFraction(ctx, component) +} + +func (r *chunkFractionRecorder) getFractions() []float64 { + r.mu.Lock() + defer r.mu.Unlock() + out := make([]float64, len(r.fractions)) + copy(out, r.fractions) + return out +} + +func TestCropImageChunks_ReportsFractions(t *testing.T) { + rec := &chunkFractionRecorder{} + ctx := rec.context(t, "Chunker:pdf") + + chunks := []schema.ChunkDoc{ + { + Text: "content 1", + Positions: jsonPositions(t, []float64{1, 0, 100, 0, 100}), + CKType: "text", + }, + { + Text: "content 2", + Positions: jsonPositions(t, []float64{1, 100, 200, 100, 200}), + CKType: "text", + }, + { + Text: "content 3", + Positions: jsonPositions(t, []float64{1, 200, 300, 200, 300}), + CKType: "text", + }, + } + + out := cropImageChunks(ctx, mockCropEngine{}, chunks) + if len(out) != 3 { + t.Fatalf("expected 3 chunks, got %d", len(out)) + } + + fractions := rec.getFractions() + // 3 tasks settled + 1 terminal 1.0 = 4 reports + if len(fractions) != 4 { + t.Fatalf("expected 4 reports, got %d (%v)", len(fractions), fractions) + } + if last := fractions[len(fractions)-1]; last != 1.0 { + t.Errorf("last fraction = %f, want 1.0", last) + } + + // Verify the multiset of reported fractions regardless of concurrent arrival order. + sorted := make([]float64, len(fractions)) + copy(sorted, fractions) + sort.Float64s(sorted) + + wantFractions := []float64{1.0 / 3.0, 2.0 / 3.0, 1.0, 1.0} + for i, want := range wantFractions { + if math.Abs(sorted[i]-want) > 1e-6 { + t.Errorf("sorted fraction[%d] = %f, want %f", i, sorted[i], want) + } + } +} + +func TestCropImageChunks_ZeroTasksReportsNothing(t *testing.T) { + rec := &chunkFractionRecorder{} + ctx := rec.context(t, "Chunker:pdf") + + chunks := []schema.ChunkDoc{ + { + Text: "no positions", + CKType: "text", + }, + } + + out := cropImageChunks(ctx, mockCropEngine{}, chunks) + if len(out) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(out)) + } + + fractions := rec.getFractions() + if len(fractions) != 0 { + t.Errorf("expected 0 reports for 0 crop tasks, got %v", fractions) + } +} + +func TestCropImageChunks_CancelledContextDoesNotReport(t *testing.T) { + rec := &chunkFractionRecorder{} + ctx, cancel := context.WithCancel(rec.context(t, "Chunker:pdf")) + cancel() // Cancel context before invocation + + chunks := []schema.ChunkDoc{ + { + Text: "content 1", + Positions: jsonPositions(t, []float64{1, 0, 100, 0, 100}), + CKType: "text", + }, + { + Text: "content 2", + Positions: jsonPositions(t, []float64{1, 100, 200, 100, 200}), + CKType: "text", + }, + } + + _ = cropImageChunks(ctx, mockCropEngine{}, chunks) + fractions := rec.getFractions() + if len(fractions) != 0 { + t.Errorf("cancelled context must not report progress, got: %v", fractions) + } +} diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index b2cd7624eb..ad4b84e13d 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -899,50 +899,6 @@ func (c *ExtractorComponent) runAutoSummary(ctx context.Context, db *gorm.DB, in return nil } -// fractionWindow is the slice of a component's 0..1 fraction budget owned by one -// sequential phase. Splitting the budget keeps the reported fraction climbing -// across phases: the run-level progress tracker keeps a high-water mark, so a -// second phase restarting from zero would freeze the bar until the component -// exits. -type fractionWindow struct{ base, span float64 } - -// report maps a phase completion ratio into the window and forwards it to the -// run-level fraction channel. The zero window reports nothing, which keeps -// disabled phases free of budget arithmetic. -func (w fractionWindow) report(ctx context.Context, ratio float64) { - if w.span <= 0 { - return - } - ratio = min(max(ratio, 0), 1) - runtime.ReportComponentFraction(ctx, w.base+w.span*ratio) -} - -// fractionSplit hands out one equal window per enabled phase, in phase order. -type fractionSplit struct { - phases int - next int -} - -func newFractionSplit(enabled ...bool) *fractionSplit { - s := &fractionSplit{} - for _, on := range enabled { - if on { - s.phases++ - } - } - return s -} - -func (s *fractionSplit) take(enabled bool) fractionWindow { - if !enabled { - return fractionWindow{} - } - span := 1 / float64(s.phases) - w := fractionWindow{base: float64(s.next) * span, span: span} - s.next++ - return w -} - // runAutoKeywordsPool dispatches keyword extraction across all chunks concurrently // using extractorPool before the tagging stage. func (c *ExtractorComponent) runAutoKeywordsPool(ctx context.Context, db *gorm.DB, in extractorInputs, w fractionWindow) error { diff --git a/internal/ingestion/component/media_dispatch.go b/internal/ingestion/component/media_dispatch.go index 044c9be8f8..10fdeda2cb 100644 --- a/internal/ingestion/component/media_dispatch.go +++ b/internal/ingestion/component/media_dispatch.go @@ -115,7 +115,7 @@ func maybeDispatchImage( text, err = extractImageText(ctx, img) } release() - parsed := dispatchParse(ctx, fileType, filename, binary, setups) + parsed := dispatchParse(ctx, fileType, filename, binary, setups, fractionWindow{}) if parsed.Err != nil { return parsed, true, parsed.Err } diff --git a/internal/ingestion/component/parser.go b/internal/ingestion/component/parser.go index 90ea5c15ab..ec8aacbdf5 100644 --- a/internal/ingestion/component/parser.go +++ b/internal/ingestion/component/parser.go @@ -523,12 +523,21 @@ func (c *ParserComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[st } } if !handledVision && !handledMedia && !handledImage && !handledAudio { - dispatched = dispatchParse(ctx, fileTypeExt, filename, binary, setups) + parsePagesOn := pdfPageProgressSupported(fileTypeExt, setups) + enhanceOn := c.enableVisionEnhancement + split := newFractionSplit(parsePagesOn, enhanceOn) + pageWin := split.take(parsePagesOn) + visionWin := split.take(enhanceOn) + + dispatched = dispatchParse(ctx, fileTypeExt, filename, binary, setups, pageWin) + if dispatched.Err == nil && ctx.Err() == nil { + pageWin.report(ctx, 1.0) + } if c.enableVisionEnhancement { // Enhancement is optional; parser-provided text and image metadata // remain available if a vision model cannot describe an image. - dispatched, _, _ = maybeDispatchVisionEnhancement(ctx, db, fileTypeExt, dispatched, inputs, c.vision) + dispatched, _, _ = maybeDispatchVisionEnhancement(ctx, db, fileTypeExt, dispatched, inputs, c.vision, visionWin) } } // Known/supported families must fail loudly when dispatch or diff --git a/internal/ingestion/component/parser_dispatch.go b/internal/ingestion/component/parser_dispatch.go index 9552a2cab3..460bdae1df 100644 --- a/internal/ingestion/component/parser_dispatch.go +++ b/internal/ingestion/component/parser_dispatch.go @@ -28,7 +28,6 @@ import ( "fmt" "strings" - "ragflow/internal/agent/runtime" "ragflow/internal/common" "ragflow/internal/ingestion/component/schema" "ragflow/internal/parser/parser" @@ -77,7 +76,7 @@ func configureParserFromSetups(p any, fileType utility.FileType, setups map[stri // re-reading setups. lib_type is no longer threaded through: the // Python dispatcher picks a single backend per family and the Go // constructors mirror that. -func dispatchParse(ctx context.Context, fileType utility.FileType, filename string, data []byte, setups map[string]schema.ParserSetup) parser.ParseResult { +func dispatchParse(ctx context.Context, fileType utility.FileType, filename string, data []byte, setups map[string]schema.ParserSetup, window fractionWindow) parser.ParseResult { if fileType == utility.FileTypeOTHER { // Unknown / unset family. The component treats the bytes // as text pages; splitIntoPages handles it. We return no @@ -97,7 +96,7 @@ func dispatchParse(ctx context.Context, fileType utility.FileType, filename stri return parser.ParseResult{Err: fmt.Errorf("parser: resolve %q: %w", fileType, err)} } configureParserFromSetups(p, fileType, setups) - attachPDFPageProgress(ctx, p) + attachPDFPageProgress(ctx, p, window) res := p.ParseWithResult(ctx, filename, data) if res.Err != nil { @@ -114,21 +113,42 @@ func dispatchParse(ctx context.Context, fileType utility.FileType, filename stri return res } +// pdfPageProgressSupported reports whether the configured parser setup produces +// in-flight page progress events (via DeepDOC OnPageDone). Remote engines +// (MinerU, PaddleOCR, Docling, etc.) and non-PDF families do not report page progress. +func pdfPageProgressSupported(fileType utility.FileType, setups map[string]schema.ParserSetup) bool { + if fileType != utility.FileTypePDF { + return false + } + var method string + if setup, ok := setups[resolveParserFamily(fileType)]; ok { + if s, ok := setup["parse_method"].(string); ok { + method = strings.ToLower(strings.TrimSpace(s)) + } + } + switch method { + case "", "deepdoc": + return true + default: + return false + } +} + // attachPDFPageProgress wires the DeepDOC page callback so a long PDF // parse reports in-flight completion instead of staying silent until // the component exits. Only the local DeepDOC backend fires it; remote // engines (MinerU, PaddleOCR, Docling, …) never call the callback and // the fraction stays at zero. -func attachPDFPageProgress(ctx context.Context, p parser.ParseResultProducer) { +func attachPDFPageProgress(ctx context.Context, p parser.ParseResultProducer, window fractionWindow) { pdfParser, ok := p.(*parser.PDFParser) - if !ok { + if !ok || window.span <= 0 { return } pdfParser.OnPageDone = func(done, total int) { if total <= 0 { return } - runtime.ReportComponentFraction(ctx, float64(done)/float64(total)) + window.report(ctx, float64(done)/float64(total)) } } diff --git a/internal/ingestion/component/parser_dispatch_progress_test.go b/internal/ingestion/component/parser_dispatch_progress_test.go index dc2e961d64..56bb6667a7 100644 --- a/internal/ingestion/component/parser_dispatch_progress_test.go +++ b/internal/ingestion/component/parser_dispatch_progress_test.go @@ -9,7 +9,9 @@ import ( "testing" "ragflow/internal/agent/runtime" + "ragflow/internal/ingestion/component/schema" "ragflow/internal/parser/parser" + "ragflow/internal/utility" ) type fractionObservation struct { @@ -32,7 +34,7 @@ func TestAttachPDFPageProgress_ReportsFraction(t *testing.T) { ctx := boundFractionContext(t, "Parser:pdf", &seen) pdfParser := parser.NewPDFParser() - attachPDFPageProgress(ctx, pdfParser) + attachPDFPageProgress(ctx, pdfParser, fractionWindow{base: 0, span: 1}) if pdfParser.OnPageDone == nil { t.Fatal("expected OnPageDone to be wired") } @@ -56,7 +58,7 @@ func TestAttachPDFPageProgress_IgnoresUnknownTotal(t *testing.T) { ctx := boundFractionContext(t, "Parser:pdf", &seen) pdfParser := parser.NewPDFParser() - attachPDFPageProgress(ctx, pdfParser) + attachPDFPageProgress(ctx, pdfParser, fractionWindow{base: 0, span: 1}) pdfParser.OnPageDone(1, 0) if len(seen) != 0 { @@ -68,9 +70,72 @@ func TestAttachPDFPageProgress_LeavesOtherParsersAlone(t *testing.T) { var seen []fractionObservation ctx := boundFractionContext(t, "Parser:txt", &seen) - attachPDFPageProgress(ctx, parser.NewTextParser()) + attachPDFPageProgress(ctx, parser.NewTextParser(), fractionWindow{base: 0, span: 1}) if len(seen) != 0 { t.Errorf("expected no fraction from a non-PDF parser, got %v", seen) } } + +func TestAttachPDFPageProgress_ZeroWindowSilent(t *testing.T) { + var seen []fractionObservation + ctx := boundFractionContext(t, "Parser:pdf", &seen) + + pdfParser := parser.NewPDFParser() + attachPDFPageProgress(ctx, pdfParser, fractionWindow{}) + if pdfParser.OnPageDone != nil { + t.Fatal("expected OnPageDone to remain nil for zero window") + } + if len(seen) != 0 { + t.Errorf("expected no observations for zero window, got %v", seen) + } +} + +func TestAttachPDFPageProgress_ReportsFractionInWindow(t *testing.T) { + var seen []fractionObservation + ctx := boundFractionContext(t, "Parser:pdf", &seen) + + pdfParser := parser.NewPDFParser() + attachPDFPageProgress(ctx, pdfParser, fractionWindow{base: 0.1, span: 0.4}) + if pdfParser.OnPageDone == nil { + t.Fatal("expected OnPageDone to be wired") + } + + pdfParser.OnPageDone(3, 12) // 0.25 -> 0.1 + 0.4*0.25 = 0.20 + pdfParser.OnPageDone(12, 12) // 1.0 -> 0.1 + 0.4*1.0 = 0.50 + + want := []fractionObservation{{"Parser:pdf", 0.2}, {"Parser:pdf", 0.5}} + if len(seen) != len(want) { + t.Fatalf("observations = %v, want %v", seen, want) + } + for i := range want { + if seen[i] != want[i] { + t.Errorf("observation %d = %v, want %v", i, seen[i], want[i]) + } + } +} + +func TestPDFPageProgressSupported(t *testing.T) { + tests := []struct { + fileType utility.FileType + method string + want bool + }{ + {utility.FileTypePDF, "", true}, + {utility.FileTypePDF, "deepdoc", true}, + {utility.FileTypePDF, "mineru", false}, + {utility.FileTypePDF, "paddleocr", false}, + {utility.FileTypePDF, "docling", false}, + {utility.FileTypePDF, "plain_text", false}, + {utility.FileTypeDOCX, "deepdoc", false}, + {utility.FileTypeMarkdown, "", false}, + } + for _, tc := range tests { + setups := map[string]schema.ParserSetup{ + "pdf": {"parse_method": tc.method}, + } + if got := pdfPageProgressSupported(tc.fileType, setups); got != tc.want { + t.Errorf("pdfPageProgressSupported(%s, method=%s) = %v, want %v", tc.fileType, tc.method, got, tc.want) + } + } +} diff --git a/internal/ingestion/component/progress_fraction.go b/internal/ingestion/component/progress_fraction.go new file mode 100644 index 0000000000..2f7c10d656 --- /dev/null +++ b/internal/ingestion/component/progress_fraction.go @@ -0,0 +1,51 @@ +package component + +import ( + "context" + + "ragflow/internal/agent/runtime" +) + +// fractionWindow is the slice of a component's 0..1 fraction budget owned by one +// sequential phase. Splitting the budget keeps the reported fraction climbing +// across phases: the run-level progress tracker keeps a high-water mark, so a +// second phase restarting from zero would freeze the bar until the component +// exits. +type fractionWindow struct{ base, span float64 } + +// report maps a phase completion ratio into the window and forwards it to the +// run-level fraction channel. The zero window reports nothing, which keeps +// disabled phases free of budget arithmetic. +func (w fractionWindow) report(ctx context.Context, ratio float64) { + if w.span <= 0 { + return + } + ratio = min(max(ratio, 0), 1) + runtime.ReportComponentFraction(ctx, w.base+w.span*ratio) +} + +// fractionSplit hands out one equal window per enabled phase, in phase order. +type fractionSplit struct { + phases int + next int +} + +func newFractionSplit(enabled ...bool) *fractionSplit { + s := &fractionSplit{} + for _, on := range enabled { + if on { + s.phases++ + } + } + return s +} + +func (s *fractionSplit) take(enabled bool) fractionWindow { + if !enabled { + return fractionWindow{} + } + span := 1 / float64(s.phases) + w := fractionWindow{base: float64(s.next) * span, span: span} + s.next++ + return w +} diff --git a/internal/ingestion/component/vision_enhancement.go b/internal/ingestion/component/vision_enhancement.go index bcdae1d68a..dbd2c8d835 100644 --- a/internal/ingestion/component/vision_enhancement.go +++ b/internal/ingestion/component/vision_enhancement.go @@ -32,6 +32,7 @@ import ( "path/filepath" "strings" "sync" + "sync/atomic" "time" "go.uber.org/zap" @@ -166,7 +167,10 @@ func maybeDispatchVisionEnhancement( dispatched parser.ParseResult, inputs map[string]any, vision visionSettings, + window fractionWindow, ) (parser.ParseResult, bool, error) { + w := window + // Only enhance successful JSON output format containing items. if dispatched.Err != nil || dispatched.OutputFormat != "json" || len(dispatched.JSON) == 0 { return dispatched, false, nil @@ -249,6 +253,8 @@ func maybeDispatchVisionEnhancement( failures := make([]string, len(items)) var wg sync.WaitGroup sem := make(chan struct{}, visionEnhancementConcurrency) + var settled atomic.Int64 + total := int64(len(items)) for slot, itemIdx := range items { if err := ctx.Err(); err != nil { break @@ -318,6 +324,8 @@ func maybeDispatchVisionEnhancement( } if resource == nil { releaseSlot() + n := settled.Add(1) + w.report(ctx, float64(n)/float64(total)) continue } if !resource.VLMDataValidated && !isUsableVisionImage(resource.VLMData) { @@ -325,12 +333,20 @@ func maybeDispatchVisionEnhancement( failures[slot] = "image payload is invalid or exceeds configured limits" } releaseSlot() + n := settled.Add(1) + w.report(ctx, float64(n)/float64(total)) continue } wg.Add(1) go func(slot int, imageData string) { defer wg.Done() defer releaseSlot() + defer func() { + if ctx.Err() == nil { + n := settled.Add(1) + w.report(ctx, float64(n)/float64(total)) + } + }() messages := buildVisionMessages(prompt, imageData) if len(messages) == 0 { @@ -352,6 +368,7 @@ func maybeDispatchVisionEnhancement( if err := ctx.Err(); err != nil { return dispatched, modified, err } + w.report(ctx, 1.0) failedItems := 0 firstFailure := "" for _, failure := range failures { diff --git a/internal/ingestion/component/vision_enhancement_progress_test.go b/internal/ingestion/component/vision_enhancement_progress_test.go new file mode 100644 index 0000000000..74626dec21 --- /dev/null +++ b/internal/ingestion/component/vision_enhancement_progress_test.go @@ -0,0 +1,180 @@ +package component + +import ( + "context" + "errors" + "math" + "sort" + "testing" + + "ragflow/internal/dao" + modelModule "ragflow/internal/entity/models" + "ragflow/internal/parser/parser" + "ragflow/internal/utility" +) + +func TestVisionEnhancement_ReportsPerFigureFractions(t *testing.T) { + invoker := &visionEnhanceCaptureInvoker{} + swapVisionGlobals(t, fakeResolver, invoker.invoke, fakePrompt) + + rec := &fractionRecorder{} + ctx := rec.context(t, "Parser:vision") + + png := visionTestPNGBase64(t) + dispatched := parser.ParseResult{ + OutputFormat: "json", + JSON: []map[string]any{ + {"text": "1", "image": png, "doc_type_kwd": "image"}, + {"text": "2", "image": png, "doc_type_kwd": "image"}, + {"text": "3", "image": png, "doc_type_kwd": "image"}, + {"text": "4", "image": png, "doc_type_kwd": "image"}, + }, + } + + win := fractionWindow{base: 0.5, span: 0.5} + _, handled, err := maybeDispatchVisionEnhancement( + ctx, dao.DB, utility.FileTypeXLSX, dispatched, + map[string]any{"tenant_id": "t1"}, + visionSettings{}, + win, + ) + if err != nil { + t.Fatalf("maybeDispatchVisionEnhancement: %v", err) + } + if !handled { + t.Fatal("handled = false, want true") + } + + fractions := rec.fractions() + if len(fractions) != 5 { + t.Fatalf("recorded fractions count = %d, want 5 (4 items + 1 close); got %v", len(fractions), fractions) + } + if last := fractions[len(fractions)-1]; last != 1.0 { + t.Errorf("last fraction = %f, want 1.0", last) + } + + // Verify the multiset of reported fractions regardless of concurrent arrival order. + sorted := make([]float64, len(fractions)) + copy(sorted, fractions) + sort.Float64s(sorted) + + wantFractions := []float64{0.625, 0.75, 0.875, 1.0, 1.0} + for i, want := range wantFractions { + if math.Abs(sorted[i]-want) > 1e-6 { + t.Errorf("sorted fraction[%d] = %f, want %f", i, sorted[i], want) + } + } +} + +func TestVisionEnhancement_SettledFractionsCountFailedFigures(t *testing.T) { + // One valid image, one invalid image that fails payload validation before VLM call, + // and one where the VLM invoker returns error. + swapVisionGlobals(t, fakeResolver, func(_ context.Context, _ modelModule.ModelDriver, _ string, _ []modelModule.Message, _ *modelModule.APIConfig) (*modelModule.ChatResponse, error) { + return nil, errors.New("invoker error") + }, fakePrompt) + + rec := &fractionRecorder{} + ctx := rec.context(t, "Parser:vision") + + dispatched := parser.ParseResult{ + OutputFormat: "json", + JSON: []map[string]any{ + {"text": "valid", "image": visionTestPNGBase64(t), "doc_type_kwd": "image"}, + {"text": "invalid", "image": "not-valid-base64", "doc_type_kwd": "image"}, + }, + } + + win := fractionWindow{base: 0.2, span: 0.8} + _, _, err := maybeDispatchVisionEnhancement( + ctx, dao.DB, utility.FileTypeXLSX, dispatched, + map[string]any{"tenant_id": "t1"}, + visionSettings{}, + win, + ) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + fractions := rec.fractions() + // 2 items + 1 close = 3 reports + if len(fractions) != 3 { + t.Fatalf("recorded fractions count = %d, want 3; got %v", len(fractions), fractions) + } + if last := fractions[len(fractions)-1]; last != 1.0 { + t.Errorf("last fraction = %f, want 1.0", last) + } + + sorted := make([]float64, len(fractions)) + copy(sorted, fractions) + sort.Float64s(sorted) + + wantFractions := []float64{0.6, 1.0, 1.0} + for i, want := range wantFractions { + if math.Abs(sorted[i]-want) > 1e-6 { + t.Errorf("sorted fraction[%d] = %f, want %f", i, sorted[i], want) + } + } +} + +func TestVisionEnhancement_ZeroWindowReportsNothing(t *testing.T) { + invoker := &visionEnhanceCaptureInvoker{} + swapVisionGlobals(t, fakeResolver, invoker.invoke, fakePrompt) + + rec := &fractionRecorder{} + ctx := rec.context(t, "Parser:vision") + + dispatched := parser.ParseResult{ + OutputFormat: "json", + JSON: []map[string]any{ + {"text": "1", "image": visionTestPNGBase64(t), "doc_type_kwd": "image"}, + }, + } + + // Zero window (default when unpassed or inactive) + _, handled, err := maybeDispatchVisionEnhancement( + ctx, dao.DB, utility.FileTypeXLSX, dispatched, + map[string]any{"tenant_id": "t1"}, + visionSettings{}, + fractionWindow{}, + ) + if err != nil { + t.Fatalf("maybeDispatchVisionEnhancement: %v", err) + } + if !handled { + t.Fatal("handled = false, want true") + } + + if len(rec.fractions()) != 0 { + t.Errorf("expected 0 fractions for zero window, got %v", rec.fractions()) + } +} + +func TestVisionEnhancement_CancelledContextReportsNothing(t *testing.T) { + invoker := &visionEnhanceCaptureInvoker{} + swapVisionGlobals(t, fakeResolver, invoker.invoke, fakePrompt) + + rec := &fractionRecorder{} + ctx, cancel := context.WithCancel(rec.context(t, "Parser:vision")) + cancel() // Cancel context before invocation + + png := visionTestPNGBase64(t) + dispatched := parser.ParseResult{ + OutputFormat: "json", + JSON: []map[string]any{ + {"text": "1", "image": png, "doc_type_kwd": "image"}, + {"text": "2", "image": png, "doc_type_kwd": "image"}, + }, + } + + win := fractionWindow{base: 0.5, span: 0.5} + _, _, _ = maybeDispatchVisionEnhancement( + ctx, dao.DB, utility.FileTypeXLSX, dispatched, + map[string]any{"tenant_id": "t1"}, + visionSettings{}, + win, + ) + + if len(rec.fractions()) != 0 { + t.Errorf("cancelled context must not report progress, got: %v", rec.fractions()) + } +} diff --git a/internal/ingestion/component/vision_enhancement_test.go b/internal/ingestion/component/vision_enhancement_test.go index ad948aa985..fd621e39cd 100644 --- a/internal/ingestion/component/vision_enhancement_test.go +++ b/internal/ingestion/component/vision_enhancement_test.go @@ -204,7 +204,7 @@ func TestVisionEnhancement_AppendsVLMToExistingText(t *testing.T) { res, handled, err := maybeDispatchVisionEnhancement( t.Context(), dao.DB, utility.FileTypeXLSX, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("maybeDispatchVisionEnhancement: %v", err) } @@ -234,7 +234,7 @@ func TestVisionEnhancement_ReportsVLMInvocationFailure(t *testing.T) { } result, _, err := maybeDispatchVisionEnhancement(t.Context(), dao.DB, utility.FileTypeDOCX, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("maybeDispatchVisionEnhancement: %v", err) } @@ -262,8 +262,7 @@ func TestVisionEnhancement_CropFailureFallsBackToInlineImage(t *testing.T) { } result, handled, err := maybeDispatchVisionEnhancement( t.Context(), dao.DB, utility.FileTypePDF, dispatched, - map[string]any{"tenant_id": "t1"}, visionSettings{}, - ) + map[string]any{"tenant_id": "t1"}, visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("maybeDispatchVisionEnhancement: %v", err) } @@ -298,8 +297,7 @@ func TestVisionEnhancement_EnhancesPDFTablesRegardlessOfOCRProvider(t *testing.T res, handled, err := maybeDispatchVisionEnhancement( t.Context(), dao.DB, utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}, - ) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("maybeDispatchVisionEnhancement: %v", err) } @@ -335,8 +333,7 @@ func TestVisionEnhancement_EnhancesTableAndCellImages(t *testing.T) { result, handled, err := maybeDispatchVisionEnhancement( t.Context(), dao.DB, utility.FileTypePDF, dispatched, - map[string]any{"tenant_id": "t1"}, visionSettings{}, - ) + map[string]any{"tenant_id": "t1"}, visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("maybeDispatchVisionEnhancement: %v", err) } @@ -365,8 +362,7 @@ func TestVisionEnhancement_EnhancesExternalPDFImage(t *testing.T) { result, handled, err := maybeDispatchVisionEnhancement( t.Context(), dao.DB, utility.FileTypePDF, dispatched, - map[string]any{"tenant_id": "t1"}, visionSettings{}, - ) + map[string]any{"tenant_id": "t1"}, visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("maybeDispatchVisionEnhancement: %v", err) } @@ -417,7 +413,7 @@ func TestVisionEnhancement_EnhancesJSONImagesAndTables(t *testing.T) { tc.fileType, dispatched, map[string]any{"tenant_id": "t1", "lang": "Japanese"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -497,7 +493,7 @@ func TestVisionEnhancement_LanguagePriority(t *testing.T) { _, handled, err := maybeDispatchVisionEnhancement( t.Context(), dao.DB, utility.FileTypePDF, dispatched, tc.inputs, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -538,7 +534,7 @@ func TestVisionEnhancement_MarkdownOutputUntouched(t *testing.T) { utility.FileTypeDOCX, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -569,7 +565,7 @@ func TestVisionEnhancement_UsesVisualPayloadRegardlessOfFileType(t *testing.T) { utility.FileTypeOTHER, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -599,7 +595,7 @@ func TestVisionEnhancement_EmptyOrNoTenantSkipped(t *testing.T) { utility.FileTypeDOCX, dispatched, map[string]any{}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil || handled { t.Errorf("handled=%v, err=%v, want false, nil for missing tenant_id", handled, err) } @@ -626,7 +622,7 @@ func TestVisionEnhancement_DispatchedErrSkipped(t *testing.T) { utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -664,7 +660,7 @@ func TestVisionEnhancement_ContextCancellation(t *testing.T) { utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if !errors.Is(err, context.Canceled) { t.Fatalf("err = %v, want context.Canceled", err) } @@ -706,7 +702,7 @@ func TestVisionEnhancement_NonStringImageFieldFiltered(t *testing.T) { utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -749,7 +745,7 @@ func TestVisionEnhancement_MoreThanConcurrencyItems(t *testing.T) { utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -790,7 +786,7 @@ func TestVisionEnhancement_PlainTextResponseNotTruncated(t *testing.T) { utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -829,7 +825,7 @@ func TestVisionEnhancement_PromptBuilderErrorSkipped(t *testing.T) { utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -868,7 +864,7 @@ func TestVisionEnhancement_ModelResolveFailureSkipped(t *testing.T) { utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -920,7 +916,7 @@ func TestVisionEnhancement_CancellationStopsSchedulingWithManyItems(t *testing.T utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if !errors.Is(err, context.Canceled) { t.Fatalf("err = %v, want context.Canceled", err) } @@ -1170,8 +1166,7 @@ func TestVisionEnhancement_GlobalModelPreferredOverPDFParser(t *testing.T) { utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{modelID: "custom-vlm@provider"}, - ) + visionSettings{modelID: "custom-vlm@provider"}, fractionWindow{}) if err != nil { t.Fatalf("maybeDispatchVisionEnhancement: %v", err) } @@ -1213,7 +1208,7 @@ func TestVisionEnhancement_InvalidImageDataSkipped(t *testing.T) { utility.FileTypePDF, dispatched, map[string]any{"tenant_id": "t1"}, - visionSettings{}) + visionSettings{}, fractionWindow{}) if err != nil { t.Fatalf("maybeDispatchVisionEnhancement: %v", err) }