feat(infra): integrate PaddleOCR's PP-StructureV3 as a document parser backend (#714)

This commit is contained in:
Lin Manhui
2025-08-13 16:37:42 +08:00
committed by GitHub
parent 708a6ed0c0
commit 6b60c07c22
30 changed files with 657 additions and 174 deletions

View File

@@ -34,10 +34,10 @@ var (
emailRegex = regexp.MustCompile(`[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}`)
)
func chunkCustom(_ context.Context, text string, config *contract.Config, opts ...parser.Option) (docs []*schema.Document, err error) {
func ChunkCustom(_ context.Context, text string, config *contract.Config, opts ...parser.Option) (docs []*schema.Document, err error) {
cs := config.ChunkingStrategy
if cs.Overlap >= cs.ChunkSize {
return nil, fmt.Errorf("[chunkCustom] invalid param, overlap >= chunk_size")
return nil, fmt.Errorf("[ChunkCustom] invalid param, overlap >= chunk_size")
}
var (

View File

@@ -39,7 +39,7 @@ func TestChunkCustom(t *testing.T) {
TrimURLAndEmail: true,
}
slices, err := chunkCustom(ctx, text, &parser.Config{ChunkingStrategy: cs})
slices, err := ChunkCustom(ctx, text, &parser.Config{ChunkingStrategy: cs})
assert.NoError(t, err)
assert.Len(t, slices, 10)

View File

@@ -24,7 +24,7 @@ import (
"github.com/coze-dev/coze-studio/backend/infra/contract/storage"
)
func putImageObject(ctx context.Context, st storage.Storage, imgExt string, uid int64, img []byte) (format string, err error) {
func PutImageObject(ctx context.Context, st storage.Storage, imgExt string, uid int64, img []byte) (format string, err error) {
secret := createSecret(uid, imgExt)
fileName := fmt.Sprintf("%d_%d_%s.%s", uid, time.Now().UnixNano(), secret, imgExt)
objectName := fmt.Sprintf("%s/%s", knowledgePrefix, fileName)

View File

@@ -41,7 +41,7 @@ type manager struct {
}
func (m *manager) GetParser(config *parser.Config) (parser.Parser, error) {
var pFn parseFn
var pFn ParseFn
if config.ParsingStrategy.HeaderLine == 0 && config.ParsingStrategy.DataStartLine == 0 {
config.ParsingStrategy.DataStartLine = 1
@@ -52,26 +52,30 @@ func (m *manager) GetParser(config *parser.Config) (parser.Parser, error) {
switch config.FileExtension {
case parser.FileExtensionPDF:
pFn = parseByPython(config, m.storage, m.ocr, goutil.GetPython3Path(), goutil.GetPythonFilePath("parse_pdf.py"))
pFn = ParseByPython(config, m.storage, m.ocr, goutil.GetPython3Path(), goutil.GetPythonFilePath("parse_pdf.py"))
case parser.FileExtensionTXT:
pFn = parseText(config)
pFn = ParseText(config)
case parser.FileExtensionMarkdown:
pFn = parseMarkdown(config, m.storage, m.ocr)
pFn = ParseMarkdown(config, m.storage, m.ocr)
case parser.FileExtensionDocx:
pFn = parseByPython(config, m.storage, m.ocr, goutil.GetPython3Path(), goutil.GetPythonFilePath("parse_docx.py"))
pFn = ParseByPython(config, m.storage, m.ocr, goutil.GetPython3Path(), goutil.GetPythonFilePath("parse_docx.py"))
case parser.FileExtensionCSV:
pFn = parseCSV(config)
pFn = ParseCSV(config)
case parser.FileExtensionXLSX:
pFn = parseXLSX(config)
pFn = ParseXLSX(config)
case parser.FileExtensionJSON:
pFn = parseJSON(config)
pFn = ParseJSON(config)
case parser.FileExtensionJsonMaps:
pFn = parseJSONMaps(config)
pFn = ParseJSONMaps(config)
case parser.FileExtensionJPG, parser.FileExtensionJPEG, parser.FileExtensionPNG:
pFn = parseImage(config, m.model)
pFn = ParseImage(config, m.model)
default:
return nil, fmt.Errorf("[Parse] document type not support, type=%s", config.FileExtension)
}
return &p{parseFn: pFn}, nil
return &Parser{ParseFn: pFn}, nil
}
func (m *manager) IsAutoAnnotationSupported() bool {
return m.model != nil
}

View File

@@ -29,7 +29,7 @@ import (
contract "github.com/coze-dev/coze-studio/backend/infra/contract/document/parser"
)
func parseCSV(config *contract.Config) parseFn {
func ParseCSV(config *contract.Config) ParseFn {
return func(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error) {
iter := &csvIterator{csv.NewReader(utfbom.SkipOnly(reader))}
return parseByRowIterator(iter, config, opts...)

View File

@@ -47,7 +47,7 @@ func TestParseCSV(t *testing.T) {
},
ChunkingStrategy: nil,
}
p1 := parseCSV(c1)
p1 := ParseCSV(c1)
docs, err := p1(ctx, r1, parser.WithExtraMeta(map[string]any{
"document_id": int64(123),
"knowledge_id": int64(456),
@@ -112,7 +112,7 @@ func TestParseCSV(t *testing.T) {
},
ChunkingStrategy: nil,
}
p2 := parseCSV(c2)
p2 := ParseCSV(c2)
docs, err = p2(ctx, r2, parser.WithExtraMeta(map[string]any{
"document_id": int64(123),
"knowledge_id": int64(456),
@@ -131,7 +131,7 @@ func TestParseCSVBadCases(t *testing.T) {
b, err := io.ReadAll(f)
assert.NoError(t, err)
pfn := parseCSV(&contract.Config{
pfn := ParseCSV(&contract.Config{
FileExtension: "csv",
ParsingStrategy: &contract.ParsingStrategy{
ExtractImage: true,
@@ -154,7 +154,7 @@ func TestParseCSVBadCases(t *testing.T) {
cols, err := document.GetDocumentColumns(resp[0])
assert.NoError(t, err)
cols[5].Nullable = false
npfn := parseCSV(&contract.Config{
npfn := ParseCSV(&contract.Config{
FileExtension: "csv",
ParsingStrategy: &contract.ParsingStrategy{
ExtractImage: true,

View File

@@ -31,7 +31,7 @@ import (
"github.com/coze-dev/coze-studio/backend/types/errno"
)
func parseImage(config *contract.Config, model chatmodel.BaseChatModel) parseFn {
func ParseImage(config *contract.Config, model chatmodel.BaseChatModel) ParseFn {
return func(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error) {
options := parser.GetCommonOptions(&parser.Options{}, opts...)
doc := &schema.Document{
@@ -76,14 +76,14 @@ func parseImage(config *contract.Config, model chatmodel.BaseChatModel) parseFn
output, err := model.Generate(ctx, []*schema.Message{input})
if err != nil {
return nil, fmt.Errorf("[parseImage] model generate failed: %w", err)
return nil, fmt.Errorf("[ParseImage] model generate failed: %w", err)
}
doc.Content = output.Content
case contract.ImageAnnotationTypeManual:
// do nothing
default:
return nil, fmt.Errorf("[parseImage] unknown image annotation type=%d", config.ParsingStrategy.ImageAnnotationType)
return nil, fmt.Errorf("[ParseImage] unknown image annotation type=%d", config.ParsingStrategy.ImageAnnotationType)
}
return []*schema.Document{doc}, nil

View File

@@ -28,7 +28,7 @@ import (
contract "github.com/coze-dev/coze-studio/backend/infra/contract/document/parser"
)
func parseJSON(config *contract.Config) parseFn {
func ParseJSON(config *contract.Config) ParseFn {
return func(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error) {
b, err := io.ReadAll(reader)
if err != nil {
@@ -41,7 +41,7 @@ func parseJSON(config *contract.Config) parseFn {
}
if len(rawSlices) == 0 {
return nil, fmt.Errorf("[parseJSON] json data is empty")
return nil, fmt.Errorf("[ParseJSON] json data is empty")
}
var header []string

View File

@@ -29,7 +29,7 @@ import (
contract "github.com/coze-dev/coze-studio/backend/infra/contract/document/parser"
)
func parseJSONMaps(config *contract.Config) parseFn {
func ParseJSONMaps(config *contract.Config) ParseFn {
return func(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error) {
b, err := io.ReadAll(reader)
if err != nil {

View File

@@ -85,7 +85,7 @@ func TestParseTableCustomContent(t *testing.T) {
},
}
pfn := parseJSONMaps(config)
pfn := ParseJSONMaps(config)
docs, err := pfn(ctx, reader, parser.WithExtraMeta(map[string]any{
"document_id": int64(123),
"knowledge_id": int64(456),

View File

@@ -55,7 +55,7 @@ func TestParseJSON(t *testing.T) {
},
ChunkingStrategy: nil,
}
pfn := parseJSON(config)
pfn := ParseJSON(config)
docs, err := pfn(context.Background(), reader, parser.WithExtraMeta(map[string]any{
"document_id": int64(123),
"knowledge_id": int64(456),
@@ -121,7 +121,7 @@ func TestParseJSONWithSchema(t *testing.T) {
},
},
}
pfn := parseJSON(config)
pfn := ParseJSON(config)
docs, err := pfn(context.Background(), reader, parser.WithExtraMeta(map[string]any{
"document_id": int64(123),
"knowledge_id": int64(456),

View File

@@ -38,7 +38,7 @@ import (
"github.com/coze-dev/coze-studio/backend/pkg/logs"
)
func parseMarkdown(config *contract.Config, storage storage.Storage, ocr ocr.OCR) parseFn {
func ParseMarkdown(config *contract.Config, storage storage.Storage, ocr ocr.OCR) ParseFn {
return func(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error) {
options := parser.GetCommonOptions(&parser.Options{}, opts...)
mdParser := goldmark.DefaultParser()
@@ -52,7 +52,7 @@ func parseMarkdown(config *contract.Config, storage storage.Storage, ocr ocr.OCR
ps := config.ParsingStrategy
if cs.ChunkType != contract.ChunkTypeCustom && cs.ChunkType != contract.ChunkTypeDefault {
return nil, fmt.Errorf("[parseMarkdown] chunk type not support, chunk type=%d", cs.ChunkType)
return nil, fmt.Errorf("[ParseMarkdown] chunk type not support, chunk type=%d", cs.ChunkType)
}
var (
@@ -173,7 +173,7 @@ func parseMarkdown(config *contract.Config, storage storage.Storage, ocr ocr.OCR
return ast.WalkStop, fmt.Errorf("failed to download image: %w", err)
}
imgSrc, err := putImageObject(ctx, storage, ext, getCreatorIDFromExtraMeta(options.ExtraMeta), img)
imgSrc, err := PutImageObject(ctx, storage, ext, GetCreatorIDFromExtraMeta(options.ExtraMeta), img)
if err != nil {
return ast.WalkStop, err
}
@@ -198,7 +198,7 @@ func parseMarkdown(config *contract.Config, storage storage.Storage, ocr ocr.OCR
pushSlice()
}
} else {
logs.CtxInfof(ctx, "[parseMarkdown] not a valid image url, skip, got=%s", imageURL)
logs.CtxInfof(ctx, "[ParseMarkdown] not a valid image url, skip, got=%s", imageURL)
}
}
}

View File

@@ -37,7 +37,7 @@ func TestParseMarkdown(t *testing.T) {
mockStorage := ms.NewMockStorage(ctrl)
mockStorage.EXPECT().PutObject(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
pfn := parseMarkdown(&contract.Config{
pfn := ParseMarkdown(&contract.Config{
FileExtension: contract.FileExtensionMarkdown,
ParsingStrategy: &contract.ParsingStrategy{
ExtractImage: true,

View File

@@ -27,7 +27,7 @@ import (
contract "github.com/coze-dev/coze-studio/backend/infra/contract/document/parser"
)
func parseText(config *contract.Config) parseFn {
func ParseText(config *contract.Config) ParseFn {
return func(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error) {
content, err := io.ReadAll(reader)
if err != nil {
@@ -36,9 +36,9 @@ func parseText(config *contract.Config) parseFn {
switch config.ChunkingStrategy.ChunkType {
case contract.ChunkTypeCustom, contract.ChunkTypeDefault:
docs, err = chunkCustom(ctx, string(content), config, opts...)
docs, err = ChunkCustom(ctx, string(content), config, opts...)
default:
return nil, fmt.Errorf("[parseText] chunk type not support, type=%d", config.ChunkingStrategy.ChunkType)
return nil, fmt.Errorf("[ParseText] chunk type not support, type=%d", config.ChunkingStrategy.ChunkType)
}
if err != nil {
return nil, err

View File

@@ -27,7 +27,7 @@ import (
contract "github.com/coze-dev/coze-studio/backend/infra/contract/document/parser"
)
func parseXLSX(config *contract.Config) parseFn {
func ParseXLSX(config *contract.Config) ParseFn {
return func(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error) {
f, err := excelize.OpenReader(reader)
if err != nil {

View File

@@ -88,7 +88,7 @@ func TestParseXLSX(t *testing.T) {
ChunkingStrategy: nil,
}
pfn := parseXLSX(config)
pfn := ParseXLSX(config)
docs, err := pfn(ctx, reader, parser.WithExtraMeta(map[string]any{
"document_id": int64(123),
"knowledge_id": int64(456),
@@ -159,7 +159,7 @@ func TestParseXLSXConvertColumnType(t *testing.T) {
ChunkingStrategy: nil,
}
pfn := parseXLSX(config)
pfn := ParseXLSX(config)
docs, err := pfn(ctx, reader, parser.WithExtraMeta(map[string]any{
"document_id": int64(123),
"knowledge_id": int64(456),

View File

@@ -24,12 +24,12 @@ import (
"github.com/cloudwego/eino/schema"
)
type p struct {
parseFn
type Parser struct {
ParseFn
}
func (p p) Parse(ctx context.Context, reader io.Reader, opts ...parser.Option) ([]*schema.Document, error) {
return p.parseFn(ctx, reader, opts...)
func (p Parser) Parse(ctx context.Context, reader io.Reader, opts ...parser.Option) ([]*schema.Document, error) {
return p.ParseFn(ctx, reader, opts...)
}
type parseFn func(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error)
type ParseFn func(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error)

View File

@@ -73,15 +73,15 @@ func (p *pyPDFTableIterator) NextRow() (row []string, end bool, err error) {
return row, false, nil
}
func parseByPython(config *contract.Config, storage storage.Storage, ocr ocr.OCR, pyPath, scriptPath string) parseFn {
func ParseByPython(config *contract.Config, storage storage.Storage, ocr ocr.OCR, pyPath, scriptPath string) ParseFn {
return func(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error) {
pr, pw, err := os.Pipe()
if err != nil {
return nil, fmt.Errorf("[parseByPython] create rpipe failed, %w", err)
return nil, fmt.Errorf("[ParseByPython] create rpipe failed, %w", err)
}
r, w, err := os.Pipe()
if err != nil {
return nil, fmt.Errorf("[parseByPython] create pipe failed: %w", err)
return nil, fmt.Errorf("[ParseByPython] create pipe failed: %w", err)
}
options := parser.GetCommonOptions(&parser.Options{ExtraMeta: map[string]any{}}, opts...)
@@ -91,13 +91,13 @@ func parseByPython(config *contract.Config, storage storage.Storage, ocr ocr.OCR
FilterPages: config.ParsingStrategy.FilterPages,
})
if err != nil {
return nil, fmt.Errorf("[parseByPython] create parse request failed, %w", err)
return nil, fmt.Errorf("[ParseByPython] create parse request failed, %w", err)
}
if _, err = pw.Write(reqb); err != nil {
return nil, fmt.Errorf("[parseByPython] write parse request bytes failed, %w", err)
return nil, fmt.Errorf("[ParseByPython] write parse request bytes failed, %w", err)
}
if err = pw.Close(); err != nil {
return nil, fmt.Errorf("[parseByPython] close write request pipe failed, %w", err)
return nil, fmt.Errorf("[ParseByPython] close write request pipe failed, %w", err)
}
cmd := exec.Command(pyPath, scriptPath)
@@ -105,31 +105,31 @@ func parseByPython(config *contract.Config, storage storage.Storage, ocr ocr.OCR
cmd.Stdout = os.Stdout
cmd.ExtraFiles = []*os.File{w, pr}
if err = cmd.Start(); err != nil {
return nil, fmt.Errorf("[parseByPython] failed to start Python script: %w", err)
return nil, fmt.Errorf("[ParseByPython] failed to start Python script: %w", err)
}
if err = w.Close(); err != nil {
return nil, fmt.Errorf("[parseByPython] failed to close write pipe: %w", err)
return nil, fmt.Errorf("[ParseByPython] failed to close write pipe: %w", err)
}
result := &pyParseResult{}
if err = json.NewDecoder(r).Decode(result); err != nil {
return nil, fmt.Errorf("[parseByPython] failed to decode result: %w", err)
return nil, fmt.Errorf("[ParseByPython] failed to decode result: %w", err)
}
if err = cmd.Wait(); err != nil {
return nil, fmt.Errorf("[parseByPython] cmd wait err: %w", err)
return nil, fmt.Errorf("[ParseByPython] cmd wait err: %w", err)
}
if result.Error != "" {
return nil, fmt.Errorf("[parseByPython] python execution failed: %s", result.Error)
return nil, fmt.Errorf("[ParseByPython] python execution failed: %s", result.Error)
}
for i, item := range result.Content {
switch item.Type {
case contentTypeText:
partDocs, err := chunkCustom(ctx, item.Content, config, opts...)
partDocs, err := ChunkCustom(ctx, item.Content, config, opts...)
if err != nil {
return nil, fmt.Errorf("[parseByPython] chunk text failed, %w", err)
return nil, fmt.Errorf("[ParseByPython] chunk text failed, %w", err)
}
docs = append(docs, partDocs...)
case contentTypeImage:
@@ -138,9 +138,9 @@ func parseByPython(config *contract.Config, storage storage.Storage, ocr ocr.OCR
}
image, err := base64.StdEncoding.DecodeString(item.Content)
if err != nil {
return nil, fmt.Errorf("[parseByPython] decode image failed, %w", err)
return nil, fmt.Errorf("[ParseByPython] decode image failed, %w", err)
}
imgSrc, err := putImageObject(ctx, storage, "png", getCreatorIDFromExtraMeta(options.ExtraMeta), image)
imgSrc, err := PutImageObject(ctx, storage, "png", GetCreatorIDFromExtraMeta(options.ExtraMeta), image)
if err != nil {
return nil, err
}
@@ -148,7 +148,7 @@ func parseByPython(config *contract.Config, storage storage.Storage, ocr ocr.OCR
if config.ParsingStrategy.ImageOCR && ocr != nil {
texts, err := ocr.FromBase64(ctx, item.Content)
if err != nil {
return nil, fmt.Errorf("[parseByPython] FromBase64 failed, %w", err)
return nil, fmt.Errorf("[ParseByPython] FromBase64 failed, %w", err)
}
label += strings.Join(texts, "\n")
}
@@ -181,15 +181,15 @@ func parseByPython(config *contract.Config, storage storage.Storage, ocr ocr.OCR
ChunkingStrategy: config.ChunkingStrategy,
}, opts...)
if err != nil {
return nil, fmt.Errorf("[parseByPython] parse table failed, %w", err)
return nil, fmt.Errorf("[ParseByPython] parse table failed, %w", err)
}
fmtTableDocs, err := formatTablesInDocument(rawTableDocs)
if err != nil {
return nil, fmt.Errorf("[parseByPython] format table failed, %w", err)
return nil, fmt.Errorf("[ParseByPython] format table failed, %w", err)
}
docs = append(docs, fmtTableDocs...)
default:
return nil, fmt.Errorf("[parseByPython] invalid content type: %s", item.Type)
return nil, fmt.Errorf("[ParseByPython] invalid content type: %s", item.Type)
}
}

View File

@@ -61,7 +61,7 @@ func getExtension(uri string) string {
return ""
}
func getCreatorIDFromExtraMeta(extraMeta map[string]any) int64 {
func GetCreatorIDFromExtraMeta(extraMeta map[string]any) int64 {
if extraMeta == nil {
return 0
}

View File

@@ -0,0 +1,91 @@
/*
* Copyright 2025 coze-dev Authors
*
* 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 ppstructure
import (
"fmt"
"github.com/coze-dev/coze-studio/backend/infra/contract/chatmodel"
"github.com/coze-dev/coze-studio/backend/infra/contract/document/ocr"
"github.com/coze-dev/coze-studio/backend/infra/contract/document/parser"
"github.com/coze-dev/coze-studio/backend/infra/contract/storage"
"github.com/coze-dev/coze-studio/backend/infra/impl/document/parser/builtin"
"github.com/coze-dev/coze-studio/backend/pkg/goutil"
)
func NewManager(apiConfig *APIConfig, ocr ocr.OCR, storage storage.Storage, imageAnnotationModel chatmodel.BaseChatModel) parser.Manager {
return &manager{
apiConfig: apiConfig,
ocr: ocr,
storage: storage,
imageAnnotationModel: imageAnnotationModel,
}
}
type manager struct {
apiConfig *APIConfig
ocr ocr.OCR
storage storage.Storage
imageAnnotationModel chatmodel.BaseChatModel
}
func (m *manager) GetParser(config *parser.Config) (parser.Parser, error) {
if config.ParsingStrategy.HeaderLine == 0 && config.ParsingStrategy.DataStartLine == 0 {
config.ParsingStrategy.DataStartLine = 1
} else if config.ParsingStrategy.HeaderLine >= config.ParsingStrategy.DataStartLine {
return nil, fmt.Errorf("[GetParser] invalid header line and data start line, header=%d, data_start=%d",
config.ParsingStrategy.HeaderLine, config.ParsingStrategy.DataStartLine)
}
var pFn builtin.ParseFn
switch config.FileExtension {
case parser.FileExtensionPDF:
fileType := 0
return &ppstructureParser{config, m.apiConfig, fileType, m.ocr, m.storage}, nil
case parser.FileExtensionTXT:
pFn = builtin.ParseText(config)
return &builtin.Parser{ParseFn: pFn}, nil
case parser.FileExtensionMarkdown:
pFn = builtin.ParseMarkdown(config, m.storage, m.ocr)
return &builtin.Parser{ParseFn: pFn}, nil
case parser.FileExtensionDocx:
pFn = builtin.ParseByPython(config, m.storage, m.ocr, goutil.GetPython3Path(), goutil.GetPythonFilePath("parse_docx.py"))
return &builtin.Parser{ParseFn: pFn}, nil
case parser.FileExtensionCSV:
pFn = builtin.ParseCSV(config)
return &builtin.Parser{ParseFn: pFn}, nil
case parser.FileExtensionXLSX:
pFn = builtin.ParseXLSX(config)
return &builtin.Parser{ParseFn: pFn}, nil
case parser.FileExtensionJSON:
pFn = builtin.ParseJSON(config)
return &builtin.Parser{ParseFn: pFn}, nil
case parser.FileExtensionJsonMaps:
pFn = builtin.ParseJSONMaps(config)
return &builtin.Parser{ParseFn: pFn}, nil
case parser.FileExtensionJPG, parser.FileExtensionJPEG, parser.FileExtensionPNG:
pFn = builtin.ParseImage(config, m.imageAnnotationModel)
return &builtin.Parser{ParseFn: pFn}, nil
default:
return nil, fmt.Errorf("[Parse] document type not support, type=%s", config.FileExtension)
}
}
func (m *manager) IsAutoAnnotationSupported() bool {
return m.imageAnnotationModel != nil
}

View File

@@ -0,0 +1,324 @@
/*
* Copyright 2025 coze-dev Authors
*
* 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 ppstructure
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"regexp"
"strings"
"github.com/cloudwego/eino/components/document/parser"
"github.com/cloudwego/eino/schema"
"github.com/coze-dev/coze-studio/backend/infra/contract/document/ocr"
contract "github.com/coze-dev/coze-studio/backend/infra/contract/document/parser"
"github.com/coze-dev/coze-studio/backend/infra/contract/storage"
"github.com/coze-dev/coze-studio/backend/infra/impl/document/parser/builtin"
)
type ppstructureParser struct {
parserConfig *contract.Config
apiConfig *APIConfig
fileType int
ocr ocr.OCR
storage storage.Storage
}
type APIConfig struct {
Client *http.Client
URL string
// see: https://paddlepaddle.github.io/PaddleX/latest/pipeline_usage/tutorials/ocr_pipelines/PP-StructureV3.html#3
UseDocOrientationClassify *bool
UseDocUnwarping *bool
UseTextlineOrientation *bool
UseSealRecognition *bool
UseFormulaRecognition *bool
UseChartRecognition *bool
UseRegionDetection *bool
LayoutThreshold *float64
LayoutNms *bool
LayoutUnclipRatio *float64
LayoutMergeBboxesMode *string
TextDetLimitSideLen *int64
TextDetLimitType *string
TextDetThresh *float64
TextDetBoxThresh *float64
TextDetUnclipRatio *float64
TextRecScoreThresh *float64
SealDetLimitSideLen *int64
SealDetLimitType *string
SealDetThresh *float64
SealDetBoxThresh *float64
SealDetUnclipRatio *float64
SealRecScoreThresh *float64
UseWiredTableCellsTransToHtml *bool
UseWirelessTableCellsTransToHtml *bool
UseTableOrientationClassify *bool
UseOcrResultsWithTableCells *bool
UseE2eWiredTableRecModel *bool
UseE2eWirelessTableRecModel *bool
}
type ppstructureResponse struct {
Result *ppstructureInferResult `json:"result"`
}
type ppstructureInferResult struct {
LayoutParsingResults []*ppstructureInnerResult `json:"layoutParsingResults"`
}
type ppstructureInnerResult struct {
Markdown *ppstructureMarkdown `json:"markdown"`
}
type ppstructureMarkdown struct {
Text *string `json:"text"`
Images map[string]string `json:"images"`
IsStart *bool `json:"isStart"`
IsEnd *bool `json:"isEnd"`
}
func (p *ppstructureParser) Parse(ctx context.Context, reader io.Reader, opts ...parser.Option) (docs []*schema.Document, err error) {
// TODO(Bobholamovic): Current chunking strategy is rather naive; we should
// implement a more sophisticated one that at least takes tables and text
// extracted from the images into consideration.
options := parser.GetCommonOptions(&parser.Options{ExtraMeta: map[string]any{}}, opts...)
fileBytes, err := io.ReadAll(reader)
if err != nil {
return nil, fmt.Errorf("[Parse] failed to read the file bytes, %w", err)
}
b64 := base64.StdEncoding.EncodeToString(fileBytes)
reqBody := p.newRequestBody(b64, p.fileType, p.parserConfig.ParsingStrategy.ExtractImage, p.parserConfig.ParsingStrategy.ExtractTable)
bodyBytes, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("[Parse] failed to serizalize the request body, %w", err)
}
req, err := http.NewRequest("POST", p.apiConfig.URL, bytes.NewReader(bodyBytes))
if err != nil {
return nil, fmt.Errorf("[Parse] failed to create a new request, %w", err)
}
req.Header.Set("Content-Type", "application/json")
resp, err := p.apiConfig.Client.Do(req)
if err != nil {
return nil, fmt.Errorf("[Parse] request failed, %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("[Parse] request failed, %w", err)
}
respBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("[Parse] failed to read the response body, %w", err)
}
var res ppstructureResponse
if err := json.Unmarshal(respBody, &res); err != nil {
return nil, fmt.Errorf("[Parse] failed to deserialize the response body, %w", err)
}
if res.Result == nil ||
res.Result.LayoutParsingResults == nil {
return nil, fmt.Errorf("[Parse] failed to get the layout parsing result, %w", err)
}
for i, item := range res.Result.LayoutParsingResults {
if item.Markdown == nil {
return nil, fmt.Errorf("[Parse] invalid response, %w", err)
}
if item.Markdown.Text == nil {
return nil, fmt.Errorf("[Parse] invalid response, %w", err)
}
should_skip := false
for _, v := range p.parserConfig.ParsingStrategy.FilterPages {
if i+1 == v {
should_skip = true
break
}
}
if should_skip {
continue
}
text := *item.Markdown.Text
// Convert the image in markdown to comments, as the image content will be added later.
pattern := `(?i)<img[^>]*>`
re := regexp.MustCompile(pattern)
// TODO(Bobholamovic): Add image numbering
text = re.ReplaceAllStringFunc(text, func(matched string) string {
return "<!-- image -->"
})
partDocs, err := builtin.ChunkCustom(ctx, text, p.parserConfig, opts...)
if err != nil {
return nil, fmt.Errorf("[Parse] chunk text failed, %w", err)
}
docs = append(docs, partDocs...)
if p.parserConfig.ParsingStrategy.ExtractImage {
if item.Markdown.Images == nil {
return nil, fmt.Errorf("[Parse] invalid response, %w", err)
}
for _, v := range item.Markdown.Images {
image, err := base64.StdEncoding.DecodeString(v)
if err != nil {
return nil, fmt.Errorf("[Parse] failed to decode an image, %w", err)
}
imgSrc, err := builtin.PutImageObject(ctx, p.storage, "png", builtin.GetCreatorIDFromExtraMeta(options.ExtraMeta), image)
if err != nil {
return nil, err
}
label := fmt.Sprintf("\n%s", imgSrc)
if p.parserConfig.ParsingStrategy.ImageOCR && p.ocr != nil {
texts, err := p.ocr.FromBase64(ctx, v)
if err != nil {
return nil, fmt.Errorf("[Parse] FromBase64 failed, %w", err)
}
label += strings.Join(texts, "\n")
}
doc := &schema.Document{
Content: label,
MetaData: map[string]any{},
}
for k, v := range options.ExtraMeta {
doc.MetaData[k] = v
}
docs = append(docs, doc)
}
}
}
return docs, nil
}
func (p *ppstructureParser) newRequestBody(file string, fileType int, extractImage bool, extractTable bool) map[string]interface{} {
payload := map[string]interface{}{
"file": file,
"fileType": fileType,
"useTableRecognition": extractTable,
"visualize": extractImage,
}
if p.apiConfig.UseDocOrientationClassify != nil {
payload["useDocOrientationClassify"] = *p.apiConfig.UseDocOrientationClassify
}
if p.apiConfig.UseDocUnwarping != nil {
payload["useDocUnwarping"] = *p.apiConfig.UseDocUnwarping
}
if p.apiConfig.UseTextlineOrientation != nil {
payload["useTextlineOrientation"] = *p.apiConfig.UseTextlineOrientation
}
if p.apiConfig.UseSealRecognition != nil {
payload["useSealRecognition"] = *p.apiConfig.UseSealRecognition
}
if p.apiConfig.UseFormulaRecognition != nil {
payload["useFormulaRecognition"] = *p.apiConfig.UseFormulaRecognition
}
if p.apiConfig.UseChartRecognition != nil {
payload["useChartRecognition"] = *p.apiConfig.UseChartRecognition
}
if p.apiConfig.UseRegionDetection != nil {
payload["useRegionDetection"] = *p.apiConfig.UseRegionDetection
}
if p.apiConfig.LayoutThreshold != nil {
payload["layoutThreshold"] = *p.apiConfig.LayoutThreshold
}
if p.apiConfig.LayoutNms != nil {
payload["layoutNms"] = *p.apiConfig.LayoutNms
}
if p.apiConfig.LayoutUnclipRatio != nil {
payload["layoutUnclipRatio"] = *p.apiConfig.LayoutUnclipRatio
}
if p.apiConfig.LayoutMergeBboxesMode != nil {
payload["layoutMergeBboxesMode"] = *p.apiConfig.LayoutMergeBboxesMode
}
if p.apiConfig.TextDetLimitSideLen != nil {
payload["textDetLimitSideLen"] = *p.apiConfig.TextDetLimitSideLen
}
if p.apiConfig.TextDetLimitType != nil {
payload["textDetLimitType"] = *p.apiConfig.TextDetLimitType
}
if p.apiConfig.TextDetThresh != nil {
payload["textDetThresh"] = *p.apiConfig.TextDetThresh
}
if p.apiConfig.TextDetBoxThresh != nil {
payload["textDetBoxThresh"] = *p.apiConfig.TextDetBoxThresh
}
if p.apiConfig.TextDetUnclipRatio != nil {
payload["textDetUnclipRatio"] = *p.apiConfig.TextDetUnclipRatio
}
if p.apiConfig.TextRecScoreThresh != nil {
payload["textRecScoreThresh"] = *p.apiConfig.TextRecScoreThresh
}
if p.apiConfig.SealDetLimitSideLen != nil {
payload["sealDetLimitSideLen"] = *p.apiConfig.SealDetLimitSideLen
}
if p.apiConfig.SealDetLimitType != nil {
payload["sealDetLimitType"] = *p.apiConfig.SealDetLimitType
}
if p.apiConfig.SealDetThresh != nil {
payload["sealDetThresh"] = *p.apiConfig.SealDetThresh
}
if p.apiConfig.SealDetBoxThresh != nil {
payload["sealDetBoxThresh"] = *p.apiConfig.SealDetBoxThresh
}
if p.apiConfig.SealDetUnclipRatio != nil {
payload["sealDetUnclipRatio"] = *p.apiConfig.SealDetUnclipRatio
}
if p.apiConfig.SealRecScoreThresh != nil {
payload["sealRecScoreThresh"] = *p.apiConfig.SealRecScoreThresh
}
if p.apiConfig.UseWiredTableCellsTransToHtml != nil {
payload["useWiredTableCellsTransToHtml"] = *p.apiConfig.UseWiredTableCellsTransToHtml
}
if p.apiConfig.UseWirelessTableCellsTransToHtml != nil {
payload["useWirelessTableCellsTransToHtml"] = *p.apiConfig.UseWirelessTableCellsTransToHtml
}
if p.apiConfig.UseTableOrientationClassify != nil {
payload["useTableOrientationClassify"] = *p.apiConfig.UseTableOrientationClassify
}
if p.apiConfig.UseOcrResultsWithTableCells != nil {
payload["useOcrResultsWithTableCells"] = *p.apiConfig.UseOcrResultsWithTableCells
}
if p.apiConfig.UseE2eWiredTableRecModel != nil {
payload["useE2eWiredTableRecModel"] = *p.apiConfig.UseE2eWiredTableRecModel
}
if p.apiConfig.UseE2eWirelessTableRecModel != nil {
payload["useE2eWirelessTableRecModel"] = *p.apiConfig.UseE2eWirelessTableRecModel
}
return payload
}