ragflow/internal/ingestion/component/schema/chunker.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.
//

package schema

import (
	"fmt"
	"math"
	"strconv"
	"strings"
)

// ChunkerFromUpstream is the shared upstream payload consumed by all four
// chunker variants (TokenChunker, TitleChunker, GroupTitleChunker,
// HierarchyTitleChunker).
//
// It mirrors the wire shape defined across two equivalent Pydantic
// schemas in the Python codebase:
//
//	rag/flow/chunker/schema.py:TokenChunkerFromUpstream
//	rag/flow/chunker/title_chunker/schema.py:TitleChunkerFromUpstream
//
// Both classes are field-identical (apart from the class name) — the
// variants share one upstream payload; their *Param structs differ.
//
// Wire shape (Pydantic):
//
//	created_time: float | None  (alias _created_time)
//	elapsed_time: float | None  (alias _elapsed_time)
//	name: str                   (required)
//	file: dict | None
//	chunks: list[dict] | None
//	output_format: Literal["json","markdown","text","html","chunks"] | None
//	json_result:      list[dict] | None  (alias "json")
//	markdown_result:  str        | None  (alias "markdown")
//	text_result:      str        | None  (alias "text")
//	html_result:      str        | None  (alias "html")
type ChunkerFromUpstream struct {
	CreatedTime *float64 `json:"_created_time,omitempty"`
	ElapsedTime *float64 `json:"_elapsed_time,omitempty"`

	// Name is the source document name. Required.
	Name string `json:"name"`

	// DocID is the originating document ID. When set, the chunker can
	// re-resolve the source PDF from storage to crop section images on
	// demand, instead of carrying the binary across the component
	// boundary.
	DocID string `json:"doc_id,omitempty"`

	// Bucket and Path are the explicit storage location of the source
	// PDF. They take precedence over DocID-driven resolution.
	Bucket string `json:"bucket,omitempty"`
	Path   string `json:"path,omitempty"`

	// File is the optional upstream file descriptor.
	File *ChunkerFileMeta `json:"file,omitempty"`

	// Chunks is the upstream chunk list, set when output_format == "chunks".
	Chunks []ChunkDoc `json:"chunks,omitempty"`

	// OutputFormat controls which of the *Result fields below is the
	// active payload. Allowed values:
	//   "json"     -> JSONResult
	//   "markdown" -> MarkdownResult
	//   "text"     -> TextResult
	//   "html"     -> HTMLResult
	//   "chunks"   -> Chunks
	OutputFormat PayloadFormat `json:"output_format,omitempty"`

	// JSONResult is the upstream structured JSON list (alias "json" in
	// Python). Set when OutputFormat == "json".
	JSONResult []ChunkDoc `json:"json,omitempty"`

	// MarkdownResult is the upstream markdown payload (alias "markdown").
	// Set when OutputFormat == "markdown".
	MarkdownResult *string `json:"markdown,omitempty"`

	// TextResult is the upstream plain-text payload (alias "text").
	// Set when OutputFormat == "text".
	TextResult *string `json:"text,omitempty"`

	// HTMLResult is the upstream HTML payload (alias "html").
	// Set when OutputFormat == "html".
	HTMLResult *string `json:"html,omitempty"`
}

// Validate enforces the only required field in ChunkerFromUpstream: Name.
// OutputFormat is not strictly required (defaults to "" and the
// component decides what to do with an empty payload), but the
// combination of `OutputFormat == "chunks"` with a non-nil Chunks is the
// happy path.
func (c *ChunkerFromUpstream) Validate() error {
	if c.Name == "" {
		return errRequiredField{Field: "name"}
	}
	if !c.OutputFormat.isKnown() {
		return errInvalidValue{Field: "output_format", Value: string(c.OutputFormat)}
	}
	switch c.OutputFormat {
	case PayloadFormatChunks:
		return nil
	case PayloadFormatJSON:
		if c.JSONResult == nil {
			return errRequiredField{Field: "json"}
		}
	case PayloadFormatMarkdown:
		if c.MarkdownResult == nil {
			return errRequiredField{Field: "markdown"}
		}
	case PayloadFormatText:
		if c.TextResult == nil {
			return errRequiredField{Field: "text"}
		}
	case PayloadFormatHTML:
		if c.HTMLResult == nil {
			return errRequiredField{Field: "html"}
		}
	}
	return nil
}

// ChunkerOutputs is the result of invoking any chunker variant. All
// four chunker components emit the same shape: a list of chunk maps
// under the "chunks" key, plus a marker output_format = "chunks".
//
// Mirrors what each chunker sets at the end of _invoke:
//
//	self.set_output("output_format", "chunks")
//	self.set_output("chunks", chunks)
//
// Unlike the Python runtime (which auto-merges every input kwarg into the
// output via ProcessBase.invoke), the Go runtime only forwards the explicit
// component output to the next node. The run-level metadata that downstream
// consumers still need (e.g. Tokenizer reads `name` for title embedding, and
// tenant_id/kb_id for embedding-model resolution) therefore lives in the
// workflow-wide CanvasState.Globals bag — seeded at pipeline start and
// published by the File component — and is read directly from ctx, not
// re-emitted by each chunker. The fields below are the chunker-owned outputs.
type ChunkerOutputs struct {
	// OutputFormat is always "chunks" on success.
	OutputFormat PayloadFormat `json:"output_format,omitempty"`

	// Chunks is the produced chunk list. Each entry is a free-form map
	// mirroring the dict shape the Python code builds (text, doc_type_kwd,
	// tk_nums, PDF_POSITIONS_KEY, mom, img_id, etc.).
	Chunks []ChunkDoc `json:"chunks,omitempty"`

	// Error is set when the component short-circuits with an error
	// message (Python: set_output("_ERROR", ...)).
	Error string `json:"_ERROR,omitempty"`
}

// ---------------------------------------------------------------------------
// TokenChunkerParam
// ---------------------------------------------------------------------------
//
// Mirrors rag/flow/chunker/token_chunker.py:TokenChunkerParam.__init__.
// All fields are user-tunable; defaults match the Python values.

type TokenChunkerParam struct {
	// DelimiterMode selects the chunking strategy.
	// Allowed value: "delimiter". The legacy "token_size" value is accepted for
	// backward compatibility and normalized to "delimiter" (it is behaviorally
	// identical at runtime; see Validate). The single-chunk "one" behavior is
	// provided by the separate OneChunker component.
	DelimiterMode string `json:"delimiter_mode"`

	// ChunkTokenSize is the target chunk size in tokens.
	ChunkTokenSize int `json:"chunk_token_size"`

	// Delimiters is the list of split tokens. Strings wrapped in
	// backticks (e.g., "`\\n`") denote user-defined regex split points.
	Delimiters []string `json:"delimiters"`

	// OverlappedPercent is the overlap percentage in [0, 90]. Mirrors
	// Python common/float_utils.py:50-58 — an integer-like float in the
	// same range so DSL templates written for the Python pipeline work
	// out of the box  A [0,1) fraction input is also
	// accepted and normalized to this scale by tokenChunkerParam.Update
	// (via normalizeOverlappedPercent).
	OverlappedPercent float64 `json:"overlapped_percent"`

	// ChildrenDelimiters is the secondary split applied to text chunks.
	ChildrenDelimiters []string `json:"children_delimiters"`

	// TableContextSize is the number of surrounding tokens to attach
	// to table chunks. 0 disables.
	TableContextSize int `json:"table_context_size"`

	// ImageContextSize is the number of surrounding tokens to attach
	// to image chunks. 0 disables.
	ImageContextSize int `json:"image_context_size"`

	// UnderCap selects the merge strategy when greedily accumulating
	// adjacent units (token_chunker UNDER_CAP vs OVER_CAP).
	//   - false (default): OVER_CAP — mirrors Python's canonical default.
	//     A chunk may exceed the token target by at most one incoming unit
	//     (a boundary overflow then closes the chunk).
	//   - true: UNDER_CAP — strictly never exceed the target; when the
	//     projected join would overflow, start a fresh chunk instead.
	// This is the wire-facing bool; the active strategy is exposed as
	// MergeStrategy so callers never have to invert it.
	UnderCap bool `json:"under_cap"`
}

// MergeStrategy selects how the TokenChunker greedily accumulates adjacent
// units into a chunk. It mirrors Python's rag/nlp/__init__.py MergeStrategy so
// the Go and Python implementations stay on the same vocabulary (OVER_CAP /
// UNDER_CAP) instead of a bare inlined bool.
type MergeStrategy int

const (
	// MergeOverCap is the canonical default (Python OVER_CAP): a chunk may
	// exceed the token target by at most one incoming unit (a boundary
	// overflow then closes the chunk).
	MergeOverCap MergeStrategy = iota
	// MergeUnderCap never exceeds the target; when the projected join would
	// overflow, a fresh chunk is started instead (Python UNDER_CAP).
	MergeUnderCap
)

// MergeStrategy reports the active merge strategy for this param. It is derived
// from UnderCap so existing configs that set "under_cap" keep working without a
// schema break, and callers read the strategy directly instead of inverting a
// bool at every call site.
func (p TokenChunkerParam) MergeStrategy() MergeStrategy {
	if p.UnderCap {
		return MergeUnderCap
	}
	return MergeOverCap
}

// Defaults returns the Python default TokenChunkerParam.
func (TokenChunkerParam) Defaults() TokenChunkerParam {
	return TokenChunkerParam{
		DelimiterMode:      "delimiter",
		ChunkTokenSize:     512,
		Delimiters:         []string{"\n"},
		OverlappedPercent:  0,
		ChildrenDelimiters: []string{},
		TableContextSize:   0,
		ImageContextSize:   0,
	}
}

// NumericFromAny normalises JSON-decoded ints to float64 so the
// schema-defaults-vs-Param-Update convention doesn't depend on the
// encoding source (yaml/toml/json all behave the same). It is shared by
// the chunker config decoders and NormalizeOverlappedPercent.
func NumericFromAny(v any) (float64, bool) {
	switch x := v.(type) {
	case float64:
		return x, true
	case float32:
		return float64(x), true
	case int:
		return float64(x), true
	case int32:
		return float64(x), true
	case int64:
		return float64(x), true
	case uint:
		return float64(x), true
	case uint32:
		return float64(x), true
	case uint64:
		return float64(x), true
	}
	return 0, false
}

// NormalizeOverlappedPercent mirrors Python common/float_utils.py:50-58
// (normalize_overlapped_percent). Python's user-facing input is a [0,1)
// fraction (token_chunker.py validates "[0, 1)"); this converts it to the
// [0,90] integer-percentage scale that the merge math (token.go:705,
// `(100-x)/100`) expects, matching Python's canonical norm. Diff Chunker-2.6.
//
// Behaviour matches Python exactly:
//   - parse like float(): numbers, or numeric strings (Python also accepts
//     strings via float()); on any failure / NaN / Inf, return 0
//     (Python: except (TypeError, ValueError, OverflowError) -> 0).
//   - 0 < v < 1 -> v *= 100   (fraction -> percentage)
//   - v = int(v)             (truncation toward zero, matches Python for
//     non-negatives)
//   - clamp to [0, 90]        (Python: max(0, min(v, 90)))
//
// Exported so TokenChunkerParam.Validate can normalize directly-constructed
// structs (not just the config-path Update), keeping both paths identical.
func NormalizeOverlappedPercent(v any) float64 {
	value, ok := NumericFromAny(v)
	if !ok {
		if s, isStr := v.(string); isStr {
			if f, err := strconv.ParseFloat(strings.TrimSpace(s), 64); err == nil {
				value, ok = f, true
			}
		}
	}
	if !ok || math.IsNaN(value) || math.IsInf(value, 0) {
		return 0
	}
	return clampOverlappedPercent(value)
}

// clampOverlappedPercent applies the [0,1) fraction -> [0,90] percent
// conversion, clamps to [0,90], then truncates to an integer (Python
// faithfulness). Clamping MUST happen before the int conversion: Go's
// float->int is implementation-defined when the value is out of int's range
// (e.g. 1e300), whereas Python's int() is arbitrary-precision, so clamping
// first keeps us aligned with Python's max(0, min(int(v), 90)). Used by the
// lenient config path (NormalizeOverlappedPercent).
func clampOverlappedPercent(value float64) float64 {
	if 0 < value && value < 1 {
		value *= 100
	}
	if value < 0 {
		value = 0
	}
	if value > 90 {
		value = 90
	}
	// Round rather than truncate so fraction inputs like 0.29 normalize
	// to 29 (user expectation) rather than 28. Mirrors Python's
	// common/float_utils.py:50-58 fix for #17418. We still clamp
	// before rounding so out-of-int-range values (e.g. 1e300) land on
	// 90, not an implementation-defined int.
	value = math.Round(value)
	return value
}

// Validate enforces the same enum/range checks the runtime component expects
// at construction time, keeping the schema and component decoder aligned.
func (p *TokenChunkerParam) Validate() error {
	// Strict overlap normalization for directly-constructed structs: scale a
	// [0,1) fraction to a [0,90] percent (matching the config path's
	// semantics) and truncate toward zero, then REJECT anything still outside
	// [0,90]. The config path (Update) clamps instead, so a fraction like 0.3
	// means 30% here too, while an out-of-range value like 95 or -5 is caught
	// rather than silently clamped — eliminating the latent footgun where a
	// fraction passed to a struct bypassed normalization.
	if v := p.OverlappedPercent; 0 < v && v < 1 {
		// Round rather than truncate so fraction inputs like 0.29 normalize
		// to 29 (user expectation) rather than 28. Mirrors Python's
		// common/float_utils.py:50-58 fix for #17418.
		p.OverlappedPercent = math.Round(v * 100)
	}
	// Backward-compat: Python removed the standalone "token_size" mode and
	// treats it as identical to "delimiter" at runtime, normalizing it in
	// check() (rag/flow/chunker/token_chunker.py). A dataset configured under
	// the old scheme may still carry delimiter_mode="token_size"; normalize it
	// here instead of rejecting it so those pipelines keep working (mirrors
	// the Python behaviour exactly).
	if p.DelimiterMode == "token_size" {
		p.DelimiterMode = "delimiter"
	}
	switch p.DelimiterMode {
	case "delimiter":
	default:
		return errInvalidValue{Field: "delimiter_mode", Value: p.DelimiterMode}
	}
	if p.ChunkTokenSize <= 0 {
		return errInvalidValue{Field: "chunk_token_size", Value: fmt.Sprintf("%d", p.ChunkTokenSize)}
	}
	if p.OverlappedPercent < 0 || p.OverlappedPercent > 90 {
		return errInvalidValue{Field: "overlapped_percent", Value: fmt.Sprintf("%v", p.OverlappedPercent)}
	}
	if p.TableContextSize < 0 {
		return errInvalidValue{Field: "table_context_size", Value: fmt.Sprintf("%d", p.TableContextSize)}
	}
	if p.ImageContextSize < 0 {
		return errInvalidValue{Field: "image_context_size", Value: fmt.Sprintf("%d", p.ImageContextSize)}
	}
	return nil
}

// ---------------------------------------------------------------------------
// TitleChunkerParam
// ---------------------------------------------------------------------------
//
// Mirrors rag/flow/chunker/title_chunker/common.py:TitleChunkerParam.
// The class also reads `self.method` (set externally — see Python
// title_chunker.py:31-37 routing on `self._param.method`). The Go port
// captures it explicitly here. The component's `check()` enforces enum
// values.

type TitleChunkerParam struct {
	// Method routes to the right title-chunker strategy.
	// Allowed values: "hierarchy", "group".
	Method string `json:"method,omitempty"`

	// Levels is a list of regex-list groups, one per hierarchy level.
	// Each group is a list of regex strings; the component picks the
	// best-matching group at runtime.
	Levels [][]string `json:"levels"`

	// Hierarchy is the heading depth used by HierarchyTitleChunker.
	// Stored as a pointer to distinguish nil (unset) from 0.
	Hierarchy *int `json:"hierarchy,omitempty"`

	// IncludeHeadingContent, when true, makes the heading text part
	// of each emitted chunk.
	IncludeHeadingContent bool `json:"include_heading_content"`

	// RootChunkAsHeading, when true, prepends the root chunk's text
	// to every emitted chunk (and drops the root chunk itself).
	RootChunkAsHeading bool `json:"root_chunk_as_heading"`
}

// Defaults returns the Python default TitleChunkerParam. `Method` is
// not initialized in the Python `__init__` (it is set externally); the
// default is left as the empty string and the component must supply it.
func (TitleChunkerParam) Defaults() TitleChunkerParam {
	return TitleChunkerParam{
		Levels:                [][]string{},
		Hierarchy:             nil,
		IncludeHeadingContent: false,
		RootChunkAsHeading:    false,
	}
}

// Validate enforces the Python `check()` invariants that are also
// expressible in pure-data terms: when Method == "hierarchy" the
// hierarchy depth and level config must be present.
func (p *TitleChunkerParam) Validate() error {
	switch p.Method {
	case "hierarchy", "group":
	case "":
		return nil
	default:
		return errInvalidValue{Field: "method", Value: p.Method}
	}
	switch p.Method {
	case "hierarchy", "group":
		if len(p.Levels) == 0 {
			return errRequiredField{Field: "levels"}
		}
	}
	if p.Method == "hierarchy" && (p.Hierarchy == nil || *p.Hierarchy <= 0) {
		return errRequiredField{Field: "hierarchy"}
	}
	return nil
}

// ---------------------------------------------------------------------------
// GroupTitleChunkerParam / HierarchyTitleChunkerParam
// ---------------------------------------------------------------------------
//
// In the Python codebase, both variants share the SAME
// `TitleChunkerParam` class — there is no per-variant param. The
// dispatch happens in title_chunker.py:31-37 by reading
// `self._param.method`.
//
// In Go we model the shared class as `TitleChunkerParam` and expose
// type aliases so component files can name the param type they
// actually use. This keeps the wire schema faithful to Python while
// giving each variant a self-documenting entry point in the registry.

type GroupTitleChunkerParam = TitleChunkerParam
type HierarchyTitleChunkerParam = TitleChunkerParam