mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-11-10 21:32:43 +08:00
feat(ingestion): add progress fraction reporting to VisionEnhancement and cropImageChunks (#20800)
This commit is contained in:
1 parent
efc1a310f1
commit
35de240c39
11 files changed
+559
-100
No files matched your search
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user