Skip to content

Commit 9b0719f

Browse files
authored
fix: Go ingestion migration batch 5 (Parser 1.1/1.7/2.11, Chunker 1.7/1.8/2.6/2.7, Tokenizer 6x fixes) (#17419)
## Summary Continuation of the Python→Go ingestion pipeline migration (File → Parser → Chunker → Extractor → Tokenizer). Fixes cover Parser, Chunker, and Tokenizer gaps identified. Fix page number (0-indexed and 1-index mixed before fix; use 1-indexed after fix) and chunk order issues. ### Parser - **Slides TCADP (1.7):** `pptx_tcadp.go` + TCADP branch in `pptx_parser.go`/`ppt_parser.go` — PowerPoint files now support `parse_method="tcadp"` via the TCADP cloud service, matching the spreadsheet-family TCADP pattern. PPT containers pass `"PPT"` as fileType (not hardcoded `"PPTX"`). - **Audio default output_format (2.11):** `defaultSetups()` audio default changed from `"text"` to `"json"`, aligning with Python `parser.py:232` and `AllowedOutputFormat["audio"]={"json"}`. - **PDF VLM enhancement (1.1):** `maybeDispatchPDFVisionEnhancement` in `pdf_vision_dispatch.go` enriches image/table items with IMAGE2TEXT model descriptions after PDF parsing, mirroring Python `enhance_media_sections_with_vision`. Semaphore fix: acquire before goroutine start to prevent unbounded goroutine creation. - **json family (2.3):** reclassified as Keep Go — `json_parser.go` is a functional enhancement, not a parity gap. - **page number:** changed from "mixed use of 1-indexed & 0-indexed" to "1-indexed" ### Chunker - **BULLET_PATTERN fallback (1.7):** 4th-level fallback in `resolveTitleLevels` (`title.go`) detects bullet/numbered-list patterns (Chinese legal, numbering, English) when outline + regex levels produce only bodyLevel. Guarded by `allBodyLevel` to never override existing structure. - **Tag/One chunker fields (1.8):** `tag.go` sets `TopInt` from source row index; `one.go` preserves `Positions`/`PDFPositions` from source items. TSV multi-line RowNum fix: tracks `contentStart` for correct row attribution. - **Overlapped_percent normalization (2.6):** `NormalizeOverlappedPercent` in `schema/chunker.go` mirrors Python `common/float_utils.py:50-58` — accepts `[0,1)` fraction or `[0,90]` percent, normalizes to canonical `[0,90]`. - **Paragraph splitting (2.7):** aligned to Python flow `naive_merge` — `CRLF` normalization, `splitKeepingDelimiter` preserves sentence delimiters, single-section merge with token-budget-governed chunking. - **chunk order:** sort by reading order ### Tokenizer - **Phantom chunk filtering (Omission 2):** `isPhantomChunk` + filter loop in `chunksFromTokenizerUpstream` skips zero-value ChunkDocs. - **Batch size env var (Omission 3):** `embeddingBatchSize()` reads `TOKENIZER_EMBEDDING_BATCH_SIZE`, defaults to 16. - **Summary empty check (Diff 5):** `TrimSpace(s) != ""` → `s != ""`, matching Python truthy check. - **chunk_order_int all paths (Diff 8):** set unconditionally before full_text/embedding branching. - **Timeout default (Diff 10):** `600s` → `60s`, matching Python `@timeout(60)`. - **Small maxTokens truncation (Diff 14):** `truncateForEmbedding` returns `""` when `maxTokens <= 10`, matching Python. ### Code review fixes - Semaphore acquire moved before goroutine in `pdf_vision_dispatch.go` (concurrency control) - Context propagation in `pptx_tcadp.go` (cancellation support) - Test resolver leak fix in `media_dispatch_test.go` (defer restore) - Migration history comments removed per AGENTS.md ## Test plan ``` bash build.sh --test ./internal/parser/parser/... ./internal/ingestion/component/... ``` ## Notes - Migration diff tracking: `docs/migration_python_go_diff.md` - Remaining gaps: Extractor component only (21 items)
1 parent 3268938 commit 9b0719f

41 files changed

Lines changed: 1542 additions & 330 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.
Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
//go:build cgo
2+
3+
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
4+
//
5+
// Licensed under the Apache License, Version 2.0 (the "License");
6+
// you may not use this file except in compliance with the License.
7+
//
8+
// http://www.apache.org/licenses/LICENSE-2.0
9+
//
10+
// Unless required by applicable law or agreed to in writing, software
11+
// distributed under the License is distributed on an "AS IS" BASIS,
12+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
// See the License for the specific language governing permissions and
14+
// limitations under the License.
15+
16+
package chunker
17+
18+
import (
19+
"context"
20+
"strings"
21+
"testing"
22+
)
23+
24+
// TestCropTitleChunks_CropsPositionedChunk pins residual A: chunks
25+
// produced by the Title/Group/Hierarchy chunkers must be cropped on demand
26+
// (mirroring the TokenChunker JSON path). The chunk shape mirrors the real
27+
// Group/Hierarchy output: it carries doc_type_kwd + positions but NO ck_type
28+
// (buildChunksFromRecordGroups never sets ck_type). cropTitleChunks must
29+
// derive ck_type from doc_type_kwd so needsCrop (pdfcrop_cgo.go:151) fires.
30+
// A text chunk with positions gets a rendered preview; a chunk without
31+
// positions is left untouched; and the derived ck_type must not leak into
32+
// the returned chunk shape.
33+
func TestCropTitleChunks_CropsPositionedChunk(t *testing.T) {
34+
ctx := context.Background()
35+
pos := jsonPositions(t, []float64{1, 10, 100, 10, 100})
36+
in := []map[string]any{
37+
{"text": "page text", "doc_type_kwd": "text", "positions": pos},
38+
{"text": "no coords", "doc_type_kwd": "text"},
39+
}
40+
41+
got := cropTitleChunks(ctx, mockCropEngine{}, in)
42+
if len(got) != 2 {
43+
t.Fatalf("len = %d, want 2", len(got))
44+
}
45+
img, _ := got[0]["image"].(string)
46+
if !strings.HasPrefix(img, "data:image/png;base64,") {
47+
t.Errorf("positioned text chunk: image = %q, want data:image/png;base64, prefix", img)
48+
}
49+
if got[1]["image"] != nil {
50+
t.Errorf("chunk without positions should not be cropped, got %v", got[1]["image"])
51+
}
52+
if _, ok := got[0]["ck_type"]; ok {
53+
t.Errorf("derived ck_type must not leak into output chunk: %v", got[0])
54+
}
55+
}
56+
57+
// TestCropTitleChunks_NilEnginePassthrough asserts the helper is a no-op
58+
// when no PDF engine is available (best-effort contract).
59+
func TestCropTitleChunks_NilEnginePassthrough(t *testing.T) {
60+
in := []map[string]any{{"text": "hello"}}
61+
got := cropTitleChunks(context.Background(), nil, in)
62+
if len(got) != 1 || got[0]["text"] != "hello" {
63+
t.Fatalf("nil engine should pass through unchanged: %v", got)
64+
}
65+
}

internal/ingestion/component/chunker/group.go

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import (
3636
"context"
3737
"encoding/json"
3838
"fmt"
39+
"log/slog"
3940
"regexp"
4041
"sort"
4142
"strings"
@@ -108,7 +109,7 @@ func buildSectionIDs(levels []int, targetLevel int) []int {
108109
// supplied inputs. Detected headings + adjacent merges happen in two
109110
// goroutines (heading detection sequential, then a fan-out over
110111
// record-buckets for the merge pass).
111-
func invokeGroup(_ context.Context, inputs map[string]any, p *titleChunkerParam) (map[string]any, error) {
112+
func invokeGroup(parentCtx context.Context, db *gorm.DB, inputs map[string]any, p *titleChunkerParam) (map[string]any, error) {
112113
records := extractLineRecords(inputs)
113114
common.Debug("chunker stage",
114115
zap.String("component", "Chunker"),
@@ -161,6 +162,21 @@ func invokeGroup(_ context.Context, inputs map[string]any, p *titleChunkerParam)
161162
zap.Int("chunks", len(chunks)),
162163
zap.Bool("plain_text", isPlainTextFormat(inputs)),
163164
)
165+
166+
// On-demand PDF preview cropping for image/table/text chunks,
167+
// mirroring the TokenChunker JSON path (token.go:513). Best-effort:
168+
// a missing or unreadable PDF simply skips cropping.
169+
if upstream, uErr := decodeChunkerFromUpstream(inputs); uErr == nil {
170+
engine, eErr := newPDFEngineFromUpstream(parentCtx, db, upstream)
171+
if eErr != nil {
172+
slog.Warn("GroupTitleChunker: could not open PDF for on-demand cropping", "err", eErr)
173+
}
174+
if engine != nil {
175+
defer engine.Close()
176+
chunks = cropTitleChunks(parentCtx, engine, chunks)
177+
}
178+
}
179+
164180
if len(chunks) == 0 {
165181
return emptyOutputs(), nil
166182
}
@@ -484,7 +500,7 @@ func (c *GroupTitleChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, in
484500
"_ERROR": "GroupTitleChunker: missing required upstream field \"name\"",
485501
}, nil
486502
}
487-
return invokeGroup(ctx, withName(inputs, name), &c.param)
503+
return invokeGroup(ctx, db, withName(inputs, name), &c.param)
488504
}
489505

490506
// init registers GroupTitleChunker under CategoryIngestion.

internal/ingestion/component/chunker/hierarchy.go

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ package chunker
4141
import (
4242
"context"
4343
"fmt"
44+
"log/slog"
4445

4546
"go.uber.org/zap"
4647
"gorm.io/gorm"
@@ -139,7 +140,7 @@ func (n *chunkNode) getPaths(paths *[][]int, titles []int, depth int, includeHea
139140
}
140141

141142
// invokeHierarchy runs the HierarchyTitleChunker strategy.
142-
func invokeHierarchy(_ context.Context, inputs map[string]any, p *titleChunkerParam) (map[string]any, error) {
143+
func invokeHierarchy(parentCtx context.Context, db *gorm.DB, inputs map[string]any, p *titleChunkerParam) (map[string]any, error) {
143144
records := extractLineRecords(inputs)
144145
common.Debug("chunker stage",
145146
zap.String("component", "Chunker"),
@@ -246,6 +247,21 @@ func invokeHierarchy(_ context.Context, inputs map[string]any, p *titleChunkerPa
246247
zap.Int("chunks", len(chunks)),
247248
zap.Bool("plain_text", isPlainTextFormat(inputs)),
248249
)
250+
251+
// On-demand PDF preview cropping for image/table/text chunks,
252+
// mirroring the TokenChunker JSON path (token.go:513). Best-effort:
253+
// a missing or unreadable PDF simply skips cropping.
254+
if upstream, uErr := decodeChunkerFromUpstream(inputs); uErr == nil {
255+
engine, eErr := newPDFEngineFromUpstream(parentCtx, db, upstream)
256+
if eErr != nil {
257+
slog.Warn("HierarchyTitleChunker: could not open PDF for on-demand cropping", "err", eErr)
258+
}
259+
if engine != nil {
260+
defer engine.Close()
261+
chunks = cropTitleChunks(parentCtx, engine, chunks)
262+
}
263+
}
264+
249265
if len(chunks) == 0 {
250266
return emptyOutputs(), nil
251267
}
@@ -302,7 +318,7 @@ func (c *HierarchyTitleChunkerComponent) Invoke(ctx context.Context, db *gorm.DB
302318
"_ERROR": "HierarchyTitleChunker: missing required upstream field \"name\"",
303319
}, nil
304320
}
305-
return invokeHierarchy(ctx, withName(inputs, name), &c.param)
321+
return invokeHierarchy(ctx, db, withName(inputs, name), &c.param)
306322
}
307323

308324
// init registers HierarchyTitleChunker under CategoryIngestion.

internal/ingestion/component/chunker/one.go

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -141,10 +141,12 @@ func emitOneFromItems(items, chunks []schema.ChunkDoc) map[string]any {
141141
return emptyOutputs()
142142
}
143143
out := schema.ChunkDoc{
144-
Text: text,
145-
DocType: docType,
146-
CKType: docType,
147-
Image: it.Image,
144+
Text: text,
145+
DocType: docType,
146+
CKType: docType,
147+
Image: it.Image,
148+
Positions: it.Positions,
149+
PDFPositions: it.PDFPositions,
148150
}
149151
return chunkOutputs([]schema.ChunkDoc{out})
150152
}
@@ -164,6 +166,11 @@ func emitOneFromItems(items, chunks []schema.ChunkDoc) map[string]any {
164166
return emptyOutputs()
165167
}
166168
out := schema.ChunkDoc{Text: merged, DocType: "text", CKType: "text"}
169+
// Multi-item merge produces a single text-only chunk mirroring Python
170+
// one.py:166-168. Per-item Positions/PDFPositions are intentionally
171+
// not carried — merging coordinates from different source items would
172+
// produce meaningless composite geometry, and the downstream
173+
// processChunkPositions would map them to incorrect pages.
167174
if img != "" {
168175
out.Image = img
169176
}

internal/ingestion/component/chunker/one_test.go

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,3 +98,46 @@ func TestOneChunker_JSONMultipleMerges(t *testing.T) {
9898
t.Errorf("image = %q, want first available media context", got)
9999
}
100100
}
101+
102+
// TestOneChunker_PreservesPositions verifies that when a single upstream
103+
// item carries PDF coordinates, the OneChunker preserves both Positions
104+
// and PDFPositions (with their coordinate values) on the output chunk.
105+
func TestOneChunker_PreservesPositions(t *testing.T) {
106+
chunks := oneChunksOf(t, map[string]any{
107+
"name": "page.pdf",
108+
"output_format": "json",
109+
"json": []map[string]any{
110+
{
111+
"text": "page text",
112+
"positions": []any{[]any{10.0, 20.0, 30.0, 40.0}},
113+
"_pdf_positions": []any{[]any{1.0, 2.0, 3.0, 4.0, 5.0}},
114+
},
115+
},
116+
})
117+
if len(chunks) != 1 {
118+
t.Fatalf("want 1 chunk, got %d", len(chunks))
119+
}
120+
assertCoordTuple(t, "positions", chunks[0]["positions"], []float64{10.0, 20.0, 30.0, 40.0})
121+
assertCoordTuple(t, "_pdf_positions", chunks[0]["_pdf_positions"], []float64{1.0, 2.0, 3.0, 4.0, 5.0})
122+
}
123+
124+
// assertCoordTuple verifies a positions/_pdf_positions field round-tripped
125+
// as a [][]float64 with the expected single-row coordinate tuple.
126+
func assertCoordTuple(t *testing.T, key string, got any, want []float64) {
127+
t.Helper()
128+
rows, ok := got.([][]float64)
129+
if !ok {
130+
t.Fatalf("%s = %v, want [][]float64 (got %T)", key, got, got)
131+
}
132+
if len(rows) != 1 {
133+
t.Fatalf("%s has %d rows, want 1", key, len(rows))
134+
}
135+
if len(rows[0]) != len(want) {
136+
t.Fatalf("%s[0] = %v, want %v (len %d vs %d)", key, rows[0], want, len(rows[0]), len(want))
137+
}
138+
for i, w := range want {
139+
if rows[0][i] != w {
140+
t.Errorf("%s[0][%d] = %v, want %v", key, i, rows[0][i], w)
141+
}
142+
}
143+
}

internal/ingestion/component/chunker/qa.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ func (c *QAChunkerComponent) invoke(_ context.Context, inputs map[string]any) (m
101101

102102
qPrefix, aPrefix := "问题:", "回答:"
103103
// Python qa.py defaults to Chinese when no language is supplied; only
104-
// an explicit "english" switches to English prefixes (diff Chunker-2.13).
104+
// an explicit "english" switches to English prefixes
105105
eng := strings.EqualFold(c.param.Lang, "english")
106106
if eng {
107107
qPrefix, aPrefix = "Question: ", "Answer: "
@@ -137,7 +137,7 @@ func (c *QAChunkerComponent) invoke(_ context.Context, inputs map[string]any) (m
137137
ContentLtks: contentLTKS,
138138
ContentSmLtks: contentSMLTKS,
139139
}
140-
// Restore metadata lost before diff Chunker-1.8: top_int (row
140+
//
141141
// index), image id + coordinates carried from the source item.
142142
if pair.RowNum >= 0 {
143143
chunk.TopInt = []int{pair.RowNum}
@@ -172,14 +172,14 @@ type qaPair struct {
172172
RowNum int
173173
// Image and positions are carried from the upstream item so the QA
174174
// chunk preserves metadata that Python sets via beAdocPdf/beAdocDocx
175-
// (diff Chunker-1.8).
175+
//
176176
Image string
177177
PDFPositions json.RawMessage
178178
Positions json.RawMessage
179179
}
180180

181181
// rmQAPrefixRe mirrors Python qa.py:241 `[\t:: ]+` — one-or-more separator
182-
// chars, so "Q:: answer" is fully stripped (diff Chunker-2.12).
182+
// chars, so "Q:: answer" is fully stripped
183183
var rmQAPrefixRe = regexp.MustCompile(`(?i)^(问题|答案|回答|user|assistant|Q|A|Question|Answer|问|答)[\t:: ]+`)
184184

185185
func rmQAPrefix(txt string) string {
@@ -424,7 +424,7 @@ func extractQAJSON(items []schema.ChunkDoc) []qaPair {
424424
}
425425
tmp := extractQAText(txt)
426426
// Preserve the source item's image id and coordinates on each
427-
// extracted pair (diff Chunker-1.8).
427+
// extracted pair
428428
for _, p := range tmp {
429429
p.Image = item.Image
430430
p.PDFPositions = item.PDFPositions

internal/ingestion/component/chunker/qa_batch2_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ func qaInvoke(t *testing.T, inputs map[string]any) []map[string]any {
4141
return chunks
4242
}
4343

44-
// TestQAChunker_DefaultLangIsChinese exercises migration diff Chunker-2.13:
44+
// TestQAChunker_DefaultLangIsChinese exercises migration :
4545
// when no language is supplied, Python defaults to Chinese prefixes
4646
// ("问题:"/"回答:"); the legacy Go code defaulted to English.
4747
func TestQAChunker_DefaultLangIsChinese(t *testing.T) {
@@ -79,7 +79,7 @@ func TestRmQAPrefixStripsMultipleSeparators(t *testing.T) {
7979
}
8080
}
8181

82-
// TestQAChunker_SetsTopInt exercises migration diff Chunker-1.8 (top_int):
82+
// TestQAChunker_SetsTopInt exercises migration (top_int):
8383
// each QA chunk must carry the source row index in `top_int`, matching
8484
// Python beAdoc(..., row_num=i).
8585
func TestQAChunker_SetsTopInt(t *testing.T) {

internal/ingestion/component/chunker/qa_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,7 @@ func TestQAChunker_PrefixSpaceSeparatorStrips(t *testing.T) {
213213
}
214214
cww, _ := chunks[0]["content_with_weight"].(string)
215215
// Python qa.py:241 uses `[\t:: ]+`, so a space is a valid separator:
216-
// a leading "A"/"Q" followed by a space is stripped (diff Chunker-2.12).
216+
// a leading "A"/"Q" followed by a space is stripped. .
217217
if cww != "Question: language model is useful\tAnswer: How does it work" {
218218
t.Fatalf("space-separator prefix not stripped: %q", cww)
219219
}

internal/ingestion/component/chunker/tag.go

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ func (c *TagChunkerComponent) invoke(_ context.Context, inputs map[string]any) (
125125
ContentLtks: contentLTKS,
126126
ContentSmLtks: contentSMLTKS,
127127
TagKwd: splitTagKwd(pair.Tags),
128+
TopInt: []int{pair.RowNum},
128129
}
129130
chunks = append(chunks, chunk)
130131
}
@@ -133,9 +134,13 @@ func (c *TagChunkerComponent) invoke(_ context.Context, inputs map[string]any) (
133134
}
134135

135136
// tagPair is a (content, tags) row extracted from the upstream payload.
137+
// RowNum is the 0-based record START line (the first physical source
138+
// line of the record), used by every extractor consistently and mapped
139+
// to Python's top_int in beAdoc(tag.py:33).
136140
type tagPair struct {
137141
Content string
138142
Tags string
143+
RowNum int
139144
}
140145

141146
// extractTagText ports tag.py:60-89 (txt) and tag.py:91-113 (csv).
@@ -157,18 +162,23 @@ func extractTagText(text string) []tagPair {
157162
func extractTagTextTab(lines []string) []tagPair {
158163
var pairs []tagPair
159164
content := ""
160-
for _, line := range lines {
165+
contentStart := -1
166+
for i, line := range lines {
161167
if strings.TrimSpace(line) == "" {
162168
continue
163169
}
170+
if contentStart < 0 {
171+
contentStart = i
172+
}
164173
parts := strings.Split(line, "\t")
165174
if len(parts) != 2 {
166175
content += "\n" + line
167176
continue
168177
}
169178
content += "\n" + parts[0]
170-
pairs = append(pairs, tagPair{Content: content, Tags: parts[1]})
179+
pairs = append(pairs, tagPair{Content: content, Tags: parts[1], RowNum: contentStart})
171180
content = ""
181+
contentStart = -1
172182
}
173183
return pairs
174184
}
@@ -202,6 +212,7 @@ func extractTagTextCSV(text string, lines []string) []tagPair {
202212
for curLine < len(lineStarts) && lineStarts[curLine] < endOff {
203213
curLine++
204214
}
215+
startLine := prevLine
205216
raw := strings.Join(lines[prevLine:curLine], "\n")
206217
prevLine = curLine
207218

@@ -210,7 +221,10 @@ func extractTagTextCSV(text string, lines []string) []tagPair {
210221
continue
211222
}
212223
content += "\n" + record[0]
213-
pairs = append(pairs, tagPair{Content: content, Tags: record[1]})
224+
// RowNum is the 0-based record START line, kept consistent with
225+
// extractTagTextTab (line i) and extractTagTable (<tr> i) so every
226+
// tag-pair source uses the same row-index convention.
227+
pairs = append(pairs, tagPair{Content: content, Tags: record[1], RowNum: startLine})
214228
content = ""
215229
}
216230
return pairs
@@ -226,7 +240,7 @@ func extractTagTable(htmlStr string) []tagPair {
226240
}
227241
rows := htmlTR.FindAllStringSubmatch(htmlStr, -1)
228242
pairs := make([]tagPair, 0, len(rows))
229-
for _, row := range rows {
243+
for i, row := range rows {
230244
cells := htmlTD.FindAllStringSubmatch(row[1], -1)
231245
var texts []string
232246
for _, cell := range cells {
@@ -237,7 +251,7 @@ func extractTagTable(htmlStr string) []tagPair {
237251
}
238252
}
239253
if len(texts) >= 2 {
240-
pairs = append(pairs, tagPair{Content: texts[0], Tags: texts[1]})
254+
pairs = append(pairs, tagPair{Content: texts[0], Tags: texts[1], RowNum: i})
241255
}
242256
}
243257
return pairs

0 commit comments

Comments
 (0)