ragflow/internal/ingestion/component/pdf_vision_dispatch_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 component

import (
	"context"
	"testing"

	"ragflow/internal/dao"
	"ragflow/internal/entity"
	modelModule "ragflow/internal/entity/models"
	"ragflow/internal/utility"

	"gorm.io/gorm"
)

func TestMaybeDispatchPDFVisionEnhancementForwardsDatasetLanguage(t *testing.T) {
	origResolver := resolveTenantModelByType
	origInvoker := visionChatInvoker
	origPrompt := figureVisionPromptBuilder
	t.Cleanup(func() {
		resolveTenantModelByType = origResolver
		visionChatInvoker = origInvoker
		figureVisionPromptBuilder = origPrompt
	})

	resolveTenantModelByType = func(context.Context, *gorm.DB, string, entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
		return &docxVisionFakeDriver{}, "pdf-vision-model", &modelModule.APIConfig{}, 0, nil
	}
	invoker := &docxVisionCaptureInvoker{}
	visionChatInvoker = invoker.invoke
	capturedLanguage := ""
	figureVisionPromptBuilder = func(_, _, language string) (string, error) {
		capturedLanguage = language
		return "describe the figure", nil
	}

	dispatched := parserDispatchResult{
		OutputFormat: "json",
		DocType:      "pdf",
		JSON: []map[string]any{
			{"text": "caption", "image": "data:image/png;base64,aW1hZ2U=", "doc_type_kwd": "image"},
		},
	}

	res, modified, err := maybeDispatchPDFVisionEnhancement(
		t.Context(),
		dao.DB,
		utility.FileTypePDF,
		dispatched,
		map[string]any{"tenant_id": "t1", "lang": "Dutch"},
	)
	if err != nil {
		t.Fatalf("maybeDispatchPDFVisionEnhancement: %v", err)
	}
	if !modified {
		t.Fatal("modified = false, want true")
	}
	if capturedLanguage != "Dutch" {
		t.Fatalf("figure prompt language = %q, want Dutch", capturedLanguage)
	}
	if got := res.JSON[0]["text"]; got != "caption\na diagram of a pipeline" {
		t.Fatalf("enhanced text = %q", got)
	}
}

// TestIsNamedPDFParseMethodWhitelistAligned verifies that the runtime
// "named parse_method" classifier agrees with (*ParserComponent).Check()'s
// PDF whitelist (parser.go:200-203):
//
//	deepdoc, plain_text, mineru, docling,
//	opendataloader, tcadp parser, paddleocr, somark
//
// Diff 2.10: a parse_method that Check() rejects must NOT be treated as a
// recognized named method by isNamedPDFParseMethod — otherwise it silently
// falls through to the CustomVLM vision path instead of failing fast at
// construction (and Python would have rejected it outright).
func TestIsNamedPDFParseMethodWhitelistAligned(t *testing.T) {
	// Values that MUST be recognized (subset of the Check() whitelist,
	// case-insensitive).
	named := []string{
		"deepdoc", "plain_text", "mineru", "docling",
		"opendataloader", "tcadp parser", "paddleocr", "somark",
		"DeepDoc", "PLAIN_TEXT", "MinerU", "DocLing",
		"OpenDataLoader", "TCADP Parser", "PaddleOCR", "SoMark",
	}
	for _, v := range named {
		if !isNamedPDFParseMethod(v) {
			t.Errorf("isNamedPDFParseMethod(%q) = false, want true (in Check() whitelist)", v)
		}
	}

	// Values that MUST NOT be recognized. These either duplicate the
	// whitelist with non-canonical spelling ("plain text"/"plaintext")
	// or are bare-family abbreviations ("tcadp") that Check() does not
	// accept, so they should be funneled to the CustomVLM path (or fail
	// construction) rather than masquerading as a named method.
	notNamed := []string{
		"plain text", "plaintext", "tcadp",
		"CustomVLM", "some_vlm", "gpt-4o",
		"", "  ",
	}
	for _, v := range notNamed {
		if isNamedPDFParseMethod(v) {
			t.Errorf("isNamedPDFParseMethod(%q) = true, want false (not in Check() whitelist)", v)
		}
	}
}

// TestIsNamedPDFParseMethodLayoutSuffixes verifies that "@"-suffixed
// layout_recognizer spellings are NOT treated as named parse methods. They
// are layout_recognizer selectors (resolved separately at
// pdf_vision_dispatch.go:62-68), and Check() rejects them as parse_method,
// so they must fall through to the CustomVLM/VLM path — consistent with the
// (*ParserComponent).Check() whitelist (parser.go:200-203).
func TestIsNamedPDFParseMethodLayoutSuffixes(t *testing.T) {
	suffixed := []string{
		"foo@mineru", "@mineru",
		"foo@paddleocr", "@paddleocr",
		"foo@somark", "@somark",
		"foo@opendataloader", "@opendataloader",
	}
	for _, v := range suffixed {
		if isNamedPDFParseMethod(v) {
			t.Errorf("isNamedPDFParseMethod(%q) = true, want false (layout_recognizer selector, not a named parse_method)", v)
		}
	}

	// An unknown suffix is also not a named method.
	if isNamedPDFParseMethod("foo@unknown") {
		t.Errorf("isNamedPDFParseMethod(%q) = true, want false", "foo@unknown")
	}
}