ragflow/internal/parser/parser/pdf_postprocess_test.go

package parser

import (
	"testing"

	deepdoctype "ragflow/internal/deepdoc/parser/type"
)

func makePDFSection(text, layout string, page int, left, right, top, bottom float64) deepdoctype.Section {
	return deepdoctype.Section{
		Text:       text,
		LayoutType: layout,
		Positions: []deepdoctype.Position{{
			PageNumbers: []int{page},
			Left:        left,
			Right:       right,
			Top:         top,
			Bottom:      bottom,
		}},
	}
}

func TestApplyPDFPostProcess_NormalizesLayoutTypes(t *testing.T) {
	result := &deepdoctype.ParseResult{
		Sections: []deepdoctype.Section{
			{Text: "a", LayoutType: ""},
			{Text: "b", LayoutType: "  "},
			{Text: "c", LayoutType: "table"},
			{Text: "d", LayoutType: "  figure  "},
		},
	}
	applyPDFPostProcess(result, pdfPostProcessOptions{})
	want := []string{"text", "text", "table", "figure"}
	for i, s := range result.Sections {
		if s.LayoutType != want[i] {
			t.Fatalf("Sections[%d].LayoutType = %q, want %q", i, s.LayoutType, want[i])
		}
	}
}

func TestApplyPDFPostProcess_AssignsDocTypeKeywords(t *testing.T) {
	result := &deepdoctype.ParseResult{
		Sections: []deepdoctype.Section{
			{Text: "a", LayoutType: "table"},
			{Text: "b", LayoutType: "figure"},
			{Text: "c", LayoutType: "text"},
			{Text: "d", LayoutType: "", Image: "abc"},
		},
	}
	applyPDFPostProcess(result, pdfPostProcessOptions{})
	// doc_type_kwd is derived from layout type only. A pre-set Image no
	// longer reclassifies a section as "image" — cropping happens lazily
	// at Markdown serialization / chunk time (see pdf_parser_common.go).
	want := []string{"table", "image", "text", "text"}
	for i, s := range result.Sections {
		if s.DocTypeKwd != want[i] {
			t.Fatalf("Sections[%d].DocTypeKwd = %q, want %q", i, s.DocTypeKwd, want[i])
		}
	}
}

func TestApplyPDFPostProcess_FlattenMediaKeepsImagesButMarksText(t *testing.T) {
	result := &deepdoctype.ParseResult{
		Sections: []deepdoctype.Section{
			{Text: "a", LayoutType: "figure", Image: "abc"},
			{Text: "b", LayoutType: "table"},
		},
	}
	applyPDFPostProcess(result, pdfPostProcessOptions{flattenMediaToText: true})
	for i, s := range result.Sections {
		if s.DocTypeKwd != "text" {
			t.Fatalf("Sections[%d].DocTypeKwd = %q, want text", i, s.DocTypeKwd)
		}
	}
	if got, want := result.Sections[0].Image, "abc"; got != want {
		t.Fatalf("Sections[0].Image = %q, want %q", got, want)
	}
}

func TestApplyPDFPostProcess_HeaderFooterFilteringIsOptional(t *testing.T) {
	result := &deepdoctype.ParseResult{
		Sections: []deepdoctype.Section{
			{Text: "header", LayoutType: "header"},
			{Text: "body", LayoutType: "text"},
		},
	}
	applyPDFPostProcess(result, pdfPostProcessOptions{})
	if len(result.Sections) != 2 {
		t.Fatalf("len(Sections) = %d, want 2 when removeHeaderFooter is false", len(result.Sections))
	}

	applyPDFPostProcess(result, pdfPostProcessOptions{removeHeaderFooter: true})
	if len(result.Sections) != 1 {
		t.Fatalf("len(Sections) = %d, want 1 when removeHeaderFooter is true", len(result.Sections))
	}
	if got, want := result.Sections[0].Text, "body"; got != want {
		t.Fatalf("remaining section = %q, want %q", got, want)
	}
}

func TestApplyPDFPostProcess_RemoveTOCByOutlines(t *testing.T) {
	result := &deepdoctype.ParseResult{
		Sections: []deepdoctype.Section{
			makePDFSection("目录", "text", 1, 50, 550, 100, 120),
			makePDFSection("章节列表", "text", 2, 50, 550, 120, 140),
			makePDFSection("正文", "text", 3, 50, 550, 100, 120),
		},
		Outlines: []deepdoctype.Outline{
			{Title: "目录", Level: 0, PageNumber: 1},
			{Title: "第一章", Level: 0, PageNumber: 3},
		},
	}
	applyPDFPostProcess(result, pdfPostProcessOptions{removeTOC: true})
	if len(result.Sections) != 1 {
		t.Fatalf("len(Sections) = %d, want 1", len(result.Sections))
	}
	if got, want := result.Sections[0].Text, "正文"; got != want {
		t.Fatalf("remaining section = %q, want %q", got, want)
	}
}

func TestApplyPDFPostProcess_ReordersMultiColumnText(t *testing.T) {
	cases := []struct {
		name      string
		pageWidth float64
		zoom      float64
	}{
		{name: "unit zoom", pageWidth: 600, zoom: 1},
		{name: "pre-normalized width", pageWidth: 200, zoom: 3},
	}
	for _, tc := range cases {
		result := &deepdoctype.ParseResult{
			Sections: []deepdoctype.Section{
				makePDFSection("right", "text", 0, 100, 166, 100, 120),
				makePDFSection("left", "text", 0, 10, 76, 100, 120),
			},
		}
		applyPDFPostProcess(result, pdfPostProcessOptions{pageWidth: tc.pageWidth, zoom: tc.zoom, enableMultiColumn: true})
		if got, want := result.Sections[0].Text, "left"; got != want {
			t.Fatalf("%s: Sections[0].Text = %q, want %q", tc.name, got, want)
		}
	}
}

// TestFilterPDFHeaderFooter_SubstringMatch pins #5: Python's remove_header_footer
// uses a substring match re.search(r"(header|footer|number)", ...) (rag/flow/parser/parser.py:754),
// while Go used an anchored exact match ^(header|footer|number)$. A layout type
// that merely CONTAINS one of those words (e.g. "page-footer") must be stripped to
// match Python, not silently kept.
func TestFilterPDFHeaderFooter_SubstringMatch(t *testing.T) {
	result := &deepdoctype.ParseResult{
		Sections: []deepdoctype.Section{
			{Text: "real header", LayoutType: "header"},
			{Text: "real footer", LayoutType: "footer"},
			{Text: "page 1", LayoutType: "number"},
			{Text: "a page-footer note", LayoutType: "page-footer"}, // composite -> substring match
			{Text: "body text", LayoutType: "text"},
		},
	}
	filterPDFHeaderFooter(result)

	kept := map[string]bool{}
	for _, s := range result.Sections {
		kept[s.LayoutType] = true
	}
	for _, lt := range []string{"header", "footer", "number"} {
		if kept[lt] {
			t.Errorf("#5 header/footer: %q should be stripped", lt)
		}
	}
	// Composite "page-footer" must be stripped by substring match (Python-equivalent).
	if kept["page-footer"] {
		t.Errorf("#5 header/footer: composite layout type %q should be stripped by substring match", "page-footer")
	}
	if !kept["text"] {
		t.Errorf("#5 header/footer: body text %q should be kept", "text")
	}
}