2026-07-15 14:53:16 +08:00
|
|
|
|
//
|
|
|
|
|
|
// 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 utility
|
|
|
|
|
|
|
Fix(go): align ingestion pipeline with Python (parser/media dispatch + PDF coordinate chain + Chunker) (#17349)
## 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.
2026-07-24 21:06:38 +08:00
|
|
|
|
import (
|
|
|
|
|
|
"reflect"
|
|
|
|
|
|
"testing"
|
|
|
|
|
|
)
|
2026-07-15 14:53:16 +08:00
|
|
|
|
|
|
|
|
|
|
// SplitKeywords - Python task_executor.run_dataflow:879
|
|
|
|
|
|
// re.split(r"[,,;;、\r\n]+", keywords) with empty filtering.
|
|
|
|
|
|
|
|
|
|
|
|
func TestSplitKeywords_Comma(t *testing.T) {
|
|
|
|
|
|
result := SplitKeywords("kw1,kw2,kw3")
|
|
|
|
|
|
if len(result) != 3 {
|
|
|
|
|
|
t.Fatalf("len = %d, want 3", len(result))
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func TestSplitKeywords_ChineseComma(t *testing.T) {
|
|
|
|
|
|
result := SplitKeywords("kw1,kw2,kw3")
|
|
|
|
|
|
if len(result) != 3 {
|
|
|
|
|
|
t.Fatalf("len = %d, want 3", len(result))
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func TestSplitKeywords_MixedDelimiters(t *testing.T) {
|
|
|
|
|
|
result := SplitKeywords("kw1,kw2;kw3")
|
|
|
|
|
|
if len(result) != 3 {
|
|
|
|
|
|
t.Fatalf("len = %d, want 3, got %v", len(result), result)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func TestSplitKeywords_FiltersEmptyStrings(t *testing.T) {
|
|
|
|
|
|
result := SplitKeywords("kw1,,kw2")
|
|
|
|
|
|
if len(result) != 2 {
|
|
|
|
|
|
t.Errorf("empty strings should be filtered: got %v", result)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
Fix(go): align ingestion pipeline with Python (parser/media dispatch + PDF coordinate chain + Chunker) (#17349)
## 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.
2026-07-24 21:06:38 +08:00
|
|
|
|
// 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)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-07-15 14:53:16 +08:00
|
|
|
|
func TestSplitKeywords_Empty(t *testing.T) {
|
|
|
|
|
|
result := SplitKeywords("")
|
|
|
|
|
|
if len(result) != 0 {
|
|
|
|
|
|
t.Errorf("got %v, want empty", result)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// SplitQuestions - splits by newline only, dropping empty lines.
|
|
|
|
|
|
|
|
|
|
|
|
func TestSplitQuestions_MultipleLines(t *testing.T) {
|
|
|
|
|
|
result := SplitQuestions("Q1\nQ2\nQ3")
|
|
|
|
|
|
if len(result) != 3 || result[0] != "Q1" || result[2] != "Q3" {
|
|
|
|
|
|
t.Fatalf("got %v, want [Q1 Q2 Q3]", result)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func TestSplitQuestions_Empty(t *testing.T) {
|
|
|
|
|
|
if result := SplitQuestions(""); result != nil {
|
|
|
|
|
|
t.Errorf("got %v, want nil", result)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func TestSplitQuestions_TrailingNewlineDropped(t *testing.T) {
|
|
|
|
|
|
// "Q1\n" would yield ["Q1", ""] with a plain split; the trailing empty
|
|
|
|
|
|
// must be filtered so the index is not seeded with a spurious [""].
|
|
|
|
|
|
result := SplitQuestions("Q1\n")
|
|
|
|
|
|
if len(result) != 1 || result[0] != "Q1" {
|
|
|
|
|
|
t.Fatalf("got %v, want [Q1]", result)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func TestSplitQuestions_BlankLinesDropped(t *testing.T) {
|
|
|
|
|
|
result := SplitQuestions("Q1\n\nQ2")
|
|
|
|
|
|
if len(result) != 2 || result[0] != "Q1" || result[1] != "Q2" {
|
|
|
|
|
|
t.Fatalf("got %v, want [Q1 Q2]", result)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func TestSplitQuestions_PreservesCommas(t *testing.T) {
|
|
|
|
|
|
// Unlike SplitKeywords, questions must NOT split on commas.
|
|
|
|
|
|
result := SplitQuestions("请问A,B和C有何区别?")
|
|
|
|
|
|
if len(result) != 1 || result[0] != "请问A,B和C有何区别?" {
|
|
|
|
|
|
t.Fatalf("got %v, want single element preserving comma", result)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|