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:
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