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

import (
	"reflect"
	"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).
//
// After the strict-cap fix, a new chunk is started only when the projected
// join exceeds the budget — so the first unit must already sit near the
// budget and the second unit must not fit alongside it.
func TestMergeByTokenSizeFromJSON_OverlapStripsTags(t *testing.T) {
	// Size a and b so:
	//   - each unit alone fits the budget (no atom-split),
	//   - the projected join exceeds the budget (forces a new chunk),
	//   - overlap+b still fits (so the overlap path is exercised).
	aText := strings.Repeat("word ", 20) + "@@1\t2.3## tail"
	bText := "body"
	aN, bN := tokenizeStr(aText), tokenizeStr(bText)
	joinedN := tokenizeStr(aText + "\n" + bText)
	// Budget just below the join so a and b cannot merge, but each alone fits.
	budget := joinedN - 1
	if budget < aN {
		budget = aN
	}
	if budget < bN {
		budget = bN
	}
	if joinedN <= budget {
		t.Fatalf("could not derive tight budget (a=%d b=%d joined=%d budget=%d)", aN, bN, joinedN, budget)
	}
	items := [][]schema.ChunkDoc{
		{
			{Text: aText, DocType: "text", CKType: "text", TKNums: intPtr(aN)},
			{Text: bText, DocType: "text", CKType: "text", TKNums: intPtr(bN)},
		},
	}
	got := mergeByTokenSizeFromJSON(items, budget, 30.0)
	merged := got[0]
	if len(merged) != 2 {
		t.Fatalf("want 2 merged chunks (overlap path), got %d (a=%d b=%d budget=%d)", len(merged), aN, bN, budget)
	}
	// 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. .
	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)
	}
	if n := tokenizeStr(merged[1].Text); n > budget {
		t.Errorf("overlap pushed second chunk over budget: tokens=%d", n)
	}
}

// TestMergeByTokenSizeFromJSON_ClampsOverlappedPct locks the review finding
// from yuzhichang (PR #17396): mergeByTokenSizeFromJSON clamps an out-of-range
// overlappedPct to [0,100] so the merge math never yields a negative/inverted
// threshold. Out-of-range values must not panic and must behave identically to
// their clamped-in-range equivalent (150 == 100, -5 == 0, and the same for
// huge magnitudes that would otherwise overflow the float->int slice index).
// clampOverlapFixture returns a fresh input for mergeByTokenSizeFromJSON.
// A new slice must be built per invocation: mergeByTokenSizeFromJSON mutates
// its perItem argument in place (token.go: perItem[idx] = merged) and returns
// the same backing array. Reusing one fixture across calls lets later calls
// reprocess already-merged chunks, and — because the result aliases the input
// — silently overwrites earlier results, making the clamp assertions vacuous
// (see code review on PR #17396).
//
// The first chunk carries TKNums 130 — just above chunk_token_size 128 — so
// the two clamp directions are both exercised: at pct=0 (threshold 128) the
// chunks stay split, while an UNCLAMPED negative pct raises the threshold
// (e.g. -5 -> 134.4) and merges them. With TKNums=100 both cases merge
// identically, so the lower-clamp assertions would pass even if the clamp
// were removed (review: coderabbitai on PR #17416).
func clampOverlapFixture() [][]schema.ChunkDoc {
	return [][]schema.ChunkDoc{
		{
			{Text: strings.Repeat("word ", 20), DocType: "text", CKType: "text", TKNums: intPtr(130)},
			{Text: "body", DocType: "text", CKType: "text", TKNums: intPtr(5)},
		},
	}
}

func TestMergeByTokenSizeFromJSON_ClampsOverlappedPct(t *testing.T) {
	at100 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 100)
	if at100 == nil || len(at100) == 0 {
		t.Fatalf("overlappedPct=100: nil/empty result")
	}
	at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150)
	atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300)
	if !reflect.DeepEqual(at100, at150) {
		t.Errorf("overlappedPct=150 should clamp to 100; output differs from 100")
	}
	if !reflect.DeepEqual(at100, atHuge) {
		t.Errorf("overlappedPct=1e300 should clamp to 100; output differs from 100")
	}

	at0 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 0)
	if at0 == nil || len(at0) == 0 {
		t.Fatalf("overlappedPct=0: nil/empty result")
	}
	atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5)
	atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300)
	if !reflect.DeepEqual(at0, atNeg) {
		t.Errorf("overlappedPct=-5 should clamp to 0; output differs from 0")
	}
	if !reflect.DeepEqual(at0, atNegHuge) {
		t.Errorf("overlappedPct=-1e300 should clamp to 0; output differs from 0")
	}
}

// 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
// covers 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)
	}
}