From 554925b5834050eff9a4b97f50ad054ef114c5ef Mon Sep 17 00:00:00 2001 From: Jack Date: Fri, 24 Jul 2026 21:06:38 +0800 Subject: [PATCH] Fix(go): align ingestion pipeline with Python (parser/media dispatch + PDF coordinate chain + Chunker) (#17349) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Aligns the Go ingestion pipeline with the Python implementation, closing several behavioral gaps found during the Python→Go migration (tracked in `docs/migration_python_go_diff.md`). Covers parser/media dispatch alignment, the PDF coordinate-chain (preview images, outline→title, chunk coordinate finalization), and the Chunker Token/QA batches below. Commits are grouped as follows. ### 1. Fix parser params (c524f450e) Fixes parser/media wiring and several dispatch gaps: - **docx/pdf vision dispatch**: correct parameter handling and VLM invocation. - **markdown vision (diff 2.5)**: also enhance items whose `doc_type_kwd` is `table`, not only `image` (parser/utils.py:181). - **media audio (diff 2.11)**: when `output_format` is `json`, carry the ASR transcription as a JSON item instead of only the `Text` field (the Invoke switch had no `json` branch and dropped it). - **email (diff 2.2)**: default `output_format` is `json` (parser.py:212), not `text`. - **tokenizer**: handle empty/whitespace-only names; trim before embedding. - **extractor**: tag-matching parameter wiring. - **split**: keyword-split regex now covers CJK/English separators. - **parser.go**: parser-param plumbing. ### 2. fix parser gap (373537da1) Image dispatch now mirrors `rag/app/picture.py:chunk()`: - Always OCR the image (PaddleOCR or local ONNX). - When OCR text is short, also call VLM (`describe`) and combine `OCR + VLM` text. - Emits a **structured JSON item** carrying the image data-URI and `doc_type_kwd:"image"`, instead of a bare `Text` string. This fixes the payload being rejected downstream by OneChunker/TokenChunker (JSON=nil). ### 3. PDF coordinate-chain fixes (55367a820, 727f8167c) Closes three items from the migration tracker in the chunker/tokenizer/task layer: - **(Chunker-1.3) `restore_pdf_text_previews`** — `needsCrop` now also returns true for `text` chunks that carry PDF positions (`pdfcrop_cgo.go`), so text blocks get a rendered preview image uploaded to storage via `imageUploadDecorator`/`ChunkImageUploader`, matching Python `restore_pdf_text_previews` + `image2id`. - **(Chunker-1.5) PDF outline → title levels** — `title.go` adds `outlineSimilarity` (rune-bigram Jaccard, mirroring `common.py:_outline_similarity`), `resolveOutlineLevels` (matches text lines to outline entries at similarity > 0.8, with a sparse guard `len(outline)/len(records) <= 0.03`), and `outlineFromInputs` (reads `file.outline`). Wired into `newLevelContext` in both `group.go` and `hierarchy.go`; falls back to the title-shape heuristic when no outline is present. - **(Tokenizer-(T)1) `finalize_pdf_chunk`** — the coordinate → `position_int`/`page_num_int`/`top_int` conversion is owned by the task layer (`processChunkPositions`→`AddPositions`), which runs *after* the tokenizer and consumes the tokenizer-owned fields. The tokenizer only preserves the raw `positions`/`_pdf_positions` (no duplicate conversion), pinned by `TestChunkDocsToMaps_PreservesPDFPositions`. ### 4. Integration test made environment-free (`internal/ingestion/task/pipeline_real_integration_test.go`) - Removed the `//go:build integration` tag so the contract tests run under the default `build.sh --test` (which does not pass `-tags integration`). - External dependencies replaced with in-memory substitutes so no MySQL/MinIO/ES is required: - MySQL → on-disk sqlite (`glebarez/sqlite`) with the needed tables auto-migrated. - MinIO → `storage.NewMemoryStorage()`. - Elasticsearch → chunks captured via `WithInsertFunc` instead of `engine.InsertChunks`/`Search`. - `requireTokenizerPool` still skips gracefully when the native tokenizer pool is unavailable; `WithLogCreateFunc(noop)` avoids depending on the operation-log table. - Added `taskChunkFieldEqualsStr` to tolerate `kb_id` being a `[]string`/`[]any` in the raw chunk payload (the search engine flattens it to a string on read). ### 5. TokenChunker alignment — Batch 1 (`internal/ingestion/component/chunker/token.go`) Closes four Chunker items from the migration tracker: - **(Chunker-2.1) sentence delimiter** — the boundary regex now also breaks on ASCII `!`/`?`. Extracted to a package-level `var sentenceDelimiter` and used in `mergeByTokenSize`, matching Python's full delimiter set. - **(Chunker-2.2) overlap tag leakage** — when a new chunk starts, its overlap prefix is taken from the previous chunk *after* `removeTag`, in both the text path (`mergeByTokenSize`) and the JSON path (`mergeByTokenSizeFromJSON`). Parser tags (`@@…##`) no longer leak into the overlap region (mirrors `nlp/__init__.py:1181`). - **(Chunker-2.11) empty-text merge** — merging a non-empty chunk into an empty previous chunk now assigns the text directly instead of being skipped (`mergeByTokenSizeFromJSON`), mirroring `token_chunker.py:236-239`. - **(Chunker-2.4) overlap token counting** — `takeFromEnd`/`takeFromStart` now count tokens exactly via `tokenizeStr` instead of the 4-bytes/token heuristic, fixing over-counting for CJK text. ### 6. QA Chunker alignment — Batch 2 (`internal/ingestion/component/chunker/qa.go` + `schema`) Closes three Chunker items from the migration tracker: - **(Chunker-2.13) default language** — an empty `lang` now defaults to Chinese prefixes (`问题：`/`回答：`) instead of English, matching `qa.py:299`. - **(Chunker-2.12) `rmQAPrefix` regex** — the separator is changed to `[\t:： ]+` (one-or-more), matching `qa.py:241`, so multiple separators (e.g. `Q:: answer`) are fully stripped. - **(Chunker-1.8 QA) missing chunk fields** — QA chunks now preserve: - `top_int` — the source row/record index, threaded through the tab/csv/markdown extractors (mirrors `qa.py` `beAdoc(..., row_num=i)`); - `image` + `doc_type_kwd:"image"`; - `_pdf_positions` / `positions` carried from the upstream JSON item. `schema.ChunkDoc` gains a `TopInt []int` field (serialized as `top_int`, registered in `UnmarshalJSON`). Note: the Tag/Table/Presentation/One chunker field gaps under 1.8 remain pending. ## Test plan - Added/updated unit tests: `pdfcrop_cgo_test.go` (`TestNeedsCrop`, `TestRestorePDFTextPreview`), `title_test.go` (`TestResolveOutlineLevels`, `TestResolveOutlineLevels_SparseGuard`, `TestNewLevelContext_OutlineBranch`, `TestOutlineFromInputs`), `tokenizer_unit_test.go` (`TestChunkDocsToMaps_PreservesPDFPositions`), `token_pdfpos_test.go`. - **Batch 1** — `token_batch1_test.go`: `TestSentenceDelimiterMatchesBangAndQuestion`, `TestMergeByTokenSizeFromJSON_OverlapStripsTags`, `TestMergeByTokenSizeFromJSON_EmptyPrevKeepsChunk`, `TestTakeFromEndRespectsTokenCount`, `TestTakeFromStartRespectsTokenCount`. - **Batch 2** — `qa_batch2_test.go`: `TestQAChunker_DefaultLangIsChinese`, `TestRmQAPrefixStripsMultipleSeparators`, `TestQAChunker_SetsTopInt`, `TestQAChunker_CarriesImageAndPositions`. Existing `qa_test.go` expectations were updated to the corrected language default / separator behavior. - `pipeline_real_integration_test.go` (`TestPipelineExecutor_Run_RealCanvasDSL_UsesGeneralPipeline`, `TestPipelineExecutor_Run_RealPDF_ProducesIndexedChunks`, `TestRunPipeline_RealPipelineOutput_ProducesIndexFields`) now runs without any external service. - `bash build.sh --test ./internal/ingestion/...` passes. - No files deleted. --- internal/ingestion/component/chunker/group.go | 96 ++++++- .../ingestion/component/chunker/group_test.go | 54 ++++ .../ingestion/component/chunker/hierarchy.go | 2 +- .../component/chunker/pdfcrop_cgo.go | 11 +- .../component/chunker/pdfcrop_cgo_test.go | 36 ++- internal/ingestion/component/chunker/qa.go | 64 ++++- .../component/chunker/qa_batch2_test.go | 147 ++++++++++ .../ingestion/component/chunker/qa_test.go | 18 +- internal/ingestion/component/chunker/title.go | 174 ++++++++++-- .../ingestion/component/chunker/title_test.go | 98 ++++++- internal/ingestion/component/chunker/token.go | 115 ++++++-- .../component/chunker/token_batch1_test.go | 128 +++++++++ .../component/chunker/token_pdfpos_test.go | 119 +++++++++ .../component/docx_vision_dispatch.go | 167 +++++------- .../component/docx_vision_dispatch_test.go | 181 +++++++++++++ internal/ingestion/component/extractor.go | 7 + .../ingestion/component/extractor_test.go | 42 +++ .../component/markdown_vision_dispatch.go | 5 +- .../ingestion/component/media_dispatch.go | 70 +++-- .../component/media_dispatch_test.go | 250 +++++++++++++++++- internal/ingestion/component/parser.go | 11 +- .../component/pdf_vision_dispatch.go | 25 +- .../component/pdf_vision_dispatch_test.go | 86 ++++++ .../ingestion/component/schema/chunk_types.go | 3 +- internal/ingestion/component/tokenizer.go | 6 +- .../ingestion/component/tokenizer_test.go | 26 ++ .../component/tokenizer_unit_test.go | 30 +++ .../template/ingestion_pipeline_picture.json | 4 +- .../task/pipeline_executor_defaults_test.go | 2 +- .../task/pipeline_real_integration_test.go | 213 +++++++-------- internal/utility/split.go | 11 +- internal/utility/split_test.go | 27 +- 32 files changed, 1895 insertions(+), 333 deletions(-) create mode 100644 internal/ingestion/component/chunker/qa_batch2_test.go create mode 100644 internal/ingestion/component/chunker/token_batch1_test.go create mode 100644 internal/ingestion/component/chunker/token_pdfpos_test.go create mode 100644 internal/ingestion/component/docx_vision_dispatch_test.go create mode 100644 internal/ingestion/component/pdf_vision_dispatch_test.go 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 5d964e5484..9a72b77a6d 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/qa.go b/internal/ingestion/component/chunker/qa.go index 4f42377c0f..4dce3f7573 100644 --- a/internal/ingestion/component/chunker/qa.go +++ b/internal/ingestion/component/chunker/qa.go @@ -29,6 +29,7 @@ package chunker import ( "context" "encoding/csv" + "encoding/json" "fmt" "html" "regexp" @@ -98,7 +99,9 @@ func (c *QAChunkerComponent) invoke(_ context.Context, inputs map[string]any) (m } qPrefix, aPrefix := "问题：", "回答：" - eng := strings.EqualFold(c.param.Lang, "english") || c.param.Lang == "" + // Python qa.py defaults to Chinese when no language is supplied; only + // an explicit "english" switches to English prefixes (diff Chunker-2.13). + eng := strings.EqualFold(c.param.Lang, "english") if eng { qPrefix, aPrefix = "Question: ", "Answer: " } @@ -133,6 +136,21 @@ func (c *QAChunkerComponent) invoke(_ context.Context, inputs map[string]any) (m ContentLtks: contentLTKS, ContentSmLtks: contentSMLTKS, } + // Restore metadata lost before diff Chunker-1.8: top_int (row + // index), image id + coordinates carried from the source item. + if pair.RowNum >= 0 { + chunk.TopInt = []int{pair.RowNum} + } + if pair.Image != "" { + chunk.Image = pair.Image + chunk.DocType = "image" + } + if len(pair.PDFPositions) > 0 { + chunk.PDFPositions = pair.PDFPositions + } + if len(pair.Positions) > 0 { + chunk.Positions = pair.Positions + } chunks = append(chunks, chunk) } @@ -148,9 +166,20 @@ func renderMarkdown(s string) string { type qaPair struct { Question string Answer string + // RowNum is the 0-based source line/record index, mapped to Python's + // top_int (qa.py beAdoc(..., row_num=i)). -1 means unset. + RowNum int + // Image and positions are carried from the upstream item so the QA + // chunk preserves metadata that Python sets via beAdocPdf/beAdocDocx + // (diff Chunker-1.8). + Image string + PDFPositions json.RawMessage + Positions json.RawMessage } -var rmQAPrefixRe = regexp.MustCompile(`(?i)^(问题|答案|回答|user|assistant|Q|A|Question|Answer|问|答)[ \t]*(?:[:：]|\t)[ \t]*`) +// rmQAPrefixRe mirrors Python qa.py:241 `[\t:： ]+` — one-or-more separator +// chars, so "Q:: answer" is fully stripped (diff Chunker-2.12). +var rmQAPrefixRe = regexp.MustCompile(`(?i)^(问题|答案|回答|user|assistant|Q|A|Question|Answer|问|答)[\t:： ]+`) func rmQAPrefix(txt string) string { return strings.TrimSpace(rmQAPrefixRe.ReplaceAllString(txt, "")) @@ -209,18 +238,19 @@ func extractQAMarkdown(md string) []qaPair { var questionStack []string var levelStack []int var answer []string + curRow := -1 codeBlock := false flushAnswer := func() { joined := strings.TrimSpace(strings.Join(answer, "\n")) if joined != "" && len(questionStack) > 0 { sumQ := strings.Join(questionStack, "\n") - pairs = append(pairs, qaPair{Question: sumQ, Answer: joined}) + pairs = append(pairs, qaPair{Question: sumQ, Answer: joined, RowNum: curRow}) } answer = nil } - for _, line := range lines { + for i, line := range lines { trimmed := strings.TrimSpace(line) if strings.HasPrefix(trimmed, "```") { codeBlock = !codeBlock @@ -239,6 +269,7 @@ func extractQAMarkdown(md string) []qaPair { flushAnswer() question := strings.TrimSpace(line[level:]) + curRow = i for len(levelStack) > 0 && level <= levelStack[len(levelStack)-1] { questionStack = questionStack[:len(questionStack)-1] @@ -273,8 +304,9 @@ func extractQAText(text string) []qaPair { func extractQATextTab(lines []string) []qaPair { var pairs []qaPair var question, answer string + var row int - for _, line := range lines { + for i, line := range lines { if strings.TrimSpace(line) == "" { continue } @@ -286,13 +318,14 @@ func extractQATextTab(lines []string) []qaPair { continue } if question != "" && answer != "" { - pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)}) + pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer), RowNum: row}) } question = parts[0] answer = parts[1] + row = i } if question != "" { - pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)}) + pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer), RowNum: row}) } return pairs } @@ -321,13 +354,16 @@ func extractQATextCSV(text string, lines []string) []qaPair { var pairs []qaPair var question, answer string + var row int prevLine := 0 + recIdx := -1 for { record, err := r.Read() if err != nil { break } + recIdx++ // Map InputOffset back to the physical lines consumed. endOff := int(r.InputOffset()) @@ -346,13 +382,14 @@ func extractQATextCSV(text string, lines []string) []qaPair { continue } if question != "" && answer != "" { - pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)}) + pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer), RowNum: row}) } question = record[0] answer = record[1] + row = recIdx } if question != "" { - pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)}) + pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer), RowNum: row}) } return pairs } @@ -385,7 +422,14 @@ func extractQAJSON(items []schema.ChunkDoc) []qaPair { continue } tmp := extractQAText(txt) - pairs = append(pairs, tmp...) + // Preserve the source item's image id and coordinates on each + // extracted pair (diff Chunker-1.8). + for _, p := range tmp { + p.Image = item.Image + p.PDFPositions = item.PDFPositions + p.Positions = item.Positions + pairs = append(pairs, p) + } } return pairs } diff --git a/internal/ingestion/component/chunker/qa_batch2_test.go b/internal/ingestion/component/chunker/qa_batch2_test.go new file mode 100644 index 0000000000..6ad8329d5f --- /dev/null +++ b/internal/ingestion/component/chunker/qa_batch2_test.go @@ -0,0 +1,147 @@ +// +// 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 ( + "context" + "strings" + "testing" +) + +// qaInvoke is a small helper that runs the QA chunker on a upstream-style +// input map and returns the produced chunks as generic maps. +func qaInvoke(t *testing.T, inputs map[string]any) []map[string]any { + t.Helper() + c, err := NewQAChunker(map[string]any{}) + if err != nil { + t.Fatalf("NewQAChunker: %v", err) + } + out, err := c.Invoke(context.Background(), inputs) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + chunks, ok := out["chunks"].([]map[string]any) + if !ok { + t.Fatalf("Invoke did not return []map chunks: %#v", out["chunks"]) + } + return chunks +} + +// TestQAChunker_DefaultLangIsChinese exercises migration diff Chunker-2.13: +// when no language is supplied, Python defaults to Chinese prefixes +// ("问题："/"回答："); the legacy Go code defaulted to English. +func TestQAChunker_DefaultLangIsChinese(t *testing.T) { + inputs := map[string]any{ + "name": "test.txt", + "output_format": "text", + "text": "What is RAG?\tRAG retrieves then generates.", + } + chunks := qaInvoke(t, inputs) + if len(chunks) == 0 { + t.Fatalf("no chunks produced") + } + cw := chunks[0]["content_with_weight"].(string) + if !contains(cw, "问题：") || !contains(cw, "回答：") { + t.Errorf("empty lang should default to Chinese prefixes, got %q", cw) + } + if contains(cw, "Question:") || contains(cw, "Answer:") { + t.Errorf("empty lang must not use English prefixes, got %q", cw) + } +} + +// TestRmQAPrefixStripsMultipleSeparators exercises migration diff +// Chunker-2.12: the prefix regex must allow one-or-more separator chars +// (Python uses `[\t:： ]+`), so "Q:: answer" is fully stripped. The legacy +// Go pattern only matched a single separator, leaving ": answer". +func TestRmQAPrefixStripsMultipleSeparators(t *testing.T) { + if got := rmQAPrefix("Q:: answer"); got != "answer" { + t.Errorf("multi-separator prefix not fully stripped: got %q", got) + } + if got := rmQAPrefix("Question: foo"); got != "foo" { + t.Errorf("single-separator prefix regression: got %q", got) + } + if got := rmQAPrefix("问：答案在此"); got != "答案在此" { + t.Errorf("CJK prefix regression: got %q", got) + } +} + +// TestQAChunker_SetsTopInt exercises migration diff Chunker-1.8 (top_int): +// each QA chunk must carry the source row index in `top_int`, matching +// Python beAdoc(..., row_num=i). +func TestQAChunker_SetsTopInt(t *testing.T) { + inputs := map[string]any{ + "name": "test.txt", + "output_format": "text", + "text": "Q1\tA1\nQ2\tA2", + } + chunks := qaInvoke(t, inputs) + if len(chunks) != 2 { + t.Fatalf("want 2 QA chunks, got %d", len(chunks)) + } + for i, c := range chunks { + raw, ok := c["top_int"] + if !ok { + t.Fatalf("chunk %d missing top_int field", i) + } + arr, ok := raw.([]any) + if !ok || len(arr) != 1 { + t.Fatalf("chunk %d top_int wrong shape: %#v", i, raw) + } + if int(arr[0].(float64)) != i { + t.Errorf("chunk %d top_int = %v, want %d", i, arr[0], i) + } + } +} + +// TestQAChunker_CarriesImageAndPositions exercises migration diff +// Chunker-1.8 (image / positions): when the upstream JSON item already +// carries an image id and pdf positions, the QA chunk must preserve them +// (Python beAdocPdf sets d["image"] and add_positions). +func TestQAChunker_CarriesImageAndPositions(t *testing.T) { + inputs := map[string]any{ + "name": "test.pdf", + "output_format": "json", + "json": []map[string]any{ + { + "text": "Q\tA", + "image": "img-42", + "_pdf_positions": [][]int{{1, 2, 3, 4, 5}}, + }, + }, + } + chunks := qaInvoke(t, inputs) + if len(chunks) != 1 { + t.Fatalf("want 1 QA chunk, got %d", len(chunks)) + } + c := chunks[0] + if c["image"] != "img-42" { + t.Errorf("QA chunk lost upstream image: %#v", c["image"]) + } + if _, ok := c["_pdf_positions"]; !ok { + t.Errorf("QA chunk lost upstream _pdf_positions") + } + // The prefix-stripped content must still be present. + cw, _ := c["content_with_weight"].(string) + if !contains(cw, "A") { + t.Errorf("QA content missing answer: %q", cw) + } +} + +// contains is a tiny helper to avoid importing strings in every test. +func contains(s, sub string) bool { + return strings.Contains(s, sub) +} diff --git a/internal/ingestion/component/chunker/qa_test.go b/internal/ingestion/component/chunker/qa_test.go index a305b00e38..4502cd00ed 100644 --- a/internal/ingestion/component/chunker/qa_test.go +++ b/internal/ingestion/component/chunker/qa_test.go @@ -39,7 +39,7 @@ func TestQAChunker_Registered(t *testing.T) { } func TestQAChunker_DelimiterTab(t *testing.T) { - comp, err := NewQAChunker(nil) + comp, err := NewQAChunker(map[string]any{"lang": "english"}) if err != nil { t.Fatal(err) } @@ -64,7 +64,7 @@ func TestQAChunker_DelimiterTab(t *testing.T) { } func TestQAChunker_DelimiterComma(t *testing.T) { - comp, err := NewQAChunker(nil) + comp, err := NewQAChunker(map[string]any{"lang": "english"}) if err != nil { t.Fatal(err) } @@ -129,7 +129,7 @@ func TestQAChunker_HTMLTable(t *testing.T) { } func TestQAChunker_RmQAPrefix(t *testing.T) { - comp, err := NewQAChunker(nil) + comp, err := NewQAChunker(map[string]any{"lang": "english"}) if err != nil { t.Fatal(err) } @@ -170,7 +170,7 @@ func TestQAChunker_Empty(t *testing.T) { } func TestQAChunker_CaseInsensitivePrefix(t *testing.T) { - comp, err := NewQAChunker(nil) + comp, err := NewQAChunker(map[string]any{"lang": "english"}) if err != nil { t.Fatal(err) } @@ -193,8 +193,8 @@ func TestQAChunker_CaseInsensitivePrefix(t *testing.T) { } } -func TestQAChunker_PrefixRequiresColonOrTab(t *testing.T) { - comp, err := NewQAChunker(nil) +func TestQAChunker_PrefixSpaceSeparatorStrips(t *testing.T) { + comp, err := NewQAChunker(map[string]any{"lang": "english"}) if err != nil { t.Fatal(err) } @@ -212,8 +212,10 @@ func TestQAChunker_PrefixRequiresColonOrTab(t *testing.T) { t.Fatalf("expected 1 chunk, got %d", len(chunks)) } cww, _ := chunks[0]["content_with_weight"].(string) - if cww != "Question: A language model is useful\tAnswer: Q How does it work" { - t.Fatalf("space-only separator should not strip prefix: %q", cww) + // Python qa.py:241 uses `[\t:： ]+`, so a space is a valid separator: + // a leading "A"/"Q" followed by a space is stripped (diff Chunker-2.12). + if cww != "Question: language model is useful\tAnswer: How does it work" { + t.Fatalf("space-separator prefix not stripped: %q", cww) } } 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..9a8523d531 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": float64(0), "page_number": float64(1)}, + map[string]any{"title": "第二章", "level": float64(1), "page_number": float64(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..e08af56daf 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 ( @@ -285,6 +285,12 @@ func (c *TokenChunkerComponent) invokeTextPayload(_ context.Context, text string return chunkOutputs(flatten(merged)) } +// sentenceDelimiter is the sentence/clause-boundary regex used to split +// oversized sections. It mirrors Python's default delimiter "\n。；！？" +// plus the English ". " fallback, and also breaks on ASCII "!" and "?" +// (diff Chunker-2.1). +var sentenceDelimiter = regexp.MustCompile(`(\n|[!?。；！？]|\.\s)`) + // mergeByTokenSize implements exact token-based chunk merging that mirrors // Python's naive_merge (rag/nlp/__init__.py:1156). It uses // tokenizeStr (= tokenizer.NumTokensFromString, cl100k_base BPE) for @@ -303,8 +309,10 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r } // Sentence/clause-boundary regex for splitting oversized sections. - // Matches Python's default delimiter "\n。；！？" plus English ". " fallback. - sentenceDelim := regexp.MustCompile(`(\n|[。；！？]|\.\s)`) + // Matches Python's default delimiter "\n。；！？" plus English ". " + // fallback. The ASCII "!" and "?" are included to match Python's + // full delimiter set (diff Chunker-2.1). + sentenceDelim := sentenceDelimiter var cks []string // chunk texts var tkns []int // token counts per chunk @@ -319,7 +327,9 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r seg := segment segTokens := tokens if overlapPct > 0 && len(cks) > 0 { - prev := cks[len(cks)-1] + // Strip parser tags before computing the overlap suffix, + // matching Python nlp/__init__.py:1181 (diff Chunker-2.2). + prev := removeTag(cks[len(cks)-1]) // Take the last overlapped_percent of the previous chunk // (in runes, matching Python's len(overlapped) * ratio). prevRunes := []rune(prev) @@ -656,23 +666,37 @@ func collectContext(chunks []schema.ChunkDoc, i, ctxTokens int, above bool) stri return strings.Join(parts, "") } -// takeFromEnd returns the last approx `tokens` worth of text (1 token -// ≈ 4 bytes is the best-effort approximation used here; python uses -// the actual tokenizer). +// takeFromEnd returns the smallest tail of text whose token count is >= +// tokens, counted exactly via tokenizeStr (diff Chunker-2.4). The previous +// 4-bytes-per-token heuristic over-counted for CJK text. func takeFromEnd(text string, tokens int) string { - bytes := tokens * 4 - if bytes >= len(text) { - return text + runes := []rune(text) + // The tail runes[i:] grows as i decreases, so the first (largest i, + // i.e. smallest tail) that meets the budget is the answer. + for i := len(runes); i > 0; i-- { + cand := string(runes[i:]) + if tokenizeStr(cand) >= tokens { + return cand + } } - return text[len(text)-bytes:] + return text } +// takeFromStart returns the smallest prefix of text whose token count is >= +// tokens, counted exactly via tokenizeStr (diff Chunker-2.4). func takeFromStart(text string, tokens int) string { - bytes := tokens * 4 - if bytes >= len(text) { - return text + runes := []rune(text) + best := text + // Prefix grows as i increases; the first (smallest) qualifying prefix + // is the answer. + for i := 1; i <= len(runes); i++ { + cand := string(runes[:i]) + if tokenizeStr(cand) >= tokens { + best = cand + break + } } - return text[:bytes] + return best } // mergeByTokenSizeFromJSON mirrors `naive_merge` at @@ -704,7 +728,10 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over // t = overlapped[overlap_cut:] + t // tnum = num_tokens_from_string(t) if len(merged) > 0 && merged[len(merged)-1].CKType == "text" && overlappedPct > 0 { - if prevText := merged[len(merged)-1].Text; prevText != "" { + // Strip parser tags before computing the overlap + // suffix, matching Python nlp/__init__.py:1181 + // (diff Chunker-2.2). + if prevText := removeTag(merged[len(merged)-1].Text); prevText != "" { runes := []rune(prevText) cut := int(float64(len(runes)) * (100 - overlappedPct*100) / 100.0) if cut < len(runes) { @@ -718,10 +745,21 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over } // Merge into the accumulated text chunk. prev := &merged[len(merged)-1] - if prev.Text != "" { + // Mirror Python token_chunker.py:236-239: when the accumulated + // chunk has empty text, assign the incoming text directly instead + // of skipping it (diff Chunker-2.11). + if prev.Text == "" { + prev.Text = ck.Text + } else { prev.Text = prev.Text + "\n" + ck.Text - prev.TKNums = intPtr(intValue(prev.TKNums) + tk) } + 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 +780,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 +797,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_batch1_test.go b/internal/ingestion/component/chunker/token_batch1_test.go new file mode 100644 index 0000000000..e9d675ea1d --- /dev/null +++ b/internal/ingestion/component/chunker/token_batch1_test.go @@ -0,0 +1,128 @@ +// +// 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 ( + "regexp" + "strings" + "testing" + + "ragflow/internal/ingestion/component/schema" +) + +// TestSentenceDelimiterMatchesBangAndQuestion exercises migration diff +// Chunker-2.1: the sentence/clause boundary regex used to split oversized +// sections must also break on ASCII "!" and "?" (Python's default delimiter +// is "\n。；！？"). The legacy Go pattern `(\n|[。；！？]|\.\s)` missed the +// ASCII variants, so English fragments like "Hi!" / "Really?" were not +// treated as boundaries. +func TestSentenceDelimiterMatchesBangAndQuestion(t *testing.T) { + // The package-level sentenceDelimiter (introduced by Fix 2.1) must + // match ASCII bang/question. + if !sentenceDelimiter.MatchString("Hi!") { + t.Errorf("sentenceDelimiter should split on '!': %q", "Hi!") + } + if !sentenceDelimiter.MatchString("Really?") { + t.Errorf("sentenceDelimiter should split on '?': %q", "Really?") + } + + // Guard: the OLD pattern must NOT match these, proving the test would + // have failed before the fix. + old := regexp.MustCompile(`(\n|[。；！？]|\.\s)`) + if old.MatchString("Hi!") || old.MatchString("Really?") { + t.Errorf("guard broken: old pattern unexpectedly matches ASCII !/?") + } +} + +// TestMergeByTokenSizeFromJSON_OverlapStripsTags exercises migration diff +// Chunker-2.2: when a new chunk is started, its overlap prefix must be taken +// from the previous chunk AFTER remove_tag, otherwise parser tags (e.g. +// "@@1\t2.3##") leak into the overlap region. Mirrors Python +// nlp/__init__.py:1181 (remove_tag applied before overlap). +func TestMergeByTokenSizeFromJSON_OverlapStripsTags(t *testing.T) { + aText := strings.Repeat("word ", 20) + "@@1\t2.3## tail" + items := [][]schema.ChunkDoc{ + { + {Text: aText, DocType: "text", CKType: "text", TKNums: intPtr(100)}, + {Text: "body", DocType: "text", CKType: "text", TKNums: intPtr(5)}, + }, + } + got := mergeByTokenSizeFromJSON(items, 128, 0.3) + merged := got[0] + if len(merged) != 2 { + t.Fatalf("want 2 merged chunks (overlap path), got %d", len(merged)) + } + // The overlap prefix is prepended to the SECOND chunk. The original + // first chunk legitimately keeps its own parser tag; only the overlap + // region (merged[1]) must be tag-free (diff Chunker-2.2). + if strings.Contains(merged[1].Text, "@@") || strings.Contains(merged[1].Text, "##") { + t.Errorf("overlap prefix leaked parser tag into chunk 1: %q", merged[1].Text) + } +} + +// TestMergeByTokenSizeFromJSON_EmptyPrevKeepsChunk exercises migration diff +// Chunker-2.11: merging a non-empty chunk into an empty previous chunk must +// assign the text directly instead of being skipped. The legacy guard +// `if prev.Text != ""` silently dropped the incoming chunk when the previous +// one had empty text. Mirrors Python token_chunker.py:236-239. +func TestMergeByTokenSizeFromJSON_EmptyPrevKeepsChunk(t *testing.T) { + items := [][]schema.ChunkDoc{ + { + {Text: "", DocType: "text", CKType: "text", TKNums: intPtr(5)}, + {Text: "keepme", DocType: "text", CKType: "text", TKNums: intPtr(5)}, + }, + } + got := mergeByTokenSizeFromJSON(items, 128, 0) + merged := got[0] + if len(merged) != 1 { + t.Fatalf("want 1 merged chunk, got %d", len(merged)) + } + if merged[0].Text != "keepme" { + t.Errorf("empty previous chunk dropped incoming text; got %q", merged[0].Text) + } +} + +// TestTakeFromEndRespectsTokenCount and TestTakeFromStartRespectsTokenCount +// exercise migration diff Chunker-2.4: takeFromEnd/takeFromStart used a +// fixed 4-bytes-per-token heuristic which badly over-counts for CJK text +// (≈3 bytes/char, 1-2 tokens/char). They must now count tokens exactly via +// tokenizeStr so the returned slice is close to the requested token budget. +func TestTakeFromEndRespectsTokenCount(t *testing.T) { + const target = 20 + s := strings.Repeat("中", 60) + got := takeFromEnd(s, target) + if !strings.HasSuffix(s, got) { + t.Fatalf("takeFromEnd result must be a suffix of input") + } + n := tokenizeStr(got) + if n < target-3 || n > target+3 { + t.Errorf("takeFromEnd(%d tokens) returned slice with %d tokens (want ~%d)", target, n, target) + } +} + +func TestTakeFromStartRespectsTokenCount(t *testing.T) { + const target = 20 + s := strings.Repeat("中", 60) + got := takeFromStart(s, target) + if !strings.HasPrefix(s, got) { + t.Fatalf("takeFromStart result must be a prefix of input") + } + n := tokenizeStr(got) + if n < target-3 || n > target+3 { + t.Errorf("takeFromStart(%d tokens) returned slice with %d tokens (want ~%d)", target, n, target) + } +} 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/docx_vision_dispatch.go b/internal/ingestion/component/docx_vision_dispatch.go index 622bb968e5..5c0ac2ad70 100644 --- a/internal/ingestion/component/docx_vision_dispatch.go +++ b/internal/ingestion/component/docx_vision_dispatch.go @@ -16,13 +16,14 @@ // DOCX vision figure dispatch: enriches the parse result with // LLM-generated descriptions of embedded images, mirroring -// Python's vision_figure_parser_docx_wrapper_naive in -// deepdoc/parser/figure_parser.py. +// Python's enhance_media_sections_with_vision in +// rag/flow/parser/utils.py (invoked from parser.py:_doc's JSON branch). // -// Unlike the PDF vision path (which replaces dispatchParse -// entirely), DOCX vision is a post-processing step: it takes -// the already-parsed markdown + extracted figures and augments -// the markdown text with vision model descriptions. +// Unlike the PDF vision path (which replaces dispatchParse entirely), +// DOCX vision is a post-processing step. It mirrors Python exactly: +// vision enrichment happens ONLY on the JSON output path, where each +// item carries a doc_type_kwd and an optional image. The markdown path +// performs no vision enrichment in Python, so it must not here either. package component @@ -31,7 +32,6 @@ import ( "fmt" "os" "path/filepath" - "sort" "strings" "sync" @@ -59,18 +59,20 @@ var ( docxVisionPromptMu sync.RWMutex ) -// maybeDispatchDOCXVision checks whether the dispatch result for a -// DOCX file contains embedded image figures and, when a vision -// model is available, enriches the markdown with AI-generated -// figure descriptions. It mirrors the Python flow: +// maybeDispatchDOCXVision enriches a DOCX parse result with vision-model +// descriptions of embedded images. It mirrors Python's +// enhance_media_sections_with_vision (rag/flow/parser/utils.py:162), which +// runs only in the JSON output branch of parser.py:_doc. // -// 1. naive_merge_docx → chunks (text + images + context) -// 2. vision_figure_parser_docx_wrapper_naive → LLM descriptions +// For each JSON item whose doc_type_kwd is "image" or "table" AND that +// carries a non-empty "image" field, the vision model describes the image +// and the description is appended to the item's text (Python: +// item["text"] = f"{text}\n{parsed_text}" if text else parsed_text). Items +// without an image (e.g. DOCX tables) are left untouched, exactly as Python +// skips them via `if item.get("image") is None: continue`. // -// The function is called AFTER dispatchParse so the normal parse -// path produces figures in dispatched.File["figures"]. -// It returns (result, handled, error). handled is true when the -// dispatched result was modified. +// The markdown output path receives no vision enrichment — Python's DOCX +// markdown branch only concatenates text and never calls the vision model. func maybeDispatchDOCXVision( ctx context.Context, fileType utility.FileType, @@ -81,11 +83,9 @@ func maybeDispatchDOCXVision( if fileType != utility.FileTypeDOCX { return dispatched, false, nil } - if dispatched.Err != nil || dispatched.OutputFormat != "markdown" { - return dispatched, false, nil - } - figs, hasFigures := extractDOCXFiguresFromDispatch(dispatched) - if !hasFigures { + // Python triggers vision enrichment only on the JSON path + // (parser.py:_doc → enhance_media_sections_with_vision). + if dispatched.Err != nil || dispatched.OutputFormat != "json" || len(dispatched.JSON) == 0 { return dispatched, false, nil } @@ -102,105 +102,74 @@ func maybeDispatchDOCXVision( return dispatched, false, nil } - descriptions := make([]string, len(figs)) + // Collect the indices of JSON items that carry an embeddable image. + type target struct { + idx int + } + var targets []target + for i, item := range dispatched.JSON { + kd, _ := item["doc_type_kwd"].(string) + if kd != "image" && kd != "table" { + continue + } + img, _ := item["image"].(string) + if img == "" { + continue + } + targets = append(targets, target{idx: i}) + } + if len(targets) == 0 { + return dispatched, false, nil + } + + descriptions := make([]string, len(targets)) var wg sync.WaitGroup sem := make(chan struct{}, docxVisionConcurrency) - for i, fig := range figs { + for slot, tg := range targets { wg.Add(1) - go func(idx int, f map[string]any) { + go func(slot int, itemIdx int) { defer wg.Done() sem <- struct{}{} defer func() { <-sem }() - imageB64, _ := f["image"].(string) - ctxAbove, _ := f["context_above"].(string) - ctxBelow, _ := f["context_below"].(string) - - if imageB64 == "" { + img, _ := dispatched.JSON[itemIdx]["image"].(string) + if img == "" { return } - - prompt, err := docxVisionPromptBuilder(ctxAbove, ctxBelow) - if err != nil { + // DOCX JSON items have no surrounding context (unlike the + // former markdown path), so use the bare figure prompt — + // matching Python's VisionFigureParser(context_size=0). + prompt, perr := docxVisionPromptBuilder("", "") + if perr != nil { return } - - messages := buildVisionMessages(prompt, imageB64) - resp, err := visionChatInvoker(ctx, driver, modelName, messages, apiConfig) - if err != nil { + messages := buildVisionMessages(prompt, img) + resp, ierr := visionChatInvoker(ctx, driver, modelName, messages, apiConfig) + if ierr != nil { return } - descriptions[idx] = extractDOCXVisionAnswer(resp) - }(i, fig) + descriptions[slot] = extractDOCXVisionAnswer(resp) + }(slot, tg.idx) } wg.Wait() - // Insert each description at the figure's position in the markdown, - // matching Python's `chunks[idx]["text"] += description`. - // Figures carry a "marker" (text immediately before the image) to - // locate the insertion point. Process in reverse order so earlier - // insertions don't shift later markers. - md := dispatched.Markdown - type indexedDesc struct { - idx int - desc string - } - var inserts []indexedDesc - for i, d := range descriptions { - if d = strings.TrimSpace(d); d == "" { + modified := false + for slot, tg := range targets { + desc := strings.TrimSpace(descriptions[slot]) + if desc == "" { continue } - if i >= len(figs) { - continue + existing, _ := dispatched.JSON[tg.idx]["text"].(string) + if existing != "" { + dispatched.JSON[tg.idx]["text"] = existing + "\n" + desc + } else { + dispatched.JSON[tg.idx]["text"] = desc } - marker, _ := figs[i]["marker"].(string) - if marker != "" { - if pos := strings.LastIndex(md, marker); pos >= 0 { - inserts = append(inserts, indexedDesc{idx: pos + len(marker), desc: d}) - continue - } - } - // Fallback: try context_above as a search anchor. - if ctx, _ := figs[i]["context_above"].(string); ctx != "" { - if pos := strings.LastIndex(md, ctx); pos >= 0 { - inserts = append(inserts, indexedDesc{idx: pos + len(ctx), desc: d}) - continue - } - } - // No anchor found — append to end. - inserts = append(inserts, indexedDesc{idx: len(md), desc: "\n\n" + d}) + modified = true } - // Sort descending by position for stable insertion. - sort.Slice(inserts, func(a, b int) bool { return inserts[a].idx > inserts[b].idx }) - for _, ins := range inserts { - desc := ins.desc - if !strings.HasPrefix(desc, "\n") { - desc = "\n\n" + desc - } - md = md[:ins.idx] + desc + md[ins.idx:] - } - dispatched.Markdown = md - return dispatched, true, nil -} - -func extractDOCXFiguresFromDispatch(dispatched parserDispatchResult) ([]map[string]any, bool) { - if dispatched.File == nil { - return nil, false - } - raw, ok := dispatched.File["figures"] - if !ok { - return nil, false - } - list, ok := raw.([]map[string]any) - if !ok { - return nil, false - } - if len(list) == 0 { - return nil, false - } - return list, true + return dispatched, modified, nil } // buildDOCXVisionPrompt loads the figure-describe prompt template diff --git a/internal/ingestion/component/docx_vision_dispatch_test.go b/internal/ingestion/component/docx_vision_dispatch_test.go new file mode 100644 index 0000000000..7554a2adac --- /dev/null +++ b/internal/ingestion/component/docx_vision_dispatch_test.go @@ -0,0 +1,181 @@ +// +// 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 component + +import ( + "context" + "sync" + "testing" + + "ragflow/internal/entity" + modelModule "ragflow/internal/entity/models" + "ragflow/internal/utility" +) + +// docxVisionFakeDriver satisfies modelModule.ModelDriver but never reaches the +// network: docxVisionCaptureInvoker intercepts the call before the driver. +type docxVisionFakeDriver struct { + modelModule.ModelDriver +} + +// docxVisionCaptureInvoker records the requested image and returns a fixed +// description, mirroring the markdown-vision test's capture driver. +type docxVisionCaptureInvoker struct { + mu sync.Mutex + images []string + captured []modelModule.Message +} + +func (c *docxVisionCaptureInvoker) invoke( + ctx context.Context, + driver modelModule.ModelDriver, + modelName string, + messages []modelModule.Message, + apiConfig *modelModule.APIConfig, +) (*modelModule.ChatResponse, error) { + c.mu.Lock() + c.captured = append(c.captured, messages...) + // Pull the data URI out of the second content part. + if parts, ok := messages[0].Content.([]interface{}); ok && len(parts) >= 2 { + if img, ok := parts[1].(map[string]any); ok { + if url, ok := img["image_url"].(map[string]any); ok { + if u, ok := url["url"].(string); ok { + c.images = append(c.images, u) + } + } + } + } + c.mu.Unlock() + ans := "a diagram of a pipeline" + return &modelModule.ChatResponse{Answer: &ans}, nil +} + +// TestMaybeDispatchDOCXVision_EnhancesJSONImages verifies Diff 2.4: DOCX vision +// enhancement must trigger on the JSON output path (like Python's +// enhance_media_sections_with_vision in parser.py:_doc) and must NOT trigger on +// the markdown path. Image items with a non-empty `image` field get their VLM +// description appended to `text`; table items (no image) and text items are +// left untouched. +func TestMaybeDispatchDOCXVision_EnhancesJSONImages(t *testing.T) { + origResolver := resolveTenantModelByType + origInvoker := visionChatInvoker + origPrompt := docxVisionPromptBuilder + defer func() { + resolveTenantModelByType = origResolver + visionChatInvoker = origInvoker + docxVisionPromptBuilder = origPrompt + }() + + resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + return &docxVisionFakeDriver{}, "docx-vision-model", &modelModule.APIConfig{}, 0, nil + } + invoker := &docxVisionCaptureInvoker{} + visionChatInvoker = invoker.invoke + docxVisionPromptBuilder = func(string, string) (string, error) { return "describe the figure", nil } + + dispatched := parserDispatchResult{ + OutputFormat: "json", + DocType: "docx", + JSON: []map[string]any{ + {"text": "Intro paragraph", "image": nil, "doc_type_kwd": "text"}, + {"text": "", "image": "aGVsbG8taW1hZ2U=", "doc_type_kwd": "image"}, + {"text": "
", "image": nil, "doc_type_kwd": "table"}, + }, + } + + res, handled, err := maybeDispatchDOCXVision( + context.Background(), + utility.FileTypeDOCX, + dispatched, + map[string]any{"tenant_id": "t1"}, + defaultSetups(), + ) + if err != nil { + t.Fatalf("maybeDispatchDOCXVision: unexpected error: %v", err) + } + if !handled { + t.Fatal("handled = false, want true (JSON image item should be enhanced)") + } + if len(res.JSON) != 3 { + t.Fatalf("JSON len = %d, want 3", len(res.JSON)) + } + if got := res.JSON[0]["text"].(string); got != "Intro paragraph" { + t.Errorf("text item text = %q, want unchanged", got) + } + // image item: VLM description appended to (empty) text. + if got, _ := res.JSON[1]["text"].(string); got != "a diagram of a pipeline" { + t.Errorf("image item text = %q, want appended VLM description", got) + } + if got, _ := res.JSON[2]["text"].(string); got != "
" { + t.Errorf("table item text = %q, want unchanged (no image)", got) + } + if len(invoker.images) != 1 { + t.Fatalf("vision invoker called %d times, want 1 (only the image item)", len(invoker.images)) + } + if want := "data:image/png;base64,aGVsbG8taW1hZ2U="; invoker.images[0] != want { + t.Errorf("vision image data URI = %q, want %q", invoker.images[0], want) + } +} + +// TestMaybeDispatchDOCXVision_JSONOnly verifies Diff 2.4: the markdown output +// path must NOT be enhanced (Python's markdown/docx branch performs no vision +// enrichment). A markdown result with embedded figures is returned untouched. +func TestMaybeDispatchDOCXVision_JSONOnly(t *testing.T) { + origResolver := resolveTenantModelByType + origInvoker := visionChatInvoker + defer func() { + resolveTenantModelByType = origResolver + visionChatInvoker = origInvoker + }() + + called := false + resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + called = true + return &docxVisionFakeDriver{}, "m", &modelModule.APIConfig{}, 0, nil + } + visionChatInvoker = func(ctx context.Context, d modelModule.ModelDriver, m string, msgs []modelModule.Message, c *modelModule.APIConfig) (*modelModule.ChatResponse, error) { + called = true + ans := "x" + return &modelModule.ChatResponse{Answer: &ans}, nil + } + + dispatched := parserDispatchResult{ + OutputFormat: "markdown", + DocType: "docx", + Markdown: "![Image](data:image/png;base64,abc)", + File: map[string]any{"figures": []map[string]any{{"image": "abc", "marker": "x"}}}, + } + + res, handled, err := maybeDispatchDOCXVision( + context.Background(), + utility.FileTypeDOCX, + dispatched, + map[string]any{"tenant_id": "t1"}, + defaultSetups(), + ) + if err != nil { + t.Fatalf("maybeDispatchDOCXVision: unexpected error: %v", err) + } + if handled { + t.Error("handled = true, want false (markdown path must not be enhanced)") + } + if called { + t.Error("vision model was resolved/invoked on the markdown path") + } + if res.Markdown != "![Image](data:image/png;base64,abc)" { + t.Errorf("markdown mutated: %q", res.Markdown) + } +} diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index 309830fc57..ac4ae73e5d 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -169,6 +169,13 @@ func NewExtractorComponent(params map[string]any) (runtime.Component, error) { } if v, ok := params["prompt"].(string); ok { p.Prompt = v + } else if v, ok := params["prompts"].(string); ok && v != "" { + // Python agent/component/llm.py:119-120 normalizes a bare-string + // prompts into [{"role":"user","content":prompts}]. Mirror that + // here so a front-end/template that emits prompts as a string + // (the graph.nodes form / dsl testdata) is not silently dropped + // by the .([]any) assertion on the list branch below. + p.Prompt = v } else if promptsRaw, ok := params["prompts"].([]any); ok && len(promptsRaw) > 0 { if first, ok := promptsRaw[0].(map[string]any); ok { if content, ok := first["content"].(string); ok { diff --git a/internal/ingestion/component/extractor_test.go b/internal/ingestion/component/extractor_test.go index 79e828e0d5..79ccb1df0e 100644 --- a/internal/ingestion/component/extractor_test.go +++ b/internal/ingestion/component/extractor_test.go @@ -570,6 +570,48 @@ func TestNewExtractorComponent_PromptsArray_PromptWins(t *testing.T) { } } +// TestNewExtractorComponent_PromptsString verifies that a bare-string +// "prompts" (the shape emitted by the front-end graph.nodes form and +// the dsl/testdata templates) is normalized into Param.Prompt, mirroring +// Python agent/component/llm.py:119-120 which coerces a string prompts +// into [{"role":"user","content":prompts}]. Without this normalization +// the string form is silently dropped (the .([]any) assertion fails). +func TestNewExtractorComponent_PromptsString(t *testing.T) { + withStubChatInvoker(t, stubResponse{Content: "ok"}) + comp, err := NewExtractorComponent(map[string]any{ + "field_name": "out", + "prompts": "Content: {TitleChunker:FlatMiceFix@chunks}", + }) + if err != nil { + t.Fatalf("NewExtractorComponent: %v", err) + } + ec := comp.(*ExtractorComponent) + want := "Content: {TitleChunker:FlatMiceFix@chunks}" + if ec.Param.Prompt != want { + t.Errorf("Prompt = %q, want %q (string prompts should be normalized)", ec.Param.Prompt, want) + } +} + +// TestNewExtractorComponent_PromptsString_PromptWins verifies that +// "prompt" (string) still takes priority over a string-form "prompts" +// when both are present, matching the prompt>prompts precedence of +// the list-form path (TestNewExtractorComponent_PromptsArray_PromptWins). +func TestNewExtractorComponent_PromptsString_PromptWins(t *testing.T) { + withStubChatInvoker(t, stubResponse{Content: "ok"}) + comp, err := NewExtractorComponent(map[string]any{ + "field_name": "out", + "prompt": "Direct prompt wins.", + "prompts": "Should be ignored.", + }) + if err != nil { + t.Fatalf("NewExtractorComponent: %v", err) + } + ec := comp.(*ExtractorComponent) + if ec.Param.Prompt != "Direct prompt wins." { + t.Errorf("Prompt = %q, want %q", ec.Param.Prompt, "Direct prompt wins.") + } +} + // TestNewExtractorComponent_SystemPromptWinsOverSysPrompt verifies // that "system_prompt" takes priority over "sys_prompt". func TestNewExtractorComponent_SystemPromptWinsOverSysPrompt(t *testing.T) { diff --git a/internal/ingestion/component/markdown_vision_dispatch.go b/internal/ingestion/component/markdown_vision_dispatch.go index 59f3793d6b..069900b7df 100644 --- a/internal/ingestion/component/markdown_vision_dispatch.go +++ b/internal/ingestion/component/markdown_vision_dispatch.go @@ -78,7 +78,10 @@ func maybeDispatchMarkdownVision( var images []imgItem for i, item := range dispatched.JSON { kd, _ := item["doc_type_kwd"].(string) - if kd != "image" { + // Diff 2.5: Python enhances both image and table items + // (parser/utils.py:181 checks {"image","table"}); only items + // carrying an image are sent to the VLM (utils.py:183). + if kd != "image" && kd != "table" { continue } img, _ := item["image"].(string) diff --git a/internal/ingestion/component/media_dispatch.go b/internal/ingestion/component/media_dispatch.go index e7a14f6721..e627c125d0 100644 --- a/internal/ingestion/component/media_dispatch.go +++ b/internal/ingestion/component/media_dispatch.go @@ -168,10 +168,17 @@ func maybeDispatchImage( } } - outputFormat, _ := setup["output_format"].(string) - if outputFormat == "" { - outputFormat = "text" - } + // The image family always emits a structured JSON item carrying the + // image attachment (data URI) and doc_type_kwd, mirroring Python + // rag/app/picture.py:71-72 (doc["image"]=img, doc["doc_type_kwd"]= + // "image"). picture.py has no "text" output mode — it always returns + // a structured doc — so output_format is hardcoded to "json" and any + // setup override is ignored. The former behavior returned a bare Text + // string, which dropped the image attachment, set doc_type to "text", + // and on the default json path produced JSON=nil so downstream + // Chunkers rejected the payload with errRequiredField{"json"}. + imageB64 := base64.StdEncoding.EncodeToString(binary) + dataURI := "data:" + imageMIME(filename) + ";base64," + imageB64 // --- Phase 2: VLM description (when OCR text is short) --- // Mirrors Python's check: if (eng and len(txt.split()) > 32) or len(txt) > 32 @@ -184,11 +191,7 @@ func maybeDispatchImage( charCount := len(ocrText) if (eng && wordCount > 32) || charCount > 32 { // OCR returned substantial text — skip VLM. - return parserDispatchResult{ - OutputFormat: outputFormat, - DocType: "image", - Text: ocrText, - }, true, nil + return imageDispatchResult(ocrText, dataURI), true, nil } } @@ -197,20 +200,12 @@ func maybeDispatchImage( if err != nil { // If VLM is unavailable but we have OCR text, return it. if ocrText != "" { - return parserDispatchResult{ - OutputFormat: outputFormat, - DocType: "image", - Text: ocrText, - }, true, nil + return imageDispatchResult(ocrText, dataURI), true, nil } return parserDispatchResult{}, true, fmt.Errorf("Parser: picture image2text model: %w", err) } - imageB64 := base64.StdEncoding.EncodeToString(binary) - mimeType := imageMIME(filename) - dataURI := "data:" + mimeType + ";base64," + imageB64 - prompt := "Describe this image in detail." // image family's contract key is system_prompt (parser.go:295), // mirroring Python parser.py:1119. Do NOT read setup["prompt"] @@ -229,11 +224,7 @@ func maybeDispatchImage( resp, err := driver.ChatWithMessages(ctx, modelName, messages, apiConfig, &modelModule.ChatConfig{Vision: &vision}, nil) if err != nil { if ocrText != "" { - return parserDispatchResult{ - OutputFormat: outputFormat, - DocType: "image", - Text: ocrText, - }, true, nil + return imageDispatchResult(ocrText, dataURI), true, nil } return parserDispatchResult{}, true, fmt.Errorf("Parser: picture describe: %w", err) @@ -253,11 +244,23 @@ func maybeDispatchImage( combined = vlmText } } + return imageDispatchResult(combined, dataURI), true, nil +} + +// imageDispatchResult builds the structured JSON payload for the image +// family: a single item carrying the combined text, the image attachment +// (data URI), and doc_type_kwd "image". Mirrors Python +// rag/app/picture.py:71-72. +func imageDispatchResult(text, dataURI string) parserDispatchResult { return parserDispatchResult{ - OutputFormat: outputFormat, + OutputFormat: "json", DocType: "image", - Text: combined, - }, true, nil + JSON: []map[string]any{{ + "text": text, + "image": dataURI, + "doc_type_kwd": "image", + }}, + } } // Audio dispatch: SPEECH2TEXT transcription --- @@ -315,6 +318,21 @@ func maybeDispatchAudio( if outputFormat == "" { outputFormat = "text" } + // Diff 2.11: when output_format is "json" the transcription must be + // carried as a JSON item. Returning it only in Text made the Invoke + // switch silently drop it (the switch has no "json" branch and the + // JSON slice was empty). Mirror the JSON-item shape used by the + // other parser branches. + if outputFormat == "json" { + return parserDispatchResult{ + OutputFormat: "json", + DocType: "audio", + JSON: []map[string]any{{ + "text": transcription, + "doc_type_kwd": "audio", + }}, + }, true, nil + } return parserDispatchResult{ OutputFormat: outputFormat, DocType: "audio", diff --git a/internal/ingestion/component/media_dispatch_test.go b/internal/ingestion/component/media_dispatch_test.go index 0b12589fad..1d684e44de 100644 --- a/internal/ingestion/component/media_dispatch_test.go +++ b/internal/ingestion/component/media_dispatch_test.go @@ -17,6 +17,7 @@ package component import ( "context" + "strings" "sync" "testing" @@ -104,8 +105,18 @@ func TestMaybeDispatchImage_UsesSystemPrompt(t *testing.T) { if !dispatched { t.Fatalf("expected dispatched=true for VISUAL file") } - if res.Text == "" { - t.Fatalf("expected non-empty combined text") + // After the output-shape fix the image branch returns JSON items + // (OutputFormat=="json"), not a bare Text field. The combined text + // now lives in JSON[0]["text"]; the legacy res.Text is no longer + // populated for the image family. + if res.OutputFormat != "json" { + t.Fatalf("OutputFormat = %q, want json (image family is always structured)", res.OutputFormat) + } + if len(res.JSON) != 1 { + t.Fatalf("JSON len = %d, want 1 (image result must be a single JSON item)", len(res.JSON)) + } + if txt, _ := res.JSON[0]["text"].(string); txt == "" { + t.Fatalf("expected non-empty combined text in JSON[0][\"text\"]") } got, ok := firstUserText(drv.captured) @@ -116,3 +127,238 @@ func TestMaybeDispatchImage_UsesSystemPrompt(t *testing.T) { t.Fatalf("VLM user text = %q, want %q (image branch must read system_prompt)", got, "自定义视觉提示") } } + +// TestMaybeDispatchImage_ReturnsJSONWithImage pins the output-shape fix: +// the image branch must return a JSON item carrying the `image` attachment +// (data URI) and `doc_type_kwd:"image"`, mirroring Python +// rag/app/picture.py:71-72. Before the fix the branch returned a bare Text +// string with JSON=nil, dropping the image attachment and (on the default +// json path) causing OneChunker/TokenChunker to reject the payload. +func TestMaybeDispatchImage_ReturnsJSONWithImage(t *testing.T) { + origResolver := resolveTenantModelByType + defer func() { resolveTenantModelByType = origResolver }() + + drv := &imagePromptCaptureDriver{} + resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + return drv, "img-model", &modelModule.APIConfig{}, 0, nil + } + + setups := defaultSetups() + res, dispatched, err := maybeDispatchImage( + context.Background(), + utility.FileTypeVISUAL, + "test.png", + []byte("not-a-real-image"), + map[string]any{"tenant_id": "t1"}, + setups, + ) + if err != nil { + t.Fatalf("maybeDispatchImage: %v", err) + } + if !dispatched { + t.Fatalf("expected dispatched=true") + } + if res.OutputFormat != "json" { + t.Fatalf("OutputFormat = %q, want json", res.OutputFormat) + } + if len(res.JSON) != 1 { + t.Fatalf("JSON len = %d, want 1", len(res.JSON)) + } + item := res.JSON[0] + if got, _ := item["doc_type_kwd"].(string); got != "image" { + t.Errorf("doc_type_kwd = %q, want \"image\"", got) + } + img, _ := item["image"].(string) + if !strings.HasPrefix(img, "data:") || !strings.Contains(img, ";base64,") { + t.Errorf("image = %q, want a data URI (data:;base64,)", img) + } + if txt, _ := item["text"].(string); txt == "" { + t.Errorf("text field empty; want non-empty combined OCR+VLM text") + } +} + +// TestMaybeDispatchImage_HardcodesJSONOutput verifies the image family +// always emits json regardless of setup["output_format"]. Python +// rag/app/picture.py:chunk() has no output_format concept — it always +// returns a structured doc. Honoring a "text" override produced a bare +// Text payload that lost the image attachment and set doc_type to "text". +func TestMaybeDispatchImage_HardcodesJSONOutput(t *testing.T) { + origResolver := resolveTenantModelByType + defer func() { resolveTenantModelByType = origResolver }() + + drv := &imagePromptCaptureDriver{} + resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + return drv, "img-model", &modelModule.APIConfig{}, 0, nil + } + + setups := defaultSetups() + setups["image"]["output_format"] = "text" // legacy/override; must be ignored + res, _, err := maybeDispatchImage( + context.Background(), + utility.FileTypeVISUAL, + "test.png", + []byte("not-a-real-image"), + map[string]any{"tenant_id": "t1"}, + setups, + ) + if err != nil { + t.Fatalf("maybeDispatchImage: %v", err) + } + if res.OutputFormat != "json" { + t.Fatalf("OutputFormat = %q, want json (image family must ignore output_format override)", res.OutputFormat) + } + if len(res.JSON) != 1 { + t.Fatalf("JSON len = %d, want 1 even when setup says text", len(res.JSON)) + } +} + +// audioTranscribeDriver is a mock ModelDriver whose TranscribeAudio returns a +// fixed transcription, so maybeDispatchAudio can be exercised without a real +// ASR provider. +type audioTranscribeDriver struct { + modelModule.ModelDriver + transcription string +} + +func (d *audioTranscribeDriver) TranscribeAudio(ctx context.Context, _ *string, _ *string, _ *modelModule.APIConfig, _ *modelModule.ASRConfig, _ *common.ModelUsage) (*modelModule.ASRResponse, error) { + return &modelModule.ASRResponse{Text: d.transcription}, nil +} + +// TestMaybeDispatchAudio_JSONCarriesTranscription pins diff 2.11: when the +// audio family's output_format is "json", the ASR transcription must be +// carried in the JSON items (not only in the Text field). Before the fix the +// branch returned Text only with an empty JSON slice, and the Invoke switch +// silently dropped the transcription because it has no "json" branch. +func TestMaybeDispatchAudio_JSONCarriesTranscription(t *testing.T) { + origResolver := resolveTenantModelByType + defer func() { resolveTenantModelByType = origResolver }() + + const want = "hello world" + drv := &audioTranscribeDriver{transcription: want} + resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + return drv, "asr-model", &modelModule.APIConfig{}, 0, nil + } + + setups := defaultSetups() + setups["audio"]["output_format"] = "json" + + res, dispatched, err := maybeDispatchAudio( + context.Background(), + utility.FileTypeAURAL, + "test.mp3", + []byte("fake-audio"), + map[string]any{"tenant_id": "t1"}, + setups, + ) + if err != nil { + t.Fatalf("maybeDispatchAudio: %v", err) + } + if !dispatched { + t.Fatalf("expected dispatched=true for AURAL file") + } + if res.OutputFormat != "json" { + t.Fatalf("OutputFormat = %q, want json", res.OutputFormat) + } + if len(res.JSON) != 1 { + t.Fatalf("JSON len = %d, want 1 (transcription must be carried as a JSON item)", len(res.JSON)) + } + if got, _ := res.JSON[0]["text"].(string); got != want { + t.Fatalf("JSON[0].text = %q, want %q", got, want) + } + if got, _ := res.JSON[0]["doc_type_kwd"].(string); got != "audio" { + t.Fatalf("JSON[0].doc_type_kwd = %q, want audio", got) + } +} + +// TestMaybeDispatchAudio_TextCarriesTranscription guards the text path: with +// output_format "text" the transcription stays in the Text field and JSON is +// empty (current default after aligning with Python parser.py:232). +func TestMaybeDispatchAudio_TextCarriesTranscription(t *testing.T) { + origResolver := resolveTenantModelByType + defer func() { resolveTenantModelByType = origResolver }() + + const want = "hello world" + drv := &audioTranscribeDriver{transcription: want} + resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + return drv, "asr-model", &modelModule.APIConfig{}, 0, nil + } + + setups := defaultSetups() + setups["audio"]["output_format"] = "text" + + res, dispatched, err := maybeDispatchAudio( + context.Background(), + utility.FileTypeAURAL, + "test.mp3", + []byte("fake-audio"), + map[string]any{"tenant_id": "t1"}, + setups, + ) + if err != nil { + t.Fatalf("maybeDispatchAudio: %v", err) + } + if !dispatched { + t.Fatalf("expected dispatched=true for AURAL file") + } + if res.OutputFormat != "text" { + t.Fatalf("OutputFormat = %q, want text", res.OutputFormat) + } + if res.Text != want { + t.Fatalf("Text = %q, want %q", res.Text, want) + } + if len(res.JSON) != 0 { + t.Fatalf("JSON len = %d, want 0 for text output", len(res.JSON)) + } +} + +// TestMaybeDispatchMarkdownVision_EnhancesTables pins diff 2.5: markdown +// vision enhancement must also process items whose doc_type_kwd is "table" +// (Python checks {"image","table"} in parser/utils.py:181), not only "image". +// Before the fix the table item was skipped and never sent to the VLM. +func TestMaybeDispatchMarkdownVision_EnhancesTables(t *testing.T) { + origResolver := resolveTenantModelByType + defer func() { resolveTenantModelByType = origResolver }() + + drv := &imagePromptCaptureDriver{} + resolveTenantModelByType = func(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { + return drv, "img-model", &modelModule.APIConfig{}, 0, nil + } + + dispatched := parserDispatchResult{ + OutputFormat: "json", + JSON: []map[string]any{ + {"doc_type_kwd": "table", "image": "base64table", "text": ""}, + }, + } + + res, handled, err := maybeDispatchMarkdownVision( + context.Background(), + utility.FileTypeMarkdown, + dispatched, + map[string]any{"tenant_id": "t1"}, + ) + if err != nil { + t.Fatalf("maybeDispatchMarkdownVision: %v", err) + } + if !handled { + t.Fatalf("expected handled=true for markdown with a table image") + } + if len(res.JSON) != 1 { + t.Fatalf("JSON len = %d, want 1", len(res.JSON)) + } + // The table item must have been sent to the VLM and its description appended. + if got, _ := res.JSON[0]["text"].(string); got != "captured" { + t.Fatalf("table item text = %q, want %q (table items must be vision-enhanced)", got, "captured") + } +} + +// TestDefaultEmailOutputFormatIsJSON pins diff 2.2: the email family default +// output_format must be "json" (matching Python parser.py:212), not "text". +// With "text" the structured email fields (from/to/subject/attachments/...) are +// flattened into a blob and lost downstream. +func TestDefaultEmailOutputFormatIsJSON(t *testing.T) { + got, _ := defaultSetups()["email"]["output_format"].(string) + if got != "json" { + t.Fatalf("email default output_format = %q, want json", got) + } +} diff --git a/internal/ingestion/component/parser.go b/internal/ingestion/component/parser.go index 3438fe0fa7..1b6f867302 100644 --- a/internal/ingestion/component/parser.go +++ b/internal/ingestion/component/parser.go @@ -302,7 +302,7 @@ func defaultSetups() map[string]schema.ParserSetup { "from", "to", "cc", "bcc", "date", "subject", "body", "attachments", "metadata", }, - "output_format": "text", + "output_format": "json", }, "audio": { "suffix": []string{ @@ -310,7 +310,7 @@ func defaultSetups() map[string]schema.ParserSetup { "aiff", "au", "midi", "wma", "realaudio", "vqf", "oggvorbis", "ape", }, - "output_format": "json", + "output_format": "text", }, "video": { "suffix": []string{"mp4", "avi", "mkv"}, @@ -478,9 +478,10 @@ func (c *ParserComponent) Invoke(ctx context.Context, inputs map[string]any) (ma dispatched = dispatchParse(ctx, fileTypeExt, filename, binary, c.Setups) dispatched = hydrateEmptyDispatchPayload(dispatched, binary) - // DOCX vision figure enhancement: enrich the markdown - // with LLM-generated descriptions of embedded images. - // Mirrors Python's vision_figure_parser_docx_wrapper_naive. + // DOCX vision figure enhancement: on the JSON output path, + // append vision-model descriptions to embedded image items + // (doc_type_kwd "image"). Mirrors Python's + // enhance_media_sections_with_vision in parser.py:_doc. dispatched, _, _ = maybeDispatchDOCXVision(ctx, fileTypeExt, dispatched, inputs, c.Setups) // Markdown vision figure enhancement: enrich parsed diff --git a/internal/ingestion/component/pdf_vision_dispatch.go b/internal/ingestion/component/pdf_vision_dispatch.go index 8354a833c9..3bd6fbf04a 100644 --- a/internal/ingestion/component/pdf_vision_dispatch.go +++ b/internal/ingestion/component/pdf_vision_dispatch.go @@ -358,16 +358,27 @@ func resolvePDFVisionModelID(setup schema.ParserSetup) (string, bool) { return "", false } +// isNamedPDFParseMethod reports whether raw is a recognized named PDF +// parse method (as opposed to a CustomVLM model name). Its membership set +// MUST stay aligned with the PDF whitelist enforced by +// (*ParserComponent).Check() (parser.go:200-203): +// +// deepdoc, plain_text, mineru, docling, +// opendataloader, tcadp parser, paddleocr, somark +// +// A parse_method that Check() rejects must not be treated as a named method +// here, otherwise it silently falls through to the CustomVLM vision path +// instead of failing fast at construction. +// +// Note: "@"-suffixed spellings such as "foo@mineru" are layout_recognizer +// selectors, not parse_method values. Check() rejects them as parse_method, +// and the MinerU layout branch is resolved from the layout_recognizer field +// separately (pdf_vision_dispatch.go:62-68), so they must NOT be recognized +// here. func isNamedPDFParseMethod(raw string) bool { method := strings.ToLower(strings.TrimSpace(raw)) - switch { - case strings.HasSuffix(method, "@paddleocr"), - strings.HasSuffix(method, "@somark"), - strings.HasSuffix(method, "@opendataloader"): - return true - } switch method { - case "deepdoc", "mineru", "plain_text", "plain text", "plaintext", "paddleocr", "docling", "opendataloader", "somark", "tcadp", "tcadp parser": + case "deepdoc", "plain_text", "mineru", "docling", "opendataloader", "tcadp parser", "paddleocr", "somark": return true } return false diff --git a/internal/ingestion/component/pdf_vision_dispatch_test.go b/internal/ingestion/component/pdf_vision_dispatch_test.go new file mode 100644 index 0000000000..de543b5db8 --- /dev/null +++ b/internal/ingestion/component/pdf_vision_dispatch_test.go @@ -0,0 +1,86 @@ +// +// 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 component + +import "testing" + +// TestIsNamedPDFParseMethodWhitelistAligned verifies that the runtime +// "named parse_method" classifier agrees with (*ParserComponent).Check()'s +// PDF whitelist (parser.go:200-203): +// +// deepdoc, plain_text, mineru, docling, +// opendataloader, tcadp parser, paddleocr, somark +// +// Diff 2.10: a parse_method that Check() rejects must NOT be treated as a +// recognized named method by isNamedPDFParseMethod — otherwise it silently +// falls through to the CustomVLM vision path instead of failing fast at +// construction (and Python would have rejected it outright). +func TestIsNamedPDFParseMethodWhitelistAligned(t *testing.T) { + // Values that MUST be recognized (subset of the Check() whitelist, + // case-insensitive). + named := []string{ + "deepdoc", "plain_text", "mineru", "docling", + "opendataloader", "tcadp parser", "paddleocr", "somark", + "DeepDoc", "PLAIN_TEXT", "MinerU", "DocLing", + "OpenDataLoader", "TCADP Parser", "PaddleOCR", "SoMark", + } + for _, v := range named { + if !isNamedPDFParseMethod(v) { + t.Errorf("isNamedPDFParseMethod(%q) = false, want true (in Check() whitelist)", v) + } + } + + // Values that MUST NOT be recognized. These either duplicate the + // whitelist with non-canonical spelling ("plain text"/"plaintext") + // or are bare-family abbreviations ("tcadp") that Check() does not + // accept, so they should be funneled to the CustomVLM path (or fail + // construction) rather than masquerading as a named method. + notNamed := []string{ + "plain text", "plaintext", "tcadp", + "CustomVLM", "some_vlm", "gpt-4o", + "", " ", + } + for _, v := range notNamed { + if isNamedPDFParseMethod(v) { + t.Errorf("isNamedPDFParseMethod(%q) = true, want false (not in Check() whitelist)", v) + } + } +} + +// TestIsNamedPDFParseMethodLayoutSuffixes verifies that "@"-suffixed +// layout_recognizer spellings are NOT treated as named parse methods. They +// are layout_recognizer selectors (resolved separately at +// pdf_vision_dispatch.go:62-68), and Check() rejects them as parse_method, +// so they must fall through to the CustomVLM/VLM path — consistent with the +// (*ParserComponent).Check() whitelist (parser.go:200-203). +func TestIsNamedPDFParseMethodLayoutSuffixes(t *testing.T) { + suffixed := []string{ + "foo@mineru", "@mineru", + "foo@paddleocr", "@paddleocr", + "foo@somark", "@somark", + "foo@opendataloader", "@opendataloader", + } + for _, v := range suffixed { + if isNamedPDFParseMethod(v) { + t.Errorf("isNamedPDFParseMethod(%q) = true, want false (layout_recognizer selector, not a named parse_method)", v) + } + } + + // An unknown suffix is also not a named method. + if isNamedPDFParseMethod("foo@unknown") { + t.Errorf("isNamedPDFParseMethod(%q) = true, want false", "foo@unknown") + } +} diff --git a/internal/ingestion/component/schema/chunk_types.go b/internal/ingestion/component/schema/chunk_types.go index b1b2001073..1cce99d2ed 100644 --- a/internal/ingestion/component/schema/chunk_types.go +++ b/internal/ingestion/component/schema/chunk_types.go @@ -122,6 +122,7 @@ type ChunkDoc struct { ContentSmLtks string `json:"content_sm_ltks,omitempty"` TagKwd []string `json:"tag_kwd,omitempty"` PageNumber *int `json:"page_number,omitempty"` + TopInt []int `json:"top_int,omitempty"` PDFPositions json.RawMessage `json:"_pdf_positions,omitempty"` Positions json.RawMessage `json:"positions,omitempty"` Extra map[string]json.RawMessage `json:"-"` @@ -142,7 +143,7 @@ func (d *ChunkDoc) UnmarshalJSON(data []byte) error { "ck_type", "tk_nums", "layout", "layout_type", "layoutno", "image", "context_above", "context_below", "questions", "keywords", "summary", "chunk_order_int", "title_tks", "title_sm_tks", "content_ltks", - "content_sm_ltks", "tag_kwd", "page_number", "_pdf_positions", "positions", + "content_sm_ltks", "tag_kwd", "page_number", "top_int", "_pdf_positions", "positions", } { delete(raw, key) } diff --git a/internal/ingestion/component/tokenizer.go b/internal/ingestion/component/tokenizer.go index dcae6b5ea9..5029772032 100644 --- a/internal/ingestion/component/tokenizer.go +++ b/internal/ingestion/component/tokenizer.go @@ -434,7 +434,11 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, em if trimmedName == "" { log.Printf("Tokenizer: empty name provided from upstream, embedding will skip title weighting") } else { - titleResults, err := encodeWithTimeout(ctx, embedder, []string{trimmedName}) + // Encode the raw name (no TrimSpace) to mirror Python + // tokenizer.py:95 which passes name verbatim to embedding. The + // empty-name guard above still uses TrimSpace, matching Python's + // `.strip()==""` check at tokenizer.py:200. + titleResults, err := encodeWithTimeout(ctx, embedder, []string{name}) if err != nil { return nil, 0, fmt.Errorf("Tokenizer: encode title: %w", err) } diff --git a/internal/ingestion/component/tokenizer_test.go b/internal/ingestion/component/tokenizer_test.go index b02b6bf9de..4883504aab 100644 --- a/internal/ingestion/component/tokenizer_test.go +++ b/internal/ingestion/component/tokenizer_test.go @@ -591,6 +591,32 @@ func TestTokenizerComponent_Embedding_EmptyNameWarnsAndUsesContentVector(t *test } } +// Python tokenizer.py:95 passes the raw name to embedding without .strip(); +// Go must match — the title embedding must receive the original name, not a +// TrimSpace'd copy. The empty-name guard still uses TrimSpace (mirroring +// Python's `.strip()==""` check at tokenizer.py:200), but the value encoded +// is the raw name. +func TestTokenizerComponent_Embedding_UsesRawNameNotTrimmed(t *testing.T) { + requireTokenizerPool(t) + c, stub := withStubEmbedder(t, 2) + + if _, err := c.Invoke(context.Background(), map[string]any{ + "name": " report.pdf ", + "output_format": "chunks", + "chunks": []map[string]any{{"text": "alpha"}}, + }); err != nil { + t.Fatalf("Invoke: %v", err) + } + if len(stub.callInputs) < 1 { + t.Fatalf("callInputs len = %d, want >= 1", len(stub.callInputs)) + } + // First call is the title embedding; it must receive the raw name with + // surrounding whitespace preserved, matching Python. + if got := stub.callInputs[0][0]; got != " report.pdf " { + t.Fatalf("title embedding input = %q, want %q (raw, not trimmed)", got, " report.pdf ") + } +} + func TestTokenizerComponent_Embedding_TruncatesByMaxTokensMinus10(t *testing.T) { requireTokenizerPool(t) c, stub := withStubEmbedder(t, 2) 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/pipeline/template/ingestion_pipeline_picture.json b/internal/ingestion/pipeline/template/ingestion_pipeline_picture.json index cca7efc839..f04f47a273 100644 --- a/internal/ingestion/pipeline/template/ingestion_pipeline_picture.json +++ b/internal/ingestion/pipeline/template/ingestion_pipeline_picture.json @@ -50,7 +50,7 @@ } }, "image": { - "output_format": "text", + "output_format": "json", "parse_method": "ocr", "preprocess": [ "main_content" @@ -231,7 +231,7 @@ "setups": [ { "fileFormat": "image", - "output_format": "text", + "output_format": "json", "parse_method": "ocr", "preprocess": [ "main_content" diff --git a/internal/ingestion/task/pipeline_executor_defaults_test.go b/internal/ingestion/task/pipeline_executor_defaults_test.go index f1e346ba7d..ceaa44fdc5 100644 --- a/internal/ingestion/task/pipeline_executor_defaults_test.go +++ b/internal/ingestion/task/pipeline_executor_defaults_test.go @@ -46,7 +46,7 @@ var builtinComponentParamsGolden = map[string]string{ "manual": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"doc\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"doc\"]}, \"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}}, \"TitleChunker:NineInsectsFind\": {\"hierarchy\": 0, \"include_heading_content\": false, \"levels\": [[\"^#[^#]\", \"^##[^#]\", \"^###[^#]\", \"^####[^#]\"], [\"第[零一二三四五六七八九十百0-9]+(分?编|部分)\", \"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"第[零一二三四五六七八九十百0-9]+条\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\"], [\"第[0-9]+章\", \"第[0-9]+节\", \"[0-9]{1,2}[\\\\. 、]\", \"[0-9]{1,2}\\\\.[0-9]{1,2}($|[^a-zA-Z/%~.-])\", \"[0-9]{1,2}\\\\.[0-9]{1,2}\\\\.[0-9]{1,2}\"], [\"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"[零一二三四五六七八九十百]+[ 、]\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\", \"[\\\\(（][0-9]{,2}[\\\\)）]\"], [\"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)\", \"Chapter (I+V?|VI*|XI|IX|X)\", \"Section [0-9]+\", \"Article [0-9]+\"]], \"method\": \"group\"}, \"Tokenizer:FunnyBalloonsGrin\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", "one": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"doc\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"doc\"]}, \"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"html\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"htm\", \"html\"]}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"OneChunker:DryDrinksVisit\": {}, \"Tokenizer:FrankWeeksListen\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", "paper": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"pdf\": {\"enable_multi_column\": true, \"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}}, \"TitleChunker:SparklySchoolsTravel\": {\"hierarchy\": 0, \"include_heading_content\": false, \"levels\": [[\"^#[^#]\", \"^##[^#]\", \"^###[^#]\", \"^####[^#]\"], [\"第[零一二三四五六七八九十百0-9]+(分?编|部分)\", \"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"第[零一二三四五六七八九十百0-9]+条\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\"], [\"第[0-9]+章\", \"第[0-9]+节\", \"[0-9]{1,2}[\\\\. 、]\", \"[0-9]{1,2}\\\\.[0-9]{1,2}($|[^a-zA-Z/%~.-])\", \"[0-9]{1,2}\\\\.[0-9]{1,2}\\\\.[0-9]{1,2}\"], [\"第[零一二三四五六七八九十百0-9]+章\", \"第[零一二三四五六七八九十百0-9]+节\", \"[零一二三四五六七八九十百]+[ 、]\", \"[\\\\(（][零一二三四五六七八九十百]+[\\\\)）]\", \"[\\\\(（][0-9]{,2}[\\\\)）]\"], [\"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)\", \"Chapter (I+V?|VI*|XI|IX|X)\", \"Section [0-9]+\", \"Article [0-9]+\"]], \"method\": \"group\"}, \"Tokenizer:GreatCarsWash\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", - "picture": "{\"File\": {}, \"Parser:ViewsCaptureLight\": {\"image\": {\"output_format\": \"text\", \"parse_method\": \"ocr\", \"preprocess\": [\"main_content\"], \"suffix\": [\"bmp\", \"gif\", \"jpeg\", \"jpg\", \"png\", \"svg\", \"tif\", \"tiff\", \"webp\"]}, \"video\": {\"output_format\": \"text\", \"preprocess\": [\"main_content\"], \"suffix\": [\"3gp\", \"3gpp\", \"avi\", \"flv\", \"mkv\", \"mov\", \"mp4\", \"mpeg\", \"mpg\", \"webm\", \"wmv\"]}}, \"TokenChunker:BrightColorsGlow\": {}, \"Tokenizer:SharpLensFocus\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", + "picture": "{\"File\": {}, \"Parser:ViewsCaptureLight\": {\"image\": {\"output_format\": \"json\", \"parse_method\": \"ocr\", \"preprocess\": [\"main_content\"], \"suffix\": [\"bmp\", \"gif\", \"jpeg\", \"jpg\", \"png\", \"svg\", \"tif\", \"tiff\", \"webp\"]}, \"video\": {\"output_format\": \"text\", \"preprocess\": [\"main_content\"], \"suffix\": [\"3gp\", \"3gpp\", \"avi\", \"flv\", \"mkv\", \"mov\", \"mp4\", \"mpeg\", \"mpg\", \"webm\", \"wmv\"]}}, \"TokenChunker:BrightColorsGlow\": {}, \"Tokenizer:SharpLensFocus\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", "presentation": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"slides\": {\"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pptx\", \"ppt\"]}}, \"Tokenizer:TallTreesDance\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"PresentationChunker:HappyHillsGlow\": {}, \"Extractor:AutoExtractDefault\": {\"field_name\": \"\", \"auto_keywords\": 0, \"auto_questions\": 0, \"llm_id\": \"\", \"auto_tags\": 0}}", "qa": "{\"File\": {}, \"Parser:HipSignsRhyme\": {\"docx\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"markdown\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"md\", \"markdown\", \"mdx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": false, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"spreadsheet\": {\"flatten_media_to_text\": false, \"output_format\": \"html\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"xls\", \"xlsx\", \"csv\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"Tokenizer:ColdCloudsDream\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}, \"QAChunker:TidyCloudsThink\": {}}", "resume": "{\"Extractor:ThreeDrinksAct\": {\"field_name\": \"metadata\", \"frequencyPenaltyEnabled\": true, \"frequency_penalty\": 0.7, \"llm_id\": \"THUDM/GLM-4.1V-9B-Thinking@SILICONFLOW\", \"maxTokensEnabled\": false, \"max_tokens\": 256, \"presencePenaltyEnabled\": true, \"presence_penalty\": 0.4, \"prompts\": [{\"content\": \"Content: {TitleChunker:FlatMiceFix@chunks}\", \"role\": \"user\"}], \"sys_prompt\": \"Act as a precise resume metadata extractor. Extract stable, chunk-supported metadata from the provided resume content.\\n\\nRules:\\n1. Use only information explicitly stated in the content. Do not infer, guess, normalize, or add missing facts.\\n2. The input may be only one chunk of a resume. Extract only what this content directly supports.\\n3. Use only these field names:\\ncandidate_name, gender, phone, email, city, location, nationality, linkedin, github, website, highest_degree, degree_levels, school_names, majors, graduation_years, work_experience_years, current_job_title, job_titles, company_names, job_experience, industries, target_job_titles, target_locations, employment_types, skills, certificates, awards, summary_tags\\n4. Ignore detailed responsibilities, project descriptions, achievement narratives, self-evaluation, and other low-value local details.\\n5. Keep values in the same language as the source text whenever possible.\\n6. Remove duplicates and keep only concise, high-value metadata.\\n7. Return only fields that are explicitly supported by the content. Do not return empty or unsupported fields.\\n\\nField guidance:\\n- highest_degree: highest explicit degree level mentioned\\n- degree_levels: all explicit degree levels mentioned\\n- school_names: explicit school, college, or university names\\n- majors: explicit fields of study\\n- graduation_years: explicit graduation years only\\n- work_experience_years: only if explicitly stated\\n- current_job_title: only if explicitly current or most recent\\n- job_titles: explicit role titles\\n- company_names: explicit employer names\\n- job_experience: concise structured work entries explicitly supported by the content, preferably including title, company, and time information when available\\n- industries: explicit industry names only\\n- target_job_titles: explicit desired roles only\\n- target_locations: explicit desired work locations only\\n- skills: concise, core, search-useful skills explicitly mentioned\\n- certificates: explicit certificate names only\\n- awards: explicit award names only\\n- summary_tags: short, high-value tags strictly supported by the content\\n\\nReturn only the extracted metadata. Do not output explanatory text.\", \"temperature\": 0.1, \"temperatureEnabled\": true, \"tenant_llm_id\": 29, \"topPEnabled\": true, \"top_p\": 0.3, \"auto_keywords\": 0, \"auto_questions\": 0, \"auto_tags\": 0}, \"File\": {}, \"Parser:HipSignsRhyme\": {\"docx\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"docx\"], \"vlm\": {}}, \"pdf\": {\"flatten_media_to_text\": true, \"output_format\": \"json\", \"parse_method\": \"DeepDOC\", \"preprocess\": [\"main_content\"], \"suffix\": [\"pdf\"], \"vlm\": {}}, \"text&code\": {\"output_format\": \"json\", \"preprocess\": [\"main_content\"], \"suffix\": [\"txt\"]}}, \"TitleChunker:FlatMiceFix\": {\"hierarchy\": 1, \"include_heading_content\": false, \"levels\": [[\"^\\\\s*(?i:(?:\\\\d+[\\\\.\\\\)]\\\\s*)?(?:EDUCATION|ACADEMIC\\\\s*BACKGROUND|ACADEMIC\\\\s*HISTORY|EDUCATIONAL\\\\s*BACKGROUND|RELEVANT\\\\s*COURSEWORK|COURSEWORK|EXPERIENCE|WORK\\\\s*EXPERIENCE|PROFESSIONAL\\\\s*EXPERIENCE|RELEVANT\\\\s*EXPERIENCE|EMPLOYMENT\\\\s*HISTORY|CAREER\\\\s*HISTORY|INTERNSHIP\\\\s*EXPERIENCE|PROJECTS|PROJECT\\\\s*EXPERIENCE|ACADEMIC\\\\s*PROJECTS|PROFESSIONAL\\\\s*PROJECTS|SKILLS|TECHNICAL\\\\s*SKILLS|CORE\\\\s*COMPETENCIES|COMPETENCIES|QUALIFICATIONS|SUMMARY\\\\s*OF\\\\s*QUALIFICATIONS|CERTIFICATIONS|LICENSES|CERTIFICATES|AWARDS|HONORS|HONOURS|ACHIEVEMENTS|PUBLICATIONS|RESEARCH|RESEARCH\\\\s*EXPERIENCE|LEADERSHIP|LEADERSHIP\\\\s*EXPERIENCE|ACTIVITIES|EXTRACURRICULAR\\\\s*ACTIVITIES|ACTIVITIES\\\\s*(?:&|AND)\\\\s*SKILLS|INVOLVEMENT|CAMPUS\\\\s*INVOLVEMENT|VOLUNTEER\\\\s*EXPERIENCE|VOLUNTEERING|COMMUNITY\\\\s*SERVICE|LANGUAGES|INTERESTS|HOBBIES|PROFILE|PROFESSIONAL\\\\s*PROFILE|SUMMARY|PROFESSIONAL\\\\s*SUMMARY|CAREER\\\\s*SUMMARY|OBJECTIVE|CAREER\\\\s*OBJECTIVE|PERSONAL\\\\s*INFORMATION|CONTACT\\\\s*INFORMATION|ADDITIONAL\\\\s*INFORMATION|TRAINING))\\\\s*[:：]?\\\\s*$\"], [\"^\\\\s*(?:\\\\d+[\\\\.、\\\\)]\\\\s*)?(?:教育背景|教育经历|学历背景|学术背景|技术背景|工作经历|工作经验|实习经历|项目经历|项目经验|科研经历|研究经历|校园经历|实践经历|专业经历|职业经历|技能|专业技能|技能特长|核心技能|技术栈|个人技能|工作技能|职业技能|技能与评价|技能与自我评价|工作技能与自我评价|职业技能与自我评价|证书|资格证书|职业资格|资质证书|获奖情况|获奖经历|荣誉|荣誉奖项|奖项|科研成果|论文发表|发表论文|领导经历|学生工作|校园活动|社团经历|活动经历|志愿经历|志愿服务|社会实践|语言能力|语言|自我评价|个人评价|自我总结|个人总结|个人优势|个人简介|个人信息|基本信息|联系方式|求职意向|应聘意向|职业目标|求职目标|兴趣爱好|兴趣特长|培训经历|其他信息|附加信息)\\\\s*[:：]?\\\\s*$\"]], \"method\": \"hierarchy\"}, \"Tokenizer:KindHandsWin\": {\"fields\": \"text\", \"filename_embd_weight\": 0.1, \"search_method\": [\"embedding\", \"full_text\"]}}", diff --git a/internal/ingestion/task/pipeline_real_integration_test.go b/internal/ingestion/task/pipeline_real_integration_test.go index 5b8178aad3..d474b7d245 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,27 +398,42 @@ 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 in-memory sqlite database and migrates +// the tables the real-pipeline contract tests seed and read. It does not touch +// the filesystem or any external MySQL. The connection pool is pinned to a +// single connection (MaxOpenConns(1)) so the in-memory database is shared +// across all statements — without that, gorm's default pool would hand each +// statement a separate connection, and ":memory:" is per-connection (so seeds +// on one connection would be invisible to reads on another). +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{ + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ Logger: gormlogger.Default.LogMode(gormlogger.Silent), }) if err != nil { - t.Fatalf("connect real mysql: %v", err) + t.Fatalf("open in-memory 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) } return db } @@ -731,3 +707,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 +} diff --git a/internal/utility/split.go b/internal/utility/split.go index d3163260e0..0334fa09d3 100644 --- a/internal/utility/split.go +++ b/internal/utility/split.go @@ -27,14 +27,15 @@ import ( // task_executor.run_dataflow:879 re.split(r"[,，;；、\r\n]+", keywords). var keywordsSplitRE = regexp.MustCompile(`[,，;；、\r\n]+`) -// nonEmpty drops empty strings from parts and returns nil if none remain. It is -// the shared tail of SplitKeywords and SplitQuestions: split by whatever -// delimiter, then prune empties and collapse an all-empty result to nil so a -// _kwd array is absent (nil) rather than [""]. +// nonEmpty drops empty and whitespace-only strings from parts and returns nil +// if none remain. It is the shared tail of SplitKeywords and SplitQuestions: +// split by whatever delimiter, then prune blanks and collapse an all-blank +// result to nil so a _kwd array is absent (nil) rather than [""]. +// Mirrors Python task_executor's `if k.strip()` filter. func nonEmpty(parts []string) []string { out := make([]string, 0, len(parts)) for _, p := range parts { - if p != "" { + if strings.TrimSpace(p) != "" { out = append(out, p) } } diff --git a/internal/utility/split_test.go b/internal/utility/split_test.go index 6225edf666..3fb4acd29c 100644 --- a/internal/utility/split_test.go +++ b/internal/utility/split_test.go @@ -16,7 +16,10 @@ package utility -import "testing" +import ( + "reflect" + "testing" +) // SplitKeywords - Python task_executor.run_dataflow:879 // re.split(r"[,，;；、\r\n]+", keywords) with empty filtering. @@ -49,6 +52,28 @@ func TestSplitKeywords_FiltersEmptyStrings(t *testing.T) { } } +// Python task_executor.run_dataflow filters split parts with `if k.strip()`, +// so whitespace-only segments (a lone space between commas, a tab, etc.) must +// be dropped — not kept as " ". Input "a, ,b" must yield ["a","b"], matching +// Python, not ["a"," ","b"]. +func TestSplitKeywords_FiltersWhitespaceOnly(t *testing.T) { + cases := []struct { + in string + want []string + }{ + {"a, ,b", []string{"a", "b"}}, + {"a，\t,b", []string{"a", "b"}}, + {" , , ", nil}, + {"kw1, ,kw2", []string{"kw1", "kw2"}}, + } + for _, c := range cases { + got := SplitKeywords(c.in) + if !reflect.DeepEqual(got, c.want) { + t.Errorf("SplitKeywords(%q) = %v, want %v", c.in, got, c.want) + } + } +} + func TestSplitKeywords_Empty(t *testing.T) { result := SplitKeywords("") if len(result) != 0 {