mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-25 01:43:27 +08:00
Implement builtin chunk method as ingestion pipeline in GO (#16822)
### Summary Implement builtin chunk mehtod as ingestion pipeline in GO
This commit is contained in:
97
internal/parser/parser/audio_parser.go
Normal file
97
internal/parser/parser/audio_parser.go
Normal file
@@ -0,0 +1,97 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// AudioParser validates and stores configuration for audio files.
|
||||
// The actual speech-to-text transcription is performed by the
|
||||
// component-layer maybeDispatchAudio, which mirrors Python's
|
||||
// Parser._audio() in rag/flow/parser/parser.py.
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// audioExtensions mirrors Python's ParserParam.Defaults() audio suffix
|
||||
// list: ["da","wave","wav","mp3","aac","flac","ogg","aiff","au","midi",
|
||||
// "wma","realaudio","vqf","oggvorbis","ape"].
|
||||
var audioExtensions = map[string]bool{
|
||||
"da": true, "wave": true, "wav": true, "mp3": true,
|
||||
"aac": true, "flac": true, "ogg": true, "aiff": true,
|
||||
"au": true, "midi": true, "wma": true, "realaudio": true,
|
||||
"vqf": true, "oggvorbis": true, "ape": true,
|
||||
}
|
||||
|
||||
// AudioParser handles audio files for transcription. The struct mirrors
|
||||
// the configuration from setups["audio"]:output_format and
|
||||
// setups["audio"].vlm.llm_id.
|
||||
type AudioParser struct {
|
||||
VLMModelID string // vlm.llm_id — identifies the speech-to-text model
|
||||
OutputFormat string
|
||||
}
|
||||
|
||||
// NewAudioParser constructs an AudioParser.
|
||||
func NewAudioParser() *AudioParser {
|
||||
return &AudioParser{}
|
||||
}
|
||||
|
||||
// ConfigureFromSetup reads audio-specific configuration from the
|
||||
// parser setup map. It extracts vlm.llm_id and output_format.
|
||||
func (p *AudioParser) ConfigureFromSetup(setup map[string]any) {
|
||||
if p == nil || setup == nil {
|
||||
return
|
||||
}
|
||||
if vlm, ok := setup["vlm"].(map[string]any); ok {
|
||||
if llmID, ok := vlm["llm_id"].(string); ok && llmID != "" {
|
||||
p.VLMModelID = llmID
|
||||
}
|
||||
}
|
||||
if v, ok := setup["output_format"].(string); ok && v != "" {
|
||||
p.OutputFormat = v
|
||||
}
|
||||
}
|
||||
|
||||
// ParseWithResult implements ParseResultProducer. It validates the
|
||||
// file extension against the audio extension whitelist. The actual
|
||||
// speech-to-text transcription happens via maybeDispatchAudio at the
|
||||
// component layer (mirrors Python's LLMBundle.transcription call).
|
||||
func (p *AudioParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
ext := strings.ToLower(filepath.Ext(filename))
|
||||
if len(ext) > 1 && ext[0] == '.' {
|
||||
ext = ext[1:]
|
||||
}
|
||||
if ext == "" || !audioExtensions[ext] {
|
||||
return ParseResult{
|
||||
Err: fmt.Errorf("audio: unsupported extension %q (filename: %s); accepted: .da/.wav/.wave/.mp3/.aac/.flac/.ogg/.aiff/.au/.midi/.wma/.realaudio/.vqf/.oggvorbis/.ape", ext, filename),
|
||||
}
|
||||
}
|
||||
|
||||
// OutputFormat and VLMModelID are consumed by maybeDispatchAudio.
|
||||
outFmt := p.OutputFormat
|
||||
if outFmt == "" {
|
||||
outFmt = "text"
|
||||
}
|
||||
|
||||
return ParseResult{
|
||||
OutputFormat: outFmt,
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"size": len(data),
|
||||
},
|
||||
}
|
||||
}
|
||||
98
internal/parser/parser/audio_parser_test.go
Normal file
98
internal/parser/parser/audio_parser_test.go
Normal file
@@ -0,0 +1,98 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAudioParser_ValidExtension(t *testing.T) {
|
||||
p := NewAudioParser()
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"output_format": "text",
|
||||
"vlm": map[string]any{
|
||||
"llm_id": "whisper-1",
|
||||
},
|
||||
})
|
||||
|
||||
result := p.ParseWithResult("test.mp3", []byte{1, 2, 3, 4})
|
||||
if result.Err != nil {
|
||||
t.Fatalf("unexpected error: %v", result.Err)
|
||||
}
|
||||
if result.OutputFormat != "text" {
|
||||
t.Errorf("output_format = %q, want text", result.OutputFormat)
|
||||
}
|
||||
if name, ok := result.File["name"].(string); !ok || name != "test.mp3" {
|
||||
t.Errorf("name = %q, want test.mp3", name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAudioParser_AllExtensions(t *testing.T) {
|
||||
for _, ext := range []string{"da", "wave", "wav", "mp3", "aac", "flac", "ogg", "aiff", "au", "midi", "wma", "realaudio", "vqf", "oggvorbis", "ape"} {
|
||||
p := NewAudioParser()
|
||||
result := p.ParseWithResult("audio."+ext, []byte{})
|
||||
if result.Err != nil {
|
||||
t.Errorf("extension .%s rejected: %v", ext, result.Err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAudioParser_InvalidExtension(t *testing.T) {
|
||||
p := NewAudioParser()
|
||||
result := p.ParseWithResult("notaudio.txt", []byte{})
|
||||
if result.Err == nil {
|
||||
t.Fatal("expected error for .txt extension")
|
||||
}
|
||||
if !strings.Contains(result.Err.Error(), "unsupported extension") {
|
||||
t.Errorf("error message should mention unsupported extension: %v", result.Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAudioParser_ConfigureVLM(t *testing.T) {
|
||||
p := NewAudioParser()
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"vlm": map[string]any{
|
||||
"llm_id": "openai-whisper",
|
||||
},
|
||||
})
|
||||
if p.VLMModelID != "openai-whisper" {
|
||||
t.Errorf("VLMModelID = %q, want openai-whisper", p.VLMModelID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAudioParser_ConfigureNilSafe(t *testing.T) {
|
||||
p := NewAudioParser()
|
||||
p.ConfigureFromSetup(nil)
|
||||
p.ConfigureFromSetup(map[string]any{})
|
||||
// Should not panic.
|
||||
result := p.ParseWithResult("test.wav", []byte{})
|
||||
if result.Err != nil {
|
||||
t.Errorf("unexpected error: %v", result.Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAudioParser_DefaultOutputFormat(t *testing.T) {
|
||||
p := NewAudioParser()
|
||||
result := p.ParseWithResult("test.flac", []byte{})
|
||||
if result.Err != nil {
|
||||
t.Fatalf("unexpected error: %v", result.Err)
|
||||
}
|
||||
if result.OutputFormat != "text" {
|
||||
t.Errorf("output_format = %q, want text (default)", result.OutputFormat)
|
||||
}
|
||||
}
|
||||
251
internal/parser/parser/csv_parser.go
Normal file
251
internal/parser/parser/csv_parser.go
Normal file
@@ -0,0 +1,251 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// CSVParser renders CSV data as HTML tables, matching the spreadsheet
|
||||
// family output_format == "html" convention from ParserParam.Defaults().
|
||||
//
|
||||
// Mirrors Python's deepdoc/parser/excel_parser.py:RAGFlowExcelParser.html():
|
||||
// - CSV data is rendered as an HTML <table> with <caption> "Data".
|
||||
// - The first row is treated as the header (<th>).
|
||||
// - Illegal control characters are replaced with spaces.
|
||||
// - Large sheets are split into chunks of chunk_rows data rows, each
|
||||
// chunk being a self-contained <table> with its own <caption> and
|
||||
// repeated header row.
|
||||
//
|
||||
// It implements the ParseResultProducer contract so the dispatch seam in
|
||||
// parser_dispatch.go routes .csv files through the structured path.
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"html"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Go port of Python's ILLEGAL_CHARACTERS_RE from
|
||||
// deepdoc/parser/excel_parser.py.
|
||||
// Pattern: [\000-\010]|[\013-\014]|[\016-\037]
|
||||
// Matches all control chars except TAB (\x09), LF (\x0A), CR (\x0D).
|
||||
var csvIllegalCharsRe = regexp.MustCompile(`[\x00-\x08]|\x0B|\x0C|[\x0E-\x1F]`)
|
||||
|
||||
const csvDefaultChunkRows = 256
|
||||
const csvSheetName = "Data"
|
||||
|
||||
// CSVParser reads RFC-4180 CSV data and emits HTML <table> payloads.
|
||||
type CSVParser struct {
|
||||
ParseMethod string
|
||||
OutputFormat string
|
||||
ChunkRows int
|
||||
TCADPAPIServer string
|
||||
TCADPAPIKey string
|
||||
TCADPTableResultType string
|
||||
TCADPMarkdownImageResponseType string
|
||||
}
|
||||
|
||||
func NewCSVParser() *CSVParser {
|
||||
return &CSVParser{
|
||||
ChunkRows: csvDefaultChunkRows,
|
||||
TCADPTableResultType: "1",
|
||||
TCADPMarkdownImageResponseType: "1",
|
||||
}
|
||||
}
|
||||
|
||||
func (p *CSVParser) String() string {
|
||||
return "CSVParser"
|
||||
}
|
||||
|
||||
func (p *CSVParser) ConfigureFromSetup(setup map[string]any) {
|
||||
if p == nil || setup == nil {
|
||||
return
|
||||
}
|
||||
if v, ok := setup["parse_method"].(string); ok && v != "" {
|
||||
p.ParseMethod = v
|
||||
}
|
||||
if v, ok := setup["output_format"].(string); ok && v != "" {
|
||||
p.OutputFormat = v
|
||||
}
|
||||
if v, ok := setup["chunk_rows"]; ok {
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
p.ChunkRows = int(n)
|
||||
case int:
|
||||
p.ChunkRows = n
|
||||
case int64:
|
||||
p.ChunkRows = int(n)
|
||||
}
|
||||
if p.ChunkRows <= 0 {
|
||||
p.ChunkRows = csvDefaultChunkRows
|
||||
}
|
||||
}
|
||||
if v, ok := setup["tcadp_apiserver"].(string); ok && v != "" {
|
||||
p.TCADPAPIServer = v
|
||||
}
|
||||
if v, ok := setup["tcadp_api_key"].(string); ok {
|
||||
p.TCADPAPIKey = v
|
||||
}
|
||||
if v, ok := setup["table_result_type"].(string); ok && v != "" {
|
||||
p.TCADPTableResultType = v
|
||||
}
|
||||
if v, ok := setup["markdown_image_response_type"].(string); ok && v != "" {
|
||||
p.TCADPMarkdownImageResponseType = v
|
||||
}
|
||||
}
|
||||
|
||||
// ParseWithResult implements ParseResultProducer. It reads CSV rows
|
||||
// and renders them as HTML <table> chunks with <caption>, header row
|
||||
// repeated per chunk, and illegal-character filtering — mirroring
|
||||
// Python's RAGFlowExcelParser.html().
|
||||
// When TCADP parse_method is configured, the file is dispatched to
|
||||
// the Tencent Cloud Document Parsing API.
|
||||
func (p *CSVParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
method := normalizeXLSXParseMethod(p.ParseMethod)
|
||||
switch method {
|
||||
case "tcadp":
|
||||
return parseSpreadsheetWithTCADP(
|
||||
filename, data, "CSV",
|
||||
p.TCADPAPIServer, p.TCADPAPIKey,
|
||||
p.TCADPTableResultType, p.TCADPMarkdownImageResponseType,
|
||||
p.OutputFormat,
|
||||
)
|
||||
case "", "csv":
|
||||
// Continue with the local CSV parser.
|
||||
default:
|
||||
return ParseResult{
|
||||
Err: fmt.Errorf("unsupported CSV parse method: %q", p.ParseMethod),
|
||||
}
|
||||
}
|
||||
|
||||
text := string(data)
|
||||
if strings.TrimSpace(text) == "" {
|
||||
return ParseResult{
|
||||
OutputFormat: "html",
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"size": len(data),
|
||||
"encoding": "utf-8",
|
||||
},
|
||||
HTML: "<table><caption>" + csvSheetName + "</caption><tr><td></td></tr></table>",
|
||||
}
|
||||
}
|
||||
|
||||
reader := csv.NewReader(strings.NewReader(text))
|
||||
reader.LazyQuotes = true
|
||||
reader.TrimLeadingSpace = true
|
||||
|
||||
records, err := reader.ReadAll()
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("csv parse: %w", err)}
|
||||
}
|
||||
|
||||
// Clean illegal control characters from all cells.
|
||||
records = cleanCSVRecords(records)
|
||||
|
||||
chunkRows := p.ChunkRows
|
||||
if chunkRows <= 0 {
|
||||
chunkRows = csvDefaultChunkRows
|
||||
}
|
||||
|
||||
return ParseResult{
|
||||
OutputFormat: "html",
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"size": len(data),
|
||||
"encoding": "utf-8",
|
||||
},
|
||||
HTML: recordsToHTMLTableChunks(records, chunkRows),
|
||||
}
|
||||
}
|
||||
|
||||
// cleanCSVRecords replaces illegal control characters in all cells
|
||||
// with a single space, matching Python's ILLEGAL_CHARACTERS_RE.
|
||||
func cleanCSVRecords(records [][]string) [][]string {
|
||||
out := make([][]string, len(records))
|
||||
for i, row := range records {
|
||||
out[i] = make([]string, len(row))
|
||||
for j, cell := range row {
|
||||
out[i][j] = csvIllegalCharsRe.ReplaceAllString(cell, " ")
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// recordsToHTMLTableChunks renders CSV records as one or more
|
||||
// self-contained HTML <table> chunks. The first row is always the
|
||||
// header (<th>). Data rows are split into chunks of chunkRows,
|
||||
// each chunk being a complete <table> with <caption> and a repeated
|
||||
// header row. Chunks are joined with newlines.
|
||||
//
|
||||
// Mirrors Python's RAGFlowExcelParser.html() chunking:
|
||||
//
|
||||
// chunks = (n_data_rows + chunk_rows - 1) // chunk_rows
|
||||
func recordsToHTMLTableChunks(records [][]string, chunkRows int) string {
|
||||
if len(records) == 0 {
|
||||
return "<table><caption>" + csvSheetName + "</caption></table>"
|
||||
}
|
||||
|
||||
// Build the header row once — repeated in every chunk.
|
||||
headerHTML := buildCSVHeaderRow(records[0])
|
||||
dataRows := records[1:]
|
||||
nData := len(dataRows)
|
||||
|
||||
if nData == 0 {
|
||||
// Only a header row exists.
|
||||
return "<table><caption>" + csvSheetName + "</caption>\n" + headerHTML + "\n</table>"
|
||||
}
|
||||
|
||||
nChunks := (nData + chunkRows - 1) / chunkRows
|
||||
var b strings.Builder
|
||||
for ci := 0; ci < nChunks; ci++ {
|
||||
start := ci * chunkRows
|
||||
end := start + chunkRows
|
||||
if end > nData {
|
||||
end = nData
|
||||
}
|
||||
|
||||
b.WriteString("<table><caption>")
|
||||
b.WriteString(csvSheetName)
|
||||
b.WriteString("</caption>\n")
|
||||
b.WriteString(headerHTML)
|
||||
|
||||
for _, row := range dataRows[start:end] {
|
||||
b.WriteString("<tr>")
|
||||
for _, cell := range row {
|
||||
b.WriteString("<td>")
|
||||
b.WriteString(html.EscapeString(strings.TrimSpace(cell)))
|
||||
b.WriteString("</td>")
|
||||
}
|
||||
b.WriteString("</tr>\n")
|
||||
}
|
||||
b.WriteString("</table>\n")
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// buildCSVHeaderRow renders the first row as an HTML <th> header row.
|
||||
func buildCSVHeaderRow(row []string) string {
|
||||
var b strings.Builder
|
||||
b.WriteString("<tr>")
|
||||
for _, cell := range row {
|
||||
b.WriteString("<th>")
|
||||
b.WriteString(html.EscapeString(strings.TrimSpace(cell)))
|
||||
b.WriteString("</th>")
|
||||
}
|
||||
b.WriteString("</tr>\n")
|
||||
return b.String()
|
||||
}
|
||||
@@ -19,23 +19,39 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
officeOxide "github.com/yfedoseev/office_oxide/go"
|
||||
)
|
||||
|
||||
type DOCXParser struct{}
|
||||
// DOCXFigure represents one embedded image plus its surrounding text
|
||||
// context, mirroring the chunk-level shape that Python's
|
||||
// naive_merge_docx produces for vision_figure_parser_docx_wrapper_naive.
|
||||
type DOCXFigure struct {
|
||||
Image string `json:"image"` // base64-encoded image bytes
|
||||
ContextAbove string `json:"context_above"` // text before the image block
|
||||
ContextBelow string `json:"context_below"` // text after the image block
|
||||
Marker string `json:"marker"` // substring to locate image position in markdown
|
||||
}
|
||||
|
||||
type DOCXParser struct {
|
||||
libType string
|
||||
}
|
||||
|
||||
func NewDOCXParser() *DOCXParser {
|
||||
return &DOCXParser{}
|
||||
}
|
||||
|
||||
// ParseWithResult captures the office_oxide ToMarkdown output
|
||||
// instead of discarding it. Returns OutputFormat="markdown" with
|
||||
// the rendered Markdown on the matching key. Mirrors the python
|
||||
// parser.py:Docx branch which uses rag.app.naive.Docx to render
|
||||
// markdown; the Go side delegates to office_oxide for now and
|
||||
// will switch to a native Markdown renderer in a follow-up.
|
||||
// and additionally extracts embedded images with their surrounding
|
||||
// text context so the downstream vision-figure dispatch can enrich
|
||||
// the markdown with LLM-generated image descriptions.
|
||||
//
|
||||
// Mirrors python naive.py: Docx() → naive_merge_docx() →
|
||||
// vision_figure_parser_docx_wrapper_naive().
|
||||
func (p *DOCXParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
doc, err := officeOxide.OpenFromBytes(data, "docx")
|
||||
if err != nil {
|
||||
@@ -48,13 +64,182 @@ func (p *DOCXParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
return ParseResult{Err: fmt.Errorf("docx to-markdown: %w", err)}
|
||||
}
|
||||
|
||||
fileMeta := map[string]any{
|
||||
"name": filename,
|
||||
"format": "docx",
|
||||
}
|
||||
|
||||
// Extract embedded images with their text context from the
|
||||
// office_oxide IR so the downstream vision dispatch can
|
||||
// enrich them. The already-opened doc handle is reused
|
||||
// (no second OpenFromBytes).
|
||||
irJSON, irErr := doc.ToIRJSON()
|
||||
var figures []DOCXFigure
|
||||
if irErr == nil {
|
||||
figures = extractDOCXFiguresFromIR(irJSON)
|
||||
}
|
||||
if len(figures) > 0 {
|
||||
figs := make([]map[string]any, 0, len(figures))
|
||||
for _, f := range figures {
|
||||
figs = append(figs, map[string]any{
|
||||
"image": f.Image,
|
||||
"context_above": f.ContextAbove,
|
||||
"context_below": f.ContextBelow,
|
||||
"marker": f.Marker,
|
||||
})
|
||||
}
|
||||
fileMeta["figures"] = figs
|
||||
}
|
||||
|
||||
return ParseResult{
|
||||
OutputFormat: "markdown",
|
||||
File: map[string]any{"name": filename, "format": "docx"},
|
||||
File: fileMeta,
|
||||
Markdown: md,
|
||||
}
|
||||
}
|
||||
|
||||
// extractDOCXFiguresFromIR parses the office_oxide IR JSON and
|
||||
// returns every embedded image block together with the plain text
|
||||
// immediately surrounding it. The context matches what Python's
|
||||
// naive_merge_docx attaches as context_above / context_below on
|
||||
// each chunk that carries an image.
|
||||
//
|
||||
// Reuses the IR already obtained from the doc handle in
|
||||
// ParseWithResult so the binary is not opened twice.
|
||||
func extractDOCXFiguresFromIR(irJSON string) []DOCXFigure {
|
||||
var ir docxIRDocument
|
||||
if err := json.Unmarshal([]byte(irJSON), &ir); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var flat []flatBlock
|
||||
for _, sec := range ir.Sections {
|
||||
for _, el := range sec.Elements {
|
||||
if el.Type == "image" {
|
||||
b64 := base64.StdEncoding.EncodeToString(el.Data)
|
||||
flat = append(flat, flatBlock{image: b64})
|
||||
continue
|
||||
}
|
||||
text := joinDOCXIRRuns(el.Content)
|
||||
flat = append(flat, flatBlock{text: text})
|
||||
}
|
||||
}
|
||||
|
||||
var figures []DOCXFigure
|
||||
for i, block := range flat {
|
||||
if block.image == "" {
|
||||
continue
|
||||
}
|
||||
fig := DOCXFigure{Image: block.image}
|
||||
|
||||
// Collect text above (backward scan up to docxContextWindow
|
||||
// chars, or until another image is hit).
|
||||
above := collectDOCXPrevText(flat, i, 512)
|
||||
fig.ContextAbove = strings.TrimSpace(above)
|
||||
|
||||
// Collect text below (forward scan up to docxContextWindow
|
||||
// chars, or until another image is hit).
|
||||
below := collectDOCXNextText(flat, i, 512)
|
||||
fig.ContextBelow = strings.TrimSpace(below)
|
||||
|
||||
// Marker: text of the immediately preceding flat block,
|
||||
// used by the vision dispatcher to locate the image position
|
||||
// in the rendered markdown for inline insertion.
|
||||
for j := i - 1; j >= 0; j-- {
|
||||
if flat[j].text != "" {
|
||||
fig.Marker = flat[j].text
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
figures = append(figures, fig)
|
||||
}
|
||||
return figures
|
||||
}
|
||||
|
||||
// --- internal types ---
|
||||
|
||||
// flatBlock is a flattened IR element used internally to collect
|
||||
// text / image context around embedded figures.
|
||||
type flatBlock struct {
|
||||
text string
|
||||
image string // base64-encoded image data (empty for non-image)
|
||||
}
|
||||
|
||||
const docxContextWindow = 512
|
||||
|
||||
func collectDOCXPrevText(flat []flatBlock, idx, maxLen int) string {
|
||||
var parts []string
|
||||
remaining := maxLen
|
||||
for i := idx - 1; i >= 0 && remaining > 0; i-- {
|
||||
if flat[i].image != "" {
|
||||
break // stop at previous image
|
||||
}
|
||||
if flat[i].text == "" {
|
||||
continue
|
||||
}
|
||||
r := []rune(flat[i].text)
|
||||
if len(r) > remaining {
|
||||
r = r[len(r)-remaining:]
|
||||
}
|
||||
parts = append([]string{string(r)}, parts...)
|
||||
remaining -= len(r)
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func collectDOCXNextText(flat []flatBlock, idx, maxLen int) string {
|
||||
var parts []string
|
||||
remaining := maxLen
|
||||
for i := idx + 1; i < len(flat) && remaining > 0; i++ {
|
||||
if flat[i].image != "" {
|
||||
break // stop at next image
|
||||
}
|
||||
if flat[i].text == "" {
|
||||
continue
|
||||
}
|
||||
r := []rune(flat[i].text)
|
||||
if len(r) > remaining {
|
||||
r = r[:remaining]
|
||||
}
|
||||
parts = append(parts, string(r))
|
||||
remaining -= len(r)
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func joinDOCXIRRuns(runs []docxIRRun) string {
|
||||
var b strings.Builder
|
||||
for _, r := range runs {
|
||||
if r.Type == "text" {
|
||||
b.WriteString(r.Text)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// --- office_oxide IR types (local copy, independent of deepdoc) ---
|
||||
|
||||
type docxIRDocument struct {
|
||||
Sections []docxIRSection `json:"sections"`
|
||||
}
|
||||
|
||||
type docxIRSection struct {
|
||||
Title string `json:"title"`
|
||||
Elements []docxIRElement `json:"elements"`
|
||||
}
|
||||
|
||||
type docxIRElement struct {
|
||||
Type string `json:"type"`
|
||||
Content []docxIRRun `json:"content"`
|
||||
Data []byte `json:"data"`
|
||||
}
|
||||
|
||||
type docxIRRun struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
func (p *DOCXParser) String() string {
|
||||
return "DOCXParser"
|
||||
}
|
||||
|
||||
379
internal/parser/parser/email_parser.go
Normal file
379
internal/parser/parser/email_parser.go
Normal file
@@ -0,0 +1,379 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
"net/mail"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/text/encoding"
|
||||
"golang.org/x/text/encoding/simplifiedchinese"
|
||||
"golang.org/x/text/transform"
|
||||
)
|
||||
|
||||
// EmailParser parses .eml (RFC 5322 email) files into structured
|
||||
// JSON or plain-text output. Mirrors Python's _email() method in
|
||||
// rag/flow/parser/parser.py.
|
||||
//
|
||||
// .msg (Outlook) files are not supported in the Go path; callers
|
||||
// receive a clear error.
|
||||
type EmailParser struct {
|
||||
fields []string
|
||||
outputFormat string
|
||||
}
|
||||
|
||||
func NewEmailParser() *EmailParser {
|
||||
return &EmailParser{}
|
||||
}
|
||||
|
||||
func (p *EmailParser) ConfigureFromSetup(setup map[string]any) {
|
||||
if p == nil || setup == nil {
|
||||
return
|
||||
}
|
||||
if v, ok := setup["output_format"].(string); ok && v != "" {
|
||||
p.outputFormat = v
|
||||
}
|
||||
if raw, ok := setup["fields"]; ok {
|
||||
switch list := raw.(type) {
|
||||
case []string:
|
||||
p.fields = list
|
||||
case []any:
|
||||
p.fields = make([]string, 0, len(list))
|
||||
for _, item := range list {
|
||||
if s, ok := item.(string); ok {
|
||||
p.fields = append(p.fields, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(p.fields) == 0 {
|
||||
p.fields = []string{"from", "to", "cc", "bcc", "date", "subject", "body", "attachments", "metadata"}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *EmailParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
ext := strings.ToLower(filepath.Ext(filename))
|
||||
if ext == ".msg" {
|
||||
return ParseResult{
|
||||
Err: fmt.Errorf("email: .msg (Outlook) files are not supported in the Go parser; use .eml format"),
|
||||
}
|
||||
}
|
||||
|
||||
emailContent := parseEML(bytes.NewReader(data), p.fields)
|
||||
|
||||
outputFormat := p.outputFormat
|
||||
if outputFormat == "" {
|
||||
outputFormat = "text"
|
||||
}
|
||||
|
||||
if outputFormat == "json" {
|
||||
emailContent["doc_type_kwd"] = "text"
|
||||
return ParseResult{
|
||||
OutputFormat: "json",
|
||||
File: map[string]any{"name": filename},
|
||||
JSON: []map[string]any{emailContent},
|
||||
}
|
||||
}
|
||||
|
||||
// Text output: flatten fields into a single string.
|
||||
var sb strings.Builder
|
||||
for k, v := range emailContent {
|
||||
switch val := v.(type) {
|
||||
case string:
|
||||
sb.WriteString(k)
|
||||
sb.WriteString(":")
|
||||
sb.WriteString(val)
|
||||
sb.WriteString("\n")
|
||||
case map[string]any:
|
||||
sb.WriteString(k)
|
||||
sb.WriteString(":{")
|
||||
for mk, mv := range val {
|
||||
if ms, ok := mv.(string); ok {
|
||||
sb.WriteString(mk)
|
||||
sb.WriteString(":")
|
||||
sb.WriteString(ms)
|
||||
sb.WriteString(", ")
|
||||
}
|
||||
}
|
||||
sb.WriteString("}\n")
|
||||
case []map[string]any:
|
||||
for _, att := range val {
|
||||
fn, _ := att["filename"].(string)
|
||||
pl, _ := att["payload"].(string)
|
||||
sb.WriteString(fn)
|
||||
sb.WriteString(":")
|
||||
sb.WriteString(pl)
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
case []string:
|
||||
sb.WriteString(strings.Join(val, "\n"))
|
||||
}
|
||||
}
|
||||
return ParseResult{
|
||||
OutputFormat: "text",
|
||||
File: map[string]any{"name": filename},
|
||||
Text: sb.String(),
|
||||
}
|
||||
}
|
||||
|
||||
// -- field set helpers --
|
||||
|
||||
func targetFieldsSet(fields []string) map[string]bool {
|
||||
m := make(map[string]bool, len(fields))
|
||||
for _, f := range fields {
|
||||
m[strings.ToLower(strings.TrimSpace(f))] = true
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// -- .eml parsing (RFC 5322 with multipart support) --
|
||||
|
||||
func parseEML(r io.Reader, fields []string) map[string]any {
|
||||
target := targetFieldsSet(fields)
|
||||
content := map[string]any{}
|
||||
|
||||
msg, err := mail.ReadMessage(r)
|
||||
if err != nil {
|
||||
content["error"] = fmt.Sprintf("email: parse error: %v", err)
|
||||
return content
|
||||
}
|
||||
|
||||
// Headers.
|
||||
meta := map[string]any{}
|
||||
for key, vals := range msg.Header {
|
||||
keyLower := strings.ToLower(key)
|
||||
val := strings.Join(vals, ", ")
|
||||
switch keyLower {
|
||||
case "from", "to", "cc", "bcc", "date", "subject":
|
||||
if target[keyLower] {
|
||||
content[keyLower] = val
|
||||
}
|
||||
default:
|
||||
meta[keyLower] = val
|
||||
}
|
||||
}
|
||||
if target["metadata"] {
|
||||
content["metadata"] = meta
|
||||
}
|
||||
|
||||
// Body and attachments — readMailBody walks all multipart parts
|
||||
// and collects text/html bodies alongside attachment parts whose
|
||||
// Content-Disposition starts with "attachment".
|
||||
needAttachments := target["attachments"]
|
||||
if target["body"] {
|
||||
contentType := msg.Header.Get("Content-Type")
|
||||
bodyText, bodyHTML, attachments := readMailBody(msg.Body, contentType, needAttachments)
|
||||
if bodyText != "" {
|
||||
content["text"] = bodyText
|
||||
}
|
||||
if bodyHTML != "" {
|
||||
content["text_html"] = bodyHTML
|
||||
}
|
||||
if needAttachments {
|
||||
content["attachments"] = attachments
|
||||
}
|
||||
} else if needAttachments {
|
||||
content["attachments"] = []map[string]any{}
|
||||
}
|
||||
|
||||
return content
|
||||
}
|
||||
|
||||
// readMailBody reads the body of an email message, handling
|
||||
// multipart/alternative, multipart/mixed, and single-part content
|
||||
// types. Returns (textBody, htmlBody, attachments).
|
||||
// When collectAttachments is true, non-text parts with Content-Disposition
|
||||
// starting with "attachment" are collected.
|
||||
func readMailBody(body io.Reader, contentType string, collectAttachments bool) (string, string, []map[string]any) {
|
||||
var attachments []map[string]any
|
||||
|
||||
mediaType, params, err := mime.ParseMediaType(contentType)
|
||||
if err != nil {
|
||||
mediaType = "text/plain"
|
||||
}
|
||||
|
||||
if !strings.HasPrefix(mediaType, "multipart/") {
|
||||
raw, _ := io.ReadAll(body)
|
||||
decoded := decodeMailPayload(raw, params["charset"])
|
||||
if mediaType == "text/html" {
|
||||
return "", decoded, attachments
|
||||
}
|
||||
return decoded, "", attachments
|
||||
}
|
||||
|
||||
boundary := params["boundary"]
|
||||
if boundary == "" {
|
||||
raw, _ := io.ReadAll(body)
|
||||
return decodeMailPayload(raw, ""), "", attachments
|
||||
}
|
||||
|
||||
mr := multipart.NewReader(body, boundary)
|
||||
var textParts, htmlParts []string
|
||||
for {
|
||||
part, err := mr.NextPart()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
partCT := part.Header.Get("Content-Type")
|
||||
partMedia, partParams, _ := mime.ParseMediaType(partCT)
|
||||
|
||||
if strings.HasPrefix(partMedia, "multipart/") {
|
||||
t, h, nestedAttachments := readMailBody(part, partCT, collectAttachments)
|
||||
if t != "" {
|
||||
textParts = append(textParts, t)
|
||||
}
|
||||
if h != "" {
|
||||
htmlParts = append(htmlParts, h)
|
||||
}
|
||||
attachments = append(attachments, nestedAttachments...)
|
||||
continue
|
||||
}
|
||||
|
||||
// Check if this part is an attachment.
|
||||
if collectAttachments && isAttachmentPart(part) {
|
||||
raw, _ := io.ReadAll(part)
|
||||
attachments = append(attachments, map[string]any{
|
||||
"filename": attachmentFilename(part),
|
||||
"payload": decodeMailPayload(raw, partParams["charset"]),
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
raw, _ := io.ReadAll(part)
|
||||
decoded := decodeMailPayload(raw, partParams["charset"])
|
||||
|
||||
switch partMedia {
|
||||
case "text/plain":
|
||||
textParts = append(textParts, decoded)
|
||||
case "text/html":
|
||||
htmlParts = append(htmlParts, decoded)
|
||||
}
|
||||
}
|
||||
return strings.Join(textParts, "\n"), strings.Join(htmlParts, "\n"), attachments
|
||||
}
|
||||
|
||||
// isAttachmentPart checks whether a multipart part should be treated as
|
||||
// an attachment (Content-Disposition starts with "attachment"). Mirrors
|
||||
// Python's check in _email().
|
||||
func isAttachmentPart(part *multipart.Part) bool {
|
||||
disp := part.Header.Get("Content-Disposition")
|
||||
if disp == "" {
|
||||
return false
|
||||
}
|
||||
dispType, _, err := mime.ParseMediaType(disp)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return strings.EqualFold(dispType, "attachment")
|
||||
}
|
||||
|
||||
// attachmentFilename extracts a filename from the part's
|
||||
// Content-Disposition or Content-Type headers.
|
||||
func attachmentFilename(part *multipart.Part) string {
|
||||
if fn := part.FileName(); fn != "" {
|
||||
return fn
|
||||
}
|
||||
ct := part.Header.Get("Content-Type")
|
||||
if ct != "" {
|
||||
_, params, _ := mime.ParseMediaType(ct)
|
||||
if name, ok := params["name"]; ok {
|
||||
return name
|
||||
}
|
||||
}
|
||||
return "attachment.bin"
|
||||
}
|
||||
|
||||
// decodeMailPayload attempts multiple charset decodings.
|
||||
// Mirrors Python's _decode_payload with fallback chain:
|
||||
// utf-8 → gb2312 → gbk → gb18030 → latin1 → utf-8 (ignore).
|
||||
func decodeMailPayload(payload []byte, charset string) string {
|
||||
if len(payload) == 0 {
|
||||
return ""
|
||||
}
|
||||
charsets := buildCharsetChain(charset)
|
||||
for _, enc := range charsets {
|
||||
if enc == "" {
|
||||
// raw bytes → already fallthrough
|
||||
return string(payload)
|
||||
}
|
||||
decoded, err := decodeWithCharset(payload, enc)
|
||||
if err == nil {
|
||||
return decoded
|
||||
}
|
||||
}
|
||||
return string(payload)
|
||||
}
|
||||
|
||||
func buildCharsetChain(declared string) []string {
|
||||
chain := make([]string, 0, 7)
|
||||
if declared != "" {
|
||||
chain = append(chain, declared)
|
||||
}
|
||||
chain = append(chain, "utf-8", "gb2312", "gbk", "gb18030", "latin1")
|
||||
return chain
|
||||
}
|
||||
|
||||
func decodeWithCharset(payload []byte, charset string) (string, error) {
|
||||
charset = strings.ToLower(strings.TrimSpace(charset))
|
||||
switch charset {
|
||||
case "utf-8", "utf8", "":
|
||||
s := string(payload)
|
||||
if !strings.ContainsRune(s, '\ufffd') {
|
||||
return s, nil
|
||||
}
|
||||
return "", fmt.Errorf("invalid utf-8")
|
||||
case "latin1", "iso-8859-1", "iso8859-1":
|
||||
runes := make([]rune, len(payload))
|
||||
for i, b := range payload {
|
||||
runes[i] = rune(b)
|
||||
}
|
||||
return string(runes), nil
|
||||
case "gb2312":
|
||||
return decodeTransform(payload, simplifiedchinese.HZGB2312.NewDecoder())
|
||||
case "gbk":
|
||||
return decodeTransform(payload, simplifiedchinese.GBK.NewDecoder())
|
||||
case "gb18030":
|
||||
return decodeTransform(payload, simplifiedchinese.GB18030.NewDecoder())
|
||||
}
|
||||
// Unknown charset: treat as latin-1.
|
||||
runes := make([]rune, len(payload))
|
||||
for i, b := range payload {
|
||||
runes[i] = rune(b)
|
||||
}
|
||||
return string(runes), nil
|
||||
}
|
||||
|
||||
func decodeTransform(payload []byte, decoder *encoding.Decoder) (string, error) {
|
||||
reader := transform.NewReader(bytes.NewReader(payload), decoder)
|
||||
decoded, err := io.ReadAll(reader)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !strings.ContainsRune(string(decoded), '\ufffd') {
|
||||
return string(decoded), nil
|
||||
}
|
||||
return "", fmt.Errorf("decode produced replacement characters")
|
||||
}
|
||||
157
internal/parser/parser/email_parser_test.go
Normal file
157
internal/parser/parser/email_parser_test.go
Normal file
@@ -0,0 +1,157 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestEmailParser_EmlJSON(t *testing.T) {
|
||||
raw := strings.Join([]string{
|
||||
"From: sender@example.com",
|
||||
"To: recipient@example.com",
|
||||
"Cc: cc@example.com",
|
||||
"Date: Mon, 07 Jul 2025 10:00:00 +0000",
|
||||
"Subject: Test Email",
|
||||
"Content-Type: text/plain; charset=utf-8",
|
||||
"X-Custom-Header: custom-value",
|
||||
"",
|
||||
"This is the body of the test email.",
|
||||
}, "\r\n")
|
||||
|
||||
p := NewEmailParser()
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"output_format": "json",
|
||||
"fields": []string{"from", "to", "cc", "date", "subject", "body", "metadata"},
|
||||
})
|
||||
|
||||
result := p.ParseWithResult("test.eml", []byte(raw))
|
||||
if result.Err != nil {
|
||||
t.Fatalf("unexpected error: %v", result.Err)
|
||||
}
|
||||
if result.OutputFormat != "json" {
|
||||
t.Fatalf("expected output_format json, got %q", result.OutputFormat)
|
||||
}
|
||||
if len(result.JSON) != 1 {
|
||||
t.Fatalf("expected 1 JSON item, got %d", len(result.JSON))
|
||||
}
|
||||
item := result.JSON[0]
|
||||
|
||||
if v, ok := item["from"].(string); !ok || v != "sender@example.com" {
|
||||
t.Errorf("from: got %q", v)
|
||||
}
|
||||
if v, ok := item["to"].(string); !ok || v != "recipient@example.com" {
|
||||
t.Errorf("to: got %q", v)
|
||||
}
|
||||
if v, ok := item["subject"].(string); !ok || v != "Test Email" {
|
||||
t.Errorf("subject: got %q", v)
|
||||
}
|
||||
if v, ok := item["text"].(string); !ok || !strings.Contains(v, "body of the test email") {
|
||||
t.Errorf("text: got %q", v)
|
||||
}
|
||||
if meta, ok := item["metadata"].(map[string]any); ok {
|
||||
if v, ok := meta["x-custom-header"].(string); !ok || v != "custom-value" {
|
||||
t.Errorf("metadata x-custom-header: got %q", v)
|
||||
}
|
||||
} else {
|
||||
t.Error("metadata missing or wrong type")
|
||||
}
|
||||
if v, ok := item["doc_type_kwd"].(string); !ok || v != "text" {
|
||||
t.Errorf("doc_type_kwd: got %q", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmailParser_EmlText(t *testing.T) {
|
||||
raw := strings.Join([]string{
|
||||
"From: sender@test.com",
|
||||
"To: recipient@test.com",
|
||||
"Subject: Hello",
|
||||
"Content-Type: text/plain; charset=utf-8",
|
||||
"",
|
||||
"Hello, world!",
|
||||
}, "\r\n")
|
||||
|
||||
p := NewEmailParser()
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"output_format": "text",
|
||||
"fields": []string{"from", "to", "subject", "body"},
|
||||
})
|
||||
|
||||
result := p.ParseWithResult("test.eml", []byte(raw))
|
||||
if result.Err != nil {
|
||||
t.Fatalf("unexpected error: %v", result.Err)
|
||||
}
|
||||
if result.OutputFormat != "text" {
|
||||
t.Fatalf("expected output_format text, got %q", result.OutputFormat)
|
||||
}
|
||||
if !strings.Contains(result.Text, "Hello, world!") {
|
||||
t.Errorf("text missing body: %q", result.Text)
|
||||
}
|
||||
if !strings.Contains(result.Text, "sender@test.com") {
|
||||
t.Errorf("text missing from: %q", result.Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmailParser_MsgNotSupported(t *testing.T) {
|
||||
p := NewEmailParser()
|
||||
result := p.ParseWithResult("test.msg", []byte{})
|
||||
if result.Err == nil {
|
||||
t.Fatal("expected error for .msg file")
|
||||
}
|
||||
if !strings.Contains(result.Err.Error(), ".msg") {
|
||||
t.Errorf("error should mention .msg: %v", result.Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmailParser_Multipart(t *testing.T) {
|
||||
boundary := "boundary123"
|
||||
raw := strings.Join([]string{
|
||||
"From: multipart@test.com",
|
||||
"To: receiver@test.com",
|
||||
"Subject: Multipart Test",
|
||||
"Content-Type: multipart/alternative; boundary=" + boundary,
|
||||
"",
|
||||
"--" + boundary,
|
||||
"Content-Type: text/plain; charset=utf-8",
|
||||
"",
|
||||
"Plain text body.",
|
||||
"--" + boundary,
|
||||
"Content-Type: text/html; charset=utf-8",
|
||||
"",
|
||||
"<p>HTML body.</p>",
|
||||
"--" + boundary + "--",
|
||||
}, "\r\n")
|
||||
|
||||
p := NewEmailParser()
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"output_format": "json",
|
||||
"fields": []string{"from", "body"},
|
||||
})
|
||||
|
||||
result := p.ParseWithResult("test.eml", []byte(raw))
|
||||
if result.Err != nil {
|
||||
t.Fatalf("unexpected error: %v", result.Err)
|
||||
}
|
||||
item := result.JSON[0]
|
||||
if v, ok := item["text"].(string); !ok || !strings.Contains(v, "Plain text body") {
|
||||
t.Errorf("text: got %q", v)
|
||||
}
|
||||
if v, ok := item["text_html"].(string); !ok || !strings.Contains(v, "HTML body") {
|
||||
t.Errorf("text_html: got %q", v)
|
||||
}
|
||||
}
|
||||
311
internal/parser/parser/epub_parser.go
Normal file
311
internal/parser/parser/epub_parser.go
Normal file
@@ -0,0 +1,311 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// EPUBParser extracts text content from EPUB files. EPUB is a ZIP
|
||||
// container with XHTML spine items. The parser:
|
||||
//
|
||||
// 1. Opens the EPUB as a ZIP archive.
|
||||
// 2. Reads META-INF/container.xml to locate the OPF manifest.
|
||||
// 3. Walks the OPF spine to collect content documents in order.
|
||||
// 4. Strips HTML tags from each content document to produce plain
|
||||
// text items, emitting one JSON item per spine entry.
|
||||
//
|
||||
// The output format is "json", matching the epub default in
|
||||
// ParserParam.Defaults().
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"encoding/xml"
|
||||
"fmt"
|
||||
"io"
|
||||
"path"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// EPUBParser extracts text from EPUB archives.
|
||||
type EPUBParser struct{}
|
||||
|
||||
func NewEPUBParser() *EPUBParser {
|
||||
return &EPUBParser{}
|
||||
}
|
||||
|
||||
func (p *EPUBParser) String() string {
|
||||
return "EPUBParser"
|
||||
}
|
||||
|
||||
// ParseWithResult implements ParseResultProducer. It extracts XHTML
|
||||
// content from the EPUB spine and emits one JSON item per spine entry
|
||||
// with {text, doc_type_kwd:"text"}.
|
||||
func (p *EPUBParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return ParseResult{
|
||||
OutputFormat: "json",
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"size": 0,
|
||||
"encoding": "utf-8",
|
||||
},
|
||||
JSON: []map[string]any{{"text": "", "doc_type_kwd": "text"}},
|
||||
}
|
||||
}
|
||||
|
||||
reader, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("epub: open zip: %w", err)}
|
||||
}
|
||||
|
||||
opfPath, err := readContainerXML(reader)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("epub: container.xml: %w", err)}
|
||||
}
|
||||
|
||||
spineItems, err := readOPFSpine(reader, opfPath)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("epub: opf: %w", err)}
|
||||
}
|
||||
|
||||
items := extractEPUBTextItems(reader, opfPath, spineItems)
|
||||
if items == nil {
|
||||
items = []map[string]any{{"text": "", "doc_type_kwd": "text"}}
|
||||
}
|
||||
return ParseResult{
|
||||
OutputFormat: "json",
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"size": len(data),
|
||||
"encoding": "utf-8",
|
||||
},
|
||||
JSON: items,
|
||||
}
|
||||
}
|
||||
|
||||
// --- container.xml parsing ---
|
||||
|
||||
type containerXML struct {
|
||||
RootFiles []struct {
|
||||
FullPath string `xml:"full-path,attr"`
|
||||
} `xml:"rootfiles>rootfile"`
|
||||
}
|
||||
|
||||
func readContainerXML(r *zip.Reader) (string, error) {
|
||||
f, err := r.Open("META-INF/container.xml")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("open container.xml: %w", err)
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
raw, err := io.ReadAll(f)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("read container.xml: %w", err)
|
||||
}
|
||||
var c containerXML
|
||||
if err := xml.Unmarshal(raw, &c); err != nil {
|
||||
return "", fmt.Errorf("parse container.xml: %w", err)
|
||||
}
|
||||
if len(c.RootFiles) == 0 {
|
||||
return "", fmt.Errorf("no rootfile in container.xml")
|
||||
}
|
||||
return c.RootFiles[0].FullPath, nil
|
||||
}
|
||||
|
||||
// --- OPF parsing ---
|
||||
|
||||
type opfPackage struct {
|
||||
Manifest struct {
|
||||
Items []struct {
|
||||
ID string `xml:"id,attr"`
|
||||
Href string `xml:"href,attr"`
|
||||
Type string `xml:"media-type,attr"`
|
||||
} `xml:"item"`
|
||||
} `xml:"manifest"`
|
||||
Spine struct {
|
||||
ItemRefs []struct {
|
||||
IDRef string `xml:"idref,attr"`
|
||||
} `xml:"itemref"`
|
||||
} `xml:"spine"`
|
||||
}
|
||||
|
||||
func readOPFSpine(r *zip.Reader, opfPath string) ([]string, error) {
|
||||
f, err := r.Open(opfPath)
|
||||
if err != nil {
|
||||
// Some EPUBs use a path with a leading slash stripped.
|
||||
clean := strings.TrimPrefix(opfPath, "/")
|
||||
if clean != opfPath {
|
||||
f, err = r.Open(clean)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open opf %q: %w", opfPath, err)
|
||||
}
|
||||
} else {
|
||||
return nil, fmt.Errorf("open opf %q: %w", opfPath, err)
|
||||
}
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
raw, err := io.ReadAll(f)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read opf: %w", err)
|
||||
}
|
||||
var pkg opfPackage
|
||||
if err := xml.Unmarshal(raw, &pkg); err != nil {
|
||||
return nil, fmt.Errorf("parse opf: %w", err)
|
||||
}
|
||||
|
||||
// Build id → href lookup.
|
||||
byID := make(map[string]string, len(pkg.Manifest.Items))
|
||||
for _, it := range pkg.Manifest.Items {
|
||||
byID[it.ID] = it.Href
|
||||
}
|
||||
|
||||
// Build spine item list (ordered hrefs).
|
||||
var hrefs []string
|
||||
for _, ref := range pkg.Spine.ItemRefs {
|
||||
if href, ok := byID[ref.IDRef]; ok {
|
||||
hrefs = append(hrefs, href)
|
||||
}
|
||||
}
|
||||
return hrefs, nil
|
||||
}
|
||||
|
||||
// --- text extraction ---
|
||||
|
||||
var (
|
||||
epubScriptRe = regexp.MustCompile(`(?is)<script[^>]*>.*?</script>`)
|
||||
epubStyleRe = regexp.MustCompile(`(?is)<style[^>]*>.*?</style>`)
|
||||
epubTagRe = regexp.MustCompile(`<[^>]+>`)
|
||||
epubEntityRe = regexp.MustCompile(`&[a-zA-Z]+;`)
|
||||
epubWSRe = regexp.MustCompile(`\s+`)
|
||||
)
|
||||
|
||||
func extractEPUBTextItems(r *zip.Reader, opfDir string, spineHrefs []string) []map[string]any {
|
||||
var items []map[string]any
|
||||
for _, href := range spineHrefs {
|
||||
text := readEPUBContentFile(r, opfDir, href)
|
||||
if strings.TrimSpace(text) == "" {
|
||||
continue
|
||||
}
|
||||
items = append(items, map[string]any{
|
||||
"text": text,
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
}
|
||||
return items
|
||||
}
|
||||
|
||||
// readEPUBContentFile resolves a spine href (relative to the OPF
|
||||
// directory) inside the ZIP, reads the raw bytes, and strips HTML to
|
||||
// return clean text.
|
||||
func readEPUBContentFile(r *zip.Reader, opfDir, href string) string {
|
||||
// Resolve href relative to the directory containing the OPF.
|
||||
// Use path.Join/Dir (slash separator) because ZIP entry paths are always POSIX.
|
||||
resolved := path.Join(path.Dir(opfDir), href)
|
||||
|
||||
var f io.ReadCloser
|
||||
var err error
|
||||
// Try the resolved path first, then raw href.
|
||||
for _, name := range []string{resolved, href} {
|
||||
if strings.TrimSpace(name) == "" {
|
||||
continue
|
||||
}
|
||||
f, err = r.Open(name)
|
||||
if err == nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
raw, err := io.ReadAll(f)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return stripHTMLTags(string(raw))
|
||||
}
|
||||
|
||||
// stripHTMLTags removes HTML markup and returns normalized plain text.
|
||||
func stripHTMLTags(s string) string {
|
||||
// Remove script/style blocks first (including their content).
|
||||
s = epubScriptRe.ReplaceAllString(s, "")
|
||||
s = epubStyleRe.ReplaceAllString(s, "")
|
||||
// Replace <br>, <p>, </p>, </div>, </tr>, </li> etc. with newlines.
|
||||
s = epubTagRe.ReplaceAllStringFunc(s, func(tag string) string {
|
||||
lower := strings.ToLower(tag)
|
||||
switch {
|
||||
case strings.HasPrefix(lower, "<br"),
|
||||
strings.HasPrefix(lower, "</p"),
|
||||
strings.HasPrefix(lower, "</div"),
|
||||
strings.HasPrefix(lower, "</h1"),
|
||||
strings.HasPrefix(lower, "</h2"),
|
||||
strings.HasPrefix(lower, "</h3"),
|
||||
strings.HasPrefix(lower, "</h4"),
|
||||
strings.HasPrefix(lower, "</h5"),
|
||||
strings.HasPrefix(lower, "</h6"),
|
||||
strings.HasPrefix(lower, "</tr"),
|
||||
strings.HasPrefix(lower, "</li"),
|
||||
strings.HasPrefix(lower, "</blockquote"):
|
||||
return "\n"
|
||||
case strings.HasPrefix(lower, "</td"),
|
||||
strings.HasPrefix(lower, "</th"):
|
||||
return " "
|
||||
}
|
||||
return ""
|
||||
})
|
||||
// Decode common XML entities.
|
||||
s = epubEntityRe.ReplaceAllStringFunc(s, xmlEntityValue)
|
||||
// Normalize whitespace.
|
||||
s = epubWSRe.ReplaceAllString(s, " ")
|
||||
return strings.TrimSpace(s)
|
||||
}
|
||||
|
||||
// xmlEntityValue decodes a few common XML named entities. A full
|
||||
// entity table isn't needed for the typical EPUB vocabulary.
|
||||
func xmlEntityValue(entity string) string {
|
||||
switch entity {
|
||||
case "&":
|
||||
return "&"
|
||||
case "<":
|
||||
return "<"
|
||||
case ">":
|
||||
return ">"
|
||||
case """:
|
||||
return "\""
|
||||
case "'":
|
||||
return "'"
|
||||
case " ":
|
||||
return " "
|
||||
case "–":
|
||||
return "–"
|
||||
case "—":
|
||||
return "—"
|
||||
case "‘":
|
||||
return "'"
|
||||
case "’":
|
||||
return "'"
|
||||
case "“":
|
||||
return "\""
|
||||
case "”":
|
||||
return "\""
|
||||
case "…":
|
||||
return "…"
|
||||
default:
|
||||
return entity
|
||||
}
|
||||
}
|
||||
@@ -52,7 +52,7 @@ func GetParserByID(parserID string) (ParseResultProducer, error) {
|
||||
case "resume":
|
||||
return NewResumePDFParser(), nil
|
||||
case "picture":
|
||||
return NewPicturePDFParser(), nil
|
||||
return NewPictureParser(), nil
|
||||
case "one":
|
||||
return NewOnePDFParser(), nil
|
||||
case "audio":
|
||||
@@ -80,7 +80,5 @@ func NewTableParser() *stubParser { return &stubParser{name: "table"} }
|
||||
func NewResumePDFParser() *stubParser { return &stubParser{name: "resume"} }
|
||||
func NewPicturePDFParser() *stubParser { return &stubParser{name: "picture"} }
|
||||
func NewOnePDFParser() *stubParser { return &stubParser{name: "one"} }
|
||||
func NewAudioParser() *stubParser { return &stubParser{name: "audio"} }
|
||||
func NewEmailParser() *stubParser { return &stubParser{name: "email"} }
|
||||
func NewTagPDFParser() *stubParser { return &stubParser{name: "tag"} }
|
||||
func NewKGPDFParser() *stubParser { return &stubParser{name: "knowledge_graph"} }
|
||||
|
||||
@@ -37,7 +37,6 @@ func TestGetParserByID_AllKnownIDs(t *testing.T) {
|
||||
{string(entity.ParserTypeResume), false},
|
||||
{string(entity.ParserTypePicture), false},
|
||||
{string(entity.ParserTypeOne), false},
|
||||
{string(entity.ParserTypeKG), false},
|
||||
{string(entity.ParserTypeTag), false},
|
||||
// Office parsers
|
||||
{string(entity.ParserTypePresentation), false},
|
||||
|
||||
228
internal/parser/parser/json_parser.go
Normal file
228
internal/parser/parser/json_parser.go
Normal file
@@ -0,0 +1,228 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// JSONParser handles .json, .jsonl, and .ldjson files. It detects
|
||||
// the payload shape:
|
||||
//
|
||||
// - JSON array: each element becomes one item.
|
||||
// - Single JSON object: emitted as one item.
|
||||
// - JSONL (newline-delimited JSON): each non-empty line is
|
||||
// independently parsed; lines that fail to parse are skipped
|
||||
// with a warning — matching the Python JsonParser behaviour.
|
||||
//
|
||||
// For JSON objects, the parser serialises the full object into a
|
||||
// `text` field so the downstream chunker receives a readable
|
||||
// representation. Arrays of objects emit one item per element.
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// JSONParser handles JSON / JSONL / LDJSON files.
|
||||
type JSONParser struct{}
|
||||
|
||||
func NewJSONParser() *JSONParser {
|
||||
return &JSONParser{}
|
||||
}
|
||||
|
||||
func (p *JSONParser) String() string {
|
||||
return "JSONParser"
|
||||
}
|
||||
|
||||
// ParseWithResult implements ParseResultProducer. It detects the JSON
|
||||
// shape (array, single object, or line-delimited) and emits one JSON
|
||||
// item per logical record, with {text, doc_type_kwd:"text"}.
|
||||
func (p *JSONParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
text := string(bytes.TrimSpace(data))
|
||||
if text == "" {
|
||||
return ParseResult{
|
||||
OutputFormat: "json",
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"size": len(data),
|
||||
"encoding": "utf-8",
|
||||
},
|
||||
JSON: []map[string]any{{"text": "", "doc_type_kwd": "text"}},
|
||||
}
|
||||
}
|
||||
|
||||
items := parseJSONContent(text)
|
||||
if items == nil {
|
||||
items = []map[string]any{{"text": "", "doc_type_kwd": "text"}}
|
||||
}
|
||||
return ParseResult{
|
||||
OutputFormat: "json",
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"size": len(data),
|
||||
"encoding": "utf-8",
|
||||
},
|
||||
JSON: items,
|
||||
}
|
||||
}
|
||||
|
||||
// parseJSONContent detects the JSON shape and dispatches.
|
||||
func parseJSONContent(text string) []map[string]any {
|
||||
trimmed := strings.TrimSpace(text)
|
||||
if trimmed == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 1. Try JSON array: [...]
|
||||
if strings.HasPrefix(trimmed, "[") {
|
||||
return parseJSONArray(trimmed)
|
||||
}
|
||||
|
||||
// 2. Try single JSON object: {...}
|
||||
if strings.HasPrefix(trimmed, "{") {
|
||||
// Prefer parsing as a single JSON object first (matches
|
||||
// Python's json.loads(txt) whole-document test). Only fall
|
||||
// back to JSONL when that fails and the content still looks
|
||||
// line-delimited.
|
||||
if err := json.Unmarshal([]byte(trimmed), &map[string]any{}); err == nil {
|
||||
return parseSingleJSONObject(trimmed)
|
||||
}
|
||||
if isJSONL(trimmed) {
|
||||
return parseJSONLines(trimmed)
|
||||
}
|
||||
return []map[string]any{{"text": trimmed, "doc_type_kwd": "text"}}
|
||||
}
|
||||
|
||||
// 3. Fallback: treat as text.
|
||||
return []map[string]any{{"text": trimmed, "doc_type_kwd": "text"}}
|
||||
}
|
||||
|
||||
// isJSONL returns true when the content looks like newline-delimited
|
||||
// JSON (multiple lines each starting with '{').
|
||||
func isJSONL(text string) bool {
|
||||
lines := strings.Split(text, "\n")
|
||||
objCount := 0
|
||||
for _, line := range lines {
|
||||
t := strings.TrimSpace(line)
|
||||
if t == "" {
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(t, "{") {
|
||||
objCount++
|
||||
if objCount >= 2 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// parseJSONArray unmarshals a JSON array and emits one item per element.
|
||||
func parseJSONArray(text string) []map[string]any {
|
||||
var arr []any
|
||||
if err := json.Unmarshal([]byte(text), &arr); err != nil {
|
||||
// Fallback: emit as plain text.
|
||||
return []map[string]any{{"text": text, "doc_type_kwd": "text"}}
|
||||
}
|
||||
if len(arr) == 0 {
|
||||
return nil
|
||||
}
|
||||
items := make([]map[string]any, 0, len(arr))
|
||||
for _, elem := range arr {
|
||||
if s, ok := elem.(string); ok {
|
||||
items = append(items, map[string]any{
|
||||
"text": s,
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
} else {
|
||||
// Re-marshal non-string elements.
|
||||
b, err := json.Marshal(elem)
|
||||
if err != nil {
|
||||
items = append(items, map[string]any{
|
||||
"text": fmt.Sprintf("%v", elem),
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
} else {
|
||||
items = append(items, map[string]any{
|
||||
"text": string(b),
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
return items
|
||||
}
|
||||
|
||||
// parseSingleJSONObject unmarshals a single JSON object and emits one item.
|
||||
func parseSingleJSONObject(text string) []map[string]any {
|
||||
var obj map[string]any
|
||||
if err := json.Unmarshal([]byte(text), &obj); err != nil {
|
||||
return []map[string]any{{"text": text, "doc_type_kwd": "text"}}
|
||||
}
|
||||
b, err := json.Marshal(obj)
|
||||
if err != nil {
|
||||
return []map[string]any{{"text": fmt.Sprintf("%v", obj), "doc_type_kwd": "text"}}
|
||||
}
|
||||
return []map[string]any{{"text": string(b), "doc_type_kwd": "text"}}
|
||||
}
|
||||
|
||||
// parseJSONLines parses line-delimited JSON. Each non-empty line is
|
||||
// independently unmarshalled; lines that fail to parse are skipped.
|
||||
func parseJSONLines(text string) []map[string]any {
|
||||
var items []map[string]any
|
||||
scanner := bufio.NewScanner(strings.NewReader(text))
|
||||
scanner.Buffer(nil, 10*1024*1024) // 10 MB max line
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
// Each line should be a JSON object.
|
||||
var obj map[string]any
|
||||
if err := json.Unmarshal([]byte(line), &obj); err != nil {
|
||||
// Also try as a string or other standalone value.
|
||||
var val any
|
||||
if err2 := json.Unmarshal([]byte(line), &val); err2 != nil {
|
||||
// Unparseable line — skip, matching Python behaviour.
|
||||
continue
|
||||
} else {
|
||||
if s, ok := val.(string); ok {
|
||||
items = append(items, map[string]any{
|
||||
"text": s,
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
} else {
|
||||
b, _ := json.Marshal(val)
|
||||
items = append(items, map[string]any{
|
||||
"text": string(b),
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
b, err := json.Marshal(obj)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
items = append(items, map[string]any{
|
||||
"text": string(b),
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
}
|
||||
return items
|
||||
}
|
||||
@@ -18,32 +18,84 @@ package parser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gomarkdown/markdown/ast"
|
||||
"github.com/gomarkdown/markdown/parser"
|
||||
)
|
||||
|
||||
type MarkdownParser struct{}
|
||||
// mdImagePattern matches markdown inline image syntax: .
|
||||
var mdImagePattern = regexp.MustCompile(`!\[[^\]]*\]\(([^)\s]+)\)`)
|
||||
|
||||
func NewMarkdownParser() *MarkdownParser {
|
||||
return &MarkdownParser{}
|
||||
// dataURIPrefix is the MIME prefix for data URI images.
|
||||
const dataURIPrefix = "data:image/"
|
||||
|
||||
// GoMarkdown is the lib_type identifier for the pure-Go markdown backend.
|
||||
const GoMarkdown = "go_markdown"
|
||||
|
||||
// ssrfAllowLoopback lets tests exercise the HTTP image fetch path against a
|
||||
// loopback httptest server. It stays false in production so loopback
|
||||
// addresses are rejected (SSRF protection).
|
||||
var ssrfAllowLoopback bool
|
||||
|
||||
type MarkdownParser struct {
|
||||
libType string
|
||||
ParseMethod string
|
||||
OutputFormat string
|
||||
VLM map[string]any
|
||||
}
|
||||
|
||||
func NewMarkdownParser(libType string) (*MarkdownParser, error) {
|
||||
switch libType {
|
||||
case GoMarkdown:
|
||||
return &MarkdownParser{
|
||||
libType: GoMarkdown,
|
||||
}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported Markdown library type: %s", libType)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *MarkdownParser) ConfigureFromSetup(setup map[string]any) {
|
||||
if p == nil || setup == nil {
|
||||
return
|
||||
}
|
||||
if v, ok := setup["parse_method"].(string); ok && v != "" {
|
||||
p.ParseMethod = v
|
||||
}
|
||||
if v, ok := setup["output_format"].(string); ok && v != "" {
|
||||
p.OutputFormat = v
|
||||
}
|
||||
if v, ok := setup["vlm"].(map[string]any); ok {
|
||||
p.VLM = v
|
||||
}
|
||||
}
|
||||
|
||||
// ParseWithResult implements ParseResultProducer (plan §6.5) and
|
||||
// returns a structured markdown payload that mirrors the Python
|
||||
// parser's `output_format == "json"` shape. Each top-level block
|
||||
// emits one item with `text` + `doc_type_kwd: "text"`. The legacy
|
||||
// debug-print path has been removed; callers consume ParseResult directly.
|
||||
// emits one item with `text` + `doc_type_kwd: "text"`. When the
|
||||
// block contains a markdown image reference (), the image
|
||||
// data is resolved and the item carries `doc_type_kwd: "image"` with
|
||||
// the base64-encoded image payload. The legacy debug-print path has
|
||||
// been removed; callers consume ParseResult directly.
|
||||
func (p *MarkdownParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
doc := markdownNew().Parse(data)
|
||||
rawText := string(data)
|
||||
|
||||
var items []map[string]any
|
||||
walkMarkdownBlocks(doc, &items)
|
||||
walkMarkdownBlocksWithImages(doc, rawText, &items)
|
||||
if items == nil {
|
||||
// No blocks emitted — surface a single empty item so the
|
||||
// downstream chunker sees a non-nil JSON slice (mirrors the
|
||||
// Python contract of always producing at least one item).
|
||||
items = []map[string]any{{"text": "", "doc_type_kwd": "text"}}
|
||||
}
|
||||
return ParseResult{
|
||||
@@ -66,55 +118,222 @@ func markdownNew() *parser.Parser {
|
||||
return parser.NewWithExtensions(extensions)
|
||||
}
|
||||
|
||||
// walkMarkdownBlocks emits one normalized item per top-level block.
|
||||
// Headings (LeafBlock / H*) are emitted with the heading text so a
|
||||
// downstream title chunker can find them; paragraphs (Paragraph)
|
||||
// emit the leaf-node text. The walker is intentionally a
|
||||
// best-effort pass — full TOC / outline handling lands with the
|
||||
// deepdoc/parser port — but it satisfies the per-block "emit
|
||||
// text+doc_type_kwd" contract enough for a Phase-1a migration
|
||||
// test to verify wire-shape parity.
|
||||
func walkMarkdownBlocks(doc ast.Node, out *[]map[string]any) {
|
||||
// walkMarkdownBlocksWithImages emits one normalized item per
|
||||
// top-level block. Headings, paragraphs, lists, and code blocks are
|
||||
// emitted with their text. When a block contains a markdown image
|
||||
// reference (), the image data is resolved via
|
||||
// resolveMarkdownImage and the item carries `doc_type_kwd: "image"`
|
||||
// together with the base64-encoded image payload.
|
||||
func walkMarkdownBlocksWithImages(doc ast.Node, rawText string, out *[]map[string]any) {
|
||||
for _, child := range doc.GetChildren() {
|
||||
var ckType string
|
||||
var docTypeKwd string
|
||||
var txt string
|
||||
|
||||
switch n := child.(type) {
|
||||
case *ast.Heading:
|
||||
*out = append(*out, map[string]any{
|
||||
"text": headingText(n),
|
||||
"doc_type_kwd": "text",
|
||||
"ck_type": "heading",
|
||||
})
|
||||
txt = headingText(n)
|
||||
ckType = "heading"
|
||||
docTypeKwd = "text"
|
||||
case *ast.Paragraph:
|
||||
*out = append(*out, map[string]any{
|
||||
"text": leafText(n),
|
||||
"doc_type_kwd": "text",
|
||||
"ck_type": "text",
|
||||
})
|
||||
txt = leafText(n)
|
||||
ckType = "text"
|
||||
docTypeKwd = "text"
|
||||
case *ast.List:
|
||||
*out = append(*out, map[string]any{
|
||||
"text": leafText(n),
|
||||
"doc_type_kwd": "text",
|
||||
"ck_type": "list",
|
||||
})
|
||||
txt = leafText(n)
|
||||
ckType = "list"
|
||||
docTypeKwd = "text"
|
||||
case *ast.CodeBlock:
|
||||
*out = append(*out, map[string]any{
|
||||
"text": leafText(n),
|
||||
"doc_type_kwd": "text",
|
||||
"ck_type": "code",
|
||||
})
|
||||
txt = leafText(n)
|
||||
ckType = "code"
|
||||
docTypeKwd = "text"
|
||||
default:
|
||||
// Block types we don't yet normalize (HTML, tables,
|
||||
// images, definitions) are best-effort: emit the leaf
|
||||
// text without a ck_type so downstream components can
|
||||
// still treat them as text chunks.
|
||||
txt := leafText(n)
|
||||
if strings.TrimSpace(txt) != "" {
|
||||
*out = append(*out, map[string]any{
|
||||
"text": txt,
|
||||
"doc_type_kwd": "text",
|
||||
})
|
||||
txt = leafText(n)
|
||||
if strings.TrimSpace(txt) == "" {
|
||||
continue
|
||||
}
|
||||
docTypeKwd = "text"
|
||||
}
|
||||
|
||||
item := map[string]any{
|
||||
"text": txt,
|
||||
"doc_type_kwd": docTypeKwd,
|
||||
}
|
||||
if ckType != "" {
|
||||
item["ck_type"] = ckType
|
||||
}
|
||||
|
||||
// Detect markdown image references in the raw source text
|
||||
// that corresponds to this block. When found, resolve the
|
||||
// image data so downstream vision enhancement can describe it.
|
||||
if imgData, imgFound := resolveMarkdownImage(txt, rawText); imgFound && imgData != "" {
|
||||
item["doc_type_kwd"] = "image"
|
||||
item["image"] = imgData
|
||||
}
|
||||
|
||||
*out = append(*out, item)
|
||||
}
|
||||
}
|
||||
|
||||
// resolveMarkdownImage extracts the first markdown image reference
|
||||
// from the given text and returns its base64-encoded data. Supports:
|
||||
// - data:image/... URIs → decoded directly
|
||||
// - http:// / https:// URLs → fetched (with basic SSRF filtering)
|
||||
//
|
||||
// Returns (base64String, true) on success, ("", false) when no image
|
||||
// is found or resolution fails.
|
||||
func resolveMarkdownImage(leafText, rawFullText string) (string, bool) {
|
||||
// Prefer matching against the raw full text to catch images
|
||||
// whose alt-text was split across markdown rendering.
|
||||
searchIn := rawFullText
|
||||
if strings.TrimSpace(searchIn) == "" {
|
||||
searchIn = leafText
|
||||
}
|
||||
matches := mdImagePattern.FindStringSubmatch(searchIn)
|
||||
if len(matches) < 2 {
|
||||
return "", false
|
||||
}
|
||||
url := matches[1]
|
||||
|
||||
if strings.HasPrefix(url, dataURIPrefix) {
|
||||
// data:image/png;base64,xxxx
|
||||
idx := strings.Index(url, "base64,")
|
||||
if idx < 0 {
|
||||
return "", false
|
||||
}
|
||||
return url[idx+len("base64,"):], true
|
||||
}
|
||||
if strings.HasPrefix(url, "http://") || strings.HasPrefix(url, "https://") {
|
||||
b64, err := fetchImageAsBase64(url)
|
||||
if err != nil {
|
||||
return "", false
|
||||
}
|
||||
return b64, true
|
||||
}
|
||||
// Local / relative paths — not fetched for security.
|
||||
return "", false
|
||||
}
|
||||
|
||||
// fetchImageAsBase64 fetches an HTTP(S) image URL and returns its
|
||||
// content as a base64-encoded string. Local/private addresses and
|
||||
// redirects to them are rejected (SSRF guard). Hostnames are resolved
|
||||
// once and the validated IP is pinned in a custom DialContext to
|
||||
// prevent DNS-rebinding TOCTOU attacks.
|
||||
func fetchImageAsBase64(rawURL string) (string, error) {
|
||||
rawURL = strings.TrimSpace(rawURL)
|
||||
|
||||
parsed, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("markdown: invalid image URL: %w", err)
|
||||
}
|
||||
|
||||
// pinned maps hostname (without port) → validated IP. The hostname
|
||||
// is resolved once per host, and the transport dials the pinned IP
|
||||
// directly instead of re-resolving DNS.
|
||||
var pinnedMu sync.Mutex
|
||||
pinned := make(map[string]net.IP)
|
||||
|
||||
pinHost := func(host string) error {
|
||||
ip, err := resolveAndValidateHost(host)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
h, _, _ := net.SplitHostPort(host)
|
||||
if h == "" {
|
||||
h = host
|
||||
}
|
||||
pinnedMu.Lock()
|
||||
pinned[h] = ip
|
||||
pinnedMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := pinHost(parsed.Host); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
transport := &http.Transport{
|
||||
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
host, port, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pinnedMu.Lock()
|
||||
ip, ok := pinned[host]
|
||||
pinnedMu.Unlock()
|
||||
if ok {
|
||||
return (&net.Dialer{}).DialContext(ctx, network, net.JoinHostPort(ip.String(), port))
|
||||
}
|
||||
return (&net.Dialer{}).DialContext(ctx, network, addr)
|
||||
},
|
||||
}
|
||||
|
||||
client := &http.Client{
|
||||
Transport: transport,
|
||||
Timeout: 30 * time.Second,
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= 5 {
|
||||
return fmt.Errorf("markdown: too many redirects")
|
||||
}
|
||||
return pinHost(req.URL.Host)
|
||||
},
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("markdown: create image request: %w", err)
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("markdown: fetch image %s: %w", rawURL, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return "", fmt.Errorf("markdown: fetch image %s: HTTP %d", rawURL, resp.StatusCode)
|
||||
}
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 32*1024*1024)) // 32 MiB cap
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("markdown: read image %s: %w", rawURL, err)
|
||||
}
|
||||
return base64.StdEncoding.EncodeToString(body), nil
|
||||
}
|
||||
|
||||
// resolveAndValidateHost resolves a host (which may include a port),
|
||||
// validates none of its IPs are internal/private, and returns the
|
||||
// first public IP for connection pinning.
|
||||
func resolveAndValidateHost(host string) (net.IP, error) {
|
||||
hostname := host
|
||||
if h, _, err := net.SplitHostPort(host); err == nil {
|
||||
hostname = h
|
||||
}
|
||||
|
||||
ip := net.ParseIP(hostname)
|
||||
if ip != nil {
|
||||
if (ip.IsLoopback() && !ssrfAllowLoopback) || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() ||
|
||||
ip.IsPrivate() || ip.IsUnspecified() {
|
||||
return nil, fmt.Errorf("markdown: rejected image URL to internal address: %s", host)
|
||||
}
|
||||
return ip, nil
|
||||
}
|
||||
|
||||
addrs, err := net.DefaultResolver.LookupIPAddr(context.Background(), hostname)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("markdown: cannot resolve image host: %s", hostname)
|
||||
}
|
||||
for _, addr := range addrs {
|
||||
ip := addr.IP
|
||||
if (ip.IsLoopback() && !ssrfAllowLoopback) || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() ||
|
||||
ip.IsPrivate() || ip.IsUnspecified() {
|
||||
return nil, fmt.Errorf("markdown: rejected image URL resolving to internal address: %s (%s)", host, ip)
|
||||
}
|
||||
}
|
||||
if len(addrs) == 0 {
|
||||
return nil, fmt.Errorf("markdown: no addresses resolved for host: %s", hostname)
|
||||
}
|
||||
return addrs[0].IP, nil
|
||||
}
|
||||
|
||||
// headingText returns the inline-text of a heading node by
|
||||
|
||||
179
internal/parser/parser/markdown_parser_test.go
Normal file
179
internal/parser/parser/markdown_parser_test.go
Normal file
@@ -0,0 +1,179 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMarkdownParser_ParseWithResult_Basic(t *testing.T) {
|
||||
p, err := NewMarkdownParser(GoMarkdown)
|
||||
if err != nil {
|
||||
t.Fatalf("NewMarkdownParser: %v", err)
|
||||
}
|
||||
md := "# Hello\n\nThis is a paragraph.\n\n* List item 1\n* List item 2\n\n```go\nfunc main() {}\n```\n"
|
||||
res := p.ParseWithResult("test.md", []byte(md))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "json"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if len(res.JSON) == 0 {
|
||||
t.Fatal("JSON is empty; want at least one item")
|
||||
}
|
||||
// Verify heading
|
||||
if got, _ := res.JSON[0]["text"].(string); got != "Hello" {
|
||||
t.Fatalf("first item text = %q, want %q", got, "Hello")
|
||||
}
|
||||
if got, _ := res.JSON[0]["ck_type"].(string); got != "heading" {
|
||||
t.Fatalf("first item ck_type = %q, want %q", got, "heading")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkdownParser_ParseWithResult_EmptyInput(t *testing.T) {
|
||||
p, _ := NewMarkdownParser(GoMarkdown)
|
||||
res := p.ParseWithResult("empty.md", []byte(""))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) != 1 {
|
||||
t.Fatalf("len(JSON) = %d, want 1", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkdownParser_ParseWithResult_ImageDataURI(t *testing.T) {
|
||||
p, _ := NewMarkdownParser(GoMarkdown)
|
||||
// 1×1 pixel transparent PNG encoded as data URI
|
||||
pixelPNG := make([]byte, 68) // minimal 1x1 PNG header
|
||||
pixelB64 := base64.StdEncoding.EncodeToString([]byte("fake-png-data"))
|
||||
md := "Some text with an image\n\n"
|
||||
_ = pixelPNG
|
||||
res := p.ParseWithResult("test.md", []byte(md))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
// Find the image item
|
||||
found := false
|
||||
for _, item := range res.JSON {
|
||||
if kd, _ := item["doc_type_kwd"].(string); kd == "image" {
|
||||
found = true
|
||||
if img, ok := item["image"].(string); !ok || img != pixelB64 {
|
||||
t.Fatalf("image data = %q, want %q", img, pixelB64)
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("no item with doc_type_kwd == 'image' found")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkdownParser_ParseWithResult_NoImage(t *testing.T) {
|
||||
p, _ := NewMarkdownParser(GoMarkdown)
|
||||
md := "# Title\n\nJust some text, no images here.\n\nMore text."
|
||||
res := p.ParseWithResult("test.md", []byte(md))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
for _, item := range res.JSON {
|
||||
if kd, _ := item["doc_type_kwd"].(string); kd == "image" {
|
||||
t.Fatal("unexpected image item in text-only markdown")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkdownParser_ConfigureFromSetup(t *testing.T) {
|
||||
p, _ := NewMarkdownParser(GoMarkdown)
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "deepdoc",
|
||||
"output_format": "json",
|
||||
"vlm": map[string]any{"llm_id": "gpt-4-vision"},
|
||||
"flatten_media_to_text": false,
|
||||
})
|
||||
if p.ParseMethod != "deepdoc" {
|
||||
t.Fatalf("ParseMethod = %q, want %q", p.ParseMethod, "deepdoc")
|
||||
}
|
||||
if p.OutputFormat != "json" {
|
||||
t.Fatalf("OutputFormat = %q, want %q", p.OutputFormat, "json")
|
||||
}
|
||||
if p.VLM == nil {
|
||||
t.Fatal("VLM is nil; want map")
|
||||
}
|
||||
if id, _ := p.VLM["llm_id"].(string); id != "gpt-4-vision" {
|
||||
t.Fatalf("VLM[llm_id] = %q, want %q", id, "gpt-4-vision")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkdownParser_ConfigureFromSetup_NilSafe(t *testing.T) {
|
||||
p, _ := NewMarkdownParser(GoMarkdown)
|
||||
p.ConfigureFromSetup(nil) // should not panic
|
||||
if p.ParseMethod != "" {
|
||||
t.Fatalf("ParseMethod should be empty after nil setup, got %q", p.ParseMethod)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMarkdownImage_DataURI(t *testing.T) {
|
||||
b64 := base64.StdEncoding.EncodeToString([]byte("fakeimage"))
|
||||
md := ""
|
||||
result, found := resolveMarkdownImage("", md)
|
||||
if !found {
|
||||
t.Fatal("expected image found for data URI")
|
||||
}
|
||||
if result != b64 {
|
||||
t.Fatalf("got %q, want %q", result, b64)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMarkdownImage_NoImage(t *testing.T) {
|
||||
_, found := resolveMarkdownImage("", "# Hello\nNo image here")
|
||||
if found {
|
||||
t.Fatal("expected no image found")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMarkdownImage_HTTPImage(t *testing.T) {
|
||||
// httptest servers bind loopback, which the SSRF guard rejects by
|
||||
// default. Allow loopback for this test so the HTTP fetch path is
|
||||
// exercised (production keeps ssrfAllowLoopback == false).
|
||||
prev := ssrfAllowLoopback
|
||||
ssrfAllowLoopback = true
|
||||
defer func() { ssrfAllowLoopback = prev }()
|
||||
|
||||
// Start a test HTTP server serving a fake PNG
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
w.Write([]byte("fake-png-bytes"))
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
md := ""
|
||||
result, found := resolveMarkdownImage("", md)
|
||||
if !found {
|
||||
t.Fatal("expected image found for HTTP URL")
|
||||
}
|
||||
expectedB64 := base64.StdEncoding.EncodeToString([]byte("fake-png-bytes"))
|
||||
if result != expectedB64 {
|
||||
t.Fatalf("got %q, want %q", result, expectedB64)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchImageAsBase64_RejectsCredentials(t *testing.T) {
|
||||
_, err := fetchImageAsBase64("https://user:pass@example.com/img.png")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for URL with credentials")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchImageAsBase64_InvalidURL(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
_, err := fetchImageAsBase64(ts.URL + "/nonexistent.png")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for 404 response")
|
||||
}
|
||||
}
|
||||
@@ -8,7 +8,7 @@ import (
|
||||
)
|
||||
|
||||
// ErrOfficeCGORequired is returned by ParseWithResult on every
|
||||
// office-parser family (DOC / DOCX / PPT / PPTX / XLS / XLSX)
|
||||
// office-parser family (DOC / DOCX / PPT / PPTX)
|
||||
// when the build is not CGO-enabled. The CGO build's
|
||||
// implementation captures the office_oxide PlainText / ToMarkdown
|
||||
// output; this stub mirrors that surface so the package compiles
|
||||
@@ -31,20 +31,6 @@ func (p *DOCParser) ParseWithResult(filename string, _ []byte) ParseResult {
|
||||
}
|
||||
}
|
||||
|
||||
func (p *XLSParser) ParseWithResult(filename string, _ []byte) ParseResult {
|
||||
return ParseResult{
|
||||
File: map[string]any{"name": filename},
|
||||
Err: fmt.Errorf("%w: xls", ErrOfficeCGORequired),
|
||||
}
|
||||
}
|
||||
|
||||
func (p *XLSXParser) ParseWithResult(filename string, _ []byte) ParseResult {
|
||||
return ParseResult{
|
||||
File: map[string]any{"name": filename},
|
||||
Err: fmt.Errorf("%w: xlsx", ErrOfficeCGORequired),
|
||||
}
|
||||
}
|
||||
|
||||
func (p *PPTParser) ParseWithResult(filename string, _ []byte) ParseResult {
|
||||
return ParseResult{
|
||||
File: map[string]any{"name": filename},
|
||||
@@ -79,26 +65,6 @@ func (p *DOCXParser) String() string {
|
||||
return "DOCXParser(no-cgo)"
|
||||
}
|
||||
|
||||
type XLSParser struct{}
|
||||
|
||||
func NewXLSParser() *XLSParser {
|
||||
return &XLSParser{}
|
||||
}
|
||||
|
||||
func (p *XLSParser) String() string {
|
||||
return "XLSParser(no-cgo)"
|
||||
}
|
||||
|
||||
type XLSXParser struct{}
|
||||
|
||||
func NewXLSXParser() *XLSXParser {
|
||||
return &XLSXParser{}
|
||||
}
|
||||
|
||||
func (p *XLSXParser) String() string {
|
||||
return "XLSXParser(no-cgo)"
|
||||
}
|
||||
|
||||
type PPTParser struct{}
|
||||
|
||||
func NewPPTParser() *PPTParser {
|
||||
|
||||
@@ -14,8 +14,6 @@ func TestOfficeParsers_ParseWithResult_NoCGO(t *testing.T) {
|
||||
}{
|
||||
{name: "docx", res: (&DOCXParser{}).ParseWithResult("a.docx", nil)},
|
||||
{name: "doc", res: (&DOCParser{}).ParseWithResult("a.doc", nil)},
|
||||
{name: "xlsx", res: (&XLSXParser{}).ParseWithResult("a.xlsx", nil)},
|
||||
{name: "xls", res: (&XLSParser{}).ParseWithResult("a.xls", nil)},
|
||||
{name: "pptx", res: (&PPTXParser{}).ParseWithResult("a.pptx", nil)},
|
||||
{name: "ppt", res: (&PPTParser{}).ParseWithResult("a.ppt", nil)},
|
||||
}
|
||||
|
||||
12
internal/parser/parser/office_parsers_shared.go
Normal file
12
internal/parser/parser/office_parsers_shared.go
Normal file
@@ -0,0 +1,12 @@
|
||||
//go:build !cgo
|
||||
|
||||
package parser
|
||||
|
||||
import "html"
|
||||
|
||||
// OfficeOxide is the lib_type identifier for office_oxide backend.
|
||||
const OfficeOxide = "office_oxide"
|
||||
|
||||
func htmlEscape(s string) string {
|
||||
return html.EscapeString(s)
|
||||
}
|
||||
433
internal/parser/parser/paddleocr.go
Normal file
433
internal/parser/parser/paddleocr.go
Normal file
@@ -0,0 +1,433 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// PaddleOCR client for the async Job API.
|
||||
// Mirrors Python's deepdoc/parser/paddleocr_parser.py:
|
||||
// - Submit job via POST /api/v2/ocr/jobs
|
||||
// - Poll with exponential backoff until "done" or "failed"
|
||||
// - Fetch result JSONL and extract text
|
||||
//
|
||||
// For image parsing, only parse_image() semantics are needed;
|
||||
// parse_pdf() is handled by the PDF vision dispatch path.
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/common"
|
||||
)
|
||||
|
||||
// imgTagPattern matches HTML <img> tags and <div> wrappers, mirroring
|
||||
// Python's _MARKDOWN_IMAGE_PATTERN in deepdoc/parser/paddleocr_parser.py.
|
||||
var imgTagPattern = regexp.MustCompile(
|
||||
`(?is)<div[^>]*>\s*<img[^>]*/>\s*</div>|<img[^>]*/>`,
|
||||
)
|
||||
|
||||
const (
|
||||
defaultPaddleOCRBaseURL = "https://paddleocr.aistudio-app.com"
|
||||
defaultPaddleOCRTimeout = 600 * time.Second
|
||||
defaultPaddleOCRAlgo = "PaddleOCR-VL"
|
||||
paddleOCRSubmitPath = "/api/v2/ocr/jobs"
|
||||
paddleOCRPollInterval = 3 * time.Second
|
||||
paddleOCRPollMultiplier = 1.5
|
||||
paddleOCRPollMaxInterval = 15 * time.Second
|
||||
)
|
||||
|
||||
// PaddleOCRClient talks to the PaddleOCR async Job API.
|
||||
type PaddleOCRClient struct {
|
||||
BaseURL string
|
||||
AccessToken string
|
||||
Algorithm string
|
||||
Timeout time.Duration
|
||||
httpClient *http.Client
|
||||
}
|
||||
|
||||
// NewPaddleOCRClientFromEnv creates a client from environment variables:
|
||||
//
|
||||
// PADDLEOCR_BASE_URL – base URL (default https://paddleocr.aistudio-app.com)
|
||||
// PADDLEOCR_ACCESS_TOKEN – bearer token
|
||||
// PADDLEOCR_ALGORITHM – algorithm name (default PaddleOCR-VL)
|
||||
func NewPaddleOCRClientFromEnv() *PaddleOCRClient {
|
||||
baseURL := common.GetEnv(common.EnvPaddleOCRBaseUrl)
|
||||
if baseURL == "" {
|
||||
baseURL = defaultPaddleOCRBaseURL
|
||||
}
|
||||
return &PaddleOCRClient{
|
||||
BaseURL: strings.TrimRight(baseURL, "/"),
|
||||
AccessToken: common.GetEnv(common.EnvPaddleOCRAccessToken),
|
||||
Algorithm: firstNonEmpty(common.GetEnv(common.EnvPaddleOCRAlgorithm), defaultPaddleOCRAlgo),
|
||||
Timeout: defaultPaddleOCRTimeout,
|
||||
httpClient: &http.Client{
|
||||
Timeout: 30 * time.Second, // per-request timeout; polling uses deadline
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Enabled reports whether the client has a usable access token.
|
||||
func (c *PaddleOCRClient) Enabled() bool {
|
||||
return c != nil && c.AccessToken != ""
|
||||
}
|
||||
|
||||
// ParseImage submits the image to PaddleOCR and returns extracted text.
|
||||
// Mirrors Python's PaddleOCRParser.parse_image().
|
||||
func (c *PaddleOCRClient) ParseImage(binary []byte, filename string) (string, error) {
|
||||
deadline := time.Now().Add(c.Timeout)
|
||||
|
||||
// Step 1: Submit job
|
||||
jobID, err := c.submitJob(binary, filename, deadline)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("paddleocr submit: %w", err)
|
||||
}
|
||||
|
||||
// Step 2: Poll until done
|
||||
resultData, err := c.pollJob(jobID, deadline)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("paddleocr poll: %w", err)
|
||||
}
|
||||
|
||||
// Step 3: Fetch result JSONL
|
||||
raw, err := c.fetchResult(resultData, deadline)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("paddleocr fetch: %w", err)
|
||||
}
|
||||
|
||||
// Step 4: Parse result
|
||||
return c.extractImageText(raw), nil
|
||||
}
|
||||
|
||||
// submitJob POSTs the image file to /api/v2/ocr/jobs and returns the job ID.
|
||||
func (c *PaddleOCRClient) submitJob(binary []byte, filename string, deadline time.Time) (string, error) {
|
||||
var buf bytes.Buffer
|
||||
w := multipart.NewWriter(&buf)
|
||||
|
||||
// Write the image file field
|
||||
fw, err := w.CreateFormFile("file", filename)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := fw.Write(binary); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Write form fields
|
||||
if err := w.WriteField("model", c.Algorithm); err != nil {
|
||||
return "", err
|
||||
}
|
||||
optionalPayload := map[string]any{"formatBlockContent": true}
|
||||
payloadBytes, _ := json.Marshal(optionalPayload)
|
||||
if err := w.WriteField("optionalPayload", string(payloadBytes)); err != nil {
|
||||
return "", err
|
||||
}
|
||||
w.Close()
|
||||
|
||||
url := c.BaseURL + paddleOCRSubmitPath
|
||||
req, err := http.NewRequest(http.MethodPost, url, &buf)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req.Header.Set("Content-Type", w.FormDataContentType())
|
||||
req.Header.Set("Client-Platform", "ragflow")
|
||||
if c.AccessToken != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+c.AccessToken)
|
||||
}
|
||||
req = req.WithContext(withDeadline(req.Context(), deadline))
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("POST %s: %w", url, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return "", fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var submitResp struct {
|
||||
Data struct {
|
||||
JobID string `json:"jobId"`
|
||||
} `json:"data"`
|
||||
JobID string `json:"jobId"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&submitResp); err != nil {
|
||||
return "", fmt.Errorf("decode submit response: %w", err)
|
||||
}
|
||||
jobID := submitResp.Data.JobID
|
||||
if jobID == "" {
|
||||
jobID = submitResp.JobID
|
||||
}
|
||||
if jobID == "" {
|
||||
return "", fmt.Errorf("no jobId in submit response")
|
||||
}
|
||||
slog.Info("paddleocr: job submitted", "jobId", jobID)
|
||||
return jobID, nil
|
||||
}
|
||||
|
||||
// pollJob polls the job status until it reaches "done" or "failed".
|
||||
// Uses exponential backoff: 3s initial, 1.5x multiplier, 15s max.
|
||||
// Returns the final poll response data.
|
||||
func (c *PaddleOCRClient) pollJob(jobID string, deadline time.Time) (map[string]any, error) {
|
||||
pollURL := fmt.Sprintf("%s/%s", c.BaseURL+paddleOCRSubmitPath, jobID)
|
||||
interval := paddleOCRPollInterval
|
||||
|
||||
for {
|
||||
if time.Now().After(deadline) {
|
||||
return nil, fmt.Errorf("job %s timed out after %v", jobID, c.Timeout)
|
||||
}
|
||||
remaining := time.Until(deadline)
|
||||
if remaining <= 0 {
|
||||
return nil, fmt.Errorf("job %s timed out", jobID)
|
||||
}
|
||||
|
||||
req, err := http.NewRequest(http.MethodGet, pollURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if c.AccessToken != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+c.AccessToken)
|
||||
}
|
||||
req.Header.Set("Client-Platform", "ragflow")
|
||||
req = req.WithContext(withDeadline(req.Context(), deadline))
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("poll %s: %w", jobID, err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return nil, fmt.Errorf("poll HTTP %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var pollResp struct {
|
||||
Data struct {
|
||||
State string `json:"state"`
|
||||
ErrorMsg string `json:"errorMsg"`
|
||||
} `json:"data"`
|
||||
State string `json:"state"`
|
||||
ErrorMsg string `json:"errorMsg"`
|
||||
}
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
if err := json.Unmarshal(bodyBytes, &pollResp); err != nil {
|
||||
return nil, fmt.Errorf("decode poll response: %w", err)
|
||||
}
|
||||
state := pollResp.Data.State
|
||||
if state == "" {
|
||||
state = pollResp.State
|
||||
}
|
||||
|
||||
switch state {
|
||||
case "done":
|
||||
slog.Info("paddleocr: job done", "jobId", jobID)
|
||||
// Return full data so caller can extract resultJsonUrl
|
||||
var fullData map[string]any
|
||||
json.Unmarshal(bodyBytes, &fullData)
|
||||
return fullData, nil
|
||||
case "failed":
|
||||
errMsg := pollResp.Data.ErrorMsg
|
||||
if errMsg == "" {
|
||||
errMsg = pollResp.ErrorMsg
|
||||
}
|
||||
return nil, fmt.Errorf("job %s failed: %s", jobID, errMsg)
|
||||
}
|
||||
|
||||
// Exponential backoff
|
||||
sleepTime := interval
|
||||
if remaining < sleepTime {
|
||||
sleepTime = remaining
|
||||
}
|
||||
time.Sleep(sleepTime)
|
||||
interval = time.Duration(float64(interval) * paddleOCRPollMultiplier)
|
||||
if interval > paddleOCRPollMaxInterval {
|
||||
interval = paddleOCRPollMaxInterval
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// fetchResult downloads and parses the JSONL result from resultJsonUrl.
|
||||
func (c *PaddleOCRClient) fetchResult(pollData map[string]any, deadline time.Time) ([]map[string]any, error) {
|
||||
data, _ := pollData["data"].(map[string]any)
|
||||
if data == nil {
|
||||
data = pollData
|
||||
}
|
||||
resultJSONURL, _ := data["resultJsonUrl"].(string)
|
||||
if resultJSONURL == "" {
|
||||
if resultURL, ok := data["resultUrl"].(map[string]any); ok {
|
||||
resultJSONURL, _ = resultURL["jsonUrl"].(string)
|
||||
}
|
||||
}
|
||||
if resultJSONURL == "" {
|
||||
return nil, fmt.Errorf("no resultJsonUrl in completion response")
|
||||
}
|
||||
|
||||
req, err := http.NewRequest(http.MethodGet, resultJSONURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if c.AccessToken != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+c.AccessToken)
|
||||
}
|
||||
req = req.WithContext(withDeadline(req.Context(), deadline))
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("fetch result: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != 200 {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return nil, fmt.Errorf("fetch result HTTP %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
content, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read result: %w", err)
|
||||
}
|
||||
|
||||
// Parse JSONL: one JSON object per line
|
||||
var lines []map[string]any
|
||||
for _, line := range strings.Split(strings.TrimSpace(string(content)), "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
var obj map[string]any
|
||||
if err := json.Unmarshal([]byte(line), &obj); err != nil {
|
||||
return nil, fmt.Errorf("parse JSONL line: %w", err)
|
||||
}
|
||||
lines = append(lines, obj)
|
||||
}
|
||||
return lines, nil
|
||||
}
|
||||
|
||||
// extractImageText extracts text from PaddleOCR JSONL results.
|
||||
// Mirrors Python's parse_image() text extraction:
|
||||
//
|
||||
// layoutParsingResults[].prunedResult.parsing_res_list[].block_content
|
||||
// Fallback: ocrResults[].prunedResult.rec_texts[]
|
||||
func (c *PaddleOCRClient) extractImageText(resultLines []map[string]any) string {
|
||||
var texts []string
|
||||
|
||||
for _, line := range resultLines {
|
||||
result, _ := line["result"].(map[string]any)
|
||||
if result == nil {
|
||||
continue
|
||||
}
|
||||
// Primary: layoutParsingResults
|
||||
if lpr, ok := result["layoutParsingResults"].([]any); ok {
|
||||
for _, lr := range lpr {
|
||||
lrMap, _ := lr.(map[string]any)
|
||||
if lrMap == nil {
|
||||
continue
|
||||
}
|
||||
pruned, _ := lrMap["prunedResult"].(map[string]any)
|
||||
if pruned == nil {
|
||||
continue
|
||||
}
|
||||
if prl, ok := pruned["parsing_res_list"].([]any); ok {
|
||||
for _, block := range prl {
|
||||
blockMap, _ := block.(map[string]any)
|
||||
if blockMap == nil {
|
||||
continue
|
||||
}
|
||||
content, _ := blockMap["block_content"].(string)
|
||||
content = strings.TrimSpace(content)
|
||||
// Remove markdown image blocks
|
||||
content = removeMarkdownImages(content)
|
||||
if content != "" {
|
||||
texts = append(texts, content)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: ocrResults for text-only models (e.g. PP-OCRv6)
|
||||
if len(texts) == 0 {
|
||||
for _, line := range resultLines {
|
||||
result, _ := line["result"].(map[string]any)
|
||||
if result == nil {
|
||||
continue
|
||||
}
|
||||
if ocr, ok := result["ocrResults"].([]any); ok {
|
||||
for _, o := range ocr {
|
||||
ocrMap, _ := o.(map[string]any)
|
||||
if ocrMap == nil {
|
||||
continue
|
||||
}
|
||||
pruned, _ := ocrMap["prunedResult"].(map[string]any)
|
||||
if pruned == nil {
|
||||
continue
|
||||
}
|
||||
if recTexts, ok := pruned["rec_texts"].([]any); ok {
|
||||
for _, t := range recTexts {
|
||||
if s, ok := t.(string); ok {
|
||||
s = strings.TrimSpace(s)
|
||||
if s != "" {
|
||||
texts = append(texts, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(texts) == 0 {
|
||||
return ""
|
||||
}
|
||||
return strings.Join(texts, "\n")
|
||||
}
|
||||
|
||||
// removeMarkdownImages strips markdown image syntax () and HTML
|
||||
// img/div blocks, mirroring Python's _remove_images_from_markdown.
|
||||
func removeMarkdownImages(md string) string {
|
||||
// Strip HTML <img> and <div><img/></div> blocks
|
||||
md = imgTagPattern.ReplaceAllString(md, "")
|
||||
return strings.TrimSpace(md)
|
||||
}
|
||||
|
||||
// withDeadline applies a deadline to a context, respecting the existing
|
||||
// deadline if it is sooner.
|
||||
func withDeadline(ctx context.Context, deadline time.Time) context.Context {
|
||||
if d, ok := ctx.Deadline(); ok && d.Before(deadline) {
|
||||
return ctx
|
||||
}
|
||||
newCtx, _ := context.WithDeadline(ctx, deadline)
|
||||
return newCtx
|
||||
}
|
||||
|
||||
func firstNonEmpty(a, b string) string {
|
||||
if a != "" {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
@@ -130,7 +130,7 @@ func (s sentinelErr) Error() string { return string(s) }
|
||||
// tiny: 1 heading, 1 paragraph, 1 unordered list item, no nested
|
||||
// formatting.
|
||||
func TestMarkdownParser_ParseWithResult(t *testing.T) {
|
||||
p := NewMarkdownParser()
|
||||
p, _ := NewMarkdownParser(GoMarkdown)
|
||||
src := []byte("# Title\n\nFirst paragraph.\n\n- Item one\n")
|
||||
res := p.ParseWithResult("doc.md", src)
|
||||
if res.Err != nil {
|
||||
|
||||
@@ -35,9 +35,11 @@ func GetParser(fileType utility.FileType) (ParseResultProducer, error) {
|
||||
case utility.FileTypePPT:
|
||||
return NewPPTParser(), nil
|
||||
case utility.FileTypeXLSX:
|
||||
return NewXLSXParser(), nil
|
||||
return NewXLSXParser("")
|
||||
case utility.FileTypeXLS:
|
||||
return NewXLSParser(), nil
|
||||
return NewXLSParser("")
|
||||
case utility.FileTypeCSV:
|
||||
return NewCSVParser(), nil
|
||||
case utility.FileTypeDOCX:
|
||||
return NewDOCXParser(), nil
|
||||
case utility.FileTypeDOC:
|
||||
@@ -47,9 +49,19 @@ func GetParser(fileType utility.FileType) (ParseResultProducer, error) {
|
||||
case utility.FileTypeHTML:
|
||||
return NewHTMLParser(), nil
|
||||
case utility.FileTypeMarkdown:
|
||||
return NewMarkdownParser(), nil
|
||||
return NewMarkdownParser(GoMarkdown)
|
||||
case utility.FileTypeTXT:
|
||||
return NewTextParser(), nil
|
||||
case utility.FileTypeEPUB:
|
||||
return NewEPUBParser(), nil
|
||||
case utility.FileTypeJSON:
|
||||
return NewJSONParser(), nil
|
||||
case utility.FileTypeEMAIL:
|
||||
return NewEmailParser(), nil
|
||||
case utility.FileTypeVISUAL:
|
||||
return NewPictureParser(), nil
|
||||
case utility.FileTypeAURAL:
|
||||
return NewAudioParser(), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported file type: %s", fileType)
|
||||
}
|
||||
|
||||
@@ -97,7 +97,7 @@ func parseMinerUMarkdownResult(filename, markdown, outputFormat string, pageCoun
|
||||
fileMeta := pdfFileMeta(filename, pageCount)
|
||||
switch strings.ToLower(strings.TrimSpace(outputFormat)) {
|
||||
case "", "json":
|
||||
mp := NewMarkdownParser()
|
||||
mp, _ := NewMarkdownParser(GoMarkdown)
|
||||
res := mp.ParseWithResult(filename, []byte(markdown))
|
||||
if res.Err != nil {
|
||||
return res
|
||||
|
||||
@@ -2,11 +2,41 @@
|
||||
|
||||
package parser
|
||||
|
||||
import "fmt"
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/ledongthuc/pdf"
|
||||
)
|
||||
|
||||
func parsePDFWithPlainText(filename string, data []byte, parser *PDFParser) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
return ParseResult{Err: fmt.Errorf("%w: %s", ErrPDFEngineUnavailable, filename)}
|
||||
|
||||
reader := bytes.NewReader(data)
|
||||
pdfReader, err := pdf.NewReader(reader, int64(len(data)))
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: plain_text open: %w", err)}
|
||||
}
|
||||
|
||||
pageCount := pdfReader.NumPage()
|
||||
items := make([]map[string]any, 0, pageCount)
|
||||
for pageNum := 1; pageNum <= pageCount; pageNum++ {
|
||||
p := pdfReader.Page(pageNum)
|
||||
if p.V.IsNull() {
|
||||
continue
|
||||
}
|
||||
text, err := p.GetPlainText(nil)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: plain_text page %d: %w", pageNum, err)}
|
||||
}
|
||||
items = append(items, map[string]any{
|
||||
"text": strings.TrimRight(text, "\n"),
|
||||
"doc_type_kwd": "text",
|
||||
"page_number": pageNum,
|
||||
})
|
||||
}
|
||||
return pdfItemsToResult(filename, items, parser.OutputFormat, pageCount)
|
||||
}
|
||||
|
||||
147
internal/parser/parser/picture_parser.go
Normal file
147
internal/parser/parser/picture_parser.go
Normal file
@@ -0,0 +1,147 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
// PictureParser validates and stores configuration for image files.
|
||||
// The actual OCR and VLM description is performed by the
|
||||
// component-layer maybeDispatchImage, which mirrors Python's
|
||||
// rag/app/picture.py:chunk() image branch.
|
||||
//
|
||||
// Python reference:
|
||||
// - Image is opened with PIL, converted to RGB
|
||||
// - PaddleOCR tried first (if layout_recognize == "@PaddleOCR")
|
||||
// - Falls back to local deepdoc.vision.OCR (ONNX text detection)
|
||||
// - If OCR text is short (≤32 chars / ≤32 words for English),
|
||||
// calls IMAGE2TEXT VLM describe() for a natural-language description
|
||||
// - Returns tokenized text with media context attached
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// imageExtensions mirrors Python's visual extensions used in
|
||||
// utility.FilenameType and the picture.py:chunk() image branch.
|
||||
// Video extensions (.mp4, .mov, ...) are excluded — those are
|
||||
// handled by maybeDispatchVideo.
|
||||
var imageExtensions = map[string]bool{
|
||||
"jpg": true, "jpeg": true, "png": true, "gif": true,
|
||||
"bmp": true, "tiff": true, "tif": true, "webp": true,
|
||||
"svg": true, "ico": true, "avif": true, "heic": true,
|
||||
"apng": true, "icon": true, "pcx": true, "tga": true,
|
||||
"exif": true, "fpx": true, "psd": true, "cdr": true,
|
||||
"pcd": true, "dxf": true, "ufo": true, "eps": true,
|
||||
"ai": true, "raw": true, "wmf": true,
|
||||
}
|
||||
|
||||
// PictureParser handles image files for OCR and VLM description.
|
||||
// Mirrors the configuration from setups["picture"]:
|
||||
// - vlm.llm_id → VLMModelID (IMAGE2TEXT model for describe)
|
||||
// - output_format
|
||||
// - image_context_size
|
||||
// - layout_recognize → ("@PaddleOCR:model_name" or empty)
|
||||
type PictureParser struct {
|
||||
VLMModelID string // vlm.llm_id — the IMAGE2TEXT model
|
||||
OutputFormat string
|
||||
ImageContextSize int // default 0
|
||||
LayoutRecognize string // layout_recognize (e.g. "@PaddleOCR")
|
||||
VideoPrompt string // for video-in-image detection (Python fallback)
|
||||
}
|
||||
|
||||
// NewPictureParser constructs a PictureParser.
|
||||
func NewPictureParser() *PictureParser {
|
||||
return &PictureParser{}
|
||||
}
|
||||
|
||||
// ConfigureFromSetup reads picture-specific configuration from the
|
||||
// parser setup map. Extracts vlm.llm_id, output_format,
|
||||
// image_context_size, layout_recognize, and video_prompt.
|
||||
func (p *PictureParser) ConfigureFromSetup(setup map[string]any) {
|
||||
if p == nil || setup == nil {
|
||||
return
|
||||
}
|
||||
if vlm, ok := setup["vlm"].(map[string]any); ok {
|
||||
if llmID, ok := vlm["llm_id"].(string); ok && llmID != "" {
|
||||
p.VLMModelID = llmID
|
||||
}
|
||||
}
|
||||
if v, ok := setup["output_format"].(string); ok && v != "" {
|
||||
p.OutputFormat = v
|
||||
}
|
||||
if v, ok := setup["image_context_size"].(float64); ok {
|
||||
p.ImageContextSize = int(v)
|
||||
}
|
||||
if v, ok := setup["layout_recognize"].(string); ok && v != "" {
|
||||
p.LayoutRecognize = v
|
||||
}
|
||||
if v, ok := setup["video_prompt"].(string); ok && v != "" {
|
||||
p.VideoPrompt = v
|
||||
}
|
||||
}
|
||||
|
||||
// ParseWithResult implements ParseResultProducer. It validates the
|
||||
// file extension against the image extension whitelist. The actual
|
||||
// OCR and VLM description happens via maybeDispatchImage at the
|
||||
// component layer (mirrors Python's picture.py:chunk()).
|
||||
func (p *PictureParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
ext := strings.ToLower(filepath.Ext(filename))
|
||||
if len(ext) > 1 && ext[0] == '.' {
|
||||
ext = ext[1:]
|
||||
}
|
||||
|
||||
// Video extensions: defer to maybeDispatchVideo. The picture.py
|
||||
// Python path handles both, but Go routes video separately.
|
||||
if ext != "" && isVideoExtension(ext) {
|
||||
return ParseResult{
|
||||
Err: fmt.Errorf("picture: video file %q should be routed through video parser, not picture", filename),
|
||||
}
|
||||
}
|
||||
|
||||
if ext == "" || !imageExtensions[ext] {
|
||||
return ParseResult{
|
||||
Err: fmt.Errorf("picture: unsupported extension %q (filename: %s); accepted: .jpg/.jpeg/.png/.gif/.bmp/.tiff/.tif/.webp/.svg/.ico/.avif/.heic/...", ext, filename),
|
||||
}
|
||||
}
|
||||
|
||||
// OutputFormat, VLMModelID, ImageContextSize, and LayoutRecognize
|
||||
// are consumed by maybeDispatchImage at the component layer.
|
||||
outFmt := p.OutputFormat
|
||||
if outFmt == "" {
|
||||
outFmt = "text"
|
||||
}
|
||||
|
||||
return ParseResult{
|
||||
OutputFormat: outFmt,
|
||||
File: map[string]any{
|
||||
"name": filename,
|
||||
"size": len(data),
|
||||
"doc_type_kwd": "image",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// isVideoExtension returns true when the extension is a video format
|
||||
// that Python's picture.py routes through the video branch.
|
||||
func isVideoExtension(ext string) bool {
|
||||
switch ext {
|
||||
case "mp4", "mov", "avi", "flv", "mpeg", "mpg",
|
||||
"webm", "wmv", "3gp", "3gpp", "mkv":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
145
internal/parser/parser/picture_parser_test.go
Normal file
145
internal/parser/parser/picture_parser_test.go
Normal file
@@ -0,0 +1,145 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package parser
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNewPictureParser(t *testing.T) {
|
||||
p := NewPictureParser()
|
||||
if p == nil {
|
||||
t.Fatal("NewPictureParser returned nil")
|
||||
}
|
||||
if p.OutputFormat != "" {
|
||||
t.Errorf("OutputFormat = %q, want empty", p.OutputFormat)
|
||||
}
|
||||
if p.VLMModelID != "" {
|
||||
t.Errorf("VLMModelID = %q, want empty", p.VLMModelID)
|
||||
}
|
||||
if p.ImageContextSize != 0 {
|
||||
t.Errorf("ImageContextSize = %d, want 0", p.ImageContextSize)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPictureParser_ConfigureFromSetup(t *testing.T) {
|
||||
p := NewPictureParser()
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"vlm": map[string]any{
|
||||
"llm_id": "gpt-4-vision",
|
||||
},
|
||||
"output_format": "text",
|
||||
"image_context_size": float64(3),
|
||||
"layout_recognize": "@PaddleOCR",
|
||||
"video_prompt": "summarize",
|
||||
})
|
||||
if p.VLMModelID != "gpt-4-vision" {
|
||||
t.Errorf("VLMModelID = %q, want gpt-4-vision", p.VLMModelID)
|
||||
}
|
||||
if p.OutputFormat != "text" {
|
||||
t.Errorf("OutputFormat = %q, want text", p.OutputFormat)
|
||||
}
|
||||
if p.ImageContextSize != 3 {
|
||||
t.Errorf("ImageContextSize = %d, want 3", p.ImageContextSize)
|
||||
}
|
||||
if p.LayoutRecognize != "@PaddleOCR" {
|
||||
t.Errorf("LayoutRecognize = %q, want @PaddleOCR", p.LayoutRecognize)
|
||||
}
|
||||
if p.VideoPrompt != "summarize" {
|
||||
t.Errorf("VideoPrompt = %q, want summarize", p.VideoPrompt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPictureParser_ParseWithResult_NilSetup(t *testing.T) {
|
||||
p := NewPictureParser()
|
||||
p.ConfigureFromSetup(nil)
|
||||
res := p.ParseWithResult("photo.png", []byte("\x89PNG"))
|
||||
if res.Err != nil {
|
||||
t.Errorf("unexpected error for nil setup: %v", res.Err)
|
||||
}
|
||||
if res.OutputFormat != "text" {
|
||||
t.Errorf("OutputFormat = %q, want text", res.OutputFormat)
|
||||
}
|
||||
file, ok := res.File["doc_type_kwd"].(string)
|
||||
if !ok || file != "image" {
|
||||
t.Errorf("doc_type_kwd = %q, want image", file)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPictureParser_ParseWithResult_ValidExtensions(t *testing.T) {
|
||||
exts := []string{"png", "jpg", "jpeg", "gif", "bmp", "tiff", "tif", "webp", "svg", "ico", "avif", "heic", "apng"}
|
||||
p := NewPictureParser()
|
||||
for _, ext := range exts {
|
||||
fn := "img." + ext
|
||||
res := p.ParseWithResult(fn, []byte{1, 2, 3})
|
||||
if res.Err != nil {
|
||||
t.Errorf("unexpected error for .%s: %v", ext, res.Err)
|
||||
}
|
||||
if res.OutputFormat != "text" {
|
||||
t.Errorf("OutputFormat = %q for .%s, want text", res.OutputFormat, ext)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPictureParser_ParseWithResult_InvalidExtension(t *testing.T) {
|
||||
p := NewPictureParser()
|
||||
res := p.ParseWithResult("sound.mp3", []byte{1})
|
||||
if res.Err == nil {
|
||||
t.Error("expected error for .mp3, got nil")
|
||||
}
|
||||
if !strings.Contains(res.Err.Error(), "mp3") {
|
||||
t.Errorf("error should mention unsupported extension, got: %v", res.Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPictureParser_ParseWithResult_VideoExtension(t *testing.T) {
|
||||
p := NewPictureParser()
|
||||
res := p.ParseWithResult("video.mp4", []byte{1})
|
||||
if res.Err == nil {
|
||||
t.Error("expected error for video .mp4, got nil")
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(res.Err.Error()), "video") {
|
||||
t.Errorf("error should mention video, got: %v", res.Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPictureParser_ParseWithResult_OutputFormat(t *testing.T) {
|
||||
p := NewPictureParser()
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"output_format": "json",
|
||||
})
|
||||
res := p.ParseWithResult("photo.jpg", []byte{1, 2, 3})
|
||||
if res.Err != nil {
|
||||
t.Fatalf("unexpected error: %v", res.Err)
|
||||
}
|
||||
if res.OutputFormat != "json" {
|
||||
t.Errorf("OutputFormat = %q, want json", res.OutputFormat)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageExtensions(t *testing.T) {
|
||||
for ext := range imageExtensions {
|
||||
if ext == "" {
|
||||
t.Error("empty extension in imageExtensions set")
|
||||
}
|
||||
// All extensions should be lowercase (no uppercase keys).
|
||||
if strings.ToLower(ext) != ext {
|
||||
t.Errorf("non-lowercase extension key: %q", ext)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,3 @@
|
||||
//go:build cgo
|
||||
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
@@ -26,21 +24,73 @@ import (
|
||||
"github.com/xuri/excelize/v2"
|
||||
)
|
||||
|
||||
type XLSParser struct{}
|
||||
type XLSParser struct {
|
||||
libType string
|
||||
ParseMethod string
|
||||
OutputFormat string
|
||||
TCADPAPIServer string
|
||||
TCADPAPIKey string
|
||||
TCADPTableResultType string
|
||||
TCADPMarkdownImageResponseType string
|
||||
}
|
||||
|
||||
func NewXLSParser() *XLSParser {
|
||||
return &XLSParser{}
|
||||
func NewXLSParser(libType string) (*XLSParser, error) {
|
||||
if libType == "" {
|
||||
libType = "excelize"
|
||||
}
|
||||
return &XLSParser{
|
||||
libType: libType,
|
||||
TCADPTableResultType: "1",
|
||||
TCADPMarkdownImageResponseType: "1",
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (p *XLSParser) String() string {
|
||||
return "XLSParser"
|
||||
}
|
||||
|
||||
// ParseWithResult delegates to excelize which handles both .xls
|
||||
// and .xlsx through the same API. The python ExcelParser falls
|
||||
// back to a similar delegation; on the Go side excelize is the
|
||||
// single library for both extensions.
|
||||
func (p *XLSParser) ConfigureFromSetup(setup map[string]any) {
|
||||
if p == nil || setup == nil {
|
||||
return
|
||||
}
|
||||
if v, ok := setup["parse_method"].(string); ok && v != "" {
|
||||
p.ParseMethod = v
|
||||
}
|
||||
if v, ok := setup["output_format"].(string); ok && v != "" {
|
||||
p.OutputFormat = v
|
||||
}
|
||||
if v, ok := setup["tcadp_apiserver"].(string); ok && v != "" {
|
||||
p.TCADPAPIServer = v
|
||||
}
|
||||
if v, ok := setup["tcadp_api_key"].(string); ok {
|
||||
p.TCADPAPIKey = v
|
||||
}
|
||||
if v, ok := setup["table_result_type"].(string); ok && v != "" {
|
||||
p.TCADPTableResultType = v
|
||||
}
|
||||
if v, ok := setup["markdown_image_response_type"].(string); ok && v != "" {
|
||||
p.TCADPMarkdownImageResponseType = v
|
||||
}
|
||||
}
|
||||
|
||||
func (p *XLSParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
method := normalizeXLSXParseMethod(p.ParseMethod)
|
||||
switch method {
|
||||
case "tcadp":
|
||||
return parseSpreadsheetWithTCADP(
|
||||
filename, data, "XLS",
|
||||
p.TCADPAPIServer, p.TCADPAPIKey,
|
||||
p.TCADPTableResultType, p.TCADPMarkdownImageResponseType,
|
||||
p.OutputFormat,
|
||||
)
|
||||
case "", "excelize":
|
||||
// Continue with the local Excelize parser.
|
||||
default:
|
||||
return ParseResult{
|
||||
Err: fmt.Errorf("unsupported XLS parse method: %q", p.ParseMethod),
|
||||
}
|
||||
}
|
||||
|
||||
f, err := excelize.OpenReader(bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("xls open: %w", err)}
|
||||
|
||||
83
internal/parser/parser/xls_tcadp.go
Normal file
83
internal/parser/parser/xls_tcadp.go
Normal file
@@ -0,0 +1,83 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
models "ragflow/internal/entity/models"
|
||||
)
|
||||
|
||||
func parseSpreadsheetWithTCADP(filename string, data []byte, fileType string, tcadpAPIServer, tcadpAPIKey, tableResultType, markdownImageResponseType string, outputFormat string) ParseResult {
|
||||
if len(data) == 0 {
|
||||
return emptyPDFResult(filename)
|
||||
}
|
||||
baseURL := strings.TrimSpace(tcadpAPIServer)
|
||||
if baseURL == "" {
|
||||
baseURL = strings.TrimSpace(os.Getenv("TCADP_APISERVER"))
|
||||
}
|
||||
if baseURL == "" {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP requires tcadp_apiserver or TCADP_APISERVER")}
|
||||
}
|
||||
apiKey := strings.TrimSpace(tcadpAPIKey)
|
||||
if apiKey == "" {
|
||||
apiKey = strings.TrimSpace(os.Getenv("TCADP_API_KEY"))
|
||||
}
|
||||
requestBody := map[string]any{
|
||||
"file_type": fileType,
|
||||
"file_base64": base64.StdEncoding.EncodeToString(data),
|
||||
"file_start_page_number": 1,
|
||||
"file_end_page_number": 1000,
|
||||
"config": map[string]any{
|
||||
"TableResultType": tableResultType,
|
||||
"MarkdownImageResponseType": markdownImageResponseType,
|
||||
},
|
||||
}
|
||||
resp, err := models.PostJSONRequest(context.Background(), models.NewDriverHTTPClient(), strings.TrimRight(baseURL, "/")+"/reconstruct_document", bearer(apiKey), requestBody)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP submit: %w", err)}
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP read submit: %w", err)}
|
||||
}
|
||||
if resp.StatusCode >= 300 {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP HTTP %d: %s", resp.StatusCode, string(raw))}
|
||||
}
|
||||
var payload struct {
|
||||
DocumentRecognizeResultURL string `json:"DocumentRecognizeResultUrl"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &payload); err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP decode submit: %w", err)}
|
||||
}
|
||||
if payload.DocumentRecognizeResultURL == "" {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP returned no DocumentRecognizeResultUrl")}
|
||||
}
|
||||
downloadReq, err := http.NewRequestWithContext(context.Background(), http.MethodGet, payload.DocumentRecognizeResultURL, nil)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP download request: %w", err)}
|
||||
}
|
||||
if auth := bearer(apiKey); auth != "" {
|
||||
downloadReq.Header.Set("Authorization", auth)
|
||||
}
|
||||
downloadResp, err := models.NewDriverHTTPClient().Do(downloadReq)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP download: %w", err)}
|
||||
}
|
||||
defer downloadResp.Body.Close()
|
||||
zipBytes, err := io.ReadAll(downloadResp.Body)
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("parser: TCADP read zip: %w", err)}
|
||||
}
|
||||
items, pageCount, err := tcadpItemsFromZip(zipBytes)
|
||||
if err != nil {
|
||||
return ParseResult{Err: err}
|
||||
}
|
||||
return pdfItemsToResult(filename, items, outputFormat, pageCount)
|
||||
}
|
||||
241
internal/parser/parser/xls_tcadp_test.go
Normal file
241
internal/parser/parser/xls_tcadp_test.go
Normal file
@@ -0,0 +1,241 @@
|
||||
package parser
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestXLSXParser_ParseWithResult_TCADPJSONIntegration(t *testing.T) {
|
||||
zipPayload := tcadpZipFixture(t)
|
||||
var server *httptest.Server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/reconstruct_document":
|
||||
_, _ = w.Write([]byte(`{"DocumentRecognizeResultUrl":"` + server.URL + `/download.zip"}`))
|
||||
case "/download.zip":
|
||||
w.Header().Set("Content-Type", "application/zip")
|
||||
_, _ = w.Write(zipPayload)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
p, err := NewXLSXParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewXLSXParser: %v", err)
|
||||
}
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "TCADP parser",
|
||||
"output_format": "json",
|
||||
"tcadp_apiserver": server.URL,
|
||||
"tcadp_api_key": "tcadp-secret",
|
||||
"table_result_type": "1",
|
||||
"markdown_image_response_type": "1",
|
||||
})
|
||||
res := p.ParseWithResult("sample.xlsx", []byte("mock xlsx content"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) < 2 {
|
||||
t.Fatalf("JSON len = %d, want >=2", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
func TestXLSParser_ParseWithResult_TCADPJSONIntegration(t *testing.T) {
|
||||
zipPayload := tcadpZipFixture(t)
|
||||
var server *httptest.Server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/reconstruct_document":
|
||||
_, _ = w.Write([]byte(`{"DocumentRecognizeResultUrl":"` + server.URL + `/download.zip"}`))
|
||||
case "/download.zip":
|
||||
_, _ = w.Write(zipPayload)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
p, err := NewXLSParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewXLSParser: %v", err)
|
||||
}
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "TCADP parser",
|
||||
"output_format": "json",
|
||||
"tcadp_apiserver": server.URL,
|
||||
"tcadp_api_key": "tcadp-secret",
|
||||
})
|
||||
res := p.ParseWithResult("sample.xls", []byte("mock xls content"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) < 2 {
|
||||
t.Fatalf("JSON len = %d, want >=2", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
func TestCSVParser_ParseWithResult_TCADPJSONIntegration(t *testing.T) {
|
||||
zipPayload := tcadpZipFixture(t)
|
||||
var server *httptest.Server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/reconstruct_document":
|
||||
_, _ = w.Write([]byte(`{"DocumentRecognizeResultUrl":"` + server.URL + `/download.zip"}`))
|
||||
case "/download.zip":
|
||||
_, _ = w.Write(zipPayload)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
p := NewCSVParser()
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "TCADP parser",
|
||||
"output_format": "json",
|
||||
"tcadp_apiserver": server.URL,
|
||||
"tcadp_api_key": "tcadp-secret",
|
||||
})
|
||||
res := p.ParseWithResult("sample.csv", []byte("a,b,c\n1,2,3"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if len(res.JSON) < 2 {
|
||||
t.Fatalf("JSON len = %d, want >=2", len(res.JSON))
|
||||
}
|
||||
}
|
||||
|
||||
func TestXLSXParser_ParseWithResult_TCADPMarkdownIntegration(t *testing.T) {
|
||||
zipPayload := tcadpZipFixture(t)
|
||||
var server *httptest.Server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/reconstruct_document":
|
||||
_, _ = w.Write([]byte(`{"DocumentRecognizeResultUrl":"` + server.URL + `/download.zip"}`))
|
||||
case "/download.zip":
|
||||
_, _ = w.Write(zipPayload)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
p, err := NewXLSXParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewXLSXParser: %v", err)
|
||||
}
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "TCADP parser",
|
||||
"output_format": "markdown",
|
||||
"tcadp_apiserver": server.URL,
|
||||
})
|
||||
res := p.ParseWithResult("sample.xlsx", []byte("mock xlsx content"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if res.Markdown == "" {
|
||||
t.Fatal("Markdown is empty; want rendered content")
|
||||
}
|
||||
}
|
||||
|
||||
func TestXLSXParser_ParseWithResult_TCADPRequiresAPIServer(t *testing.T) {
|
||||
p, err := NewXLSXParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewXLSXParser: %v", err)
|
||||
}
|
||||
p.ConfigureFromSetup(map[string]any{"parse_method": "TCADP parser"})
|
||||
res := p.ParseWithResult("sample.xlsx", []byte("mock xlsx content"))
|
||||
if res.Err == nil {
|
||||
t.Fatal("expected error about tcadp_apiserver, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestXLSParser_ParseWithResult_TCADPRequiresAPIServer(t *testing.T) {
|
||||
p, err := NewXLSParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewXLSParser: %v", err)
|
||||
}
|
||||
p.ConfigureFromSetup(map[string]any{"parse_method": "TCADP parser"})
|
||||
res := p.ParseWithResult("sample.xls", []byte("mock xls content"))
|
||||
if res.Err == nil {
|
||||
t.Fatal("expected error about tcadp_apiserver, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCSVParser_ParseWithResult_TCADPRequiresAPIServer(t *testing.T) {
|
||||
p := NewCSVParser()
|
||||
p.ConfigureFromSetup(map[string]any{"parse_method": "TCADP parser"})
|
||||
res := p.ParseWithResult("sample.csv", []byte("mock csv content"))
|
||||
if res.Err == nil {
|
||||
t.Fatal("expected error about tcadp_apiserver, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestXLSXParser_ParseWithResult_InvalidXLSXHandled(t *testing.T) {
|
||||
p, err := NewXLSXParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewXLSXParser: %v", err)
|
||||
}
|
||||
res := p.ParseWithResult("sample.xlsx", []byte("not a valid xlsx"))
|
||||
if res.Err == nil {
|
||||
t.Fatal("expected error for invalid xlsx, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestXLSParser_ParseWithResult_InvalidXLSHandled(t *testing.T) {
|
||||
p, err := NewXLSParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewXLSParser: %v", err)
|
||||
}
|
||||
res := p.ParseWithResult("sample.xls", []byte("not a valid xls"))
|
||||
if res.Err == nil {
|
||||
t.Fatal("expected error for invalid xls, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCSVParser_ParseWithResult_DefaultCSVBehavior(t *testing.T) {
|
||||
p := NewCSVParser()
|
||||
res := p.ParseWithResult("sample.csv", []byte("a,b\n1,2"))
|
||||
if res.Err != nil {
|
||||
t.Fatalf("ParseWithResult: %v", res.Err)
|
||||
}
|
||||
if got, want := res.OutputFormat, "html"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if res.HTML == "" {
|
||||
t.Fatal("HTML is empty; want rendered table")
|
||||
}
|
||||
}
|
||||
|
||||
func TestXLSXParser_ConfigureFromSetup_TCADP(t *testing.T) {
|
||||
p, err := NewXLSXParser("")
|
||||
if err != nil {
|
||||
t.Fatalf("NewXLSXParser: %v", err)
|
||||
}
|
||||
p.ConfigureFromSetup(map[string]any{
|
||||
"parse_method": "TCADP parser",
|
||||
"output_format": "json",
|
||||
"tcadp_apiserver": "https://tcadp.example.com",
|
||||
"tcadp_api_key": "secret",
|
||||
"table_result_type": "2",
|
||||
"markdown_image_response_type": "2",
|
||||
})
|
||||
if got, want := p.ParseMethod, "TCADP parser"; got != want {
|
||||
t.Fatalf("ParseMethod = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := p.OutputFormat, "json"; got != want {
|
||||
t.Fatalf("OutputFormat = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := p.TCADPAPIServer, "https://tcadp.example.com"; got != want {
|
||||
t.Fatalf("TCADPAPIServer = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := p.TCADPTableResultType, "2"; got != want {
|
||||
t.Fatalf("TCADPTableResultType = %q, want %q", got, want)
|
||||
}
|
||||
if got, want := p.TCADPMarkdownImageResponseType, "2"; got != want {
|
||||
t.Fatalf("TCADPMarkdownImageResponseType = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,3 @@
|
||||
//go:build cgo
|
||||
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
@@ -26,23 +24,81 @@ import (
|
||||
"github.com/xuri/excelize/v2"
|
||||
)
|
||||
|
||||
type XLSXParser struct{}
|
||||
type XLSXParser struct {
|
||||
libType string
|
||||
ParseMethod string
|
||||
OutputFormat string
|
||||
TCADPAPIServer string
|
||||
TCADPAPIKey string
|
||||
TCADPTableResultType string
|
||||
TCADPMarkdownImageResponseType string
|
||||
}
|
||||
|
||||
func NewXLSXParser() *XLSXParser {
|
||||
return &XLSXParser{}
|
||||
func NewXLSXParser(libType string) (*XLSXParser, error) {
|
||||
if libType == "" {
|
||||
libType = "excelize"
|
||||
}
|
||||
return &XLSXParser{
|
||||
libType: libType,
|
||||
TCADPTableResultType: "1",
|
||||
TCADPMarkdownImageResponseType: "1",
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (p *XLSXParser) String() string {
|
||||
return "XLSXParser"
|
||||
}
|
||||
|
||||
// ParseWithResult renders the spreadsheet as HTML — the python
|
||||
// ExcelParser shape. Each sheet becomes a <table> with row /
|
||||
// column structure preserved; sheet names become <h3> headings
|
||||
// so a downstream title chunker can pick them up. Implementation
|
||||
// uses excelize (already in go.mod) instead of office_oxide's
|
||||
// PlainText/ToMarkdown so cell boundaries survive the round-trip.
|
||||
func (p *XLSXParser) ConfigureFromSetup(setup map[string]any) {
|
||||
if p == nil || setup == nil {
|
||||
return
|
||||
}
|
||||
if v, ok := setup["parse_method"].(string); ok && v != "" {
|
||||
p.ParseMethod = v
|
||||
}
|
||||
if v, ok := setup["output_format"].(string); ok && v != "" {
|
||||
p.OutputFormat = v
|
||||
}
|
||||
if v, ok := setup["tcadp_apiserver"].(string); ok && v != "" {
|
||||
p.TCADPAPIServer = v
|
||||
}
|
||||
if v, ok := setup["tcadp_api_key"].(string); ok {
|
||||
p.TCADPAPIKey = v
|
||||
}
|
||||
if v, ok := setup["table_result_type"].(string); ok && v != "" {
|
||||
p.TCADPTableResultType = v
|
||||
}
|
||||
if v, ok := setup["markdown_image_response_type"].(string); ok && v != "" {
|
||||
p.TCADPMarkdownImageResponseType = v
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeXLSXParseMethod(raw string) string {
|
||||
method := strings.ToLower(strings.TrimSpace(raw))
|
||||
if method == "tcadp parser" {
|
||||
return "tcadp"
|
||||
}
|
||||
return method
|
||||
}
|
||||
|
||||
func (p *XLSXParser) ParseWithResult(filename string, data []byte) ParseResult {
|
||||
method := normalizeXLSXParseMethod(p.ParseMethod)
|
||||
switch method {
|
||||
case "tcadp":
|
||||
return parseSpreadsheetWithTCADP(
|
||||
filename, data, "XLSX",
|
||||
p.TCADPAPIServer, p.TCADPAPIKey,
|
||||
p.TCADPTableResultType, p.TCADPMarkdownImageResponseType,
|
||||
p.OutputFormat,
|
||||
)
|
||||
case "", "excelize":
|
||||
// Continue with the local Excelize parser.
|
||||
default:
|
||||
return ParseResult{
|
||||
Err: fmt.Errorf("unsupported XLSX parse method: %q", p.ParseMethod),
|
||||
}
|
||||
}
|
||||
|
||||
f, err := excelize.OpenReader(bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return ParseResult{Err: fmt.Errorf("xlsx open: %w", err)}
|
||||
|
||||
Reference in New Issue
Block a user