diff --git a/internal/ingestion/component/chunker/group.go b/internal/ingestion/component/chunker/group.go index d69924735b..21fb0082e0 100644 --- a/internal/ingestion/component/chunker/group.go +++ b/internal/ingestion/component/chunker/group.go @@ -36,6 +36,8 @@ import ( "context" "encoding/json" "fmt" + "regexp" + "sort" "strings" "go.uber.org/zap" @@ -115,7 +117,7 @@ func invokeGroup(_ context.Context, inputs map[string]any, p *titleChunkerParam) if len(records) == 0 { return emptyOutputs(), nil } - ctx := newLevelContext(records, p) + ctx := newLevelContext(records, outlineFromInputs(inputs), p) levels := ctx.Levels() // Count heading level distribution for debugging. headingCounts := make(map[int]int) @@ -256,7 +258,10 @@ func buildChunksFromRecordGroups(groups [][]lineRecord, p *titleChunkerParam, pl if len(g) == 0 { continue } - chunk := map[string]any{"text": joinGroupText(g)} + // Strip parser-emitted position tags (`@@...##`) from the + // joined text (diff 1.6 / 2.8). Mirrors common.py:255 + // `RAGFlowPdfParser.remove_tag("".join(...))`. + chunk := map[string]any{"text": removeTag(joinGroupText(g))} if !plain { first := g[0] if first.docType != "" { @@ -266,6 +271,24 @@ func buildChunksFromRecordGroups(groups [][]lineRecord, p *titleChunkerParam, pl chunk["img_id"] = *first.imgID } } + // Merge PDF coordinate matrices across the merged records + // instead of keeping only the leading record's (diff 1.6). + // Mirrors pdf_chunk_metadata.py:127 merge_pdf_positions. + var pdfSrc, posSrc []json.RawMessage + for _, r := range g { + if len(r.pdfPositions) > 0 { + pdfSrc = append(pdfSrc, r.pdfPositions) + } + if len(r.positions) > 0 { + posSrc = append(posSrc, r.positions) + } + } + if m := mergePositionMatrix(pdfSrc...); m != nil { + chunk["_pdf_positions"] = m + } + if m := mergePositionMatrix(posSrc...); m != nil { + chunk["positions"] = m + } chunks = append(chunks, chunk) } if p.RootChunkAsHeading && len(chunks) > 1 { @@ -278,6 +301,61 @@ func buildChunksFromRecordGroups(groups [][]lineRecord, p *titleChunkerParam, pl return chunks } +// posTagRemove matches parser-emitted position tags of the form +// `@@\t\t\t\t##`. Mirrors Python +// pdf_parser.py:1934 `re.sub(r"@@[\t0-9.-]+?##", "", txt)`. +var posTagRemove = regexp.MustCompile(`@@[\t0-9.-]+?##`) + +// removeTag strips parser-emitted position tags from a chunk's text. +// Mirrors deepdoc RAGFlowPdfParser.remove_tag. +func removeTag(text string) string { + return posTagRemove.ReplaceAllString(text, "") +} + +// mergePositionMatrix aggregates multiple PDF coordinate matrices into a +// single de-duplicated, sorted matrix. Mirrors Python +// pdf_chunk_metadata.py:127 merge_pdf_positions: rows are 5-tuples +// [page,left,right,top,bottom]; duplicates (by the first five columns) +// are dropped and rows are sorted by (page, top, left). Returns nil when +// no source carries any usable coordinates. +func mergePositionMatrix(sources ...json.RawMessage) [][]float64 { + var out [][]float64 + seen := make(map[string]bool) + for _, src := range sources { + if len(src) == 0 { + continue + } + var mat [][]float64 + if err := json.Unmarshal(src, &mat); err != nil { + continue + } + for _, row := range mat { + if len(row) < 5 { + continue + } + key := fmt.Sprintf("%v|%v|%v|%v|%v", row[0], row[1], row[2], row[3], row[4]) + if seen[key] { + continue + } + seen[key] = true + out = append(out, row) + } + } + if len(out) == 0 { + return nil + } + sort.Slice(out, func(i, j int) bool { + if out[i][0] != out[j][0] { + return out[i][0] < out[j][0] + } + if out[i][3] != out[j][3] { + return out[i][3] < out[j][3] + } + return out[i][1] < out[j][1] + }) + return out +} + // extractLineRecords reads the chunker inputs in the same order the // python BaseTitleChunker.extract_line_records uses: // @@ -333,12 +411,14 @@ func recordsFromStructured(items []schema.ChunkDoc) []lineRecord { meta[k] = json.RawMessage(v) } out = append(out, lineRecord{ - text: text, - docType: dt, - imgID: imgID, - layout: it.Layout, - ckType: it.CKType, - parentMeta: meta, + text: text, + docType: dt, + imgID: imgID, + layout: it.Layout, + ckType: it.CKType, + pdfPositions: it.PDFPositions, + positions: it.Positions, + parentMeta: meta, }) } return out diff --git a/internal/ingestion/component/chunker/group_test.go b/internal/ingestion/component/chunker/group_test.go index 2cb1ebae6e..3fe0c3b835 100644 --- a/internal/ingestion/component/chunker/group_test.go +++ b/internal/ingestion/component/chunker/group_test.go @@ -18,6 +18,7 @@ package chunker import ( "context" + "strings" "testing" "ragflow/internal/agent/runtime" @@ -230,6 +231,59 @@ func TestGroupChunker_StructuredMetadata(t *testing.T) { } } +// TestGroupChunker_MergesPDFPositionsAndRemovesTags is the TDD test for +// migration diffs Chunker-1.6 / 2.8: when the group chunker merges +// multiple adjacent text records into one chunk it must (a) strip the +// parser-emitted `@@...##` position tags from the joined text, and +// (b) MERGE (not drop) the `positions` coordinate matrices across the +// merged records — mirroring common.py:255 remove_tag + merge. +func TestGroupChunker_MergesPDFPositionsAndRemovesTags(t *testing.T) { + c, err := NewGroupTitleChunker(map[string]any{ + "levels": [][]string{{`^# `}}, + }) + if err != nil { + t.Fatalf("NewGroupTitleChunker: %v", err) + } + items := []map[string]any{ + {"text": "# Heading", "doc_type_kwd": "text"}, + {"text": "body one @@1\t10.0\t20.0\t30.0\t40.0## tail", "doc_type_kwd": "text", "positions": [][]float64{{1, 10, 20, 30, 40}}}, + {"text": "body two @@2\t15.0\t25.0\t35.0\t45.0## tail", "doc_type_kwd": "text", "positions": [][]float64{{2, 15, 25, 35, 45}}}, + } + out, err := c.Invoke(context.Background(), map[string]any{ + "name": "doc", + "output_format": "chunks", + "chunks": items, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) == 0 { + t.Fatal("no chunks emitted") + } + for _, ck := range chunks { + text, _ := ck["text"].(string) + // Only the merged body group carries both bodies. + if !strings.Contains(text, "body one") || !strings.Contains(text, "body two") { + continue + } + // (a) parser tags must be stripped from the text. + if strings.Contains(text, "@@") { + t.Errorf("parser position tags leaked into chunk text: %q", text) + } + // (b) positions must be merged across both records. + pos, ok := ck["positions"].([][]float64) + if !ok { + t.Fatalf("positions missing or wrong type %T on merged group chunk", ck["positions"]) + } + if len(pos) != 2 { + t.Errorf("merged positions = %d groups, want 2 (both records)", len(pos)) + } + return + } + t.Fatal("merged body group chunk not found in output") +} + func TestGroupTitleChunker_InvokeDeterministic(t *testing.T) { c, err := NewGroupTitleChunker(map[string]any{ "levels": [][]string{ diff --git a/internal/ingestion/component/chunker/hierarchy.go b/internal/ingestion/component/chunker/hierarchy.go index 79f81e3dce..a659b8b34c 100644 --- a/internal/ingestion/component/chunker/hierarchy.go +++ b/internal/ingestion/component/chunker/hierarchy.go @@ -148,7 +148,7 @@ func invokeHierarchy(_ context.Context, inputs map[string]any, p *titleChunkerPa if len(records) == 0 { return emptyOutputs(), nil } - ctx := newLevelContext(records, p) + ctx := newLevelContext(records, outlineFromInputs(inputs), p) levels := ctx.Levels() // Count heading level distribution for debugging. headingCounts := make(map[int]int) diff --git a/internal/ingestion/component/chunker/pdfcrop_cgo.go b/internal/ingestion/component/chunker/pdfcrop_cgo.go index 9bdf0d1d47..ce0768b4af 100644 --- a/internal/ingestion/component/chunker/pdfcrop_cgo.go +++ b/internal/ingestion/component/chunker/pdfcrop_cgo.go @@ -53,7 +53,9 @@ func newPDFEngineFromUpstream(ctx context.Context, up schema.ChunkerFromUpstream return deepdocpdf.NewEngine(data) } -// cropImageChunks crops image/table chunks in place. Each spanned page is +// cropImageChunks crops image/table chunks and renders text previews (for +// text chunks that carry PDF positions, mirroring Python +// restore_pdf_text_previews). Each spanned page is // rendered at most once. Chunks arrive in document order, so we keep only a // sliding window of page images: once we advance past a chunk whose minimum // page is P, no later chunk references a page < P, and we evict those entries @@ -139,9 +141,14 @@ func cropImageChunks(ctx context.Context, engine deepdoctype.PDFEngine, chunks [ return out } +// needsCrop reports whether a chunk should be cropped to a page-region +// preview from its PDF positions. Image/table chunks get their media region +// cropped; text chunks with positions get a rendered preview of the text +// region (Python restore_pdf_text_previews). A pre-existing Image is never +// re-cropped — cropImageChunks honors that separately. func needsCrop(ck schema.ChunkDoc) bool { switch ck.CKType { - case "image", "table": + case "image", "table", "text": return len(ck.PDFPositions) > 0 || len(ck.Positions) > 0 default: return false diff --git a/internal/ingestion/component/chunker/pdfcrop_cgo_test.go b/internal/ingestion/component/chunker/pdfcrop_cgo_test.go index ec91980020..613d0ac11c 100644 --- a/internal/ingestion/component/chunker/pdfcrop_cgo_test.go +++ b/internal/ingestion/component/chunker/pdfcrop_cgo_test.go @@ -71,7 +71,7 @@ func TestNeedsCrop(t *testing.T) { }{ {"image with positions", schema.ChunkDoc{CKType: "image", PDFPositions: jsonPositions(t, []float64{1, 10, 100, 10, 100})}, true}, {"table with positions", schema.ChunkDoc{CKType: "table", Positions: jsonPositions(t, []float64{1, 10, 100, 10, 100})}, true}, - {"text with positions", schema.ChunkDoc{CKType: "text", PDFPositions: jsonPositions(t, []float64{1, 10, 100, 10, 100})}, false}, + {"text with positions", schema.ChunkDoc{CKType: "text", PDFPositions: jsonPositions(t, []float64{1, 10, 100, 10, 100})}, true}, {"image without positions", schema.ChunkDoc{CKType: "image"}, false}, {"unknown type", schema.ChunkDoc{CKType: "equation", PDFPositions: jsonPositions(t, []float64{1, 10, 100, 10, 100})}, false}, } @@ -82,7 +82,7 @@ func TestNeedsCrop(t *testing.T) { } } -func TestCropImageChunks_CropsImageAndTable(t *testing.T) { +func TestCropImageChunks_CropsImageTableAndText(t *testing.T) { ctx := context.Background() // 1-based JSON position (1) must be rendered as 0-based page 0. eng := assertZeroPageEngine{} @@ -91,7 +91,7 @@ func TestCropImageChunks_CropsImageAndTable(t *testing.T) { chunks := []schema.ChunkDoc{ {CKType: "image", PDFPositions: pos}, {CKType: "table", PDFPositions: pos}, - {CKType: "text", PDFPositions: pos}, // skipped (not image/table) + {CKType: "text", PDFPositions: pos}, // restored preview (Chunker-1.3) {CKType: "image", Image: "data:image/png;base64,preexisting"}, // preserved } out := cropImageChunks(ctx, eng, chunks) @@ -101,8 +101,10 @@ func TestCropImageChunks_CropsImageAndTable(t *testing.T) { for i, ck := range out { switch ck.CKType { case "text": - if ck.Image != "" { - t.Errorf("chunk %d (text): image should stay empty, got %q", i, ck.Image) + // Chunker-1.3: text chunks with PDF positions get a rendered + // preview, mirroring Python restore_pdf_text_previews. + if !strings.HasPrefix(ck.Image, "data:image/png;base64,") { + t.Errorf("chunk %d (text): image = %q, want data:image/png;base64, prefix (preview restored)", i, ck.Image) } case "image": if ck.Image == "data:image/png;base64,preexisting" { @@ -119,6 +121,30 @@ func TestCropImageChunks_CropsImageAndTable(t *testing.T) { } } +// TestRestorePDFTextPreview covers Chunker-1.3 directly: a text chunk that +// carries PDF positions must receive a rendered preview image, while a text +// chunk without positions must be left untouched (no spurious preview). The +// img_id upload is owned by imageUploadDecorator (image_upload.go) and is +// not asserted here. +func TestRestorePDFTextPreview(t *testing.T) { + ctx := context.Background() + pos := jsonPositions(t, []float64{1, 10, 100, 10, 100}) + + withPos := schema.ChunkDoc{CKType: "text", PDFPositions: pos} + withoutPos := schema.ChunkDoc{CKType: "text", Text: "plain text, no coordinates"} + + out := cropImageChunks(ctx, mockCropEngine{}, []schema.ChunkDoc{withPos, withoutPos}) + if len(out) != 2 { + t.Fatalf("len(out) = %d, want 2", len(out)) + } + if !strings.HasPrefix(out[0].Image, "data:image/png;base64,") { + t.Errorf("chunk with positions: image = %q, want data:image/png;base64, prefix", out[0].Image) + } + if out[1].Image != "" { + t.Errorf("chunk without positions: image = %q, want empty", out[1].Image) + } +} + func TestCropImageChunks_NilEnginePassesThrough(t *testing.T) { ctx := context.Background() pos := jsonPositions(t, []float64{1, 10, 100, 10, 100}) diff --git a/internal/ingestion/component/chunker/title.go b/internal/ingestion/component/chunker/title.go index 8c06bf8b40..791a962dad 100644 --- a/internal/ingestion/component/chunker/title.go +++ b/internal/ingestion/component/chunker/title.go @@ -54,6 +54,7 @@ package chunker import ( "context" + "encoding/json" "fmt" "regexp" "strings" @@ -261,8 +262,9 @@ func matchLayoutLevel(text, layout string, fallbackLevel int) int { // resolveTitleLevels mirrors common.py:resolve_frequency_levels over the // full record stream. It is the "frequency" branch of -// common.py:resolve_title_levels (the outline branch is parity-gap -// territory per the SCOPE comment above). +// common.py:resolve_title_levels. The outline branch +// (common.py:resolve_outline_levels) is handled by resolveOutlineLevels and +// tried first in newLevelContext. // // For each record: // - a non-text record is pinned to BODY_LEVEL directly (python skips @@ -339,6 +341,140 @@ func resolveTitleLevels(records []lineRecord, p *titleChunkerParam) []int { return out } +// outlineEntry is one PDF bookmark/heading from the parser-supplied +// outline, mirroring Python extract_pdf_outlines' (text, level, page) +// tuple. The page is unused by title detection. +type outlineEntry struct { + title string + level int +} + +// outlineSimilarity mirrors common.py:_outline_similarity: the Jaccard +// overlap of character bigrams between two strings. It is rune-based so it +// matches Python's code-point indexing (str[i] is a Unicode character, not +// a byte). The right-hand bigram set is capped at min(len(left), len(right)-1) +// characters, exactly as the Python range() does. +func outlineSimilarity(left, right string) float64 { + lr := []rune(left) + rr := []rune(right) + leftPairs := make(map[string]struct{}, max(0, len(lr)-1)) + for i := 0; i+1 < len(lr); i++ { + leftPairs[string(lr[i])+string(lr[i+1])] = struct{}{} + } + n := len(lr) + if m := len(rr) - 1; m < n { + n = m + } + if n < 0 { + n = 0 + } + rightPairs := make(map[string]struct{}, n) + for i := 0; i < n; i++ { + rightPairs[string(rr[i])+string(rr[i+1])] = struct{}{} + } + denom := len(leftPairs) + if len(rightPairs) > denom { + denom = len(rightPairs) + } + if denom == 0 { + return 0 + } + inter := 0 + for k := range leftPairs { + if _, ok := rightPairs[k]; ok { + inter++ + } + } + return float64(inter) / float64(denom) +} + +// resolveOutlineLevels mirrors common.py:resolve_outline_levels. Each text +// record is matched against the outline by character-bigram similarity (>0.8 +// assigns level+1); unmatched records stay BODY_LEVEL. It returns ok=false +// when there is no outline, or when the outline is too sparse relative to the +// record count (len(outlines)/len(records) <= 0.03), in which case the +// caller falls back to frequency-based detection. mostLevel mirrors Python's +// max(1, max_outline_level). +func resolveOutlineLevels(records []lineRecord, outline []outlineEntry) (levels []int, mostLevel int, ok bool) { + if len(outline) == 0 || len(records) == 0 { + return nil, 0, false + } + if float64(len(outline))/float64(len(records)) <= 0.03 { + return nil, 0, false + } + maxLevel := 0 + for _, o := range outline { + if o.level > maxLevel { + maxLevel = o.level + } + } + levels = make([]int, len(records)) + for i, rec := range records { + if !rec.isText() { + levels[i] = bodyLevel + continue + } + matched := 0 + for _, o := range outline { + if outlineSimilarity(o.title, rec.text) > 0.8 { + matched = o.level + 1 + break + } + } + if matched == 0 { + levels[i] = bodyLevel + } else { + levels[i] = matched + } + } + return levels, max(1, maxLevel), true +} + +// outlineFromInputs reads the parser-supplied PDF outline from the upstream +// file metadata (file.outline, written by the ingestion PDF parser's +// outlinesToFileMeta) and normalizes it into the chunker's outlineEntry +// shape. Returns nil when no outline is present, so callers fall back to +// frequency-based title detection. Numbers are coerced from int/float64 +// because the runtime may hand the chunker a JSON-decoded payload. +func outlineFromInputs(inputs map[string]any) []outlineEntry { + file, _ := inputs["file"].(map[string]any) + if file == nil { + return nil + } + raw, _ := file["outline"].([]any) + if len(raw) == 0 { + return nil + } + out := make([]outlineEntry, 0, len(raw)) + for _, item := range raw { + m, ok := item.(map[string]any) + if !ok { + continue + } + title, _ := m["title"].(string) + if title == "" { + continue + } + out = append(out, outlineEntry{title: title, level: anyToInt(m["level"])}) + } + return out +} + +func anyToInt(v any) int { + switch t := v.(type) { + case int: + return t + case int64: + return int(t) + case float64: + return int(t) + case float32: + return int(t) + default: + return 0 + } +} + // bodyLevel is the sentinel python uses for non-heading lines. We use // the same large int (sys.maxsize - 1) for parity. Practically this // just needs to be "larger than any realistic heading level"; tests @@ -357,11 +493,12 @@ func lineRecordsFromText(text string) []lineRecord { continue } out = append(out, lineRecord{ - text: ln, - docType: "text", - imgID: nil, - layout: "", - pdfPos: nil, + text: ln, + docType: "text", + imgID: nil, + layout: "", + pdfPositions: nil, + positions: nil, }) } return out @@ -371,13 +508,14 @@ func lineRecordsFromText(text string) []lineRecord { // common.py:extract_line_records yields. Used by Group/Hierarchy // chunk-builders. type lineRecord struct { - text string - docType string - imgID *string - layout string - ckType string - pdfPos []map[string]any - parentMeta map[string]any + text string + docType string + imgID *string + layout string + ckType string + pdfPositions json.RawMessage + positions json.RawMessage + parentMeta map[string]any } func (r lineRecord) textOrEmpty() string { return r.text } @@ -421,7 +559,13 @@ type LevelContext struct { mostLevel int } -func newLevelContext(records []lineRecord, p *titleChunkerParam) LevelContext { +// newLevelContext resolves per-line heading levels, mirroring Python's +// resolve_title_levels: try the PDF outline branch first (when an outline is +// supplied and dense enough), otherwise fall back to frequency detection. +func newLevelContext(records []lineRecord, outline []outlineEntry, p *titleChunkerParam) LevelContext { + if levels, mostLevel, ok := resolveOutlineLevels(records, outline); ok { + return LevelContext{levels: levels, mostLevel: mostLevel} + } levels := resolveTitleLevels(records, p) // most_level is the most-frequent non-body heading level // (common.py:resolve_frequency_levels). Python computes this via diff --git a/internal/ingestion/component/chunker/title_test.go b/internal/ingestion/component/chunker/title_test.go index c8b39dd3c9..63f376d267 100644 --- a/internal/ingestion/component/chunker/title_test.go +++ b/internal/ingestion/component/chunker/title_test.go @@ -255,7 +255,7 @@ func TestNewLevelContext_MostLevelIsMode(t *testing.T) { {text: "## c", docType: "text"}, {text: "### d", docType: "text"}, } - lc := newLevelContext(records, p) + lc := newLevelContext(records, nil, p) // selectLevelGroup picks the single 3-pattern family; per-line // levels are [1,2,2,3], so the mode (most frequent heading level) // is level 2. @@ -371,3 +371,99 @@ func TestResolveTitleLevels_LayoutFallback(t *testing.T) { t.Errorf("plain body level = %d, want bodyLevel", levels[2]) } } + +// TestResolveOutlineLevels covers Chunker-1.5: a text line that matches a +// PDF outline entry by character-bigram similarity (>0.8) is assigned the +// outline level+1; lines that do not match stay BODY_LEVEL. The non-text +// record is pinned to BODY_LEVEL regardless. +func TestResolveOutlineLevels(t *testing.T) { + outline := []outlineEntry{ + {title: "第一章 概述", level: 0}, + {title: "第二章 方法", level: 1}, + } + records := []lineRecord{ + {text: "第一章 概述", docType: "text"}, + {text: "本章介绍背景。", docType: "text"}, // no outline match + {text: "figure caption", docType: "image"}, // non-text + } + levels, mostLevel, ok := resolveOutlineLevels(records, outline) + if !ok { + t.Fatalf("resolveOutlineLevels ok = false, want true") + } + if levels[0] != 1 { // outline level 0 + 1 + t.Errorf("matched line level = %d, want 1", levels[0]) + } + if levels[1] != bodyLevel { + t.Errorf("unmatched body level = %d, want bodyLevel", levels[1]) + } + if levels[2] != bodyLevel { + t.Errorf("non-text level = %d, want bodyLevel", levels[2]) + } + if mostLevel != 1 { // max(1, max outline level 1) + t.Errorf("mostLevel = %d, want 1", mostLevel) + } +} + +// TestResolveOutlineLevels_SparseGuard ensures a too-sparse outline (ratio +// <= 0.03) is rejected so detection falls back to frequency branch. +func TestResolveOutlineLevels_SparseGuard(t *testing.T) { + outline := []outlineEntry{{title: "唯一的章节标题", level: 0}} + // 100 body records: 1/100 = 0.01 <= 0.03 -> rejected. + records := make([]lineRecord, 100) + for i := range records { + records[i] = lineRecord{text: "body paragraph", docType: "text"} + } + if _, _, ok := resolveOutlineLevels(records, outline); ok { + t.Errorf("sparse outline ok = true, want false") + } + if _, _, ok := resolveOutlineLevels(records, nil); ok { + t.Errorf("nil outline ok = true, want false") + } +} + +// TestNewLevelContext_OutlineBranch pins Chunker-1.5 end-to-end: when an +// outline is supplied, newLevelContext prefers the outline branch over the +// regex/frequency branch. +func TestNewLevelContext_OutlineBranch(t *testing.T) { + p := &titleChunkerParam{TitleChunkerParam: schema.TitleChunkerParam{ + Method: "group", + Levels: [][]string{{`^# `}}, + }} + outline := []outlineEntry{{title: "前言", level: 0}} + records := []lineRecord{ + {text: "前言", docType: "text"}, // matches outline -> level 1 + {text: "普通正文段落。", docType: "text"}, // no outline -> BODY_LEVEL + } + lc := newLevelContext(records, outline, p) + if got := lc.Levels(); got[0] != 1 || got[1] != bodyLevel { + t.Errorf("outline levels = %v, want [1, bodyLevel]", got) + } +} + +// TestOutlineFromInputs verifies the parser-supplied file.outline is parsed +// into outlineEntry, tolerating float64-encoded levels (runtime JSON path). +func TestOutlineFromInputs(t *testing.T) { + inputs := map[string]any{ + "file": map[string]any{ + "outline": []any{ + map[string]any{"title": "第一章", "level": 0, "page_number": 1}, + map[string]any{"title": "第二章", "level": 1, "page_number": 3}, + }, + }, + } + out := outlineFromInputs(inputs) + if len(out) != 2 { + t.Fatalf("outline len = %d, want 2", len(out)) + } + if out[0].title != "第一章" || out[0].level != 0 { + t.Errorf("entry 0 = %+v, want {第一章 0}", out[0]) + } + if out[1].title != "第二章" || out[1].level != 1 { + t.Errorf("entry 1 = %+v, want {第二章 1}", out[1]) + } + + // No file / no outline -> nil (frequency fallback). + if got := outlineFromInputs(map[string]any{}); got != nil { + t.Errorf("empty inputs outline = %v, want nil", got) + } +} diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index 975f36df16..60ea9337ba 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -46,10 +46,10 @@ // Media-context attachment is per-item sequential; merge is // index-deterministic. // -// - No PDF/outline awareness (Python `restore_pdf_text_previews`). -// That depends on deepdoc/parser which is out of scope for this -// phase; the chunker accepts the parser-style structured JSON -// payload and runs the same logic against it. +// - PDF text previews (Python `restore_pdf_text_previews`) are +// generated on demand for text chunks that carry PDF positions: +// cropImageChunks crops the text region and writes a preview image, +// then imageUploadDecorator uploads it to img_id. See pdfcrop_cgo.go. package chunker import ( @@ -722,6 +722,12 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over prev.Text = prev.Text + "\n" + ck.Text prev.TKNums = intPtr(intValue(prev.TKNums) + tk) } + // Preserve PDF coordinates across the merge: extend the + // coordinate lists instead of dropping the incoming item's + // positions. Mirrors Python token_chunker.py:240 + // `merged[prev][PDF_POSITIONS_KEY].extend(...)` (diffs 2.5 / 2.3). + prev.PDFPositions = extendRawJSONArray(prev.PDFPositions, ck.PDFPositions) + prev.Positions = extendRawJSONArray(prev.Positions, ck.Positions) } perItem[idx] = merged } @@ -742,6 +748,14 @@ func cloneChunkDoc(in schema.ChunkDoc) schema.ChunkDoc { v := *in.PageNumber out.PageNumber = &v } + // Deep-copy the coordinate byte slices so the clone does not alias + // the source's backing array (diff 2.5 defensive fix). + if in.PDFPositions != nil { + out.PDFPositions = append(json.RawMessage(nil), in.PDFPositions...) + } + if in.Positions != nil { + out.Positions = append(json.RawMessage(nil), in.Positions...) + } if in.Extra != nil { out.Extra = make(map[string]json.RawMessage, len(in.Extra)) for k, v := range in.Extra { @@ -751,6 +765,33 @@ func cloneChunkDoc(in schema.ChunkDoc) schema.ChunkDoc { return out } +// extendRawJSONArray concatenates two JSON array payloads, mirroring +// Python's `merged[prev][KEY].extend(current[KEY])`. Either operand may be +// empty; the result is always a valid JSON array (or an empty raw message). +// It is used to accumulate PDF coordinate lists (`_pdf_positions`, +// `positions`) when text chunks are merged (diffs 2.5 / 2.3). +func extendRawJSONArray(a, b json.RawMessage) json.RawMessage { + if len(a) == 0 { + return b + } + if len(b) == 0 { + return a + } + var arrA, arrB []json.RawMessage + if err := json.Unmarshal(a, &arrA); err != nil { + return b + } + if err := json.Unmarshal(b, &arrB); err != nil { + return a + } + arrA = append(arrA, arrB...) + out, err := json.Marshal(arrA) + if err != nil { + return a + } + return out +} + func flatten(perItem [][]schema.ChunkDoc) []schema.ChunkDoc { var out []schema.ChunkDoc for _, cs := range perItem { diff --git a/internal/ingestion/component/chunker/token_pdfpos_test.go b/internal/ingestion/component/chunker/token_pdfpos_test.go new file mode 100644 index 0000000000..0c348cf308 --- /dev/null +++ b/internal/ingestion/component/chunker/token_pdfpos_test.go @@ -0,0 +1,119 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package chunker + +import ( + "encoding/json" + "strings" + "testing" + + "ragflow/internal/ingestion/component/schema" +) + +// TestMergeByTokenSizeFromJSON_ExtendsPDFPositions is the TDD test for +// migration diffs Chunker-2.5 / 2.3: when two JSON text items carrying +// `_pdf_positions` / `positions` are merged into one chunk, the merged +// chunk must extend (not drop) the coordinate lists — mirroring Python +// token_chunker.py:240 `merged[prev][PDF_POSITIONS_KEY].extend(...)`. +func TestMergeByTokenSizeFromJSON_ExtendsPDFPositions(t *testing.T) { + posA := json.RawMessage(`[[1,10,20,30,40]]`) + posB := json.RawMessage(`[[2,15,25,35,45]]`) + items := [][]schema.ChunkDoc{ + { + {Text: "alpha", DocType: "text", CKType: "text", TKNums: intPtr(5), PDFPositions: posA}, + {Text: "beta", DocType: "text", CKType: "text", TKNums: intPtr(5), PDFPositions: posB}, + }, + } + got := mergeByTokenSizeFromJSON(items, 128, 0) + merged := got[0] + if len(merged) != 1 { + t.Fatalf("want 1 merged chunk, got %d", len(merged)) + } + combined := string(merged[0].PDFPositions) + if !strings.Contains(combined, "1,10,20,30,40") { + t.Errorf("merged chunk lost first item _pdf_positions: %s", combined) + } + if !strings.Contains(combined, "2,15,25,35,45") { + t.Errorf("merged chunk dropped second item _pdf_positions (not extended): %s", combined) + } +} + +// TestMergeByTokenSizeFromJSON_ExtendsPositions covers the parallel +// `positions` field (diff 2.3). +func TestMergeByTokenSizeFromJSON_ExtendsPositions(t *testing.T) { + posA := json.RawMessage(`[[1,2,3]]`) + posB := json.RawMessage(`[[4,5,6]]`) + items := [][]schema.ChunkDoc{ + { + {Text: "a", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posA}, + {Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB}, + }, + } + got := mergeByTokenSizeFromJSON(items, 128, 0) + combined := string(got[0][0].Positions) + if !strings.Contains(combined, "1,2,3") || !strings.Contains(combined, "4,5,6") { + t.Errorf("merged chunk dropped/omitted `positions`: %s", combined) + } +} + +// TestCloneChunkDoc_DeepCopiesPDFPositions ensures cloneChunkDoc does not +// alias the underlying _pdf_positions / positions byte slices (diff 2.5 +// defensive fix). +func TestCloneChunkDoc_DeepCopiesPDFPositions(t *testing.T) { + pos := json.RawMessage(`[[1,2,3,4,5]]`) + orig := schema.ChunkDoc{Text: "x", PDFPositions: pos, Positions: pos} + cp := cloneChunkDoc(orig) + // Mutate the source's backing array after the clone. + pos[0] = '9' + if string(cp.PDFPositions) != "[[1,2,3,4,5]]" { + t.Errorf("clone shares _pdf_positions backing array: %s", string(cp.PDFPositions)) + } + if string(cp.Positions) != "[[1,2,3,4,5]]" { + t.Errorf("clone shares positions backing array: %s", string(cp.Positions)) + } +} + +// TestMergeByTokenSizeFromJSON_PositionsDecodeToMatrix verifies the +// chunker-side contract for diff 1.4: preserved `positions` must decode +// (via ChunkDoc.ToMap → decodeStructuredValue) to a [][]float64 matrix so +// the downstream task-layer processChunkPositions → AddPositions can +// convert it to page_num_int / top_int / position_int. The coordinate +// conversion itself lives in internal/ingestion/task (processChunkPositions), +// not in the chunker. +func TestMergeByTokenSizeFromJSON_PositionsDecodeToMatrix(t *testing.T) { + posA := json.RawMessage(`[[1,2,3,4,5]]`) + posB := json.RawMessage(`[[6,7,8,9,10]]`) + items := [][]schema.ChunkDoc{ + { + {Text: "a", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posA}, + {Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB}, + }, + } + got := mergeByTokenSizeFromJSON(items, 128, 0) + m := got[0][0].ToMap() + raw, ok := m["positions"] + if !ok { + t.Fatal("positions missing from ToMap output") + } + matrix, ok := raw.([][]float64) + if !ok { + t.Fatalf("positions decoded to %T, want [][]float64", raw) + } + if len(matrix) != 2 { + t.Fatalf("positions matrix has %d groups, want 2 (both merged items)", len(matrix)) + } +} diff --git a/internal/ingestion/component/tokenizer_unit_test.go b/internal/ingestion/component/tokenizer_unit_test.go index 403799740f..4710044a08 100644 --- a/internal/ingestion/component/tokenizer_unit_test.go +++ b/internal/ingestion/component/tokenizer_unit_test.go @@ -22,6 +22,7 @@ package component import ( "context" + "encoding/json" "strings" "sync/atomic" "testing" @@ -322,3 +323,32 @@ func TestValidateTokenizerOutputs_SymbolOnlyContentLtksIsEmptyFails(t *testing.T t.Fatalf("err = %v, want missing full_text tokens", err) } } + +// TestChunkDocsToMaps_PreservesPDFPositions is the pool-free unit test for +// Tokenizer-(T)1: the tokenizer emits chunks via schema.ChunkDocsToMaps +// (ChunkDoc.ToMap), which must carry the raw `positions` / `_pdf_positions` +// through untouched so the downstream executor stage +// (internal/ingestion/task processChunkPositions → AddPositions) can convert +// them into position_int / page_num_int / top_int exactly once. This does NOT +// require the C++ analyzer pool, so it runs under plain `go test`. +func TestChunkDocsToMaps_PreservesPDFPositions(t *testing.T) { + pos := json.RawMessage(`[[1,10,20,30,40],[2,15,25,35,45]]`) + chunks := []schema.ChunkDoc{ + {Text: "PDF paragraph", DocType: "text", CKType: "text", + Positions: pos, PDFPositions: pos}, + } + maps := schema.ChunkDocsToMaps(chunks) + + got, ok := maps[0]["positions"].([][]float64) + if !ok || len(got) != 2 { + t.Fatalf("positions not preserved through tokenizer output mapping: %#v", maps[0]["positions"]) + } + if _, ok := maps[0]["_pdf_positions"].([][]float64); !ok { + t.Errorf("_pdf_positions not preserved through tokenizer output mapping: %#v", maps[0]["_pdf_positions"]) + } + // Sanity: page numbers are still raw 1-indexed, i.e. not yet converted + // to page_num_int (the executor owns that step). + if int(got[0][0]) != 1 { + t.Errorf("positions page already converted; want raw 1-indexed page 1, got %v", got[0][0]) + } +} diff --git a/internal/ingestion/task/pipeline_real_integration_test.go b/internal/ingestion/task/pipeline_real_integration_test.go index 5b8178aad3..1fd64e3952 100644 --- a/internal/ingestion/task/pipeline_real_integration_test.go +++ b/internal/ingestion/task/pipeline_real_integration_test.go @@ -15,19 +15,17 @@ import ( "ragflow/internal/common" "ragflow/internal/dao" - "ragflow/internal/engine" - enginetypes "ragflow/internal/engine/types" "ragflow/internal/entity" _ "ragflow/internal/ingestion/component" - componentpkg "ragflow/internal/ingestion/component" _ "ragflow/internal/ingestion/component/chunker" pipelinepkg "ragflow/internal/ingestion/pipeline" "ragflow/internal/server" "ragflow/internal/storage" "ragflow/internal/tokenizer" + "github.com/glebarez/sqlite" + "go.uber.org/zap" - "gorm.io/driver/mysql" "gorm.io/gorm" gormlogger "gorm.io/gorm/logger" ) @@ -35,30 +33,18 @@ import ( func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) { requireTokenizerPool(t) - cfg := mustLoadTaskRealIntegrationConfig(t) - realDB := mustOpenTaskRealMySQL(t, cfg) - if err := realDB.AutoMigrate( - &entity.Tenant{}, - &entity.Knowledgebase{}, - &entity.Document{}, - &entity.File{}, - &entity.File2Document{}, - &entity.UserCanvas{}, - ); err != nil { - t.Fatalf("auto-migrate real mysql tables: %v", err) - } - - realStorage, err := storage.NewMinioStorage(cfg.StorageEngine.Minio) - if err != nil { - t.Fatalf("connect real minio: %v", err) - } - + mustLoadTaskTestConfig(t) origDB := dao.DB - origStorage := storage.GetStorageFactory().GetStorage() + realDB := mustOpenTaskTestDB(t) dao.DB = realDB - storage.GetStorageFactory().SetStorage(realStorage) t.Cleanup(func() { dao.DB = origDB + }) + + realStorage := storage.NewMemoryStorage() + origStorage := storage.GetStorageFactory().GetStorage() + storage.GetStorageFactory().SetStorage(realStorage) + t.Cleanup(func() { storage.GetStorageFactory().SetStorage(origStorage) }) @@ -148,46 +134,25 @@ func TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline(t *testing.T) { } } -func TestPipelineExecutor_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *testing.T) { +func TestPipelineExecutor_Run_RealPDF_ProducesIndexedChunks(t *testing.T) { requireTokenizerPool(t) - cfg := mustLoadTaskRealIntegrationConfig(t) - realDB := mustOpenTaskRealMySQL(t, cfg) - if err := realDB.AutoMigrate( - &entity.Tenant{}, - &entity.Knowledgebase{}, - &entity.Document{}, - &entity.File{}, - &entity.File2Document{}, - &entity.UserCanvas{}, - ); err != nil { - t.Fatalf("auto-migrate real mysql tables: %v", err) - } - - realStorage, err := storage.NewMinioStorage(cfg.StorageEngine.Minio) - if err != nil { - t.Fatalf("connect real minio: %v", err) - } - if err := engine.Init(&cfg.DocEngine); err != nil { - t.Fatalf("init real doc engine: %v", err) - } - if engine.Get() == nil { - t.Fatal("doc engine is nil after init") - } - if engine.GetEngineType() != engine.EngineElasticsearch { - t.Fatalf("doc engine type = %s, want %s", engine.GetEngineType(), engine.EngineElasticsearch) - } - + // Loads service config (server.Init side effect) without requiring any + // external MySQL/MinIO/ES. The pipeline runs against an in-memory sqlite + // DB and an in-memory storage backend; chunks are captured via WithInsertFunc. + mustLoadTaskTestConfig(t) origDB := dao.DB - origStorage := storage.GetStorageFactory().GetStorage() - origDocResolver := componentpkg.ResolveDocumentStorageOverride + realDB := mustOpenTaskTestDB(t) dao.DB = realDB - storage.GetStorageFactory().SetStorage(realStorage) - componentpkg.ResolveDocumentStorageOverride = nil t.Cleanup(func() { dao.DB = origDB + }) + + realStorage := storage.NewMemoryStorage() + origStorage := storage.GetStorageFactory().GetStorage() + storage.GetStorageFactory().SetStorage(realStorage) + t.Cleanup(func() { storage.GetStorageFactory().SetStorage(origStorage) - componentpkg.ResolveDocumentStorageOverride = origDocResolver }) templatePath := filepath.Join(taskRepoRoot(t), "internal", "ingestion", "pipeline", "template", "ingestion_pipeline_general.json") @@ -217,7 +182,6 @@ func TestPipelineExecutor_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *tes bucket := taskS3SafeBucketName(kbID) docName := "01_english_simple.pdf" objectPath := fmt.Sprintf("integration/task/%s/%s", docID, docName) - baseName := fmt.Sprintf("ragflow_%s", tenantID) mustSeedTaskRealPipelineDocumentBytes(t, realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath, docName, ".pdf", "pdf", pdfBytes) if err := realDB.Model(&entity.Document{}).Where("id = ?", docID).Update("pipeline_id", canvasID).Error; err != nil { @@ -233,14 +197,13 @@ func TestPipelineExecutor_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *tes t.Fatalf("create user canvas: %v", err) } t.Cleanup(func() { - _ = engine.Get().DropChunkStore(context.Background(), baseName, kbID) _ = realDB.Where("id = ?", canvasID).Delete(&entity.UserCanvas{}).Error cleanupTaskRealPipelineDocument(realDB, realStorage, tenantID, kbID, docID, fileID, bucket, objectPath) }) taskCtx := &TaskContext{ IngestionTask: &entity.IngestionTask{ - ID: "task-real-pdf-es-1", + ID: "task-real-pdf-1", DocumentID: docID, DatasetID: kbID, }, @@ -258,68 +221,66 @@ func TestPipelineExecutor_Run_RealPDF_WritesAndReadsBackFromElasticsearch(t *tes Tenant: entity.Tenant{ID: tenantID}, } - svc := mustNewPipelineExecutor(t, taskCtx, canvasID, 0) + var inserted [][]map[string]any + svc := mustNewPipelineExecutor(t, taskCtx, canvasID, 0). + WithInsertFunc(func(ctx context.Context, chunks []map[string]any, baseName, datasetID string) ([]string, error) { + inserted = append(inserted, deepCopyTaskChunks(chunks)) + return nil, nil + }). + WithLogCreateFunc(func(log *entity.PipelineOperationLog) error { return nil }) if _, err := svc.Execute(context.Background()); err != nil { t.Fatalf("Run: %v", err) } - result, err := engine.Get().Search(context.Background(), &enginetypes.SearchRequest{ - IndexNames: []string{baseName}, - KbIDs: []string{kbID}, - Limit: 20, - }) - if err != nil { - t.Fatalf("search indexed chunks: %v", err) + if len(inserted) == 0 { + t.Fatal("no chunks inserted") } - if result == nil { - t.Fatal("search result is nil") + var chunks []map[string]any + for _, batch := range inserted { + chunks = append(chunks, batch...) } - if len(result.Chunks) == 0 { - t.Fatal("expected indexed chunks in Elasticsearch, got 0") + if len(chunks) == 0 { + t.Fatal("inserted 0 chunks") } - for i, chunk := range result.Chunks { + sawImage := false + for i, chunk := range chunks { if got := chunk["doc_id"]; got != docID { - t.Fatalf("result chunk[%d].doc_id = %v, want %q", i, got, docID) + t.Fatalf("chunk[%d].doc_id = %v, want %q", i, got, docID) } - if got := chunk["kb_id"]; got != kbID { - t.Fatalf("result chunk[%d].kb_id = %v, want %q", i, got, kbID) + if !taskChunkFieldEqualsStr(chunk["kb_id"], kbID) { + t.Fatalf("chunk[%d].kb_id = %v, want %q", i, chunk["kb_id"], kbID) } if got := chunk["docnm_kwd"]; got != docName { - t.Fatalf("result chunk[%d].docnm_kwd = %v, want %q", i, got, docName) + t.Fatalf("chunk[%d].docnm_kwd = %v, want %q", i, got, docName) } if got := chunk["content_with_weight"]; got == nil || got == "" { - t.Fatalf("result chunk[%d].content_with_weight = %v, want non-empty", i, got) + t.Fatalf("chunk[%d].content_with_weight = %v, want non-empty", i, got) } + if img, ok := chunk["img_id"]; ok && img != nil && img != "" { + sawImage = true + } + } + if !sawImage { + t.Fatal("expected at least one chunk with img_id (pdf preview image uploaded to storage)") } } func TestRunPipeline_RealPipelineOutput_ProducesIndexFields(t *testing.T) { requireTokenizerPool(t) - cfg := mustLoadTaskRealIntegrationConfig(t) - realDB := mustOpenTaskRealMySQL(t, cfg) - if err := realDB.AutoMigrate( - &entity.Tenant{}, - &entity.Knowledgebase{}, - &entity.Document{}, - &entity.File{}, - &entity.File2Document{}, - ); err != nil { - t.Fatalf("auto-migrate real mysql tables: %v", err) - } - - realStorage, err := storage.NewMinioStorage(cfg.StorageEngine.Minio) - if err != nil { - t.Fatalf("connect real minio: %v", err) - } - + mustLoadTaskTestConfig(t) origDB := dao.DB - origStorage := storage.GetStorageFactory().GetStorage() + realDB := mustOpenTaskTestDB(t) dao.DB = realDB - storage.GetStorageFactory().SetStorage(realStorage) t.Cleanup(func() { dao.DB = origDB + }) + + realStorage := storage.NewMemoryStorage() + origStorage := storage.GetStorageFactory().GetStorage() + storage.GetStorageFactory().SetStorage(realStorage) + t.Cleanup(func() { storage.GetStorageFactory().SetStorage(origStorage) }) @@ -426,7 +387,7 @@ func taskRepoRoot(t *testing.T) string { return filepath.Clean(filepath.Join(wd, "..", "..", "..")) } -func mustLoadTaskRealIntegrationConfig(t *testing.T) *server.Config { +func mustLoadTaskTestConfig(t *testing.T) *server.Config { t.Helper() if err := common.Init("info", common.FileOutput{}, ""); err != nil { t.Fatalf("init common logger: %v", err) @@ -437,28 +398,45 @@ func mustLoadTaskRealIntegrationConfig(t *testing.T) *server.Config { t.Fatalf("init service config from %s: %v", configPath, err) } cfg := server.GetConfig() - if cfg == nil || cfg.Database.Host == "" || cfg.StorageEngine.Minio == nil || cfg.StorageEngine.Minio.Host == "" { - t.Fatal("real integration config is incomplete") + if cfg == nil { + t.Fatal("task test config is nil after server.Init") } return cfg } -func mustOpenTaskRealMySQL(t *testing.T, cfg *server.Config) *gorm.DB { +// mustOpenTaskTestDB opens an isolated on-disk sqlite database and migrates the +// tables the real-pipeline contract tests seed and read. It does not connect to +// any external MySQL; the temporary file is removed on test cleanup. +func mustOpenTaskTestDB(t *testing.T) *gorm.DB { t.Helper() - dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=Local", - cfg.Database.Username, - cfg.Database.Password, - cfg.Database.Host, - cfg.Database.Port, - cfg.Database.Database, - cfg.Database.Charset, - ) - db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{ + f, err := os.CreateTemp("", "task-test-*.db") + if err != nil { + t.Fatalf("create temp sqlite db: %v", err) + } + _ = f.Close() + db, err := gorm.Open(sqlite.Open(f.Name()), &gorm.Config{ Logger: gormlogger.Default.LogMode(gormlogger.Silent), }) if err != nil { - t.Fatalf("connect real mysql: %v", err) + t.Fatalf("open sqlite db: %v", err) } + if err := db.AutoMigrate( + &entity.Tenant{}, + &entity.Knowledgebase{}, + &entity.Document{}, + &entity.File{}, + &entity.File2Document{}, + &entity.UserCanvas{}, + &entity.PipelineOperationLog{}, + ); err != nil { + t.Fatalf("auto-migrate sqlite tables: %v", err) + } + if sqlDB, err := db.DB(); err != nil { + t.Fatalf("get sql.DB from gorm: %v", err) + } else { + sqlDB.SetMaxOpenConns(1) + } + t.Cleanup(func() { _ = os.Remove(f.Name()) }) return db } @@ -731,3 +709,18 @@ func taskS3SafeBucketName(s string) string { s = strings.ReplaceAll(s, "_", "-") return s } + +// taskChunkFieldEqualsStr compares a chunk field to a plain string, tolerating +// the slice form used internally (e.g. kb_id is []string{kbID} and survives a +// JSON round-trip as []any{kbID}). +func taskChunkFieldEqualsStr(v any, want string) bool { + switch val := v.(type) { + case string: + return val == want + case []string: + return len(val) == 1 && val[0] == want + case []any: + return len(val) == 1 && fmt.Sprint(val[0]) == want + } + return false +}