mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-07 12:00:44 +08:00
Align Go parser backends and PDF pipeline with Python (#16676)
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:
@@ -123,6 +123,9 @@ func (p *PaddleOCRLocalModel) OCRFile(modelName *string, content []byte, fileURL
|
||||
"file": base64Str,
|
||||
"fileType": fileType,
|
||||
}
|
||||
if ocrConfig != nil && strings.TrimSpace(ocrConfig.Algorithm) != "" {
|
||||
reqData["algorithm"] = ocrConfig.Algorithm
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqData)
|
||||
if err != nil {
|
||||
@@ -138,6 +141,9 @@ func (p *PaddleOCRLocalModel) OCRFile(modelName *string, content []byte, fileURL
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if auth := BearerAuth(apiConfig); auth != "" {
|
||||
req.Header.Set("Authorization", auth)
|
||||
}
|
||||
|
||||
resp, err := p.baseModel.httpClient.Do(req)
|
||||
if err != nil {
|
||||
|
||||
@@ -196,6 +196,7 @@ type TTSConfig struct {
|
||||
}
|
||||
|
||||
type OCRConfig struct {
|
||||
Algorithm string
|
||||
}
|
||||
|
||||
type ParseFileConfig struct {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"])
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
433
internal/ingestion/component/pdf_vision_dispatch.go
Normal 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: ®ion, 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: ®ion, 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
46
internal/ingestion/component/pdf_vision_dispatch_cgo.go
Normal 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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -17,6 +17,7 @@
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
@@ -33,6 +34,8 @@ import (
|
||||
_ "ragflow/internal/ingestion/component/chunker"
|
||||
"ragflow/internal/storage"
|
||||
"ragflow/internal/tokenizer"
|
||||
|
||||
"github.com/signintech/gopdf"
|
||||
)
|
||||
|
||||
type fixedEmbedder struct{}
|
||||
@@ -220,6 +223,102 @@ func TestPipelineRun_TemplateOne_RealComponents(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipelineRun_TemplateOne_RealComponents_PDFDeepDocChunking(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
t.Setenv("DEEPDOC_URL", "")
|
||||
t.Setenv("OSSDEEPDOC_URL", "")
|
||||
|
||||
templatePath := filepath.Join(repoRootFromPipelineTest(t), "agent", "templates", "ingestion_pipeline_one.json")
|
||||
templateBytes, err := os.ReadFile(templatePath)
|
||||
if err != nil {
|
||||
t.Fatalf("read template: %v", err)
|
||||
}
|
||||
terminalIDs := terminalComponentIDsFromTemplate(t, templateBytes)
|
||||
if len(terminalIDs) != 1 || terminalIDs[0] != "Tokenizer:FrankWeeksListen" {
|
||||
t.Fatalf("terminal ids = %v, want [Tokenizer:FrankWeeksListen]", terminalIDs)
|
||||
}
|
||||
|
||||
fixture := loadTemplatePipelinePDFFixture(t)
|
||||
mem := withRealTemplateDeps(t)
|
||||
const (
|
||||
bucket = "test-bucket"
|
||||
path = "fixtures/template-one.pdf"
|
||||
)
|
||||
docID := seedTemplateDocumentBytes(t, mem, fixture.Name, bucket, path, fixture.Bytes)
|
||||
|
||||
pipe, err := NewPipelineFromDSL(templateBytes, "template-one-pdf-real")
|
||||
if err != nil {
|
||||
t.Fatalf("NewPipelineFromDSL: %v", err)
|
||||
}
|
||||
out, err := pipe.Run(context.Background(), map[string]any{
|
||||
"doc_id": docID,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
payload := terminalPayloadFromRunOutput(t, out, terminalIDs[0])
|
||||
if got := payload["output_format"]; got != "chunks" {
|
||||
t.Fatalf("output_format = %v, want chunks", got)
|
||||
}
|
||||
chunks, ok := payload["chunks"].([]map[string]any)
|
||||
if !ok || len(chunks) != 1 {
|
||||
t.Fatalf("chunks = %T/%v, want 1 merged chunk", payload["chunks"], payload["chunks"])
|
||||
}
|
||||
chunkText, _ := chunks[0]["text"].(string)
|
||||
assertNormalizedContainsAll(t, chunkText, fixture.ExpectContains...)
|
||||
if _, ok := chunks[0]["content_ltks"].(string); !ok || chunks[0]["content_ltks"] == "" {
|
||||
t.Fatalf("chunks[0].content_ltks missing or empty: %v", chunks[0]["content_ltks"])
|
||||
}
|
||||
if _, ok := chunks[0]["content_sm_ltks"].(string); !ok || chunks[0]["content_sm_ltks"] == "" {
|
||||
t.Fatalf("chunks[0].content_sm_ltks missing or empty: %v", chunks[0]["content_sm_ltks"])
|
||||
}
|
||||
vec := floatSliceFromAny(t, chunks[0]["q_4_vec"])
|
||||
trimmedChunkText := strings.TrimSpace(chunkText)
|
||||
if len(vec) != 4 || vec[0] != float64(len(trimmedChunkText)) {
|
||||
t.Fatalf("chunks[0].q_4_vec = %v, want first=%v", vec, float64(len(trimmedChunkText)))
|
||||
}
|
||||
if got := payload["embedding_token_consumption"]; got != tokenizer.NumTokensFromString(trimmedChunkText) {
|
||||
t.Fatalf("embedding_token_consumption = %v, want %d", got, tokenizer.NumTokensFromString(trimmedChunkText))
|
||||
}
|
||||
|
||||
state := stateFromRunOutput(t, out)
|
||||
parserState, ok := state["Parser:HipSignsRhyme"]
|
||||
if !ok {
|
||||
t.Fatal("missing Parser:HipSignsRhyme state")
|
||||
}
|
||||
if got := parserState["output_format"]; got != "json" {
|
||||
t.Fatalf("parser output_format = %v, want json", got)
|
||||
}
|
||||
fileState, ok := parserState["file"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("parser file = %T, want map[string]any", parserState["file"])
|
||||
}
|
||||
if got := fileState["name"]; got != fixture.Name {
|
||||
t.Fatalf("parser file.name = %v, want %q", got, fixture.Name)
|
||||
}
|
||||
if got := fileState["page_count"]; got != fixture.PageCount {
|
||||
t.Fatalf("parser file.page_count = %v, want %d", got, fixture.PageCount)
|
||||
}
|
||||
jsonItems, ok := parserState["json"].([]map[string]any)
|
||||
if !ok || len(jsonItems) == 0 {
|
||||
t.Fatalf("parser json = %T/%v, want non-empty []map[string]any", parserState["json"], parserState["json"])
|
||||
}
|
||||
parserJoined := joinJSONItemTexts(jsonItems)
|
||||
assertNormalizedContainsAll(t, parserJoined, fixture.ExpectContains...)
|
||||
|
||||
chunkerState, ok := state["TokenChunker:DryDrinksVisit"]
|
||||
if !ok {
|
||||
t.Fatal("missing TokenChunker:DryDrinksVisit state")
|
||||
}
|
||||
chunkerChunks, ok := chunkerState["chunks"].([]map[string]any)
|
||||
if !ok || len(chunkerChunks) != 1 {
|
||||
t.Fatalf("chunker chunks = %T/%v, want 1 item", chunkerState["chunks"], chunkerState["chunks"])
|
||||
}
|
||||
if got := chunkerChunks[0]["text"]; got != chunkText {
|
||||
t.Fatalf("chunker chunk text != tokenizer chunk text:\nchunker=%q\ntokenizer=%q", got, chunkText)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipelineRun_TemplateManual_RealComponents(t *testing.T) {
|
||||
requireTokenizerPool(t)
|
||||
|
||||
@@ -750,7 +849,12 @@ var registerTemplateDocumentRef func(docID string, ref componentpkg.DocumentStor
|
||||
|
||||
func seedTemplateDocument(t *testing.T, stg storage.Storage, name, bucket, path, content string) string {
|
||||
t.Helper()
|
||||
if err := stg.Put(bucket, path, []byte(content)); err != nil {
|
||||
return seedTemplateDocumentBytes(t, stg, name, bucket, path, []byte(content))
|
||||
}
|
||||
|
||||
func seedTemplateDocumentBytes(t *testing.T, stg storage.Storage, name, bucket, path string, content []byte) string {
|
||||
t.Helper()
|
||||
if err := stg.Put(bucket, path, content); err != nil {
|
||||
t.Fatalf("seed storage: %v", err)
|
||||
}
|
||||
if registerTemplateDocumentRef == nil {
|
||||
@@ -765,6 +869,88 @@ func seedTemplateDocument(t *testing.T, stg storage.Storage, name, bucket, path,
|
||||
return docID
|
||||
}
|
||||
|
||||
type templatePDFFixture struct {
|
||||
Name string
|
||||
Bytes []byte
|
||||
PageCount int
|
||||
ExpectContains []string
|
||||
}
|
||||
|
||||
func loadTemplatePipelinePDFFixture(t *testing.T) templatePDFFixture {
|
||||
t.Helper()
|
||||
data, err := generateTemplatePipelinePDF()
|
||||
if err != nil {
|
||||
t.Fatalf("generate pdf fixture: %v", err)
|
||||
}
|
||||
return templatePDFFixture{
|
||||
Name: "generated-6pages.pdf",
|
||||
Bytes: data,
|
||||
PageCount: 6,
|
||||
ExpectContains: []string{
|
||||
"Pipeline PDF Fixture",
|
||||
"Page 1 explains why deepdoc parsing matters for chunking.",
|
||||
"Page 3 ensures the parser crosses page boundaries correctly.",
|
||||
"Page 6 confirms the tokenizer sees one merged chunk at the end.",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func generateTemplatePipelinePDF() ([]byte, error) {
|
||||
fontPath, err := findTemplatePDFFont()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pages := []string{
|
||||
"Pipeline PDF Fixture\nPage 1 explains why deepdoc parsing matters for chunking.",
|
||||
"Pipeline PDF Fixture\nPage 2 keeps a second page in the document for integration coverage.",
|
||||
"Pipeline PDF Fixture\nPage 3 ensures the parser crosses page boundaries correctly.",
|
||||
"Pipeline PDF Fixture\nPage 4 adds enough content to resemble a real multi-page document.",
|
||||
"Pipeline PDF Fixture\nPage 5 verifies the chunker does not drop later pages.",
|
||||
"Pipeline PDF Fixture\nPage 6 confirms the tokenizer sees one merged chunk at the end.",
|
||||
}
|
||||
|
||||
pdf := &gopdf.GoPdf{}
|
||||
pdf.Start(gopdf.Config{PageSize: *gopdf.PageSizeA4})
|
||||
if err := pdf.AddTTFFont("fixture", fontPath); err != nil {
|
||||
return nil, fmt.Errorf("AddTTFFont: %w", err)
|
||||
}
|
||||
for _, pageText := range pages {
|
||||
pdf.AddPage()
|
||||
if err := pdf.SetFont("fixture", "", 16); err != nil {
|
||||
return nil, fmt.Errorf("SetFont title: %w", err)
|
||||
}
|
||||
pdf.SetXY(56, 72)
|
||||
parts := strings.Split(pageText, "\n")
|
||||
for idx, line := range parts {
|
||||
if idx == 1 {
|
||||
if err := pdf.SetFont("fixture", "", 12); err != nil {
|
||||
return nil, fmt.Errorf("SetFont body: %w", err)
|
||||
}
|
||||
pdf.SetXY(56, 104)
|
||||
}
|
||||
if err := pdf.Text(line); err != nil {
|
||||
return nil, fmt.Errorf("Text(%q): %w", line, err)
|
||||
}
|
||||
pdf.Br(20)
|
||||
}
|
||||
}
|
||||
return bytes.Clone(pdf.GetBytesPdf()), nil
|
||||
}
|
||||
|
||||
func findTemplatePDFFont() (string, error) {
|
||||
candidates := []string{
|
||||
"/usr/share/fonts/truetype/LiberationSerif-Regular.ttf",
|
||||
"/usr/share/fonts/truetype/DejaVuSerif.ttf",
|
||||
"/usr/share/fonts/truetype/DejaVuSans.ttf",
|
||||
}
|
||||
for _, candidate := range candidates {
|
||||
if _, err := os.Stat(candidate); err == nil {
|
||||
return candidate, nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("no usable TTF font found for generated PDF fixture")
|
||||
}
|
||||
|
||||
func assertMetadataContainsString(t *testing.T, metadata map[string]any, key, want string) {
|
||||
t.Helper()
|
||||
raw, ok := metadata[key]
|
||||
@@ -897,6 +1083,33 @@ func floatSliceFromAny(t *testing.T, v any) []float64 {
|
||||
}
|
||||
}
|
||||
|
||||
func joinJSONItemTexts(items []map[string]any) string {
|
||||
var parts []string
|
||||
for _, item := range items {
|
||||
text, _ := item["text"].(string)
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
parts = append(parts, text)
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func assertNormalizedContainsAll(t *testing.T, got string, wantSubstrings ...string) {
|
||||
t.Helper()
|
||||
normalizedGot := normalizeTestText(got)
|
||||
for _, want := range wantSubstrings {
|
||||
if !strings.Contains(normalizedGot, normalizeTestText(want)) {
|
||||
t.Fatalf("normalized text %q does not contain %q", normalizedGot, normalizeTestText(want))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeTestText(s string) string {
|
||||
return strings.Join(strings.Fields(strings.ReplaceAll(s, "\u00ad", "")), " ")
|
||||
}
|
||||
|
||||
func terminalPayloadFromRunOutput(t *testing.T, out map[string]any, terminalID string) map[string]any {
|
||||
t.Helper()
|
||||
if out == nil {
|
||||
|
||||
64
internal/parser/parser/docx_parser_cgo_test.go
Normal file
64
internal/parser/parser/docx_parser_cgo_test.go
Normal file
@@ -0,0 +1,64 @@
|
||||
//go:build cgo
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDOCXParser_ParseWithResult_CGOMinimalDocument(t *testing.T) {
|
||||
p, err := NewDOCXParser(OfficeOxide)
|
||||
if err != nil {
|
||||
t.Fatalf("NewDOCXParser: %v", err)
|
||||
}
|
||||
data := minimalDOCX(t, "Hello from DOCX parser")
|
||||
res := p.ParseWithResult("sample.docx", data)
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "markdown"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if res.Markdown == "" {
|
||||
t.Fatal("Markdown is empty; want parsed content")
|
||||
}
|
||||
}
|
||||
|
||||
func minimalDOCX(t *testing.T, text string) []byte {
|
||||
t.Helper()
|
||||
var buf bytes.Buffer
|
||||
zw := zip.NewWriter(&buf)
|
||||
writeZipFile(t, zw, "[Content_Types].xml", `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
|
||||
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
|
||||
<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>
|
||||
<Default Extension="xml" ContentType="application/xml"/>
|
||||
<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
|
||||
</Types>`)
|
||||
writeZipFile(t, zw, "_rels/.rels", `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
|
||||
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
|
||||
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/>
|
||||
</Relationships>`)
|
||||
writeZipFile(t, zw, "word/document.xml", `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
|
||||
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
|
||||
<w:body>
|
||||
<w:p><w:r><w:t>`+text+`</w:t></w:r></w:p>
|
||||
</w:body>
|
||||
</w:document>`)
|
||||
if err := zw.Close(); err != nil {
|
||||
t.Fatalf("zip close: %v", err)
|
||||
}
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func writeZipFile(t *testing.T, zw *zip.Writer, name, body string) {
|
||||
t.Helper()
|
||||
w, err := zw.Create(name)
|
||||
if err != nil {
|
||||
t.Fatalf("create zip entry %s: %v", name, err)
|
||||
}
|
||||
if _, err := w.Write([]byte(body)); err != nil {
|
||||
t.Fatalf("write zip entry %s: %v", name, err)
|
||||
}
|
||||
}
|
||||
32
internal/parser/parser/office_parsers_nocgo_test.go
Normal file
32
internal/parser/parser/office_parsers_nocgo_test.go
Normal file
@@ -0,0 +1,32 @@
|
||||
//go:build !cgo
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestOfficeParsers_ParseWithResult_NoCGO(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
res ParseResult
|
||||
}{
|
||||
{name: "docx", res: (&DOCXParser{}).ParseWithResult("a.docx", nil)},
|
||||
{name: "doc", res: (&DOCParser{}).ParseWithResult("a.doc", nil)},
|
||||
{name: "xlsx", res: (&XLSXParser{}).ParseWithResult("a.xlsx", nil)},
|
||||
{name: "xls", res: (&XLSParser{}).ParseWithResult("a.xls", nil)},
|
||||
{name: "pptx", res: (&PPTXParser{}).ParseWithResult("a.pptx", nil)},
|
||||
{name: "ppt", res: (&PPTParser{}).ParseWithResult("a.ppt", nil)},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if tc.res.Err == nil {
|
||||
t.Fatal("want ErrOfficeCGORequired, got nil")
|
||||
}
|
||||
if !errors.Is(tc.res.Err, ErrOfficeCGORequired) {
|
||||
t.Fatalf("err = %v, want wraps ErrOfficeCGORequired", tc.res.Err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
34
internal/parser/parser/office_routes_cgo_test.go
Normal file
34
internal/parser/parser/office_routes_cgo_test.go
Normal file
@@ -0,0 +1,34 @@
|
||||
//go:build cgo
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"ragflow/internal/utility"
|
||||
)
|
||||
|
||||
func TestGetParser_RoutesOfficeFamilies(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
fileType utility.FileType
|
||||
}{
|
||||
{name: "docx", fileType: utility.FileTypeDOCX},
|
||||
{name: "doc", fileType: utility.FileTypeDOC},
|
||||
{name: "pptx", fileType: utility.FileTypePPTX},
|
||||
{name: "ppt", fileType: utility.FileTypePPT},
|
||||
{name: "xlsx", fileType: utility.FileTypeXLSX},
|
||||
{name: "xls", fileType: utility.FileTypeXLS},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
p, err := GetParser(tc.fileType, map[string]string{"lib_type": OfficeOxide})
|
||||
if err != nil {
|
||||
t.Fatalf("GetParser(%q): %v", tc.fileType, err)
|
||||
}
|
||||
if _, ok := p.(ParseResultProducer); !ok {
|
||||
t.Fatalf("GetParser(%q) returned %T, want ParseResultProducer", tc.fileType, p)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -12,10 +12,36 @@ import (
|
||||
)
|
||||
|
||||
func (p *PDFParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
if err := p.validateParseMethod(); err != nil {
|
||||
return ParseResult{Err: err}
|
||||
}
|
||||
switch normalizePDFParseMethod(p.ParseMethod) {
|
||||
case "plain_text":
|
||||
return parsePDFWithPlainText(filename, data, p)
|
||||
case "mineru":
|
||||
return parsePDFWithMinerU(filename, data, p)
|
||||
case "paddleocr":
|
||||
return parsePDFWithPaddleOCR(filename, data, p)
|
||||
case "docling":
|
||||
return parsePDFWithDocling(filename, data, p)
|
||||
case "opendataloader":
|
||||
return parsePDFWithOpenDataLoader(filename, data, p)
|
||||
case "somark":
|
||||
return parsePDFWithSoMark(filename, data, p)
|
||||
case "tcadp":
|
||||
return parsePDFWithTCADP(filename, data, p)
|
||||
}
|
||||
cfg := deepdoctype.DefaultParserConfig()
|
||||
cfg.SkipOCR = false
|
||||
parser := deepdocpdf.NewParser(cfg)
|
||||
res := parsePDFWithDeepDoc(context.Background(), filename, data, parser.Parse)
|
||||
res := parsePDFWithDeepDocOptions(context.Background(), filename, data, pdfPostProcessOptions{
|
||||
outputFormat: p.OutputFormat,
|
||||
zoom: cfg.Zoom,
|
||||
enableMultiColumn: p.EnableMultiColumn,
|
||||
flattenMediaToText: p.FlattenMediaToText,
|
||||
removeTOC: p.RemoveTOC,
|
||||
removeHeaderFooter: p.RemoveHeaderFooter,
|
||||
}, parser.Parse)
|
||||
if res.Err != nil && errors.Is(res.Err, deepdocpdf.ErrNoPDFData) {
|
||||
return ParseResult{Err: fmt.Errorf("%w: %s", ErrPDFEngineUnavailable, filename)}
|
||||
}
|
||||
|
||||
@@ -21,8 +21,11 @@ import (
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"image"
|
||||
"os"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
deepdocpdf "ragflow/internal/deepdoc/parser/pdf"
|
||||
"ragflow/internal/deepdoc/parser/pdf/inference"
|
||||
@@ -36,17 +39,84 @@ import (
|
||||
// is compiled behind `//go:build cgo`.
|
||||
var ErrPDFEngineUnavailable = errors.New("parser: PDF backend unavailable in this build")
|
||||
|
||||
var supportedPDFParseMethods = map[string]struct{}{
|
||||
"": {},
|
||||
"deepdoc": {},
|
||||
"plain_text": {},
|
||||
"mineru": {},
|
||||
"paddleocr": {},
|
||||
"docling": {},
|
||||
"opendataloader": {},
|
||||
"somark": {},
|
||||
"tcadp": {},
|
||||
}
|
||||
|
||||
type PDFParser struct {
|
||||
ParserType string // DeepDoc, PaddleOCR, MinerU
|
||||
Model string // DeepDoc@buildin@ragflow
|
||||
LibType string // pdf_oxide, used by DeepDoc
|
||||
|
||||
FlattenMediaToText bool
|
||||
RemoveTOC bool
|
||||
RemoveHeaderFooter bool
|
||||
EnableMultiColumn bool
|
||||
OutputFormat string
|
||||
ParseMethod string
|
||||
MinerUAPIServer string
|
||||
MinerUAPIKey string
|
||||
MinerUBackend string
|
||||
MinerUPollTimeout time.Duration
|
||||
PaddleOCRBaseURL string
|
||||
PaddleOCRAPIKey string
|
||||
PaddleOCRAlgorithm string
|
||||
DoclingServerURL string
|
||||
DoclingAPIKey string
|
||||
OpenDataLoaderAPIServer string
|
||||
OpenDataLoaderAPIKey string
|
||||
OpenDataLoaderTimeout int
|
||||
OpenDataLoaderHybrid string
|
||||
OpenDataLoaderImageOutput string
|
||||
OpenDataLoaderSanitize *bool
|
||||
SoMarkBaseURL string
|
||||
SoMarkAPIKey string
|
||||
SoMarkImageFormat string
|
||||
SoMarkFormulaFormat string
|
||||
SoMarkTableFormat string
|
||||
SoMarkCSFormat string
|
||||
SoMarkEnableTextCrossPage bool
|
||||
SoMarkEnableTableCrossPage bool
|
||||
SoMarkEnableTitleLevelRecognition bool
|
||||
SoMarkEnableInlineImage bool
|
||||
SoMarkEnableTableImage bool
|
||||
SoMarkEnableImageUnderstanding bool
|
||||
SoMarkKeepHeaderFooter bool
|
||||
TCADPAPIServer string
|
||||
TCADPAPIKey string
|
||||
TCADPTableResultType string
|
||||
TCADPMarkdownImageResponseType string
|
||||
}
|
||||
|
||||
func NewPDFParser() *PDFParser {
|
||||
return &PDFParser{
|
||||
ParserType: "DeepDoc",
|
||||
Model: "DeepDoc@buildin@ragflow",
|
||||
LibType: "pdf_oxide",
|
||||
ParserType: "DeepDoc",
|
||||
Model: "DeepDoc@buildin@ragflow",
|
||||
LibType: "pdf_oxide",
|
||||
ParseMethod: "deepdoc",
|
||||
OutputFormat: "json",
|
||||
MinerUBackend: "pipeline",
|
||||
MinerUPollTimeout: minerUPollTimeout,
|
||||
PaddleOCRAlgorithm: "PaddleOCR-VL",
|
||||
OpenDataLoaderTimeout: 600,
|
||||
SoMarkBaseURL: "https://somark.tech/api/v1",
|
||||
SoMarkImageFormat: "url",
|
||||
SoMarkFormulaFormat: "latex",
|
||||
SoMarkTableFormat: "html",
|
||||
SoMarkCSFormat: "image",
|
||||
SoMarkEnableInlineImage: true,
|
||||
SoMarkEnableTableImage: true,
|
||||
SoMarkEnableImageUnderstanding: true,
|
||||
TCADPTableResultType: "1",
|
||||
TCADPMarkdownImageResponseType: "1",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -54,6 +124,161 @@ func (p *PDFParser) String() string {
|
||||
return "PDFParser"
|
||||
}
|
||||
|
||||
func (p *PDFParser) ConfigureFromSetup(setup map[string]any) {
|
||||
if p == nil || setup == nil {
|
||||
return
|
||||
}
|
||||
if v, ok := setup["flatten_media_to_text"].(bool); ok {
|
||||
p.FlattenMediaToText = v
|
||||
}
|
||||
if v, ok := setup["remove_toc"].(bool); ok {
|
||||
p.RemoveTOC = v
|
||||
}
|
||||
if v, ok := setup["remove_header_footer"].(bool); ok {
|
||||
p.RemoveHeaderFooter = v
|
||||
}
|
||||
if v, ok := setup["enable_multi_column"].(bool); ok {
|
||||
p.EnableMultiColumn = v
|
||||
}
|
||||
if v, ok := setup["parse_method"].(string); ok && v != "" {
|
||||
p.ParseMethod = v
|
||||
}
|
||||
if v, ok := setup["mineru_apiserver"].(string); ok && v != "" {
|
||||
p.MinerUAPIServer = v
|
||||
}
|
||||
if v, ok := setup["mineru_api_key"].(string); ok {
|
||||
p.MinerUAPIKey = v
|
||||
}
|
||||
if v, ok := setup["mineru_backend"].(string); ok && v != "" {
|
||||
p.MinerUBackend = v
|
||||
}
|
||||
if v, ok := setup["mineru_timeout_seconds"].(int); ok && v > 0 {
|
||||
p.MinerUPollTimeout = time.Duration(v) * time.Second
|
||||
}
|
||||
if v, ok := setup["mineru_timeout_seconds"].(float64); ok && v > 0 {
|
||||
p.MinerUPollTimeout = time.Duration(v * float64(time.Second))
|
||||
}
|
||||
if v, ok := setup["output_format"].(string); ok && v != "" {
|
||||
p.OutputFormat = v
|
||||
}
|
||||
if v, ok := setup["paddleocr_base_url"].(string); ok && v != "" {
|
||||
p.PaddleOCRBaseURL = v
|
||||
}
|
||||
if v, ok := setup["paddleocr_api_key"].(string); ok {
|
||||
p.PaddleOCRAPIKey = v
|
||||
}
|
||||
if v, ok := setup["paddleocr_algorithm"].(string); ok && v != "" {
|
||||
p.PaddleOCRAlgorithm = v
|
||||
}
|
||||
if v, ok := setup["docling_server_url"].(string); ok && v != "" {
|
||||
p.DoclingServerURL = v
|
||||
}
|
||||
if v, ok := setup["docling_api_key"].(string); ok {
|
||||
p.DoclingAPIKey = v
|
||||
}
|
||||
if v, ok := setup["opendataloader_apiserver"].(string); ok && v != "" {
|
||||
p.OpenDataLoaderAPIServer = v
|
||||
}
|
||||
if v, ok := setup["opendataloader_api_key"].(string); ok {
|
||||
p.OpenDataLoaderAPIKey = v
|
||||
}
|
||||
if v, ok := setup["opendataloader_timeout"].(int); ok && v > 0 {
|
||||
p.OpenDataLoaderTimeout = v
|
||||
}
|
||||
if v, ok := setup["opendataloader_timeout"].(float64); ok && v > 0 {
|
||||
p.OpenDataLoaderTimeout = int(v)
|
||||
}
|
||||
if v, ok := setup["hybrid"].(string); ok && v != "" {
|
||||
p.OpenDataLoaderHybrid = v
|
||||
}
|
||||
if v, ok := setup["image_output"].(string); ok && v != "" {
|
||||
p.OpenDataLoaderImageOutput = v
|
||||
}
|
||||
if v, ok := setup["sanitize"].(bool); ok {
|
||||
p.OpenDataLoaderSanitize = &v
|
||||
}
|
||||
if v, ok := setup["somark_base_url"].(string); ok && v != "" {
|
||||
p.SoMarkBaseURL = v
|
||||
}
|
||||
if v, ok := setup["somark_api_key"].(string); ok {
|
||||
p.SoMarkAPIKey = v
|
||||
}
|
||||
if v, ok := setup["somark_image_format"].(string); ok && v != "" {
|
||||
p.SoMarkImageFormat = v
|
||||
}
|
||||
if v, ok := setup["somark_formula_format"].(string); ok && v != "" {
|
||||
p.SoMarkFormulaFormat = v
|
||||
}
|
||||
if v, ok := setup["somark_table_format"].(string); ok && v != "" {
|
||||
p.SoMarkTableFormat = v
|
||||
}
|
||||
if v, ok := setup["somark_cs_format"].(string); ok && v != "" {
|
||||
p.SoMarkCSFormat = v
|
||||
}
|
||||
if v, ok := setup["somark_enable_text_cross_page"].(bool); ok {
|
||||
p.SoMarkEnableTextCrossPage = v
|
||||
}
|
||||
if v, ok := setup["somark_enable_table_cross_page"].(bool); ok {
|
||||
p.SoMarkEnableTableCrossPage = v
|
||||
}
|
||||
if v, ok := setup["somark_enable_title_level_recognition"].(bool); ok {
|
||||
p.SoMarkEnableTitleLevelRecognition = v
|
||||
}
|
||||
if v, ok := setup["somark_enable_inline_image"].(bool); ok {
|
||||
p.SoMarkEnableInlineImage = v
|
||||
}
|
||||
if v, ok := setup["somark_enable_table_image"].(bool); ok {
|
||||
p.SoMarkEnableTableImage = v
|
||||
}
|
||||
if v, ok := setup["somark_enable_image_understanding"].(bool); ok {
|
||||
p.SoMarkEnableImageUnderstanding = v
|
||||
}
|
||||
if v, ok := setup["somark_keep_header_footer"].(bool); ok {
|
||||
p.SoMarkKeepHeaderFooter = v
|
||||
}
|
||||
if v, ok := setup["tcadp_apiserver"].(string); ok && v != "" {
|
||||
p.TCADPAPIServer = v
|
||||
}
|
||||
if v, ok := setup["tcadp_api_key"].(string); ok {
|
||||
p.TCADPAPIKey = v
|
||||
}
|
||||
if v, ok := setup["table_result_type"].(string); ok && v != "" {
|
||||
p.TCADPTableResultType = v
|
||||
}
|
||||
if v, ok := setup["markdown_image_response_type"].(string); ok && v != "" {
|
||||
p.TCADPMarkdownImageResponseType = v
|
||||
}
|
||||
}
|
||||
|
||||
func normalizePDFParseMethod(raw string) string {
|
||||
method := strings.ToLower(strings.TrimSpace(raw))
|
||||
switch {
|
||||
case strings.HasSuffix(method, "@mineru"):
|
||||
return "mineru"
|
||||
case strings.HasSuffix(method, "@paddleocr"):
|
||||
return "paddleocr"
|
||||
case strings.HasSuffix(method, "@somark"):
|
||||
return "somark"
|
||||
case strings.HasSuffix(method, "@opendataloader"):
|
||||
return "opendataloader"
|
||||
}
|
||||
switch method {
|
||||
case "plaintext":
|
||||
return "plain_text"
|
||||
case "tcadp parser":
|
||||
return "tcadp"
|
||||
}
|
||||
return method
|
||||
}
|
||||
|
||||
func (p *PDFParser) validateParseMethod() error {
|
||||
method := normalizePDFParseMethod(p.ParseMethod)
|
||||
if _, ok := supportedPDFParseMethods[method]; ok {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("parser: unsupported PDF parse_method %q (Go currently supports: deepdoc, plain_text, mineru, paddleocr, docling, opendataloader, somark, tcadp; tenant-resolved custom IMAGE2TEXT/VLM model names are not supported in the Go parser layer)", p.ParseMethod)
|
||||
}
|
||||
|
||||
func emptyPDFResult(filename string) ParseResult {
|
||||
return ParseResult{
|
||||
OutputFormat: "json",
|
||||
@@ -85,10 +310,22 @@ func deepDocAnalyzerFromEnv() deepdoctype.DocAnalyzer {
|
||||
}
|
||||
|
||||
func pdfParseResultToJSON(filename string, parsed *deepdoctype.ParseResult) ParseResult {
|
||||
return pdfParseResultToJSONWithOptions(filename, parsed, pdfPostProcessOptions{})
|
||||
}
|
||||
|
||||
func pdfParseResultToJSONWithOptions(filename string, parsed *deepdoctype.ParseResult, opts pdfPostProcessOptions) ParseResult {
|
||||
if parsed == nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: nil DeepDOC PDF result for %s", filename)}
|
||||
}
|
||||
items := pdflayout.SectionsToJSON(parsed.Sections)
|
||||
processed := *parsed
|
||||
processed.Sections = append([]deepdoctype.Section(nil), parsed.Sections...)
|
||||
processed.Outlines = append([]deepdoctype.Outline(nil), parsed.Outlines...)
|
||||
if opts.enableMultiColumn && opts.pageWidth <= 0 {
|
||||
opts.pageWidth = firstPDFPageWidth(processed.PageImages, opts.zoom)
|
||||
}
|
||||
applyPDFPostProcess(&processed, opts)
|
||||
|
||||
items := pdflayout.SectionsToJSON(processed.Sections)
|
||||
if len(items) == 0 {
|
||||
items = []map[string]any{{"text": "", "doc_type_kwd": "text"}}
|
||||
}
|
||||
@@ -96,9 +333,14 @@ func pdfParseResultToJSON(filename string, parsed *deepdoctype.ParseResult) Pars
|
||||
if layoutType, _ := items[i]["layout_type"].(string); layoutType != "" {
|
||||
items[i]["layout"] = layoutType
|
||||
}
|
||||
if _, ok := items[i]["page_number"]; !ok {
|
||||
items[i]["page_number"] = firstPageNumber(items[i]["_pdf_positions"])
|
||||
if normalized := normalizePDFPositions(items[i]["_pdf_positions"]); len(normalized) > 0 {
|
||||
items[i]["_pdf_positions"] = normalized
|
||||
items[i]["positions"] = normalized
|
||||
if _, ok := items[i]["page_number"]; !ok {
|
||||
items[i]["page_number"] = firstPageNumber(normalized)
|
||||
}
|
||||
}
|
||||
normalizePDFDocType(items[i])
|
||||
if img, _ := items[i]["image"].(string); img != "" {
|
||||
items[i]["image"] = "data:image/png;base64," + img
|
||||
}
|
||||
@@ -107,13 +349,36 @@ func pdfParseResultToJSON(filename string, parsed *deepdoctype.ParseResult) Pars
|
||||
OutputFormat: "json",
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"page_count": len(parsed.PageImages),
|
||||
"outline": outlinesToFileMeta(parsed.Outlines),
|
||||
"page_count": len(processed.PageImages),
|
||||
"outline": outlinesToFileMeta(processed.Outlines),
|
||||
},
|
||||
JSON: items,
|
||||
}
|
||||
}
|
||||
|
||||
func pdfParseResultToMarkdownWithOptions(filename string, parsed *deepdoctype.ParseResult, opts pdfPostProcessOptions) ParseResult {
|
||||
if parsed == nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: nil DeepDOC PDF result for %s", filename)}
|
||||
}
|
||||
processed := *parsed
|
||||
processed.Sections = append([]deepdoctype.Section(nil), parsed.Sections...)
|
||||
processed.Outlines = append([]deepdoctype.Outline(nil), parsed.Outlines...)
|
||||
if opts.enableMultiColumn && opts.pageWidth <= 0 {
|
||||
opts.pageWidth = firstPDFPageWidth(processed.PageImages, opts.zoom)
|
||||
}
|
||||
applyPDFPostProcess(&processed, opts)
|
||||
|
||||
return ParseResult{
|
||||
OutputFormat: "markdown",
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"page_count": len(processed.PageImages),
|
||||
"outline": outlinesToFileMeta(processed.Outlines),
|
||||
},
|
||||
Markdown: sectionsToMarkdown(processed.Sections),
|
||||
}
|
||||
}
|
||||
|
||||
func outlinesToFileMeta(outlines []deepdoctype.Outline) []map[string]any {
|
||||
if len(outlines) == 0 {
|
||||
return []map[string]any{}
|
||||
@@ -134,11 +399,7 @@ func firstPageNumber(raw any) int {
|
||||
if !ok || len(positions) == 0 || len(positions[0]) == 0 {
|
||||
return 0
|
||||
}
|
||||
pages, ok := positions[0][0].([]any)
|
||||
if !ok || len(pages) == 0 {
|
||||
return 0
|
||||
}
|
||||
switch v := pages[0].(type) {
|
||||
switch v := positions[0][0].(type) {
|
||||
case int:
|
||||
return v
|
||||
case int64:
|
||||
@@ -163,7 +424,136 @@ func inlinePNGDataURL(raw string) string {
|
||||
return "data:image/png;base64," + raw
|
||||
}
|
||||
|
||||
func sectionsToMarkdown(sections []deepdoctype.Section) string {
|
||||
var b strings.Builder
|
||||
for _, section := range sections {
|
||||
layoutType := strings.TrimSpace(section.LayoutType)
|
||||
if layoutType == deepdoctype.LayoutTypeTitle {
|
||||
b.WriteString("\n## ")
|
||||
}
|
||||
if layoutType == deepdoctype.LayoutTypeFigure && section.Image != "" {
|
||||
b.WriteString("\n
|
||||
b.WriteString(inlinePNGDataURL(section.Image))
|
||||
b.WriteString(")")
|
||||
continue
|
||||
}
|
||||
b.WriteString(section.Text)
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func firstPDFPageWidth(pageImages map[int]image.Image, zoom float64) float64 {
|
||||
if len(pageImages) == 0 {
|
||||
return 0
|
||||
}
|
||||
if zoom <= 0 {
|
||||
zoom = deepdoctype.DefaultParserConfig().Zoom
|
||||
}
|
||||
pages := make([]int, 0, len(pageImages))
|
||||
for page := range pageImages {
|
||||
pages = append(pages, page)
|
||||
}
|
||||
sort.Ints(pages)
|
||||
img := pageImages[pages[0]]
|
||||
if img == nil {
|
||||
return 0
|
||||
}
|
||||
return float64(img.Bounds().Dx()) / zoom
|
||||
}
|
||||
|
||||
func normalizePDFPositions(raw any) [][]any {
|
||||
positions, ok := raw.([][]any)
|
||||
if !ok || len(positions) == 0 {
|
||||
return nil
|
||||
}
|
||||
normalized := make([][]any, 0, len(positions))
|
||||
for _, pos := range positions {
|
||||
if len(pos) < 5 {
|
||||
continue
|
||||
}
|
||||
pageNumber, ok := normalizePDFPageNumber(pos[0])
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
left, lok := numericAny(pos[1])
|
||||
right, rok := numericAny(pos[2])
|
||||
top, tok := numericAny(pos[3])
|
||||
bottom, bok := numericAny(pos[4])
|
||||
if !lok || !rok || !tok || !bok {
|
||||
continue
|
||||
}
|
||||
normalized = append(normalized, []any{pageNumber, left, right, top, bottom})
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
func normalizePDFPageNumber(raw any) (int, bool) {
|
||||
switch v := raw.(type) {
|
||||
case int:
|
||||
if v <= 0 {
|
||||
return v + 1, true
|
||||
}
|
||||
return v, true
|
||||
case int64:
|
||||
return normalizePDFPageNumber(int(v))
|
||||
case float64:
|
||||
return normalizePDFPageNumber(int(v))
|
||||
case []any:
|
||||
if len(v) == 0 {
|
||||
return 0, false
|
||||
}
|
||||
return normalizePDFPageNumber(v[len(v)-1])
|
||||
case []int:
|
||||
if len(v) == 0 {
|
||||
return 0, false
|
||||
}
|
||||
return normalizePDFPageNumber(v[len(v)-1])
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
func numericAny(raw any) (float64, bool) {
|
||||
switch v := raw.(type) {
|
||||
case int:
|
||||
return float64(v), true
|
||||
case int64:
|
||||
return float64(v), true
|
||||
case float64:
|
||||
return v, true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
func normalizePDFDocType(item map[string]any) {
|
||||
if item == nil {
|
||||
return
|
||||
}
|
||||
if docType, _ := item["doc_type_kwd"].(string); docType != "" {
|
||||
return
|
||||
}
|
||||
layoutType, _ := item["layout_type"].(string)
|
||||
switch layoutType {
|
||||
case "table":
|
||||
item["doc_type_kwd"] = "table"
|
||||
case "figure", "image":
|
||||
item["doc_type_kwd"] = "image"
|
||||
default:
|
||||
if img, _ := item["image"].(string); img != "" {
|
||||
item["doc_type_kwd"] = "image"
|
||||
return
|
||||
}
|
||||
item["doc_type_kwd"] = "text"
|
||||
}
|
||||
}
|
||||
|
||||
func parsePDFWithDeepDoc(ctx context.Context, filename string, data []byte, parseFn func(context.Context, []byte, deepdoctype.DocAnalyzer) (*deepdoctype.ParseResult, error)) ParseResult {
|
||||
return parsePDFWithDeepDocOptions(ctx, filename, data, pdfPostProcessOptions{}, parseFn)
|
||||
}
|
||||
|
||||
func parsePDFWithDeepDocOptions(ctx context.Context, filename string, data []byte, opts pdfPostProcessOptions, parseFn func(context.Context, []byte, deepdoctype.DocAnalyzer) (*deepdoctype.ParseResult, error)) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
@@ -171,7 +561,15 @@ func parsePDFWithDeepDoc(ctx context.Context, filename string, data []byte, pars
|
||||
if err != nil {
|
||||
return ParseResult{Err: err}
|
||||
}
|
||||
res := pdfParseResultToJSON(filename, parsed)
|
||||
var res ParseResult
|
||||
switch strings.ToLower(strings.TrimSpace(opts.outputFormat)) {
|
||||
case "", "json":
|
||||
res = pdfParseResultToJSONWithOptions(filename, parsed, opts)
|
||||
case "markdown":
|
||||
res = pdfParseResultToMarkdownWithOptions(filename, parsed, opts)
|
||||
default:
|
||||
return ParseResult{Err: fmt.Errorf("parser: unsupported PDF output_format %q", opts.outputFormat)}
|
||||
}
|
||||
for i := range res.JSON {
|
||||
if img, _ := res.JSON[i]["image"].(string); img != "" {
|
||||
res.JSON[i]["image"] = inlinePNGDataURL(img)
|
||||
|
||||
347
internal/parser/parser/pdf_parser_common_test.go
Normal file
347
internal/parser/parser/pdf_parser_common_test.go
Normal file
@@ -0,0 +1,347 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
deepdoctype "ragflow/internal/deepdoc/parser/type"
|
||||
)
|
||||
|
||||
func TestPDFParseResultToJSON_NormalizesCoreFields(t *testing.T) {
|
||||
parsed := &deepdoctype.ParseResult{
|
||||
Sections: []deepdoctype.Section{
|
||||
{
|
||||
Text: "Title block",
|
||||
LayoutType: deepdoctype.LayoutTypeTitle,
|
||||
Positions: []deepdoctype.Position{
|
||||
{
|
||||
PageNumbers: []int{0},
|
||||
Left: 10,
|
||||
Right: 20,
|
||||
Top: 30,
|
||||
Bottom: 40,
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Text: "Figure caption",
|
||||
LayoutType: deepdoctype.LayoutTypeFigure,
|
||||
Image: "aGVsbG8=",
|
||||
Positions: []deepdoctype.Position{
|
||||
{
|
||||
PageNumbers: []int{1},
|
||||
Left: 1,
|
||||
Right: 2,
|
||||
Top: 3,
|
||||
Bottom: 4,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Outlines: []deepdoctype.Outline{
|
||||
{Title: "Intro", Level: 1, PageNumber: 2},
|
||||
},
|
||||
}
|
||||
|
||||
res := pdfParseResultToJSON("sample.pdf", parsed)
|
||||
if res.Err != nil {
|
||||
t.Fatalf("pdfParseResultToJSON: %v", res.Err)
|
||||
}
|
||||
if res.OutputFormat != "json" {
|
||||
t.Fatalf("OutputFormat = %q, want json", res.OutputFormat)
|
||||
}
|
||||
if got, want := res.File["name"], "sample.pdf"; got != want {
|
||||
t.Fatalf("File.name = %v, want %v", got, want)
|
||||
}
|
||||
outline, ok := res.File["outline"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("File.outline type = %T, want []map[string]any", res.File["outline"])
|
||||
}
|
||||
if len(outline) != 1 || outline[0]["page_number"] != 2 {
|
||||
t.Fatalf("File.outline = %+v, want page_number 2", outline)
|
||||
}
|
||||
if len(res.JSON) != 2 {
|
||||
t.Fatalf("JSON len = %d, want 2", len(res.JSON))
|
||||
}
|
||||
if got, want := res.JSON[0]["layout"], "title"; got != want {
|
||||
t.Fatalf("JSON[0].layout = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[0]["page_number"], 1; got != want {
|
||||
t.Fatalf("JSON[0].page_number = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[0]["doc_type_kwd"], "text"; got != want {
|
||||
t.Fatalf("JSON[0].doc_type_kwd = %v, want %v", got, want)
|
||||
}
|
||||
pdfPositions, ok := res.JSON[0]["_pdf_positions"].([][]any)
|
||||
if !ok {
|
||||
t.Fatalf("JSON[0]._pdf_positions type = %T, want [][]any", res.JSON[0]["_pdf_positions"])
|
||||
}
|
||||
if len(pdfPositions) != 1 || pdfPositions[0][0] != 1 {
|
||||
t.Fatalf("JSON[0]._pdf_positions = %+v, want canonical 1-based positions", pdfPositions)
|
||||
}
|
||||
if got := res.JSON[0]["positions"]; got == nil {
|
||||
t.Fatal("JSON[0].positions missing after normalization")
|
||||
}
|
||||
if got, want := res.JSON[1]["doc_type_kwd"], "image"; got != want {
|
||||
t.Fatalf("JSON[1].doc_type_kwd = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[1]["page_number"], 1; got != want {
|
||||
t.Fatalf("JSON[1].page_number = %v, want %v", got, want)
|
||||
}
|
||||
secondPDFPositions, ok := res.JSON[1]["_pdf_positions"].([][]any)
|
||||
if !ok {
|
||||
t.Fatalf("JSON[1]._pdf_positions type = %T, want [][]any", res.JSON[1]["_pdf_positions"])
|
||||
}
|
||||
if len(secondPDFPositions) != 1 || secondPDFPositions[0][0] != 1 {
|
||||
t.Fatalf("JSON[1]._pdf_positions = %+v, want canonical 1-based positions", secondPDFPositions)
|
||||
}
|
||||
if got, want := res.JSON[1]["image"], "data:image/png;base64,aGVsbG8="; got != want {
|
||||
t.Fatalf("JSON[1].image = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParseResultToJSON_PreservesPositivePageNumbers(t *testing.T) {
|
||||
parsed := &deepdoctype.ParseResult{
|
||||
Sections: []deepdoctype.Section{
|
||||
{
|
||||
Text: "Already one-based",
|
||||
LayoutType: deepdoctype.LayoutTypeTable,
|
||||
Positions: []deepdoctype.Position{
|
||||
{
|
||||
PageNumbers: []int{3},
|
||||
Left: 10,
|
||||
Right: 20,
|
||||
Top: 30,
|
||||
Bottom: 40,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
res := pdfParseResultToJSON("one-based.pdf", parsed)
|
||||
if got, want := res.JSON[0]["page_number"], 3; got != want {
|
||||
t.Fatalf("JSON[1].page_number = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[0]["doc_type_kwd"], "table"; got != want {
|
||||
t.Fatalf("JSON[0].doc_type_kwd = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParseResultToJSON_EmptySectionsStillEmitPlaceholder(t *testing.T) {
|
||||
res := pdfParseResultToJSON("empty.pdf", &deepdoctype.ParseResult{})
|
||||
if res.Err != nil {
|
||||
t.Fatalf("pdfParseResultToJSON: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("JSON len = %d, want 1", len(res.JSON))
|
||||
}
|
||||
if got, want := res.JSON[0]["doc_type_kwd"], "text"; got != want {
|
||||
t.Fatalf("JSON[0].doc_type_kwd = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParseResultToJSON_DefaultKeepsHeaderFooterLikePython(t *testing.T) {
|
||||
parsed := &deepdoctype.ParseResult{
|
||||
Sections: []deepdoctype.Section{
|
||||
{Text: "Header", LayoutType: "header"},
|
||||
{
|
||||
Text: "Body",
|
||||
LayoutType: "",
|
||||
Positions: []deepdoctype.Position{{
|
||||
PageNumbers: []int{0},
|
||||
Left: 10,
|
||||
Right: 20,
|
||||
Top: 30,
|
||||
Bottom: 40,
|
||||
}},
|
||||
},
|
||||
{Text: "Footer", LayoutType: "footer"},
|
||||
},
|
||||
}
|
||||
|
||||
res := pdfParseResultToJSON("filtered.pdf", parsed)
|
||||
if res.Err != nil {
|
||||
t.Fatalf("pdfParseResultToJSON: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) != 3 {
|
||||
t.Fatalf("JSON len = %d, want 3", len(res.JSON))
|
||||
}
|
||||
if got, want := res.JSON[0]["text"], "Header"; got != want {
|
||||
t.Fatalf("JSON[0].text = %v, want %v", got, want)
|
||||
}
|
||||
if got, want := res.JSON[1]["text"], "Body"; got != want {
|
||||
t.Fatalf("JSON[1].text = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParseResultToJSONWithOptions_FiltersHeaderFooterWhenEnabled(t *testing.T) {
|
||||
parsed := &deepdoctype.ParseResult{
|
||||
Sections: []deepdoctype.Section{
|
||||
{Text: "Header", LayoutType: "header"},
|
||||
{
|
||||
Text: "Body",
|
||||
LayoutType: "",
|
||||
Positions: []deepdoctype.Position{{
|
||||
PageNumbers: []int{0},
|
||||
Left: 10,
|
||||
Right: 20,
|
||||
Top: 30,
|
||||
Bottom: 40,
|
||||
}},
|
||||
},
|
||||
{Text: "Footer", LayoutType: "footer"},
|
||||
},
|
||||
}
|
||||
|
||||
res := pdfParseResultToJSONWithOptions("filtered.pdf", parsed, pdfPostProcessOptions{removeHeaderFooter: true})
|
||||
if res.Err != nil {
|
||||
t.Fatalf("pdfParseResultToJSONWithOptions: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("JSON len = %d, want 1", len(res.JSON))
|
||||
}
|
||||
if got, want := res.JSON[0]["text"], "Body"; got != want {
|
||||
t.Fatalf("JSON[0].text = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParseResultToJSONWithOptions_RemovesTOCByOutline(t *testing.T) {
|
||||
parsed := &deepdoctype.ParseResult{
|
||||
Sections: []deepdoctype.Section{
|
||||
{
|
||||
Text: "Contents",
|
||||
LayoutType: "text",
|
||||
Positions: []deepdoctype.Position{{
|
||||
PageNumbers: []int{1},
|
||||
Left: 10,
|
||||
Right: 20,
|
||||
Top: 30,
|
||||
Bottom: 40,
|
||||
}},
|
||||
},
|
||||
{
|
||||
Text: "Body",
|
||||
LayoutType: "text",
|
||||
Positions: []deepdoctype.Position{{
|
||||
PageNumbers: []int{3},
|
||||
Left: 10,
|
||||
Right: 20,
|
||||
Top: 30,
|
||||
Bottom: 40,
|
||||
}},
|
||||
},
|
||||
},
|
||||
Outlines: []deepdoctype.Outline{
|
||||
{Title: "目录", Level: 0, PageNumber: 1},
|
||||
{Title: "Chapter 1", Level: 0, PageNumber: 3},
|
||||
},
|
||||
}
|
||||
|
||||
res := pdfParseResultToJSONWithOptions("toc.pdf", parsed, pdfPostProcessOptions{removeTOC: true})
|
||||
if res.Err != nil {
|
||||
t.Fatalf("pdfParseResultToJSONWithOptions: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("JSON len = %d, want 1", len(res.JSON))
|
||||
}
|
||||
if got, want := res.JSON[0]["text"], "Body"; got != want {
|
||||
t.Fatalf("JSON[0].text = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ConfigureFromSetup(t *testing.T) {
|
||||
p := NewPDFParser()
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "deepdoc",
|
||||
"output_format": "markdown",
|
||||
"enable_multi_column": true,
|
||||
"flatten_media_to_text": true,
|
||||
"remove_toc": true,
|
||||
"remove_header_footer": true,
|
||||
})
|
||||
if got, want := p.OutputFormat, "markdown"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if !p.EnableMultiColumn {
|
||||
t.Fatal("EnableMultiColumn = false, want true")
|
||||
}
|
||||
if got, want := p.ParseMethod, "deepdoc"; got != want {
|
||||
t.Fatalf("ParseMethod = %q, want %q", got, want)
|
||||
}
|
||||
if !p.FlattenMediaToText {
|
||||
t.Fatal("FlattenMediaToText = false, want true")
|
||||
}
|
||||
if !p.RemoveTOC {
|
||||
t.Fatal("RemoveTOC = false, want true")
|
||||
}
|
||||
if !p.RemoveHeaderFooter {
|
||||
t.Fatal("RemoveHeaderFooter = false, want true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParseResultToMarkdownWithOptions_RendersLikePython(t *testing.T) {
|
||||
parsed := &deepdoctype.ParseResult{
|
||||
Sections: []deepdoctype.Section{
|
||||
{Text: "Title", LayoutType: deepdoctype.LayoutTypeTitle},
|
||||
{Text: "Figure", LayoutType: deepdoctype.LayoutTypeFigure, Image: "aGVsbG8="},
|
||||
{Text: "Body", LayoutType: deepdoctype.LayoutTypeText},
|
||||
},
|
||||
}
|
||||
|
||||
res := pdfParseResultToMarkdownWithOptions("sample.pdf", parsed, pdfPostProcessOptions{})
|
||||
if res.Err != nil {
|
||||
t.Fatalf("pdfParseResultToMarkdownWithOptions: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "markdown"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if res.Markdown == "" {
|
||||
t.Fatal("Markdown is empty; want rendered content")
|
||||
}
|
||||
if !strings.Contains(res.Markdown, "## Title") {
|
||||
t.Fatalf("Markdown = %q, want title heading", res.Markdown)
|
||||
}
|
||||
if !strings.Contains(res.Markdown, "") {
|
||||
t.Fatalf("Markdown = %q, want inline image", res.Markdown)
|
||||
}
|
||||
if !strings.Contains(res.Markdown, "Body") {
|
||||
t.Fatalf("Markdown = %q, want body text", res.Markdown)
|
||||
}
|
||||
if len(res.JSON) != 0 {
|
||||
t.Fatalf("JSON len = %d, want 0 for markdown output", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ValidateParseMethod(t *testing.T) {
|
||||
p := NewPDFParser()
|
||||
if err := p.validateParseMethod(); err != nil {
|
||||
t.Fatalf("default validateParseMethod: %v", err)
|
||||
}
|
||||
|
||||
p.ConfigureFromSetup(map[string]any{"parse_method": "PaddleOCR"})
|
||||
if err := p.validateParseMethod(); err != nil {
|
||||
t.Fatalf("validateParseMethod(PaddleOCR): %v", err)
|
||||
}
|
||||
|
||||
p.ConfigureFromSetup(map[string]any{"parse_method": "tenant@provider@SoMark"})
|
||||
if err := p.validateParseMethod(); err != nil {
|
||||
t.Fatalf("validateParseMethod(tenant@provider@SoMark): %v", err)
|
||||
}
|
||||
|
||||
if got, want := normalizePDFParseMethod("tenant@provider@OpenDataLoader"), "opendataloader"; got != want {
|
||||
t.Fatalf("normalizePDFParseMethod(OpenDataLoader suffix) = %q, want %q", got, want)
|
||||
}
|
||||
|
||||
p.ConfigureFromSetup(map[string]any{"parse_method": "CustomVLM"})
|
||||
err := p.validateParseMethod()
|
||||
if err == nil {
|
||||
t.Fatal("validateParseMethod: want error for unsupported parse_method, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "parse_method") {
|
||||
t.Fatalf("validateParseMethod error = %q, want parse_method context", err.Error())
|
||||
}
|
||||
if !strings.Contains(err.Error(), "IMAGE2TEXT") {
|
||||
t.Fatalf("validateParseMethod error = %q, want IMAGE2TEXT/VLM guidance", err.Error())
|
||||
}
|
||||
}
|
||||
249
internal/parser/parser/pdf_parser_docling.go
Normal file
249
internal/parser/parser/pdf_parser_docling.go
Normal file
@@ -0,0 +1,249 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
models "ragflow/internal/entity/models"
|
||||
)
|
||||
|
||||
type doclingChunk struct {
|
||||
Text string `json:"text"`
|
||||
Chunk *struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"chunk"`
|
||||
}
|
||||
|
||||
type doclingDocument struct {
|
||||
MDContent string `json:"md_content"`
|
||||
TextContent string `json:"text_content"`
|
||||
JSONContent map[string]any `json:"json_content"`
|
||||
}
|
||||
|
||||
type doclingResult struct {
|
||||
Document *doclingDocument `json:"document"`
|
||||
Result *doclingDocument `json:"result"`
|
||||
}
|
||||
|
||||
type doclingResponse struct {
|
||||
Document *doclingDocument `json:"document"`
|
||||
Documents []doclingDocument `json:"documents"`
|
||||
Results []doclingResult `json:"results"`
|
||||
}
|
||||
|
||||
func parsePDFWithDocling(filename string, data []byte, parser *PDFParser) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
serverURL := strings.TrimSpace(parser.DoclingServerURL)
|
||||
if serverURL == "" {
|
||||
serverURL = strings.TrimSpace(os.Getenv("DOCLING_SERVER_URL"))
|
||||
}
|
||||
if serverURL == "" {
|
||||
return ParseResult{Err: fmt.Errorf("parser: Docling requires docling_server_url or DOCLING_SERVER_URL")}
|
||||
}
|
||||
apiKey := strings.TrimSpace(parser.DoclingAPIKey)
|
||||
if apiKey == "" {
|
||||
apiKey = strings.TrimSpace(os.Getenv("DOCLING_API_KEY"))
|
||||
}
|
||||
|
||||
baseURL := strings.TrimRight(serverURL, "/")
|
||||
auth := ""
|
||||
if apiKey != "" {
|
||||
auth = "Bearer " + apiKey
|
||||
}
|
||||
encoded := base64.StdEncoding.EncodeToString(data)
|
||||
payloads := []struct {
|
||||
endpoint string
|
||||
body func() map[string]any
|
||||
chunked bool
|
||||
}{
|
||||
{
|
||||
endpoint: "/v1/convert/source",
|
||||
chunked: true,
|
||||
body: func() map[string]any { return doclingChunkedPayload(filename, encoded, false) },
|
||||
},
|
||||
{
|
||||
endpoint: "/v1alpha/convert/source",
|
||||
chunked: true,
|
||||
body: func() map[string]any { return doclingChunkedPayload(filename, encoded, true) },
|
||||
},
|
||||
{
|
||||
endpoint: "/v1/convert/source",
|
||||
chunked: false,
|
||||
body: func() map[string]any { return doclingStandardPayload(filename, encoded, false) },
|
||||
},
|
||||
{
|
||||
endpoint: "/v1alpha/convert/source",
|
||||
chunked: false,
|
||||
body: func() map[string]any { return doclingStandardPayload(filename, encoded, true) },
|
||||
},
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for _, candidate := range payloads {
|
||||
url := baseURL + candidate.endpoint
|
||||
resp, err := models.PostJSONRequest(context.Background(), models.NewDriverHTTPClient(), url, auth, candidate.body())
|
||||
if err != nil {
|
||||
lastErr = fmt.Errorf("%s: %w", candidate.endpoint, err)
|
||||
continue
|
||||
}
|
||||
body, readErr := func() ([]byte, error) {
|
||||
defer resp.Body.Close()
|
||||
return io.ReadAll(resp.Body)
|
||||
}()
|
||||
if readErr != nil {
|
||||
lastErr = fmt.Errorf("%s: read response: %w", candidate.endpoint, readErr)
|
||||
continue
|
||||
}
|
||||
if resp.StatusCode >= 300 {
|
||||
if candidate.chunked {
|
||||
lastErr = fmt.Errorf("%s: HTTP %d", candidate.endpoint, resp.StatusCode)
|
||||
continue
|
||||
}
|
||||
lastErr = fmt.Errorf("%s: HTTP %d %s", candidate.endpoint, resp.StatusCode, string(body))
|
||||
continue
|
||||
}
|
||||
if candidate.chunked {
|
||||
if res, ok := parseDoclingChunkedResult(filename, body, parser.OutputFormat); ok {
|
||||
return res
|
||||
}
|
||||
lastErr = fmt.Errorf("%s: chunked response contained no usable text", candidate.endpoint)
|
||||
continue
|
||||
}
|
||||
if res, ok := parseDoclingStandardResult(filename, body, parser.OutputFormat); ok {
|
||||
return res
|
||||
}
|
||||
lastErr = fmt.Errorf("%s: standard response contained no parsed document", candidate.endpoint)
|
||||
}
|
||||
if lastErr == nil {
|
||||
lastErr = fmt.Errorf("Docling remote convert failed")
|
||||
}
|
||||
return ParseResult{Err: fmt.Errorf("parser: Docling convert: %w", lastErr)}
|
||||
}
|
||||
|
||||
func doclingStandardPayload(filename string, encoded string, alpha bool) map[string]any {
|
||||
source := map[string]any{"filename": filename, "base64_string": encoded}
|
||||
options := map[string]any{"from_formats": []string{"pdf"}, "to_formats": []string{"json", "md", "text"}}
|
||||
if alpha {
|
||||
return map[string]any{
|
||||
"options": options,
|
||||
"file_sources": []map[string]any{source},
|
||||
}
|
||||
}
|
||||
source["kind"] = "file"
|
||||
return map[string]any{
|
||||
"options": options,
|
||||
"sources": []map[string]any{source},
|
||||
}
|
||||
}
|
||||
|
||||
func doclingChunkedPayload(filename string, encoded string, alpha bool) map[string]any {
|
||||
payload := doclingStandardPayload(filename, encoded, alpha)
|
||||
payload["options"] = map[string]any{
|
||||
"from_formats": []string{"pdf"},
|
||||
"to_formats": []string{"json", "md", "text"},
|
||||
"do_chunking": true,
|
||||
"chunking_options": map[string]any{
|
||||
"max_tokens": 512,
|
||||
"overlap": 50,
|
||||
"tokenizer": "sentencepiece",
|
||||
},
|
||||
}
|
||||
return payload
|
||||
}
|
||||
|
||||
func parseDoclingChunkedResult(filename string, body []byte, outputFormat string) (ParseResult, bool) {
|
||||
var chunks []doclingChunk
|
||||
if err := json.Unmarshal(body, &chunks); err != nil {
|
||||
var wrapped struct {
|
||||
Results []doclingChunk `json:"results"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &wrapped); err != nil {
|
||||
return ParseResult{}, false
|
||||
}
|
||||
chunks = wrapped.Results
|
||||
}
|
||||
texts := make([]string, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
text := strings.TrimSpace(chunk.Text)
|
||||
if text == "" && chunk.Chunk != nil {
|
||||
text = strings.TrimSpace(chunk.Chunk.Text)
|
||||
}
|
||||
if text != "" {
|
||||
texts = append(texts, text)
|
||||
}
|
||||
}
|
||||
if len(texts) == 0 {
|
||||
return ParseResult{}, false
|
||||
}
|
||||
pageCount := 0
|
||||
if len(texts) > 0 {
|
||||
pageCount = len(texts)
|
||||
}
|
||||
return doclingTextsToResult(filename, texts, outputFormat, pageCount), true
|
||||
}
|
||||
|
||||
func parseDoclingStandardResult(filename string, body []byte, outputFormat string) (ParseResult, bool) {
|
||||
var payload doclingResponse
|
||||
if err := json.Unmarshal(body, &payload); err != nil {
|
||||
return ParseResult{}, false
|
||||
}
|
||||
docs := make([]doclingDocument, 0, 1+len(payload.Documents)+len(payload.Results))
|
||||
if payload.Document != nil {
|
||||
docs = append(docs, *payload.Document)
|
||||
}
|
||||
docs = append(docs, payload.Documents...)
|
||||
for _, result := range payload.Results {
|
||||
switch {
|
||||
case result.Document != nil:
|
||||
docs = append(docs, *result.Document)
|
||||
case result.Result != nil:
|
||||
docs = append(docs, *result.Result)
|
||||
}
|
||||
}
|
||||
for _, doc := range docs {
|
||||
if md := strings.TrimSpace(doc.MDContent); md != "" {
|
||||
return parseMinerUMarkdownResult(filename, md, outputFormat, len(docs)), true
|
||||
}
|
||||
if txt := strings.TrimSpace(doc.TextContent); txt != "" {
|
||||
return doclingTextsToResult(filename, []string{txt}, outputFormat, len(docs)), true
|
||||
}
|
||||
if md, _ := doc.JSONContent["md_content"].(string); strings.TrimSpace(md) != "" {
|
||||
return parseMinerUMarkdownResult(filename, md, outputFormat, len(docs)), true
|
||||
}
|
||||
}
|
||||
return ParseResult{}, false
|
||||
}
|
||||
|
||||
func doclingTextsToResult(filename string, texts []string, outputFormat string, pageCount int) ParseResult {
|
||||
fileMeta := pdfFileMeta(filename, pageCount)
|
||||
switch strings.ToLower(strings.TrimSpace(outputFormat)) {
|
||||
case "", "json":
|
||||
items := make([]map[string]any, 0, len(texts))
|
||||
for _, text := range texts {
|
||||
items = append(items, map[string]any{
|
||||
"text": text,
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
}
|
||||
return ParseResult{
|
||||
OutputFormat: "json",
|
||||
File: fileMeta,
|
||||
JSON: items,
|
||||
}
|
||||
case "markdown":
|
||||
return ParseResult{
|
||||
OutputFormat: "markdown",
|
||||
File: fileMeta,
|
||||
Markdown: strings.Join(texts, "\n\n"),
|
||||
}
|
||||
default:
|
||||
return ParseResult{Err: fmt.Errorf("parser: unsupported PDF output_format %q", outputFormat)}
|
||||
}
|
||||
}
|
||||
161
internal/parser/parser/pdf_parser_docling_test.go
Normal file
161
internal/parser/parser/pdf_parser_docling_test.go
Normal file
@@ -0,0 +1,161 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPDFParser_ParseWithResult_DoclingChunkedMarkdownIntegration(t *testing.T) {
|
||||
var requestCount atomic.Int32
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
current := requestCount.Add(1)
|
||||
if got, want := r.Header.Get("Authorization"), "Bearer doc-secret"; got != want {
|
||||
t.Errorf("Authorization = %q, want %q", got, want)
|
||||
return
|
||||
}
|
||||
if r.Method != http.MethodPost || r.URL.Path != "/v1/convert/source" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Errorf("Decode: %v", err)
|
||||
return
|
||||
}
|
||||
options, _ := body["options"].(map[string]any)
|
||||
if got, want := options["do_chunking"], true; got != want {
|
||||
t.Errorf("do_chunking = %#v, want %#v", got, want)
|
||||
return
|
||||
}
|
||||
chunkingOptions, _ := options["chunking_options"].(map[string]any)
|
||||
if got, want := chunkingOptions["tokenizer"], "sentencepiece"; got != want {
|
||||
t.Errorf("chunking_options.tokenizer = %#v, want %#v", got, want)
|
||||
return
|
||||
}
|
||||
sources, _ := body["sources"].([]any)
|
||||
if len(sources) != 1 {
|
||||
t.Errorf("sources len = %d, want 1", len(sources))
|
||||
return
|
||||
}
|
||||
source, _ := sources[0].(map[string]any)
|
||||
raw, err := base64.StdEncoding.DecodeString(source["base64_string"].(string))
|
||||
if err != nil {
|
||||
t.Errorf("DecodeString: %v", err)
|
||||
return
|
||||
}
|
||||
if got := string(raw); !strings.HasPrefix(got, "%PDF") {
|
||||
t.Errorf("uploaded file = %q, want PDF bytes", got)
|
||||
return
|
||||
}
|
||||
if current != 1 {
|
||||
t.Errorf("request count = %d, want first chunked request only", current)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`[{"text":"Chunk A"},{"chunk":{"text":"Chunk B"}}]`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "Docling",
|
||||
"output_format": "markdown",
|
||||
"docling_server_url": server.URL,
|
||||
"docling_api_key": "doc-secret",
|
||||
})
|
||||
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "markdown"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := res.Markdown, "Chunk A\n\nChunk B"; got != want {
|
||||
t.Fatalf("Markdown = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_DoclingFallbackToStandardJSONIntegration(t *testing.T) {
|
||||
var requestCount atomic.Int32
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
current := requestCount.Add(1)
|
||||
switch current {
|
||||
case 1:
|
||||
if r.URL.Path != "/v1/convert/source" {
|
||||
t.Errorf("request 1 path = %q, want /v1/convert/source", r.URL.Path)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusUnprocessableEntity)
|
||||
_, _ = w.Write([]byte(`{"detail":"chunking unsupported"}`))
|
||||
case 2:
|
||||
if r.URL.Path != "/v1alpha/convert/source" {
|
||||
t.Errorf("request 2 path = %q, want /v1alpha/convert/source", r.URL.Path)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusUnprocessableEntity)
|
||||
_, _ = w.Write([]byte(`{"detail":"chunking unsupported"}`))
|
||||
case 3:
|
||||
if r.URL.Path != "/v1/convert/source" {
|
||||
t.Errorf("request 3 path = %q, want /v1/convert/source", r.URL.Path)
|
||||
return
|
||||
}
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Errorf("Decode: %v", err)
|
||||
return
|
||||
}
|
||||
options, _ := body["options"].(map[string]any)
|
||||
if _, exists := options["do_chunking"]; exists {
|
||||
t.Errorf("standard fallback payload unexpectedly contains do_chunking: %#v", options)
|
||||
return
|
||||
}
|
||||
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()
|
||||
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "Docling",
|
||||
"output_format": "json",
|
||||
"docling_server_url": server.URL,
|
||||
})
|
||||
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "json"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if len(res.JSON) == 0 {
|
||||
t.Fatal("JSON is empty; want markdown-normalized items")
|
||||
}
|
||||
if got, want := requestCount.Load(), int32(3); got != want {
|
||||
t.Fatalf("requestCount = %d, want %d", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_DoclingRequiresServerURL(t *testing.T) {
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{"parse_method": "Docling"})
|
||||
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err == nil {
|
||||
t.Fatal("ParseWithResult: want error when docling_server_url is missing, got nil")
|
||||
}
|
||||
if !strings.Contains(res.Err.Error(), "docling_server_url") {
|
||||
t.Fatalf("error = %q, want docling_server_url context", res.Err.Error())
|
||||
}
|
||||
}
|
||||
88
internal/parser/parser/pdf_parser_fixture_cgo_test.go
Normal file
88
internal/parser/parser/pdf_parser_fixture_cgo_test.go
Normal file
@@ -0,0 +1,88 @@
|
||||
//go:build cgo
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPDFParser_ParseWithResult_CGOFixture(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)
|
||||
}
|
||||
pdf := NewPDFParser()
|
||||
res := pdf.ParseWithResult("Doc1.pdf", data)
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "json"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if len(res.JSON) == 0 {
|
||||
t.Fatal("JSON is empty; want at least 1 parsed item")
|
||||
}
|
||||
if got := res.File["page_count"]; got == nil {
|
||||
t.Fatal("File.page_count missing")
|
||||
}
|
||||
if positions, ok := res.JSON[0]["_pdf_positions"].([][]any); ok && len(positions) == 0 {
|
||||
t.Fatal("JSON[0]._pdf_positions is empty; want normalized positions for fixture text")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_CGOFixtureMarkdown(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)
|
||||
}
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{"output_format": "markdown"})
|
||||
|
||||
res := pdf.ParseWithResult("Doc1.pdf", data)
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "markdown"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if res.Markdown == "" {
|
||||
t.Fatal("Markdown is empty; want rendered content")
|
||||
}
|
||||
if len(res.JSON) != 0 {
|
||||
t.Fatalf("JSON len = %d, want 0 for markdown output", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_CGOFixturePlainText(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)
|
||||
}
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{"parse_method": "plain_text"})
|
||||
res := pdf.ParseWithResult("Doc1.pdf", data)
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "json"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if len(res.JSON) == 0 {
|
||||
t.Fatal("JSON is empty; want page text items")
|
||||
}
|
||||
if got, _ := res.JSON[0]["text"].(string); strings.TrimSpace(got) == "" {
|
||||
t.Fatal("plain_text first page is empty")
|
||||
}
|
||||
}
|
||||
119
internal/parser/parser/pdf_parser_mineru.go
Normal file
119
internal/parser/parser/pdf_parser_mineru.go
Normal file
@@ -0,0 +1,119 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
models "ragflow/internal/entity/models"
|
||||
)
|
||||
|
||||
const minerUPollTimeout = 30 * time.Second
|
||||
const minerUPollInterval = 200 * time.Millisecond
|
||||
|
||||
func parsePDFWithMinerU(filename string, data []byte, parser *PDFParser) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
apiServer := strings.TrimSpace(parser.MinerUAPIServer)
|
||||
if apiServer == "" {
|
||||
apiServer = strings.TrimSpace(os.Getenv("MINERU_APISERVER"))
|
||||
}
|
||||
if apiServer == "" {
|
||||
return ParseResult{Err: fmt.Errorf("parser: MinerU requires mineru_apiserver or MINERU_APISERVER")}
|
||||
}
|
||||
apiKey := parser.MinerUAPIKey
|
||||
if strings.TrimSpace(apiKey) == "" {
|
||||
apiKey = strings.TrimSpace(os.Getenv("MINERU_API_KEY"))
|
||||
}
|
||||
backend := strings.TrimSpace(parser.MinerUBackend)
|
||||
if backend == "" {
|
||||
backend = strings.TrimSpace(os.Getenv("MINERU_BACKEND"))
|
||||
}
|
||||
if backend == "" {
|
||||
backend = "pipeline"
|
||||
}
|
||||
timeout := parser.MinerUPollTimeout
|
||||
if timeout <= 0 {
|
||||
timeout = minerUPollTimeout
|
||||
}
|
||||
|
||||
driver := models.NewMinerLocalUModel(
|
||||
map[string]string{"default": apiServer},
|
||||
models.URLSuffix{DocumentParse: "file_parse", Task: "tasks"},
|
||||
)
|
||||
apiConfig := &models.APIConfig{
|
||||
BaseURL: &apiServer,
|
||||
}
|
||||
if apiKey != "" {
|
||||
apiConfig.ApiKey = &apiKey
|
||||
}
|
||||
|
||||
task, err := driver.ParseFile(&backend, data, nil, apiConfig, &models.ParseFileConfig{})
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: MinerU submit: %w", err)}
|
||||
}
|
||||
content, err := pollMinerUTask(driver, task.TaskID, apiConfig, timeout)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: MinerU result: %w", err)}
|
||||
}
|
||||
pageCount := 0
|
||||
if strings.TrimSpace(content) != "" {
|
||||
pageCount = 1
|
||||
}
|
||||
return parseMinerUMarkdownResult(filename, content, parser.OutputFormat, pageCount)
|
||||
}
|
||||
|
||||
func pollMinerUTask(driver *models.MinerULocalModel, taskID string, apiConfig *models.APIConfig, timeout time.Duration) (string, error) {
|
||||
if timeout <= 0 {
|
||||
timeout = minerUPollTimeout
|
||||
}
|
||||
deadline := time.Now().Add(timeout)
|
||||
var lastErr error
|
||||
for {
|
||||
task, err := driver.ShowTask(taskID, apiConfig)
|
||||
if err == nil {
|
||||
for _, segment := range task.Segments {
|
||||
if strings.TrimSpace(segment.Content) != "" {
|
||||
return segment.Content, nil
|
||||
}
|
||||
}
|
||||
lastErr = fmt.Errorf("empty MinerU task content")
|
||||
} else {
|
||||
lastErr = err
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
if lastErr == nil {
|
||||
lastErr = fmt.Errorf("timed out waiting for MinerU task %s", taskID)
|
||||
}
|
||||
return "", lastErr
|
||||
}
|
||||
time.Sleep(minerUPollInterval)
|
||||
}
|
||||
}
|
||||
|
||||
func parseMinerUMarkdownResult(filename, markdown, outputFormat string, pageCount int) ParseResult {
|
||||
fileMeta := pdfFileMeta(filename, pageCount)
|
||||
switch strings.ToLower(strings.TrimSpace(outputFormat)) {
|
||||
case "", "json":
|
||||
mp, err := NewMarkdownParser(GoMarkdown)
|
||||
if err != nil {
|
||||
return ParseResult{Err: err}
|
||||
}
|
||||
res := mp.ParseWithResult(filename, []byte(markdown))
|
||||
if res.Err != nil {
|
||||
return res
|
||||
}
|
||||
res.File = fileMeta
|
||||
return res
|
||||
case "markdown":
|
||||
return ParseResult{
|
||||
OutputFormat: "markdown",
|
||||
File: fileMeta,
|
||||
Markdown: markdown,
|
||||
}
|
||||
default:
|
||||
return ParseResult{Err: fmt.Errorf("parser: unsupported PDF output_format %q", outputFormat)}
|
||||
}
|
||||
}
|
||||
161
internal/parser/parser/pdf_parser_mineru_test.go
Normal file
161
internal/parser/parser/pdf_parser_mineru_test.go
Normal file
@@ -0,0 +1,161 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPDFParser_ParseWithResult_MinerUMarkdownIntegration(t *testing.T) {
|
||||
var submitCalled atomic.Bool
|
||||
var resultCalled atomic.Bool
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/file_parse":
|
||||
submitCalled.Store(true)
|
||||
if got, want := r.Header.Get("Authorization"), "Bearer secret"; got != want {
|
||||
t.Errorf("Authorization = %q, want %q", got, want)
|
||||
return
|
||||
}
|
||||
reader, err := r.MultipartReader()
|
||||
if err != nil {
|
||||
t.Errorf("MultipartReader: %v", err)
|
||||
return
|
||||
}
|
||||
var backend string
|
||||
var fileSeen bool
|
||||
for {
|
||||
part, err := reader.NextPart()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
t.Errorf("NextPart: %v", err)
|
||||
return
|
||||
}
|
||||
if part.FormName() == "backend" {
|
||||
body, _ := io.ReadAll(part)
|
||||
backend = string(body)
|
||||
}
|
||||
if part.FormName() == "files" {
|
||||
fileSeen = true
|
||||
body, _ := io.ReadAll(part)
|
||||
if !strings.HasPrefix(string(body), "%PDF") {
|
||||
t.Errorf("uploaded file = %q, want PDF bytes", string(body))
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
if backend != "pipeline" {
|
||||
t.Errorf("backend = %q, want pipeline", backend)
|
||||
return
|
||||
}
|
||||
if !fileSeen {
|
||||
t.Error("multipart upload missing files part")
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"data":{"task_id":"task-1"}}`))
|
||||
case r.Method == http.MethodGet && r.URL.Path == "/tasks/task-1/result":
|
||||
resultCalled.Store(true)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"results":{"doc":{"md_content":"# Title\n\nBody paragraph.\n"}}}`))
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "MinerU",
|
||||
"output_format": "markdown",
|
||||
"mineru_apiserver": server.URL,
|
||||
"mineru_api_key": "secret",
|
||||
"mineru_backend": "pipeline",
|
||||
})
|
||||
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if !submitCalled.Load() || !resultCalled.Load() {
|
||||
t.Fatalf("submit/result called = %v/%v, want true/true", submitCalled.Load(), resultCalled.Load())
|
||||
}
|
||||
if got, want := res.OutputFormat, "markdown"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := res.Markdown, "# Title\n\nBody paragraph.\n"; got != want {
|
||||
t.Fatalf("Markdown = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := res.File["name"], "sample.pdf"; got != want {
|
||||
t.Fatalf("File.name = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_MinerUJSONIntegration(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-2"}}`))
|
||||
case r.Method == http.MethodGet && r.URL.Path == "/tasks/task-2/result":
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"results":{"doc":{"md_content":"# Title\n\nBody paragraph.\n"}}}`))
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "MinerU",
|
||||
"output_format": "json",
|
||||
"mineru_apiserver": server.URL,
|
||||
})
|
||||
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "json"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if len(res.JSON) == 0 {
|
||||
t.Fatal("JSON is empty; want markdown-normalized items")
|
||||
}
|
||||
if got := res.JSON[0]["text"]; got == nil {
|
||||
t.Fatal("JSON[0].text missing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_MinerURequiresAPIServer(t *testing.T) {
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{"parse_method": "MinerU"})
|
||||
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err == nil {
|
||||
t.Fatal("ParseWithResult: want error when mineru_apiserver is missing, got nil")
|
||||
}
|
||||
if !strings.Contains(res.Err.Error(), "mineru_apiserver") {
|
||||
t.Fatalf("error = %q, want mineru_apiserver context", res.Err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestMinerUUploadShape(t *testing.T) {
|
||||
body := &strings.Builder{}
|
||||
writer := multipart.NewWriter(body)
|
||||
_ = writer.WriteField("backend", "pipeline")
|
||||
if err := writer.Close(); err != nil {
|
||||
t.Fatalf("Close: %v", err)
|
||||
}
|
||||
if !strings.Contains(body.String(), "pipeline") {
|
||||
t.Fatalf("multipart body = %q, want backend field", body.String())
|
||||
}
|
||||
}
|
||||
@@ -7,6 +7,25 @@ import (
|
||||
)
|
||||
|
||||
func (p *PDFParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
if err := p.validateParseMethod(); err != nil {
|
||||
return ParseResult{Err: err}
|
||||
}
|
||||
switch normalizePDFParseMethod(p.ParseMethod) {
|
||||
case "plain_text":
|
||||
return parsePDFWithPlainText(filename, data, p)
|
||||
case "mineru":
|
||||
return parsePDFWithMinerU(filename, data, p)
|
||||
case "paddleocr":
|
||||
return parsePDFWithPaddleOCR(filename, data, p)
|
||||
case "docling":
|
||||
return parsePDFWithDocling(filename, data, p)
|
||||
case "opendataloader":
|
||||
return parsePDFWithOpenDataLoader(filename, data, p)
|
||||
case "somark":
|
||||
return parsePDFWithSoMark(filename, data, p)
|
||||
case "tcadp":
|
||||
return parsePDFWithTCADP(filename, data, p)
|
||||
}
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
|
||||
231
internal/parser/parser/pdf_parser_opendataloader.go
Normal file
231
internal/parser/parser/pdf_parser_opendataloader.go
Normal file
@@ -0,0 +1,231 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
models "ragflow/internal/entity/models"
|
||||
)
|
||||
|
||||
func parsePDFWithOpenDataLoader(filename string, data []byte, parser *PDFParser) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
baseURL := strings.TrimSpace(parser.OpenDataLoaderAPIServer)
|
||||
if baseURL == "" {
|
||||
baseURL = strings.TrimSpace(os.Getenv("OPENDATALOADER_APISERVER"))
|
||||
}
|
||||
if baseURL == "" {
|
||||
return ParseResult{Err: fmt.Errorf("parser: OpenDataLoader requires opendataloader_apiserver or OPENDATALOADER_APISERVER")}
|
||||
}
|
||||
apiKey := strings.TrimSpace(parser.OpenDataLoaderAPIKey)
|
||||
if apiKey == "" {
|
||||
apiKey = strings.TrimSpace(os.Getenv("OPENDATALOADER_API_KEY"))
|
||||
}
|
||||
|
||||
bodyReader, contentType, err := openDataLoaderMultipart(filename, data, parser)
|
||||
if err != nil {
|
||||
return ParseResult{Err: err}
|
||||
}
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, strings.TrimRight(baseURL, "/")+"/file_parse", bodyReader)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: OpenDataLoader request: %w", err)}
|
||||
}
|
||||
req.Header.Set("Content-Type", contentType)
|
||||
if apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
resp, err := models.NewDriverHTTPClient().Do(req)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: OpenDataLoader submit: %w", err)}
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: OpenDataLoader read: %w", err)}
|
||||
}
|
||||
if resp.StatusCode >= 300 {
|
||||
return ParseResult{Err: fmt.Errorf("parser: OpenDataLoader HTTP %d: %s", resp.StatusCode, string(raw))}
|
||||
}
|
||||
var payload struct {
|
||||
JSONDoc any `json:"json_doc"`
|
||||
MDText string `json:"md_text"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &payload); err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: OpenDataLoader decode: %w", err)}
|
||||
}
|
||||
if payload.JSONDoc != nil {
|
||||
items := openDataLoaderItems(payload.JSONDoc)
|
||||
if len(items) > 0 {
|
||||
return pdfItemsToResult(filename, items, parser.OutputFormat, openDataLoaderPageCount(payload.JSONDoc))
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(payload.MDText) != "" {
|
||||
return parseMinerUMarkdownResult(filename, payload.MDText, parser.OutputFormat, 1)
|
||||
}
|
||||
return ParseResult{Err: fmt.Errorf("parser: OpenDataLoader returned no parsed content")}
|
||||
}
|
||||
|
||||
func openDataLoaderMultipart(filename string, data []byte, parser *PDFParser) (io.Reader, string, error) {
|
||||
var body strings.Builder
|
||||
writer := multipart.NewWriter(&body)
|
||||
part, err := writer.CreateFormFile("file", filename)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("parser: OpenDataLoader create form file: %w", err)
|
||||
}
|
||||
if _, err := part.Write(data); err != nil {
|
||||
return nil, "", fmt.Errorf("parser: OpenDataLoader write PDF: %w", err)
|
||||
}
|
||||
if parser.OpenDataLoaderHybrid != "" {
|
||||
_ = writer.WriteField("hybrid", parser.OpenDataLoaderHybrid)
|
||||
}
|
||||
if parser.OpenDataLoaderImageOutput != "" {
|
||||
_ = writer.WriteField("image_output", parser.OpenDataLoaderImageOutput)
|
||||
}
|
||||
if parser.OpenDataLoaderSanitize != nil {
|
||||
if *parser.OpenDataLoaderSanitize {
|
||||
_ = writer.WriteField("sanitize", "true")
|
||||
} else {
|
||||
_ = writer.WriteField("sanitize", "false")
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
return nil, "", fmt.Errorf("parser: OpenDataLoader finalize form: %w", err)
|
||||
}
|
||||
return strings.NewReader(body.String()), writer.FormDataContentType(), nil
|
||||
}
|
||||
|
||||
func openDataLoaderItems(root any) []map[string]any {
|
||||
items := []map[string]any{}
|
||||
var walk func(any)
|
||||
walk = func(node any) {
|
||||
switch v := node.(type) {
|
||||
case map[string]any:
|
||||
if item := openDataLoaderNodeToItem(v); item != nil {
|
||||
items = append(items, item)
|
||||
}
|
||||
for _, child := range v {
|
||||
walk(child)
|
||||
}
|
||||
case []any:
|
||||
for _, child := range v {
|
||||
walk(child)
|
||||
}
|
||||
}
|
||||
}
|
||||
walk(root)
|
||||
return items
|
||||
}
|
||||
|
||||
func openDataLoaderNodeToItem(el map[string]any) map[string]any {
|
||||
rawType, _ := el["type"].(string)
|
||||
t := strings.ToLower(strings.TrimSpace(rawType))
|
||||
if t == "" {
|
||||
return nil
|
||||
}
|
||||
text := strings.TrimSpace(stringValue(el["content"]))
|
||||
if text == "" {
|
||||
text = strings.TrimSpace(stringValue(el["text"]))
|
||||
}
|
||||
switch t {
|
||||
case "table":
|
||||
html := strings.TrimSpace(stringValue(el["html"]))
|
||||
if html == "" {
|
||||
html = strings.TrimSpace(stringValue(el["html_content"]))
|
||||
}
|
||||
if html != "" {
|
||||
text = html
|
||||
}
|
||||
if text == "" {
|
||||
text = openDataLoaderCellsText(el["cells"])
|
||||
}
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
return map[string]any{"text": text, "doc_type_kwd": "table", "layout": "table"}
|
||||
case "image", "picture", "figure":
|
||||
if text == "" {
|
||||
text = "[Image]"
|
||||
}
|
||||
return map[string]any{"text": text, "doc_type_kwd": "image", "layout": "figure"}
|
||||
case "formula", "equation":
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
return map[string]any{"text": text, "doc_type_kwd": "text", "layout": "equation"}
|
||||
default:
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
layout := "text"
|
||||
if t == "title" || t == "heading" {
|
||||
layout = "title"
|
||||
}
|
||||
return map[string]any{"text": text, "doc_type_kwd": "text", "layout": layout}
|
||||
}
|
||||
}
|
||||
|
||||
func openDataLoaderCellsText(raw any) string {
|
||||
cells, ok := raw.([]any)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
rows := make(map[int][]string)
|
||||
maxRow := -1
|
||||
for _, cellRaw := range cells {
|
||||
cell, ok := cellRaw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
row := int(numberValue(cell["row"]))
|
||||
if row == 0 {
|
||||
row = int(numberValue(cell["row_index"]))
|
||||
}
|
||||
rows[row] = append(rows[row], stringValue(cell["content"]))
|
||||
if row > maxRow {
|
||||
maxRow = row
|
||||
}
|
||||
}
|
||||
if len(rows) == 0 || maxRow < 0 {
|
||||
return ""
|
||||
}
|
||||
parts := make([]string, 0, len(rows))
|
||||
for i := 0; i <= maxRow; i++ {
|
||||
if cols, ok := rows[i]; ok {
|
||||
parts = append(parts, strings.Join(cols, " | "))
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func openDataLoaderPageCount(root any) int {
|
||||
pages := collectPDFPageNumbers(root)
|
||||
if len(pages) > 0 {
|
||||
return len(pages)
|
||||
}
|
||||
if root != nil {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func stringValue(v any) string {
|
||||
s, _ := v.(string)
|
||||
return s
|
||||
}
|
||||
|
||||
func numberValue(v any) float64 {
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return n
|
||||
case int:
|
||||
return float64(n)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
152
internal/parser/parser/pdf_parser_opendataloader_test.go
Normal file
152
internal/parser/parser/pdf_parser_opendataloader_test.go
Normal file
@@ -0,0 +1,152 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPDFParser_ParseWithResult_OpenDataLoaderJSONIntegration(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
|
||||
}
|
||||
if got, want := r.Header.Get("Authorization"), "Bearer odl-secret"; got != want {
|
||||
t.Errorf("Authorization = %q, want %q", got, want)
|
||||
return
|
||||
}
|
||||
reader, err := r.MultipartReader()
|
||||
if err != nil {
|
||||
t.Errorf("MultipartReader: %v", err)
|
||||
return
|
||||
}
|
||||
seenHybrid := false
|
||||
seenSanitize := false
|
||||
for {
|
||||
part, err := reader.NextPart()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
t.Errorf("NextPart: %v", err)
|
||||
return
|
||||
}
|
||||
switch part.FormName() {
|
||||
case "file":
|
||||
body, _ := io.ReadAll(part)
|
||||
if !strings.HasPrefix(string(body), "%PDF") {
|
||||
t.Errorf("uploaded file = %q, want PDF bytes", string(body))
|
||||
return
|
||||
}
|
||||
case "hybrid":
|
||||
body, _ := io.ReadAll(part)
|
||||
seenHybrid = string(body) == "docling-fast"
|
||||
case "sanitize":
|
||||
body, _ := io.ReadAll(part)
|
||||
seenSanitize = string(body) == "true"
|
||||
}
|
||||
}
|
||||
if !seenHybrid || !seenSanitize {
|
||||
t.Errorf("multipart fields missing: hybrid=%v sanitize=%v", seenHybrid, seenSanitize)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"json_doc":{"type":"title","content":"ODL Title","children":[{"type":"paragraph","content":"ODL Body"},{"type":"table","html":"<table><tr><td>a</td></tr></table>"}]}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
sanitize := true
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "OpenDataLoader",
|
||||
"output_format": "json",
|
||||
"opendataloader_apiserver": server.URL,
|
||||
"opendataloader_api_key": "odl-secret",
|
||||
"hybrid": "docling-fast",
|
||||
"sanitize": sanitize,
|
||||
})
|
||||
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "json"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if len(res.JSON) < 2 {
|
||||
t.Fatalf("JSON len = %d, want >=2", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_OpenDataLoaderMarkdownFallback(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if 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()
|
||||
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "OpenDataLoader",
|
||||
"output_format": "markdown",
|
||||
"opendataloader_apiserver": server.URL,
|
||||
})
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if !strings.Contains(res.Markdown, "ODL Title") {
|
||||
t.Fatalf("Markdown = %q, want ODL Title", res.Markdown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_OpenDataLoaderRequiresAPIServer(t *testing.T) {
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{"parse_method": "OpenDataLoader"})
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err == nil || !strings.Contains(res.Err.Error(), "opendataloader_apiserver") {
|
||||
t.Fatalf("error = %v, want opendataloader_apiserver context", res.Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenDataLoaderItems_TableCellsFallback(t *testing.T) {
|
||||
root := map[string]any{
|
||||
"type": "table",
|
||||
"cells": []any{
|
||||
map[string]any{"row": 0, "content": "a"},
|
||||
map[string]any{"row": 0, "content": "b"},
|
||||
},
|
||||
}
|
||||
items := openDataLoaderItems(root)
|
||||
if len(items) != 1 {
|
||||
t.Fatalf("items len = %d, want 1", len(items))
|
||||
}
|
||||
if got, _ := json.Marshal(items[0]); !strings.Contains(string(got), "a | b") {
|
||||
t.Fatalf("item = %s, want row text", string(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenDataLoaderItems_TableCellsFallbackSparseRows(t *testing.T) {
|
||||
root := map[string]any{
|
||||
"type": "table",
|
||||
"cells": []any{
|
||||
map[string]any{"row": 0, "content": "a"},
|
||||
map[string]any{"row": 5, "content": "z"},
|
||||
},
|
||||
}
|
||||
items := openDataLoaderItems(root)
|
||||
if len(items) != 1 {
|
||||
t.Fatalf("items len = %d, want 1", len(items))
|
||||
}
|
||||
if got, _ := json.Marshal(items[0]); !strings.Contains(string(got), "z") {
|
||||
t.Fatalf("item = %s, want sparse row text", string(got))
|
||||
}
|
||||
}
|
||||
62
internal/parser/parser/pdf_parser_paddleocr.go
Normal file
62
internal/parser/parser/pdf_parser_paddleocr.go
Normal file
@@ -0,0 +1,62 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
models "ragflow/internal/entity/models"
|
||||
)
|
||||
|
||||
func parsePDFWithPaddleOCR(filename string, data []byte, parser *PDFParser) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
baseURL := strings.TrimSpace(parser.PaddleOCRBaseURL)
|
||||
if baseURL == "" {
|
||||
baseURL = strings.TrimSpace(os.Getenv("PADDLEOCR_BASE_URL"))
|
||||
}
|
||||
if baseURL == "" {
|
||||
baseURL = strings.TrimSpace(os.Getenv("PADDLEOCR_API_URL"))
|
||||
}
|
||||
if baseURL == "" {
|
||||
return ParseResult{Err: fmt.Errorf("parser: PaddleOCR requires paddleocr_base_url or PADDLEOCR_BASE_URL")}
|
||||
}
|
||||
apiKey := parser.PaddleOCRAPIKey
|
||||
if strings.TrimSpace(apiKey) == "" {
|
||||
apiKey = strings.TrimSpace(os.Getenv("PADDLEOCR_ACCESS_TOKEN"))
|
||||
}
|
||||
algorithm := strings.TrimSpace(parser.PaddleOCRAlgorithm)
|
||||
if algorithm == "" {
|
||||
algorithm = strings.TrimSpace(os.Getenv("PADDLEOCR_ALGORITHM"))
|
||||
}
|
||||
if algorithm == "" {
|
||||
algorithm = "PaddleOCR-VL"
|
||||
}
|
||||
|
||||
driver := models.NewPaddleOCRLocalModel(
|
||||
map[string]string{"default": baseURL},
|
||||
models.URLSuffix{OCR: "layout-parsing"},
|
||||
)
|
||||
apiConfig := &models.APIConfig{
|
||||
BaseURL: &baseURL,
|
||||
}
|
||||
if apiKey != "" {
|
||||
apiConfig.ApiKey = &apiKey
|
||||
}
|
||||
|
||||
resp, err := driver.OCRFile(&algorithm, data, &filename, apiConfig, &models.OCRConfig{
|
||||
Algorithm: algorithm,
|
||||
})
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: PaddleOCR OCRFile: %w", err)}
|
||||
}
|
||||
if resp == nil || resp.Text == nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: PaddleOCR returned empty text")}
|
||||
}
|
||||
pageCount := 1
|
||||
if resp.Text != nil && strings.TrimSpace(*resp.Text) == "" {
|
||||
pageCount = 0
|
||||
}
|
||||
return parseMinerUMarkdownResult(filename, *resp.Text, parser.OutputFormat, pageCount)
|
||||
}
|
||||
127
internal/parser/parser/pdf_parser_paddleocr_test.go
Normal file
127
internal/parser/parser/pdf_parser_paddleocr_test.go
Normal file
@@ -0,0 +1,127 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPDFParser_ParseWithResult_PaddleOCRMarkdownIntegration(t *testing.T) {
|
||||
var called atomic.Bool
|
||||
|
||||
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
|
||||
}
|
||||
called.Store(true)
|
||||
if got, want := r.Header.Get("Authorization"), "Bearer paddle-secret"; got != want {
|
||||
t.Errorf("Authorization = %q, want %q", got, want)
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
File string `json:"file"`
|
||||
FileType int `json:"fileType"`
|
||||
Algorithm string `json:"algorithm"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Errorf("Decode: %v", err)
|
||||
return
|
||||
}
|
||||
if got, want := body.FileType, 0; got != want {
|
||||
t.Errorf("fileType = %d, want %d for PDF", got, want)
|
||||
return
|
||||
}
|
||||
if got, want := body.Algorithm, "PaddleOCR-VL"; got != want {
|
||||
t.Errorf("algorithm = %q, want %q", got, want)
|
||||
return
|
||||
}
|
||||
raw, err := base64.StdEncoding.DecodeString(body.File)
|
||||
if err != nil {
|
||||
t.Errorf("DecodeString: %v", err)
|
||||
return
|
||||
}
|
||||
if got := string(raw); !strings.HasPrefix(got, "%PDF") {
|
||||
t.Errorf("uploaded file = %q, want PDF bytes", got)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"errorCode":0,"result":{"layoutParsingResults":[{"markdown":{"text":"# Paddle Title\n\nBody paragraph.\n"}}]}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "PaddleOCR",
|
||||
"output_format": "markdown",
|
||||
"paddleocr_base_url": server.URL,
|
||||
"paddleocr_api_key": "paddle-secret",
|
||||
})
|
||||
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if !called.Load() {
|
||||
t.Fatal("PaddleOCR server was not called")
|
||||
}
|
||||
if got, want := res.OutputFormat, "markdown"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := res.Markdown, "# Paddle Title\n\nBody paragraph."; got != want {
|
||||
t.Fatalf("Markdown = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := res.File["name"], "sample.pdf"; got != want {
|
||||
t.Fatalf("File.name = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_PaddleOCRJSONIntegration(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
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"errorCode":0,"result":{"layoutParsingResults":[{"markdown":{"text":"# Paddle Title\n\nBody paragraph.\n"}}]}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "PaddleOCR",
|
||||
"output_format": "json",
|
||||
"paddleocr_base_url": server.URL,
|
||||
})
|
||||
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "json"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if len(res.JSON) == 0 {
|
||||
t.Fatal("JSON is empty; want markdown-normalized items")
|
||||
}
|
||||
if got := res.JSON[0]["text"]; got == nil {
|
||||
t.Fatal("JSON[0].text missing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_PaddleOCRRequiresBaseURL(t *testing.T) {
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{"parse_method": "PaddleOCR"})
|
||||
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err == nil {
|
||||
t.Fatal("ParseWithResult: want error when paddleocr_base_url is missing, got nil")
|
||||
}
|
||||
if !strings.Contains(res.Err.Error(), "paddleocr_base_url") {
|
||||
t.Fatalf("error = %q, want paddleocr_base_url context", res.Err.Error())
|
||||
}
|
||||
}
|
||||
38
internal/parser/parser/pdf_parser_plaintext_cgo.go
Normal file
38
internal/parser/parser/pdf_parser_plaintext_cgo.go
Normal file
@@ -0,0 +1,38 @@
|
||||
//go:build cgo
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"ragflow/internal/deepdoc/parser/pdf/pdfoxide"
|
||||
)
|
||||
|
||||
func parsePDFWithPlainText(filename string, data []byte, parser *PDFParser) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
doc, err := pdfoxide.OpenBytes(data)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: plain_text open: %w", err)}
|
||||
}
|
||||
defer doc.Close()
|
||||
|
||||
pageCount, err := doc.PageCount()
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: plain_text page count: %w", err)}
|
||||
}
|
||||
items := make([]map[string]any, 0, pageCount)
|
||||
for page := 0; page < pageCount; page++ {
|
||||
text, err := doc.GetPageText(page)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: plain_text page %d: %w", page+1, err)}
|
||||
}
|
||||
items = append(items, map[string]any{
|
||||
"text": text,
|
||||
"doc_type_kwd": "text",
|
||||
"page_number": page + 1,
|
||||
})
|
||||
}
|
||||
return pdfItemsToResult(filename, items, parser.OutputFormat, pageCount)
|
||||
}
|
||||
12
internal/parser/parser/pdf_parser_plaintext_nocgo.go
Normal file
12
internal/parser/parser/pdf_parser_plaintext_nocgo.go
Normal file
@@ -0,0 +1,12 @@
|
||||
//go:build !cgo
|
||||
|
||||
package parser
|
||||
|
||||
import "fmt"
|
||||
|
||||
func parsePDFWithPlainText(filename string, data []byte, parser *PDFParser) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
return ParseResult{Err: fmt.Errorf("%w: %s", ErrPDFEngineUnavailable, filename)}
|
||||
}
|
||||
109
internal/parser/parser/pdf_parser_remote_common.go
Normal file
109
internal/parser/parser/pdf_parser_remote_common.go
Normal file
@@ -0,0 +1,109 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func pdfFileMeta(filename string, pageCount int) map[string]any {
|
||||
if pageCount < 0 {
|
||||
pageCount = 0
|
||||
}
|
||||
return map[string]any{
|
||||
"name": filename,
|
||||
"page_count": pageCount,
|
||||
"outline": []map[string]any{},
|
||||
}
|
||||
}
|
||||
|
||||
func pdfItemsToResult(filename string, items []map[string]any, outputFormat string, pageCount int) ParseResult {
|
||||
if len(items) == 0 {
|
||||
items = []map[string]any{{"text": "", "doc_type_kwd": "text"}}
|
||||
}
|
||||
pageCount = normalizePDFPageCount(pageCount, items)
|
||||
switch strings.ToLower(strings.TrimSpace(outputFormat)) {
|
||||
case "", "json":
|
||||
return ParseResult{
|
||||
OutputFormat: "json",
|
||||
File: pdfFileMeta(filename, pageCount),
|
||||
JSON: items,
|
||||
}
|
||||
case "markdown":
|
||||
var b strings.Builder
|
||||
for _, item := range items {
|
||||
text, _ := item["text"].(string)
|
||||
layout, _ := item["layout"].(string)
|
||||
if strings.TrimSpace(text) == "" {
|
||||
continue
|
||||
}
|
||||
if layout == "title" && !strings.HasPrefix(strings.TrimSpace(text), "#") {
|
||||
b.WriteString("## ")
|
||||
}
|
||||
b.WriteString(text)
|
||||
b.WriteString("\n\n")
|
||||
}
|
||||
return ParseResult{
|
||||
OutputFormat: "markdown",
|
||||
File: pdfFileMeta(filename, pageCount),
|
||||
Markdown: strings.TrimRight(b.String(), "\n"),
|
||||
}
|
||||
default:
|
||||
return ParseResult{Err: fmt.Errorf("parser: unsupported PDF output_format %q", outputFormat)}
|
||||
}
|
||||
}
|
||||
|
||||
func normalizePDFPageCount(fallback int, items []map[string]any) int {
|
||||
if inferred := inferPDFPageCountFromItems(items); inferred > fallback {
|
||||
return inferred
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func inferPDFPageCountFromItems(items []map[string]any) int {
|
||||
pages := map[int]struct{}{}
|
||||
for _, item := range items {
|
||||
for page := range collectPDFPageNumbers(item) {
|
||||
pages[page] = struct{}{}
|
||||
}
|
||||
}
|
||||
return len(pages)
|
||||
}
|
||||
|
||||
func collectPDFPageNumbers(raw any) map[int]struct{} {
|
||||
pages := map[int]struct{}{}
|
||||
var walk func(any)
|
||||
walk = func(node any) {
|
||||
switch v := node.(type) {
|
||||
case map[string]any:
|
||||
for _, key := range []string{"page_number", "page_num", "page_no", "page_index", "page_idx", "page"} {
|
||||
if page := int(numberValue(v[key])); page > 0 {
|
||||
pages[page] = struct{}{}
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"_pdf_positions", "positions"} {
|
||||
if positions, ok := v[key].([][]any); ok {
|
||||
for _, pos := range positions {
|
||||
if len(pos) > 0 {
|
||||
if page := int(numberValue(pos[0])); page > 0 {
|
||||
pages[page] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, child := range v {
|
||||
walk(child)
|
||||
}
|
||||
case []any:
|
||||
for _, child := range v {
|
||||
walk(child)
|
||||
}
|
||||
case []map[string]any:
|
||||
for _, child := range v {
|
||||
walk(child)
|
||||
}
|
||||
}
|
||||
}
|
||||
walk(raw)
|
||||
return pages
|
||||
}
|
||||
257
internal/parser/parser/pdf_parser_somark.go
Normal file
257
internal/parser/parser/pdf_parser_somark.go
Normal file
@@ -0,0 +1,257 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
models "ragflow/internal/entity/models"
|
||||
)
|
||||
|
||||
func parsePDFWithSoMark(filename string, data []byte, parser *PDFParser) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
baseURL := strings.TrimSpace(parser.SoMarkBaseURL)
|
||||
if baseURL == "" {
|
||||
baseURL = strings.TrimSpace(os.Getenv("SOMARK_BASE_URL"))
|
||||
}
|
||||
if baseURL == "" {
|
||||
return ParseResult{Err: fmt.Errorf("parser: SoMark requires somark_base_url or SOMARK_BASE_URL")}
|
||||
}
|
||||
apiKey := parser.SoMarkAPIKey
|
||||
if strings.TrimSpace(apiKey) == "" {
|
||||
apiKey = strings.TrimSpace(os.Getenv("SOMARK_API_KEY"))
|
||||
}
|
||||
taskID, err := soMarkSubmit(strings.TrimRight(baseURL, "/"), filename, data, parser, apiKey)
|
||||
if err != nil {
|
||||
return ParseResult{Err: err}
|
||||
}
|
||||
result, err := soMarkPoll(strings.TrimRight(baseURL, "/"), taskID, apiKey)
|
||||
if err != nil {
|
||||
return ParseResult{Err: err}
|
||||
}
|
||||
items, pageCount := soMarkItems(result, parser.SoMarkKeepHeaderFooter)
|
||||
if len(items) == 0 {
|
||||
return ParseResult{Err: fmt.Errorf("parser: SoMark returned no usable blocks")}
|
||||
}
|
||||
return pdfItemsToResult(filename, items, parser.OutputFormat, pageCount)
|
||||
}
|
||||
|
||||
func soMarkSubmit(baseURL, filename string, data []byte, parser *PDFParser, apiKey string) (string, error) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
part, err := writer.CreateFormFile("file", filename)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("parser: SoMark create file: %w", err)
|
||||
}
|
||||
if _, err := part.Write(data); err != nil {
|
||||
return "", fmt.Errorf("parser: SoMark write file: %w", err)
|
||||
}
|
||||
_ = writer.WriteField("output_formats", "json")
|
||||
elementFormats, _ := json.Marshal(map[string]any{
|
||||
"image": envOrDefault("SOMARK_IMAGE_FORMAT", parser.SoMarkImageFormat, "url"),
|
||||
"formula": envOrDefault("SOMARK_FORMULA_FORMAT", parser.SoMarkFormulaFormat, "latex"),
|
||||
"table": envOrDefault("SOMARK_TABLE_FORMAT", parser.SoMarkTableFormat, "html"),
|
||||
"cs": envOrDefault("SOMARK_CS_FORMAT", parser.SoMarkCSFormat, "image"),
|
||||
})
|
||||
featureConfig, _ := json.Marshal(map[string]any{
|
||||
"enable_text_cross_page": envOrBool("SOMARK_ENABLE_TEXT_CROSS_PAGE", parser.SoMarkEnableTextCrossPage),
|
||||
"enable_table_cross_page": envOrBool("SOMARK_ENABLE_TABLE_CROSS_PAGE", parser.SoMarkEnableTableCrossPage),
|
||||
"enable_title_level_recognition": envOrBool("SOMARK_ENABLE_TITLE_LEVEL_RECOGNITION", parser.SoMarkEnableTitleLevelRecognition),
|
||||
"enable_inline_image": envOrBool("SOMARK_ENABLE_INLINE_IMAGE", parser.SoMarkEnableInlineImage),
|
||||
"enable_table_image": envOrBool("SOMARK_ENABLE_TABLE_IMAGE", parser.SoMarkEnableTableImage),
|
||||
"enable_image_understanding": envOrBool("SOMARK_ENABLE_IMAGE_UNDERSTANDING", parser.SoMarkEnableImageUnderstanding),
|
||||
"keep_header_footer": envOrBool("SOMARK_KEEP_HEADER_FOOTER", parser.SoMarkKeepHeaderFooter),
|
||||
})
|
||||
_ = writer.WriteField("element_formats", string(elementFormats))
|
||||
_ = writer.WriteField("feature_config", string(featureConfig))
|
||||
if apiKey != "" {
|
||||
_ = writer.WriteField("api_key", apiKey)
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
return "", fmt.Errorf("parser: SoMark finalize form: %w", err)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, baseURL+"/parse/async", &body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("parser: SoMark request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
resp, err := models.NewDriverHTTPClient().Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("parser: SoMark submit: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("parser: SoMark read submit: %w", err)
|
||||
}
|
||||
var payload struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Data struct {
|
||||
TaskID string `json:"task_id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &payload); err != nil {
|
||||
return "", fmt.Errorf("parser: SoMark decode submit: %w", err)
|
||||
}
|
||||
if payload.Code != 0 {
|
||||
return "", fmt.Errorf("parser: SoMark submit business error code=%d message=%s", payload.Code, payload.Message)
|
||||
}
|
||||
if payload.Data.TaskID == "" {
|
||||
return "", fmt.Errorf("parser: SoMark submit returned no task_id")
|
||||
}
|
||||
return payload.Data.TaskID, nil
|
||||
}
|
||||
|
||||
func soMarkPoll(baseURL, taskID, apiKey string) (map[string]any, error) {
|
||||
form := url.Values{"task_id": {taskID}}
|
||||
if apiKey != "" {
|
||||
form.Set("api_key", apiKey)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, baseURL+"/parse/async_check", strings.NewReader(form.Encode()))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parser: SoMark poll request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
resp, err := models.NewDriverHTTPClient().Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parser: SoMark poll: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parser: SoMark read poll: %w", err)
|
||||
}
|
||||
var payload struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Data map[string]any `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &payload); err != nil {
|
||||
return nil, fmt.Errorf("parser: SoMark decode poll: %w", err)
|
||||
}
|
||||
if payload.Code != 0 {
|
||||
return nil, fmt.Errorf("parser: SoMark poll business error code=%d message=%s", payload.Code, payload.Message)
|
||||
}
|
||||
status, _ := payload.Data["status"].(string)
|
||||
if status != "SUCCESS" {
|
||||
return nil, fmt.Errorf("parser: SoMark task %s status=%s", taskID, status)
|
||||
}
|
||||
result, _ := payload.Data["result"].(map[string]any)
|
||||
if result == nil {
|
||||
return nil, fmt.Errorf("parser: SoMark SUCCESS without result")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func soMarkItems(result map[string]any, keepHeaderFooter bool) ([]map[string]any, int) {
|
||||
outputs, _ := result["outputs"].(map[string]any)
|
||||
jsonPayload, _ := outputs["json"].(map[string]any)
|
||||
pages, _ := jsonPayload["pages"].([]any)
|
||||
items := make([]map[string]any, 0)
|
||||
for _, pageRaw := range pages {
|
||||
page, ok := pageRaw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
blocks, _ := page["blocks"].([]any)
|
||||
for _, blockRaw := range blocks {
|
||||
block, ok := blockRaw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if item := soMarkBlockToItem(block, keepHeaderFooter); item != nil {
|
||||
items = append(items, item)
|
||||
}
|
||||
}
|
||||
}
|
||||
return items, len(pages)
|
||||
}
|
||||
|
||||
func soMarkBlockToItem(block map[string]any, keepHeaderFooter bool) map[string]any {
|
||||
blockType := strings.ToLower(strings.TrimSpace(stringValue(block["type"])))
|
||||
switch blockType {
|
||||
case "cate", "cate_item", "blank":
|
||||
return nil
|
||||
case "header", "footer":
|
||||
if !keepHeaderFooter {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
content := strings.TrimSpace(stringValue(block["content"]))
|
||||
switch blockType {
|
||||
case "figure", "cs", "qrcode", "stamp":
|
||||
if content == "" {
|
||||
content = "[Image]"
|
||||
}
|
||||
return map[string]any{"text": content, "doc_type_kwd": "image", "layout": "figure"}
|
||||
case "table":
|
||||
if content == "" {
|
||||
return nil
|
||||
}
|
||||
return map[string]any{"text": content, "doc_type_kwd": "table", "layout": "table"}
|
||||
case "equation":
|
||||
if content == "" {
|
||||
return nil
|
||||
}
|
||||
return map[string]any{"text": content, "doc_type_kwd": "text", "layout": "equation"}
|
||||
case "title":
|
||||
level := int(numberValue(block["title_level"]))
|
||||
if level < 1 || level > 6 {
|
||||
level = 1
|
||||
}
|
||||
if content == "" {
|
||||
return nil
|
||||
}
|
||||
return map[string]any{"text": strings.Repeat("#", level) + " " + content, "doc_type_kwd": "text", "layout": "title"}
|
||||
default:
|
||||
if content == "" {
|
||||
return nil
|
||||
}
|
||||
return map[string]any{"text": content, "doc_type_kwd": "text", "layout": "text"}
|
||||
}
|
||||
}
|
||||
|
||||
func envOrDefault(envKey, configured, fallback string) string {
|
||||
if configured != "" {
|
||||
return configured
|
||||
}
|
||||
if raw := strings.TrimSpace(os.Getenv(envKey)); raw != "" {
|
||||
return raw
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func envOrBool(envKey string, configured bool) bool {
|
||||
if raw := strings.TrimSpace(os.Getenv(envKey)); raw != "" {
|
||||
switch strings.ToLower(raw) {
|
||||
case "1", "true", "yes", "on":
|
||||
return true
|
||||
case "0", "false", "no", "off":
|
||||
return false
|
||||
}
|
||||
}
|
||||
return configured
|
||||
}
|
||||
|
||||
func urlEncoded(values map[string]string) string {
|
||||
parts := make([]string, 0, len(values))
|
||||
for k, v := range values {
|
||||
if strings.TrimSpace(v) == "" {
|
||||
if k == "api_key" {
|
||||
continue
|
||||
}
|
||||
}
|
||||
parts = append(parts, fmt.Sprintf("%s=%s", k, v))
|
||||
}
|
||||
return strings.Join(parts, "&")
|
||||
}
|
||||
140
internal/parser/parser/pdf_parser_somark_test.go
Normal file
140
internal/parser/parser/pdf_parser_somark_test.go
Normal file
@@ -0,0 +1,140 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPDFParser_ParseWithResult_SoMarkJSONIntegration(t *testing.T) {
|
||||
var submitSeen bool
|
||||
var pollSeen bool
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/parse/async":
|
||||
submitSeen = true
|
||||
reader, err := r.MultipartReader()
|
||||
if err != nil {
|
||||
t.Errorf("MultipartReader: %v", err)
|
||||
return
|
||||
}
|
||||
fields := map[string]string{}
|
||||
for {
|
||||
part, err := reader.NextPart()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
t.Errorf("NextPart: %v", err)
|
||||
return
|
||||
}
|
||||
body, _ := io.ReadAll(part)
|
||||
fields[part.FormName()] = string(body)
|
||||
}
|
||||
if fields["api_key"] != "somark-secret" {
|
||||
t.Errorf("api_key = %q, want somark-secret", fields["api_key"])
|
||||
return
|
||||
}
|
||||
if !strings.Contains(fields["element_formats"], "image") {
|
||||
t.Errorf("element_formats = %q", fields["element_formats"])
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"code":0,"data":{"task_id":"task-1"}}`))
|
||||
case "/parse/async_check":
|
||||
pollSeen = true
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
if !strings.Contains(string(body), "task_id=task-1") {
|
||||
t.Errorf("poll body = %q, want task_id", string(body))
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"code":0,"data":{"status":"SUCCESS","result":{"outputs":{"json":{"pages":[{"blocks":[{"type":"title","content":"SoMark Title","title_level":2},{"type":"figure","content":"Figure caption"},{"type":"table","content":"<table><tr><td>x</td></tr></table>"}]}]}}}}}`))
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "SoMark",
|
||||
"output_format": "json",
|
||||
"somark_base_url": server.URL,
|
||||
"somark_api_key": "somark-secret",
|
||||
})
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if !submitSeen || !pollSeen {
|
||||
t.Fatalf("submit/poll seen = %v/%v, want true/true", submitSeen, pollSeen)
|
||||
}
|
||||
if len(res.JSON) < 3 {
|
||||
t.Fatalf("JSON len = %d, want >=3", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_SoMarkMarkdownIntegration(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-2"}}`))
|
||||
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()
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "SoMark",
|
||||
"output_format": "markdown",
|
||||
"somark_base_url": server.URL,
|
||||
})
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if !strings.Contains(res.Markdown, "SoMark Title") {
|
||||
t.Fatalf("Markdown = %q, want title", res.Markdown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_SoMarkRequiresBaseURL(t *testing.T) {
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{"parse_method": "SoMark"})
|
||||
pdf.SoMarkBaseURL = ""
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err == nil || !strings.Contains(res.Err.Error(), "somark_base_url") {
|
||||
t.Fatalf("error = %v, want somark_base_url context", res.Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSoMarkBlockToItem_DropsHeaderByDefault(t *testing.T) {
|
||||
if item := soMarkBlockToItem(map[string]any{"type": "header", "content": "x"}, false); item != nil {
|
||||
t.Fatalf("item = %#v, want nil", item)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSoMarkSubmitMultipartShape(t *testing.T) {
|
||||
var form multipart.Form
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_ = r.ParseMultipartForm(1 << 20)
|
||||
form = *r.MultipartForm
|
||||
_, _ = w.Write([]byte(`{"code":0,"data":{"task_id":"task-3"}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
taskID, err := soMarkSubmit(server.URL, "sample.pdf", []byte("%PDF"), NewPDFParser(), "key")
|
||||
if err != nil {
|
||||
t.Fatalf("soMarkSubmit: %v", err)
|
||||
}
|
||||
if taskID != "task-3" {
|
||||
t.Fatalf("taskID = %q, want task-3", taskID)
|
||||
}
|
||||
if got := form.Value["api_key"][0]; got != "key" {
|
||||
t.Fatalf("api_key = %q, want key", got)
|
||||
}
|
||||
}
|
||||
205
internal/parser/parser/pdf_parser_tcadp.go
Normal file
205
internal/parser/parser/pdf_parser_tcadp.go
Normal file
@@ -0,0 +1,205 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
models "ragflow/internal/entity/models"
|
||||
)
|
||||
|
||||
func parsePDFWithTCADP(filename string, data []byte, parser *PDFParser) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
baseURL := strings.TrimSpace(parser.TCADPAPIServer)
|
||||
if baseURL == "" {
|
||||
baseURL = strings.TrimSpace(os.Getenv("TCADP_APISERVER"))
|
||||
}
|
||||
if baseURL == "" {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP requires tcadp_apiserver or TCADP_APISERVER")}
|
||||
}
|
||||
apiKey := strings.TrimSpace(parser.TCADPAPIKey)
|
||||
if apiKey == "" {
|
||||
apiKey = strings.TrimSpace(os.Getenv("TCADP_API_KEY"))
|
||||
}
|
||||
requestBody := map[string]any{
|
||||
"file_type": "PDF",
|
||||
"file_base64": base64.StdEncoding.EncodeToString(data),
|
||||
"file_start_page_number": 1,
|
||||
"file_end_page_number": 1000,
|
||||
"config": map[string]any{
|
||||
"TableResultType": parser.TCADPTableResultType,
|
||||
"MarkdownImageResponseType": parser.TCADPMarkdownImageResponseType,
|
||||
},
|
||||
}
|
||||
resp, err := models.PostJSONRequest(context.Background(), models.NewDriverHTTPClient(), strings.TrimRight(baseURL, "/")+"/reconstruct_document", bearer(apiKey), requestBody)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP submit: %w", err)}
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP read submit: %w", err)}
|
||||
}
|
||||
if resp.StatusCode >= 300 {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP HTTP %d: %s", resp.StatusCode, string(raw))}
|
||||
}
|
||||
var payload struct {
|
||||
DocumentRecognizeResultURL string `json:"DocumentRecognizeResultUrl"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &payload); err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP decode submit: %w", err)}
|
||||
}
|
||||
if payload.DocumentRecognizeResultURL == "" {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP returned no DocumentRecognizeResultUrl")}
|
||||
}
|
||||
downloadReq, err := http.NewRequestWithContext(context.Background(), http.MethodGet, payload.DocumentRecognizeResultURL, nil)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP download request: %w", err)}
|
||||
}
|
||||
if auth := bearer(apiKey); auth != "" {
|
||||
downloadReq.Header.Set("Authorization", auth)
|
||||
}
|
||||
downloadResp, err := models.NewDriverHTTPClient().Do(downloadReq)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP download: %w", err)}
|
||||
}
|
||||
defer downloadResp.Body.Close()
|
||||
zipBytes, err := io.ReadAll(downloadResp.Body)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP read zip: %w", err)}
|
||||
}
|
||||
items, pageCount, err := tcadpItemsFromZip(zipBytes)
|
||||
if err != nil {
|
||||
return ParseResult{Err: err}
|
||||
}
|
||||
return pdfItemsToResult(filename, items, parser.OutputFormat, pageCount)
|
||||
}
|
||||
|
||||
func tcadpItemsFromZip(zipBytes []byte) ([]map[string]any, int, error) {
|
||||
reader, err := zip.NewReader(bytes.NewReader(zipBytes), int64(len(zipBytes)))
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("parser: TCADP zip: %w", err)
|
||||
}
|
||||
items := make([]map[string]any, 0)
|
||||
pageCount := 0
|
||||
for _, file := range reader.File {
|
||||
if strings.HasSuffix(file.Name, ".md") {
|
||||
rc, err := file.Open()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
body, err := io.ReadAll(rc)
|
||||
rc.Close()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
items = append(items, map[string]any{"text": strings.TrimSpace(string(body)), "doc_type_kwd": "text", "layout": "text"})
|
||||
if strings.TrimSpace(string(body)) != "" && pageCount == 0 {
|
||||
pageCount = 1
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !strings.HasSuffix(file.Name, ".json") {
|
||||
continue
|
||||
}
|
||||
rc, err := file.Open()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
var raw any
|
||||
err = json.NewDecoder(rc).Decode(&raw)
|
||||
rc.Close()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
items = append(items, tcadpAnyToItems(raw)...)
|
||||
if pages := collectPDFPageNumbers(raw); len(pages) > pageCount {
|
||||
pageCount = len(pages)
|
||||
}
|
||||
}
|
||||
if len(items) == 0 {
|
||||
return nil, 0, fmt.Errorf("parser: TCADP zip contained no supported content")
|
||||
}
|
||||
return items, pageCount, nil
|
||||
}
|
||||
|
||||
func tcadpAnyToItems(raw any) []map[string]any {
|
||||
switch v := raw.(type) {
|
||||
case []any:
|
||||
items := make([]map[string]any, 0)
|
||||
for _, item := range v {
|
||||
items = append(items, tcadpAnyToItems(item)...)
|
||||
}
|
||||
return items
|
||||
case map[string]any:
|
||||
text := strings.TrimSpace(stringValue(v["content"]))
|
||||
contentType := strings.ToLower(strings.TrimSpace(stringValue(v["type"])))
|
||||
switch contentType {
|
||||
case "table":
|
||||
if text == "" {
|
||||
text = tcadpTableRowsText(v["table_data"])
|
||||
}
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
return []map[string]any{{"text": text, "doc_type_kwd": "table", "layout": "table"}}
|
||||
case "image":
|
||||
caption := strings.TrimSpace(stringValue(v["caption"]))
|
||||
if caption == "" {
|
||||
caption = "[Image]"
|
||||
}
|
||||
return []map[string]any{{"text": caption, "doc_type_kwd": "image", "layout": "figure"}}
|
||||
case "equation":
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
return []map[string]any{{"text": "$$" + text + "$$", "doc_type_kwd": "text", "layout": "equation"}}
|
||||
default:
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
return []map[string]any{{"text": text, "doc_type_kwd": "text", "layout": "text"}}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func tcadpTableRowsText(raw any) string {
|
||||
table, ok := raw.(map[string]any)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
rows, ok := table["rows"].([]any)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
lines := make([]string, 0, len(rows))
|
||||
for _, rowRaw := range rows {
|
||||
row, ok := rowRaw.([]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
cols := make([]string, 0, len(row))
|
||||
for _, col := range row {
|
||||
cols = append(cols, stringValue(col))
|
||||
}
|
||||
lines = append(lines, strings.Join(cols, " | "))
|
||||
}
|
||||
return strings.Join(lines, "\n")
|
||||
}
|
||||
|
||||
func bearer(apiKey string) string {
|
||||
if strings.TrimSpace(apiKey) == "" {
|
||||
return ""
|
||||
}
|
||||
return "Bearer " + apiKey
|
||||
}
|
||||
104
internal/parser/parser/pdf_parser_tcadp_test.go
Normal file
104
internal/parser/parser/pdf_parser_tcadp_test.go
Normal file
@@ -0,0 +1,104 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPDFParser_ParseWithResult_TCADPJSONIntegration(t *testing.T) {
|
||||
zipPayload := tcadpZipFixture(t)
|
||||
var server *httptest.Server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/reconstruct_document":
|
||||
if got, want := r.Header.Get("Authorization"), "Bearer tcadp-secret"; got != want {
|
||||
t.Errorf("Authorization = %q, want %q", got, want)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"DocumentRecognizeResultUrl":"` + server.URL + `/download.zip"}`))
|
||||
case "/download.zip":
|
||||
w.Header().Set("Content-Type", "application/zip")
|
||||
_, _ = w.Write(zipPayload)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "TCADP parser",
|
||||
"output_format": "json",
|
||||
"tcadp_apiserver": server.URL,
|
||||
"tcadp_api_key": "tcadp-secret",
|
||||
})
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) < 2 {
|
||||
t.Fatalf("JSON len = %d, want >=2", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_TCADPMarkdownIntegration(t *testing.T) {
|
||||
zipPayload := tcadpZipFixture(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()
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "TCADP parser",
|
||||
"output_format": "markdown",
|
||||
"tcadp_apiserver": server.URL,
|
||||
})
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if !strings.Contains(res.Markdown, "Hello TCADP") {
|
||||
t.Fatalf("Markdown = %q, want fixture text", res.Markdown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFParser_ParseWithResult_TCADPRequiresAPIServer(t *testing.T) {
|
||||
pdf := NewPDFParser()
|
||||
pdf.ConfigureFromSetup(map[string]any{"parse_method": "TCADP parser"})
|
||||
res := pdf.ParseWithResult("sample.pdf", []byte("%PDF-1.4\nmock"))
|
||||
if res.Err == nil || !strings.Contains(res.Err.Error(), "tcadp_apiserver") {
|
||||
t.Fatalf("error = %v, want tcadp_apiserver context", res.Err)
|
||||
}
|
||||
}
|
||||
|
||||
func tcadpZipFixture(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"))
|
||||
f2, err := writer.Create("blocks.json")
|
||||
if err != nil {
|
||||
t.Fatalf("Create json: %v", err)
|
||||
}
|
||||
_, _ = f2.Write([]byte(`[{"type":"table","table_data":{"rows":[["a","b"]]}}]`))
|
||||
if err := writer.Close(); err != nil {
|
||||
t.Fatalf("Close zip: %v", err)
|
||||
}
|
||||
return buf.Bytes()
|
||||
}
|
||||
207
internal/parser/parser/pdf_postprocess.go
Normal file
207
internal/parser/parser/pdf_postprocess.go
Normal file
@@ -0,0 +1,207 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"math"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
deepdoctype "ragflow/internal/deepdoc/parser/type"
|
||||
)
|
||||
|
||||
var pdfHeaderFooterPattern = regexp.MustCompile(`(?i)^(header|footer|number)$`)
|
||||
var pdfTOCTitlePattern = regexp.MustCompile(`(?i)^(contents|目录|目次|table of contents|致谢|acknowledge)$`)
|
||||
|
||||
type pdfPostProcessOptions struct {
|
||||
outputFormat string
|
||||
pageWidth float64
|
||||
zoom float64
|
||||
enableMultiColumn bool
|
||||
flattenMediaToText bool
|
||||
removeTOC bool
|
||||
removeHeaderFooter bool
|
||||
}
|
||||
|
||||
func applyPDFPostProcess(result *deepdoctype.ParseResult, opts pdfPostProcessOptions) {
|
||||
if result == nil {
|
||||
return
|
||||
}
|
||||
if opts.enableMultiColumn && opts.pageWidth > 0 {
|
||||
reorderPDFMultiColumn(result, opts.pageWidth, opts.zoom)
|
||||
}
|
||||
if opts.removeTOC {
|
||||
removePDFTOCByOutlines(result, result.Outlines)
|
||||
}
|
||||
normalizePDFLayoutTypes(result)
|
||||
if opts.removeHeaderFooter {
|
||||
filterPDFHeaderFooter(result)
|
||||
}
|
||||
assignPDFDocTypeKeywords(result, opts.flattenMediaToText)
|
||||
}
|
||||
|
||||
func normalizePDFLayoutTypes(result *deepdoctype.ParseResult) {
|
||||
for i := range result.Sections {
|
||||
layoutType := strings.TrimSpace(result.Sections[i].LayoutType)
|
||||
if layoutType == "" {
|
||||
layoutType = deepdoctype.LayoutTypeText
|
||||
}
|
||||
result.Sections[i].LayoutType = layoutType
|
||||
}
|
||||
}
|
||||
|
||||
func filterPDFHeaderFooter(result *deepdoctype.ParseResult) {
|
||||
filtered := result.Sections[:0]
|
||||
for _, s := range result.Sections {
|
||||
if pdfHeaderFooterPattern.MatchString(strings.TrimSpace(s.LayoutType)) {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, s)
|
||||
}
|
||||
result.Sections = filtered
|
||||
}
|
||||
|
||||
func assignPDFDocTypeKeywords(result *deepdoctype.ParseResult, flatten bool) {
|
||||
for i := range result.Sections {
|
||||
section := &result.Sections[i]
|
||||
if flatten {
|
||||
section.DocTypeKwd = "text"
|
||||
continue
|
||||
}
|
||||
switch strings.TrimSpace(section.LayoutType) {
|
||||
case deepdoctype.LayoutTypeTable:
|
||||
section.DocTypeKwd = "table"
|
||||
case deepdoctype.LayoutTypeFigure:
|
||||
section.DocTypeKwd = "image"
|
||||
default:
|
||||
if section.Image != "" {
|
||||
section.DocTypeKwd = "image"
|
||||
} else {
|
||||
section.DocTypeKwd = "text"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func removePDFTOCByOutlines(result *deepdoctype.ParseResult, outlines []deepdoctype.Outline) {
|
||||
if result == nil || len(outlines) == 0 {
|
||||
return
|
||||
}
|
||||
tocPage, contentPage := findPDFTOCPageRange(outlines)
|
||||
if contentPage <= tocPage {
|
||||
return
|
||||
}
|
||||
filtered := result.Sections[:0]
|
||||
for _, s := range result.Sections {
|
||||
page := firstSectionPage(s)
|
||||
if page >= tocPage && page < contentPage {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, s)
|
||||
}
|
||||
result.Sections = filtered
|
||||
}
|
||||
|
||||
func findPDFTOCPageRange(outlines []deepdoctype.Outline) (tocPage, contentPage int) {
|
||||
outer:
|
||||
for i, o := range outlines {
|
||||
title := strings.TrimSpace(o.Title)
|
||||
if idx := strings.Index(title, "@@"); idx >= 0 {
|
||||
title = strings.TrimSpace(title[:idx])
|
||||
}
|
||||
if !pdfTOCTitlePattern.MatchString(strings.ToLower(title)) {
|
||||
continue
|
||||
}
|
||||
tocPage = o.PageNumber
|
||||
for _, next := range outlines[i+1:] {
|
||||
if next.Level != o.Level {
|
||||
continue
|
||||
}
|
||||
nextTitle := strings.TrimSpace(next.Title)
|
||||
if idx := strings.Index(nextTitle, "@@"); idx >= 0 {
|
||||
nextTitle = strings.TrimSpace(nextTitle[:idx])
|
||||
}
|
||||
if pdfTOCTitlePattern.MatchString(strings.ToLower(nextTitle)) {
|
||||
continue
|
||||
}
|
||||
contentPage = next.PageNumber
|
||||
break outer
|
||||
}
|
||||
break
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func reorderPDFMultiColumn(result *deepdoctype.ParseResult, pageWidth, _ float64) {
|
||||
if result == nil || len(result.Sections) < 2 {
|
||||
return
|
||||
}
|
||||
|
||||
var widths []float64
|
||||
for _, s := range result.Sections {
|
||||
if strings.TrimSpace(s.LayoutType) != deepdoctype.LayoutTypeText || len(s.Positions) == 0 {
|
||||
continue
|
||||
}
|
||||
width := s.Positions[0].Right - s.Positions[0].Left
|
||||
if width > 0 {
|
||||
widths = append(widths, width)
|
||||
}
|
||||
}
|
||||
if len(widths) == 0 {
|
||||
return
|
||||
}
|
||||
sort.Float64s(widths)
|
||||
medianWidth := widths[len(widths)/2]
|
||||
if medianWidth >= pageWidth/2 {
|
||||
return
|
||||
}
|
||||
|
||||
sort.Slice(result.Sections, func(i, j int) bool {
|
||||
pi, pj := firstSectionPage(result.Sections[i]), firstSectionPage(result.Sections[j])
|
||||
if pi != pj {
|
||||
return pi < pj
|
||||
}
|
||||
xi, xj := firstSectionLeft(result.Sections[i]), firstSectionLeft(result.Sections[j])
|
||||
if math.Abs(xi-xj) > 1e-6 {
|
||||
return xi < xj
|
||||
}
|
||||
return firstSectionTop(result.Sections[i]) < firstSectionTop(result.Sections[j])
|
||||
})
|
||||
|
||||
threshold := medianWidth / 2
|
||||
for i := len(result.Sections) - 1; i >= 1; i-- {
|
||||
for j := i - 1; j >= 0; j-- {
|
||||
if firstSectionPage(result.Sections[j]) != firstSectionPage(result.Sections[j+1]) {
|
||||
continue
|
||||
}
|
||||
if math.Abs(firstSectionLeft(result.Sections[j])-firstSectionLeft(result.Sections[j+1])) >= threshold {
|
||||
continue
|
||||
}
|
||||
if firstSectionTop(result.Sections[j+1]) < firstSectionTop(result.Sections[j]) {
|
||||
result.Sections[j], result.Sections[j+1] = result.Sections[j+1], result.Sections[j]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func firstSectionPage(s deepdoctype.Section) int {
|
||||
for _, p := range s.Positions {
|
||||
for _, pn := range p.PageNumbers {
|
||||
return pn
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func firstSectionLeft(s deepdoctype.Section) float64 {
|
||||
for _, p := range s.Positions {
|
||||
return p.Left
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func firstSectionTop(s deepdoctype.Section) float64 {
|
||||
for _, p := range s.Positions {
|
||||
return p.Top
|
||||
}
|
||||
return 0
|
||||
}
|
||||
140
internal/parser/parser/pdf_postprocess_test.go
Normal file
140
internal/parser/parser/pdf_postprocess_test.go
Normal file
@@ -0,0 +1,140 @@
|
||||
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{})
|
||||
want := []string{"table", "image", "text", "image"}
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user