mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-30 20:49:21 +08:00
## 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)
783 lines
24 KiB
Go
783 lines
24 KiB
Go
//
|
||
// 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.
|
||
//
|
||
|
||
// SCOPE (honest) for title.go:
|
||
//
|
||
// - TitleChunker is the dispatcher for the three Go chunks: this
|
||
// file holds the TitleChunker variant itself, which mirrors
|
||
// python `rag/flow/chunker/title_chunker/title_chunker.py` — it
|
||
// holds no business logic of its own; the actual chunking
|
||
// happens in group.go / hierarchy.go depending on the
|
||
// `method` param.
|
||
//
|
||
// TitleChunker.Invoke dispatches synchronously to group.go or
|
||
// hierarchy.go.
|
||
//
|
||
// - HEADING DETECTION PARITY:
|
||
// The Python side uses two heading-detection strategies in
|
||
// `common.py:resolve_title_levels`:
|
||
// (1) PDF outlines (extract_pdf_outlines, requires deepdoc/parser)
|
||
// (2) Regex families (the user's `levels` param) with a layout-hint
|
||
// fallback (layout field matches section/title/head).
|
||
// The Go port ships strategy (2) IN FULL, including the
|
||
// match_layout_level fallback ported from common.py (Gap C closed):
|
||
// a non-regex text record whose layout flags it as a
|
||
// section/title/head and whose text passes not_title is promoted to
|
||
// fallback_level = len(selected_group) + 1. Strategy (1) —
|
||
// PDF-outline detection — still requires deepdoc/parser binary
|
||
// access and remains the only parity gap.
|
||
//
|
||
// - The Go port has NO hardcoded BULLET_PATTERN fallback. Heading
|
||
// detection relies entirely on the user-supplied `levels` param
|
||
// (which templates carry as comprehensive regex families) paired
|
||
// with the layout-hint fallback. A PDF without matching levels
|
||
// produces BODY_LEVEL-only records — Python's BULLET_PATTERN-based
|
||
// tree_merge / hierarchical_merge would still find structure.
|
||
//
|
||
// - GROUP-TITLE and HIERARCHY-TITLE are separate Go files
|
||
// (`group.go`, `hierarchy.go`); they share the resolve_levels
|
||
// path with TitleChunker but their build_chunks logic differs.
|
||
package chunker
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"regexp"
|
||
"strings"
|
||
"unicode/utf8"
|
||
|
||
"ragflow/internal/agent/runtime"
|
||
"ragflow/internal/ingestion/component/globals"
|
||
"ragflow/internal/ingestion/component/schema"
|
||
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
const ComponentNameTitleChunker = "TitleChunker"
|
||
|
||
type titleChunkerParam struct {
|
||
schema.TitleChunkerParam
|
||
}
|
||
|
||
func (p *titleChunkerParam) Update(conf map[string]any) {
|
||
if conf == nil {
|
||
return
|
||
}
|
||
if v, ok := conf["method"].(string); ok {
|
||
p.TitleChunkerParam.Method = v
|
||
}
|
||
if v, ok := conf["levels"].([]any); ok {
|
||
p.TitleChunkerParam.Levels = parseLevels(v)
|
||
} else if v, ok := conf["levels"].([][]string); ok {
|
||
p.TitleChunkerParam.Levels = v
|
||
}
|
||
if v, ok := schema.NumericFromAny(conf["hierarchy"]); ok {
|
||
n := int(v)
|
||
p.TitleChunkerParam.Hierarchy = &n
|
||
}
|
||
if v, ok := conf["include_heading_content"].(bool); ok {
|
||
p.TitleChunkerParam.IncludeHeadingContent = v
|
||
}
|
||
if v, ok := conf["root_chunk_as_heading"].(bool); ok {
|
||
p.TitleChunkerParam.RootChunkAsHeading = v
|
||
}
|
||
}
|
||
|
||
// parseLevels accepts a [[string]] representation — the natural
|
||
// JSON form — and rebroadcasts the strings verbatim. Stricter
|
||
// validation lives in schema.TitleChunkerParam.Validate.
|
||
func parseLevels(in []any) [][]string {
|
||
out := make([][]string, 0, len(in))
|
||
for _, lvl := range in {
|
||
group, ok := lvl.([]any)
|
||
if !ok {
|
||
continue
|
||
}
|
||
row := make([]string, 0, len(group))
|
||
for _, pat := range group {
|
||
if s, ok := pat.(string); ok && s != "" {
|
||
row = append(row, s)
|
||
}
|
||
}
|
||
if len(row) > 0 {
|
||
out = append(out, row)
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
func defaultsTitle() titleChunkerParam {
|
||
return titleChunkerParam{TitleChunkerParam: schema.TitleChunkerParam{}.Defaults()}
|
||
}
|
||
|
||
// selectLevelGroup mirrors common.py:select_level_group. Returns the
|
||
// regex-list family with the highest hit count across the input
|
||
// lines.
|
||
func selectLevelGroup(lines []string, rawLevels [][]string) []string {
|
||
if len(rawLevels) == 0 {
|
||
return nil
|
||
}
|
||
hits := make([]int, len(rawLevels))
|
||
for i, group := range rawLevels {
|
||
for _, line := range lines {
|
||
text := trim(line)
|
||
if text == "" {
|
||
continue
|
||
}
|
||
// Mirror common.py:select_level_group — a bullet-looking
|
||
// line cannot count as a heading hit for any family.
|
||
if notBullet(text) {
|
||
continue
|
||
}
|
||
for _, pattern := range group {
|
||
if re := compileLevelPattern(pattern); re != nil {
|
||
if re.MatchString(text) {
|
||
hits[i]++
|
||
break
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
bestIdx := -1
|
||
best := 0
|
||
for i, h := range hits {
|
||
if h > best {
|
||
best = h
|
||
bestIdx = i
|
||
}
|
||
}
|
||
if bestIdx < 0 {
|
||
return nil
|
||
}
|
||
group := make([]string, 0, len(rawLevels[bestIdx]))
|
||
for _, p := range rawLevels[bestIdx] {
|
||
if p != "" {
|
||
group = append(group, p)
|
||
}
|
||
}
|
||
return group
|
||
}
|
||
|
||
// matchRegexLevel mirrors common.py:match_regex_level. Levels are
|
||
// 1-indexed; "" or no-match returns 0 (BODY).
|
||
func matchRegexLevel(text string, group []string) int {
|
||
stripped := trim(text)
|
||
if stripped == "" {
|
||
return 0
|
||
}
|
||
// Mirror common.py:match_regex_level — a bullet-looking line is not
|
||
// a heading regardless of the regex family.
|
||
if notBullet(stripped) {
|
||
return 0
|
||
}
|
||
for lvl, pattern := range group {
|
||
re := compileLevelPattern(pattern)
|
||
if re != nil && re.MatchString(stripped) {
|
||
return lvl + 1
|
||
}
|
||
}
|
||
return 0
|
||
}
|
||
|
||
// notBulletPatterns mirrors rag/nlp.not_bullet. A line matching any of
|
||
// these looks like a numbered/bulleted list item rather than a section
|
||
// heading, so it must not be promoted to a title level by the regex
|
||
// families.
|
||
var notBulletPatterns = []*regexp.Regexp{
|
||
regexp.MustCompile(`^0`),
|
||
regexp.MustCompile(`^[0-9]+ +[0-9~个只-]`),
|
||
regexp.MustCompile(`^[0-9]+\.{2,}`),
|
||
regexp.MustCompile(`^[0-9]+(\.[0-9]+){2,}[的中]`),
|
||
}
|
||
|
||
// notBullet mirrors rag/nlp.not_bullet: true when the line looks like a
|
||
// numbered/bulleted list entry rather than a genuine section heading.
|
||
func notBullet(s string) bool {
|
||
for _, re := range notBulletPatterns {
|
||
if re.MatchString(s) {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// notTitlePatterns mirrors rag/nlp.not_title:
|
||
// - notTitleException: `第...条` is always a title (never "not a title").
|
||
// - notTitlePunct: a body line typically carries one of these punctuation
|
||
// marks, so their presence flags the line as body text.
|
||
var (
|
||
notTitleException = regexp.MustCompile(`^第[零一二三四五六七八九十百0-9]+条`)
|
||
notTitlePunct = regexp.MustCompile(`[,;,。;!!]`)
|
||
layoutHeadingRe = regexp.MustCompile(`(?i)(section|title|head)`)
|
||
bulletLayoutRe = regexp.MustCompile(`(?i)(title|head)`)
|
||
numericOnly = regexp.MustCompile(`^[0-9]+$`)
|
||
)
|
||
|
||
// notTitle mirrors rag/nlp.not_title. Returns true when the line looks
|
||
// like body text rather than a section heading, so layout-based heading
|
||
// detection must skip it.
|
||
func notTitle(s string) bool {
|
||
if notTitleException.MatchString(s) {
|
||
return false
|
||
}
|
||
if len(strings.Fields(s)) > 12 || (!strings.Contains(s, " ") && utf8.RuneCountInString(s) >= 32) {
|
||
return true
|
||
}
|
||
return notTitlePunct.MatchString(s)
|
||
}
|
||
|
||
// beforeAt mirrors python's `text.split("@")[0]` (used by
|
||
// match_layout_level / not_title): the part of the line before the first
|
||
// "@" separator, whitespace-trimmed.
|
||
func beforeAt(s string) string {
|
||
if i := strings.Index(s, "@"); i >= 0 {
|
||
return strings.TrimSpace(s[:i])
|
||
}
|
||
return strings.TrimSpace(s)
|
||
}
|
||
|
||
// matchLayoutLevel mirrors common.py:match_layout_level. When the
|
||
// record's layout field flags it as a section/title/head and the text is
|
||
// title-like (not not_title), the record is promoted to `fallback_level`
|
||
// (common.py:resolve_frequency_levels sets this to len(level_group) + 1).
|
||
// Otherwise it stays BODY_LEVEL.
|
||
func matchLayoutLevel(text, layout string, fallbackLevel int) int {
|
||
if layoutHeadingRe.MatchString(layout) && !notTitle(beforeAt(text)) {
|
||
return fallbackLevel
|
||
}
|
||
return bodyLevel
|
||
}
|
||
|
||
// 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
|
||
// (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
|
||
// regex/layout detection for doc_type_kwd != "text");
|
||
// - a text record first tries the selected regex group (match_regex_level);
|
||
// on no match it falls back to match_layout_level (layout hint), else
|
||
// BODY_LEVEL.
|
||
//
|
||
// fallback_level is len(selected_group) + 1, exactly as in python.
|
||
// isColonTitle mirrors Python make_colon_as_title intent: returns true
|
||
// when a line ends with colon, has sentence-ending punctuation before
|
||
// the colon, and the text between them is at least 32 runes long.
|
||
func isColonTitle(s string) bool {
|
||
if s == "" {
|
||
return false
|
||
}
|
||
if !strings.HasSuffix(s, ":") && !strings.HasSuffix(s, ":") {
|
||
return false
|
||
}
|
||
body := strings.TrimSuffix(s, ":")
|
||
body = strings.TrimSuffix(body, ":")
|
||
if body == s {
|
||
return false
|
||
}
|
||
lastPunct := strings.LastIndexAny(body, "。?!!?;;.")
|
||
if lastPunct < 0 {
|
||
return false
|
||
}
|
||
// Use rune-aware slicing: body[lastPunct+1] would be wrong for CJK
|
||
// punctuation (e.g. "。" is 3 bytes in UTF-8). Decode the rune at
|
||
// lastPunct to get the correct byte width and skip the full rune.
|
||
_, runeLen := utf8.DecodeRuneInString(body[lastPunct:])
|
||
between := strings.TrimSpace(body[lastPunct+runeLen:])
|
||
return utf8.RuneCountInString(between) >= 32
|
||
}
|
||
|
||
func resolveTitleLevels(records []lineRecord, p *titleChunkerParam) []int {
|
||
lines := make([]string, len(records))
|
||
for i, r := range records {
|
||
lines[i] = r.text
|
||
}
|
||
group := selectLevelGroup(lines, p.Levels)
|
||
fallbackLevel := len(group) + 1
|
||
out := make([]int, len(records))
|
||
for i, rec := range records {
|
||
if !rec.isText() {
|
||
out[i] = bodyLevel
|
||
continue
|
||
}
|
||
// Python tree_merge short/numeric line filter:
|
||
// sections = [s for s in sections if len(s.split("@")[0].strip()) > 1
|
||
// and not re.match(r"[0-9]+$", s.split("@")[0].strip())]
|
||
if text := beforeAt(rec.text); utf8.RuneCountInString(text) <= 1 || numericOnly.MatchString(text) {
|
||
out[i] = bodyLevel
|
||
continue
|
||
}
|
||
if lvl := matchRegexLevel(rec.text, group); lvl != 0 {
|
||
out[i] = lvl
|
||
continue
|
||
}
|
||
if rec.ckType == "heading" {
|
||
out[i] = fallbackLevel
|
||
continue
|
||
}
|
||
// Python make_colon_as_title: promote lines ending with colon
|
||
// that have sentence-ending punctuation before it and at least
|
||
// 32 chars between the punctuation and the colon.
|
||
if text := beforeAt(rec.text); isColonTitle(text) {
|
||
out[i] = fallbackLevel
|
||
continue
|
||
}
|
||
out[i] = matchLayoutLevel(rec.text, rec.layout, fallbackLevel)
|
||
}
|
||
|
||
// Fallback 4: bullet-pattern detection (Chunker-1.7).
|
||
// When frequency-based detection assigns bodyLevel to every record
|
||
// (no regex matched a single line), use Python's BULLET_PATTERN
|
||
// heuristic (rag/nlp/__init__.py:303-320 bullets_category +
|
||
// title_frequency) to recover structure from numbered/bulleted
|
||
// list entries. This only fires when user levels are empty or
|
||
// matched nothing — never overrides an existing level assignment.
|
||
if len(group) == 0 || allBodyLevel(out) {
|
||
if bull := bulletsCategory(records); bull >= 0 {
|
||
bulletsSize := len(bulletPatterns[bull])
|
||
for i, rec := range records {
|
||
if out[i] != bodyLevel || !rec.isText() {
|
||
continue
|
||
}
|
||
trimmed := strings.TrimSpace(rec.text)
|
||
matched := false
|
||
for j, pat := range bulletPatterns[bull] {
|
||
if pat.MatchString(trimmed) && !notBullet(trimmed) {
|
||
out[i] = j + 1
|
||
matched = true
|
||
break
|
||
}
|
||
}
|
||
if matched {
|
||
continue
|
||
}
|
||
// layout-title fallback: bulletsSize + 1 (mirrors
|
||
// Python len(BULLET_PATTERN[bull]) + 1).
|
||
if bulletLayoutRe.MatchString(rec.layout) && !notTitle(beforeAt(rec.text)) {
|
||
out[i] = bulletsSize + 1
|
||
}
|
||
}
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
// allBodyLevel returns true when every level equals the bodyLevel
|
||
// sentinel. Used by the bullet fallback guard to detect the "no
|
||
// structure found" state where BULLET_PATTERN should kick in.
|
||
func allBodyLevel(levels []int) bool {
|
||
for _, l := range levels {
|
||
if l < bodyLevel {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
|
||
// bulletPatterns mirrors Python BULLET_PATTERN (rag/nlp/__init__.py:258).
|
||
// Group 4 (markdown headings) is omitted — DSL levels cover that case.
|
||
// Each group is depth-ordered: index 0 is the topmost level.
|
||
var bulletPatterns = [][]*regexp.Regexp{
|
||
// Group 0 — Chinese legal (编/章/节/条)
|
||
{
|
||
regexp.MustCompile(`^第[零一二三四五六七八九十百0-9]+(分?编|部分)`),
|
||
regexp.MustCompile(`^第[零一二三四五六七八九十百0-9]+章`),
|
||
regexp.MustCompile(`^第[零一二三四五六七八九十百0-9]+节`),
|
||
regexp.MustCompile(`^第[零一二三四五六七八九十百0-9]+条`),
|
||
regexp.MustCompile(`^[\((][零一二三四五六七八九十百]+[\))]`),
|
||
},
|
||
// Group 1 — Numbering (1., 1.1, 1.1.1)
|
||
{
|
||
regexp.MustCompile(`^第[0-9]+章`),
|
||
regexp.MustCompile(`^第[0-9]+节`),
|
||
regexp.MustCompile(`^[0-9]{0,2}[\. 、]`),
|
||
regexp.MustCompile(`^[0-9]{0,2}\.[0-9]{0,2}[^a-zA-Z/%~-]`),
|
||
regexp.MustCompile(`^[0-9]{0,2}\.[0-9]{0,2}\.[0-9]{0,2}`),
|
||
regexp.MustCompile(`^[0-9]{0,2}\.[0-9]{0,2}\.[0-9]{0,2}\.[0-9]{0,2}`),
|
||
},
|
||
// Group 2 — Chinese numbering (一、, (一))
|
||
{
|
||
regexp.MustCompile(`^第[零一二三四五六七八九十百0-9]+章`),
|
||
regexp.MustCompile(`^第[零一二三四五六七八九十百0-9]+节`),
|
||
regexp.MustCompile(`^[零一二三四五六七八九十百]+[ 、]`),
|
||
regexp.MustCompile(`^[\((][零一二三四五六七八九十百]+[\))]`),
|
||
regexp.MustCompile(`^[\((][0-9]{0,2}[\))]`),
|
||
},
|
||
// Group 3 — English legal
|
||
{
|
||
regexp.MustCompile(`^PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)`),
|
||
regexp.MustCompile(`^Chapter (I+V?|VI*|XI|IX|X)`),
|
||
regexp.MustCompile(`^Section [0-9]+`),
|
||
regexp.MustCompile(`^Article [0-9]+`),
|
||
},
|
||
}
|
||
|
||
// bulletsCategory mirrors Python bullets_category (rag/nlp/__init__.py:303).
|
||
// Counts hits per bullet-pattern group across all text records,
|
||
// excluding not-bullet lines. Returns the group index with the highest
|
||
// hit count, or -1 if no group matched any line.
|
||
func bulletsCategory(records []lineRecord) int {
|
||
hits := make([]int, len(bulletPatterns))
|
||
for grpIdx, group := range bulletPatterns {
|
||
for _, rec := range records {
|
||
if !rec.isText() {
|
||
continue
|
||
}
|
||
txt := strings.TrimSpace(rec.text)
|
||
for _, pat := range group {
|
||
if pat.MatchString(txt) && !notBullet(txt) {
|
||
hits[grpIdx]++
|
||
break
|
||
}
|
||
}
|
||
}
|
||
}
|
||
best, bestHits := -1, 0
|
||
for i, h := range hits {
|
||
if h > bestHits {
|
||
bestHits = h
|
||
best = i
|
||
}
|
||
}
|
||
return best
|
||
}
|
||
|
||
// 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
|
||
// only check relative ordering.
|
||
const bodyLevel = 1<<31 - 1
|
||
|
||
// lineRecords mirrors common.py:extract_line_records' markdown/text/html
|
||
// branch. Returns one record per non-empty input line.
|
||
func lineRecordsFromText(text string) []lineRecord {
|
||
if text == "" {
|
||
return nil
|
||
}
|
||
out := make([]lineRecord, 0)
|
||
for _, ln := range strings.Split(text, "\n") {
|
||
if trim(ln) == "" {
|
||
continue
|
||
}
|
||
out = append(out, lineRecord{
|
||
text: ln,
|
||
docType: "text",
|
||
imgID: nil,
|
||
layout: "",
|
||
pdfPositions: nil,
|
||
positions: nil,
|
||
})
|
||
}
|
||
return out
|
||
}
|
||
|
||
// lineRecord is the internal common shape — same fields as
|
||
// 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
|
||
pdfPositions json.RawMessage
|
||
positions json.RawMessage
|
||
parentMeta map[string]any
|
||
}
|
||
|
||
func (r lineRecord) textOrEmpty() string { return r.text }
|
||
func (r lineRecord) isText() bool { return r.docType == "text" }
|
||
|
||
// trim mirrors python's str.strip(): remove leading/trailing Unicode
|
||
// whitespace. Used for emptiness checks and regex matching so the Go
|
||
// port matches python's `text.strip()` exactly.
|
||
func trim(s string) string {
|
||
return strings.TrimSpace(s)
|
||
}
|
||
|
||
// compileLevelPattern returns a compiled regex for `pattern`. Returns
|
||
// nil on empty/error to skip the entry — same effect as a regex
|
||
// mismatch in python (the row falls through to body).
|
||
//
|
||
// Python's match_regex_level / select_level_group use `re.match`, which
|
||
// anchors at the START of the string (not the end). We mirror that by
|
||
// prepending `^` (unless the caller already anchored) so Go's
|
||
// MatchString behaves exactly like re.match.
|
||
func compileLevelPattern(pattern string) *regexp.Regexp {
|
||
if pattern == "" {
|
||
return nil
|
||
}
|
||
anchored := pattern
|
||
if !strings.HasPrefix(pattern, "^") {
|
||
anchored = "^(?:" + pattern + ")"
|
||
}
|
||
re, err := regexp.Compile(anchored)
|
||
if err != nil {
|
||
return nil
|
||
}
|
||
return re
|
||
}
|
||
|
||
// LevelContext groups the level-detection artefacts (per-line levels
|
||
// + the most-common heading level) so the strategy implementations
|
||
// don't recompute.
|
||
type LevelContext struct {
|
||
levels []int
|
||
mostLevel int
|
||
}
|
||
|
||
// 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
|
||
// Counter(levels).most_common() over the heading levels only. Walk
|
||
// levels in input order so ties resolve to the first-encountered
|
||
// level, matching python's insertion-order tie-break.
|
||
counts := make(map[int]int)
|
||
for _, lvl := range levels {
|
||
if lvl < bodyLevel {
|
||
counts[lvl]++
|
||
}
|
||
}
|
||
most := 0
|
||
best := 0
|
||
for _, lvl := range levels {
|
||
if c := counts[lvl]; c > best {
|
||
best = c
|
||
most = lvl
|
||
}
|
||
}
|
||
return LevelContext{levels: levels, mostLevel: most}
|
||
}
|
||
|
||
func (lc LevelContext) Levels() []int {
|
||
out := make([]int, len(lc.levels))
|
||
copy(out, lc.levels)
|
||
return out
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Component implementations
|
||
// ---------------------------------------------------------------------------
|
||
|
||
// TitleChunkerComponent dispatches based on `param.method`. Heading
|
||
// detection is shared (resolveTitleLevels); the actual chunk-build
|
||
// logic lives in group.go / hierarchy.go.
|
||
type TitleChunkerComponent struct {
|
||
name string
|
||
param titleChunkerParam
|
||
}
|
||
|
||
// NewTitleChunker constructs a TitleChunker from the DSL param map.
|
||
// Errors here surface as canvas compile failures.
|
||
func NewTitleChunker(params map[string]any) (runtime.Component, error) {
|
||
p := defaultsTitle()
|
||
p.Method = "group" // default to group
|
||
p.Update(params)
|
||
if err := p.TitleChunkerParam.Validate(); err != nil {
|
||
return nil, fmt.Errorf("TitleChunker: %w", err)
|
||
}
|
||
return &TitleChunkerComponent{
|
||
name: ComponentNameTitleChunker,
|
||
param: p,
|
||
}, nil
|
||
}
|
||
|
||
// Inputs is exposed so callers can introspect.
|
||
func (c *TitleChunkerComponent) Inputs() map[string]string { return ChunkerInputs }
|
||
|
||
// Outputs is exposed so callers can introspect.
|
||
func (c *TitleChunkerComponent) Outputs() map[string]string { return ChunkerOutputs }
|
||
|
||
// Invoke delegates to the chosen strategy (group or hierarchy).
|
||
func (c *TitleChunkerComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map[string]any) (map[string]any, error) {
|
||
if inputs == nil {
|
||
inputs = map[string]any{}
|
||
}
|
||
// `name` is read from the workflow-wide Globals bag (seeded at
|
||
// pipeline start, published by the File component), not from the
|
||
// upstream output map.
|
||
name := globals.GlobalOrInput(ctx, inputs, "name", "")
|
||
if name == "" {
|
||
return map[string]any{
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{},
|
||
"_ERROR": "TitleChunker: missing required upstream field \"name\"",
|
||
}, nil
|
||
}
|
||
switch c.param.Method {
|
||
case "hierarchy":
|
||
return invokeHierarchy(ctx, db, inputs, &c.param)
|
||
case "group":
|
||
return invokeGroup(ctx, db, inputs, &c.param)
|
||
default:
|
||
return map[string]any{
|
||
"output_format": "chunks",
|
||
"chunks": []map[string]any{},
|
||
"_ERROR": fmt.Sprintf("TitleChunker: unsupported method %q", c.param.Method),
|
||
}, nil
|
||
}
|
||
}
|
||
|
||
// init registers TitleChunker under CategoryIngestion.
|
||
func init() {
|
||
MustRegisterChunker(ComponentNameTitleChunker)
|
||
}
|