feat(ingestion): add progress fraction reporting to VisionEnhancement and cropImageChunks (#20800)

This commit is contained in:
jay77721 authored and GitHub committed 2026-10-12 13:47:54 +08:00
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)
}
}
-44
View File
@@ -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
}
+11 -2
View File
@@ -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)
}