mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-30 20:49:21 +08:00
327 lines
10 KiB
Go
327 lines
10 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.
|
|
//
|
|
|
|
package chunker
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
|
|
"ragflow/internal/agent/runtime"
|
|
"ragflow/internal/ingestion/component/schema"
|
|
)
|
|
|
|
func TestGroupTitleChunker_Registered(t *testing.T) {
|
|
factory, cat, meta, ok := runtime.DefaultRegistry.Lookup("GroupTitleChunker")
|
|
if !ok {
|
|
t.Fatal("GroupTitleChunker: registry miss")
|
|
}
|
|
if cat != runtime.CategoryIngestion {
|
|
t.Errorf("category = %q, want %q", cat, runtime.CategoryIngestion)
|
|
}
|
|
if factory == nil {
|
|
t.Error("factory is nil")
|
|
}
|
|
if len(meta.Inputs) == 0 {
|
|
t.Errorf("inputs metadata is empty")
|
|
}
|
|
if len(meta.Outputs) == 0 {
|
|
t.Errorf("outputs metadata is empty")
|
|
}
|
|
}
|
|
|
|
func TestGroupTitleChunker_InputsOutputs_NonEmpty(t *testing.T) {
|
|
_, _, meta, ok := runtime.DefaultRegistry.Lookup("GroupTitleChunker")
|
|
if !ok {
|
|
t.Fatal("registry miss")
|
|
}
|
|
if len(meta.Inputs) == 0 {
|
|
t.Error("inputs metadata is empty")
|
|
}
|
|
if len(meta.Outputs) == 0 {
|
|
t.Error("outputs metadata is empty")
|
|
}
|
|
}
|
|
|
|
func TestGroupTitleChunker_NewRejectsHierarchyWithoutHierarchyParam(t *testing.T) {
|
|
// Although method defaults to "group", the levels branch still
|
|
// fires and rejects an empty levels config.
|
|
if _, err := NewGroupTitleChunker(map[string]any{}); err == nil {
|
|
t.Fatal("expected error for empty levels, got nil")
|
|
}
|
|
}
|
|
|
|
// TestExtractLineRecords_MarkdownKey verifies extractLineRecords reads
|
|
// the "markdown" payload key. Before this fix extractLineRecords only
|
|
// looked at "text"/"content" and silently returned nil for parser
|
|
// output that carried output_format="markdown".
|
|
func TestExtractLineRecords_MarkdownKey(t *testing.T) {
|
|
records := extractLineRecords(map[string]any{
|
|
"output_format": "markdown",
|
|
"markdown": "line1\nline2\nline3",
|
|
})
|
|
if len(records) != 3 {
|
|
t.Fatalf("got %d records, want 3", len(records))
|
|
}
|
|
for i, r := range records {
|
|
if r.textOrEmpty() == "" {
|
|
t.Errorf("record[%d]: empty text", i)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestExtractLineRecords_HTMLKey verifies extractLineRecords reads the
|
|
// "html" payload key, same safety-net rationale as TestMarkdownKey.
|
|
func TestExtractLineRecords_HTMLKey(t *testing.T) {
|
|
records := extractLineRecords(map[string]any{
|
|
"output_format": "html",
|
|
"html": "<p>first</p>\n<p>second</p>",
|
|
})
|
|
if len(records) != 2 {
|
|
t.Fatalf("got %d records, want 2", len(records))
|
|
}
|
|
}
|
|
|
|
// TestExtractLineRecords_TextKeyStillWorks ensures the existing "text"
|
|
// key path is not broken by the addition of "markdown"/"html".
|
|
func TestExtractLineRecords_TextKeyStillWorks(t *testing.T) {
|
|
records := extractLineRecords(map[string]any{
|
|
"text": "hello\nworld",
|
|
})
|
|
if len(records) != 2 {
|
|
t.Fatalf("got %d records, want 2", len(records))
|
|
}
|
|
}
|
|
|
|
func TestGroupTitleChunker_InvokeEmptyInput(t *testing.T) {
|
|
c, err := NewGroupTitleChunker(map[string]any{
|
|
"levels": [][]string{{`^# `}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewGroupTitleChunker: %v", err)
|
|
}
|
|
out, err := c.Invoke(t.Context(), nil, map[string]any{})
|
|
if err != nil {
|
|
t.Fatalf("Invoke: %v", err)
|
|
}
|
|
if got, want := out["output_format"], "chunks"; got != want {
|
|
t.Errorf("output_format = %v, want %v", got, want)
|
|
}
|
|
chunks, _ := out["chunks"].([]map[string]any)
|
|
if len(chunks) != 0 {
|
|
t.Errorf("chunks = %d, want 0", len(chunks))
|
|
}
|
|
}
|
|
|
|
func TestGroupTitleChunker_Headings_ASCII(t *testing.T) {
|
|
c, err := NewGroupTitleChunker(map[string]any{
|
|
"levels": [][]string{
|
|
{`^# `},
|
|
{`^## `},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewGroupTitleChunker: %v", err)
|
|
}
|
|
input := "# Heading One\nBody line under H1.\nAnother body line.\n# Heading Two\nBody under second H1."
|
|
out, err := c.Invoke(t.Context(), nil, map[string]any{
|
|
"name": "doc.md",
|
|
"text": input,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Invoke: %v", err)
|
|
}
|
|
chunks, _ := out["chunks"].([]map[string]any)
|
|
if len(chunks) == 0 {
|
|
t.Fatal("chunks: want >=1, got 0")
|
|
}
|
|
}
|
|
|
|
func TestGroupTitleChunker_RootChunkAsHeading_StillSingleGroup(t *testing.T) {
|
|
// When the input is single-group, the root-as-heading branch
|
|
// doesn't reduce the count below 1.
|
|
c, err := NewGroupTitleChunker(map[string]any{
|
|
"levels": [][]string{{`^# `}},
|
|
"root_chunk_as_heading": true,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewGroupTitleChunker: %v", err)
|
|
}
|
|
out, err := c.Invoke(t.Context(), nil, map[string]any{
|
|
"name": "doc.md",
|
|
"text": "Body without any heading here.\nMore body.",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Invoke: %v", err)
|
|
}
|
|
chunks, _ := out["chunks"].([]map[string]any)
|
|
if len(chunks) == 0 {
|
|
t.Fatal("chunks: want >=1, got 0")
|
|
}
|
|
}
|
|
|
|
// TestResolveGroupTargetLevel_UsesMostLevelDirectly pins Gap F: when
|
|
// `hierarchy` is unset the group target level is `most_level` DIRECTLY,
|
|
// not resolve_target_level (which would re-rank the distinct heading
|
|
// levels and pick the wrong depth when levels are not contiguous from
|
|
// 1). With levels {2,2,3} most_level is 2, but resolve_target_level
|
|
// would return 3.
|
|
func TestResolveGroupTargetLevel_UsesMostLevelDirectly(t *testing.T) {
|
|
levels := []int{2, 2, 3, bodyLevel, bodyLevel}
|
|
pUnset := &titleChunkerParam{TitleChunkerParam: schema.TitleChunkerParam{Method: "group"}}
|
|
if got := resolveGroupTargetLevel(levels, pUnset, 2); got != 2 {
|
|
t.Errorf("unset hierarchy: got %d, want 2 (most_level directly)", got)
|
|
}
|
|
h := 2
|
|
pSet := &titleChunkerParam{TitleChunkerParam: schema.TitleChunkerParam{Method: "group", Hierarchy: &h}}
|
|
if got := resolveGroupTargetLevel(levels, pSet, 2); got != 3 {
|
|
t.Errorf("hierarchy=2: got %d, want 3 (resolve_target_level)", got)
|
|
}
|
|
}
|
|
|
|
// TestGroupChunker_StructuredMetadata pins Gap E: for a structured
|
|
// (output_format=chunks) payload, non-text records keep their
|
|
// doc_type_kwd and img_id on the emitted chunk.
|
|
func TestGroupChunker_StructuredMetadata(t *testing.T) {
|
|
c, err := NewGroupTitleChunker(map[string]any{
|
|
"levels": [][]string{{`^# `}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewGroupTitleChunker: %v", err)
|
|
}
|
|
items := []map[string]any{
|
|
{"text": "# Heading", "doc_type_kwd": "text"},
|
|
{"text": "body line", "doc_type_kwd": "text"},
|
|
{"text": "an image caption", "doc_type_kwd": "image", "img_id": "img-9"},
|
|
}
|
|
out, err := c.Invoke(t.Context(), nil, map[string]any{
|
|
"name": "doc",
|
|
"output_format": "chunks",
|
|
"chunks": items,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Invoke: %v", err)
|
|
}
|
|
chunks, _ := out["chunks"].([]map[string]any)
|
|
found := false
|
|
for _, ck := range chunks {
|
|
if dt, _ := ck["doc_type_kwd"].(string); dt == "image" {
|
|
found = true
|
|
if ck["img_id"] != "img-9" {
|
|
t.Errorf("image chunk img_id = %v, want img-9", ck["img_id"])
|
|
}
|
|
}
|
|
}
|
|
if !found {
|
|
t.Fatal("no image chunk emitted")
|
|
}
|
|
}
|
|
|
|
// TestGroupChunker_MergesPDFPositionsAndRemovesTags is the TDD test for
|
|
// migration diffs Chunker-1.6 / 2.8: when the group chunker merges
|
|
// multiple adjacent text records into one chunk it must (a) strip the
|
|
// parser-emitted `@@...##` position tags from the joined text, and
|
|
// (b) MERGE (not drop) the `positions` coordinate matrices across the
|
|
// merged records — mirroring common.py:255 remove_tag + merge.
|
|
func TestGroupChunker_MergesPDFPositionsAndRemovesTags(t *testing.T) {
|
|
c, err := NewGroupTitleChunker(map[string]any{
|
|
"levels": [][]string{{`^# `}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewGroupTitleChunker: %v", err)
|
|
}
|
|
items := []map[string]any{
|
|
{"text": "# Heading", "doc_type_kwd": "text"},
|
|
{"text": "body one @@1\t10.0\t20.0\t30.0\t40.0## tail", "doc_type_kwd": "text", "positions": [][]float64{{1, 10, 20, 30, 40}}},
|
|
{"text": "body two @@2\t15.0\t25.0\t35.0\t45.0## tail", "doc_type_kwd": "text", "positions": [][]float64{{2, 15, 25, 35, 45}}},
|
|
}
|
|
out, err := c.Invoke(t.Context(), nil, map[string]any{
|
|
"name": "doc",
|
|
"output_format": "chunks",
|
|
"chunks": items,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Invoke: %v", err)
|
|
}
|
|
chunks, _ := out["chunks"].([]map[string]any)
|
|
if len(chunks) == 0 {
|
|
t.Fatal("no chunks emitted")
|
|
}
|
|
for _, ck := range chunks {
|
|
text, _ := ck["text"].(string)
|
|
// Only the merged body group carries both bodies.
|
|
if !strings.Contains(text, "body one") || !strings.Contains(text, "body two") {
|
|
continue
|
|
}
|
|
// (a) parser tags must be stripped from the text.
|
|
if strings.Contains(text, "@@") {
|
|
t.Errorf("parser position tags leaked into chunk text: %q", text)
|
|
}
|
|
// (b) positions must be merged across both records.
|
|
pos, ok := ck["positions"].([][]float64)
|
|
if !ok {
|
|
t.Fatalf("positions missing or wrong type %T on merged group chunk", ck["positions"])
|
|
}
|
|
if len(pos) != 2 {
|
|
t.Errorf("merged positions = %d groups, want 2 (both records)", len(pos))
|
|
}
|
|
return
|
|
}
|
|
t.Fatal("merged body group chunk not found in output")
|
|
}
|
|
|
|
func TestGroupTitleChunker_InvokeDeterministic(t *testing.T) {
|
|
c, err := NewGroupTitleChunker(map[string]any{
|
|
"levels": [][]string{
|
|
{`^# `},
|
|
{`^## `},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewGroupTitleChunker: %v", err)
|
|
}
|
|
inputs := map[string]any{
|
|
"name": "doc.md",
|
|
"text": "# A\nbody a1\nbody a2\n# B\nbody b1\nbody b2",
|
|
}
|
|
var firstLen int
|
|
var firstTexts []string
|
|
for run := 0; run < 10; run++ {
|
|
out, err := c.Invoke(t.Context(), nil, inputs)
|
|
if err != nil {
|
|
t.Fatalf("Invoke run %d: %v", run, err)
|
|
}
|
|
chunks, _ := out["chunks"].([]map[string]any)
|
|
texts := make([]string, len(chunks))
|
|
for i, ck := range chunks {
|
|
texts[i], _ = ck["text"].(string)
|
|
}
|
|
if run == 0 {
|
|
firstLen = len(chunks)
|
|
firstTexts = texts
|
|
continue
|
|
}
|
|
if firstLen != len(chunks) {
|
|
t.Fatalf("run %d: chunk count changed (%d vs %d)", run, len(chunks), firstLen)
|
|
}
|
|
for i := range chunks {
|
|
if firstTexts[i] != texts[i] {
|
|
t.Fatalf("run %d: chunk[%d] text changed", run, i)
|
|
}
|
|
}
|
|
}
|
|
}
|