WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,176 @@
|
||||
// Package anydoc converts office documents to Markdown inside the Go process,
|
||||
// without the Python docreader service.
|
||||
//
|
||||
// The conversion itself is the anydoc Rust library, linked as a static archive
|
||||
// through cgo (see third_party/anydoc-go). That archive needs a Rust toolchain
|
||||
// to produce, so it is opt-in: only builds tagged `anydoc` link it. Every other
|
||||
// build compiles the stub in backend_stub.go, where Available reports false and
|
||||
// the engine registry hides the engine with a reason the UI can show.
|
||||
//
|
||||
// The exported surface is deliberately narrow — bytes in, Markdown and embedded
|
||||
// images out — so the backend can be replaced (a WASI runtime, a future
|
||||
// upstream Go module) without touching the parser code that calls it.
|
||||
package anydoc
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ImageDir is the markdown path prefix for extracted images. The image
|
||||
// resolver matches references by this path and swaps them for storage URLs.
|
||||
const ImageDir = "images/"
|
||||
|
||||
// Result is one converted document.
|
||||
type Result struct {
|
||||
// Markdown is GitHub-Flavored Markdown for the whole document.
|
||||
// When WithAssets is set, embedded images appear in place as
|
||||
// `` via anydoc's official serializer.
|
||||
Markdown string
|
||||
// Assets are the images embedded in the document, in document order.
|
||||
// Always empty for PDF, which anydoc renders straight to Markdown
|
||||
// without a document model.
|
||||
Assets []Asset
|
||||
// AssetsError explains why Assets is empty when images were asked for.
|
||||
// The document-model parse is what both extracts images and places them
|
||||
// in the markdown; when it fails, Markdown still comes from the text-only
|
||||
// renderer so the conversion succeeds and the caller decides whether to
|
||||
// log or ignore the loss.
|
||||
AssetsError error
|
||||
}
|
||||
|
||||
// Asset is one image embedded in a document.
|
||||
type Asset struct {
|
||||
// ID is the document-model asset index, as referenced by in-place image
|
||||
// sources. It is stable for a single conversion.
|
||||
ID uint64
|
||||
// Name is a generated, extension-carrying file name ("image-1.png"):
|
||||
// embedded assets have no name of their own inside the container.
|
||||
Name string
|
||||
// MediaType is the IANA media type the container declared.
|
||||
MediaType string
|
||||
// Data is the raw image bytes.
|
||||
Data []byte
|
||||
// Alt is the image's alternative text as the document author wrote it,
|
||||
// empty when the document carries none.
|
||||
Alt string
|
||||
// Section is the text of the nearest heading above the image, empty when
|
||||
// the image sits before the first heading.
|
||||
Section string
|
||||
}
|
||||
|
||||
// Options tunes a single conversion.
|
||||
type Options struct {
|
||||
// Format names the parser explicitly ("docx", "csv", ...). Empty means
|
||||
// the format is detected from the content, which every format except
|
||||
// CSV carries a signature for.
|
||||
Format string
|
||||
// WithAssets extracts embedded images and renders them in place in the
|
||||
// markdown using anydoc's official GFM serializer (asset images are
|
||||
// rewritten to `images/image-N.ext` links first). Ignored for PDF.
|
||||
WithAssets bool
|
||||
}
|
||||
|
||||
// ErrUnavailable is returned by Convert when the binding is not linked into
|
||||
// this build. Callers that can fall back to another engine should check for it
|
||||
// with errors.Is.
|
||||
var ErrUnavailable = errors.New("anydoc: binding not built into this binary")
|
||||
|
||||
// supportedExtensions are the file types anydoc converts, mapped to the format
|
||||
// name its parser selector uses. The list is static so that engine metadata
|
||||
// (file types shown in the UI) is identical whether or not the binding is
|
||||
// linked in.
|
||||
var supportedExtensions = map[string]string{
|
||||
"doc": "doc",
|
||||
"docx": "docx",
|
||||
"docm": "docx",
|
||||
"odt": "odt",
|
||||
"rtf": "rtf",
|
||||
"ppt": "ppt",
|
||||
"pptx": "pptx",
|
||||
"pptm": "pptx",
|
||||
"odp": "odp",
|
||||
"xls": "xlsx",
|
||||
"xlsx": "xlsx",
|
||||
"xlsm": "xlsx",
|
||||
"ods": "ods",
|
||||
"epub": "epub",
|
||||
"csv": "csv",
|
||||
"pdf": "pdf",
|
||||
}
|
||||
|
||||
// SupportedFileTypes returns the extensions anydoc can convert, sorted so the
|
||||
// engine list is stable across restarts.
|
||||
func SupportedFileTypes() []string {
|
||||
types := make([]string, 0, len(supportedExtensions))
|
||||
for ext := range supportedExtensions {
|
||||
types = append(types, ext)
|
||||
}
|
||||
slices.Sort(types)
|
||||
return types
|
||||
}
|
||||
|
||||
// FormatForFile resolves the anydoc format name for a file type or file name.
|
||||
// ok is false for anything anydoc does not convert.
|
||||
func FormatForFile(fileType, fileName string) (format string, ok bool) {
|
||||
ext := normalizeExt(fileType)
|
||||
if ext == "" {
|
||||
ext = normalizeExt(filepath.Ext(fileName))
|
||||
}
|
||||
format, ok = supportedExtensions[ext]
|
||||
return format, ok
|
||||
}
|
||||
|
||||
// Supports reports whether anydoc converts this file type, regardless of
|
||||
// whether the binding is linked into this build.
|
||||
func Supports(fileType, fileName string) bool {
|
||||
_, ok := FormatForFile(fileType, fileName)
|
||||
return ok
|
||||
}
|
||||
|
||||
// Available reports whether conversions can actually run in this binary.
|
||||
func Available() bool { return backendAvailable() }
|
||||
|
||||
// UnavailableReason explains, for the UI, why Available is false. It returns
|
||||
// "" when the binding is available.
|
||||
func UnavailableReason() string { return backendUnavailableReason() }
|
||||
|
||||
// Version reports the anydoc release the linked binding was built from, or ""
|
||||
// when no binding is linked.
|
||||
func Version() string { return backendVersion() }
|
||||
|
||||
// Convert turns document bytes into Markdown. It returns ErrUnavailable when
|
||||
// the binding is not linked into this build.
|
||||
func Convert(data []byte, opts Options) (*Result, error) {
|
||||
if len(data) == 0 {
|
||||
return nil, fmt.Errorf("anydoc: empty document")
|
||||
}
|
||||
if !backendAvailable() {
|
||||
return nil, fmt.Errorf("%w: %s", ErrUnavailable, backendUnavailableReason())
|
||||
}
|
||||
// PDF has no document model, so asset extraction would only produce an
|
||||
// error the caller has to special-case. Drop the request instead.
|
||||
if opts.Format == "pdf" {
|
||||
opts.WithAssets = false
|
||||
}
|
||||
return backendConvert(data, opts)
|
||||
}
|
||||
|
||||
// PDFNeedsOCR reports whether a Convert error means the PDF has no usable
|
||||
// text layer. Scanned exam papers fail this way; the caller should fall back
|
||||
// to an engine that can rasterize pages for OCR.
|
||||
func PDFNeedsOCR(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
msg := strings.ToLower(err.Error())
|
||||
return strings.Contains(msg, "ocr is required") ||
|
||||
strings.Contains(msg, "no extractable text")
|
||||
}
|
||||
|
||||
func normalizeExt(s string) string {
|
||||
return strings.ToLower(strings.TrimPrefix(strings.TrimSpace(s), "."))
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package anydoc
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFormatForFileResolvesTypeAndName(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
fileType string
|
||||
fileName string
|
||||
wantFormat string
|
||||
wantOK bool
|
||||
}{
|
||||
{name: "file type", fileType: "docx", wantFormat: "docx", wantOK: true},
|
||||
{name: "dotted file type", fileType: ".DOCX", wantFormat: "docx", wantOK: true},
|
||||
{name: "macro format maps to base format", fileType: "pptm", wantFormat: "pptx", wantOK: true},
|
||||
{name: "legacy excel maps to xlsx parser", fileType: "xls", wantFormat: "xlsx", wantOK: true},
|
||||
{name: "falls back to file name", fileName: "report.odt", wantFormat: "odt", wantOK: true},
|
||||
{name: "unsupported type", fileType: "png", wantOK: false},
|
||||
{name: "nothing to go on", wantOK: false},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
format, ok := FormatForFile(tc.fileType, tc.fileName)
|
||||
if ok != tc.wantOK {
|
||||
t.Fatalf("FormatForFile(%q, %q) ok = %v, want %v", tc.fileType, tc.fileName, ok, tc.wantOK)
|
||||
}
|
||||
if format != tc.wantFormat {
|
||||
t.Fatalf("FormatForFile(%q, %q) = %q, want %q", tc.fileType, tc.fileName, format, tc.wantFormat)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The engine list must look the same whether or not the converter is linked
|
||||
// in, so that the UI can show the file types with a "not built in" reason.
|
||||
func TestSupportedFileTypesIsStableAndSorted(t *testing.T) {
|
||||
types := SupportedFileTypes()
|
||||
if !slices.IsSorted(types) {
|
||||
t.Fatalf("SupportedFileTypes is not sorted: %v", types)
|
||||
}
|
||||
for _, want := range []string{"docx", "pptx", "xlsx", "epub", "pdf", "csv"} {
|
||||
if !slices.Contains(types, want) {
|
||||
t.Errorf("SupportedFileTypes does not include %q: %v", want, types)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFNeedsOCR(t *testing.T) {
|
||||
if PDFNeedsOCR(nil) {
|
||||
t.Fatal("nil error was treated as needing OCR")
|
||||
}
|
||||
if !PDFNeedsOCR(errors.New("PDF has no extractable text (Scanned, 5 pages): OCR is required")) {
|
||||
t.Error("scanned-PDF error was not detected")
|
||||
}
|
||||
if PDFNeedsOCR(errors.New("not a readable zip archive")) {
|
||||
t.Error("unrelated error was treated as needing OCR")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertRejectsEmptyInput(t *testing.T) {
|
||||
if _, err := Convert(nil, Options{Format: "docx"}); err == nil {
|
||||
t.Fatal("Convert(nil) succeeded, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
// Without the build tag there is no converter, and callers must be able to
|
||||
// detect that so they can fall back to another engine.
|
||||
func TestConvertReportsUnavailableBackend(t *testing.T) {
|
||||
if Available() {
|
||||
t.Skip("converter is linked into this build")
|
||||
}
|
||||
_, err := Convert([]byte("id,name\n1,a\n"), Options{Format: "csv"})
|
||||
if !errors.Is(err, ErrUnavailable) {
|
||||
t.Fatalf("Convert error = %v, want ErrUnavailable", err)
|
||||
}
|
||||
if UnavailableReason() == "" {
|
||||
t.Error("UnavailableReason is empty while the converter is unavailable")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
//go:build anydoc && cgo
|
||||
|
||||
package anydoc
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"mime"
|
||||
"strings"
|
||||
|
||||
upstream "github.com/firecrawl/anydoc/go"
|
||||
)
|
||||
|
||||
// This build links the anydoc Rust archive through the upstream cgo bindings
|
||||
// (vendored under third_party/anydoc-go until they are published as a module).
|
||||
// Conversion runs in-process: no subprocess, no HTTP hop, no Python.
|
||||
|
||||
func backendAvailable() bool { return true }
|
||||
|
||||
func backendUnavailableReason() string { return "" }
|
||||
|
||||
func backendVersion() string { return upstream.Version }
|
||||
|
||||
func backendConvert(data []byte, opts Options) (*Result, error) {
|
||||
format, err := upstreamFormat(opts.Format)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if !opts.WithAssets {
|
||||
markdown, err := upstream.ToMarkdownBytes(data, format)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("anydoc: markdown conversion failed: %w", err)
|
||||
}
|
||||
return &Result{Markdown: markdown}, nil
|
||||
}
|
||||
|
||||
// Official GFM serializer after rewriting Asset images to External URLs
|
||||
// (`images/image-N.ext`). ToDocument still supplies the image bytes.
|
||||
document, err := upstream.ToDocument(data, format)
|
||||
if err != nil {
|
||||
markdown, mdErr := upstream.ToMarkdownBytes(data, format)
|
||||
if mdErr != nil {
|
||||
return nil, fmt.Errorf("anydoc: markdown conversion failed: %w", mdErr)
|
||||
}
|
||||
return &Result{
|
||||
Markdown: markdown,
|
||||
AssetsError: fmt.Errorf("anydoc: image extraction failed: %w", err),
|
||||
}, nil
|
||||
}
|
||||
markdown, err := upstream.ToMarkdownWithAssetLinks(data, format)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("anydoc: markdown conversion failed: %w", err)
|
||||
}
|
||||
return &Result{Markdown: markdown, Assets: collectAssets(document)}, nil
|
||||
}
|
||||
|
||||
// upstreamFormat maps our format name onto the binding's constant. An empty
|
||||
// name means "detect from content".
|
||||
func upstreamFormat(name string) (*upstream.Format, error) {
|
||||
if name == "" {
|
||||
return nil, nil
|
||||
}
|
||||
format, ok := upstream.FormatFromExtension(name)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("anydoc: unsupported format %q", name)
|
||||
}
|
||||
return &format, nil
|
||||
}
|
||||
|
||||
// collectAssets names each embedded image after its position in the document,
|
||||
// because containers store assets without a usable file name of their own, and
|
||||
// attaches the alt text and section the image appeared under.
|
||||
func collectAssets(document *upstream.Document) []Asset {
|
||||
if document == nil || len(document.Assets) == 0 {
|
||||
return nil
|
||||
}
|
||||
placements := placeAssets(document.Blocks)
|
||||
assets := make([]Asset, 0, len(document.Assets))
|
||||
for i, asset := range document.Assets {
|
||||
if len(asset.Data) == 0 {
|
||||
continue
|
||||
}
|
||||
placement := placements[uint64(i)]
|
||||
assets = append(assets, Asset{
|
||||
ID: uint64(i),
|
||||
Name: fmt.Sprintf("image-%d%s", i+1, extensionFor(asset.MediaType)),
|
||||
MediaType: asset.MediaType,
|
||||
Data: asset.Data,
|
||||
Alt: placement.alt,
|
||||
Section: placement.section,
|
||||
})
|
||||
}
|
||||
return assets
|
||||
}
|
||||
|
||||
// placement is where in the document an asset was referenced from.
|
||||
type placement struct {
|
||||
alt string
|
||||
section string
|
||||
}
|
||||
|
||||
// placeAssets walks the document body and records, per asset, the alt text of
|
||||
// the image that referenced it and the heading it sat under. The walk
|
||||
// recurses into lists, tables, and quotes, since an image can be nested
|
||||
// anywhere a paragraph can.
|
||||
func placeAssets(blocks []upstream.Block) map[uint64]placement {
|
||||
placements := make(map[uint64]placement)
|
||||
section := ""
|
||||
walkBlocks(blocks, §ion, placements)
|
||||
return placements
|
||||
}
|
||||
|
||||
func walkBlocks(blocks []upstream.Block, section *string, placements map[uint64]placement) {
|
||||
for _, block := range blocks {
|
||||
switch block.Kind {
|
||||
case "heading":
|
||||
*section = inlineText(block.Content)
|
||||
walkInlines(block.Content, *section, placements)
|
||||
case "paragraph":
|
||||
walkInlines(block.Content, *section, placements)
|
||||
case "block_quote":
|
||||
walkBlocks(block.Blocks, section, placements)
|
||||
case "list":
|
||||
if block.List != nil {
|
||||
for _, item := range block.List.Items {
|
||||
walkBlocks(item.Blocks, section, placements)
|
||||
}
|
||||
}
|
||||
case "table":
|
||||
if block.Table != nil {
|
||||
for _, row := range block.Table.Grid {
|
||||
for _, slot := range row {
|
||||
if slot.Cell != nil {
|
||||
walkBlocks(slot.Cell.Blocks, section, placements)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func walkInlines(inlines []upstream.Inline, section string, placements map[uint64]placement) {
|
||||
for _, inline := range inlines {
|
||||
switch inline.Kind {
|
||||
case "image":
|
||||
if inline.Source == nil || inline.Source.Kind != "asset" || inline.Source.AssetID == nil {
|
||||
continue
|
||||
}
|
||||
alt := ""
|
||||
if inline.Alt != nil {
|
||||
alt = strings.TrimSpace(*inline.Alt)
|
||||
}
|
||||
// First reference wins: the same asset can be placed more than
|
||||
// once, and the earliest placement is the one a reader would
|
||||
// associate it with.
|
||||
if _, seen := placements[*inline.Source.AssetID]; !seen {
|
||||
placements[*inline.Source.AssetID] = placement{alt: alt, section: section}
|
||||
}
|
||||
case "link":
|
||||
walkInlines(inline.Content, section, placements)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// inlineText flattens inline content to plain text, which is all a heading is
|
||||
// needed for here.
|
||||
func inlineText(inlines []upstream.Inline) string {
|
||||
var text strings.Builder
|
||||
for _, inline := range inlines {
|
||||
switch inline.Kind {
|
||||
case "text":
|
||||
if inline.Text != nil {
|
||||
text.WriteString(*inline.Text)
|
||||
}
|
||||
case "link":
|
||||
text.WriteString(inlineText(inline.Content))
|
||||
}
|
||||
}
|
||||
return strings.TrimSpace(text.String())
|
||||
}
|
||||
|
||||
func extensionFor(mediaType string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(mediaType)) {
|
||||
case "image/jpeg", "image/jpg":
|
||||
return ".jpg"
|
||||
case "image/png":
|
||||
return ".png"
|
||||
case "image/gif":
|
||||
return ".gif"
|
||||
case "image/webp":
|
||||
return ".webp"
|
||||
case "image/bmp":
|
||||
return ".bmp"
|
||||
case "image/tiff":
|
||||
return ".tiff"
|
||||
case "image/svg+xml":
|
||||
return ".svg"
|
||||
}
|
||||
// Anything else (EMF/WMF drawings, unknown object payloads) keeps
|
||||
// whatever extension the media type registry knows, and .bin otherwise,
|
||||
// so downstream storage never writes an extension-less blob.
|
||||
if exts, err := mime.ExtensionsByType(mediaType); err == nil && len(exts) > 0 {
|
||||
return exts[0]
|
||||
}
|
||||
return ".bin"
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
//go:build !anydoc || !cgo
|
||||
|
||||
package anydoc
|
||||
|
||||
// The default build links no converter: the anydoc archive is a Rust build
|
||||
// artifact, and requiring a Rust toolchain to compile WeKnora would be a steep
|
||||
// price for one optional engine. Everything still compiles and the engine
|
||||
// registry simply reports the engine as unavailable.
|
||||
|
||||
const unavailableReason = "anydoc engine not built into this binary " +
|
||||
"(rebuild with `make build-anydoc` / `-tags anydoc`; " +
|
||||
"Docker images need `--build-arg WITH_ANYDOC=1`)"
|
||||
|
||||
func backendAvailable() bool { return false }
|
||||
|
||||
func backendUnavailableReason() string { return unavailableReason }
|
||||
|
||||
func backendVersion() string { return "" }
|
||||
|
||||
func backendConvert(_ []byte, _ Options) (*Result, error) {
|
||||
// Convert checks Available first, so this is only reachable if a caller
|
||||
// bypasses it.
|
||||
return nil, ErrUnavailable
|
||||
}
|
||||
@@ -0,0 +1,379 @@
|
||||
//go:build anydoc && cgo
|
||||
|
||||
package anydoc
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// These tests exercise the linked Rust converter, so they only build with the
|
||||
// `anydoc` tag. Run them with:
|
||||
//
|
||||
// scripts/build-anydoc-lib.sh && go test -tags anydoc ./internal/infrastructure/docparser/...
|
||||
|
||||
// SupportedFileTypes is a static list, so that the engine can advertise its
|
||||
// file types in builds that link no converter. This is the check that it still
|
||||
// agrees with what the converter actually accepts, in either direction.
|
||||
func TestSupportedFileTypesMatchTheConverter(t *testing.T) {
|
||||
for _, fileType := range SupportedFileTypes() {
|
||||
format, ok := FormatForFile(fileType, "")
|
||||
if !ok {
|
||||
t.Errorf("%q is advertised but maps to no format", fileType)
|
||||
continue
|
||||
}
|
||||
if _, err := upstreamFormat(format); err != nil {
|
||||
t.Errorf("%q maps to format %q, which the converter rejects: %v", fileType, format, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertCSV(t *testing.T) {
|
||||
result, err := Convert([]byte("quarter,widgets\nQ1,12\nQ2,15\n"), Options{Format: "csv"})
|
||||
if err != nil {
|
||||
t.Fatalf("Convert: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "| quarter | widgets |") {
|
||||
t.Fatalf("expected a markdown table header, got:\n%s", result.Markdown)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "| Q1 | 12 |") {
|
||||
t.Fatalf("expected the first data row, got:\n%s", result.Markdown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertDocxWithEmbeddedImage(t *testing.T) {
|
||||
document := buildDocx(t)
|
||||
|
||||
result, err := Convert(document, Options{Format: "docx", WithAssets: true})
|
||||
if err != nil {
|
||||
t.Fatalf("Convert: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "# Quarterly report") {
|
||||
t.Fatalf("expected the heading, got:\n%s", result.Markdown)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "Widgets shipped on time.") {
|
||||
t.Fatalf("expected the body paragraph, got:\n%s", result.Markdown)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "Closing remarks.") {
|
||||
t.Fatalf("expected the trailing paragraph, got:\n%s", result.Markdown)
|
||||
}
|
||||
if len(result.Assets) != 1 {
|
||||
t.Fatalf("got %d assets, want 1", len(result.Assets))
|
||||
}
|
||||
asset := result.Assets[0]
|
||||
if asset.Name != "image-1.png" {
|
||||
t.Errorf("asset name = %q, want image-1.png", asset.Name)
|
||||
}
|
||||
if !bytes.Equal(asset.Data, onePixelPNG()) {
|
||||
t.Errorf("asset data does not round-trip the embedded image")
|
||||
}
|
||||
if asset.Alt != "Shipping chart" {
|
||||
t.Errorf("asset alt = %q, want %q", asset.Alt, "Shipping chart")
|
||||
}
|
||||
if asset.Section != "Quarterly report" {
|
||||
t.Errorf("asset section = %q, want %q", asset.Section, "Quarterly report")
|
||||
}
|
||||
imageLink := ""
|
||||
if !strings.Contains(result.Markdown, imageLink) {
|
||||
t.Fatalf("expected the image in place, got:\n%s", result.Markdown)
|
||||
}
|
||||
before := strings.Index(result.Markdown, "Widgets shipped on time.")
|
||||
at := strings.Index(result.Markdown, imageLink)
|
||||
after := strings.Index(result.Markdown, "Closing remarks.")
|
||||
if before < 0 || at < 0 || after < 0 || !(before < at && at < after) {
|
||||
t.Fatalf("image is not between the surrounding paragraphs:\n%s", result.Markdown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertTextlessPDFNeedsOCR(t *testing.T) {
|
||||
_, err := Convert(textlessPDF(), Options{Format: "pdf"})
|
||||
if err == nil {
|
||||
t.Fatal("Convert succeeded on a PDF with no text, want an error")
|
||||
}
|
||||
if !PDFNeedsOCR(err) {
|
||||
t.Fatalf("Convert error = %v, want an OCR-required error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The ABI reports error messages through a thread-local slot that a second
|
||||
// call reads, so a goroutine resumed on another OS thread would report an
|
||||
// empty or someone else's message. The binding pins the goroutine for the
|
||||
// pair; this is the regression test for that.
|
||||
func TestErrorDetailSurvivesConcurrency(t *testing.T) {
|
||||
const conversions = 4000
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var empty, crossed atomic.Int64
|
||||
for i := 0; i < conversions; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
// docx and xlsx failures carry distinctive messages, so a detail
|
||||
// naming the other format proves it came from the wrong thread.
|
||||
format, want, other := "docx", "not a readable zip archive", "unreadable workbook"
|
||||
if i%2 == 1 {
|
||||
format, want, other = "xlsx", "unreadable workbook", "not a readable zip archive"
|
||||
}
|
||||
_, err := Convert([]byte(fmt.Sprintf("garbage %d", i)), Options{Format: format})
|
||||
if err == nil {
|
||||
t.Errorf("Convert(garbage) succeeded")
|
||||
return
|
||||
}
|
||||
detail := err.Error()
|
||||
switch {
|
||||
case strings.Contains(detail, other):
|
||||
crossed.Add(1)
|
||||
case !strings.Contains(detail, want):
|
||||
empty.Add(1)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if empty.Load() > 0 || crossed.Load() > 0 {
|
||||
t.Fatalf("error details lost or crossed between goroutines: lost=%d crossed=%d (of %d)",
|
||||
empty.Load(), crossed.Load(), conversions)
|
||||
}
|
||||
}
|
||||
|
||||
// Detection reads the container itself, so a document whose format is not
|
||||
// named still converts.
|
||||
func TestConvertDetectsFormatFromContent(t *testing.T) {
|
||||
result, err := Convert(buildDocx(t), Options{})
|
||||
if err != nil {
|
||||
t.Fatalf("Convert: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "# Quarterly report") {
|
||||
t.Fatalf("expected the heading, got:\n%s", result.Markdown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConvertRejectsGarbage(t *testing.T) {
|
||||
if _, err := Convert([]byte("not a document at all"), Options{Format: "docx"}); err == nil {
|
||||
t.Fatal("Convert succeeded on garbage input, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
// PDF has no document model, so WithAssets has to be dropped rather than turn
|
||||
// a working conversion into an error.
|
||||
func TestConvertPDFIgnoresAssetRequest(t *testing.T) {
|
||||
result, err := Convert(minimalPDF(), Options{Format: "pdf", WithAssets: true})
|
||||
if err != nil {
|
||||
t.Fatalf("Convert: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.Markdown, "Shipping summary") {
|
||||
t.Fatalf("expected the page text, got:\n%s", result.Markdown)
|
||||
}
|
||||
if len(result.Assets) != 0 {
|
||||
t.Errorf("got %d assets for a PDF, want 0", len(result.Assets))
|
||||
}
|
||||
}
|
||||
|
||||
// A PDF whose catalog holds a deeply nested array is the shape of
|
||||
// RUSTSEC-2026-0187: lopdf 0.41 recursed per nesting level and killed the
|
||||
// process with an uncatchable stack overflow, which no Go-side or Rust-side
|
||||
// recover can contain. The pinned dependency in third_party/anydoc-go bounds
|
||||
// the recursion, so this must come back as an ordinary error — and if a
|
||||
// dependency bump ever reintroduces it, this test crashes rather than passing
|
||||
// quietly.
|
||||
func TestDeeplyNestedPDFFailsWithoutKillingTheProcess(t *testing.T) {
|
||||
const depth = 50000
|
||||
nested := strings.Repeat("[", depth) + strings.Repeat("]", depth)
|
||||
|
||||
_, err := Convert(nestedPDF(nested), Options{Format: "pdf"})
|
||||
if err == nil {
|
||||
t.Fatal("Convert succeeded on a PDF with no text, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
// A show-text operator with no operand made the page-classification scan walk
|
||||
// back over the whole content stream looking for one, so a stream of bare
|
||||
// `] TJ` tokens cost time quadratic in its length: before pdf-inspector 1.14.2
|
||||
// this input took 26 seconds of CPU, and a 2 MB one took nearly two minutes.
|
||||
// Unbounded work is the half of the hostile-PDF problem that guarded() cannot
|
||||
// catch — it never panics, it just holds the core — so the bound is pinned
|
||||
// here. The budget is deliberately far above the ~7ms a bounded lookback needs:
|
||||
// what it has to distinguish is linear from quadratic, not fast from slow.
|
||||
func TestDetectorLookbackStaysLinear(t *testing.T) {
|
||||
const (
|
||||
operators = 200000
|
||||
budget = 15 * time.Second
|
||||
)
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
// The result is irrelevant: the page carries one real text run, so
|
||||
// this converts either way. Only how long it takes is under test.
|
||||
_, _ = Convert(unmatchedShowTextPDF(operators), Options{Format: "pdf"})
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(budget):
|
||||
t.Fatalf("converting a PDF with %d unmatched TJ operators took over %s; "+
|
||||
"the detector's operand lookback is no longer bounded", operators, budget)
|
||||
}
|
||||
}
|
||||
|
||||
// buildDocx writes the smallest OOXML package that carries a heading, a
|
||||
// paragraph, and one embedded image.
|
||||
func buildDocx(t *testing.T) []byte {
|
||||
t.Helper()
|
||||
|
||||
parts := map[string]string{
|
||||
"[Content_Types].xml": `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<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"/>
|
||||
<Default Extension="png" ContentType="image/png"/>
|
||||
<Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
|
||||
<Override PartName="/word/styles.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.styles+xml"/>
|
||||
</Types>`,
|
||||
"_rels/.rels": `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<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>`,
|
||||
"word/styles.xml": `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<w:styles xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
|
||||
<w:style w:type="paragraph" w:styleId="Heading1">
|
||||
<w:name w:val="heading 1"/>
|
||||
<w:pPr><w:outlineLvl w:val="0"/></w:pPr>
|
||||
</w:style>
|
||||
</w:styles>`,
|
||||
"word/_rels/document.xml.rels": `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
|
||||
<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/styles" Target="styles.xml"/>
|
||||
<Relationship Id="rId10" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/image" Target="media/image1.png"/>
|
||||
</Relationships>`,
|
||||
"word/document.xml": `<?xml version="1.0" encoding="UTF-8"?>
|
||||
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"
|
||||
xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships"
|
||||
xmlns:wp="http://schemas.openxmlformats.org/drawingml/2006/wordprocessingDrawing"
|
||||
xmlns:a="http://schemas.openxmlformats.org/drawingml/2006/main"
|
||||
xmlns:pic="http://schemas.openxmlformats.org/drawingml/2006/picture">
|
||||
<w:body>
|
||||
<w:p><w:pPr><w:pStyle w:val="Heading1"/></w:pPr><w:r><w:t>Quarterly report</w:t></w:r></w:p>
|
||||
<w:p><w:r><w:t>Widgets shipped on time.</w:t></w:r></w:p>
|
||||
<w:p><w:r><w:drawing><wp:inline><wp:docPr id="1" name="Chart" descr="Shipping chart"/>
|
||||
<a:graphic><a:graphicData uri="http://schemas.openxmlformats.org/drawingml/2006/picture">
|
||||
<pic:pic><pic:nvPicPr><pic:cNvPr id="1" name="Chart"/><pic:cNvPicPr/></pic:nvPicPr>
|
||||
<pic:blipFill><a:blip r:embed="rId10"/></pic:blipFill>
|
||||
<pic:spPr/></pic:pic>
|
||||
</a:graphicData></a:graphic>
|
||||
</wp:inline></w:drawing></w:r></w:p>
|
||||
<w:p><w:r><w:t>Closing remarks.</w:t></w:r></w:p>
|
||||
</w:body>
|
||||
</w:document>`,
|
||||
}
|
||||
|
||||
var buf bytes.Buffer
|
||||
archive := zip.NewWriter(&buf)
|
||||
for name, content := range parts {
|
||||
writer, err := archive.Create(name)
|
||||
if err != nil {
|
||||
t.Fatalf("create %s: %v", name, err)
|
||||
}
|
||||
if _, err := writer.Write([]byte(content)); err != nil {
|
||||
t.Fatalf("write %s: %v", name, err)
|
||||
}
|
||||
}
|
||||
image, err := archive.Create("word/media/image1.png")
|
||||
if err != nil {
|
||||
t.Fatalf("create image part: %v", err)
|
||||
}
|
||||
if _, err := image.Write(onePixelPNG()); err != nil {
|
||||
t.Fatalf("write image part: %v", err)
|
||||
}
|
||||
if err := archive.Close(); err != nil {
|
||||
t.Fatalf("close archive: %v", err)
|
||||
}
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
// onePixelPNG is a 1x1 transparent PNG: the smallest thing a container will
|
||||
// accept as an image part.
|
||||
func onePixelPNG() []byte {
|
||||
return []byte{
|
||||
0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a,
|
||||
0x00, 0x00, 0x00, 0x0d, 'I', 'H', 'D', 'R',
|
||||
0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,
|
||||
0x08, 0x06, 0x00, 0x00, 0x00, 0x1f, 0x15, 0xc4, 0x89,
|
||||
0x00, 0x00, 0x00, 0x0a, 'I', 'D', 'A', 'T',
|
||||
0x78, 0x9c, 0x63, 0x00, 0x01, 0x00, 0x00, 0x05, 0x00, 0x01,
|
||||
0x0d, 0x0a, 0x2d, 0xb4,
|
||||
0x00, 0x00, 0x00, 0x00, 'I', 'E', 'N', 'D', 0xae, 0x42, 0x60, 0x82,
|
||||
}
|
||||
}
|
||||
|
||||
// minimalPDF is a one-page PDF with a single text run, written by hand so the
|
||||
// test carries no binary fixture.
|
||||
func minimalPDF() []byte {
|
||||
content := "BT /F1 12 Tf 20 100 Td (Shipping summary) Tj ET\n"
|
||||
return writePDF([]string{
|
||||
"<< /Type /Catalog /Pages 2 0 R >>",
|
||||
"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
|
||||
"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 200 200] /Contents 4 0 R " +
|
||||
"/Resources << /Font << /F1 5 0 R >> >> >>",
|
||||
fmt.Sprintf("<< /Length %d >>\nstream\n%sendstream", len(content), content),
|
||||
"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
|
||||
})
|
||||
}
|
||||
|
||||
// nestedPDF is a one-page PDF carrying the given nested value in its catalog.
|
||||
func nestedPDF(nested string) []byte {
|
||||
return writePDF([]string{
|
||||
"<< /Type /Catalog /Pages 2 0 R /Nested " + nested + " >>",
|
||||
"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
|
||||
"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 200 200] >>",
|
||||
})
|
||||
}
|
||||
|
||||
// unmatchedShowTextPDF is a one-page PDF whose content stream shows text once
|
||||
// and then repeats a TJ operator whose array operand is never opened.
|
||||
func unmatchedShowTextPDF(operators int) []byte {
|
||||
content := "BT /F1 12 Tf 20 100 Td (Shipping summary) Tj ET\n" +
|
||||
strings.Repeat("] TJ\n", operators)
|
||||
return writePDF([]string{
|
||||
"<< /Type /Catalog /Pages 2 0 R >>",
|
||||
"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
|
||||
"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 200 200] /Contents 4 0 R " +
|
||||
"/Resources << /Font << /F1 5 0 R >> >> >>",
|
||||
fmt.Sprintf("<< /Length %d >>\nstream\n%sendstream", len(content), content),
|
||||
"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
|
||||
})
|
||||
}
|
||||
|
||||
// textlessPDF is a one-page PDF with no text stream: the shape of a scanned
|
||||
// exam paper as far as pdf-inspector is concerned.
|
||||
func textlessPDF() []byte {
|
||||
return writePDF([]string{
|
||||
"<< /Type /Catalog /Pages 2 0 R >>",
|
||||
"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
|
||||
"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 200 200] >>",
|
||||
})
|
||||
}
|
||||
|
||||
func writePDF(objects []string) []byte {
|
||||
var pdf bytes.Buffer
|
||||
offsets := make([]int, len(objects))
|
||||
pdf.WriteString("%PDF-1.4\n")
|
||||
for i, object := range objects {
|
||||
offsets[i] = pdf.Len()
|
||||
fmt.Fprintf(&pdf, "%d 0 obj\n%s\nendobj\n", i+1, object)
|
||||
}
|
||||
|
||||
xref := pdf.Len()
|
||||
fmt.Fprintf(&pdf, "xref\n0 %d\n0000000000 65535 f \n", len(objects)+1)
|
||||
for _, offset := range offsets {
|
||||
fmt.Fprintf(&pdf, "%010d 00000 n \n", offset)
|
||||
}
|
||||
fmt.Fprintf(&pdf, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(objects)+1, xref)
|
||||
return pdf.Bytes()
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser/anydoc"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// AnydocReader converts office documents to markdown in this process, using
|
||||
// the anydoc converter instead of the Python docreader service.
|
||||
//
|
||||
// Embedded images are rendered in place from the document model. PDFs with no
|
||||
// text layer (scanned exam papers) cannot be converted by anydoc; when a
|
||||
// fallback reader is configured (typically the builtin docreader) those files
|
||||
// are handed off so pages can be rasterized for OCR.
|
||||
type AnydocReader struct {
|
||||
// extractImages controls whether embedded images are parsed and placed
|
||||
// in the markdown. Off halves parse time for text-only indexing.
|
||||
extractImages bool
|
||||
// fallback is used when a PDF has no extractable text. Nil means the
|
||||
// conversion error is returned to the caller.
|
||||
fallback interfaces.DocReader
|
||||
}
|
||||
|
||||
// NewAnydocReader builds a reader. The "anydoc_extract_images" override turns
|
||||
// off image extraction. fallback may be nil.
|
||||
func NewAnydocReader(overrides map[string]string, fallback interfaces.DocReader) *AnydocReader {
|
||||
return &AnydocReader{
|
||||
extractImages: !isFalsey(overrides["anydoc_extract_images"]),
|
||||
fallback: fallback,
|
||||
}
|
||||
}
|
||||
|
||||
// Read converts the document carried by the request.
|
||||
func (r *AnydocReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
if req.URL != "" && len(req.FileContent) == 0 {
|
||||
return nil, fmt.Errorf("anydoc engine reads uploaded documents, not URLs")
|
||||
}
|
||||
|
||||
format, ok := anydoc.FormatForFile(req.FileType, req.FileName)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("anydoc engine does not support file type %q", fileTypeOf(req))
|
||||
}
|
||||
|
||||
converted, err := anydoc.Convert(req.FileContent, anydoc.Options{
|
||||
Format: format,
|
||||
WithAssets: r.extractImages,
|
||||
})
|
||||
if format == "pdf" && (anydoc.PDFNeedsOCR(err) || pdfMarkdownIsEmpty(converted, err)) {
|
||||
return r.readScannedPDF(ctx, req, err)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("anydoc conversion failed for %q: %w", req.FileName, err)
|
||||
}
|
||||
if converted.AssetsError != nil {
|
||||
logger.Warnf(ctx, "[anydoc] %q parsed as text but its images were dropped: %v",
|
||||
req.FileName, converted.AssetsError)
|
||||
}
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: converted.Markdown,
|
||||
ImageRefs: imageRefsFromAssets(converted.Assets),
|
||||
Metadata: map[string]string{
|
||||
"parser": AnydocEngineName,
|
||||
"anydoc_version": anydoc.Version(),
|
||||
"source_format": format,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func pdfMarkdownIsEmpty(converted *anydoc.Result, err error) bool {
|
||||
if err != nil || converted == nil {
|
||||
return false
|
||||
}
|
||||
return strings.TrimSpace(converted.Markdown) == ""
|
||||
}
|
||||
|
||||
func (r *AnydocReader) readScannedPDF(ctx context.Context, req *types.ReadRequest, convertErr error) (*types.ReadResult, error) {
|
||||
if !r.canFallback() {
|
||||
if convertErr != nil {
|
||||
return nil, fmt.Errorf("anydoc conversion failed for %q: %w", req.FileName, convertErr)
|
||||
}
|
||||
return nil, fmt.Errorf("anydoc conversion failed for %q: PDF has no extractable text; OCR is required", req.FileName)
|
||||
}
|
||||
|
||||
logger.Infof(ctx, "[anydoc] %q has no text layer, falling back to builtin for scanned-PDF OCR", req.FileName)
|
||||
fallbackReq := *req
|
||||
fallbackReq.ParserEngine = BuiltinEngineName
|
||||
result, err := r.fallback.Read(ctx, &fallbackReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("anydoc scanned-PDF fallback failed for %q: %w", req.FileName, err)
|
||||
}
|
||||
if result == nil {
|
||||
return nil, fmt.Errorf("anydoc scanned-PDF fallback returned no result for %q", req.FileName)
|
||||
}
|
||||
if result.Metadata == nil {
|
||||
result.Metadata = map[string]string{}
|
||||
}
|
||||
if result.Metadata["parser"] == "" {
|
||||
result.Metadata["parser"] = BuiltinEngineName
|
||||
}
|
||||
result.Metadata["anydoc_fallback"] = "scanned_pdf"
|
||||
if result.Metadata["image_source_type"] == "" {
|
||||
result.Metadata["image_source_type"] = "scanned_pdf"
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (r *AnydocReader) canFallback() bool {
|
||||
if r.fallback == nil {
|
||||
return false
|
||||
}
|
||||
if connected, ok := r.fallback.(interface{ IsConnected() bool }); ok {
|
||||
return connected.IsConnected()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// imageRefsFromAssets turns extracted assets into ImageRefs that match the
|
||||
// in-place markdown links renderMarkdown wrote.
|
||||
func imageRefsFromAssets(assets []anydoc.Asset) []types.ImageRef {
|
||||
if len(assets) == 0 {
|
||||
return nil
|
||||
}
|
||||
refs := make([]types.ImageRef, 0, len(assets))
|
||||
for _, asset := range assets {
|
||||
ref := anydoc.ImageDir + asset.Name
|
||||
refs = append(refs, types.ImageRef{
|
||||
Filename: asset.Name,
|
||||
OriginalRef: ref,
|
||||
MimeType: asset.MediaType,
|
||||
ImageData: asset.Data,
|
||||
})
|
||||
}
|
||||
return refs
|
||||
}
|
||||
|
||||
// fileTypeOf reports the request's file type, falling back to the file name's
|
||||
// extension the way every reader in this package does.
|
||||
func fileTypeOf(req *types.ReadRequest) string {
|
||||
if ft := strings.TrimPrefix(strings.ToLower(strings.TrimSpace(req.FileType)), "."); ft != "" {
|
||||
return ft
|
||||
}
|
||||
return strings.TrimPrefix(strings.ToLower(filepath.Ext(req.FileName)), ".")
|
||||
}
|
||||
|
||||
func isFalsey(value string) bool {
|
||||
switch strings.ToLower(strings.TrimSpace(value)) {
|
||||
case "false", "0", "no", "off":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser/anydoc"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestAnydocReaderRejectsUnsupportedFileType(t *testing.T) {
|
||||
_, err := NewAnydocReader(nil, nil).Read(context.Background(), &types.ReadRequest{
|
||||
FileName: "photo.png",
|
||||
FileType: "png",
|
||||
FileContent: []byte("not really a png"),
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("Read succeeded for an unsupported file type, want an error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "png") {
|
||||
t.Errorf("error does not name the file type: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnydocReaderRejectsURLs(t *testing.T) {
|
||||
_, err := NewAnydocReader(nil, nil).Read(context.Background(), &types.ReadRequest{
|
||||
URL: "https://example.com/report.docx",
|
||||
FileType: "docx",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("Read succeeded for a URL request, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAnydocReaderHonoursImageExtractionOverride(t *testing.T) {
|
||||
if !NewAnydocReader(nil, nil).extractImages {
|
||||
t.Error("image extraction is off by default, want on")
|
||||
}
|
||||
if NewAnydocReader(map[string]string{"anydoc_extract_images": "false"}, nil).extractImages {
|
||||
t.Error("image extraction stayed on after the override turned it off")
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageRefsFromAssetsMatchInPlaceMarkdownPaths(t *testing.T) {
|
||||
refs := imageRefsFromAssets([]anydoc.Asset{
|
||||
{Name: "image-1.png", MediaType: "image/png", Data: []byte("first")},
|
||||
{Name: "image-2.jpg", MediaType: "image/jpeg", Data: []byte("second")},
|
||||
})
|
||||
if len(refs) != 2 {
|
||||
t.Fatalf("got %d image refs, want 2", len(refs))
|
||||
}
|
||||
if refs[0].OriginalRef != "images/image-1.png" {
|
||||
t.Errorf("first OriginalRef = %q, want images/image-1.png", refs[0].OriginalRef)
|
||||
}
|
||||
if refs[1].MimeType != "image/jpeg" {
|
||||
t.Errorf("second mime type = %q, want image/jpeg", refs[1].MimeType)
|
||||
}
|
||||
if len(refs[0].ImageData) == 0 {
|
||||
t.Error("image ref carries no bytes")
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageRefsFromAssetsEmpty(t *testing.T) {
|
||||
if refs := imageRefsFromAssets(nil); refs != nil {
|
||||
t.Errorf("got %d image refs, want none", len(refs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnydocReaderFallsBackForScannedPDF(t *testing.T) {
|
||||
fallback := &stubDocReader{result: &types.ReadResult{
|
||||
MarkdownContent: "",
|
||||
ImageRefs: []types.ImageRef{{
|
||||
Filename: "page_1.jpg",
|
||||
OriginalRef: "images/page_1.jpg",
|
||||
ImageData: []byte("jpeg"),
|
||||
}},
|
||||
Metadata: map[string]string{"image_source_type": "scanned_pdf"},
|
||||
}}
|
||||
result, err := NewAnydocReader(nil, fallback).readScannedPDF(context.Background(), &types.ReadRequest{
|
||||
FileName: "exam.pdf",
|
||||
FileType: "pdf",
|
||||
}, errors.New("PDF has no extractable text (Scanned, 5 pages): OCR is required"))
|
||||
if err != nil {
|
||||
t.Fatalf("readScannedPDF: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.MarkdownContent, "images/page_1.jpg") {
|
||||
t.Errorf("fallback markdown missing page image:\n%s", result.MarkdownContent)
|
||||
}
|
||||
if result.Metadata["anydoc_fallback"] != "scanned_pdf" {
|
||||
t.Errorf("anydoc_fallback = %q, want scanned_pdf", result.Metadata["anydoc_fallback"])
|
||||
}
|
||||
if result.Metadata["image_source_type"] != "scanned_pdf" {
|
||||
t.Errorf("image_source_type = %q, want scanned_pdf", result.Metadata["image_source_type"])
|
||||
}
|
||||
if fallback.got == nil || fallback.got.ParserEngine != BuiltinEngineName {
|
||||
t.Errorf("fallback request engine = %v, want builtin", fallback.got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnydocReaderScannedPDFWithoutFallbackKeepsError(t *testing.T) {
|
||||
_, err := NewAnydocReader(nil, nil).readScannedPDF(context.Background(), &types.ReadRequest{
|
||||
FileName: "exam.pdf",
|
||||
}, errors.New("PDF has no extractable text: OCR is required"))
|
||||
if err == nil {
|
||||
t.Fatal("readScannedPDF succeeded without a fallback, want an error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "OCR is required") {
|
||||
t.Errorf("error does not mention OCR: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnydocReaderDisconnectedFallbackKeepsError(t *testing.T) {
|
||||
fallback := &disconnectedDocReader{}
|
||||
_, err := NewAnydocReader(nil, fallback).readScannedPDF(context.Background(), &types.ReadRequest{
|
||||
FileName: "exam.pdf",
|
||||
}, errors.New("PDF has no extractable text: OCR is required"))
|
||||
if err == nil {
|
||||
t.Fatal("readScannedPDF succeeded with a disconnected fallback, want an error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "OCR is required") {
|
||||
t.Errorf("error does not mention OCR: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
type stubDocReader struct {
|
||||
got *types.ReadRequest
|
||||
result *types.ReadResult
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *stubDocReader) Read(_ context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
s.got = req
|
||||
return s.result, s.err
|
||||
}
|
||||
|
||||
type disconnectedDocReader struct{}
|
||||
|
||||
func (disconnectedDocReader) Read(context.Context, *types.ReadRequest) (*types.ReadResult, error) {
|
||||
return nil, errors.New("should not be called")
|
||||
}
|
||||
|
||||
func (disconnectedDocReader) IsConnected() bool { return false }
|
||||
@@ -0,0 +1,223 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// simpleFormats lists file extensions that Go can handle without the Python service.
|
||||
var simpleFormats = map[string]bool{
|
||||
"md": true, "markdown": true,
|
||||
"txt": true, "text": true,
|
||||
"csv": true,
|
||||
"json": true,
|
||||
}
|
||||
|
||||
var imageFormats = map[string]bool{
|
||||
"jpg": true, "jpeg": true, "png": true, "gif": true,
|
||||
"bmp": true, "tiff": true, "webp": true,
|
||||
}
|
||||
|
||||
var audioFormats = map[string]bool{
|
||||
"mp3": true, "wav": true, "m4a": true, "flac": true, "ogg": true,
|
||||
}
|
||||
|
||||
func init() {
|
||||
for k := range imageFormats {
|
||||
simpleFormats[k] = true
|
||||
}
|
||||
for k := range audioFormats {
|
||||
simpleFormats[k] = true
|
||||
}
|
||||
}
|
||||
|
||||
// IsSimpleFormat returns true if the file type can be handled by the Go SimpleFormatReader.
|
||||
func IsSimpleFormat(fileType string) bool {
|
||||
return simpleFormats[strings.ToLower(strings.TrimPrefix(fileType, "."))]
|
||||
}
|
||||
|
||||
// SimpleFormatReader handles simple file formats and images directly in Go,
|
||||
// bypassing the Python docreader service.
|
||||
type SimpleFormatReader struct{}
|
||||
|
||||
// Read reads simple format files and returns markdown.
|
||||
func (b *SimpleFormatReader) Read(_ context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
ft := strings.ToLower(strings.TrimPrefix(req.FileType, "."))
|
||||
if ft == "" {
|
||||
ft = strings.TrimPrefix(strings.ToLower(filepath.Ext(req.FileName)), ".")
|
||||
}
|
||||
|
||||
switch {
|
||||
case ft == "md" || ft == "markdown":
|
||||
return &types.ReadResult{MarkdownContent: string(req.FileContent)}, nil
|
||||
case ft == "txt" || ft == "text":
|
||||
return &types.ReadResult{MarkdownContent: string(req.FileContent)}, nil
|
||||
case ft == "csv":
|
||||
md, err := csvToMarkdown(req.FileContent)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("csv conversion failed: %w", err)
|
||||
}
|
||||
return &types.ReadResult{MarkdownContent: md}, nil
|
||||
case ft == "json":
|
||||
md, err := jsonToMarkdown(req.FileContent)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("json conversion failed: %w", err)
|
||||
}
|
||||
return &types.ReadResult{MarkdownContent: md}, nil
|
||||
case imageFormats[ft]:
|
||||
return imageToResult(req.FileName, req.FileContent), nil
|
||||
case audioFormats[ft]:
|
||||
return audioToResult(req.FileName, req.FileContent), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported simple format: %s", ft)
|
||||
}
|
||||
}
|
||||
|
||||
// imageToResult wraps a standalone image as a markdown image reference with
|
||||
// the raw bytes in ImageRefs, matching Python ImageParser behaviour.
|
||||
func imageToResult(fileName string, data []byte) *types.ReadResult {
|
||||
if fileName == "" {
|
||||
fileName = "image.png"
|
||||
}
|
||||
refPath := "images/" + fileName
|
||||
// Encode spaces so the markdown URL is valid and matches the regex in ResolveAndStore.
|
||||
safeRef := strings.ReplaceAll(refPath, " ", "%20")
|
||||
mime := http.DetectContentType(data)
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: fmt.Sprintf("", fileName, safeRef),
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: fileName,
|
||||
OriginalRef: safeRef,
|
||||
MimeType: mime,
|
||||
ImageData: data,
|
||||
IsOriginal: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// IsImageFormat returns true if the file type is a recognized image format.
|
||||
func IsImageFormat(fileType string) bool {
|
||||
return imageFormats[strings.ToLower(strings.TrimPrefix(fileType, "."))]
|
||||
}
|
||||
|
||||
// IsAudioFormat returns true if the file type is a recognized audio format.
|
||||
func IsAudioFormat(fileType string) bool {
|
||||
return audioFormats[strings.ToLower(strings.TrimPrefix(fileType, "."))]
|
||||
}
|
||||
|
||||
// audioToResult wraps a standalone audio file. The actual transcription is
|
||||
// handled by the ASR model in the knowledge service pipeline. Here we just
|
||||
// return a placeholder markdown with the raw bytes preserved for upstream
|
||||
// processing.
|
||||
func audioToResult(fileName string, data []byte) *types.ReadResult {
|
||||
if fileName == "" {
|
||||
fileName = "audio.mp3"
|
||||
}
|
||||
// Return a placeholder; the knowledge service will replace this with
|
||||
// the ASR transcription result.
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: fmt.Sprintf("[Audio file: %s]", fileName),
|
||||
IsAudio: true,
|
||||
AudioData: data,
|
||||
}
|
||||
}
|
||||
|
||||
// ensureOriginalImageRef checks whether the input file is an image and, if the
|
||||
// returned markdown does not already contain a markdown image reference for it,
|
||||
// prepends one and appends the raw bytes to imageRefs. This guarantees that
|
||||
// when MinerU OCRs a standalone image, the downstream chunks still carry the
|
||||
// original image link for retrieval display.
|
||||
func ensureOriginalImageRef(req *types.ReadRequest, mdContent string, imageRefs []types.ImageRef) (string, []types.ImageRef) {
|
||||
ft := strings.ToLower(strings.TrimPrefix(req.FileType, "."))
|
||||
if ft == "" {
|
||||
ft = strings.TrimPrefix(strings.ToLower(filepath.Ext(req.FileName)), ".")
|
||||
}
|
||||
if !imageFormats[ft] {
|
||||
return mdContent, imageRefs
|
||||
}
|
||||
if len(req.FileContent) == 0 {
|
||||
return mdContent, imageRefs
|
||||
}
|
||||
|
||||
fileName := req.FileName
|
||||
if fileName == "" {
|
||||
fileName = "image." + ft
|
||||
}
|
||||
refPath := "images/" + fileName
|
||||
|
||||
if strings.Contains(mdContent, refPath) {
|
||||
return mdContent, imageRefs
|
||||
}
|
||||
|
||||
imgLine := fmt.Sprintf("", fileName, refPath)
|
||||
if strings.TrimSpace(mdContent) == "" {
|
||||
mdContent = imgLine
|
||||
} else {
|
||||
mdContent = imgLine + "\n\n" + mdContent
|
||||
}
|
||||
|
||||
mime := http.DetectContentType(req.FileContent)
|
||||
imageRefs = append(imageRefs, types.ImageRef{
|
||||
Filename: fileName,
|
||||
OriginalRef: refPath,
|
||||
MimeType: mime,
|
||||
ImageData: req.FileContent,
|
||||
IsOriginal: true,
|
||||
})
|
||||
|
||||
return mdContent, imageRefs
|
||||
}
|
||||
|
||||
func csvToMarkdown(data []byte) (string, error) {
|
||||
reader := csv.NewReader(strings.NewReader(string(data)))
|
||||
reader.LazyQuotes = true
|
||||
reader.TrimLeadingSpace = true
|
||||
|
||||
records, err := reader.ReadAll()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(records) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
|
||||
// Header row
|
||||
header := records[0]
|
||||
sb.WriteString("| ")
|
||||
sb.WriteString(strings.Join(header, " | "))
|
||||
sb.WriteString(" |\n")
|
||||
|
||||
// Separator
|
||||
sb.WriteString("|")
|
||||
for range header {
|
||||
sb.WriteString(" --- |")
|
||||
}
|
||||
sb.WriteString("\n")
|
||||
|
||||
// Data rows
|
||||
for _, row := range records[1:] {
|
||||
sb.WriteString("| ")
|
||||
// Pad row if shorter than header
|
||||
cells := make([]string, len(header))
|
||||
for i := range cells {
|
||||
if i < len(row) {
|
||||
cells[i] = row[i]
|
||||
}
|
||||
}
|
||||
sb.WriteString(strings.Join(cells, " | "))
|
||||
sb.WriteString(" |\n")
|
||||
}
|
||||
|
||||
return sb.String(), nil
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
// Package docparser turns uploaded files and URLs into Markdown plus the
|
||||
// images they contain, which is the input the chunking and multimodal stages
|
||||
// downstream expect.
|
||||
//
|
||||
// A parse request names a parser engine, and every engine is one entry in the
|
||||
// registry:
|
||||
//
|
||||
// - engines.go declares the catalog — each engine's name, description,
|
||||
// supported file types, availability probe, and the reader that does the
|
||||
// work.
|
||||
// - engine_registry.go turns an engine name into a reader (NewReader) and
|
||||
// merges the local catalog with the engines the Python docreader reports
|
||||
// over its ListEngines RPC (ListAllEngines).
|
||||
//
|
||||
// Readers differ in where the parsing happens: in this process
|
||||
// (SimpleFormatReader, AnydocReader), in the docreader service over gRPC or
|
||||
// HTTP (GRPCDocumentReader, HTTPDocumentReader), or in a remote API (MinerU,
|
||||
// PaddleOCR-VL, WeKnora Cloud). They all return types.ReadResult, so the rest
|
||||
// of the package — image resolution and storage, table normalization — is
|
||||
// shared regardless of which engine ran.
|
||||
package docparser
|
||||
@@ -0,0 +1,151 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// EngineRegistration is what every locally registered parser engine provides:
|
||||
// the metadata the engine list shows, and the reader that does the parsing.
|
||||
// Remote-only engines (e.g. markitdown) live in the Python docreader, are
|
||||
// discovered through its ListEngines RPC, and never register here — the
|
||||
// registry routes them to the docreader client by default.
|
||||
type EngineRegistration interface {
|
||||
Name() string
|
||||
Description() string
|
||||
FileTypes(docreaderConnected bool) []string
|
||||
CheckAvailable(docreaderConnected bool, overrides map[string]string) (available bool, reason string)
|
||||
// NewReader builds the reader for one parse request. Returning an error
|
||||
// means this engine cannot serve the request (missing credentials,
|
||||
// unreachable service); the caller reports it rather than silently
|
||||
// parsing with something else.
|
||||
NewReader(ctx context.Context, deps ReaderDeps) (interfaces.DocReader, error)
|
||||
}
|
||||
|
||||
// ReaderDeps carries everything an engine may need to build its reader but
|
||||
// cannot construct itself: tenant configuration, tenant credentials, and the
|
||||
// shared docreader connection.
|
||||
type ReaderDeps struct {
|
||||
// Overrides holds tenant-level engine configuration (service endpoints,
|
||||
// API keys), as produced by ParserEngineConfig.ToOverridesMap.
|
||||
Overrides map[string]string
|
||||
// Remote is the docreader client. Nil when the service is not connected.
|
||||
Remote interfaces.DocReader
|
||||
// WeKnoraCloudCredentials resolves the tenant's WeKnora Cloud
|
||||
// credentials. It is a function rather than a value because resolving
|
||||
// them can hit the database, which most engines never need. Nil, or a
|
||||
// nil return, means the tenant has not configured them.
|
||||
WeKnoraCloudCredentials func(ctx context.Context) *types.WeKnoraCloudCredentials
|
||||
}
|
||||
|
||||
// localEngines holds all locally registered parser engines, in registration
|
||||
// order — which is also the order the engine list is shown in.
|
||||
var localEngines []EngineRegistration
|
||||
|
||||
// RegisterEngine adds an engine to the local registry. Called from init().
|
||||
func RegisterEngine(e EngineRegistration) {
|
||||
localEngines = append(localEngines, e)
|
||||
}
|
||||
|
||||
// lookupEngine returns the locally registered engine with this name.
|
||||
func lookupEngine(name string) (EngineRegistration, bool) {
|
||||
for _, engine := range localEngines {
|
||||
if engine.Name() == name {
|
||||
return engine, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// NewReader builds the reader for an engine.
|
||||
//
|
||||
// An empty engine name means "no explicit choice": simple formats are handled
|
||||
// in Go and everything else goes to the docreader service. An unknown name is
|
||||
// routed to the docreader too, so engines that only exist in the Python
|
||||
// service keep working without a Go-side registration.
|
||||
func NewReader(
|
||||
ctx context.Context, engine, fileType string, isURL bool, deps ReaderDeps,
|
||||
) (interfaces.DocReader, error) {
|
||||
if registration, ok := lookupEngine(engine); ok {
|
||||
return registration.NewReader(ctx, deps)
|
||||
}
|
||||
if engine == "" && !isURL && IsSimpleFormat(fileType) {
|
||||
return &SimpleFormatReader{}, nil
|
||||
}
|
||||
return remoteReader(deps)
|
||||
}
|
||||
|
||||
// remoteReader returns the docreader client, or an error when the service is
|
||||
// not connected — a nil interface value here would panic at the call site.
|
||||
func remoteReader(deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
if deps.Remote == nil {
|
||||
return nil, errNotConnected
|
||||
}
|
||||
return deps.Remote, nil
|
||||
}
|
||||
|
||||
// ListAllEngines returns the merged engine list: locally registered engines
|
||||
// plus engines discovered from the remote docreader via ListEngines RPC.
|
||||
//
|
||||
// Merge rules:
|
||||
// - Local engines are always included, with Go-side availability checks.
|
||||
// - For a remote engine whose name matches a local one, the remote's
|
||||
// file_types and description take precedence (the remote service is
|
||||
// authoritative for its own capabilities).
|
||||
// - Remote engines not present locally are appended as-is, enabling
|
||||
// auto-discovery of newly added docreader engines without Go changes.
|
||||
func ListAllEngines(
|
||||
docreaderConnected bool, overrides map[string]string, remoteEngines []types.ParserEngineInfo,
|
||||
) []types.ParserEngineInfo {
|
||||
remoteMap := make(map[string]types.ParserEngineInfo, len(remoteEngines))
|
||||
for _, re := range remoteEngines {
|
||||
remoteMap[re.Name] = re
|
||||
}
|
||||
|
||||
seen := make(map[string]bool, len(localEngines))
|
||||
result := make([]types.ParserEngineInfo, 0, len(localEngines)+len(remoteEngines))
|
||||
|
||||
for _, e := range localEngines {
|
||||
name := e.Name()
|
||||
seen[name] = true
|
||||
|
||||
fileTypes := e.FileTypes(docreaderConnected)
|
||||
description := e.Description()
|
||||
|
||||
if re, ok := remoteMap[name]; ok {
|
||||
if len(re.FileTypes) > 0 {
|
||||
fileTypes = re.FileTypes
|
||||
}
|
||||
if re.Description != "" {
|
||||
description = re.Description
|
||||
}
|
||||
}
|
||||
|
||||
available, reason := e.CheckAvailable(docreaderConnected, overrides)
|
||||
result = append(result, types.ParserEngineInfo{
|
||||
Name: name,
|
||||
Description: description,
|
||||
FileTypes: fileTypes,
|
||||
Available: available,
|
||||
UnavailableReason: reason,
|
||||
})
|
||||
}
|
||||
|
||||
for _, re := range remoteEngines {
|
||||
if seen[re.Name] {
|
||||
continue
|
||||
}
|
||||
result = append(result, re)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// errEngineUnavailable reports an engine that is registered but cannot run for
|
||||
// this tenant or this build.
|
||||
func errEngineUnavailable(engine, reason string) error {
|
||||
return fmt.Errorf("parser engine %q is unavailable: %s", engine, reason)
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser/anydoc"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// stubRemote stands in for the docreader client. NewReader only ever hands it
|
||||
// back, so the embedded interface is never called — and would panic if routing
|
||||
// ever sent a parse here by mistake.
|
||||
type stubRemote struct{ interfaces.DocReader }
|
||||
|
||||
func TestNewReaderRoutesByEngine(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
remote := &stubRemote{}
|
||||
deps := ReaderDeps{Remote: remote}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
engine string
|
||||
fileType string
|
||||
isURL bool
|
||||
want any
|
||||
}{
|
||||
{name: "simple engine", engine: SimpleEngineName, fileType: "md", want: &SimpleFormatReader{}},
|
||||
{name: "builtin engine goes remote", engine: BuiltinEngineName, fileType: "md", want: remote},
|
||||
{name: "unset engine handles simple formats in Go", fileType: "csv", want: &SimpleFormatReader{}},
|
||||
{name: "unset engine sends complex formats to docreader", fileType: "docx", want: remote},
|
||||
{name: "default pptx engine is remote markitdown", engine: "markitdown", fileType: "pptx", want: remote},
|
||||
{name: "URLs always go to docreader", fileType: "md", isURL: true, want: remote},
|
||||
{name: "docreader-only engines fall through", engine: "markitdown", fileType: "docx", want: remote},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
reader, err := NewReader(ctx, tc.engine, tc.fileType, tc.isURL, deps)
|
||||
if err != nil {
|
||||
t.Fatalf("NewReader: %v", err)
|
||||
}
|
||||
if _, isSimple := tc.want.(*SimpleFormatReader); isSimple {
|
||||
if _, ok := reader.(*SimpleFormatReader); !ok {
|
||||
t.Fatalf("reader = %T, want *SimpleFormatReader", reader)
|
||||
}
|
||||
return
|
||||
}
|
||||
if reader != tc.want {
|
||||
t.Fatalf("reader = %T, want the docreader client", reader)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A disconnected docreader has to surface as an error, not as a nil reader the
|
||||
// caller would dereference.
|
||||
func TestNewReaderReportsDisconnectedDocReader(t *testing.T) {
|
||||
if _, err := NewReader(context.Background(), BuiltinEngineName, "docx", false, ReaderDeps{}); err == nil {
|
||||
t.Fatal("NewReader succeeded without a docreader connection, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewReaderRequiresWeKnoraCloudCredentials(t *testing.T) {
|
||||
_, err := NewReader(context.Background(), WeKnoraCloudEngineName, "docx", false, ReaderDeps{
|
||||
WeKnoraCloudCredentials: func(context.Context) *types.WeKnoraCloudCredentials { return nil },
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("NewReader succeeded without credentials, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
// The anydoc engine is only linked into builds tagged `anydoc`; everywhere
|
||||
// else it must be listed as unavailable and refuse to build a reader, so a
|
||||
// knowledge base configured for it fails loudly instead of parsing with
|
||||
// something else.
|
||||
func TestAnydocEngineFollowsBuildAvailability(t *testing.T) {
|
||||
reader, err := NewReader(context.Background(), AnydocEngineName, "docx", false, ReaderDeps{})
|
||||
|
||||
if anydoc.Available() {
|
||||
if err != nil {
|
||||
t.Fatalf("NewReader: %v", err)
|
||||
}
|
||||
if _, ok := reader.(*AnydocReader); !ok {
|
||||
t.Fatalf("reader = %T, want *AnydocReader", reader)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatal("NewReader succeeded without the converter linked in, want an error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAllEnginesIncludesAnydoc(t *testing.T) {
|
||||
for _, engine := range ListAllEngines(true, nil, nil) {
|
||||
if engine.Name != AnydocEngineName {
|
||||
continue
|
||||
}
|
||||
if engine.Available != anydoc.Available() {
|
||||
t.Errorf("anydoc availability = %v, want %v", engine.Available, anydoc.Available())
|
||||
}
|
||||
if !engine.Available && engine.UnavailableReason == "" {
|
||||
t.Error("anydoc is unavailable without a reason to show")
|
||||
}
|
||||
if len(engine.FileTypes) == 0 {
|
||||
t.Error("anydoc lists no file types")
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatal("anydoc engine not found in the engine list")
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser/anydoc"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestListAllEnginesBuiltinIncludesDocumentFormats(t *testing.T) {
|
||||
engines := ListAllEngines(true, nil, nil)
|
||||
for _, engine := range engines {
|
||||
if engine.Name != "builtin" {
|
||||
continue
|
||||
}
|
||||
if !engine.Available {
|
||||
t.Fatalf("builtin engine is unavailable: %s", engine.UnavailableReason)
|
||||
}
|
||||
|
||||
fileTypes := make(map[string]bool, len(engine.FileTypes))
|
||||
for _, fileType := range engine.FileTypes {
|
||||
fileTypes[fileType] = true
|
||||
}
|
||||
for _, want := range []string{"html", "htm", "xmind", "ppt", "pptx"} {
|
||||
if !fileTypes[want] {
|
||||
t.Errorf("builtin engine file types do not include %q: %v", want, engine.FileTypes)
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
t.Fatal("builtin engine not found")
|
||||
}
|
||||
|
||||
func TestDefaultParserEnginePrefersAnydocWhenLinked(t *testing.T) {
|
||||
cases := []string{"pptx", "ppt", "pdf", "docx"}
|
||||
if anydoc.Available() {
|
||||
for _, ft := range cases {
|
||||
if got := types.DefaultParserEngine(ft); got != AnydocEngineName {
|
||||
t.Errorf("DefaultParserEngine(%s) = %q, want anydoc when the binding is linked", ft, got)
|
||||
}
|
||||
}
|
||||
if got := types.DefaultParserEngine("csv"); got != "" {
|
||||
t.Errorf("DefaultParserEngine(csv) = %q, want empty so the Go simple reader stays default", got)
|
||||
}
|
||||
return
|
||||
}
|
||||
if got := types.DefaultParserEngine("pptx"); got != "markitdown" {
|
||||
t.Fatalf("DefaultParserEngine(pptx) = %q, want markitdown when anydoc is unavailable", got)
|
||||
}
|
||||
if got := types.DefaultParserEngine("pdf"); got != "" {
|
||||
t.Fatalf("DefaultParserEngine(pdf) = %q, want empty when anydoc is unavailable", got)
|
||||
}
|
||||
if got := types.DefaultParserEngine("docx"); got != "" {
|
||||
t.Fatalf("DefaultParserEngine(docx) = %q, want empty when anydoc is unavailable", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,287 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/infrastructure/docparser/anydoc"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// Engine names, as stored in knowledge base parser rules and shown in the UI.
|
||||
const (
|
||||
// BuiltinEngineName is the DocReader (Python) parser suite.
|
||||
BuiltinEngineName = "builtin"
|
||||
// SimpleEngineName is Go-native handling of text formats and images.
|
||||
SimpleEngineName = "simple"
|
||||
// AnydocEngineName is the in-process anydoc office-document converter.
|
||||
AnydocEngineName = "anydoc"
|
||||
// WeKnoraCloudEngineName is the hosted WeKnora Cloud document reader.
|
||||
WeKnoraCloudEngineName = "weknoracloud"
|
||||
// MinerUEngineName is a self-hosted MinerU service.
|
||||
MinerUEngineName = "mineru"
|
||||
// MinerUCloudEngineName is the MinerU Cloud API.
|
||||
MinerUCloudEngineName = "mineru_cloud"
|
||||
// PaddleOCRVLEngineName is a self-hosted PaddleOCR-VL pipeline.
|
||||
PaddleOCRVLEngineName = "paddleocr_vl"
|
||||
// PaddleOCRVLCloudEngineName is the PaddleOCR-VL AI Studio cloud API.
|
||||
PaddleOCRVLCloudEngineName = "paddleocr_vl_cloud"
|
||||
)
|
||||
|
||||
func init() {
|
||||
types.SetPreferParserEngine(preferAnydocWhenAvailable)
|
||||
RegisterEngine(&builtinEngine{})
|
||||
RegisterEngine(&simpleEngine{})
|
||||
RegisterEngine(&anydocEngine{})
|
||||
RegisterEngine(&weKnoraCloudEngine{})
|
||||
RegisterEngine(&mineruEngine{})
|
||||
RegisterEngine(&mineruCloudEngine{})
|
||||
RegisterEngine(&paddleOCRVLEngine{})
|
||||
RegisterEngine(&paddleOCRVLCloudEngine{})
|
||||
}
|
||||
|
||||
// preferAnydocWhenAvailable is the type-level default override: when the
|
||||
// anydoc binding is linked and converts this file type, use it instead of
|
||||
// builtin or the markitdown fallback. Simple formats stay on the Go reader.
|
||||
func preferAnydocWhenAvailable(fileType string) string {
|
||||
if IsSimpleFormat(fileType) {
|
||||
return ""
|
||||
}
|
||||
if anydoc.Available() && anydoc.Supports(fileType, "") {
|
||||
return AnydocEngineName
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// builtin — DocReader-backed parser for complex document formats.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type builtinEngine struct{}
|
||||
|
||||
func (e *builtinEngine) Name() string { return BuiltinEngineName }
|
||||
|
||||
func (e *builtinEngine) Description() string { return "DocReader built-in parser engine" }
|
||||
|
||||
func (e *builtinEngine) FileTypes(_ bool) []string {
|
||||
return []string{
|
||||
"docx", "doc", "pdf", "md", "markdown", "xlsx", "xls",
|
||||
"pptx", "ppt", "epub",
|
||||
"html", "htm", "mhtml", "xmind",
|
||||
"jpg", "jpeg", "png", "gif", "bmp", "tiff", "webp",
|
||||
"mp3", "wav", "m4a", "flac", "ogg",
|
||||
}
|
||||
}
|
||||
|
||||
func (e *builtinEngine) CheckAvailable(docreaderConnected bool, _ map[string]string) (bool, string) {
|
||||
if docreaderConnected {
|
||||
return true, ""
|
||||
}
|
||||
return false, "DocReader service not connected"
|
||||
}
|
||||
|
||||
// NewReader returns the docreader client. Selecting "builtin" explicitly means
|
||||
// the docreader is wanted even for formats the simple reader could handle.
|
||||
func (e *builtinEngine) NewReader(_ context.Context, deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
return remoteReader(deps)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// simple — Go handles md/txt/csv/json natively, no external service needed.
|
||||
// Distinct from docreader's "builtin", which uses Python libraries for complex
|
||||
// formats (docx, pdf, ...).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type simpleEngine struct{}
|
||||
|
||||
func (e *simpleEngine) Name() string { return SimpleEngineName }
|
||||
|
||||
func (e *simpleEngine) Description() string {
|
||||
return "Simple format & image parsing (no external service required)"
|
||||
}
|
||||
|
||||
func (e *simpleEngine) FileTypes(_ bool) []string {
|
||||
return []string{
|
||||
"md", "markdown", "txt", "csv", "json",
|
||||
"jpg", "jpeg", "png", "gif", "bmp", "tiff", "webp",
|
||||
"mp3", "wav", "m4a", "flac", "ogg",
|
||||
}
|
||||
}
|
||||
|
||||
func (e *simpleEngine) CheckAvailable(_ bool, _ map[string]string) (bool, string) {
|
||||
return true, ""
|
||||
}
|
||||
|
||||
func (e *simpleEngine) NewReader(_ context.Context, _ ReaderDeps) (interfaces.DocReader, error) {
|
||||
return &SimpleFormatReader{}, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// anydoc — office documents converted in this process, no external service.
|
||||
// Only present in builds that link the converter (see the anydoc package).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type anydocEngine struct{}
|
||||
|
||||
func (e *anydocEngine) Name() string { return AnydocEngineName }
|
||||
|
||||
func (e *anydocEngine) Description() string {
|
||||
return "anydoc in-process office document converter (no external service required)"
|
||||
}
|
||||
|
||||
func (e *anydocEngine) FileTypes(_ bool) []string { return anydoc.SupportedFileTypes() }
|
||||
|
||||
func (e *anydocEngine) CheckAvailable(_ bool, _ map[string]string) (bool, string) {
|
||||
if anydoc.Available() {
|
||||
return true, ""
|
||||
}
|
||||
return false, anydoc.UnavailableReason()
|
||||
}
|
||||
|
||||
func (e *anydocEngine) NewReader(_ context.Context, deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
if !anydoc.Available() {
|
||||
return nil, errEngineUnavailable(AnydocEngineName, anydoc.UnavailableReason())
|
||||
}
|
||||
return NewAnydocReader(deps.Overrides, deps.Remote), nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// weknoracloud — Tenant-scoped WeKnoraCloud docreader with signed requests.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type weKnoraCloudEngine struct{}
|
||||
|
||||
func (e *weKnoraCloudEngine) Name() string { return WeKnoraCloudEngineName }
|
||||
|
||||
func (e *weKnoraCloudEngine) Description() string { return "WeKnoraCloud document reader" }
|
||||
|
||||
func (e *weKnoraCloudEngine) FileTypes(_ bool) []string {
|
||||
return []string{"docx", "doc", "pdf", "md", "markdown", "xlsx", "xls", "pptx", "ppt"}
|
||||
}
|
||||
|
||||
func (e *weKnoraCloudEngine) CheckAvailable(_ bool, overrides map[string]string) (bool, string) {
|
||||
if overrides["weknoracloud_app_id"] != "" {
|
||||
return true, ""
|
||||
}
|
||||
return false, "WeKnora Cloud credentials not configured. Go to Settings → WeKnora Cloud to set up."
|
||||
}
|
||||
|
||||
func (e *weKnoraCloudEngine) NewReader(
|
||||
ctx context.Context, deps ReaderDeps,
|
||||
) (interfaces.DocReader, error) {
|
||||
if deps.WeKnoraCloudCredentials == nil {
|
||||
return nil, errEngineUnavailable(WeKnoraCloudEngineName, "no credential resolver configured")
|
||||
}
|
||||
creds := deps.WeKnoraCloudCredentials(ctx)
|
||||
if creds == nil {
|
||||
return nil, errEngineUnavailable(WeKnoraCloudEngineName, "tenant credentials not configured")
|
||||
}
|
||||
return NewWeKnoraCloudSignedDocumentReader(creds.AppID, creds.AppSecret)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mineru — Go-native, calls self-hosted MinerU API directly
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type mineruEngine struct{}
|
||||
|
||||
func (e *mineruEngine) Name() string { return MinerUEngineName }
|
||||
|
||||
func (e *mineruEngine) Description() string { return "MinerU self-hosted service" }
|
||||
|
||||
func (e *mineruEngine) FileTypes(_ bool) []string {
|
||||
return []string{"pdf", "jpg", "jpeg", "png", "bmp", "tiff", "doc", "docx", "ppt", "pptx"}
|
||||
}
|
||||
|
||||
func (e *mineruEngine) CheckAvailable(_ bool, overrides map[string]string) (bool, string) {
|
||||
endpoint := strings.TrimSpace(overrides["mineru_endpoint"])
|
||||
if endpoint == "" {
|
||||
return false, "MinerU service not configured"
|
||||
}
|
||||
return PingMinerU(endpoint)
|
||||
}
|
||||
|
||||
func (e *mineruEngine) NewReader(_ context.Context, deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
return NewMinerUReader(deps.Overrides), nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mineru_cloud — Go-native, calls MinerU Cloud API directly
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type mineruCloudEngine struct{}
|
||||
|
||||
func (e *mineruCloudEngine) Name() string { return MinerUCloudEngineName }
|
||||
|
||||
func (e *mineruCloudEngine) Description() string { return "MinerU Cloud API" }
|
||||
|
||||
func (e *mineruCloudEngine) FileTypes(_ bool) []string {
|
||||
return []string{"pdf", "jpg", "jpeg", "png", "bmp", "tiff", "doc", "docx", "ppt", "pptx"}
|
||||
}
|
||||
|
||||
func (e *mineruCloudEngine) CheckAvailable(_ bool, overrides map[string]string) (bool, string) {
|
||||
apiKey := strings.TrimSpace(overrides["mineru_api_key"])
|
||||
if apiKey == "" {
|
||||
return false, "MinerU API Key not configured"
|
||||
}
|
||||
return PingMinerUCloud(apiKey)
|
||||
}
|
||||
|
||||
func (e *mineruCloudEngine) NewReader(_ context.Context, deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
return NewMinerUCloudReader(deps.Overrides), nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// paddleocr_vl — Go-native, calls a self-hosted PaddleOCR-VL pipeline service
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type paddleOCRVLEngine struct{}
|
||||
|
||||
func (e *paddleOCRVLEngine) Name() string { return PaddleOCRVLEngineName }
|
||||
|
||||
func (e *paddleOCRVLEngine) Description() string { return "PaddleOCR-VL self-hosted service" }
|
||||
|
||||
func (e *paddleOCRVLEngine) FileTypes(_ bool) []string {
|
||||
return []string{"pdf", "jpg", "jpeg", "png", "bmp", "tiff"}
|
||||
}
|
||||
|
||||
func (e *paddleOCRVLEngine) CheckAvailable(_ bool, overrides map[string]string) (bool, string) {
|
||||
endpoint := strings.TrimSpace(overrides["paddleocr_vl_endpoint"])
|
||||
if endpoint == "" {
|
||||
return false, "PaddleOCR-VL service not configured"
|
||||
}
|
||||
return PingPaddleOCRVL(endpoint)
|
||||
}
|
||||
|
||||
func (e *paddleOCRVLEngine) NewReader(_ context.Context, deps ReaderDeps) (interfaces.DocReader, error) {
|
||||
return NewPaddleOCRVLReader(deps.Overrides), nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// paddleocr_vl_cloud — Go-native, calls the PaddleOCR-VL AI Studio cloud API
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type paddleOCRVLCloudEngine struct{}
|
||||
|
||||
func (e *paddleOCRVLCloudEngine) Name() string { return PaddleOCRVLCloudEngineName }
|
||||
|
||||
func (e *paddleOCRVLCloudEngine) Description() string { return "PaddleOCR-VL Cloud API" }
|
||||
|
||||
func (e *paddleOCRVLCloudEngine) FileTypes(_ bool) []string {
|
||||
return []string{"pdf", "jpg", "jpeg", "png", "bmp", "tiff"}
|
||||
}
|
||||
|
||||
func (e *paddleOCRVLCloudEngine) CheckAvailable(_ bool, overrides map[string]string) (bool, string) {
|
||||
token := strings.TrimSpace(overrides["paddleocr_vl_cloud_token"])
|
||||
if token == "" {
|
||||
return false, "PaddleOCR-VL Cloud Token not configured"
|
||||
}
|
||||
return PingPaddleOCRVLCloud(token)
|
||||
}
|
||||
|
||||
func (e *paddleOCRVLCloudEngine) NewReader(
|
||||
_ context.Context, deps ReaderDeps,
|
||||
) (interfaces.DocReader, error) {
|
||||
return NewPaddleOCRVLCloudReader(deps.Overrides), nil
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
docclient "github.com/Tencent/WeKnora/docreader/client"
|
||||
"github.com/Tencent/WeKnora/docreader/proto"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/codes"
|
||||
"google.golang.org/grpc/resolver"
|
||||
"google.golang.org/grpc/status"
|
||||
)
|
||||
|
||||
func getMaxMessageSize() int {
|
||||
return int(utils.GetMaxFileSizeMB() * 1024 * 1024)
|
||||
}
|
||||
|
||||
// GRPCDocumentReader implements DocumentReader over gRPC.
|
||||
type GRPCDocumentReader struct {
|
||||
mu sync.RWMutex
|
||||
conn *grpc.ClientConn
|
||||
client proto.DocReaderClient
|
||||
addr string
|
||||
}
|
||||
|
||||
func NewGRPCDocumentReader(addr string) (*GRPCDocumentReader, error) {
|
||||
p := &GRPCDocumentReader{}
|
||||
if addr != "" {
|
||||
if err := p.connect(addr); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
func (p *GRPCDocumentReader) connect(addr string) error {
|
||||
authConfig := docclient.LoadAuthConfigFromEnv()
|
||||
opts, err := authConfig.BuildDialOptions(getMaxMessageSize())
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to build docreader dial options: %w", err)
|
||||
}
|
||||
if authConfig.TLSEnabled {
|
||||
logger.Infof(context.Background(), "TLS enabled for docreader gRPC client")
|
||||
}
|
||||
if authConfig.AuthToken != "" {
|
||||
logger.Infof(context.Background(),
|
||||
"Token authentication enabled for docreader gRPC client (TLS=%v)",
|
||||
authConfig.TLSEnabled,
|
||||
)
|
||||
}
|
||||
|
||||
resolver.SetDefaultScheme("dns")
|
||||
|
||||
start := time.Now()
|
||||
conn, err := grpc.Dial("dns:///"+addr, opts...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to connect to docreader: %w", err)
|
||||
}
|
||||
logger.Infof(context.Background(), "Connected to docreader in %v", time.Since(start))
|
||||
|
||||
p.conn = conn
|
||||
p.client = proto.NewDocReaderClient(conn)
|
||||
p.addr = addr
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *GRPCDocumentReader) Reconnect(addr string) error {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
if p.conn != nil {
|
||||
_ = p.conn.Close()
|
||||
p.conn = nil
|
||||
p.client = nil
|
||||
p.addr = ""
|
||||
}
|
||||
return p.connect(addr)
|
||||
}
|
||||
|
||||
func (p *GRPCDocumentReader) IsConnected() bool {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
return p.conn != nil
|
||||
}
|
||||
|
||||
func (p *GRPCDocumentReader) Close() error {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.conn != nil {
|
||||
return p.conn.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
var errNotConnected = fmt.Errorf("docreader service not connected")
|
||||
|
||||
func (p *GRPCDocumentReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
p.mu.RLock()
|
||||
client := p.client
|
||||
p.mu.RUnlock()
|
||||
if client == nil {
|
||||
return nil, errNotConnected
|
||||
}
|
||||
|
||||
protoReq := &proto.ReadRequest{
|
||||
FileContent: req.FileContent,
|
||||
FileName: req.FileName,
|
||||
FileType: req.FileType,
|
||||
Url: req.URL,
|
||||
Title: req.Title,
|
||||
RequestId: req.RequestID,
|
||||
Config: &proto.ReadConfig{
|
||||
ParserEngine: req.ParserEngine,
|
||||
ParserEngineOverrides: req.ParserEngineOverrides,
|
||||
},
|
||||
}
|
||||
|
||||
// Use the streaming RPC so documents with many page images (large scanned
|
||||
// PDFs) are not capped by the unary message-size limit. The meta frame
|
||||
// arrives first, followed by one frame per image.
|
||||
result, err := p.readStream(ctx, client, protoReq)
|
||||
if err != nil {
|
||||
// An older docreader build may not implement ReadStream. Fall back to
|
||||
// the unary Read RPC so a version-skewed deployment still parses
|
||||
// documents (small/medium docs only — the unary path remains capped by
|
||||
// the gRPC message-size limit, which is exactly what streaming avoids).
|
||||
if status.Code(err) == codes.Unimplemented {
|
||||
logger.Warnf(ctx, "docreader ReadStream unimplemented, falling back to unary Read: %v", err)
|
||||
return p.readUnary(ctx, client, protoReq)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// readStream consumes the server-streaming ReadStream RPC: one meta frame
|
||||
// followed by one frame per image. Errors are returned verbatim so the caller
|
||||
// can inspect the gRPC status code (e.g. Unimplemented) for fallback.
|
||||
func (p *GRPCDocumentReader) readStream(
|
||||
ctx context.Context, client proto.DocReaderClient, protoReq *proto.ReadRequest,
|
||||
) (*types.ReadResult, error) {
|
||||
stream, err := client.ReadStream(ctx, protoReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gRPC ReadStream failed: %w", err)
|
||||
}
|
||||
|
||||
result := &types.ReadResult{}
|
||||
gotMeta := false
|
||||
for {
|
||||
frame, recvErr := stream.Recv()
|
||||
if recvErr == io.EOF {
|
||||
break
|
||||
}
|
||||
if recvErr != nil {
|
||||
return nil, fmt.Errorf("gRPC ReadStream recv failed: %w", recvErr)
|
||||
}
|
||||
|
||||
if meta := frame.GetMeta(); meta != nil {
|
||||
gotMeta = true
|
||||
result.MarkdownContent = meta.GetMarkdownContent()
|
||||
result.ImageDirPath = meta.GetImageDirPath()
|
||||
result.Metadata = meta.GetMetadata()
|
||||
result.Error = meta.GetError()
|
||||
if n := meta.GetImageCount(); n > 0 {
|
||||
result.ImageRefs = make([]types.ImageRef, 0, n)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if img := frame.GetImage(); img != nil {
|
||||
result.ImageRefs = append(result.ImageRefs, types.ImageRef{
|
||||
Filename: img.GetFilename(),
|
||||
OriginalRef: img.GetOriginalRef(),
|
||||
MimeType: img.GetMimeType(),
|
||||
StorageKey: img.GetStorageKey(),
|
||||
ImageData: img.GetImageData(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if !gotMeta {
|
||||
return nil, fmt.Errorf("gRPC ReadStream returned no metadata frame")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// readUnary calls the legacy unary Read RPC. Used only as a compatibility
|
||||
// fallback when the connected docreader does not implement ReadStream.
|
||||
func (p *GRPCDocumentReader) readUnary(
|
||||
ctx context.Context, client proto.DocReaderClient, protoReq *proto.ReadRequest,
|
||||
) (*types.ReadResult, error) {
|
||||
resp, err := client.Read(ctx, protoReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gRPC Read failed: %w", err)
|
||||
}
|
||||
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: resp.GetMarkdownContent(),
|
||||
ImageDirPath: resp.GetImageDirPath(),
|
||||
Metadata: resp.GetMetadata(),
|
||||
Error: resp.GetError(),
|
||||
}
|
||||
if refs := resp.GetImageRefs(); len(refs) > 0 {
|
||||
result.ImageRefs = make([]types.ImageRef, 0, len(refs))
|
||||
for _, img := range refs {
|
||||
result.ImageRefs = append(result.ImageRefs, types.ImageRef{
|
||||
Filename: img.GetFilename(),
|
||||
OriginalRef: img.GetOriginalRef(),
|
||||
MimeType: img.GetMimeType(),
|
||||
StorageKey: img.GetStorageKey(),
|
||||
ImageData: img.GetImageData(),
|
||||
})
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (p *GRPCDocumentReader) ListEngines(ctx context.Context, overrides map[string]string) ([]types.ParserEngineInfo, error) {
|
||||
p.mu.RLock()
|
||||
client := p.client
|
||||
p.mu.RUnlock()
|
||||
if client == nil {
|
||||
return nil, errNotConnected
|
||||
}
|
||||
|
||||
resp, err := client.ListEngines(ctx, &proto.ListEnginesRequest{ConfigOverrides: overrides})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("gRPC ListEngines failed: %w", err)
|
||||
}
|
||||
|
||||
result := make([]types.ParserEngineInfo, 0, len(resp.GetEngines()))
|
||||
for _, e := range resp.GetEngines() {
|
||||
result = append(result, types.ParserEngineInfo{
|
||||
Name: e.GetName(),
|
||||
Description: e.GetDescription(),
|
||||
FileTypes: e.GetFileTypes(),
|
||||
Available: e.GetAvailable(),
|
||||
UnavailableReason: e.GetUnavailableReason(),
|
||||
})
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
)
|
||||
|
||||
// stringOr returns val (trimmed) if non-empty, otherwise fallback.
|
||||
func stringOr(val, fallback string) string {
|
||||
val = strings.TrimSpace(val)
|
||||
if val == "" {
|
||||
return fallback
|
||||
}
|
||||
return val
|
||||
}
|
||||
|
||||
// parseBoolOr parses a truthy string ("true","1","yes"), returning fallback on empty.
|
||||
func parseBoolOr(val string, fallback bool) bool {
|
||||
val = strings.ToLower(strings.TrimSpace(val))
|
||||
if val == "" {
|
||||
return fallback
|
||||
}
|
||||
return val == "true" || val == "1" || val == "yes"
|
||||
}
|
||||
|
||||
// firstNonEmpty returns the first non-empty string, or "" if all are empty.
|
||||
func firstNonEmpty(vals ...string) string {
|
||||
for _, v := range vals {
|
||||
if v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// sleepCtx sleeps for d but returns early if ctx is cancelled.
|
||||
func sleepCtx(ctx context.Context, d time.Duration) {
|
||||
t := time.NewTimer(d)
|
||||
defer t.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-t.C:
|
||||
}
|
||||
}
|
||||
|
||||
// logResponseStructure recursively logs the structure of an API response,
|
||||
// truncating large string values. label identifies the subsystem (e.g. "MinerU").
|
||||
func logResponseStructure(label string, obj interface{}, prefix string) {
|
||||
switch v := obj.(type) {
|
||||
case map[string]interface{}:
|
||||
keys := make([]string, 0, len(v))
|
||||
for k := range v {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
logger.Infof(context.Background(), "[%s] %s = {object with %d keys: %s}", label, prefix, len(v), strings.Join(keys, ", "))
|
||||
for _, key := range keys {
|
||||
val := v[key]
|
||||
path := prefix + "." + key
|
||||
switch inner := val.(type) {
|
||||
case map[string]interface{}:
|
||||
logResponseStructure(label, inner, path)
|
||||
case []interface{}:
|
||||
logger.Infof(context.Background(), "[%s] %s = [array with %d items]", label, path, len(inner))
|
||||
if len(inner) > 0 {
|
||||
logger.Infof(context.Background(), "[%s] %s[0] type=%T", label, path, inner[0])
|
||||
if len(inner) <= 3 {
|
||||
for i, item := range inner {
|
||||
logResponseStructure(label, item, fmt.Sprintf("%s[%d]", path, i))
|
||||
}
|
||||
} else {
|
||||
logResponseStructure(label, inner[0], path+"[0]")
|
||||
logger.Infof(context.Background(), "[%s] ... and %d more items in %s", label, len(inner)-1, path)
|
||||
}
|
||||
}
|
||||
case string:
|
||||
if len(inner) > 200 {
|
||||
logger.Infof(context.Background(), "[%s] %s = string(%d chars): %.200s...", label, path, len(inner), inner)
|
||||
} else {
|
||||
logger.Infof(context.Background(), "[%s] %s = %q", label, path, inner)
|
||||
}
|
||||
case float64:
|
||||
logger.Infof(context.Background(), "[%s] %s = %v (number)", label, path, inner)
|
||||
case bool:
|
||||
logger.Infof(context.Background(), "[%s] %s = %v (bool)", label, path, inner)
|
||||
case nil:
|
||||
logger.Infof(context.Background(), "[%s] %s = null", label, path)
|
||||
default:
|
||||
logger.Infof(context.Background(), "[%s] %s = %v (%T)", label, path, val, val)
|
||||
}
|
||||
}
|
||||
case []interface{}:
|
||||
logger.Infof(context.Background(), "[%s] %s = [array with %d items]", label, prefix, len(v))
|
||||
if len(v) > 0 {
|
||||
if len(v) <= 3 {
|
||||
for i, item := range v {
|
||||
logResponseStructure(label, item, fmt.Sprintf("%s[%d]", prefix, i))
|
||||
}
|
||||
} else {
|
||||
logResponseStructure(label, v[0], prefix+"[0]")
|
||||
logger.Infof(context.Background(), "[%s] ... and %d more items in %s", label, len(v)-1, prefix)
|
||||
}
|
||||
}
|
||||
default:
|
||||
logger.Infof(context.Background(), "[%s] %s = %v (%T)", label, prefix, v, v)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
"github.com/JohannesKaufmann/html-to-markdown/v2/converter"
|
||||
"github.com/JohannesKaufmann/html-to-markdown/v2/plugin/base"
|
||||
"github.com/JohannesKaufmann/html-to-markdown/v2/plugin/commonmark"
|
||||
"github.com/JohannesKaufmann/html-to-markdown/v2/plugin/table"
|
||||
)
|
||||
|
||||
var (
|
||||
// htmlTableBlockPattern matches a single (non-nested) <table>...</table>
|
||||
// block, the form OCR/layout engines such as PaddleOCR-VL emit tables in.
|
||||
htmlTableBlockPattern = regexp.MustCompile(`(?is)<table\b[^>]*>.*?</table>`)
|
||||
|
||||
// htmlLayoutAttrPattern matches presentational HTML attributes that carry
|
||||
// no semantic value (text-align styles, CSS classes, sizing). Structural
|
||||
// attributes like rowspan/colspan are intentionally excluded.
|
||||
htmlLayoutAttrPattern = regexp.MustCompile(
|
||||
`(?is)\s+(?:style|class|align|valign|width|height|bgcolor)\s*=\s*(?:"[^"]*"|'[^']*'|[^\s>]+)`,
|
||||
)
|
||||
|
||||
// htmlSpanAttrPattern detects rowspan/colspan, which Markdown tables cannot
|
||||
// represent; such tables keep their HTML form (attributes stripped) instead.
|
||||
htmlSpanAttrPattern = regexp.MustCompile(`(?i)\b(?:row|col)span\b`)
|
||||
|
||||
// markdownTableSeparatorPattern matches the |---|---| delimiter row that a
|
||||
// valid GFM table must contain.
|
||||
markdownTableSeparatorPattern = regexp.MustCompile(`(?m)^\s*\|?\s*:?-+:?\s*(?:\|\s*:?-+:?\s*)+\|?\s*$`)
|
||||
)
|
||||
|
||||
// normalizeHTMLTables rewrites inline HTML <table> blocks embedded in OCR
|
||||
// markdown output. PaddleOCR-VL emits tables as HTML with per-cell text-align
|
||||
// styles, which (1) waste tokens on layout markup and (2) are not recognized
|
||||
// by the chunker's table-protection logic, so large tables get split mid-row.
|
||||
//
|
||||
// Each table block is converted to a GFM Markdown table when possible. Tables
|
||||
// that use rowspan/colspan (which Markdown cannot express) fall back to having
|
||||
// their presentational attributes stripped so they stay intact as HTML.
|
||||
func normalizeHTMLTables(md string) string {
|
||||
if !strings.Contains(strings.ToLower(md), "<table") {
|
||||
return md
|
||||
}
|
||||
|
||||
conv := converter.NewConverter(
|
||||
converter.WithPlugins(
|
||||
base.NewBasePlugin(),
|
||||
commonmark.NewCommonmarkPlugin(),
|
||||
table.NewTablePlugin(),
|
||||
),
|
||||
)
|
||||
|
||||
return htmlTableBlockPattern.ReplaceAllStringFunc(md, func(block string) string {
|
||||
if htmlSpanAttrPattern.MatchString(block) {
|
||||
return stripHTMLLayoutAttrs(block)
|
||||
}
|
||||
converted, err := conv.ConvertString(block)
|
||||
if err != nil {
|
||||
return stripHTMLLayoutAttrs(block)
|
||||
}
|
||||
converted = strings.TrimSpace(converted)
|
||||
if converted == "" || !markdownTableSeparatorPattern.MatchString(converted) {
|
||||
return stripHTMLLayoutAttrs(block)
|
||||
}
|
||||
// Pad with blank lines so the Markdown table is a standalone block that
|
||||
// the chunker recognizes (and protects) as a table.
|
||||
return "\n\n" + converted + "\n\n"
|
||||
})
|
||||
}
|
||||
|
||||
// stripHTMLLayoutAttrs removes presentational attributes from an HTML fragment
|
||||
// while preserving structural attributes (rowspan/colspan) and text content.
|
||||
func stripHTMLLayoutAttrs(html string) string {
|
||||
return htmlLayoutAttrPattern.ReplaceAllString(html, "")
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNormalizeHTMLTables_ConvertsStyledTableToMarkdown(t *testing.T) {
|
||||
// Mirrors PaddleOCR-VL output: a data table where every cell carries a
|
||||
// text-align style that wastes tokens.
|
||||
input := `# 报告
|
||||
|
||||
<table><tr><td style="text-align:center;">指标</td><td style="text-align:center;">数值</td></tr>` +
|
||||
`<tr><td style="text-align:center;">营收</td><td style="text-align:right;">10亿</td></tr>` +
|
||||
`<tr><td style="text-align:center;">利润</td><td style="text-align:right;">2.3亿</td></tr></table>
|
||||
|
||||
结尾。`
|
||||
|
||||
got := normalizeHTMLTables(input)
|
||||
|
||||
if strings.Contains(got, "<table") {
|
||||
t.Fatalf("expected HTML table to be converted away, got:\n%s", got)
|
||||
}
|
||||
if strings.Contains(got, "text-align") {
|
||||
t.Fatalf("expected style attributes removed, got:\n%s", got)
|
||||
}
|
||||
if !markdownTableSeparatorPattern.MatchString(got) {
|
||||
t.Fatalf("expected a Markdown table separator row, got:\n%s", got)
|
||||
}
|
||||
for _, want := range []string{"指标", "数值", "营收", "10亿", "利润", "2.3亿"} {
|
||||
if !strings.Contains(got, want) {
|
||||
t.Fatalf("expected converted table to retain %q, got:\n%s", want, got)
|
||||
}
|
||||
}
|
||||
if !strings.Contains(got, "# 报告") || !strings.Contains(got, "结尾。") {
|
||||
t.Fatalf("expected surrounding markdown to be preserved, got:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeHTMLTables_StripsAttrsOnSpanTables(t *testing.T) {
|
||||
// rowspan/colspan cannot be expressed in Markdown, so the table stays HTML
|
||||
// but its presentational attributes are stripped.
|
||||
input := `<table><tr><td colspan="2" style="text-align:center;" class="hdr">合计</td></tr>` +
|
||||
`<tr><td style="text-align:left;">A</td><td width="80">B</td></tr></table>`
|
||||
|
||||
got := normalizeHTMLTables(input)
|
||||
|
||||
if !strings.Contains(got, "<table") {
|
||||
t.Fatalf("expected span table to remain HTML, got:\n%s", got)
|
||||
}
|
||||
if !strings.Contains(got, `colspan="2"`) {
|
||||
t.Fatalf("expected colspan to be preserved, got:\n%s", got)
|
||||
}
|
||||
for _, banned := range []string{"text-align", "class=", "width="} {
|
||||
if strings.Contains(got, banned) {
|
||||
t.Fatalf("expected %q to be stripped, got:\n%s", banned, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeHTMLTables_NoTableUnchanged(t *testing.T) {
|
||||
input := "# 标题\n\n普通段落,没有表格。\n\n| a | b |\n| --- | --- |\n| 1 | 2 |"
|
||||
if got := normalizeHTMLTables(input); got != input {
|
||||
t.Fatalf("expected content without HTML tables to be unchanged, got:\n%s", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
const (
|
||||
PathRead = "/read"
|
||||
PathListEngines = "/list-engines"
|
||||
)
|
||||
|
||||
// --- JSON DTOs ---
|
||||
|
||||
type httpReadConfig struct {
|
||||
ParserEngine string `json:"parser_engine,omitempty"`
|
||||
ParserEngineOverrides map[string]string `json:"parser_engine_overrides,omitempty"`
|
||||
}
|
||||
|
||||
type httpReadRequest struct {
|
||||
FileContent string `json:"file_content,omitempty"` // base64
|
||||
FileName string `json:"file_name,omitempty"`
|
||||
FileType string `json:"file_type,omitempty"`
|
||||
URL string `json:"url,omitempty"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Config *httpReadConfig `json:"config,omitempty"`
|
||||
RequestID string `json:"request_id,omitempty"`
|
||||
}
|
||||
|
||||
type httpImageRef struct {
|
||||
Filename string `json:"filename"`
|
||||
OriginalRef string `json:"original_ref"`
|
||||
MimeType string `json:"mime_type"`
|
||||
StorageKey string `json:"storage_key,omitempty"`
|
||||
ImageData []byte `json:"image_data,omitempty"`
|
||||
}
|
||||
|
||||
type httpReadResponse struct {
|
||||
MarkdownContent string `json:"markdown_content"`
|
||||
ImageRefs []httpImageRef `json:"image_refs,omitempty"`
|
||||
ImageDirPath string `json:"image_dir_path,omitempty"`
|
||||
Metadata map[string]string `json:"metadata,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// HTTPDocumentReader implements DocumentReader over HTTP/JSON.
|
||||
type HTTPDocumentReader struct {
|
||||
mu sync.RWMutex
|
||||
baseURL string
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
func NewHTTPDocumentReader(baseURL string) (*HTTPDocumentReader, error) {
|
||||
baseURL = strings.TrimSuffix(strings.TrimSpace(baseURL), "/")
|
||||
if baseURL != "" {
|
||||
if err := secutils.ValidateURLForSSRF(baseURL); err != nil {
|
||||
return nil, fmt.Errorf("docreader address failed SSRF validation: %w", err)
|
||||
}
|
||||
}
|
||||
clientCfg := secutils.DefaultSSRFSafeHTTPClientConfig()
|
||||
clientCfg.Timeout = 5 * time.Minute
|
||||
p := &HTTPDocumentReader{
|
||||
baseURL: baseURL,
|
||||
client: secutils.NewSSRFSafeHTTPClient(clientCfg),
|
||||
}
|
||||
if p.baseURL != "" {
|
||||
logger.Infof(context.Background(), "INFO: HTTP docreader base URL: %s", p.baseURL)
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) base() string {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
return p.baseURL
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) Reconnect(addr string) error {
|
||||
addr = strings.TrimSuffix(strings.TrimSpace(addr), "/")
|
||||
if addr != "" {
|
||||
if err := secutils.ValidateURLForSSRF(addr); err != nil {
|
||||
return fmt.Errorf("docreader address failed SSRF validation: %w", err)
|
||||
}
|
||||
}
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
p.baseURL = addr
|
||||
logger.Infof(context.Background(), "INFO: HTTP docreader base URL set to %s", p.baseURL)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) IsConnected() bool {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
return p.baseURL != ""
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) Close() error { return nil }
|
||||
|
||||
type httpListEnginesRequest struct {
|
||||
ConfigOverrides map[string]string `json:"config_overrides,omitempty"`
|
||||
}
|
||||
|
||||
type httpParserEngineInfo struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
FileTypes []string `json:"file_types"`
|
||||
Available bool `json:"available"`
|
||||
UnavailableReason string `json:"unavailable_reason,omitempty"`
|
||||
}
|
||||
|
||||
type httpListEnginesResponse struct {
|
||||
Engines []httpParserEngineInfo `json:"engines"`
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) ListEngines(ctx context.Context, overrides map[string]string) ([]types.ParserEngineInfo, error) {
|
||||
base := p.base()
|
||||
if base == "" {
|
||||
return nil, errNotConnected
|
||||
}
|
||||
if err := secutils.ValidateURLForSSRF(base); err != nil {
|
||||
return nil, fmt.Errorf("docreader address failed SSRF validation: %w", err)
|
||||
}
|
||||
|
||||
body := httpListEnginesRequest{ConfigOverrides: overrides}
|
||||
jsonBody, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http marshal list-engines request: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, base+PathListEngines, bytes.NewReader(jsonBody))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http new request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := p.client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http list-engines failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
respBytes, _ := io.ReadAll(resp.Body)
|
||||
return nil, fmt.Errorf("http list-engines status %d: %s", resp.StatusCode, string(respBytes))
|
||||
}
|
||||
|
||||
var out httpListEnginesResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
||||
return nil, fmt.Errorf("http decode list-engines response: %w", err)
|
||||
}
|
||||
|
||||
result := make([]types.ParserEngineInfo, 0, len(out.Engines))
|
||||
for _, e := range out.Engines {
|
||||
result = append(result, types.ParserEngineInfo{
|
||||
Name: e.Name,
|
||||
Description: e.Description,
|
||||
FileTypes: e.FileTypes,
|
||||
Available: e.Available,
|
||||
UnavailableReason: e.UnavailableReason,
|
||||
})
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func fromHTTPReadResponse(resp *httpReadResponse) *types.ReadResult {
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: resp.MarkdownContent,
|
||||
ImageDirPath: resp.ImageDirPath,
|
||||
Metadata: resp.Metadata,
|
||||
Error: resp.Error,
|
||||
}
|
||||
for _, ref := range resp.ImageRefs {
|
||||
result.ImageRefs = append(result.ImageRefs, types.ImageRef{
|
||||
Filename: ref.Filename,
|
||||
OriginalRef: ref.OriginalRef,
|
||||
MimeType: ref.MimeType,
|
||||
StorageKey: ref.StorageKey,
|
||||
ImageData: ref.ImageData,
|
||||
})
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (p *HTTPDocumentReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
base := p.base()
|
||||
if base == "" {
|
||||
return nil, errNotConnected
|
||||
}
|
||||
if err := secutils.ValidateURLForSSRF(base); err != nil {
|
||||
return nil, fmt.Errorf("docreader address failed SSRF validation: %w", err)
|
||||
}
|
||||
|
||||
body := httpReadRequest{
|
||||
FileName: req.FileName,
|
||||
FileType: req.FileType,
|
||||
URL: req.URL,
|
||||
Title: req.Title,
|
||||
RequestID: req.RequestID,
|
||||
Config: &httpReadConfig{
|
||||
ParserEngine: req.ParserEngine,
|
||||
ParserEngineOverrides: req.ParserEngineOverrides,
|
||||
},
|
||||
}
|
||||
if len(req.FileContent) > 0 {
|
||||
body.FileContent = base64.StdEncoding.EncodeToString(req.FileContent)
|
||||
}
|
||||
|
||||
jsonBody, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http marshal read request: %w", err)
|
||||
}
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, base+PathRead, bytes.NewReader(jsonBody))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http new request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.ContentLength = int64(len(jsonBody))
|
||||
|
||||
resp, err := p.client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("http read failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
return nil, fmt.Errorf("http read status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
var out httpReadResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
||||
return nil, fmt.Errorf("http decode read response: %w", err)
|
||||
}
|
||||
return fromHTTPReadResponse(&out), nil
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,80 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestResolveAndStoreRelativeHTMLImages(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: `<table><tr><td><img src="images/profile.jpg" alt="profile"/></td></tr></table>`,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "profile.jpg",
|
||||
OriginalRef: "images/profile.jpg",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 stored image but got %d", len(imgs))
|
||||
}
|
||||
if len(svc.saved) != 1 {
|
||||
t.Fatalf("expected SaveBytes to be called once but got %d", len(svc.saved))
|
||||
}
|
||||
if strings.Contains(out, `src="images/profile.jpg"`) {
|
||||
t.Fatalf("expected relative html img src to be replaced, got: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `src="local://test/`) {
|
||||
t.Fatalf("expected stored file url in html img src, got: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveAndStoreDedupsSameImageRefVariants ensures that when the same
|
||||
// MinerU image is referenced under multiple path forms (e.g. "images/foo.png"
|
||||
// and "./images/foo.png"), the bytes are uploaded to object storage only once.
|
||||
func TestResolveAndStoreDedupsSameImageRefVariants(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: strings.Join([]string{
|
||||
"",
|
||||
`<img src="./images/foo.png"/>`,
|
||||
}, "\n"),
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "foo.png",
|
||||
OriginalRef: "images/foo.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
{
|
||||
Filename: "foo.png",
|
||||
OriginalRef: "./images/foo.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
_, _, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(svc.saved) != 1 {
|
||||
t.Fatalf("expected SaveBytes to be called once for duplicate refs, got %d", len(svc.saved))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,600 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"image/png"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// createTestPNG generates a minimal PNG image with the given dimensions.
|
||||
func createTestPNG(w, h int) []byte {
|
||||
img := image.NewRGBA(image.Rect(0, 0, w, h))
|
||||
for y := 0; y < h; y++ {
|
||||
for x := 0; x < w; x++ {
|
||||
img.Set(x, y, color.RGBA{R: 128, G: 128, B: 128, A: 255})
|
||||
}
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
_ = png.Encode(&buf, img)
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func TestIsIconImage(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data []byte
|
||||
expect bool
|
||||
}{
|
||||
{
|
||||
name: "tiny bytes (< 512B)",
|
||||
data: make([]byte, 256),
|
||||
expect: true,
|
||||
},
|
||||
{
|
||||
name: "small icon 32x32",
|
||||
data: createTestPNG(32, 32),
|
||||
expect: true,
|
||||
},
|
||||
{
|
||||
name: "small icon 48x48",
|
||||
data: createTestPNG(48, 48),
|
||||
expect: true,
|
||||
},
|
||||
{
|
||||
name: "borderline 64x64",
|
||||
data: createTestPNG(64, 64),
|
||||
expect: false,
|
||||
},
|
||||
{
|
||||
name: "normal image 200x150",
|
||||
data: createTestPNG(200, 150),
|
||||
expect: false,
|
||||
},
|
||||
{
|
||||
name: "wide but short 200x30",
|
||||
data: createTestPNG(200, 30),
|
||||
expect: false,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := isIconImage(tt.data)
|
||||
if got != tt.expect {
|
||||
t.Errorf("isIconImage() = %v, want %v (data len=%d)", got, tt.expect, len(tt.data))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// captureSaveBytes implements interfaces.FileService for tests; only SaveBytes records data.
|
||||
type captureSaveBytes struct {
|
||||
saved [][]byte
|
||||
urls []string
|
||||
}
|
||||
|
||||
func (c *captureSaveBytes) CheckConnectivity(context.Context) error { return nil }
|
||||
|
||||
func (c *captureSaveBytes) SaveFile(context.Context, *multipart.FileHeader, uint64, string) (string, error) {
|
||||
return "", fmt.Errorf("not implemented")
|
||||
}
|
||||
|
||||
func (c *captureSaveBytes) SaveBytes(_ context.Context, data []byte, _ uint64, fileName string, _ bool) (string, error) {
|
||||
c.saved = append(c.saved, append([]byte(nil), data...))
|
||||
u := "local://test/" + fileName
|
||||
c.urls = append(c.urls, u)
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func (c *captureSaveBytes) GetFile(context.Context, string) (io.ReadCloser, error) {
|
||||
return nil, fmt.Errorf("not implemented")
|
||||
}
|
||||
|
||||
func (c *captureSaveBytes) GetFileURL(context.Context, string) (string, error) { return "", nil }
|
||||
|
||||
func (c *captureSaveBytes) DeleteFile(context.Context, string) error { return nil }
|
||||
|
||||
func (c *captureSaveBytes) CopyFile(context.Context, string, uint64, string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
var _ interfaces.FileService = (*captureSaveBytes)(nil)
|
||||
|
||||
func TestResolveDataURIImages(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := "pre  post"
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveDataURIImages(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("got %d images", len(imgs))
|
||||
}
|
||||
if len(svc.saved) != 1 || !bytes.Equal(svc.saved[0], png) {
|
||||
t.Fatal("SaveBytes payload mismatch")
|
||||
}
|
||||
if !strings.Contains(out, "local://test/") || strings.Contains(out, "data:image") {
|
||||
t.Fatalf("markdown: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveDataURIImages_CaseInsensitive(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := ""
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveDataURIImages(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || !strings.Contains(out, "local://test/") {
|
||||
t.Fatalf("imgs=%d out=%q", len(imgs), out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveDataURIImages_AltTextWithBracket(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := `pre ![C:\Users\Frank\RichOle\YAJ]7D)SW_W6.png](data:image/png;base64,` + b64 + `) post`
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveDataURIImages(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image but got %d; alt text with ] was not handled", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "base64") {
|
||||
t.Fatalf("base64 content still present in output: %s", out[:min(200, len(out))])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveDataURIImages_XEmfMimeType(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := ``
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveDataURIImages(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image for x-emf mime type but got %d", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "data:image") {
|
||||
t.Fatalf("data URI still present for x-emf: %s", out[:min(200, len(out))])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveHTMLDataURIImages(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := `text <img src="data:image/png;base64,` + b64 + `" alt="test"> more`
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveHTMLDataURIImages(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image from HTML img tag but got %d", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "<img") || strings.Contains(out, "base64") {
|
||||
t.Fatalf("HTML img tag or base64 still present: %s", out[:min(200, len(out))])
|
||||
}
|
||||
if !strings.Contains(out, " {
|
||||
t.Fatalf("expected markdown image replacement, got: %s", out[:min(200, len(out))])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveBareBase64Content_BareDataURI(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
md := `some text data:image/png;base64,` + b64 + ` more text`
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveBareBase64Content(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image from bare data URI but got %d", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "base64") {
|
||||
t.Fatalf("base64 content still present: %s", out[:min(200, len(out))])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveBareBase64Content_InsideBrokenMarkdownRef(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
// Simulate a broken markdown ref that the primary regex already handled
|
||||
// If primary regex handles it, this won't reach the bare handler
|
||||
// But test the bare handler's behavior when prev char is '('
|
||||
md := `text (data:image/png;base64,` + b64 + `) more`
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveBareBase64Content(context.Background(), md, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image but got %d", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "base64") {
|
||||
t.Fatalf("base64 still present: %s", out[:min(200, len(out))])
|
||||
}
|
||||
// When preceded by '(', should replace with just URL (not wrapped in ![image]())
|
||||
if strings.Contains(out, "![image]") {
|
||||
t.Fatalf("should not wrap in ![image]() when inside parentheses: %s", out[:min(200, len(out))])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStore_MultipleFormats(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
// Markdown with: standard data URI, alt text with ], x-emf mime type
|
||||
md := strings.Join([]string{
|
||||
`# Document`,
|
||||
``,
|
||||
`![C:\path\file]name.png](data:image/png;base64,` + b64 + `)`,
|
||||
``,
|
||||
}, "\n\n")
|
||||
|
||||
result := &types.ReadResult{MarkdownContent: md}
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 3 {
|
||||
t.Fatalf("expected 3 images but got %d", len(imgs))
|
||||
}
|
||||
if strings.Contains(out, "base64") {
|
||||
t.Fatalf("base64 content still in output after ResolveAndStore")
|
||||
}
|
||||
if strings.Contains(out, "data:image") {
|
||||
t.Fatalf("data URI still in output after ResolveAndStore")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageWithTitle(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: ``,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected 1 image but got %d", len(imgs))
|
||||
}
|
||||
if len(svc.saved) != 1 {
|
||||
t.Fatalf("expected SaveBytes to be called once but got %d", len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, ` || !strings.Contains(out, ` "图片")`) {
|
||||
t.Fatalf("markdown image title was not preserved around stored URL: %s", out)
|
||||
}
|
||||
if strings.Contains(out, "images/test.png") {
|
||||
t.Fatalf("original image path was not replaced: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageWithSingleQuotedTitle(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: ``,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || len(svc.saved) != 1 {
|
||||
t.Fatalf("expected one stored image, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, ` || !strings.Contains(out, ` '图片')`) {
|
||||
t.Fatalf("markdown image title was not preserved around stored URL: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageWithSpacedFilename(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: ``,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "第 1 页.png",
|
||||
OriginalRef: "images/第 1 页.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || len(svc.saved) != 1 {
|
||||
t.Fatalf("expected one stored image, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, ` {
|
||||
t.Fatalf("spaced filename path was not replaced: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageTitleContainingRightParen(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: ` 图片")`,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || len(svc.saved) != 1 {
|
||||
t.Fatalf("expected one stored image, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, ` || !strings.Contains(out, ` "阶段 1) 图片")`) {
|
||||
t.Fatalf("right-paren title was not preserved around stored URL: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageWithMultilineTitle(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: "",
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || len(svc.saved) != 1 {
|
||||
t.Fatalf("expected one stored image, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, " || !strings.Contains(out, "\n \"图片说明\")") {
|
||||
t.Fatalf("multiline title was not preserved around stored URL: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreMarkdownImageWithAngleDestination(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: `.png> "阶段 1) 图片")`,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "第 1 页 (测试).png",
|
||||
OriginalRef: "images/第 1 页 (测试).png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(imgs) != 1 || len(svc.saved) != 1 {
|
||||
t.Fatalf("expected one stored image, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if !strings.Contains(out, ` || !strings.Contains(out, `> "阶段 1) 图片")`) {
|
||||
t.Fatalf("angle destination wrapper or title was not preserved: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreDuplicateReferencesStoreOnce(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: strings.Join([]string{
|
||||
``,
|
||||
``,
|
||||
}, "\n\n"),
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(svc.saved) != 1 {
|
||||
t.Fatalf("expected duplicate references to save once but got %d", len(svc.saved))
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected one stored image record but got %d", len(imgs))
|
||||
}
|
||||
if strings.Count(out, "local://test/") != 2 {
|
||||
t.Fatalf("expected both markdown references to be replaced: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `"第一处"`) || !strings.Contains(out, `"第二处"`) {
|
||||
t.Fatalf("duplicate reference titles were not preserved: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreUnknownReferenceUnchanged(t *testing.T) {
|
||||
png := createTestPNG(200, 150)
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: ``,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: "test.png",
|
||||
OriginalRef: "images/test.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: png,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(svc.saved) != 0 || len(imgs) != 0 {
|
||||
t.Fatalf("unknown reference should not be saved, got imgs=%d saved=%d", len(imgs), len(svc.saved))
|
||||
}
|
||||
if out != result.MarkdownContent {
|
||||
t.Fatalf("unknown reference changed: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreSharedMHTMLContract(t *testing.T) {
|
||||
type contractImage struct {
|
||||
Filename string `json:"filename"`
|
||||
OriginalRef string `json:"original_ref"`
|
||||
MimeType string `json:"mime_type"`
|
||||
ImageDataBase64 string `json:"image_data_base64"`
|
||||
}
|
||||
var contract struct {
|
||||
MarkdownContent string `json:"markdown_content"`
|
||||
Images []contractImage `json:"images"`
|
||||
}
|
||||
|
||||
contractPath := filepath.Join("..", "..", "..", "testdata", "mhtml", "titled-image-contract.json")
|
||||
raw, err := os.ReadFile(contractPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := json.Unmarshal(raw, &contract); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(contract.Images) != 1 {
|
||||
t.Fatalf("expected one contract image but got %d", len(contract.Images))
|
||||
}
|
||||
|
||||
imageContract := contract.Images[0]
|
||||
imageData, err := base64.StdEncoding.DecodeString(imageContract.ImageDataBase64)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
result := &types.ReadResult{
|
||||
MarkdownContent: contract.MarkdownContent,
|
||||
ImageRefs: []types.ImageRef{
|
||||
{
|
||||
Filename: imageContract.Filename,
|
||||
OriginalRef: imageContract.OriginalRef,
|
||||
MimeType: imageContract.MimeType,
|
||||
ImageData: imageData,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
svc := &captureSaveBytes{}
|
||||
r := NewImageResolver()
|
||||
out, imgs, err := r.ResolveAndStore(context.Background(), result, svc, 1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(svc.saved) != 1 {
|
||||
t.Fatalf("expected SaveBytes once but got %d", len(svc.saved))
|
||||
}
|
||||
if !bytes.Equal(svc.saved[0], imageData) {
|
||||
t.Fatal("stored image bytes do not match contract fixture")
|
||||
}
|
||||
if len(imgs) != 1 {
|
||||
t.Fatalf("expected one stored image record but got %d", len(imgs))
|
||||
}
|
||||
if imgs[0].OriginalRef != imageContract.OriginalRef || imgs[0].MimeType != imageContract.MimeType {
|
||||
t.Fatalf("unexpected stored image metadata: %#v", imgs[0])
|
||||
}
|
||||
if strings.Contains(out, imageContract.OriginalRef) {
|
||||
t.Fatalf("original image path remained in output: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, ` {
|
||||
t.Fatalf("stored URL missing from markdown: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `"阶段 1) 图片"`) {
|
||||
t.Fatalf("image title was not preserved: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, "| 赛季制建立 | BP、Rank |\n\n {
|
||||
t.Fatalf("table-to-image block boundary changed: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `"阶段 1) 图片")`+"\n\n高机动性身法与独特枪械反馈") {
|
||||
t.Fatalf("image-to-body block boundary changed: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, "alpha \nbeta") {
|
||||
t.Fatalf("markdown hard break changed: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, "```\nline1\n\n\nline2\n```") {
|
||||
t.Fatalf("fenced code blank lines changed: %s", out)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// defaultJSONChunkSize is the target chunk size in bytes for JSON semantic
|
||||
// splitting. Approximates 512 tokens (1 token ≈ 3-4 bytes for mixed content).
|
||||
const defaultJSONChunkSize = 1536
|
||||
|
||||
// minJSONChunkSize is the minimum chunk size. A new chunk is only started
|
||||
// when the current chunk has reached at least this size.
|
||||
var minJSONChunkSize = defaultJSONChunkSize - 200
|
||||
|
||||
// jsonToMarkdown converts raw JSON bytes into markdown text
|
||||
//
|
||||
// Key properties:
|
||||
// - Every output chunk is a **valid JSON object** (not a fragment).
|
||||
// - Nested paths from root to leaf are **fully preserved** in each chunk.
|
||||
// - Arrays are converted to index-keyed dicts so the algorithm is uniform.
|
||||
// - Small objects that fit within maxChunkSize are kept intact (not split).
|
||||
// - The output is a series of fenced ```json code blocks separated by \n\n,
|
||||
// which the downstream text chunker can split at block boundaries.
|
||||
func jsonToMarkdown(data []byte) (string, error) {
|
||||
data = trimBOM(data)
|
||||
if len(data) == 0 {
|
||||
return "", fmt.Errorf("empty JSON content")
|
||||
}
|
||||
if !json.Valid(data) {
|
||||
return "", fmt.Errorf("invalid JSON content")
|
||||
}
|
||||
|
||||
var parsed interface{}
|
||||
if err := json.Unmarshal(data, &parsed); err != nil {
|
||||
return "", fmt.Errorf("failed to parse JSON: %w", err)
|
||||
}
|
||||
|
||||
// Normalize: convert top-level arrays to index-keyed dicts
|
||||
normalized := listToDictPreprocess(parsed)
|
||||
|
||||
// If the whole thing fits in one chunk, just format it
|
||||
wholeSize := jsonSize(normalized)
|
||||
if wholeSize <= defaultJSONChunkSize {
|
||||
formatted := formatValue(normalized)
|
||||
return wrapCodeBlock(formatted), nil
|
||||
}
|
||||
|
||||
// Recursive split
|
||||
chunks := recursiveJSONSplit(normalized, nil, nil)
|
||||
|
||||
// Convert each chunk dict to a fenced code block
|
||||
blocks := make([]string, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
if len(chunk) == 0 {
|
||||
continue
|
||||
}
|
||||
blocks = append(blocks, wrapCodeBlock(formatValue(chunk)))
|
||||
}
|
||||
|
||||
if len(blocks) == 0 {
|
||||
return wrapCodeBlock(formatValue(normalized)), nil
|
||||
}
|
||||
return strings.Join(blocks, "\n\n"), nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// RecursiveJsonSplitter core algorithm
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// recursiveJSONSplit splits a JSON dict into a list of JSON dicts,
|
||||
// each fitting within defaultJSONChunkSize. It preserves the full nested
|
||||
// path from root to each leaf by using setNestedDict.
|
||||
//
|
||||
// This is a Go port of LangChain's RecursiveJsonSplitter._json_split.
|
||||
func recursiveJSONSplit(
|
||||
data interface{},
|
||||
currentPath []string,
|
||||
chunks []map[string]interface{},
|
||||
) []map[string]interface{} {
|
||||
if chunks == nil {
|
||||
chunks = []map[string]interface{}{{}}
|
||||
}
|
||||
|
||||
dict, ok := data.(map[string]interface{})
|
||||
if !ok {
|
||||
// Scalar or already-processed value: place it at the current path
|
||||
if len(currentPath) > 0 && len(chunks) > 0 {
|
||||
setNestedDict(chunks[len(chunks)-1], currentPath, data)
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
|
||||
// Sort keys for deterministic output
|
||||
keys := sortedKeys(dict)
|
||||
|
||||
for _, key := range keys {
|
||||
value := dict[key]
|
||||
newPath := append(append([]string{}, currentPath...), key)
|
||||
|
||||
// Measure sizes
|
||||
chunkSize := jsonSize(chunks[len(chunks)-1])
|
||||
itemSize := jsonSize(map[string]interface{}{key: value})
|
||||
remaining := defaultJSONChunkSize - chunkSize
|
||||
|
||||
if itemSize <= remaining {
|
||||
// Item fits in the current chunk — add it preserving the path
|
||||
setNestedDict(chunks[len(chunks)-1], newPath, value)
|
||||
} else {
|
||||
// Item doesn't fit
|
||||
if chunkSize >= minJSONChunkSize {
|
||||
// Current chunk is big enough, start a new one
|
||||
chunks = append(chunks, map[string]interface{}{})
|
||||
}
|
||||
|
||||
// Check if the value itself is a dict/list that can be recursed into
|
||||
normalized := listToDictPreprocess(value)
|
||||
if subDict, isDict := normalized.(map[string]interface{}); isDict && canSplitDict(subDict) {
|
||||
// Recurse into the sub-object
|
||||
chunks = recursiveJSONSplit(subDict, newPath, chunks)
|
||||
} else {
|
||||
// Cannot split further (scalar or single-key dict) — place as-is
|
||||
setNestedDict(chunks[len(chunks)-1], newPath, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return chunks
|
||||
}
|
||||
|
||||
// setNestedDict sets a value in a nested dict structure, creating
|
||||
// intermediate dicts as needed. This preserves the full JSON path.
|
||||
//
|
||||
// Example: setNestedDict(d, ["config","db","host"], "localhost")
|
||||
// produces: {"config": {"db": {"host": "localhost"}}}
|
||||
func setNestedDict(d map[string]interface{}, path []string, value interface{}) {
|
||||
if len(path) == 0 {
|
||||
return
|
||||
}
|
||||
current := d
|
||||
for _, key := range path[:len(path)-1] {
|
||||
next, ok := current[key]
|
||||
if !ok {
|
||||
next = map[string]interface{}{}
|
||||
current[key] = next
|
||||
}
|
||||
if nextDict, ok := next.(map[string]interface{}); ok {
|
||||
current = nextDict
|
||||
} else {
|
||||
// Path conflict (existing value is not a dict) — overwrite
|
||||
newDict := map[string]interface{}{}
|
||||
current[key] = newDict
|
||||
current = newDict
|
||||
}
|
||||
}
|
||||
current[path[len(path)-1]] = value
|
||||
}
|
||||
|
||||
// listToDictPreprocess recursively converts JSON arrays to index-keyed
|
||||
// dicts so the splitter can treat everything uniformly.
|
||||
//
|
||||
// Example: ["a","b","c"] → {"0":"a", "1":"b", "2":"c"}
|
||||
func listToDictPreprocess(data interface{}) interface{} {
|
||||
switch v := data.(type) {
|
||||
case map[string]interface{}:
|
||||
result := make(map[string]interface{}, len(v))
|
||||
for k, val := range v {
|
||||
result[k] = listToDictPreprocess(val)
|
||||
}
|
||||
return result
|
||||
case []interface{}:
|
||||
result := make(map[string]interface{}, len(v))
|
||||
for i, item := range v {
|
||||
result[fmt.Sprintf("%d", i)] = listToDictPreprocess(item)
|
||||
}
|
||||
return result
|
||||
default:
|
||||
return data
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// jsonSize returns the serialized JSON byte length of a value.
|
||||
func jsonSize(v interface{}) int {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return len(b)
|
||||
}
|
||||
|
||||
// formatValue formats a JSON value with indentation.
|
||||
func formatValue(v interface{}) string {
|
||||
b, err := json.MarshalIndent(v, "", " ")
|
||||
if err != nil {
|
||||
b, _ = json.Marshal(v)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// indentJSON formats raw JSON bytes with indentation.
|
||||
func indentJSON(data []byte) (string, error) {
|
||||
var buf bytes.Buffer
|
||||
if err := json.Indent(&buf, data, "", " "); err != nil {
|
||||
return string(data), err
|
||||
}
|
||||
return buf.String(), nil
|
||||
}
|
||||
|
||||
// wrapCodeBlock wraps content in a fenced JSON code block.
|
||||
func wrapCodeBlock(content string) string {
|
||||
return "```json\n" + content + "\n```"
|
||||
}
|
||||
|
||||
// trimBOM removes a UTF-8 BOM prefix if present.
|
||||
func trimBOM(data []byte) []byte {
|
||||
if len(data) >= 3 && data[0] == 0xEF && data[1] == 0xBB && data[2] == 0xBF {
|
||||
return data[3:]
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
// canSplitDict returns true if a dict can be meaningfully split.
|
||||
// A dict with multiple keys can be split by distributing keys across chunks.
|
||||
// A dict with a single key can still be split if its value is a splittable dict.
|
||||
func canSplitDict(d map[string]interface{}) bool {
|
||||
if len(d) > 1 {
|
||||
return true
|
||||
}
|
||||
if len(d) == 1 {
|
||||
for _, v := range d {
|
||||
if sub, ok := v.(map[string]interface{}); ok && len(sub) > 1 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// sortedKeys returns the keys of a map in sorted order.
|
||||
// When all keys are numeric strings (from array-to-dict conversion),
|
||||
// sorts numerically so "2" comes before "10".
|
||||
func sortedKeys(m map[string]interface{}) []string {
|
||||
keys := make([]string, 0, len(m))
|
||||
allNumeric := true
|
||||
for k := range m {
|
||||
keys = append(keys, k)
|
||||
if allNumeric {
|
||||
if _, err := strconv.Atoi(k); err != nil {
|
||||
allNumeric = false
|
||||
}
|
||||
}
|
||||
}
|
||||
if allNumeric {
|
||||
sort.Slice(keys, func(i, j int) bool {
|
||||
ni, _ := strconv.Atoi(keys[i])
|
||||
nj, _ := strconv.Atoi(keys[j])
|
||||
return ni < nj
|
||||
})
|
||||
} else {
|
||||
sort.Strings(keys)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
@@ -0,0 +1,705 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Helper: visual test reporter
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func visualReport(t *testing.T, label string, input string, result string, err error, checks []checkResult) {
|
||||
t.Helper()
|
||||
|
||||
const width = 72
|
||||
bar := strings.Repeat("─", width)
|
||||
thickBar := strings.Repeat("━", width)
|
||||
|
||||
t.Logf("\n%s", thickBar)
|
||||
t.Logf("📋 TEST: %s", label)
|
||||
t.Logf("%s", bar)
|
||||
|
||||
// ── Input ──
|
||||
inputPreview := input
|
||||
if len(inputPreview) > 200 {
|
||||
inputPreview = inputPreview[:200] + fmt.Sprintf("... (%d bytes total)", len(input))
|
||||
}
|
||||
t.Logf("📥 INPUT (%d bytes):", len(input))
|
||||
for _, line := range strings.Split(inputPreview, "\n") {
|
||||
t.Logf(" %s", line)
|
||||
}
|
||||
t.Logf("%s", bar)
|
||||
|
||||
// ── Output / Error ──
|
||||
if err != nil {
|
||||
t.Logf("❌ ERROR: %v", err)
|
||||
} else {
|
||||
blockCount := strings.Count(result, "```json")
|
||||
totalLen := len(result)
|
||||
t.Logf("📤 OUTPUT: %d code block(s), %d bytes total", blockCount, totalLen)
|
||||
t.Logf("%s", bar)
|
||||
|
||||
blocks := splitBlocks(result)
|
||||
for i, blk := range blocks {
|
||||
preview := blk
|
||||
lines := strings.Split(preview, "\n")
|
||||
if len(lines) > 10 {
|
||||
preview = strings.Join(lines[:4], "\n") +
|
||||
fmt.Sprintf("\n ... (%d lines omitted) ...\n", len(lines)-8) +
|
||||
strings.Join(lines[len(lines)-4:], "\n")
|
||||
}
|
||||
t.Logf(" 📦 Block #%d (%d bytes):", i+1, len(blk))
|
||||
for _, line := range strings.Split(preview, "\n") {
|
||||
t.Logf(" │ %s", line)
|
||||
}
|
||||
}
|
||||
}
|
||||
t.Logf("%s", bar)
|
||||
|
||||
// ── Checks ──
|
||||
allPass := true
|
||||
for _, c := range checks {
|
||||
icon := "✅"
|
||||
if !c.pass {
|
||||
icon = "💥"
|
||||
allPass = false
|
||||
}
|
||||
t.Logf("%s %s", icon, c.desc)
|
||||
if !c.pass {
|
||||
t.Errorf("FAIL: %s", c.desc)
|
||||
}
|
||||
}
|
||||
if allPass {
|
||||
t.Logf("🎉 ALL CHECKS PASSED")
|
||||
}
|
||||
t.Logf("%s\n", thickBar)
|
||||
}
|
||||
|
||||
// splitBlocks extracts the content between ```json and ``` fences.
|
||||
func splitBlocks(result string) []string {
|
||||
var blocks []string
|
||||
rest := result
|
||||
for {
|
||||
start := strings.Index(rest, "```json\n")
|
||||
if start < 0 {
|
||||
break
|
||||
}
|
||||
rest = rest[start+len("```json\n"):]
|
||||
end := strings.Index(rest, "\n```")
|
||||
if end < 0 {
|
||||
blocks = append(blocks, rest)
|
||||
break
|
||||
}
|
||||
blocks = append(blocks, rest[:end])
|
||||
rest = rest[end+len("\n```"):]
|
||||
}
|
||||
return blocks
|
||||
}
|
||||
|
||||
type checkResult struct {
|
||||
desc string
|
||||
pass bool
|
||||
}
|
||||
|
||||
func check(desc string, pass bool) checkResult { return checkResult{desc, pass} }
|
||||
|
||||
// blockIsValidJSON checks that each code block inside the result is valid JSON.
|
||||
func allBlocksValidJSON(result string) bool {
|
||||
for _, blk := range splitBlocks(result) {
|
||||
if !json.Valid([]byte(blk)) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 1: Small inputs — kept intact, NOT split
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestJsonToMarkdown_SmallObject(t *testing.T) {
|
||||
input := `{"name": "test", "version": "1.0"}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Small Object → single block, kept intact", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains key 'name'", strings.Contains(result, `"name"`)),
|
||||
check("contains key 'version'", strings.Contains(result, `"version"`)),
|
||||
check(fmt.Sprintf("exactly 1 block (got %d)", blocks), blocks == 1),
|
||||
check("both keys in SAME block (not split apart)",
|
||||
strings.Contains(result, `"name"`) && strings.Contains(result, `"version"`) && blocks == 1),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_SmallArray(t *testing.T) {
|
||||
input := `[{"id": 1, "name": "a"}, {"id": 2, "name": "b"}]`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Small Array → single block, kept intact", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("exactly 1 block (got %d)", blocks), blocks == 1),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_EmptyObject(t *testing.T) {
|
||||
input := `{}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Empty Object {}", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains '{}'", strings.Contains(result, "{}")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_EmptyArray(t *testing.T) {
|
||||
input := `[]`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Empty Array []", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains '{}'", strings.Contains(result, "{}")), // [] → {} after list-to-dict
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_PrimitiveString(t *testing.T) {
|
||||
input := `"hello world"`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Primitive String", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains 'hello world'", strings.Contains(result, "hello world")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_PrimitiveNumber(t *testing.T) {
|
||||
input := `42`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Primitive Number", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains '42'", strings.Contains(result, "42")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_PrimitiveBoolean(t *testing.T) {
|
||||
input := `true`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Primitive Boolean", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains 'true'", strings.Contains(result, "true")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_PrimitiveNull(t *testing.T) {
|
||||
input := `null`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Primitive Null", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains 'null'", strings.Contains(result, "null")),
|
||||
})
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 2: Large data — smart recursive splitting
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestJsonToMarkdown_LargeObject(t *testing.T) {
|
||||
obj := make(map[string]interface{})
|
||||
for i := 0; i < 50; i++ {
|
||||
obj[fmt.Sprintf("key_%02d", i)] = strings.Repeat("value", 50)
|
||||
}
|
||||
data, _ := json.Marshal(obj)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Large Object (50 keys) → multiple blocks", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("multiple blocks (got %d)", blocks), blocks >= 2),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_LargeArray(t *testing.T) {
|
||||
arr := make([]interface{}, 100)
|
||||
for i := range arr {
|
||||
arr[i] = map[string]interface{}{
|
||||
"id": i,
|
||||
"name": strings.Repeat("name", 20),
|
||||
"description": strings.Repeat("desc", 30),
|
||||
}
|
||||
}
|
||||
data, _ := json.Marshal(arr)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Large Array (100 elements) → multiple blocks", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("multiple blocks (got %d)", blocks), blocks >= 2),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_RecursiveSplitLargeKey(t *testing.T) {
|
||||
inner := make(map[string]interface{})
|
||||
for i := 0; i < 30; i++ {
|
||||
inner[fmt.Sprintf("sub_%02d", i)] = strings.Repeat("val", 50)
|
||||
}
|
||||
obj := map[string]interface{}{"bigkey": inner}
|
||||
data, _ := json.Marshal(obj)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Recursive Split: 1 big key with 30 sub-keys", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("multiple blocks (got %d)", blocks), blocks >= 2),
|
||||
check("path preserved: chunks contain 'bigkey'", strings.Contains(result, "bigkey")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_PathPreservation(t *testing.T) {
|
||||
// Core test: verify that nested paths are preserved in each chunk
|
||||
inner := make(map[string]interface{})
|
||||
for i := 0; i < 40; i++ {
|
||||
inner[fmt.Sprintf("field_%02d", i)] = strings.Repeat("data", 60)
|
||||
}
|
||||
obj := map[string]interface{}{
|
||||
"config": map[string]interface{}{
|
||||
"database": inner,
|
||||
},
|
||||
}
|
||||
data, _ := json.Marshal(obj)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := splitBlocks(result)
|
||||
// Every block should contain the full path "config" → "database"
|
||||
allHavePath := true
|
||||
for _, blk := range blocks {
|
||||
if !strings.Contains(blk, `"config"`) || !strings.Contains(blk, `"database"`) {
|
||||
allHavePath = false
|
||||
break
|
||||
}
|
||||
}
|
||||
visualReport(t, "Path Preservation: config.database.* across chunks", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("multiple blocks (got %d)", len(blocks)), len(blocks) >= 2),
|
||||
check("ALL blocks preserve path: config → database", allHavePath),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_MixedArrayElements(t *testing.T) {
|
||||
arr := []interface{}{
|
||||
map[string]interface{}{"id": 1},
|
||||
map[string]interface{}{"id": 2},
|
||||
map[string]interface{}{"id": 3, "data": strings.Repeat("x", 2000)},
|
||||
map[string]interface{}{"id": 4},
|
||||
}
|
||||
data, _ := json.Marshal(arr)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Mixed Array: 3 small + 1 large element", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("at least 2 blocks (got %d)", blocks), blocks >= 2),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_DeepNested(t *testing.T) {
|
||||
input := `{"l1": {"l2": {"l3": {"data": "` + strings.Repeat("x", 2000) + `"}}}}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Deep Nested (3 levels, large leaf)", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("block count ≥ 1 (got %d)", blocks), blocks >= 1),
|
||||
})
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 3: Error Handling
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestJsonToMarkdown_InvalidJSON(t *testing.T) {
|
||||
input := `{invalid json}`
|
||||
_, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Invalid JSON → error", input, "", err, []checkResult{
|
||||
check("returns error", err != nil),
|
||||
check("error mentions 'invalid JSON'",
|
||||
err != nil && strings.Contains(err.Error(), "invalid JSON")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_EmptyInput(t *testing.T) {
|
||||
input := ""
|
||||
_, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Empty Input → error", input, "", err, []checkResult{
|
||||
check("returns error", err != nil),
|
||||
check("error mentions 'empty'",
|
||||
err != nil && strings.Contains(err.Error(), "empty")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_WhitespaceOnly(t *testing.T) {
|
||||
input := " \n\t "
|
||||
_, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Whitespace-Only Input → error", input, "", err, []checkResult{
|
||||
check("returns error", err != nil),
|
||||
})
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 4: Edge Cases & Encoding
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestJsonToMarkdown_BOM(t *testing.T) {
|
||||
raw := append([]byte{0xEF, 0xBB, 0xBF}, []byte(`{"key": "value"}`)...)
|
||||
result, err := jsonToMarkdown(raw)
|
||||
visualReport(t, "UTF-8 BOM prefix → stripped", string(raw), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("contains key after BOM removal", strings.Contains(result, `"key"`)),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_UnicodeContent(t *testing.T) {
|
||||
input := `{"名称": "WeKnora 知识库", "描述": "支持中文 JSON 🎉", "emoji": "🚀"}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Unicode / Chinese / Emoji", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains Chinese", strings.Contains(result, "知识库")),
|
||||
check("contains emoji", strings.Contains(result, "🚀")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_SpecialCharsInStrings(t *testing.T) {
|
||||
input := `{"html": "<div class=\"test\">hello</div>", "path": "C:\\Users\\test", "newline": "line1\nline2"}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Special characters (HTML, backslash, newline)", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_NullValues(t *testing.T) {
|
||||
input := `{"name": "test", "value": null, "list": [null, 1, null]}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Object with null values", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains 'null'", strings.Contains(result, "null")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_MixedValueTypes(t *testing.T) {
|
||||
input := `{"str": "hello", "num": 42, "float": 3.14, "bool": true, "null_val": null, "arr": [1,2], "obj": {"a":1}}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Mixed value types in one object", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains string", strings.Contains(result, "hello")),
|
||||
check("contains number", strings.Contains(result, "42")),
|
||||
check("contains float", strings.Contains(result, "3.14")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_NestedArrayOfArrays(t *testing.T) {
|
||||
input := `[[1, 2, 3], [4, 5, 6], [7, 8, 9]]`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Nested Array of Arrays (matrix)", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_LargePrimitiveString(t *testing.T) {
|
||||
bigStr := strings.Repeat("abcdefghij", 300)
|
||||
input := fmt.Sprintf(`"%s"`, bigStr)
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
visualReport(t, "Large Primitive String (3000 chars) → fallback", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("content preserved", strings.Contains(result, "abcdefghij")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_SingleKeyObject(t *testing.T) {
|
||||
input := `{"only_key": "only_value"}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Single-Key Object", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("exactly 1 block (got %d)", blocks), blocks == 1),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_SingleElementArray(t *testing.T) {
|
||||
input := `[{"id": 1, "name": "solo"}]`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Single-Element Array", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("exactly 1 block (got %d)", blocks), blocks == 1),
|
||||
})
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 5: Realistic Scenarios
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestJsonToMarkdown_RealisticConfig(t *testing.T) {
|
||||
input := `{
|
||||
"app": {"name": "WeKnora", "version": "2.0.0", "debug": false},
|
||||
"database": {"host": "localhost", "port": 5432, "name": "weknora_db", "pool_size": 10},
|
||||
"redis": {"host": "localhost", "port": 6379},
|
||||
"features": {"json_upload": true, "multimodel": true, "graph_extraction": false}
|
||||
}`
|
||||
result, err := jsonToMarkdown([]byte(input))
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Realistic: App Config (fits in 1 block)", input, result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains 'WeKnora'", strings.Contains(result, "WeKnora")),
|
||||
check("contains 'database'", strings.Contains(result, "database")),
|
||||
check(fmt.Sprintf("single block (fits, got %d)", blocks), blocks == 1),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_RealisticAPIResponse(t *testing.T) {
|
||||
users := make([]interface{}, 40)
|
||||
for i := range users {
|
||||
users[i] = map[string]interface{}{
|
||||
"id": i + 1,
|
||||
"username": fmt.Sprintf("user_%03d", i+1),
|
||||
"email": fmt.Sprintf("user%03d@example.com", i+1),
|
||||
"role": "editor",
|
||||
"active": i%3 != 0,
|
||||
}
|
||||
}
|
||||
apiResp := map[string]interface{}{
|
||||
"status": "ok",
|
||||
"total": 40,
|
||||
"page": 1,
|
||||
"data": users,
|
||||
}
|
||||
data, _ := json.Marshal(apiResp)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := strings.Count(result, "```json")
|
||||
visualReport(t, "Realistic: API Response (40 user records)", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check("contains user data", strings.Contains(result, "user_001")),
|
||||
check(fmt.Sprintf("multiple blocks (got %d)", blocks), blocks >= 2),
|
||||
})
|
||||
}
|
||||
|
||||
func TestJsonToMarkdown_RealisticLargeNestedConfig(t *testing.T) {
|
||||
// A config where one section is huge and must be recursively split,
|
||||
// while other sections are small
|
||||
rules := make(map[string]interface{})
|
||||
for i := 0; i < 30; i++ {
|
||||
rules[fmt.Sprintf("rule_%02d", i)] = map[string]interface{}{
|
||||
"pattern": strings.Repeat("pattern", 10),
|
||||
"action": "allow",
|
||||
"priority": i,
|
||||
}
|
||||
}
|
||||
config := map[string]interface{}{
|
||||
"version": "3.0",
|
||||
"metadata": map[string]interface{}{"author": "admin", "updated": "2026-03-24"},
|
||||
"firewall": map[string]interface{}{
|
||||
"enabled": true,
|
||||
"rules": rules,
|
||||
},
|
||||
}
|
||||
data, _ := json.Marshal(config)
|
||||
|
||||
result, err := jsonToMarkdown(data)
|
||||
blocks := splitBlocks(result)
|
||||
// Each block should still have the path "firewall" → "rules" where applicable
|
||||
visualReport(t, "Realistic: Nested config with 30 firewall rules", string(data), result, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("all blocks are valid JSON", allBlocksValidJSON(result)),
|
||||
check(fmt.Sprintf("multiple blocks needed (got %d)", len(blocks)), len(blocks) >= 2),
|
||||
check("contains 'firewall'", strings.Contains(result, "firewall")),
|
||||
check("contains 'version'", strings.Contains(result, "version")),
|
||||
})
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 6: Integration — SimpleFormatReader
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestSimpleFormatReader_JSON(t *testing.T) {
|
||||
reader := &SimpleFormatReader{}
|
||||
input := `{"test": "integration", "count": 42}`
|
||||
req := &types.ReadRequest{
|
||||
FileName: "config.json",
|
||||
FileType: "json",
|
||||
FileContent: []byte(input),
|
||||
}
|
||||
|
||||
result, err := reader.Read(context.Background(), req)
|
||||
|
||||
var resultStr string
|
||||
if result != nil {
|
||||
resultStr = result.MarkdownContent
|
||||
}
|
||||
visualReport(t, "SimpleFormatReader.Read() with .json file", input, resultStr, err, []checkResult{
|
||||
check("no error", err == nil),
|
||||
check("result not nil", result != nil),
|
||||
check("markdown contains 'test'",
|
||||
result != nil && strings.Contains(result.MarkdownContent, "test")),
|
||||
check("markdown has code block",
|
||||
result != nil && strings.Contains(result.MarkdownContent, "```json")),
|
||||
check("output is valid JSON inside block",
|
||||
result != nil && allBlocksValidJSON(result.MarkdownContent)),
|
||||
})
|
||||
}
|
||||
|
||||
func TestSimpleFormatReader_JSON_Invalid(t *testing.T) {
|
||||
reader := &SimpleFormatReader{}
|
||||
req := &types.ReadRequest{
|
||||
FileName: "bad.json",
|
||||
FileType: "json",
|
||||
FileContent: []byte(`{not valid json}`),
|
||||
}
|
||||
|
||||
_, err := reader.Read(context.Background(), req)
|
||||
visualReport(t, "SimpleFormatReader with invalid JSON → error", `{not valid json}`, "", err, []checkResult{
|
||||
check("returns error", err != nil),
|
||||
check("error mentions 'json'",
|
||||
err != nil && strings.Contains(strings.ToLower(err.Error()), "json")),
|
||||
})
|
||||
}
|
||||
|
||||
func TestIsSimpleFormat_JSON(t *testing.T) {
|
||||
cases := []struct {
|
||||
input string
|
||||
want bool
|
||||
}{
|
||||
{"json", true}, {"JSON", true}, {".json", true}, {"Json", true},
|
||||
{"txt", true}, {"csv", true},
|
||||
{"pdf", false}, {"docx", false},
|
||||
}
|
||||
|
||||
t.Logf("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
|
||||
t.Logf("📋 TEST: IsSimpleFormat() recognizes JSON")
|
||||
t.Logf("─────────────────────────────────────────────────────────")
|
||||
allPass := true
|
||||
for _, tc := range cases {
|
||||
got := IsSimpleFormat(tc.input)
|
||||
icon := "✅"
|
||||
if got != tc.want {
|
||||
icon = "💥"
|
||||
allPass = false
|
||||
t.Errorf("IsSimpleFormat(%q) = %v, want %v", tc.input, got, tc.want)
|
||||
}
|
||||
t.Logf("%s IsSimpleFormat(%q) = %v (want %v)", icon, tc.input, got, tc.want)
|
||||
}
|
||||
if allPass {
|
||||
t.Logf("🎉 ALL CHECKS PASSED")
|
||||
}
|
||||
t.Logf("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n")
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// Group 7: list-to-dict preprocessing
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
func TestListToDictPreprocess(t *testing.T) {
|
||||
// Simulate what json.Unmarshal produces: numbers become float64
|
||||
input := []interface{}{"a", "b", map[string]interface{}{"nested": []interface{}{float64(1), float64(2)}}}
|
||||
result := listToDictPreprocess(input)
|
||||
|
||||
dict, ok := result.(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("expected map, got %T", result)
|
||||
}
|
||||
|
||||
t.Logf("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
|
||||
t.Logf("📋 TEST: listToDictPreprocess")
|
||||
t.Logf("─────────────────────────────────────────────────────────")
|
||||
t.Logf("📥 INPUT: [\"a\", \"b\", {\"nested\": [1, 2]}]")
|
||||
t.Logf("📤 OUTPUT: %s", formatValue(dict))
|
||||
t.Logf("─────────────────────────────────────────────────────────")
|
||||
|
||||
checks := []checkResult{
|
||||
check("top level is dict", ok),
|
||||
check(`key "0" = "a"`, dict["0"] == "a"),
|
||||
check(`key "1" = "b"`, dict["1"] == "b"),
|
||||
}
|
||||
// Check nested array also converted
|
||||
if nested, ok2 := dict["2"].(map[string]interface{}); ok2 {
|
||||
if innerNested, ok3 := nested["nested"].(map[string]interface{}); ok3 {
|
||||
checks = append(checks, check(`nested array converted to dict`, innerNested["0"] == float64(1)))
|
||||
} else {
|
||||
checks = append(checks, check(`nested array converted to dict`, false))
|
||||
}
|
||||
}
|
||||
allPass := true
|
||||
for _, c := range checks {
|
||||
icon := "✅"
|
||||
if !c.pass {
|
||||
icon = "💥"
|
||||
allPass = false
|
||||
t.Errorf("FAIL: %s", c.desc)
|
||||
}
|
||||
t.Logf("%s %s", icon, c.desc)
|
||||
}
|
||||
if allPass {
|
||||
t.Logf("🎉 ALL CHECKS PASSED")
|
||||
}
|
||||
t.Logf("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n")
|
||||
}
|
||||
|
||||
func TestSetNestedDict(t *testing.T) {
|
||||
d := map[string]interface{}{}
|
||||
setNestedDict(d, []string{"config", "database", "host"}, "localhost")
|
||||
|
||||
t.Logf("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
|
||||
t.Logf("📋 TEST: setNestedDict path preservation")
|
||||
t.Logf("─────────────────────────────────────────────────────────")
|
||||
t.Logf("📥 path: [\"config\", \"database\", \"host\"]")
|
||||
t.Logf("📥 value: \"localhost\"")
|
||||
t.Logf("📤 result: %s", formatValue(d))
|
||||
t.Logf("─────────────────────────────────────────────────────────")
|
||||
|
||||
// Verify the structure
|
||||
config, ok1 := d["config"].(map[string]interface{})
|
||||
db, ok2 := map[string]interface{}{}, false
|
||||
host := ""
|
||||
if ok1 {
|
||||
db, ok2 = config["database"].(map[string]interface{})
|
||||
}
|
||||
if ok2 {
|
||||
host, _ = db["host"].(string)
|
||||
}
|
||||
|
||||
checks := []checkResult{
|
||||
check(`d["config"] is dict`, ok1),
|
||||
check(`d["config"]["database"] is dict`, ok2),
|
||||
check(`d["config"]["database"]["host"] = "localhost"`, host == "localhost"),
|
||||
}
|
||||
allPass := true
|
||||
for _, c := range checks {
|
||||
icon := "✅"
|
||||
if !c.pass {
|
||||
icon = "💥"
|
||||
allPass = false
|
||||
t.Errorf("FAIL: %s", c.desc)
|
||||
}
|
||||
t.Logf("%s %s", icon, c.desc)
|
||||
}
|
||||
if allPass {
|
||||
t.Logf("🎉 ALL CHECKS PASSED")
|
||||
}
|
||||
t.Logf("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n")
|
||||
}
|
||||
@@ -0,0 +1,262 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
type markdownImageTargetSpan struct {
|
||||
TargetStart int
|
||||
TargetEnd int
|
||||
}
|
||||
|
||||
func scanMarkdownImageTargets(markdown string) []markdownImageTargetSpan {
|
||||
var spans []markdownImageTargetSpan
|
||||
for i := 0; i+1 < len(markdown); i++ {
|
||||
if markdown[i] != '!' || markdown[i+1] != '[' || isEscaped(markdown, i) {
|
||||
continue
|
||||
}
|
||||
|
||||
altEnd := findMarkdownImageAltEnd(markdown, i+2)
|
||||
if altEnd == -1 {
|
||||
continue
|
||||
}
|
||||
|
||||
targetStart := altEnd + 2
|
||||
targetEnd, ok := findMarkdownImageTargetEnd(markdown, targetStart)
|
||||
if !ok {
|
||||
i = altEnd
|
||||
continue
|
||||
}
|
||||
spans = append(spans, markdownImageTargetSpan{
|
||||
TargetStart: targetStart,
|
||||
TargetEnd: targetEnd,
|
||||
})
|
||||
i = targetEnd
|
||||
}
|
||||
return spans
|
||||
}
|
||||
|
||||
func findMarkdownImageAltEnd(markdown string, start int) int {
|
||||
for i := start; i+1 < len(markdown); i++ {
|
||||
if markdown[i] == ']' && markdown[i+1] == '(' && !isEscaped(markdown, i) {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
func findMarkdownImageTargetEnd(markdown string, start int) (int, bool) {
|
||||
parenDepth := 1
|
||||
inAngleDestination := false
|
||||
seenNonSpace := false
|
||||
var inQuote byte
|
||||
|
||||
for i := start; i < len(markdown); i++ {
|
||||
ch := markdown[i]
|
||||
if ch == '\\' {
|
||||
i++
|
||||
continue
|
||||
}
|
||||
|
||||
if !seenNonSpace && !isMarkdownSpace(ch) {
|
||||
seenNonSpace = true
|
||||
if ch == '<' {
|
||||
inAngleDestination = true
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
if inAngleDestination {
|
||||
if ch == '>' {
|
||||
inAngleDestination = false
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if inQuote != 0 {
|
||||
if ch == inQuote {
|
||||
inQuote = 0
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if (ch == '"' || ch == '\'') && i > start && isMarkdownSpace(markdown[i-1]) {
|
||||
inQuote = ch
|
||||
continue
|
||||
}
|
||||
|
||||
switch ch {
|
||||
case '(':
|
||||
parenDepth++
|
||||
case ')':
|
||||
parenDepth--
|
||||
if parenDepth == 0 {
|
||||
return i, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func splitMarkdownImageTarget(
|
||||
raw string,
|
||||
refMap map[string]types.ImageRef,
|
||||
) (path string, pathStart int, pathEnd int, ok bool) {
|
||||
start, end := trimMarkdownSpaceBounds(raw, 0, len(raw))
|
||||
if start >= end {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
|
||||
trimmed := raw[start:end]
|
||||
if _, found := refMap[trimmed]; found {
|
||||
return trimmed, start, end, true
|
||||
}
|
||||
|
||||
if raw[start] == '<' {
|
||||
return splitAngleMarkdownImageTarget(raw, start, end, refMap)
|
||||
}
|
||||
|
||||
titleStart, found := parseMarkdownImageTitleSuffix(trimmed)
|
||||
if !found {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
|
||||
pathEnd = start + titleStart
|
||||
for pathEnd > start && isMarkdownSpace(raw[pathEnd-1]) {
|
||||
pathEnd--
|
||||
}
|
||||
if pathEnd == start {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
|
||||
path = raw[start:pathEnd]
|
||||
if _, found := refMap[path]; !found {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
return path, start, pathEnd, true
|
||||
}
|
||||
|
||||
func splitAngleMarkdownImageTarget(
|
||||
raw string,
|
||||
start int,
|
||||
end int,
|
||||
refMap map[string]types.ImageRef,
|
||||
) (path string, pathStart int, pathEnd int, ok bool) {
|
||||
closeIdx := -1
|
||||
for i := start + 1; i < end; i++ {
|
||||
if raw[i] == '>' && !isEscaped(raw, i) {
|
||||
closeIdx = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if closeIdx == -1 {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
|
||||
path = raw[start+1 : closeIdx]
|
||||
if _, found := refMap[path]; !found {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
if !isEmptyOrMarkdownImageTitleSuffix(raw[closeIdx+1 : end]) {
|
||||
return "", 0, 0, false
|
||||
}
|
||||
return path, start + 1, closeIdx, true
|
||||
}
|
||||
|
||||
func parseMarkdownImageTitleSuffix(raw string) (titleStart int, ok bool) {
|
||||
_, end := trimMarkdownSpaceBounds(raw, 0, len(raw))
|
||||
if end == 0 {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
switch raw[end-1] {
|
||||
case '"', '\'':
|
||||
quote := raw[end-1]
|
||||
for i := end - 2; i >= 0; i-- {
|
||||
if raw[i] != quote || isEscaped(raw, i) {
|
||||
continue
|
||||
}
|
||||
if i == 0 || !isMarkdownSpace(raw[i-1]) {
|
||||
return 0, false
|
||||
}
|
||||
if markdownTitleHasBlankLine(raw[i+1 : end-1]) {
|
||||
return 0, false
|
||||
}
|
||||
return i, true
|
||||
}
|
||||
case ')':
|
||||
depth := 0
|
||||
for i := end - 2; i >= 0; i-- {
|
||||
if isEscaped(raw, i) {
|
||||
continue
|
||||
}
|
||||
switch raw[i] {
|
||||
case ')':
|
||||
depth++
|
||||
case '(':
|
||||
if depth == 0 {
|
||||
if i == 0 || !isMarkdownSpace(raw[i-1]) {
|
||||
return 0, false
|
||||
}
|
||||
if markdownTitleHasBlankLine(raw[i+1 : end-1]) {
|
||||
return 0, false
|
||||
}
|
||||
return i, true
|
||||
}
|
||||
depth--
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func isEmptyOrMarkdownImageTitleSuffix(raw string) bool {
|
||||
start, end := trimMarkdownSpaceBounds(raw, 0, len(raw))
|
||||
if start == end {
|
||||
return true
|
||||
}
|
||||
titleStart, ok := parseMarkdownImageTitleSuffix(raw[:end])
|
||||
return ok && isAllMarkdownSpace(raw[:titleStart])
|
||||
}
|
||||
|
||||
func markdownTitleHasBlankLine(title string) bool {
|
||||
for _, line := range strings.Split(title, "\n") {
|
||||
if strings.Trim(line, " \t") == "" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isAllMarkdownSpace(raw string) bool {
|
||||
for i := 0; i < len(raw); i++ {
|
||||
if !isMarkdownSpace(raw[i]) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func trimMarkdownSpaceBounds(raw string, start int, end int) (int, int) {
|
||||
for start < end && isMarkdownSpace(raw[start]) {
|
||||
start++
|
||||
}
|
||||
for end > start && isMarkdownSpace(raw[end-1]) {
|
||||
end--
|
||||
}
|
||||
return start, end
|
||||
}
|
||||
|
||||
func isMarkdownSpace(b byte) bool {
|
||||
return b == ' ' || b == '\t' || b == '\n' || b == '\r'
|
||||
}
|
||||
|
||||
func isEscaped(s string, pos int) bool {
|
||||
backslashes := 0
|
||||
for i := pos - 1; i >= 0 && s[i] == '\\'; i-- {
|
||||
backslashes++
|
||||
}
|
||||
return backslashes%2 == 1
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestScanMarkdownImageTargets(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
targets []string
|
||||
}{
|
||||
{
|
||||
name: "without title",
|
||||
input: `before  after`,
|
||||
targets: []string{`images/a.png`},
|
||||
},
|
||||
{
|
||||
name: "double quoted title",
|
||||
input: ``,
|
||||
targets: []string{`images/a.png "图"`},
|
||||
},
|
||||
{
|
||||
name: "single quoted title",
|
||||
input: ``,
|
||||
targets: []string{`images/a.png '图'`},
|
||||
},
|
||||
{
|
||||
name: "parenthesized title",
|
||||
input: `)`,
|
||||
targets: []string{`images/a.png (图)`},
|
||||
},
|
||||
{
|
||||
name: "title containing right paren",
|
||||
input: ` 结果")`,
|
||||
targets: []string{`images/a.png "阶段 1) 结果"`},
|
||||
},
|
||||
{
|
||||
name: "title containing both parens",
|
||||
input: ` 结果')`,
|
||||
targets: []string{`images/a.png '阶段 (1) 结果'`},
|
||||
},
|
||||
{
|
||||
name: "escaped quote in title",
|
||||
input: ``,
|
||||
targets: []string{`images/a.png "阶段 \"1\""`},
|
||||
},
|
||||
{
|
||||
name: "multiline title",
|
||||
input: "",
|
||||
targets: []string{"images/a.png\n \"多行 title\""},
|
||||
},
|
||||
{
|
||||
name: "spaced path",
|
||||
input: ``,
|
||||
targets: []string{`images/第 1 页.png`},
|
||||
},
|
||||
{
|
||||
name: "path containing balanced parens",
|
||||
input: `.png "title")`,
|
||||
targets: []string{`images/a_(1).png "title"`},
|
||||
},
|
||||
{
|
||||
name: "angle destination",
|
||||
input: ``,
|
||||
targets: []string{`<images/a b.png> "title"`},
|
||||
},
|
||||
{
|
||||
name: "malformed image is skipped",
|
||||
input: ` after `,
|
||||
targets: []string{`images/b.png`},
|
||||
},
|
||||
{
|
||||
name: "escaped image marker is skipped",
|
||||
input: `\ `,
|
||||
targets: []string{`images/b.png`},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
spans := scanMarkdownImageTargets(tt.input)
|
||||
if len(spans) != len(tt.targets) {
|
||||
t.Fatalf("got %d spans, want %d: %#v", len(spans), len(tt.targets), spans)
|
||||
}
|
||||
for i, span := range spans {
|
||||
got := tt.input[span.TargetStart:span.TargetEnd]
|
||||
if got != tt.targets[i] {
|
||||
t.Fatalf("target %d = %q, want %q", i, got, tt.targets[i])
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitMarkdownImageTarget(t *testing.T) {
|
||||
refMap := map[string]types.ImageRef{
|
||||
"images/a.png": {OriginalRef: "images/a.png"},
|
||||
"images/第 1 页.png": {OriginalRef: "images/第 1 页.png"},
|
||||
"images/a_(1).png": {OriginalRef: "images/a_(1).png"},
|
||||
"images/第 1 页 (测试).gif": {OriginalRef: "images/第 1 页 (测试).gif"},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
raw string
|
||||
wantPath string
|
||||
wantText string
|
||||
wantOK bool
|
||||
}{
|
||||
{
|
||||
name: "without title",
|
||||
raw: `images/a.png`,
|
||||
wantPath: `images/a.png`,
|
||||
wantText: `local://stored`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "double quoted title",
|
||||
raw: `images/a.png "图"`,
|
||||
wantPath: `images/a.png`,
|
||||
wantText: `local://stored "图"`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "single quoted title",
|
||||
raw: `images/a.png '图'`,
|
||||
wantPath: `images/a.png`,
|
||||
wantText: `local://stored '图'`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "parenthesized title",
|
||||
raw: `images/a.png (阶段 (1) 结果)`,
|
||||
wantPath: `images/a.png`,
|
||||
wantText: `local://stored (阶段 (1) 结果)`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "multiline title",
|
||||
raw: "images/a.png\n \"多行 title\"",
|
||||
wantPath: `images/a.png`,
|
||||
wantText: "local://stored\n \"多行 title\"",
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "path with spaces wins without title",
|
||||
raw: `images/第 1 页.png`,
|
||||
wantPath: `images/第 1 页.png`,
|
||||
wantText: `local://stored`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "path with balanced parens",
|
||||
raw: `images/a_(1).png "title"`,
|
||||
wantPath: `images/a_(1).png`,
|
||||
wantText: `local://stored "title"`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "angle destination replaces only inner path",
|
||||
raw: `<images/第 1 页 (测试).gif> "阶段 1) 图片"`,
|
||||
wantPath: `images/第 1 页 (测试).gif`,
|
||||
wantText: `<local://stored> "阶段 1) 图片"`,
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "unknown reference",
|
||||
raw: `images/missing.png "title"`,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "blank line in title is invalid",
|
||||
raw: "images/a.png \"line1\n\nline2\"",
|
||||
wantOK: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
path, pathStart, pathEnd, ok := splitMarkdownImageTarget(tt.raw, refMap)
|
||||
if ok != tt.wantOK {
|
||||
t.Fatalf("ok = %v, want %v", ok, tt.wantOK)
|
||||
}
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if path != tt.wantPath {
|
||||
t.Fatalf("path = %q, want %q", path, tt.wantPath)
|
||||
}
|
||||
replaced := tt.raw[:pathStart] + "local://stored" + tt.raw[pathEnd:]
|
||||
if replaced != tt.wantText {
|
||||
t.Fatalf("replacement = %q, want %q", replaced, tt.wantText)
|
||||
}
|
||||
if strings.Contains(replaced, tt.wantPath) {
|
||||
t.Fatalf("replacement still contains original path: %q", replaced)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,517 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultPollInterval = 3 * time.Second
|
||||
defaultCloudTimeout = 600 * time.Second
|
||||
defaultBaseURL = "https://mineru.net/api/v4"
|
||||
)
|
||||
|
||||
// MinerUCloudReader calls the MinerU Cloud API (mineru.net) to read/convert documents.
|
||||
// Flow: POST /file-urls/batch → PUT file → poll GET /extract-results/batch/{batch_id}.
|
||||
type MinerUCloudReader struct {
|
||||
apiKey string
|
||||
baseURL string
|
||||
model string
|
||||
formulaEnable bool
|
||||
tableEnable bool
|
||||
ocrEnable bool
|
||||
language string
|
||||
}
|
||||
|
||||
// NewMinerUCloudReader creates a reader from ParserEngineOverrides.
|
||||
func NewMinerUCloudReader(overrides map[string]string) *MinerUCloudReader {
|
||||
return &MinerUCloudReader{
|
||||
apiKey: strings.TrimSpace(overrides["mineru_api_key"]),
|
||||
baseURL: defaultBaseURL,
|
||||
model: stringOr(overrides["mineru_cloud_model"], "pipeline"),
|
||||
formulaEnable: parseBoolOr(overrides["mineru_cloud_enable_formula"], true),
|
||||
tableEnable: parseBoolOr(overrides["mineru_cloud_enable_table"], true),
|
||||
ocrEnable: parseBoolOr(overrides["mineru_cloud_enable_ocr"], true),
|
||||
language: stringOr(overrides["mineru_cloud_language"], "ch"),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *MinerUCloudReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
if c.apiKey == "" {
|
||||
return &types.ReadResult{Error: "MinerU Cloud API key is not configured"}, nil
|
||||
}
|
||||
|
||||
content := req.FileContent
|
||||
if len(content) == 0 {
|
||||
return &types.ReadResult{Error: "no file content provided"}, nil
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[MinerUCloud] Parsing file=%s size=%d via %s", req.FileName, len(content), c.baseURL)
|
||||
|
||||
ext := filepath.Ext(req.FileName)
|
||||
if ext == "" && req.FileType != "" {
|
||||
ext = "." + req.FileType
|
||||
}
|
||||
if ext == "" {
|
||||
ext = ".pdf"
|
||||
}
|
||||
fileName := strings.TrimSuffix(req.FileName, ext) + ext
|
||||
if fileName == ext {
|
||||
fileName = "document" + ext
|
||||
}
|
||||
|
||||
batchID, uploadURL, err := c.applyUploadURLs(ctx, fileName, ext)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("MinerU Cloud apply upload URLs: %w", err)
|
||||
}
|
||||
|
||||
if err := c.uploadFile(ctx, uploadURL, content); err != nil {
|
||||
return nil, fmt.Errorf("MinerU Cloud file upload: %w", err)
|
||||
}
|
||||
|
||||
mdContent, imageRefs, err := c.pollBatchResult(ctx, batchID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("MinerU Cloud poll: %w", err)
|
||||
}
|
||||
|
||||
mdContent, imageRefs = ensureOriginalImageRef(req, mdContent, imageRefs)
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: mdContent,
|
||||
ImageRefs: imageRefs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// --- batch upload API ---
|
||||
|
||||
type batchApplyResponse struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
Data struct {
|
||||
BatchID string `json:"batch_id"`
|
||||
FileURLs []string `json:"file_urls"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
func (c *MinerUCloudReader) applyUploadURLs(ctx context.Context, fileName, ext string) (string, string, error) {
|
||||
modelVersion := c.model
|
||||
if strings.ToLower(ext) == ".html" {
|
||||
modelVersion = "MinerU-HTML"
|
||||
}
|
||||
|
||||
payload := map[string]interface{}{
|
||||
"files": []map[string]string{{"name": fileName, "data_id": uuid.New().String()}},
|
||||
"model_version": modelVersion,
|
||||
"is_ocr": c.ocrEnable,
|
||||
"enable_formula": c.formulaEnable,
|
||||
"enable_table": c.tableEnable,
|
||||
"language": c.language,
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("marshal payload: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/file-urls/batch", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 30 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("HTTP request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
return "", "", fmt.Errorf("API status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result batchApplyResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return "", "", fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
if result.Code != 0 {
|
||||
return "", "", fmt.Errorf("API error: %s", result.Msg)
|
||||
}
|
||||
if len(result.Data.FileURLs) == 0 {
|
||||
return "", "", fmt.Errorf("API returned no file_urls")
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[MinerUCloud] batch apply ok: batch_id=%s, urls=%d", result.Data.BatchID, len(result.Data.FileURLs))
|
||||
return result.Data.BatchID, result.Data.FileURLs[0], nil
|
||||
}
|
||||
|
||||
func (c *MinerUCloudReader) uploadFile(ctx context.Context, uploadURL string, content []byte) error {
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPut, uploadURL, bytes.NewReader(content))
|
||||
if err != nil {
|
||||
return fmt.Errorf("create PUT request: %w", err)
|
||||
}
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 120 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return fmt.Errorf("PUT upload: %w", err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
if resp.StatusCode >= 300 {
|
||||
return fmt.Errorf("PUT upload status %d", resp.StatusCode)
|
||||
}
|
||||
logger.Infof(context.Background(), "[MinerUCloud] file uploaded, status=%d", resp.StatusCode)
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- polling ---
|
||||
|
||||
type batchPollResponse struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
Data struct {
|
||||
ExtractResult json.RawMessage `json:"extract_result"` // can be object or array
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
type extractResultItem struct {
|
||||
State string `json:"state"`
|
||||
FileName string `json:"file_name"`
|
||||
Markdown string `json:"markdown"`
|
||||
Content string `json:"content"`
|
||||
Text string `json:"text"`
|
||||
ErrMsg string `json:"err_msg"`
|
||||
Progress struct {
|
||||
ExtractedPages int `json:"extracted_pages"`
|
||||
TotalPages int `json:"total_pages"`
|
||||
} `json:"extract_progress"`
|
||||
FullZipURL string `json:"full_zip_url"`
|
||||
}
|
||||
|
||||
func (c *MinerUCloudReader) pollBatchResult(ctx context.Context, batchID string) (string, []types.ImageRef, error) {
|
||||
deadline := time.Now().Add(defaultCloudTimeout)
|
||||
pollCount := 0
|
||||
headers := map[string]string{
|
||||
"Authorization": "Bearer " + c.apiKey,
|
||||
}
|
||||
|
||||
for time.Now().Before(deadline) {
|
||||
// Bail out promptly on caller cancellation instead of spinning:
|
||||
// fetchBatchStatus fails immediately and sleepCtx returns at once on a
|
||||
// cancelled ctx, so without this guard the loop busy-hammers the cloud
|
||||
// API and floods logs until the deadline.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
pollCount++
|
||||
|
||||
items, err := c.fetchBatchStatus(ctx, batchID, headers)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[MinerUCloud] poll #%d failed: %v", pollCount, err)
|
||||
sleepCtx(ctx, defaultPollInterval)
|
||||
continue
|
||||
}
|
||||
|
||||
if len(items) == 0 {
|
||||
if pollCount <= 3 || pollCount%10 == 0 {
|
||||
logger.Infof(context.Background(), "[MinerUCloud] poll #%d: extract_result empty, retrying", pollCount)
|
||||
}
|
||||
sleepCtx(ctx, defaultPollInterval)
|
||||
continue
|
||||
}
|
||||
|
||||
item := items[0]
|
||||
state := strings.ToLower(item.State)
|
||||
|
||||
if pollCount == 1 || pollCount%10 == 0 || state == "done" || state == "failed" {
|
||||
logger.Infof(context.Background(), "[MinerUCloud] poll #%d: file=%s state=%s pages=%d/%d",
|
||||
pollCount, item.FileName, state, item.Progress.ExtractedPages, item.Progress.TotalPages)
|
||||
}
|
||||
|
||||
if state == "failed" {
|
||||
return "", nil, fmt.Errorf("MinerU Cloud task failed: %s", item.ErrMsg)
|
||||
}
|
||||
|
||||
if state == "done" {
|
||||
return c.extractDoneResult(ctx, &item)
|
||||
}
|
||||
|
||||
sleepCtx(ctx, defaultPollInterval)
|
||||
}
|
||||
|
||||
return "", nil, fmt.Errorf("MinerU Cloud task timed out after %d polls", pollCount)
|
||||
}
|
||||
|
||||
func (c *MinerUCloudReader) fetchBatchStatus(ctx context.Context, batchID string, headers map[string]string) ([]extractResultItem, error) {
|
||||
url := fmt.Sprintf("%s/extract-results/batch/%s", c.baseURL, batchID)
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for k, v := range headers {
|
||||
httpReq.Header.Set(k, v)
|
||||
}
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 30 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read poll response body: %w", err)
|
||||
}
|
||||
|
||||
var pollResp batchPollResponse
|
||||
if err := json.Unmarshal(respBody, &pollResp); err != nil {
|
||||
return nil, fmt.Errorf("decode poll response: %w", err)
|
||||
}
|
||||
if pollResp.Code != 0 {
|
||||
return nil, fmt.Errorf("poll error code=%d msg=%s", pollResp.Code, pollResp.Msg)
|
||||
}
|
||||
|
||||
if len(pollResp.Data.ExtractResult) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Dump the raw extract_result JSON for debugging
|
||||
rawExtract := string(pollResp.Data.ExtractResult)
|
||||
if len(rawExtract) > 4000 {
|
||||
logger.Infof(context.Background(), "[MinerUCloud] Raw extract_result (truncated to 4000 chars): %s ...", rawExtract[:4000])
|
||||
} else {
|
||||
logger.Infof(context.Background(), "[MinerUCloud] Raw extract_result: %s", rawExtract)
|
||||
}
|
||||
|
||||
// Pretty-print the structure to reveal all available fields
|
||||
var rawObj interface{}
|
||||
if err := json.Unmarshal(pollResp.Data.ExtractResult, &rawObj); err == nil {
|
||||
logResponseStructure("MinerUCloud", rawObj, "extract_result")
|
||||
}
|
||||
|
||||
// The extract_result can be either a single object or an array
|
||||
var items []extractResultItem
|
||||
if pollResp.Data.ExtractResult[0] == '[' {
|
||||
if err := json.Unmarshal(pollResp.Data.ExtractResult, &items); err != nil {
|
||||
return nil, fmt.Errorf("decode extract_result array: %w", err)
|
||||
}
|
||||
} else {
|
||||
var single extractResultItem
|
||||
if err := json.Unmarshal(pollResp.Data.ExtractResult, &single); err != nil {
|
||||
return nil, fmt.Errorf("decode extract_result object: %w", err)
|
||||
}
|
||||
items = []extractResultItem{single}
|
||||
}
|
||||
|
||||
return items, nil
|
||||
}
|
||||
|
||||
// extractDoneResult extracts markdown and images from a completed batch item.
|
||||
// Prefers inline markdown/content fields; falls back to downloading full_zip_url.
|
||||
func (c *MinerUCloudReader) extractDoneResult(_ context.Context, item *extractResultItem) (string, []types.ImageRef, error) {
|
||||
text := firstNonEmpty(item.Markdown, item.Content, item.Text)
|
||||
if text != "" {
|
||||
logger.Infof(context.Background(), "[MinerUCloud] parsed (inline), length=%d", len(text))
|
||||
return text, nil, nil
|
||||
}
|
||||
|
||||
if item.FullZipURL == "" {
|
||||
return "", nil, fmt.Errorf("MinerU Cloud state=done but no markdown/content or full_zip_url")
|
||||
}
|
||||
|
||||
md, imageRefs, err := downloadAndExtractZip(item.FullZipURL)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("extract zip: %w", err)
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[MinerUCloud] parsed (zip), markdown=%d chars, images=%d", len(md), len(imageRefs))
|
||||
return md, imageRefs, nil
|
||||
}
|
||||
|
||||
// --- ZIP handling ---
|
||||
|
||||
var imgRefPattern = regexp.MustCompile(`!\[[^\]]*\]\(([^)]+)\)`)
|
||||
|
||||
func downloadAndExtractZip(zipURL string) (string, []types.ImageRef, error) {
|
||||
if err := utils.ValidateURLForSSRF(zipURL); err != nil {
|
||||
return "", nil, fmt.Errorf("zip URL blocked by SSRF check: %v", err)
|
||||
}
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 120 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Get(zipURL)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("download zip: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", nil, fmt.Errorf("download zip status %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
zipData, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("read zip body: %w", err)
|
||||
}
|
||||
|
||||
zr, err := zip.NewReader(bytes.NewReader(zipData), int64(len(zipData)))
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("open zip: %w", err)
|
||||
}
|
||||
|
||||
// Find .md files
|
||||
var mdFiles []string
|
||||
entries := make(map[string]*zip.File)
|
||||
for _, f := range zr.File {
|
||||
entries[f.Name] = f
|
||||
if strings.HasSuffix(f.Name, ".md") {
|
||||
mdFiles = append(mdFiles, f.Name)
|
||||
}
|
||||
}
|
||||
if len(mdFiles) == 0 {
|
||||
return "", nil, fmt.Errorf("no .md file found in zip")
|
||||
}
|
||||
sort.Slice(mdFiles, func(i, j int) bool {
|
||||
di, dj := strings.Count(mdFiles[i], "/"), strings.Count(mdFiles[j], "/")
|
||||
if di != dj {
|
||||
return di < dj
|
||||
}
|
||||
return mdFiles[i] < mdFiles[j]
|
||||
})
|
||||
|
||||
mdText, err := readZipEntry(entries[mdFiles[0]])
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("read md file: %w", err)
|
||||
}
|
||||
|
||||
mdDir := filepath.Dir(mdFiles[0])
|
||||
|
||||
// Extract referenced images
|
||||
var imageRefs []types.ImageRef
|
||||
seen := map[string]bool{}
|
||||
for _, match := range imgRefPattern.FindAllStringSubmatch(mdText, -1) {
|
||||
imgPath := match[1]
|
||||
if strings.HasPrefix(imgPath, "http://") || strings.HasPrefix(imgPath, "https://") || strings.HasPrefix(imgPath, "data:") {
|
||||
continue
|
||||
}
|
||||
if seen[imgPath] {
|
||||
continue
|
||||
}
|
||||
seen[imgPath] = true
|
||||
|
||||
resolved := resolveInZip(imgPath, mdDir, entries)
|
||||
if resolved == nil {
|
||||
logger.Errorf(context.Background(), "[MinerUCloud] image not found in zip: %s", imgPath)
|
||||
continue
|
||||
}
|
||||
|
||||
imgData, err := readZipEntryBytes(resolved)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[MinerUCloud] failed to read zip image %s: %v", imgPath, err)
|
||||
continue
|
||||
}
|
||||
|
||||
ext := strings.ToLower(filepath.Ext(resolved.Name))
|
||||
if ext == "" {
|
||||
ext = ".png"
|
||||
}
|
||||
mimeType := mime.TypeByExtension(ext)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/png"
|
||||
}
|
||||
|
||||
imageRefs = append(imageRefs, types.ImageRef{
|
||||
Filename: filepath.Base(resolved.Name),
|
||||
OriginalRef: imgPath,
|
||||
MimeType: mimeType,
|
||||
ImageData: imgData,
|
||||
})
|
||||
}
|
||||
|
||||
return mdText, imageRefs, nil
|
||||
}
|
||||
|
||||
func resolveInZip(imgPath, mdDir string, entries map[string]*zip.File) *zip.File {
|
||||
normalized := strings.ReplaceAll(imgPath, "\\", "/")
|
||||
if f, ok := entries[normalized]; ok {
|
||||
return f
|
||||
}
|
||||
if mdDir != "" && mdDir != "." {
|
||||
joined := mdDir + "/" + normalized
|
||||
if f, ok := entries[joined]; ok {
|
||||
return f
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func readZipEntry(f *zip.File) (string, error) {
|
||||
rc, err := f.Open()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer rc.Close()
|
||||
data, err := io.ReadAll(rc)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
func readZipEntryBytes(f *zip.File) ([]byte, error) {
|
||||
rc, err := f.Open()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rc.Close()
|
||||
return io.ReadAll(rc)
|
||||
}
|
||||
|
||||
// PingMinerUCloud checks if the MinerU Cloud API is reachable with the given API key.
|
||||
func PingMinerUCloud(apiKey string) (bool, string) {
|
||||
apiKey = strings.TrimSpace(apiKey)
|
||||
if apiKey == "" {
|
||||
return false, "未配置 MinerU Cloud API Key"
|
||||
}
|
||||
|
||||
targetURL := defaultBaseURL + "/file-urls/batch"
|
||||
payload := []byte(`{"files":[],"model_version":"pipeline"}`)
|
||||
req, err := http.NewRequest(http.MethodPost, targetURL, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return false, fmt.Sprintf("构建请求失败: %v", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: 10 * time.Second,
|
||||
MaxRedirects: 5,
|
||||
})
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return false, fmt.Sprintf("MinerU Cloud 不可达: %v", err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
|
||||
if resp.StatusCode == 401 || resp.StatusCode == 403 {
|
||||
return false, "MinerU Cloud API Key 无效"
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
@@ -0,0 +1,472 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
const mineruTimeout = 1000 * time.Second // large docs can take a while
|
||||
|
||||
var (
|
||||
b64DataURIPattern = regexp.MustCompile(`^data:image/(\w+);base64,(.+)$`)
|
||||
minerUFileTypePattern = regexp.MustCompile(`^[a-z0-9]+$`)
|
||||
)
|
||||
|
||||
// MinerUReader calls a self-hosted MinerU API to read/convert documents.
|
||||
type MinerUReader struct {
|
||||
endpoint string
|
||||
backend string // "pipeline", "vlm-*", "hybrid-*"
|
||||
vlmServerURL string // vLLM server URL for vlm-http-client / hybrid-http-client
|
||||
formulaEnable bool
|
||||
tableEnable bool
|
||||
parseMethod string
|
||||
language string
|
||||
}
|
||||
|
||||
// NewMinerUReader creates a reader from ParserEngineOverrides.
|
||||
func NewMinerUReader(overrides map[string]string) *MinerUReader {
|
||||
var legacyOCREnabled *bool
|
||||
if raw, ok := overrides["mineru_enable_ocr"]; ok {
|
||||
value := parseBoolOr(raw, true)
|
||||
legacyOCREnabled = &value
|
||||
}
|
||||
|
||||
c := &MinerUReader{
|
||||
endpoint: strings.TrimRight(overrides["mineru_endpoint"], "/"),
|
||||
backend: stringOr(overrides["mineru_model"], "pipeline"),
|
||||
vlmServerURL: overrides["mineru_vlm_server_url"],
|
||||
formulaEnable: parseBoolOr(overrides["mineru_enable_formula"], true),
|
||||
tableEnable: parseBoolOr(overrides["mineru_enable_table"], true),
|
||||
parseMethod: types.ResolveMinerUParseMethod(overrides["mineru_parse_method"], legacyOCREnabled),
|
||||
language: stringOr(overrides["mineru_language"], "ch"),
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *MinerUReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
if c.endpoint == "" {
|
||||
return &types.ReadResult{Error: "MinerU endpoint is not configured"}, nil
|
||||
}
|
||||
if err := validateMinerUOutboundURL(c.endpoint); err != nil {
|
||||
return &types.ReadResult{Error: err.Error()}, nil
|
||||
}
|
||||
if c.vlmServerURL != "" {
|
||||
if err := validateMinerUOutboundURL(c.vlmServerURL); err != nil {
|
||||
return &types.ReadResult{Error: err.Error()}, nil
|
||||
}
|
||||
}
|
||||
|
||||
content := req.FileContent
|
||||
if len(content) == 0 {
|
||||
return &types.ReadResult{Error: "no file content provided"}, nil
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[MinerU] Parsing file=%s size=%d via %s", req.FileName, len(content), c.endpoint)
|
||||
|
||||
mdContent, imagesB64, err := c.callFileParse(ctx, content, req.FileName, req.FileType)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("MinerU file_parse: %w", err)
|
||||
}
|
||||
|
||||
// MinerU already returns markdown with embedded HTML blocks (e.g. <table>, <details>).
|
||||
// Re-running the whole document through html-to-markdown corrupts valid markdown
|
||||
// by escaping headings and image syntax. Only apply narrow compatibility fixes.
|
||||
mdContent = normalizeMinerUMarkdown(mdContent)
|
||||
|
||||
// Process images: decode base64, build ImageRef list, replace refs in markdown
|
||||
imageRefs, mdContent := c.processImages(mdContent, imagesB64)
|
||||
|
||||
mdContent, imageRefs = ensureOriginalImageRef(req, mdContent, imageRefs)
|
||||
|
||||
logger.Infof(context.Background(), "[MinerU] Parsed successfully, markdown=%d chars, images=%d", len(mdContent), len(imageRefs))
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: mdContent,
|
||||
ImageRefs: imageRefs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
type mineruFileEntry struct {
|
||||
MDContent string `json:"md_content"`
|
||||
Images map[string]string `json:"images"` // path -> "data:image/png;base64,..." or raw base64
|
||||
}
|
||||
|
||||
func minerUCleanFileType(fileType string) string {
|
||||
cleanType := strings.ToLower(strings.TrimSpace(fileType))
|
||||
cleanType = strings.TrimPrefix(cleanType, ".")
|
||||
if minerUFileTypePattern.MatchString(cleanType) {
|
||||
return cleanType
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func minerUUploadFileName(fileName, fileType string) string {
|
||||
cleanName := strings.TrimSpace(fileName)
|
||||
cleanName = strings.ReplaceAll(cleanName, `\`, "/")
|
||||
cleanName = path.Base(cleanName)
|
||||
cleanName = strings.Map(func(r rune) rune {
|
||||
if r == '\r' || r == '\n' || r == 0 {
|
||||
return -1
|
||||
}
|
||||
return r
|
||||
}, cleanName)
|
||||
cleanName = strings.TrimSpace(cleanName)
|
||||
if cleanName != "" && cleanName != "." && cleanName != ".." && cleanName != "/" {
|
||||
if filepath.Ext(cleanName) == "" {
|
||||
if cleanType := minerUCleanFileType(fileType); cleanType != "" {
|
||||
return cleanName + "." + cleanType
|
||||
}
|
||||
}
|
||||
return cleanName
|
||||
}
|
||||
|
||||
if cleanType := minerUCleanFileType(fileType); cleanType != "" {
|
||||
return "document." + cleanType
|
||||
}
|
||||
return "document"
|
||||
}
|
||||
|
||||
// minerUResultStem mirrors MinerU's upload.stem: the basename without extension.
|
||||
func minerUResultStem(uploadFileName string) string {
|
||||
stem := strings.TrimSuffix(path.Base(uploadFileName), filepath.Ext(uploadFileName))
|
||||
if stem == "" || stem == "." {
|
||||
return ""
|
||||
}
|
||||
return stem
|
||||
}
|
||||
|
||||
func minerUResultLookupKeys(uploadFileName string) []string {
|
||||
keys := make([]string, 0, 3)
|
||||
if stem := minerUResultStem(uploadFileName); stem != "" {
|
||||
keys = append(keys, stem)
|
||||
}
|
||||
return append(keys, "document", "files")
|
||||
}
|
||||
|
||||
func parseMinerUFileParseResponse(respBody []byte, uploadFileName string) (string, map[string]string, string, error) {
|
||||
var envelope struct {
|
||||
Results map[string]mineruFileEntry `json:"results"`
|
||||
}
|
||||
if err := json.Unmarshal(respBody, &envelope); err != nil {
|
||||
return "", nil, "", fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
if len(envelope.Results) == 0 {
|
||||
return "", nil, "", nil
|
||||
}
|
||||
|
||||
for _, key := range minerUResultLookupKeys(uploadFileName) {
|
||||
if entry, ok := envelope.Results[key]; ok {
|
||||
if entry.MDContent != "" || len(entry.Images) > 0 {
|
||||
return entry.MDContent, entry.Images, key, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for key, entry := range envelope.Results {
|
||||
if entry.MDContent != "" || len(entry.Images) > 0 {
|
||||
return entry.MDContent, entry.Images, key, nil
|
||||
}
|
||||
}
|
||||
return "", nil, "", nil
|
||||
}
|
||||
|
||||
func (c *MinerUReader) callFileParse(
|
||||
ctx context.Context,
|
||||
content []byte,
|
||||
fileName string,
|
||||
fileType string,
|
||||
) (string, map[string]string, error) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
|
||||
// Form fields
|
||||
fields := map[string]string{
|
||||
"return_md": "true",
|
||||
"return_images": "true",
|
||||
"table_enable": fmt.Sprintf("%v", c.tableEnable),
|
||||
"formula_enable": fmt.Sprintf("%v", c.formulaEnable),
|
||||
"parse_method": c.parseMethod,
|
||||
"start_page_id": "0",
|
||||
"end_page_id": "99999",
|
||||
"backend": c.backend,
|
||||
"response_format_zip": "false",
|
||||
"return_middle_json": "false",
|
||||
"return_model_output": "false",
|
||||
"return_content_list": "true",
|
||||
}
|
||||
if c.language != "" {
|
||||
fields["lang_list"] = c.language
|
||||
}
|
||||
if c.vlmServerURL != "" && (strings.HasPrefix(c.backend, "vlm-http-client") || strings.HasPrefix(c.backend, "hybrid-http-client")) {
|
||||
fields["server_url"] = c.vlmServerURL
|
||||
}
|
||||
for k, v := range fields {
|
||||
_ = writer.WriteField(k, v)
|
||||
}
|
||||
|
||||
uploadFileName := minerUUploadFileName(fileName, fileType)
|
||||
|
||||
// File part
|
||||
part, err := writer.CreateFormFile("files", uploadFileName)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("create form file: %w", err)
|
||||
}
|
||||
if _, err := part.Write(content); err != nil {
|
||||
return "", nil, fmt.Errorf("write file content: %w", err)
|
||||
}
|
||||
writer.Close()
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint+"/file_parse", &body)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: mineruTimeout,
|
||||
MaxRedirects: 5,
|
||||
})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("HTTP request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
return "", nil, fmt.Errorf("MinerU API status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("read response body: %w", err)
|
||||
}
|
||||
|
||||
// Dump raw response for debugging (truncate if too large)
|
||||
rawStr := string(respBody)
|
||||
if len(rawStr) > 4000 {
|
||||
logger.Infof(context.Background(), "[MinerU] Raw response (truncated to 4000 chars): %s ...", rawStr[:4000])
|
||||
} else {
|
||||
logger.Infof(context.Background(), "[MinerU] Raw response: %s", rawStr)
|
||||
}
|
||||
|
||||
// Also pretty-print the top-level structure (without large base64 blobs)
|
||||
var rawMap map[string]interface{}
|
||||
if err := json.Unmarshal(respBody, &rawMap); err == nil {
|
||||
c.logMinerUResponseStructure(rawMap, "")
|
||||
}
|
||||
|
||||
mdContent, imagesB64, resultKey, err := parseMinerUFileParseResponse(respBody, uploadFileName)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
if resultKey != "" {
|
||||
logger.Infof(context.Background(), "[MinerU] Using response path: results.%s", resultKey)
|
||||
return mdContent, imagesB64, nil
|
||||
}
|
||||
|
||||
logger.Errorf(context.Background(), "[MinerU] Response has no markdown/images under results")
|
||||
return "", nil, nil
|
||||
}
|
||||
|
||||
// processImages decodes base64 images from MinerU response and returns ImageRef list.
|
||||
// It also replaces image references in the markdown content.
|
||||
func (c *MinerUReader) processImages(mdContent string, imagesB64 map[string]string) ([]types.ImageRef, string) {
|
||||
var refs []types.ImageRef
|
||||
|
||||
for ipath, b64Str := range imagesB64 {
|
||||
matchedRefs := mineruImageOriginalRefs(mdContent, ipath)
|
||||
if len(matchedRefs) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
var imgBytes []byte
|
||||
var ext string
|
||||
|
||||
if m := b64DataURIPattern.FindStringSubmatch(b64Str); len(m) == 3 {
|
||||
ext = m[1]
|
||||
decoded, err := base64.StdEncoding.DecodeString(m[2])
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[MinerU] Failed to decode base64 image %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
imgBytes = decoded
|
||||
} else {
|
||||
// raw base64 without data URI prefix
|
||||
decoded, err := base64.StdEncoding.DecodeString(b64Str)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[MinerU] Failed to decode raw base64 image %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
imgBytes = decoded
|
||||
ext = strings.TrimPrefix(filepath.Ext(ipath), ".")
|
||||
if ext == "" {
|
||||
ext = "png"
|
||||
}
|
||||
}
|
||||
|
||||
mimeType := mime.TypeByExtension("." + ext)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/png"
|
||||
}
|
||||
|
||||
for _, originalRef := range matchedRefs {
|
||||
refs = append(refs, types.ImageRef{
|
||||
Filename: ipath,
|
||||
OriginalRef: originalRef,
|
||||
MimeType: mimeType,
|
||||
ImageData: imgBytes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return refs, mdContent
|
||||
}
|
||||
|
||||
// logMinerUResponseStructure logs the structure of the MinerU API response.
|
||||
func (c *MinerUReader) logMinerUResponseStructure(obj interface{}, prefix string) {
|
||||
logResponseStructure("MinerU", obj, prefix)
|
||||
}
|
||||
|
||||
// validateMinerUOutboundURL rejects MinerU endpoints that would reach private
|
||||
// or otherwise restricted hosts when parsed or probed from the app server.
|
||||
func validateMinerUOutboundURL(rawURL string) error {
|
||||
rawURL = strings.TrimSpace(rawURL)
|
||||
if rawURL == "" {
|
||||
return nil
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(rawURL); err != nil {
|
||||
return fmt.Errorf("MinerU URL blocked by SSRF check: %v", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// PingMinerU checks if the self-hosted MinerU service is reachable.
|
||||
func PingMinerU(endpoint string) (bool, string) {
|
||||
endpoint = strings.TrimRight(endpoint, "/")
|
||||
if endpoint == "" {
|
||||
return false, "未配置 MinerU 端点"
|
||||
}
|
||||
if err := validateMinerUOutboundURL(endpoint); err != nil {
|
||||
return false, err.Error()
|
||||
}
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: 5 * time.Second,
|
||||
MaxRedirects: 5,
|
||||
})
|
||||
resp, err := client.Get(endpoint + "/docs")
|
||||
if err != nil {
|
||||
return false, fmt.Sprintf("MinerU 服务不可达: %v", err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode >= 400 {
|
||||
return false, fmt.Sprintf("MinerU 服务返回状态 %d", resp.StatusCode)
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
|
||||
// escapedImageSyntaxPattern matches markdown image references whose `[` was
|
||||
// over-escaped to `\[` by html-to-markdown. The URL group mirrors the
|
||||
// downstream image-extraction regex so escapes are only stripped for actual
|
||||
// image references.
|
||||
var escapedImageSyntaxPattern = regexp.MustCompile(`!\\\[(.*?)\\?\]\(([^()\n]*(?:\([^)]*\)[^()\n]*)*)\)`)
|
||||
|
||||
// escapedHeadingPattern restores markdown headings that were over-escaped to
|
||||
// \# Heading. We only touch line-leading heading markers to avoid rewriting
|
||||
// body text that intentionally contains escaped # characters.
|
||||
var escapedHeadingPattern = regexp.MustCompile(`(?m)^\\(#{1,6})(\s+)`)
|
||||
|
||||
// unescapeMarkdownImageSyntax restores `` from html-to-markdown's
|
||||
// over-escaped `!\[alt\](url)` form. Without this, the downstream image regex
|
||||
// in ImageResolver fails to match and images are never persisted.
|
||||
func unescapeMarkdownImageSyntax(content string) string {
|
||||
return escapedImageSyntaxPattern.ReplaceAllString(content, "")
|
||||
}
|
||||
|
||||
func normalizeEscapedMarkdownHeadings(content string) string {
|
||||
return escapedHeadingPattern.ReplaceAllString(content, `$1$2`)
|
||||
}
|
||||
|
||||
func normalizeMinerUMarkdown(content string) string {
|
||||
content = unescapeMarkdownImageSyntax(content)
|
||||
content = normalizeEscapedMarkdownHeadings(content)
|
||||
return content
|
||||
}
|
||||
|
||||
func mineruImageOriginalRefs(mdContent, imagePath string) []string {
|
||||
normalizedTarget := normalizeMinerUImagePath(imagePath)
|
||||
if normalizedTarget == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
referenced := extractImageRefsFromContent(mdContent)
|
||||
seen := make(map[string]struct{}, len(referenced))
|
||||
var matched []string
|
||||
for _, refPath := range referenced {
|
||||
if normalizeMinerUImagePath(refPath) != normalizedTarget {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[refPath]; ok {
|
||||
continue
|
||||
}
|
||||
matched = append(matched, refPath)
|
||||
seen[refPath] = struct{}{}
|
||||
}
|
||||
|
||||
return matched
|
||||
}
|
||||
|
||||
// imgMarkdownPatternAllowSpaces matches markdown image syntax while allowing
|
||||
// spaces in the URL group, so that paths like "images/第 1 页.jpg" produced by
|
||||
// MinerU on Chinese documents are still detected as image references.
|
||||
var imgMarkdownPatternAllowSpaces = regexp.MustCompile(
|
||||
`!\[(.*?)\]\(([^()\n]*(?:\([^)]*\)[^()\n]*)*)\)`,
|
||||
)
|
||||
|
||||
func extractImageRefsFromContent(content string) []string {
|
||||
var refs []string
|
||||
|
||||
for _, match := range imgMarkdownPatternAllowSpaces.FindAllStringSubmatch(content, -1) {
|
||||
if len(match) >= 3 {
|
||||
refs = append(refs, strings.TrimSpace(match[2]))
|
||||
}
|
||||
}
|
||||
for _, match := range imgHTMLSrc.FindAllStringSubmatch(content, -1) {
|
||||
if len(match) >= 3 {
|
||||
refs = append(refs, match[2])
|
||||
}
|
||||
}
|
||||
|
||||
return refs
|
||||
}
|
||||
|
||||
func normalizeMinerUImagePath(p string) string {
|
||||
p = strings.TrimSpace(p)
|
||||
if p == "" {
|
||||
return ""
|
||||
}
|
||||
if decoded, err := url.PathUnescape(p); err == nil && decoded != "" {
|
||||
p = decoded
|
||||
}
|
||||
p = strings.TrimPrefix(p, "./")
|
||||
p = strings.TrimPrefix(p, "/")
|
||||
p = strings.TrimPrefix(p, "images/")
|
||||
return p
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestValidateMinerUOutboundURL_RejectsLoopback(t *testing.T) {
|
||||
err := validateMinerUOutboundURL("http://127.0.0.1:8080")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "SSRF")
|
||||
}
|
||||
|
||||
func TestPingMinerU_RejectsPrivateEndpoint(t *testing.T) {
|
||||
ok, msg := PingMinerU("http://127.0.0.1:8080")
|
||||
assert.False(t, ok)
|
||||
assert.Contains(t, msg, "SSRF")
|
||||
}
|
||||
@@ -0,0 +1,316 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
func TestNewMinerUReaderResolvesParseMethod(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
overrides map[string]string
|
||||
want string
|
||||
}{
|
||||
{name: "default is auto", overrides: map[string]string{}, want: "auto"},
|
||||
{name: "legacy enabled becomes auto", overrides: map[string]string{"mineru_enable_ocr": "true"}, want: "auto"},
|
||||
{name: "legacy disabled becomes text", overrides: map[string]string{"mineru_enable_ocr": "false"}, want: "txt"},
|
||||
{name: "explicit method wins", overrides: map[string]string{"mineru_parse_method": "ocr", "mineru_enable_ocr": "false"}, want: "ocr"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
reader := NewMinerUReader(tt.overrides)
|
||||
if reader.parseMethod != tt.want {
|
||||
t.Fatalf("parseMethod = %q, want %q", reader.parseMethod, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMinerUUploadFileName(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
fileName string
|
||||
fileType string
|
||||
want string
|
||||
}{
|
||||
{name: "original basename", fileName: "reports/Illustrator 研报.pdf", fileType: "pdf", want: "Illustrator 研报.pdf"},
|
||||
{
|
||||
name: "strip multipart controls",
|
||||
fileName: "reports/\r\nIllustrator\x00 export.pdf",
|
||||
fileType: "pdf",
|
||||
want: "Illustrator export.pdf",
|
||||
},
|
||||
{name: "append extension from type", fileName: "report", fileType: "pdf", want: "report.pdf"},
|
||||
{name: "fallback from bare type", fileType: "pdf", want: "document.pdf"},
|
||||
{name: "fallback from dotted uppercase type", fileType: ".PDF", want: "document.pdf"},
|
||||
{name: "fallback from path-only name", fileName: "/", fileType: "pdf", want: "document.pdf"},
|
||||
{name: "legacy fallback without metadata", want: "document"},
|
||||
{name: "reject unsafe fallback type", fileType: "../pdf", want: "document"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := minerUUploadFileName(tt.fileName, tt.fileType); got != tt.want {
|
||||
t.Fatalf("minerUUploadFileName() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseMinerUFileParseResponse(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
uploadFileName string
|
||||
response map[string]any
|
||||
wantMarkdown string
|
||||
wantResultKey string
|
||||
}{
|
||||
{
|
||||
name: "match upload stem",
|
||||
uploadFileName: "Illustrator 研报.pdf",
|
||||
response: map[string]any{
|
||||
"results": map[string]any{
|
||||
"Illustrator 研报": map[string]any{"md_content": "from-stem", "images": map[string]any{}},
|
||||
},
|
||||
},
|
||||
wantMarkdown: "from-stem",
|
||||
wantResultKey: "Illustrator 研报",
|
||||
},
|
||||
{
|
||||
name: "legacy document key",
|
||||
uploadFileName: "document",
|
||||
response: map[string]any{
|
||||
"results": map[string]any{
|
||||
"document": map[string]any{"md_content": "legacy", "images": map[string]any{}},
|
||||
},
|
||||
},
|
||||
wantMarkdown: "legacy",
|
||||
wantResultKey: "document",
|
||||
},
|
||||
{
|
||||
name: "files variant",
|
||||
uploadFileName: "document.pdf",
|
||||
response: map[string]any{
|
||||
"results": map[string]any{
|
||||
"files": map[string]any{"md_content": "from-files", "images": map[string]any{}},
|
||||
},
|
||||
},
|
||||
wantMarkdown: "from-files",
|
||||
wantResultKey: "files",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
respBody, err := json.Marshal(tt.response)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal response: %v", err)
|
||||
}
|
||||
gotMarkdown, _, gotKey, err := parseMinerUFileParseResponse(respBody, tt.uploadFileName)
|
||||
if err != nil {
|
||||
t.Fatalf("parseMinerUFileParseResponse() error: %v", err)
|
||||
}
|
||||
if gotMarkdown != tt.wantMarkdown {
|
||||
t.Fatalf("markdown = %q, want %q", gotMarkdown, tt.wantMarkdown)
|
||||
}
|
||||
if gotKey != tt.wantResultKey {
|
||||
t.Fatalf("result key = %q, want %q", gotKey, tt.wantResultKey)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMinerUReaderPreservesMultipartFilename(t *testing.T) {
|
||||
utils.SetSSRFWhitelistFromRaw("127.0.0.1")
|
||||
t.Cleanup(func() {
|
||||
utils.SetSSRFWhitelistFromRaw("")
|
||||
})
|
||||
|
||||
const fileContent = "%PDF-1.7 synthetic test content"
|
||||
tests := []struct {
|
||||
name string
|
||||
fileName string
|
||||
fileType string
|
||||
want string
|
||||
}{
|
||||
{name: "original basename", fileName: "reports/Illustrator 研报.pdf", fileType: "pdf", want: "Illustrator 研报.pdf"},
|
||||
{
|
||||
name: "strip multipart controls",
|
||||
fileName: "reports/\r\nIllustrator\x00 export.pdf",
|
||||
fileType: "pdf",
|
||||
want: "Illustrator export.pdf",
|
||||
},
|
||||
{name: "fallback from bare type", fileType: "pdf", want: "document.pdf"},
|
||||
{name: "fallback from dotted uppercase type", fileType: ".PDF", want: "document.pdf"},
|
||||
{name: "fallback from path-only name", fileName: "/", fileType: "pdf", want: "document.pdf"},
|
||||
{name: "legacy fallback without metadata", want: "document"},
|
||||
{name: "reject unsafe fallback type", fileType: "../pdf", want: "document"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
type capturedRequest struct {
|
||||
method string
|
||||
path string
|
||||
fileName string
|
||||
content string
|
||||
err error
|
||||
}
|
||||
|
||||
captured := make(chan capturedRequest, 1)
|
||||
uploadName := minerUUploadFileName(tt.fileName, tt.fileType)
|
||||
resultStem := minerUResultStem(uploadName)
|
||||
responseBody, err := json.Marshal(map[string]any{
|
||||
"results": map[string]any{
|
||||
resultStem: map[string]any{
|
||||
"md_content": "ok",
|
||||
"images": map[string]any{},
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("marshal response: %v", err)
|
||||
}
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
got := capturedRequest{method: r.Method, path: r.URL.Path}
|
||||
file, header, err := r.FormFile("files")
|
||||
if err != nil {
|
||||
got.err = err
|
||||
} else {
|
||||
got.fileName = header.Filename
|
||||
content, readErr := io.ReadAll(file)
|
||||
got.content = string(content)
|
||||
got.err = readErr
|
||||
_ = file.Close()
|
||||
}
|
||||
captured <- got
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write(responseBody)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
reader := NewMinerUReader(map[string]string{"mineru_endpoint": server.URL})
|
||||
result, err := reader.Read(context.Background(), &types.ReadRequest{
|
||||
FileContent: []byte(fileContent),
|
||||
FileName: tt.fileName,
|
||||
FileType: tt.fileType,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Read returned error: %v", err)
|
||||
}
|
||||
if result.MarkdownContent != "ok" {
|
||||
t.Fatalf("MarkdownContent = %q, want %q", result.MarkdownContent, "ok")
|
||||
}
|
||||
|
||||
request := <-captured
|
||||
if request.err != nil {
|
||||
t.Fatalf("capture multipart request: %v", request.err)
|
||||
}
|
||||
if request.method != http.MethodPost {
|
||||
t.Errorf("method = %q, want %q", request.method, http.MethodPost)
|
||||
}
|
||||
if request.path != "/file_parse" {
|
||||
t.Errorf("path = %q, want %q", request.path, "/file_parse")
|
||||
}
|
||||
if request.fileName != tt.want {
|
||||
t.Errorf("multipart filename = %q, want %q", request.fileName, tt.want)
|
||||
}
|
||||
if request.content != fileContent {
|
||||
t.Errorf("multipart content = %q, want %q", request.content, fileContent)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeMinerUMarkdownPreservesMarkdownAndHTML(t *testing.T) {
|
||||
input := strings.Join([]string{
|
||||
"# Heading",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
`<details><summary>text_image</summary>caption</details>`,
|
||||
"",
|
||||
`<table><tr><td><img src="images/profile.jpg"/></td></tr></table>`,
|
||||
}, "\n")
|
||||
|
||||
got := normalizeMinerUMarkdown(input)
|
||||
|
||||
if !strings.Contains(got, "# Heading") {
|
||||
t.Fatalf("expected heading to stay intact, got: %q", got)
|
||||
}
|
||||
if strings.Contains(got, `\# Heading`) {
|
||||
t.Fatalf("expected heading to avoid escaped form, got: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "") {
|
||||
t.Fatalf("expected markdown image syntax to stay intact, got: %q", got)
|
||||
}
|
||||
if strings.Contains(got, `!\[](images/cover.jpg)`) {
|
||||
t.Fatalf("expected markdown image syntax to avoid escaped form, got: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, `<details><summary>text_image</summary>caption</details>`) {
|
||||
t.Fatalf("expected details/summary block to be preserved, got: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, `<img src="images/profile.jpg"/>`) {
|
||||
t.Fatalf("expected html img tag to be preserved, got: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessImagesKeepsReferencedVariants(t *testing.T) {
|
||||
reader := &MinerUReader{}
|
||||
mdContent := strings.Join([]string{
|
||||
"",
|
||||
`<img src="./images/profile.jpg"/>`,
|
||||
``,
|
||||
}, "\n")
|
||||
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
images := map[string]string{
|
||||
"cover.jpg": "data:image/png;base64," + b64,
|
||||
"profile.jpg": "data:image/png;base64," + b64,
|
||||
"plain.jpg": "data:image/png;base64," + b64,
|
||||
}
|
||||
|
||||
refs, gotMarkdown := reader.processImages(mdContent, images)
|
||||
|
||||
if gotMarkdown != mdContent {
|
||||
t.Fatalf("processImages should not rewrite markdown content")
|
||||
}
|
||||
if len(refs) != 3 {
|
||||
t.Fatalf("expected 3 image refs, got %d", len(refs))
|
||||
}
|
||||
}
|
||||
|
||||
// TestProcessImagesMatchesPathsWithSpaces guards against a regression where
|
||||
// MinerU image filenames containing spaces (common on Chinese documents,
|
||||
// e.g. "images/第 1 页.jpg") would be silently dropped because the markdown
|
||||
// regex used to extract refs disallowed whitespace inside the URL group.
|
||||
func TestProcessImagesMatchesPathsWithSpaces(t *testing.T) {
|
||||
reader := &MinerUReader{}
|
||||
mdContent := ""
|
||||
|
||||
png := createTestPNG(200, 150)
|
||||
b64 := base64.StdEncoding.EncodeToString(png)
|
||||
images := map[string]string{
|
||||
"第 1 页.jpg": "data:image/png;base64," + b64,
|
||||
}
|
||||
|
||||
refs, _ := reader.processImages(mdContent, images)
|
||||
if len(refs) != 1 {
|
||||
t.Fatalf("expected 1 image ref for path with spaces, got %d", len(refs))
|
||||
}
|
||||
if refs[0].OriginalRef != "images/第 1 页.jpg" {
|
||||
t.Fatalf("unexpected OriginalRef: %q", refs[0].OriginalRef)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,365 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
const (
|
||||
paddleOCRVLCloudDefaultBaseURL = "https://paddleocr.aistudio-app.com/api/v2/ocr/jobs"
|
||||
paddleOCRVLCloudDefaultModel = "PaddleOCR-VL-1.6"
|
||||
paddleOCRVLCloudPollInterval = 5 * time.Second
|
||||
paddleOCRVLCloudTimeout = 600 * time.Second
|
||||
)
|
||||
|
||||
// PaddleOCRVLCloudReader calls the PaddleOCR-VL AI Studio cloud API.
|
||||
// Flow: POST /jobs (multipart) → poll GET /jobs/{id} → download result JSONL,
|
||||
// then fetch each referenced image URL.
|
||||
type PaddleOCRVLCloudReader struct {
|
||||
token string
|
||||
baseURL string
|
||||
model string
|
||||
useSeal bool
|
||||
useChart bool
|
||||
}
|
||||
|
||||
// NewPaddleOCRVLCloudReader creates a reader from ParserEngineOverrides.
|
||||
func NewPaddleOCRVLCloudReader(overrides map[string]string) *PaddleOCRVLCloudReader {
|
||||
return &PaddleOCRVLCloudReader{
|
||||
token: strings.TrimSpace(overrides["paddleocr_vl_cloud_token"]),
|
||||
baseURL: strings.TrimRight(stringOr(overrides["paddleocr_vl_cloud_base_url"], paddleOCRVLCloudDefaultBaseURL), "/"),
|
||||
model: stringOr(overrides["paddleocr_vl_cloud_model"], paddleOCRVLCloudDefaultModel),
|
||||
useSeal: parseBoolOr(overrides["paddleocr_vl_cloud_use_seal_recognition"], true),
|
||||
useChart: parseBoolOr(overrides["paddleocr_vl_cloud_use_chart_recognition"], false),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLCloudReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
if c.token == "" {
|
||||
return &types.ReadResult{Error: "PaddleOCR-VL Cloud token is not configured"}, nil
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(c.baseURL); err != nil {
|
||||
return &types.ReadResult{Error: fmt.Sprintf("PaddleOCR-VL Cloud base URL blocked by SSRF policy: %v", err)}, nil
|
||||
}
|
||||
|
||||
content := req.FileContent
|
||||
if len(content) == 0 {
|
||||
return &types.ReadResult{Error: "no file content provided"}, nil
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL Cloud] Parsing file=%s size=%d model=%s",
|
||||
req.FileName, len(content), c.model)
|
||||
|
||||
jobID, err := c.submitJob(ctx, req, content)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("PaddleOCR-VL Cloud submit: %w", err)
|
||||
}
|
||||
|
||||
jsonlURL, err := c.pollJob(ctx, jobID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("PaddleOCR-VL Cloud poll: %w", err)
|
||||
}
|
||||
|
||||
mdContent, imagesURL, err := c.fetchResults(jsonlURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("PaddleOCR-VL Cloud fetch results: %w", err)
|
||||
}
|
||||
|
||||
mdContent = normalizeHTMLTables(mdContent)
|
||||
|
||||
imageRefs := c.downloadImages(mdContent, imagesURL)
|
||||
mdContent, imageRefs = ensureOriginalImageRef(req, mdContent, imageRefs)
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL Cloud] Parsed successfully, markdown=%d chars, images=%d",
|
||||
len(mdContent), len(imageRefs))
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: mdContent,
|
||||
ImageRefs: imageRefs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLCloudReader) optionalPayload() map[string]interface{} {
|
||||
// Shared with the self-hosted engine so both produce identical output.
|
||||
return paddleOCRVLRecognitionParams(c.useSeal, c.useChart)
|
||||
}
|
||||
|
||||
// --- job submit ---
|
||||
|
||||
type paddleOCRVLCloudSubmitResponse struct {
|
||||
Data struct {
|
||||
JobID string `json:"jobId"`
|
||||
} `json:"data"`
|
||||
ErrorCode int `json:"errorCode"`
|
||||
ErrorMsg string `json:"errorMsg"`
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLCloudReader) submitJob(ctx context.Context, req *types.ReadRequest, content []byte) (string, error) {
|
||||
optional, err := json.Marshal(c.optionalPayload())
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("marshal optionalPayload: %w", err)
|
||||
}
|
||||
|
||||
fileName := req.FileName
|
||||
if fileName == "" {
|
||||
ext := strings.TrimPrefix(req.FileType, ".")
|
||||
if ext == "" {
|
||||
ext = "pdf"
|
||||
}
|
||||
fileName = "document." + ext
|
||||
}
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("model", c.model)
|
||||
_ = writer.WriteField("optionalPayload", string(optional))
|
||||
part, err := writer.CreateFormFile("file", filepath.Base(fileName))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create form file: %w", err)
|
||||
}
|
||||
if _, err := part.Write(content); err != nil {
|
||||
return "", fmt.Errorf("write file content: %w", err)
|
||||
}
|
||||
writer.Close()
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL, &body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Authorization", "bearer "+c.token)
|
||||
httpReq.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 60 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("HTTP request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("API status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result paddleOCRVLCloudSubmitResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return "", fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
if result.Data.JobID == "" {
|
||||
return "", fmt.Errorf("API returned no jobId: %s", string(respBody))
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL Cloud] job submitted: jobId=%s", result.Data.JobID)
|
||||
return result.Data.JobID, nil
|
||||
}
|
||||
|
||||
// --- polling ---
|
||||
|
||||
type paddleOCRVLCloudPollResponse struct {
|
||||
Data struct {
|
||||
State string `json:"state"`
|
||||
ErrorMsg string `json:"errorMsg"`
|
||||
ExtractProgress struct {
|
||||
TotalPages int `json:"totalPages"`
|
||||
ExtractedPages int `json:"extractedPages"`
|
||||
} `json:"extractProgress"`
|
||||
ResultURL struct {
|
||||
JSONURL string `json:"jsonUrl"`
|
||||
} `json:"resultUrl"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLCloudReader) pollJob(ctx context.Context, jobID string) (string, error) {
|
||||
deadline := time.Now().Add(paddleOCRVLCloudTimeout)
|
||||
pollCount := 0
|
||||
url := c.baseURL + "/" + jobID
|
||||
|
||||
for time.Now().Before(deadline) {
|
||||
// Bail out promptly when the caller cancels (task cancelled / timed
|
||||
// out) instead of spinning: client.Do would fail immediately and
|
||||
// sleepCtx returns at once on a cancelled ctx, so without this guard
|
||||
// the loop busy-hammers the cloud API and floods logs until deadline.
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
pollCount++
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create poll request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Authorization", "bearer "+c.token)
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 30 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] poll #%d failed: %v", pollCount, err)
|
||||
sleepCtx(ctx, paddleOCRVLCloudPollInterval)
|
||||
continue
|
||||
}
|
||||
respBody, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] poll #%d status %d: %s", pollCount, resp.StatusCode, string(respBody))
|
||||
sleepCtx(ctx, paddleOCRVLCloudPollInterval)
|
||||
continue
|
||||
}
|
||||
|
||||
var pollResp paddleOCRVLCloudPollResponse
|
||||
if err := json.Unmarshal(respBody, &pollResp); err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] poll #%d decode error: %v", pollCount, err)
|
||||
sleepCtx(ctx, paddleOCRVLCloudPollInterval)
|
||||
continue
|
||||
}
|
||||
|
||||
state := strings.ToLower(pollResp.Data.State)
|
||||
if pollCount == 1 || pollCount%6 == 0 || state == "done" || state == "failed" {
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL Cloud] poll #%d: state=%s pages=%d/%d",
|
||||
pollCount, state, pollResp.Data.ExtractProgress.ExtractedPages, pollResp.Data.ExtractProgress.TotalPages)
|
||||
}
|
||||
|
||||
switch state {
|
||||
case "done":
|
||||
if pollResp.Data.ResultURL.JSONURL == "" {
|
||||
return "", fmt.Errorf("state=done but no jsonUrl")
|
||||
}
|
||||
return pollResp.Data.ResultURL.JSONURL, nil
|
||||
case "failed":
|
||||
return "", fmt.Errorf("task failed: %s", pollResp.Data.ErrorMsg)
|
||||
}
|
||||
|
||||
sleepCtx(ctx, paddleOCRVLCloudPollInterval)
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("task timed out after %d polls", pollCount)
|
||||
}
|
||||
|
||||
// --- result parsing ---
|
||||
|
||||
type paddleOCRVLCloudResultLine struct {
|
||||
Result struct {
|
||||
LayoutParsingResults []struct {
|
||||
Markdown struct {
|
||||
Text string `json:"text"`
|
||||
Images map[string]string `json:"images"`
|
||||
} `json:"markdown"`
|
||||
} `json:"layoutParsingResults"`
|
||||
} `json:"result"`
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLCloudReader) fetchResults(jsonlURL string) (string, map[string]string, error) {
|
||||
if err := utils.ValidateURLForSSRF(jsonlURL); err != nil {
|
||||
return "", nil, fmt.Errorf("jsonl URL blocked by SSRF check: %v", err)
|
||||
}
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 120 * time.Second, MaxRedirects: 5})
|
||||
resp, err := client.Get(jsonlURL)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("download jsonl: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", nil, fmt.Errorf("download jsonl status %d", resp.StatusCode)
|
||||
}
|
||||
data, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("read jsonl body: %w", err)
|
||||
}
|
||||
|
||||
texts := make([]string, 0)
|
||||
images := make(map[string]string)
|
||||
for _, line := range strings.Split(strings.TrimSpace(string(data)), "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
var parsed paddleOCRVLCloudResultLine
|
||||
if err := json.Unmarshal([]byte(line), &parsed); err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] skip malformed jsonl line: %v", err)
|
||||
continue
|
||||
}
|
||||
for _, p := range parsed.Result.LayoutParsingResults {
|
||||
if t := strings.TrimSpace(p.Markdown.Text); t != "" {
|
||||
texts = append(texts, p.Markdown.Text)
|
||||
}
|
||||
for path, u := range p.Markdown.Images {
|
||||
if _, ok := images[path]; !ok {
|
||||
images[path] = u
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL Cloud] fetched %d page(s), images=%d", len(texts), len(images))
|
||||
return strings.Join(texts, "\n\n"), images, nil
|
||||
}
|
||||
|
||||
// downloadImages fetches each referenced image URL and builds ImageRef entries.
|
||||
func (c *PaddleOCRVLCloudReader) downloadImages(mdContent string, imagesURL map[string]string) []types.ImageRef {
|
||||
var refs []types.ImageRef
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{Timeout: 60 * time.Second, MaxRedirects: 5})
|
||||
|
||||
for ipath, u := range imagesURL {
|
||||
matchedRefs := mineruImageOriginalRefs(mdContent, ipath)
|
||||
if len(matchedRefs) == 0 {
|
||||
continue
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(u); err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] image URL blocked %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
resp, err := client.Get(u)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] download image %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
imgBytes, err := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if err != nil || resp.StatusCode != http.StatusOK {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL Cloud] read image %s status=%d err=%v", ipath, resp.StatusCode, err)
|
||||
continue
|
||||
}
|
||||
|
||||
ext := strings.TrimPrefix(filepath.Ext(ipath), ".")
|
||||
if ext == "" {
|
||||
ext = "png"
|
||||
}
|
||||
mimeType := mime.TypeByExtension("." + ext)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/png"
|
||||
}
|
||||
|
||||
for _, originalRef := range matchedRefs {
|
||||
refs = append(refs, types.ImageRef{
|
||||
Filename: ipath,
|
||||
OriginalRef: originalRef,
|
||||
MimeType: mimeType,
|
||||
ImageData: imgBytes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return refs
|
||||
}
|
||||
|
||||
// PingPaddleOCRVLCloud checks whether the cloud token is present (the API has
|
||||
// no lightweight health endpoint, so we only validate configuration here).
|
||||
func PingPaddleOCRVLCloud(token string) (bool, string) {
|
||||
if strings.TrimSpace(token) == "" {
|
||||
return false, "未配置 PaddleOCR-VL Cloud Token"
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
@@ -0,0 +1,292 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
const paddleOCRVLTimeout = 1000 * time.Second // large scanned PDFs can take a while
|
||||
|
||||
// PaddleOCRVLReader calls a self-hosted PaddleOCR-VL pipeline service
|
||||
// (the full document-parsing API, not the bare VLM inference server).
|
||||
//
|
||||
// Flow: POST {endpoint}/layout-parsing with base64 file → synchronous JSON
|
||||
// response containing per-page markdown + inline base64 images.
|
||||
type PaddleOCRVLReader struct {
|
||||
endpoint string
|
||||
useSeal bool
|
||||
useChart bool
|
||||
}
|
||||
|
||||
// NewPaddleOCRVLReader creates a reader from ParserEngineOverrides.
|
||||
func NewPaddleOCRVLReader(overrides map[string]string) *PaddleOCRVLReader {
|
||||
return &PaddleOCRVLReader{
|
||||
endpoint: strings.TrimRight(overrides["paddleocr_vl_endpoint"], "/"),
|
||||
useSeal: parseBoolOr(overrides["paddleocr_vl_use_seal_recognition"], true),
|
||||
useChart: parseBoolOr(overrides["paddleocr_vl_use_chart_recognition"], false),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
if c.endpoint == "" {
|
||||
return &types.ReadResult{Error: "PaddleOCR-VL endpoint is not configured"}, nil
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(c.endpoint); err != nil {
|
||||
return &types.ReadResult{Error: fmt.Sprintf("PaddleOCR-VL endpoint blocked by SSRF policy: %v", err)}, nil
|
||||
}
|
||||
|
||||
content := req.FileContent
|
||||
if len(content) == 0 {
|
||||
return &types.ReadResult{Error: "no file content provided"}, nil
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL] Parsing file=%s size=%d via %s",
|
||||
req.FileName, len(content), c.endpoint)
|
||||
|
||||
mdContent, imagesB64, err := c.callLayoutParsing(ctx, req, content)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("PaddleOCR-VL layout-parsing: %w", err)
|
||||
}
|
||||
|
||||
// PaddleOCR-VL renders tables as styled HTML (per-cell text-align), which
|
||||
// wastes tokens and defeats the chunker's table-protection logic. Convert
|
||||
// them to Markdown tables (or strip layout attributes when conversion is
|
||||
// not possible) before downstream processing.
|
||||
mdContent = normalizeHTMLTables(mdContent)
|
||||
|
||||
imageRefs, mdContent := c.processImages(mdContent, imagesB64)
|
||||
mdContent, imageRefs = ensureOriginalImageRef(req, mdContent, imageRefs)
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL] Parsed successfully, markdown=%d chars, images=%d",
|
||||
len(mdContent), len(imageRefs))
|
||||
|
||||
return &types.ReadResult{
|
||||
MarkdownContent: mdContent,
|
||||
ImageRefs: imageRefs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// paddleOCRVLRecognitionParams returns the recognition / page-restructuring
|
||||
// parameters shared by the self-hosted (/layout-parsing, top-level body) and
|
||||
// cloud (optionalPayload) request bodies. Keeping both identical ensures the
|
||||
// self-hosted engine reproduces the cloud output: cross-page table merging,
|
||||
// multi-level heading reconstruction, header/footer stripping, and the same
|
||||
// sampling / resolution settings used by the AI Studio service.
|
||||
func paddleOCRVLRecognitionParams(useSeal, useChart bool) map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"markdownIgnoreLabels": []string{
|
||||
"header", "header_image", "footer", "footer_image",
|
||||
"number", "footnote", "aside_text",
|
||||
},
|
||||
"useDocOrientationClassify": false,
|
||||
"useDocUnwarping": false,
|
||||
"useLayoutDetection": true,
|
||||
"useChartRecognition": useChart,
|
||||
"useSealRecognition": useSeal,
|
||||
"useOcrForImageBlock": false,
|
||||
"mergeTables": true,
|
||||
"relevelTitles": true,
|
||||
"restructurePages": true,
|
||||
"layoutShapeMode": "auto",
|
||||
"promptLabel": "ocr",
|
||||
"layoutNms": true,
|
||||
"repetitionPenalty": 1,
|
||||
"temperature": 0,
|
||||
"topP": 1,
|
||||
"minPixels": 147384,
|
||||
"maxPixels": 2822400,
|
||||
}
|
||||
}
|
||||
|
||||
// fileTypeCode maps a request to the PaddleOCR-VL fileType field:
|
||||
// 0 = PDF, 1 = image (including TIFF).
|
||||
func fileTypeCode(req *types.ReadRequest) int {
|
||||
ft := strings.ToLower(strings.TrimPrefix(req.FileType, "."))
|
||||
if ft == "" {
|
||||
ft = strings.TrimPrefix(strings.ToLower(filepath.Ext(req.FileName)), ".")
|
||||
}
|
||||
if ft == "pdf" {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
// paddleOCRVLResponse mirrors the relevant fields of the PaddleX serving
|
||||
// /layout-parsing response. The service returns one entry per page.
|
||||
type paddleOCRVLResponse struct {
|
||||
ErrorCode int `json:"errorCode"`
|
||||
ErrorMsg string `json:"errorMsg"`
|
||||
Result struct {
|
||||
LayoutParsingResults []struct {
|
||||
Markdown struct {
|
||||
Text string `json:"text"`
|
||||
Images map[string]string `json:"images"`
|
||||
} `json:"markdown"`
|
||||
} `json:"layoutParsingResults"`
|
||||
} `json:"result"`
|
||||
}
|
||||
|
||||
func (c *PaddleOCRVLReader) callLayoutParsing(
|
||||
ctx context.Context, req *types.ReadRequest, content []byte,
|
||||
) (string, map[string]string, error) {
|
||||
payload := paddleOCRVLRecognitionParams(c.useSeal, c.useChart)
|
||||
payload["file"] = base64.StdEncoding.EncodeToString(content)
|
||||
payload["fileType"] = fileTypeCode(req)
|
||||
payload["visualize"] = false
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("marshal payload: %w", err)
|
||||
}
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(
|
||||
ctx, http.MethodPost, c.endpoint+"/layout-parsing", bytes.NewReader(body),
|
||||
)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: paddleOCRVLTimeout,
|
||||
MaxRedirects: 5,
|
||||
})
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("HTTP request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", nil, fmt.Errorf("read response body: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", nil, fmt.Errorf("PaddleOCR-VL API status %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result paddleOCRVLResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return "", nil, fmt.Errorf("decode response: %w", err)
|
||||
}
|
||||
if result.ErrorCode != 0 {
|
||||
return "", nil, fmt.Errorf("PaddleOCR-VL error %d: %s", result.ErrorCode, result.ErrorMsg)
|
||||
}
|
||||
|
||||
pages := result.Result.LayoutParsingResults
|
||||
if len(pages) == 0 {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL] response has no layoutParsingResults")
|
||||
return "", nil, nil
|
||||
}
|
||||
|
||||
// Merge per-page markdown and image dicts into one document.
|
||||
texts := make([]string, 0, len(pages))
|
||||
images := make(map[string]string)
|
||||
for _, p := range pages {
|
||||
if t := strings.TrimSpace(p.Markdown.Text); t != "" {
|
||||
texts = append(texts, p.Markdown.Text)
|
||||
}
|
||||
for path, data := range p.Markdown.Images {
|
||||
if _, ok := images[path]; !ok {
|
||||
images[path] = data
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.Infof(context.Background(), "[PaddleOCR-VL] parsed %d page(s), images=%d", len(pages), len(images))
|
||||
return strings.Join(texts, "\n\n"), images, nil
|
||||
}
|
||||
|
||||
// processImages decodes the inline base64 images returned by PaddleOCR-VL and
|
||||
// builds ImageRef entries, matching them against references in the markdown.
|
||||
func (c *PaddleOCRVLReader) processImages(
|
||||
mdContent string, imagesB64 map[string]string,
|
||||
) ([]types.ImageRef, string) {
|
||||
var refs []types.ImageRef
|
||||
|
||||
for ipath, b64Str := range imagesB64 {
|
||||
matchedRefs := mineruImageOriginalRefs(mdContent, ipath)
|
||||
if len(matchedRefs) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
var imgBytes []byte
|
||||
var ext string
|
||||
if m := b64DataURIPattern.FindStringSubmatch(b64Str); len(m) == 3 {
|
||||
ext = m[1]
|
||||
decoded, err := base64.StdEncoding.DecodeString(m[2])
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL] decode base64 image %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
imgBytes = decoded
|
||||
} else {
|
||||
decoded, err := base64.StdEncoding.DecodeString(b64Str)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[PaddleOCR-VL] decode raw base64 image %s: %v", ipath, err)
|
||||
continue
|
||||
}
|
||||
imgBytes = decoded
|
||||
ext = strings.TrimPrefix(filepath.Ext(ipath), ".")
|
||||
if ext == "" {
|
||||
ext = "png"
|
||||
}
|
||||
}
|
||||
|
||||
mimeType := mime.TypeByExtension("." + ext)
|
||||
if mimeType == "" {
|
||||
mimeType = "image/png"
|
||||
}
|
||||
|
||||
for _, originalRef := range matchedRefs {
|
||||
refs = append(refs, types.ImageRef{
|
||||
Filename: ipath,
|
||||
OriginalRef: originalRef,
|
||||
MimeType: mimeType,
|
||||
ImageData: imgBytes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return refs, mdContent
|
||||
}
|
||||
|
||||
// PingPaddleOCRVL checks whether a self-hosted PaddleOCR-VL service is reachable.
|
||||
func PingPaddleOCRVL(endpoint string) (bool, string) {
|
||||
endpoint = strings.TrimRight(endpoint, "/")
|
||||
if endpoint == "" {
|
||||
return false, "未配置 PaddleOCR-VL 端点"
|
||||
}
|
||||
if err := utils.ValidateURLForSSRF(endpoint); err != nil {
|
||||
return false, fmt.Sprintf("PaddleOCR-VL 端点未通过 SSRF 校验: %v", err)
|
||||
}
|
||||
client := utils.NewSSRFSafeHTTPClient(utils.SSRFSafeHTTPClientConfig{
|
||||
Timeout: 5 * time.Second,
|
||||
MaxRedirects: 5,
|
||||
})
|
||||
// The pipeline only exposes POST /layout-parsing; an empty GET should still
|
||||
// produce a routed HTTP response (e.g. 404/405) when the service is up.
|
||||
resp, err := client.Get(endpoint + "/layout-parsing")
|
||||
if err != nil {
|
||||
return false, fmt.Sprintf("PaddleOCR-VL 服务不可达: %v", err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode >= 500 {
|
||||
return false, fmt.Sprintf("PaddleOCR-VL 服务返回状态 %d", resp.StatusCode)
|
||||
}
|
||||
return true, ""
|
||||
}
|
||||
@@ -0,0 +1,406 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
// remoteImageServer serves a valid PNG for every request and counts the hits so
|
||||
// that budget and dedup behaviour can be asserted directly.
|
||||
func remoteImageServer(t *testing.T) (*httptest.Server, *int64) {
|
||||
t.Helper()
|
||||
var hits int64
|
||||
png := createTestPNG(200, 200)
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt64(&hits, 1)
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(png)
|
||||
}))
|
||||
t.Cleanup(ts.Close)
|
||||
return ts, &hits
|
||||
}
|
||||
|
||||
func allowLocalhost(t *testing.T) {
|
||||
t.Helper()
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
secutils.ResetSSRFWhitelistForTest()
|
||||
}
|
||||
|
||||
func resolve(t *testing.T, markdown string) (string, []StoredImage, *mockFileService) {
|
||||
t.Helper()
|
||||
fSvc := &mockFileService{}
|
||||
updated, images, err := NewImageResolver().ResolveRemoteImages(
|
||||
context.Background(), markdown, fSvc, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
return updated, images, fSvc
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_HTMLTagRewritesSrcAndKeepsTag(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
markdown := fmt.Sprintf(
|
||||
`<p><img class="screenshot" src="%s/a.png" alt="login screen" width="600"></p>`, ts.URL)
|
||||
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected 1 stored image, got %d", len(images))
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("original URL still present: %s", updated)
|
||||
}
|
||||
if !strings.Contains(updated, images[0].ServingURL) {
|
||||
t.Errorf("serving URL missing from output: %s", updated)
|
||||
}
|
||||
// The surrounding markup is the reason these documents use HTML at all.
|
||||
for _, want := range []string{`<img`, `class="screenshot"`, `alt="login screen"`, `width="600"`, `</p>`} {
|
||||
if !strings.Contains(updated, want) {
|
||||
t.Errorf("tag structure lost, %q missing from: %s", want, updated)
|
||||
}
|
||||
}
|
||||
if strings.Contains(updated, "\n\n<img src=\"%s/html.png\">", ts.URL, ts.URL)
|
||||
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
|
||||
if len(images) != 2 {
|
||||
t.Fatalf("expected 2 stored images, got %d", len(images))
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("a remote URL survived: %s", updated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_IndependentBudgetPerSyntax(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
// Adding HTML support must not reduce how many Markdown images a document
|
||||
// already gets resolved: each syntax carries its own budget.
|
||||
var b strings.Builder
|
||||
for i := 0; i < maxRemoteImages; i++ {
|
||||
fmt.Fprintf(&b, "\n\n", i, ts.URL, i)
|
||||
}
|
||||
for i := 0; i < 5; i++ {
|
||||
fmt.Fprintf(&b, "<img src=\"%s/html%d.png\">\n\n", ts.URL, i)
|
||||
}
|
||||
|
||||
updated, images, _ := resolve(t, b.String())
|
||||
|
||||
if len(images) != maxRemoteImages+5 {
|
||||
t.Fatalf("expected %d stored images, got %d", maxRemoteImages+5, len(images))
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("a remote URL survived: %s", updated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_HTMLSrcNormalization(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
t.Run("uppercase scheme", func(t *testing.T) {
|
||||
upper := strings.Replace(ts.URL, "http://", "HTTP://", 1)
|
||||
updated, images, _ := resolve(t, fmt.Sprintf(`<img SRC="%s/up.png">`, upper))
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected uppercase scheme to be fetched, got %d", len(images))
|
||||
}
|
||||
// The scheme is normalized because every fetcher downstream compares it
|
||||
// byte-for-byte; an image counted as resolved but unfetchable is worse
|
||||
// than one left alone.
|
||||
if strings.Contains(images[0].OriginalRef, "HTTP://") {
|
||||
t.Errorf("scheme was not normalized: %q", images[0].OriginalRef)
|
||||
}
|
||||
if strings.Contains(updated, "HTTP://") {
|
||||
t.Errorf("document still carries the uppercase scheme: %s", updated)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("markdown uppercase scheme is left alone", func(t *testing.T) {
|
||||
// Markdown targets are handed through unchanged, so a document that
|
||||
// upstream leaves untouched stays untouched.
|
||||
upper := strings.Replace(ts.URL, "http://", "HTTP://", 1)
|
||||
markdown := fmt.Sprintf("", upper)
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected no markdown behaviour change, got %d image(s)", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("markdown content changed:\n got %q\nwant %q", updated, markdown)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("padded src", func(t *testing.T) {
|
||||
_, images, _ := resolve(t, fmt.Sprintf(`<img src=" %s/pad.png ">`, ts.URL))
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected padded src to be fetched, got %d", len(images))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("entity encoded query", func(t *testing.T) {
|
||||
var got string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
got = r.URL.RawQuery
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(createTestPNG(200, 200))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, images, _ := resolve(t, fmt.Sprintf(`<img src="%s/q.png?a=1&b=2">`, srv.URL))
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected 1 stored image, got %d", len(images))
|
||||
}
|
||||
if got != "a=1&b=2" {
|
||||
t.Errorf("entities were not decoded before the request: raw query %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_HTMLSrcFormsOutOfScope(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, hits := remoteImageServer(t)
|
||||
|
||||
cases := map[string]string{
|
||||
"unquoted src": fmt.Sprintf(`<img src=%s/unquoted.png>`, ts.URL),
|
||||
"srcset only": fmt.Sprintf(`<img srcset="%s/set.png 2x">`, ts.URL),
|
||||
"relative src": `<img src="./images/local.png">`,
|
||||
"data uri src": `<img src="data:image/png;base64,AAAA">`,
|
||||
"provider src": `<img src="local://images/x.png">`,
|
||||
"non-img tag": fmt.Sprintf(`<video src="%s/v.mp4">`, ts.URL),
|
||||
}
|
||||
|
||||
for name, markdown := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
before := atomic.LoadInt64(hits)
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected no stored images, got %d", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("content changed:\n got %q\nwant %q", updated, markdown)
|
||||
}
|
||||
if got := atomic.LoadInt64(hits) - before; got != 0 {
|
||||
t.Errorf("expected no fetch, got %d", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_HyphenatedAttributeIsNotSrc(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
// A lazy-loading tag carries a placeholder in src and the real image in
|
||||
// data-src. Treating data-src as src would capture whichever came first and
|
||||
// leave the other unvisited, so only src counts.
|
||||
markdown := fmt.Sprintf(`<img data-src="/placeholder.svg" src="%s/real.png">`, ts.URL)
|
||||
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected the src image to be resolved, got %d", len(images))
|
||||
}
|
||||
if images[0].OriginalRef != ts.URL+"/real.png" {
|
||||
t.Errorf("expected src to be the fetched reference, got %s", images[0].OriginalRef)
|
||||
}
|
||||
if !strings.Contains(updated, `data-src="/placeholder.svg"`) {
|
||||
t.Errorf("data-src must be left alone: %s", updated)
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("src was not rewritten: %s", updated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_HyphenatedAttributeOnlyIsNotResolved(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, hits := remoteImageServer(t)
|
||||
|
||||
// The other half of the same trade-off, pinned deliberately. Requiring
|
||||
// whitespace before src means a tag carrying only a hyphenated *-src is no
|
||||
// longer treated as an image at all — the previous boundary did resolve it.
|
||||
// Recovering that needs an attribute parser rather than a regex; capturing a
|
||||
// lazy-loading placeholder while the real src goes unvisited is the worse of
|
||||
// the two failures.
|
||||
for _, attr := range []string{"data-src", "ng-src", "data-original-src"} {
|
||||
t.Run(attr, func(t *testing.T) {
|
||||
before := atomic.LoadInt64(hits)
|
||||
markdown := fmt.Sprintf(`<img %s="%s/only.png">`, attr, ts.URL)
|
||||
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected no stored images, got %d", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("content changed:\n got %q\nwant %q", updated, markdown)
|
||||
}
|
||||
if got := atomic.LoadInt64(hits) - before; got != 0 {
|
||||
t.Errorf("expected no fetch, got %d", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_WhitelistedHostAgreesWithDocument(t *testing.T) {
|
||||
var hits int64
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
atomic.AddInt64(&hits, 1)
|
||||
if got := r.URL.RawQuery; got != "a=1&b=2" {
|
||||
t.Errorf("expected entities decoded for the request, got raw query %q", got)
|
||||
}
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(createTestPNG(200, 200))
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
t.Setenv("IMAGE_HOST_KEEP_URL", strings.TrimPrefix(ts.URL, "http://"))
|
||||
secutils.ResetSSRFWhitelistForTest()
|
||||
|
||||
markdown := fmt.Sprintf(`<img src=" %s/q.png?a=1&b=2 " alt="kept">`, ts.URL)
|
||||
|
||||
updated, images, fSvc := resolve(t, markdown)
|
||||
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected 1 stored image, got %d", len(images))
|
||||
}
|
||||
if atomic.LoadInt64(&hits) != 1 {
|
||||
t.Fatalf("expected the image to be downloaded once, got %d", hits)
|
||||
}
|
||||
if len(fSvc.saved) != 0 {
|
||||
t.Errorf("a whitelisted image must not be uploaded, got %d uploads", len(fSvc.saved))
|
||||
}
|
||||
|
||||
serving := images[0].ServingURL
|
||||
// ServingURL carries two obligations that have to agree: later stages locate
|
||||
// the image by searching the document for it, and the multimodal stage fetches
|
||||
// that same string.
|
||||
if !strings.Contains(updated, serving) {
|
||||
t.Errorf("ServingURL %q is not present in the document %q", serving, updated)
|
||||
}
|
||||
if serving != ts.URL+"/q.png?a=1&b=2" {
|
||||
t.Errorf("ServingURL must be the normalized fetchable URL, got %q", serving)
|
||||
}
|
||||
if strings.TrimSpace(serving) != serving {
|
||||
t.Errorf("ServingURL must not carry padding: %q", serving)
|
||||
}
|
||||
if strings.Contains(serving, "&") {
|
||||
t.Errorf("ServingURL must not carry HTML entities: %q", serving)
|
||||
}
|
||||
if !strings.Contains(updated, `alt="kept"`) {
|
||||
t.Errorf("the rest of the tag must survive: %s", updated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_CodeExampleIsRewritten(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
// Accepted limitation: this resolver does not parse Markdown structure, so a
|
||||
// tag shown as a documentation example inside a fence is rewritten like any
|
||||
// other. Recognising code spans reliably needs a Markdown parser; a
|
||||
// hand-rolled scanner risks the opposite and far worse failure of treating a
|
||||
// document's real screenshots as examples.
|
||||
markdown := fmt.Sprintf("```html\n<img src=\"%s/example.png\">\n```", ts.URL)
|
||||
|
||||
updated, images, _ := resolve(t, markdown)
|
||||
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected the example to be resolved, got %d", len(images))
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("expected the example to be rewritten: %s", updated)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_CancelledContextStops(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, hits := remoteImageServer(t)
|
||||
|
||||
var b strings.Builder
|
||||
for i := 0; i < 5; i++ {
|
||||
fmt.Fprintf(&b, "<img src=\"%s/c%d.png\">\n\n", ts.URL, i)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
_, images, err := NewImageResolver().ResolveRemoteImages(ctx, b.String(), &mockFileService{}, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected no images after cancellation, got %d", len(images))
|
||||
}
|
||||
if got := atomic.LoadInt64(hits); got != 0 {
|
||||
t.Errorf("expected no fetch after cancellation, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAndStoreThenResolveRemoteImages(t *testing.T) {
|
||||
allowLocalhost(t)
|
||||
ts, _ := remoteImageServer(t)
|
||||
|
||||
// Mirrors the real ingest order, where ResolveAndStore runs first and
|
||||
// ResolveRemoteImages then works on its output. This is the only place the
|
||||
// two HTML resolvers can interfere with each other.
|
||||
markdown := fmt.Sprintf(
|
||||
"<img src=\"images/relative.png\">\n\n<img src=\"%s/remote.png\">", ts.URL)
|
||||
|
||||
fSvc := &mockFileService{}
|
||||
resolver := NewImageResolver()
|
||||
|
||||
afterStore, storeImages, err := resolver.ResolveAndStore(context.Background(), &types.ReadResult{
|
||||
MarkdownContent: markdown,
|
||||
ImageRefs: []types.ImageRef{{
|
||||
OriginalRef: "images/relative.png",
|
||||
Filename: "relative.png",
|
||||
MimeType: "image/png",
|
||||
ImageData: createTestPNG(200, 200),
|
||||
}},
|
||||
}, fSvc, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveAndStore: %v", err)
|
||||
}
|
||||
if len(storeImages) != 1 {
|
||||
t.Fatalf("expected the relative image stored, got %d", len(storeImages))
|
||||
}
|
||||
|
||||
afterRemote, remoteImages, err := resolver.ResolveRemoteImages(
|
||||
context.Background(), afterStore, fSvc, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveRemoteImages: %v", err)
|
||||
}
|
||||
if len(remoteImages) != 1 {
|
||||
t.Fatalf("expected the remote image stored, got %d", len(remoteImages))
|
||||
}
|
||||
if strings.Contains(afterRemote, ts.URL) {
|
||||
t.Errorf("remote URL survived: %s", afterRemote)
|
||||
}
|
||||
if strings.Contains(afterRemote, "images/relative.png") {
|
||||
t.Errorf("relative reference was not rewritten by the earlier pass: %s", afterRemote)
|
||||
}
|
||||
if n := strings.Count(afterRemote, "<img"); n != 2 {
|
||||
t.Errorf("expected both tags preserved, found %d", n)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,264 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// mockFileService is a minimal FileService implementation for testing.
|
||||
type mockFileService struct {
|
||||
saved []savedEntry
|
||||
}
|
||||
|
||||
type savedEntry struct {
|
||||
Data []byte
|
||||
TenantID uint64
|
||||
FileName string
|
||||
}
|
||||
|
||||
func (m *mockFileService) CheckConnectivity(ctx context.Context) error { return nil }
|
||||
func (m *mockFileService) SaveFile(ctx context.Context, file *multipart.FileHeader, tenantID uint64, knowledgeID string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
func (m *mockFileService) SaveBytes(ctx context.Context, data []byte, tenantID uint64, fileName string, temp bool) (string, error) {
|
||||
m.saved = append(m.saved, savedEntry{Data: data, TenantID: tenantID, FileName: fileName})
|
||||
return fmt.Sprintf("local://images/%s", fileName), nil
|
||||
}
|
||||
func (m *mockFileService) GetFile(ctx context.Context, filePath string) (io.ReadCloser, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockFileService) GetFileURL(ctx context.Context, filePath string) (string, error) {
|
||||
return filePath, nil
|
||||
}
|
||||
func (m *mockFileService) DeleteFile(ctx context.Context, filePath string) error { return nil }
|
||||
func (m *mockFileService) CopyFile(ctx context.Context, srcPath string, tenantID uint64, knowledgeID string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_NormalDownload(t *testing.T) {
|
||||
// Whitelist localhost for this test so the test server is reachable
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
|
||||
// Create a test HTTP server that serves a real PNG image.
|
||||
pngData := createTestPNG(200, 200)
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write(pngData)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
markdown := fmt.Sprintf("# Hello\n\n\n\nSome text", ts.URL)
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 1 {
|
||||
t.Fatalf("expected 1 stored image, got %d", len(images))
|
||||
}
|
||||
|
||||
// URL should have been replaced.
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("original URL should have been replaced in markdown, got: %s", updated)
|
||||
}
|
||||
if !strings.Contains(updated, "local://images/") {
|
||||
t.Errorf("expected local:// URL in markdown, got: %s", updated)
|
||||
}
|
||||
|
||||
// Verify saved data.
|
||||
if len(fSvc.saved) != 1 {
|
||||
t.Fatalf("expected 1 saved entry, got %d", len(fSvc.saved))
|
||||
}
|
||||
if fSvc.saved[0].TenantID != 42 {
|
||||
t.Errorf("expected tenantID 42, got %d", fSvc.saved[0].TenantID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_SSRFBlocked(t *testing.T) {
|
||||
// URLs pointing to private IPs should be blocked by SSRF check.
|
||||
markdown := "\n\n"
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// Both images should be left unchanged (SSRF blocked).
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected 0 stored images (SSRF blocked), got %d", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("markdown should be unchanged when SSRF blocked")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_NonImageContentType(t *testing.T) {
|
||||
// Whitelist localhost for this test so the test server is reachable
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("<html>not an image</html>"))
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
markdown := fmt.Sprintf("", ts.URL)
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected 0 images for non-image content type, got %d", len(images))
|
||||
}
|
||||
// Original URL should be preserved.
|
||||
if !strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("original URL should be preserved for non-image content")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_ProviderSchemeSkipped(t *testing.T) {
|
||||
markdown := "\n"
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected 0 images for provider:// URLs, got %d", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("markdown should be unchanged for provider:// URLs")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_MultipleImages(t *testing.T) {
|
||||
// Whitelist localhost for this test so the test server is reachable
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
|
||||
pngData := createTestPNG(256, 256)
|
||||
callCount := 0
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
callCount++
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write(pngData)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
markdown := fmt.Sprintf("\n\ntext\n\n\n\n",
|
||||
ts.URL, ts.URL, ts.URL)
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 10)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 3 {
|
||||
t.Fatalf("expected 3 stored images, got %d", len(images))
|
||||
}
|
||||
if callCount != 3 {
|
||||
t.Errorf("expected 3 HTTP requests, got %d", callCount)
|
||||
}
|
||||
if strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("all original URLs should have been replaced")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_NoImages(t *testing.T) {
|
||||
markdown := "# Just text\n\nNo images here."
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected 0 images, got %d", len(images))
|
||||
}
|
||||
if updated != markdown {
|
||||
t.Errorf("markdown should be unchanged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveRemoteImages_Server404(t *testing.T) {
|
||||
// Whitelist localhost for this test so the test server is reachable
|
||||
t.Setenv("SSRF_WHITELIST", "127.0.0.1,localhost")
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
markdown := fmt.Sprintf("", ts.URL)
|
||||
|
||||
resolver := NewImageResolver()
|
||||
fSvc := &mockFileService{}
|
||||
|
||||
updated, images, err := resolver.ResolveRemoteImages(context.Background(), markdown, fSvc, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if len(images) != 0 {
|
||||
t.Errorf("expected 0 images for 404, got %d", len(images))
|
||||
}
|
||||
// Original URL preserved on failure.
|
||||
if !strings.Contains(updated, ts.URL) {
|
||||
t.Errorf("original URL should be preserved on download failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtFromURLPath(t *testing.T) {
|
||||
tests := []struct {
|
||||
url string
|
||||
expect string
|
||||
}{
|
||||
{"https://example.com/photo.jpg", ".jpg"},
|
||||
{"https://example.com/photo.JPEG", ".jpeg"},
|
||||
{"https://example.com/photo.png?v=2", ""}, // query param — path.Ext won't catch it cleanly but that's ok
|
||||
{"https://example.com/photo.gif", ".gif"},
|
||||
{"https://example.com/photo.webp", ".webp"},
|
||||
{"https://example.com/photo.bmp", ".bmp"},
|
||||
{"https://example.com/photo.svg", ".svg"},
|
||||
{"https://example.com/photo.pdf", ""},
|
||||
{"https://example.com/noext", ""},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.url, func(t *testing.T) {
|
||||
got := extFromURLPath(tt.url)
|
||||
if got != tt.expect {
|
||||
t.Errorf("extFromURLPath(%q) = %q, want %q", tt.url, got, tt.expect)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
package docparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/models/utils"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
weKnoraCloudReaderBaseURL = "https://weknora.weixin.qq.com/api/v1/doc"
|
||||
)
|
||||
|
||||
// WeKnoraCloudSignedDocumentReader implements the docreader HTTP protocol with WeKnoraCloud signing.
|
||||
type WeKnoraCloudSignedDocumentReader struct {
|
||||
appID string
|
||||
apiKey string
|
||||
client *http.Client
|
||||
initialPollInterval time.Duration
|
||||
maxPollInterval time.Duration
|
||||
pollTimeout time.Duration
|
||||
}
|
||||
|
||||
func NewWeKnoraCloudSignedDocumentReader(appID, apiKey string) (*WeKnoraCloudSignedDocumentReader, error) {
|
||||
if appID == "" {
|
||||
return nil, fmt.Errorf("WeKnoraCloud appID is required")
|
||||
}
|
||||
if apiKey == "" {
|
||||
return nil, fmt.Errorf("WeKnoraCloud apiKey is required")
|
||||
}
|
||||
clientCfg := secutils.DefaultSSRFSafeHTTPClientConfig()
|
||||
clientCfg.Timeout = 500 * time.Minute
|
||||
return &WeKnoraCloudSignedDocumentReader{
|
||||
appID: appID,
|
||||
apiKey: apiKey,
|
||||
initialPollInterval: 500 * time.Millisecond,
|
||||
maxPollInterval: 10 * time.Second,
|
||||
pollTimeout: 20 * time.Minute,
|
||||
client: secutils.NewSSRFSafeHTTPClient(clientCfg),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) Reconnect(addr string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) IsConnected() bool { return true }
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) ListEngines(ctx context.Context, overrides map[string]string) ([]types.ParserEngineInfo, error) {
|
||||
return []types.ParserEngineInfo{{
|
||||
Name: WeKnoraCloudEngineName,
|
||||
Description: "WeKnoraCloud signed docreader",
|
||||
FileTypes: []string{"docx", "doc", "pdf", "md", "markdown", "xlsx", "xls", "pptx", "ppt"},
|
||||
Available: true,
|
||||
}}, nil
|
||||
}
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) Read(ctx context.Context, req *types.ReadRequest) (*types.ReadResult, error) {
|
||||
logger.Infof(ctx, "[WeKnoraCloud] read start file=%q type=%q engine=%q hasURL=%v contentLen=%d requestID=%q",
|
||||
req.FileName, req.FileType, req.ParserEngine, strings.TrimSpace(req.URL) != "", len(req.FileContent), req.RequestID)
|
||||
body := httpReadRequest{
|
||||
FileName: req.FileName,
|
||||
FileType: req.FileType,
|
||||
URL: req.URL,
|
||||
Title: req.Title,
|
||||
RequestID: req.RequestID,
|
||||
Config: &httpReadConfig{
|
||||
ParserEngine: req.ParserEngine,
|
||||
ParserEngineOverrides: req.ParserEngineOverrides,
|
||||
},
|
||||
}
|
||||
if len(req.FileContent) > 0 {
|
||||
body.FileContent = base64.StdEncoding.EncodeToString(req.FileContent)
|
||||
}
|
||||
jsonBody, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] marshal read request: %v", err)
|
||||
return nil, fmt.Errorf("http marshal read request: %w", err)
|
||||
}
|
||||
httpReq, err := p.newSignedRequest(ctx, http.MethodPost, weKnoraCloudReaderBaseURL+"/reader", jsonBody)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] signed read request: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
resp, err := p.client.Do(httpReq)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] http read request failed: %v", err)
|
||||
return nil, fmt.Errorf("http read failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusAccepted && resp.StatusCode != http.StatusOK {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] http read unexpected status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
return nil, fmt.Errorf("http read status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
var submit weKnoraCloudAsyncSubmitResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&submit); err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] decode read submit response: %v", err)
|
||||
return nil, fmt.Errorf("http decode read submit response: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(submit.TaskID) == "" {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] submit response missing task_id (status=%q message=%q)", submit.Status, submit.Message)
|
||||
return nil, fmt.Errorf("weknoracloud docreader submit response missing task_id")
|
||||
}
|
||||
logger.Infof(ctx, "[WeKnoraCloud] task submitted task_id=%s file=%q type=%q", submit.TaskID, req.FileName, req.FileType)
|
||||
return p.pollTaskResult(ctx, submit.TaskID)
|
||||
}
|
||||
|
||||
type weKnoraCloudAsyncSubmitResponse struct {
|
||||
TaskID string `json:"task_id"`
|
||||
Status string `json:"status"`
|
||||
Message string `json:"message"`
|
||||
CreatedAt int64 `json:"created_at"`
|
||||
}
|
||||
|
||||
type weKnoraCloudAsyncTaskResponse struct {
|
||||
TaskID string `json:"task_id"`
|
||||
Status string `json:"status"`
|
||||
Message string `json:"message"`
|
||||
Progress float64 `json:"progress"`
|
||||
Result *httpReadResponse `json:"result"`
|
||||
Error string `json:"error"`
|
||||
CreatedAt int64 `json:"created_at"`
|
||||
UpdatedAt int64 `json:"updated_at"`
|
||||
}
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) pollTaskResult(ctx context.Context, taskID string) (*types.ReadResult, error) {
|
||||
pollCtx := ctx
|
||||
if _, ok := ctx.Deadline(); !ok && p.pollTimeout > 0 {
|
||||
var cancel context.CancelFunc
|
||||
pollCtx, cancel = context.WithTimeout(ctx, p.pollTimeout)
|
||||
defer cancel()
|
||||
}
|
||||
statusURL := weKnoraCloudReaderBaseURL + "/" + taskID
|
||||
currentInterval := p.initialPollInterval
|
||||
for {
|
||||
httpReq, err := p.newSignedRequest(pollCtx, http.MethodGet, statusURL, nil)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] poll signed request task_id=%s: %v", taskID, err)
|
||||
return nil, err
|
||||
}
|
||||
resp, err := p.client.Do(httpReq)
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] http poll task_id=%s failed: %v", taskID, err)
|
||||
return nil, fmt.Errorf("http poll task failed: %w", err)
|
||||
}
|
||||
var taskResp weKnoraCloudAsyncTaskResponse
|
||||
func() {
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
err = fmt.Errorf("http poll task status %d: %s", resp.StatusCode, string(bodyBytes))
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] poll task_id=%s status %d: %s", taskID, resp.StatusCode, string(bodyBytes))
|
||||
return
|
||||
}
|
||||
if decodeErr := json.NewDecoder(resp.Body).Decode(&taskResp); decodeErr != nil {
|
||||
err = fmt.Errorf("http decode task response: %w", decodeErr)
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] poll task_id=%s decode response: %v", taskID, decodeErr)
|
||||
}
|
||||
}()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
switch taskResp.Status {
|
||||
case "completed":
|
||||
if taskResp.Result == nil {
|
||||
logger.Infof(ctx, "[WeKnoraCloud] task_id=%s completed with no result payload", taskID)
|
||||
return &types.ReadResult{}, nil
|
||||
}
|
||||
res := fromHTTPReadResponse(taskResp.Result)
|
||||
if res.Error != "" {
|
||||
logger.Errorf(ctx, "[WeKnoraCloud] task_id=%s completed with result.error: %s", taskID, res.Error)
|
||||
} else {
|
||||
logger.Debugf(ctx, "[WeKnoraCloud] task_id=%s completed ok markdownLen=%d", taskID, len(res.MarkdownContent))
|
||||
}
|
||||
return res, nil
|
||||
case "failed":
|
||||
if taskResp.Error != "" {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] task_id=%s failed: %s", taskID, taskResp.Error)
|
||||
return nil, fmt.Errorf("weknoracloud docreader task failed: %s", taskResp.Error)
|
||||
}
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] task_id=%s failed: %s", taskID, taskResp.Message)
|
||||
return nil, fmt.Errorf("weknoracloud docreader task failed: %s", taskResp.Message)
|
||||
case "cancelled":
|
||||
if taskResp.Error != "" {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] task_id=%s cancelled: %s", taskID, taskResp.Error)
|
||||
return nil, fmt.Errorf("weknoracloud docreader task cancelled: %s", taskResp.Error)
|
||||
}
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] task_id=%s cancelled", taskID)
|
||||
return nil, fmt.Errorf("weknoracloud docreader task cancelled")
|
||||
}
|
||||
if err := pollCtx.Err(); err != nil {
|
||||
logger.Errorf(ctx, "[WeKnoraCloud] poll task_id=%s aborted before sleep: %v", taskID, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Exponential backoff: multiply by 1.5 each time, cap at maxPollInterval
|
||||
select {
|
||||
case <-pollCtx.Done():
|
||||
logger.Errorf(ctx, "[WeKnoraCloud] poll task_id=%s stopped: %v", taskID, pollCtx.Err())
|
||||
return nil, pollCtx.Err()
|
||||
case <-time.After(currentInterval):
|
||||
// Update interval for next iteration
|
||||
nextInterval := time.Duration(float64(currentInterval) * 1.5)
|
||||
if nextInterval > p.maxPollInterval {
|
||||
nextInterval = p.maxPollInterval
|
||||
}
|
||||
currentInterval = nextInterval
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *WeKnoraCloudSignedDocumentReader) newSignedRequest(ctx context.Context, method, url string, body []byte) (*http.Request, error) {
|
||||
requestID := uuid.New().String()
|
||||
if len(body) == 0 {
|
||||
body = []byte("{}")
|
||||
}
|
||||
httpReq, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
logger.Errorf(context.Background(), "[WeKnoraCloud] http new request %s %s: %v", method, url, err)
|
||||
return nil, fmt.Errorf("http new request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.ContentLength = int64(len(body))
|
||||
for k, v := range utils.Sign(p.appID, p.apiKey, requestID, string(body)) {
|
||||
httpReq.Header.Set(k, v)
|
||||
}
|
||||
return httpReq, nil
|
||||
}
|
||||
Reference in New Issue
Block a user