ragflow/internal/ingestion/component/chunker/textcode_parity_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 (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"testing"

	"ragflow/internal/parser/parser"
)

const (
	textCodeDelimiter = "\n!?;。；！？"
	textCodeTokenSize = 128
)

type textCodeGolden struct {
	Meta   map[string]any   `json:"meta"`
	Chunks []map[string]any `json:"chunks"`
}

// TestTextParserTokenChunkerParity covers the production text&code handoff:
// TextParser.ParseWithResult emits structured JSON items after splitting on
// the text&code delimiters while keeping matched delimiters in the item text;
// TokenChunker owns the final token-budget merge. The Python golden is
// generated by the same TxtParser path used by parser._code; comparison
// follows the parser alignment contract and ignores representation-only chunk
// boundaries and whitespace.
func TestTextParserTokenChunkerParity(t *testing.T) {
	cases := []struct {
		name       string
		samplePath string
		goldenPath string
		minChunks  int
	}{
		{
			name:       "en",
			samplePath: "../../../parser/parser/testdata/textcode.sample.en.txt",
			goldenPath: "testdata/textcode/python.en.golden.json",
			minChunks:  1,
		},
		{
			name:       "zh",
			samplePath: "../../../parser/parser/testdata/textcode.sample.zh.txt",
			goldenPath: "testdata/textcode/python.zh.golden.json",
			minChunks:  1,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			sample, err := os.ReadFile(tc.samplePath)
			if err != nil {
				t.Fatalf("read sample: %v", err)
			}
			goldenRaw, err := os.ReadFile(tc.goldenPath)
			if err != nil {
				t.Fatalf("read golden: %v", err)
			}
			var golden textCodeGolden
			if err := json.Unmarshal(goldenRaw, &golden); err != nil {
				t.Fatalf("decode golden: %v", err)
			}
			if len(golden.Chunks) == 0 {
				t.Fatal("golden has no chunks")
			}

			parsed := parser.NewTextParser().ParseWithResult(
				t.Context(), tc.samplePath, sample,
			)
			if parsed.Err != nil {
				t.Fatalf("TextParser.ParseWithResult: %v", parsed.Err)
			}

			component, err := NewTokenChunker(map[string]any{
				"chunk_token_size": float64(textCodeTokenSize),
				// Parser already consumed the text&code delimiters. An empty
				// list makes this invocation exercise only the downstream merge.
				"delimiters": []string{},
			})
			if err != nil {
				t.Fatalf("NewTokenChunker: %v", err)
			}
			out, err := component.Invoke(context.Background(), nil, map[string]any{
				"name":          tc.samplePath,
				"output_format": "json",
				"json":          parsed.JSON,
			})
			if err != nil {
				t.Fatalf("TokenChunker.Invoke: %v", err)
			}
			if msg, ok := out["_ERROR"].(string); ok && msg != "" {
				t.Fatalf("TokenChunker returned _ERROR: %s", msg)
			}
			got, ok := out["chunks"].([]map[string]any)
			if !ok {
				t.Fatalf("TokenChunker chunks has type %T", out["chunks"])
			}
			if len(got) < tc.minChunks {
				t.Fatalf("final chunk count = %d, want at least %d", len(got), tc.minChunks)
			}

			wantText := normalizeTextCodeGolden(golden.Chunks, textCodeDelimiter)
			gotText := normalizeTextCodeGolden(got, textCodeDelimiter)
			if gotText != wantText {
				t.Fatalf("text&code final chunks diverged:\n%s", textCodeDiff(gotText, wantText))
			}
		})
	}
}

func normalizeTextCodeGolden(items []map[string]any, delimiter string) string {
	parts := make([]string, 0, len(items))
	for _, item := range items {
		text, _ := item["text"].(string)
		var normalized strings.Builder
		for _, r := range text {
			if strings.ContainsRune(delimiter, r) {
				normalized.WriteByte(' ')
			} else {
				normalized.WriteRune(r)
			}
		}
		if text := strings.Join(strings.Fields(normalized.String()), " "); text != "" {
			parts = append(parts, text)
		}
	}
	return strings.Join(parts, " ")
}

func textCodeDiff(got, want string) string {
	const max = 2000
	if len(got) > max {
		got = got[:max] + "...(truncated)"
	}
	if len(want) > max {
		want = want[:max] + "...(truncated)"
	}
	return fmt.Sprintf("--- GO ---\n%s\n--- PYTHON ---\n%s", got, want)
}