zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-21 07:01:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

Align Go parser backends and PDF pipeline with Python (#16676)

Browse Source 
Ports remaining Go parser wiring and PDF backends, adds tenant-aware VLM
dispatch, aligns post-processing with Python, and adds end-to-end
pipeline coverage with a generated six-page PDF.
This commit is contained in:
Zhichang Yu
2026-07-06 19:50:54 +08:00
committed by GitHub
parent 3044283442
commit 55af1d70f3
36 changed files with 4918 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
internal/ingestion/component/parser.go
View File

@@ -84,6 +84,7 @@ import (
"ragflow/internal/agent/runtime"
"ragflow/internal/ingestion/component/schema"
"ragflow/internal/utility"
)
const ComponentNameParser = "Parser"
@@ -291,8 +292,19 @@ func (c *ParserComponent) Invoke(ctx context.Context, inputs map[string]any) (ma
}
}
dispatched := dispatchParse(fileTypeExt, filename, binary, c.Param.Setups)
dispatched = hydrateEmptyDispatchPayload(dispatched, binary)
dispatched, handledVision, visionErr := maybeDispatchPDFVision(fileTypeExt, filename, binary, inputs, c.Param.Setups)
if visionErr != nil {
return nil, visionErr
}
if !handledVision {
dispatched = dispatchParse(fileTypeExt, filename, binary, c.Param.Setups)
dispatched = hydrateEmptyDispatchPayload(dispatched, binary)
}
// Known/supported families must fail loudly when dispatch or
// parsing breaks. Only unknown families keep the raw-text fallback.
if dispatched.Err != nil && fileTypeExt != utility.FileTypeOTHER {
return nil, dispatched.Err
}
// 3. Build the legacy `pages` slice. When the dispatch path
// produced a JSON payload, we re-shape it into the page

28
internal/ingestion/component/parser_dispatch.go
View File

@@ -55,6 +55,30 @@ type parserDispatchResult struct {
Err error
}
type parserSetupConfigurer interface {
ConfigureFromSetup(setup map[string]any)
}
func resolveParserFamily(fileType utility.FileType) string {
if family := pythonFamilyName(string(fileType)); family != "" {
return family
}
return string(fileType)
}
func configureParserFromSetups(p any, fileType utility.FileType, setups map[string]schema.ParserSetup) {
cfg, ok := p.(parserSetupConfigurer)
if !ok {
return
}
family := resolveParserFamily(fileType)
setup, ok := setups[family]
if !ok {
return
}
cfg.ConfigureFromSetup(map[string]any(setup))
}
// resolveOutputFormat picks the wire format for this run. The
// Python side asks the setup, then checks the value is in
// allowed_output_format[fileType]. We mirror that exact sequence:
@@ -108,7 +132,8 @@ func resolveOutputFormat(family string, setups map[string]schema.ParserSetup, al
// tell the difference between "explicit OCR" and "default DeepDOC"
// without re-reading setups.
func resolveLibType(fileType utility.FileType, setups map[string]schema.ParserSetup) (libType, parseMethod string) {
setup, ok := setups[string(fileType)]
family := resolveParserFamily(fileType)
setup, ok := setups[family]
if !ok {
return "", ""
}
@@ -149,6 +174,7 @@ func dispatchParse(fileType utility.FileType, filename string, data []byte, setu
if err != nil {
return parserDispatchResult{Err: fmt.Errorf("Parser: resolve %q: %w", fileType, err)}
}
configureParserFromSetups(p, fileType, setups)
res := p.ParseWithResult(filename, data)
if res.Err != nil {

99
internal/ingestion/component/parser_dispatch_pdf_vision_cgo_test.go Normal file
View File

@@ -0,0 +1,99 @@
//go:build cgo
package component
import (
"context"
"os"
"path/filepath"
"strings"
"sync/atomic"
"testing"
modelModule "ragflow/internal/entity/models"
"ragflow/internal/ingestion/component/schema"
)
func TestDispatch_PDFVisionJSON_RealPDFFixture(t *testing.T) {
origPromptLoader := pdfVisionPromptLoader
origResolver := pdfVisionModelResolver
origInvoker := pdfVisionChatInvoker
t.Cleanup(func() {
pdfVisionPromptLoader = origPromptLoader
pdfVisionModelResolver = origResolver
pdfVisionChatInvoker = origInvoker
})
pdfVisionPromptLoader = func(name string) (string, error) {
return "Describe page {{ page }}.", nil
}
pdfVisionModelResolver = func(tenantID string, modelID string) (modelModule.ModelDriver, string, *modelModule.APIConfig, error) {
if tenantID != "tenant-vision" || modelID != "CustomVLM" {
t.Fatalf("resolver got tenant/model %q/%q", tenantID, modelID)
}
return nil, "fixture-vlm", nil, nil
}
var callCount atomic.Int32
pdfVisionChatInvoker = func(_ modelModule.ModelDriver, modelName string, messages []modelModule.Message, _ *modelModule.APIConfig) (*modelModule.ChatResponse, error) {
if modelName != "fixture-vlm" {
t.Fatalf("modelName = %q, want fixture-vlm", modelName)
}
callCount.Add(1)
content, ok := messages[0].Content.([]interface{})
if !ok || len(content) != 2 {
t.Fatalf("messages[0].Content = %#v, want multimodal payload", messages[0].Content)
}
promptBlock, ok := content[0].(map[string]any)
if !ok {
t.Fatalf("prompt block = %T, want map[string]any", content[0])
}
prompt, _ := promptBlock["text"].(string)
answer := "Recognized " + prompt + "\n\n--- Page ---"
return &modelModule.ChatResponse{Answer: &answer}, nil
}
path := filepath.Join("..", "..", "..", "test", "benchmark", "test_docs", "Doc1.pdf")
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("ReadFile(%s): %v", path, err)
}
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "CustomVLM"
param.Setups["pdf"]["output_format"] = "json"
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": data,
"file_type": "pdf",
"name": "Doc1.pdf",
"tenant_id": "tenant-vision",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
file, ok := out["file"].(map[string]any)
if !ok {
t.Fatalf("file metadata missing: %T", out["file"])
}
pageCount, ok := file["page_count"].(int)
if !ok || pageCount < 1 {
t.Fatalf("file.page_count = %#v, want positive int", file["page_count"])
}
if got := int(callCount.Load()); got != pageCount {
t.Fatalf("vision call count = %d, want %d", got, pageCount)
}
jsonItems, ok := out["json"].([]map[string]any)
if !ok || len(jsonItems) == 0 {
t.Fatalf("json payload missing or empty: %T", out["json"])
}
if got, _ := jsonItems[0]["text"].(string); !strings.Contains(got, "Recognized Describe page 1.") {
t.Fatalf("json[0].text = %q, want rendered vision answer", got)
}
if positions, ok := jsonItems[0]["_pdf_positions"].([][]any); !ok || len(positions) == 0 {
t.Fatalf("json[0]._pdf_positions = %#v, want normalized page positions", jsonItems[0]["_pdf_positions"])
}
}

554
internal/ingestion/component/parser_dispatch_test.go
View File

@@ -26,13 +26,31 @@
package component
import (
"archive/zip"
"bytes"
"context"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"reflect"
"strings"
"testing"
modelModule "ragflow/internal/entity/models"
"ragflow/internal/ingestion/component/schema"
"ragflow/internal/utility"
)
type captureSetupConfigurer struct {
setup map[string]any
}
func (c *captureSetupConfigurer) ConfigureFromSetup(setup map[string]any) {
c.setup = setup
}
// TestDispatch_OutputFormatValidation_Allowed is the happy-path
// pin: a Markdown file with output_format=json passes the
// allowed_output_format check and runs the structured dispatch.
@@ -64,6 +82,11 @@ func TestDispatch_OutputFormatValidation_Allowed(t *testing.T) {
if !ok || len(pages) == 0 {
t.Errorf("pages slice missing or empty: %T", out["pages"])
}
if ok && len(pages) > 0 {
if got, _ := pages[0]["text"].(string); !strings.Contains(got, "Title") {
t.Errorf("pages[0].text = %q, want content containing Title", got)
}
}
// File metadata is carried through dispatch.
if fm, ok := out["file"].(map[string]any); !ok || fm["name"] != "doc.md" {
t.Errorf("file metadata missing or wrong: %+v", out["file"])
@@ -130,22 +153,23 @@ func TestDispatch_TextPageMode_NoFileType(t *testing.T) {
}
}
// TestDispatch_TextPageMode_PDFInput pins the current text-page
// behavior when PDF bytes are not successfully parsed into a
// structured payload in this test environment.
func TestDispatch_TextPageMode_PDFInput(t *testing.T) {
// TestDispatch_SupportedFamilyFailure_HardErrors pins the agreed
// migration rule: once a supported family is identified, parser
// resolution/execution failures must surface as errors instead of
// silently degrading to text-page mode.
func TestDispatch_SupportedFamilyFailure_HardErrors(t *testing.T) {
param := schema.ParserParam{}.Defaults()
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
_, err := c.Invoke(context.Background(), map[string]any{
"binary": []byte("PDF payload as bytes (not a real PDF — stub test)\n"),
"file_type": "pdf",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
if err == nil {
t.Fatal("Invoke: want error for supported family parse failure, got nil")
}
if got, want := out["output_format"], "text"; got != want {
t.Errorf("output_format = %v, want %v (PDF input stayed in text-page mode)", got, want)
if !strings.Contains(err.Error(), "pdf") {
t.Errorf("error %q must mention pdf", err.Error())
}
}
@@ -251,3 +275,515 @@ func TestResolveOutputFormat_DefaultsAndWhitelist(t *testing.T) {
})
}
}
func TestConfigureParserFromSetups_UsesPythonFamilySetup(t *testing.T) {
setups := schema.ParserParam{}.Defaults().Setups
got := &captureSetupConfigurer{}
configureParserFromSetups(got, utility.FileTypePDF, setups)
want := map[string]any(setups["pdf"])
if !reflect.DeepEqual(got.setup, want) {
t.Fatalf("ConfigureFromSetup got %+v, want %+v", got.setup, want)
}
}
func TestDispatch_PDFMarkdown_UsesConfiguredOutputFormat(t *testing.T) {
t.Setenv("DEEPDOC_URL", "")
t.Setenv("OSSDEEPDOC_URL", "")
path := filepath.Join("..", "..", "..", "test", "benchmark", "test_docs", "Doc1.pdf")
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("ReadFile(%s): %v", path, err)
}
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["output_format"] = "markdown"
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": data,
"file_type": "pdf",
"name": "Doc1.pdf",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
if got, want := out["output_format"], "markdown"; got != want {
t.Fatalf("output_format = %v, want %v", got, want)
}
md, ok := out["markdown"].(string)
if !ok || md == "" {
t.Fatalf("markdown payload missing or empty: %T", out["markdown"])
}
if _, ok := out["json"]; ok {
t.Fatalf("json payload must be absent for markdown output: %+v", out["json"])
}
}
func TestDispatch_PDFPlainText_UsesConfiguredBackend(t *testing.T) {
path := filepath.Join("..", "..", "..", "test", "benchmark", "test_docs", "Doc1.pdf")
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("ReadFile(%s): %v", path, err)
}
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "plain_text"
param.Setups["pdf"]["output_format"] = "json"
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": data,
"file_type": "pdf",
"name": "Doc1.pdf",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
jsonItems, ok := out["json"].([]map[string]any)
if !ok || len(jsonItems) == 0 {
t.Fatalf("json payload missing or empty: %T", out["json"])
}
if got, _ := jsonItems[0]["text"].(string); strings.TrimSpace(got) == "" {
t.Fatalf("json first item text = %q, want non-empty", got)
}
}
func TestDispatch_PDFUnsupportedParseMethod_HardErrors(t *testing.T) {
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "CustomVLM"
c := &ParserComponent{Param: param}
_, err := c.Invoke(context.Background(), map[string]any{
"binary": []byte("%PDF-1.4"),
"file_type": "pdf",
"name": "bad.pdf",
})
if err == nil {
t.Fatal("Invoke: want error for unsupported PDF parse_method, got nil")
}
if !strings.Contains(err.Error(), "parse_method") || !strings.Contains(err.Error(), "tenant_id") {
t.Fatalf("error = %q, want parse_method + tenant_id context", err.Error())
}
}
func TestDispatch_PDFVisionJSON_UsesTenantAwareModel(t *testing.T) {
origPromptLoader := pdfVisionPromptLoader
origRenderer := pdfVisionPageRenderer
origResolver := pdfVisionModelResolver
origInvoker := pdfVisionChatInvoker
t.Cleanup(func() {
pdfVisionPromptLoader = origPromptLoader
pdfVisionPageRenderer = origRenderer
pdfVisionModelResolver = origResolver
pdfVisionChatInvoker = origInvoker
})
var prompts []string
pdfVisionPromptLoader = func(name string) (string, error) {
if name != "vision_llm_describe_prompt" {
return "", fmt.Errorf("unexpected prompt %q", name)
}
return "Describe page {{ page }}.", nil
}
pdfVisionPageRenderer = func(_ []byte) ([]pdfVisionPage, error) {
return []pdfVisionPage{
{PageNumber: 1, WidthPts: 100, HeightPts: 200, ImageURL: "data:image/png;base64,aaa"},
{PageNumber: 2, WidthPts: 120, HeightPts: 240, ImageURL: "data:image/png;base64,bbb"},
}, nil
}
pdfVisionModelResolver = func(tenantID string, modelID string) (modelModule.ModelDriver, string, *modelModule.APIConfig, error) {
if tenantID != "tenant-1" || modelID != "CustomVLM" {
return nil, "", nil, fmt.Errorf("resolver got tenant/model %q/%q", tenantID, modelID)
}
return nil, "resolved-vlm", nil, nil
}
pdfVisionChatInvoker = func(_ modelModule.ModelDriver, modelName string, messages []modelModule.Message, _ *modelModule.APIConfig) (*modelModule.ChatResponse, error) {
if modelName != "resolved-vlm" {
return nil, fmt.Errorf("modelName = %q, want resolved-vlm", modelName)
}
if len(messages) != 1 {
return nil, fmt.Errorf("messages len = %d, want 1", len(messages))
}
content, ok := messages[0].Content.([]interface{})
if !ok || len(content) != 2 {
return nil, fmt.Errorf("content = %#v, want multimodal prompt+image", messages[0].Content)
}
block, ok := content[0].(map[string]any)
if !ok {
return nil, fmt.Errorf("content[0] = %T, want map[string]any", content[0])
}
prompt, _ := block["text"].(string)
prompts = append(prompts, prompt)
answer := "Transcribed " + prompt
return &modelModule.ChatResponse{Answer: &answer}, nil
}
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "CustomVLM"
param.Setups["pdf"]["output_format"] = "json"
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": []byte("%PDF-1.4"),
"file_type": "pdf",
"name": "vision.pdf",
"tenant_id": "tenant-1",
})
if err == nil {
jsonItems, ok := out["json"].([]map[string]any)
if !ok || len(jsonItems) != 2 {
t.Fatalf("json payload = %#v, want 2 items", out["json"])
}
if got, want := jsonItems[0]["page_number"], 1; got != want {
t.Fatalf("json[0].page_number = %v, want %v", got, want)
}
if positions, ok := jsonItems[0]["_pdf_positions"].([][]any); !ok || len(positions) != 1 {
t.Fatalf("json[0]._pdf_positions = %#v, want one normalized page box", jsonItems[0]["_pdf_positions"])
}
if file, ok := out["file"].(map[string]any); !ok || file["parse_method"] != "CustomVLM" || file["page_count"] != 2 {
t.Fatalf("file metadata = %#v, want parse_method/page_count", out["file"])
}
if len(prompts) != 2 || !strings.Contains(prompts[0], "page 1") || !strings.Contains(prompts[1], "page 2") {
t.Fatalf("prompts = %#v, want rendered page-specific prompts", prompts)
}
return
}
t.Fatalf("Invoke: %v", err)
}
func TestDispatch_PDFVisionJSON_PreservesEmptyPages(t *testing.T) {
origPromptLoader := pdfVisionPromptLoader
origRenderer := pdfVisionPageRenderer
origResolver := pdfVisionModelResolver
origInvoker := pdfVisionChatInvoker
t.Cleanup(func() {
pdfVisionPromptLoader = origPromptLoader
pdfVisionPageRenderer = origRenderer
pdfVisionModelResolver = origResolver
pdfVisionChatInvoker = origInvoker
})
pdfVisionPromptLoader = func(string) (string, error) { return "Describe page {{ page }}.", nil }
pdfVisionPageRenderer = func(_ []byte) ([]pdfVisionPage, error) {
return []pdfVisionPage{
{PageNumber: 1, WidthPts: 100, HeightPts: 200, ImageURL: "data:image/png;base64,aaa"},
{PageNumber: 2, WidthPts: 120, HeightPts: 240, ImageURL: "data:image/png;base64,bbb"},
}, nil
}
pdfVisionModelResolver = func(string, string) (modelModule.ModelDriver, string, *modelModule.APIConfig, error) {
return nil, "resolved-vlm", nil, nil
}
call := 0
pdfVisionChatInvoker = func(_ modelModule.ModelDriver, _ string, _ []modelModule.Message, _ *modelModule.APIConfig) (*modelModule.ChatResponse, error) {
call++
answer := ""
if call == 1 {
answer = "First page"
}
return &modelModule.ChatResponse{Answer: &answer}, nil
}
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "CustomVLM"
param.Setups["pdf"]["output_format"] = "json"
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": []byte("%PDF-1.4"),
"file_type": "pdf",
"name": "vision.pdf",
"tenant_id": "tenant-1",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
jsonItems, ok := out["json"].([]map[string]any)
if !ok || len(jsonItems) != 2 {
t.Fatalf("json payload = %#v, want 2 items", out["json"])
}
if got := jsonItems[1]["text"]; got != "" {
t.Fatalf("json[1].text = %#v, want empty string placeholder", got)
}
}
func TestDispatch_PDFMinerUMarkdown_UsesConfiguredBackend(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch {
case r.Method == http.MethodPost && r.URL.Path == "/file_parse":
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"data":{"task_id":"task-3"}}`))
case r.Method == http.MethodGet && r.URL.Path == "/tasks/task-3/result":
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"results":{"doc":{"md_content":"# Title\n\nBody\n"}}}`))
default:
http.NotFound(w, r)
}
}))
defer server.Close()
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "MinerU"
param.Setups["pdf"]["output_format"] = "markdown"
param.Setups["pdf"]["mineru_apiserver"] = server.URL
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": []byte("%PDF-1.4"),
"file_type": "pdf",
"name": "sample.pdf",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
if got, want := out["output_format"], "markdown"; got != want {
t.Fatalf("output_format = %v, want %v", got, want)
}
md, ok := out["markdown"].(string)
if !ok || !strings.Contains(md, "Title") {
t.Fatalf("markdown payload = %#v, want Title content", out["markdown"])
}
}
func TestDispatch_PDFPaddleOCRMarkdown_UsesConfiguredBackend(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost || r.URL.Path != "/layout-parsing" {
http.NotFound(w, r)
return
}
if got, want := r.Header.Get("Authorization"), "Bearer paddle-secret"; got != want {
t.Errorf("Authorization = %q, want %q", got, want)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"errorCode":0,"result":{"layoutParsingResults":[{"markdown":{"text":"# Paddle Title\n\nPaddle body.\n"}}]}}`))
}))
defer server.Close()
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "PaddleOCR"
param.Setups["pdf"]["output_format"] = "markdown"
param.Setups["pdf"]["paddleocr_base_url"] = server.URL
param.Setups["pdf"]["paddleocr_api_key"] = "paddle-secret"
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": []byte("%PDF-1.4"),
"file_type": "pdf",
"name": "sample.pdf",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
if got, want := out["output_format"], "markdown"; got != want {
t.Fatalf("output_format = %v, want %v", got, want)
}
md, ok := out["markdown"].(string)
if !ok || !strings.Contains(md, "Paddle Title") {
t.Fatalf("markdown payload = %#v, want Paddle Title content", out["markdown"])
}
}
func TestDispatch_PDFDoclingMarkdown_UsesConfiguredBackend(t *testing.T) {
var requestCount int
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requestCount++
if got, want := r.Header.Get("Authorization"), "Bearer doc-secret"; got != want {
t.Errorf("Authorization = %q, want %q", got, want)
return
}
switch {
case r.Method == http.MethodPost && r.URL.Path == "/v1/convert/source" && requestCount == 1:
w.WriteHeader(http.StatusUnprocessableEntity)
_, _ = w.Write([]byte(`{"detail":"chunking unsupported"}`))
case r.Method == http.MethodPost && r.URL.Path == "/v1alpha/convert/source" && requestCount == 2:
w.WriteHeader(http.StatusUnprocessableEntity)
_, _ = w.Write([]byte(`{"detail":"chunking unsupported"}`))
case r.Method == http.MethodPost && r.URL.Path == "/v1/convert/source" && requestCount == 3:
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"document":{"md_content":"# Docling Title\n\nDocling body.\n"}}`))
default:
http.NotFound(w, r)
}
}))
defer server.Close()
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "Docling"
param.Setups["pdf"]["output_format"] = "markdown"
param.Setups["pdf"]["docling_server_url"] = server.URL
param.Setups["pdf"]["docling_api_key"] = "doc-secret"
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": []byte("%PDF-1.4"),
"file_type": "pdf",
"name": "sample.pdf",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
if got, want := out["output_format"], "markdown"; got != want {
t.Fatalf("output_format = %v, want %v", got, want)
}
md, ok := out["markdown"].(string)
if !ok || !strings.Contains(md, "Docling Title") {
t.Fatalf("markdown payload = %#v, want Docling Title content", out["markdown"])
}
if got, want := requestCount, 3; got != want {
t.Fatalf("requestCount = %d, want %d", got, want)
}
}
func TestDispatch_PDFOpenDataLoaderMarkdown_UsesConfiguredBackend(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost || r.URL.Path != "/file_parse" {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"json_doc":null,"md_text":"# ODL Title\n\nODL body.\n"}`))
}))
defer server.Close()
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "OpenDataLoader"
param.Setups["pdf"]["output_format"] = "markdown"
param.Setups["pdf"]["opendataloader_apiserver"] = server.URL
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": []byte("%PDF-1.4"),
"file_type": "pdf",
"name": "sample.pdf",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
md, ok := out["markdown"].(string)
if !ok || !strings.Contains(md, "ODL Title") {
t.Fatalf("markdown payload = %#v, want ODL Title", out["markdown"])
}
}
func TestDispatch_PDFSoMarkMarkdown_UsesConfiguredBackend(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/parse/async":
_, _ = w.Write([]byte(`{"code":0,"data":{"task_id":"task-4"}}`))
case "/parse/async_check":
_, _ = w.Write([]byte(`{"code":0,"data":{"status":"SUCCESS","result":{"outputs":{"json":{"pages":[{"blocks":[{"type":"title","content":"SoMark Title","title_level":1},{"type":"text","content":"Body"}]}]}}}}}`))
default:
http.NotFound(w, r)
}
}))
defer server.Close()
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "SoMark"
param.Setups["pdf"]["output_format"] = "markdown"
param.Setups["pdf"]["somark_base_url"] = server.URL
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": []byte("%PDF-1.4"),
"file_type": "pdf",
"name": "sample.pdf",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
md, ok := out["markdown"].(string)
if !ok || !strings.Contains(md, "SoMark Title") {
t.Fatalf("markdown payload = %#v, want SoMark Title", out["markdown"])
}
}
func TestDispatch_PDFTCADPMarkdown_UsesConfiguredBackend(t *testing.T) {
zipPayload := tcadpZipFixtureForComponent(t)
var server *httptest.Server
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/reconstruct_document":
_, _ = w.Write([]byte(`{"DocumentRecognizeResultUrl":"` + server.URL + `/download.zip"}`))
case "/download.zip":
_, _ = w.Write(zipPayload)
default:
http.NotFound(w, r)
}
}))
defer server.Close()
param := schema.ParserParam{}.Defaults()
param.Setups["pdf"]["parse_method"] = "TCADP parser"
param.Setups["pdf"]["output_format"] = "markdown"
param.Setups["pdf"]["tcadp_apiserver"] = server.URL
c := &ParserComponent{Param: param}
out, err := c.Invoke(context.Background(), map[string]any{
"binary": []byte("%PDF-1.4"),
"file_type": "pdf",
"name": "sample.pdf",
})
if err != nil {
t.Fatalf("Invoke: %v", err)
}
md, ok := out["markdown"].(string)
if !ok || !strings.Contains(md, "Hello TCADP") {
t.Fatalf("markdown payload = %#v, want Hello TCADP", out["markdown"])
}
}
func tcadpZipFixtureForComponent(t *testing.T) []byte {
t.Helper()
var buf bytes.Buffer
writer := zip.NewWriter(&buf)
f1, err := writer.Create("result.md")
if err != nil {
t.Fatalf("Create md: %v", err)
}
_, _ = f1.Write([]byte("Hello TCADP"))
if err := writer.Close(); err != nil {
t.Fatalf("Close zip: %v", err)
}
return buf.Bytes()
}
func TestResolveLibType_UsesOwningFamilySetup(t *testing.T) {
setups := schema.ParserParam{}.Defaults().Setups
setups["slides"]["lib_type"] = "office_oxide"
setups["slides"]["parse_method"] = "deepdoc"
setups["spreadsheet"]["lib_type"] = "office_oxide"
setups["spreadsheet"]["parse_method"] = "deepdoc"
cases := []struct {
name string
fileType utility.FileType
wantLibType string
wantParseMethod string
}{
{
name: "pptx resolves from slides family",
fileType: utility.FileTypePPTX,
wantLibType: "office_oxide",
wantParseMethod: "deepdoc",
},
{
name: "xlsx resolves from spreadsheet family",
fileType: utility.FileTypeXLSX,
wantLibType: "office_oxide",
wantParseMethod: "deepdoc",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
gotLibType, gotParseMethod := resolveLibType(tc.fileType, setups)
if gotLibType != tc.wantLibType || gotParseMethod != tc.wantParseMethod {
t.Fatalf("resolveLibType(%q) = (%q, %q), want (%q, %q)",
tc.fileType, gotLibType, gotParseMethod, tc.wantLibType, tc.wantParseMethod)
}
})
}
}

27
internal/ingestion/component/parser_test.go
View File

@@ -450,3 +450,30 @@ func TestParserComponent_Invoke_PageSizeHint(t *testing.T) {
}
}
}
// TestParseBatch_IsFormatAgnostic pins the batching contract:
// parseBatch does not resolve parsers or inspect file families. It
// only wraps already-prepared page bytes into schema.Page items.
func TestParseBatch_IsFormatAgnostic(t *testing.T) {
got, err := parseBatch(context.Background(), [][]byte{
[]byte("first page from dispatch"),
[]byte("<table>second page from html dispatch</table>"),
})
if err != nil {
t.Fatalf("parseBatch: %v", err)
}
if len(got) != 2 {
t.Fatalf("len(got) = %d, want 2", len(got))
}
if got[0]["text"] != "first page from dispatch" {
t.Fatalf("got[0][text] = %v, want first page from dispatch", got[0]["text"])
}
if got[1]["text"] != "<table>second page from html dispatch</table>" {
t.Fatalf("got[1][text] = %v, want HTML payload preserved verbatim", got[1]["text"])
}
for i := range got {
if got[i]["doc_type_kwd"] != "text" {
t.Fatalf("got[%d][doc_type_kwd] = %v, want text", i, got[i]["doc_type_kwd"])
}
}
}

433
internal/ingestion/component/pdf_vision_dispatch.go Normal file
View File

@@ -0,0 +1,433 @@
package component
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
"sync"
"ragflow/internal/dao"
"ragflow/internal/entity"
modelModule "ragflow/internal/entity/models"
"ragflow/internal/ingestion/component/schema"
"ragflow/internal/utility"
"gorm.io/gorm"
)
type pdfVisionPage struct {
PageNumber int
WidthPts float64
HeightPts float64
ImageURL string
}
var (
pdfVisionPromptLoader = loadPDFVisionPrompt
pdfVisionPageRenderer = defaultRenderPDFVisionPages
pdfVisionModelResolver = defaultPDFVisionModelResolver
pdfVisionChatInvoker = defaultPDFVisionChatInvoker
)
var (
pdfVisionPromptCache = make(map[string]string)
pdfVisionPromptCacheMu sync.RWMutex
pdfVisionPromptsBase string
pdfVisionPromptsOnce sync.Once
)
func maybeDispatchPDFVision(
fileType utility.FileType,
filename string,
binary []byte,
inputs map[string]any,
setups map[string]schema.ParserSetup,
) (parserDispatchResult, bool, error) {
if fileType != utility.FileTypePDF {
return parserDispatchResult{}, false, nil
}
setup, ok := setups["pdf"]
if !ok {
return parserDispatchResult{}, false, nil
}
modelID, useVision := resolvePDFVisionModelID(setup)
if !useVision {
return parserDispatchResult{}, false, nil
}
tenantID := getStringOr(inputs, "tenant_id", "")
if tenantID == "" {
return parserDispatchResult{}, true, fmt.Errorf(
`Parser: pdf parse_method %q requires tenant_id to resolve IMAGE2TEXT model`, modelID)
}
res, err := dispatchPDFVision(filename, binary, tenantID, modelID, setup)
if err != nil {
return parserDispatchResult{}, true, err
}
return res, true, nil
}
func resolvePDFVisionModelID(setup schema.ParserSetup) (string, bool) {
if setup == nil {
return "", false
}
if raw, ok := setup["parse_method"].(string); ok {
method := strings.TrimSpace(raw)
if method != "" && !isNamedPDFParseMethod(method) {
return method, true
}
}
if raw, ok := setup["layout_recognizer"].(string); ok {
method := strings.TrimSpace(raw)
if method == "" || strings.EqualFold(method, "plain text") || strings.EqualFold(method, "plaintext") {
return "", false
}
if !isNamedPDFParseMethod(method) {
return method, true
}
}
return "", false
}
func isNamedPDFParseMethod(raw string) bool {
method := strings.ToLower(strings.TrimSpace(raw))
switch {
case strings.HasSuffix(method, "@mineru"),
strings.HasSuffix(method, "@paddleocr"),
strings.HasSuffix(method, "@somark"),
strings.HasSuffix(method, "@opendataloader"):
return true
}
switch method {
case "deepdoc", "plain_text", "plaintext", "mineru", "paddleocr", "docling", "opendataloader", "somark", "tcadp", "tcadp parser":
return true
}
return false
}
func dispatchPDFVision(
filename string,
binary []byte,
tenantID string,
modelID string,
setup schema.ParserSetup,
) (parserDispatchResult, error) {
renderedPages, err := pdfVisionPageRenderer(binary)
if err != nil {
return parserDispatchResult{}, fmt.Errorf("Parser: pdf vision render: %w", err)
}
driver, resolvedModelName, apiConfig, err := pdfVisionModelResolver(tenantID, modelID)
if err != nil {
return parserDispatchResult{}, fmt.Errorf("Parser: pdf vision model %q: %w", modelID, err)
}
promptTemplate, err := pdfVisionPromptLoader("vision_llm_describe_prompt")
if err != nil {
return parserDispatchResult{}, fmt.Errorf("Parser: load vision prompt: %w", err)
}
items := make([]map[string]any, 0, len(renderedPages))
markdownParts := make([]string, 0, len(renderedPages))
for _, page := range renderedPages {
prompt := renderPDFVisionPrompt(promptTemplate, page.PageNumber)
resp, err := pdfVisionChatInvoker(driver, resolvedModelName, buildPDFVisionMessages(prompt, page.ImageURL), apiConfig)
if err != nil {
return parserDispatchResult{}, fmt.Errorf("Parser: pdf vision page %d: %w", page.PageNumber, err)
}
text := extractPDFVisionAnswer(resp)
positions := [][]any{{page.PageNumber, 0.0, page.WidthPts, 0.0, page.HeightPts}}
items = append(items, map[string]any{
"text": text,
"doc_type_kwd": "text",
"page_number": page.PageNumber,
"_pdf_positions": positions,
"positions": positions,
})
if text != "" {
markdownParts = append(markdownParts, text)
}
}
outputFormat := "json"
if v, ok := setup["output_format"].(string); ok && strings.TrimSpace(v) != "" {
outputFormat = strings.ToLower(strings.TrimSpace(v))
}
fileMeta := map[string]any{
"name": filename,
"page_count": len(renderedPages),
"outline": []map[string]any{},
"parse_method": modelID,
}
switch outputFormat {
case "json":
return parserDispatchResult{
OutputFormat: "json",
File: fileMeta,
JSON: items,
}, nil
case "markdown":
return parserDispatchResult{
OutputFormat: "markdown",
File: fileMeta,
Markdown: strings.TrimSpace(strings.Join(markdownParts, "\n\n")),
}, nil
default:
return parserDispatchResult{}, fmt.Errorf("Parser: unsupported PDF output_format %q for vision parse_method %q", outputFormat, modelID)
}
}
func buildPDFVisionMessages(prompt string, imageURL string) []modelModule.Message {
return []modelModule.Message{{
Role: "user",
Content: []interface{}{
map[string]any{"type": "text", "text": prompt},
map[string]any{"type": "image_url", "image_url": map[string]any{"url": imageURL}},
},
}}
}
func extractPDFVisionAnswer(resp *modelModule.ChatResponse) string {
if resp == nil || resp.Answer == nil {
return ""
}
return strings.TrimSpace(*resp.Answer)
}
func defaultPDFVisionModelResolver(
tenantID string,
modelID string,
) (modelModule.ModelDriver, string, *modelModule.APIConfig, error) {
if strings.TrimSpace(modelID) == "" {
driver, modelName, apiConfig, _, err := resolveTenantModelByType(tenantID, entity.ModelTypeImage2Text)
return driver, modelName, apiConfig, err
}
driver, modelName, apiConfig, _, err := resolveModelConfigFromProviderInstance(tenantID, entity.ModelTypeImage2Text, modelID)
return driver, modelName, apiConfig, err
}
func defaultPDFVisionChatInvoker(
driver modelModule.ModelDriver,
modelName string,
messages []modelModule.Message,
apiConfig *modelModule.APIConfig,
) (*modelModule.ChatResponse, error) {
vision := true
return driver.ChatWithMessages(modelName, messages, apiConfig, &modelModule.ChatConfig{Vision: &vision})
}
func loadPDFVisionPrompt(name string) (string, error) {
pdfVisionPromptCacheMu.RLock()
if cached, ok := pdfVisionPromptCache[name]; ok {
pdfVisionPromptCacheMu.RUnlock()
return cached, nil
}
pdfVisionPromptCacheMu.RUnlock()
baseDir, err := pdfVisionPromptsBaseDir()
if err != nil {
return "", err
}
promptPath := filepath.Join(baseDir, "rag", "prompts", fmt.Sprintf("%s.md", name))
content, err := os.ReadFile(promptPath)
if err != nil {
return "", fmt.Errorf("prompt file %q not found: %w", name, err)
}
cached := strings.TrimSpace(string(content))
pdfVisionPromptCacheMu.Lock()
pdfVisionPromptCache[name] = cached
pdfVisionPromptCacheMu.Unlock()
return cached, nil
}
func pdfVisionPromptsBaseDir() (string, error) {
var initErr error
pdfVisionPromptsOnce.Do(func() {
cwd, err := os.Getwd()
if err != nil {
initErr = err
return
}
for dir := cwd; dir != "/" && dir != "."; dir = filepath.Dir(dir) {
if _, err := os.Stat(filepath.Join(dir, "rag", "prompts")); err == nil {
pdfVisionPromptsBase = dir
return
}
next := filepath.Dir(dir)
if next == dir {
break
}
}
pdfVisionPromptsBase = "/ragflow"
})
if initErr != nil {
return "", initErr
}
return pdfVisionPromptsBase, nil
}
func renderPDFVisionPrompt(template string, page int) string {
rendered := strings.ReplaceAll(template, "{{ page }}", fmt.Sprintf("%d", page))
rendered = strings.ReplaceAll(rendered, "{{page}}", fmt.Sprintf("%d", page))
return rendered
}
type tenantModelExtra struct {
MaxTokens *int `json:"max_tokens"`
}
func resolveTenantModelByType(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
tenantDAO := dao.NewTenantDAO()
tenant, err := tenantDAO.GetByID(tenantID)
if err != nil {
return nil, "", nil, 0, err
}
var modelID string
switch modelType {
case entity.ModelTypeChat:
modelID = tenant.LLMID
case entity.ModelTypeEmbedding:
modelID = tenant.EmbdID
case entity.ModelTypeRerank:
modelID = tenant.RerankID
case entity.ModelTypeSpeech2Text:
modelID = tenant.ASRID
case entity.ModelTypeImage2Text:
modelID = tenant.Img2TxtID
case entity.ModelTypeTTS:
modelID = tenant.TTSID
case entity.ModelTypeOCR:
modelID = tenant.OCRID
default:
return nil, "", nil, 0, fmt.Errorf("invalid model type: %s", modelType)
}
if modelID == "" {
return nil, "", nil, 0, fmt.Errorf("no default %s model is set", modelType)
}
return resolveModelConfigFromProviderInstance(tenantID, modelType, modelID)
}
func resolveModelConfigFromProviderInstance(tenantID string, modelType entity.ModelType, modelName string) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) {
pureModelName, instanceName, providerName, err := parseCompositeModelName(modelName)
if err != nil {
return nil, "", nil, 0, err
}
providerDAO := dao.NewTenantModelProviderDAO()
instanceDAO := dao.NewTenantModelInstanceDAO()
modelDAO := dao.NewTenantModelDAO()
provider, err := providerDAO.GetByTenantIDAndProviderName(tenantID, providerName)
if err != nil {
return nil, "", nil, 0, fmt.Errorf("provider %q lookup failed: %w", providerName, err)
}
instance, err := instanceDAO.GetByProviderIDAndInstanceName(provider.ID, instanceName)
if err != nil {
return nil, "", nil, 0, fmt.Errorf("instance %q lookup failed: %w", instanceName, err)
}
apiKey := instance.APIKey
var extra map[string]string
_ = json.Unmarshal([]byte(instance.Extra), &extra)
region := extra["region"]
baseURL := extra["base_url"]
modelObj, modelErr := modelDAO.GetByProviderIDAndInstanceIDAndModelTypeAndModelName(
provider.ID, instance.ID, string(modelType), pureModelName,
)
switch {
case modelErr == nil:
if modelObj.Status == "inactive" {
return nil, "", nil, 0, fmt.Errorf("model %q is disabled", modelName)
}
providerInfo := dao.GetModelProviderManager().FindProvider(providerName)
if providerInfo == nil {
return nil, "", nil, 0, fmt.Errorf("provider %q driver not found", providerName)
}
driver, err := newModelDriverForBaseURLLocal(providerInfo.ModelDriver, providerName, region, baseURL)
if err != nil {
return nil, "", nil, 0, err
}
maxTokens := 0
if mi, _ := dao.GetModelProviderManager().GetModelByName(providerName, pureModelName); mi != nil && mi.MaxTokens != nil {
maxTokens = *mi.MaxTokens
}
if modelObj != nil && strings.TrimSpace(modelObj.Extra) != "" {
var tenantExtra tenantModelExtra
if err := json.Unmarshal([]byte(modelObj.Extra), &tenantExtra); err != nil {
return nil, "", nil, 0, err
}
if tenantExtra.MaxTokens != nil && *tenantExtra.MaxTokens > 0 {
maxTokens = *tenantExtra.MaxTokens
}
}
apiConfig := &modelModule.APIConfig{ApiKey: &apiKey, Region: &region, BaseURL: &baseURL}
return driver, modelObj.ModelName, apiConfig, maxTokens, nil
case !errorsIsRecordNotFound(modelErr):
return nil, "", nil, 0, fmt.Errorf("model %q lookup failed: %w", modelName, modelErr)
}
targetFactoryName := providerName
if region == "intl" && strings.EqualFold(providerName, "siliconflow") {
targetFactoryName = "siliconflow_intl"
}
targetProvider := dao.GetModelProviderManager().FindProvider(targetFactoryName)
if targetProvider == nil {
return nil, "", nil, 0, fmt.Errorf("model provider config not found: %s", providerName)
}
var llmInfo *modelModule.Model
for i := range targetProvider.Models {
if strings.EqualFold(targetProvider.Models[i].Name, pureModelName) {
llmInfo = targetProvider.Models[i]
break
}
}
if llmInfo == nil {
return nil, "", nil, 0, fmt.Errorf("model config not found: %s", modelName)
}
driver, err := newModelDriverForBaseURLLocal(targetProvider.ModelDriver, providerName, region, baseURL)
if err != nil {
return nil, "", nil, 0, err
}
apiConfig := &modelModule.APIConfig{ApiKey: &apiKey, Region: &region, BaseURL: &baseURL}
maxTokens := 0
if llmInfo.MaxTokens != nil {
maxTokens = *llmInfo.MaxTokens
}
return driver, llmInfo.Name, apiConfig, maxTokens, nil
}
func parseCompositeModelName(compositeName string) (modelName, instanceName, providerName string, err error) {
parts := strings.Split(compositeName, "@")
switch len(parts) {
case 3:
return parts[0], parts[1], parts[2], nil
case 2:
return parts[0], "default", parts[1], nil
case 1:
return parts[0], "", "", fmt.Errorf("provider name missing in model name: %s", compositeName)
default:
return "", "", "", fmt.Errorf("invalid model name format: %s", compositeName)
}
}
func newModelDriverForBaseURLLocal(driver modelModule.ModelDriver, providerName, region, baseURL string) (modelModule.ModelDriver, error) {
if driver == nil {
return nil, fmt.Errorf("provider %s driver not found", providerName)
}
if strings.TrimSpace(baseURL) == "" {
return driver, nil
}
baseURLByRegion := map[string]string{region: baseURL}
if region == "" {
baseURLByRegion["default"] = baseURL
}
newDriver := driver.NewInstance(baseURLByRegion)
if newDriver == nil {
return nil, fmt.Errorf("provider %s does not support custom base_url", providerName)
}
return newDriver, nil
}
func errorsIsRecordNotFound(err error) bool {
return err != nil && (err == gorm.ErrRecordNotFound || strings.Contains(err.Error(), gorm.ErrRecordNotFound.Error()))
}

46
internal/ingestion/component/pdf_vision_dispatch_cgo.go Normal file
View File

@@ -0,0 +1,46 @@
//go:build cgo
package component
import (
"bytes"
"encoding/base64"
"fmt"
"image/png"
deepdocpdf "ragflow/internal/deepdoc/parser/pdf"
)
const pdfVisionZoom = 3.0
func defaultRenderPDFVisionPages(binary []byte) ([]pdfVisionPage, error) {
engine, err := deepdocpdf.NewEngine(binary)
if err != nil {
return nil, err
}
defer engine.Close()
pageCount, err := engine.PageCount()
if err != nil {
return nil, err
}
pages := make([]pdfVisionPage, 0, pageCount)
for pageIdx := 0; pageIdx < pageCount; pageIdx++ {
img, err := deepdocpdf.RenderPageToImage(engine, pageIdx)
if err != nil {
return nil, fmt.Errorf("page %d: %w", pageIdx+1, err)
}
var buf bytes.Buffer
if err := png.Encode(&buf, img); err != nil {
return nil, fmt.Errorf("page %d encode png: %w", pageIdx+1, err)
}
bounds := img.Bounds()
pages = append(pages, pdfVisionPage{
PageNumber: pageIdx + 1,
WidthPts: float64(bounds.Dx()) / pdfVisionZoom,
HeightPts: float64(bounds.Dy()) / pdfVisionZoom,
ImageURL: "data:image/png;base64," + base64.StdEncoding.EncodeToString(buf.Bytes()),
})
}
return pages, nil
}

9
internal/ingestion/component/pdf_vision_dispatch_nocgo.go Normal file
View File

@@ -0,0 +1,9 @@
//go:build !cgo
package component
import "fmt"
func defaultRenderPDFVisionPages(_ []byte) ([]pdfVisionPage, error) {
return nil, fmt.Errorf("tenant-aware PDF IMAGE2TEXT backend requires cgo rendering support")
}