mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-12 14:45:42 +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:
Zhichang Yu
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
|
||||
}
|
||||
Reference in New Issue
Block a user