xinmi/internal/application/service/digital_human_upload.go

package service

import (
	"context"
	"fmt"
	"path/filepath"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/Tencent/WeKnora/internal/infrastructure/docparser"
	"github.com/Tencent/WeKnora/internal/types"
)

const DigitalHumanFileLimit = 20 << 20
const DigitalHumanUploadLimit = 50 << 20
const DigitalHumanFileCountLimit = 99

type DigitalHumanUpload struct {
	Name string
	Data []byte
}

func DigitalHumanUploadSupported(name string) bool {
	switch strings.ToLower(filepath.Ext(name)) {
	case ".md", ".markdown", ".txt", ".srt", ".vtt", ".doc", ".docx", ".pdf", ".rtf", ".odt", ".ppt", ".pptx", ".csv", ".json", ".html", ".htm":
		return true
	}
	return false
}

func (s *DigitalHumanService) readDigitalHumanUpload(ctx context.Context, file DigitalHumanUpload) (types.DigitalHumanSource, error) {
	name := filepath.Base(file.Name)
	source := types.DigitalHumanSource{Name: name}
	if !DigitalHumanUploadSupported(name) || len(file.Data) == 0 || len(file.Data) > DigitalHumanFileLimit {
		return source, digitalHumanInputError("材料格式不支持、内容为空或单份超过20MB：" + name)
	}
	fileType := strings.TrimPrefix(strings.ToLower(filepath.Ext(name)), ".")
	if fileType == "srt" || fileType == "vtt" || fileType == "txt" || fileType == "md" || fileType == "markdown" {
		if !utf8.Valid(file.Data) {
			return source, digitalHumanInputError("文本或字幕文件须采用 UTF-8 编码：" + name)
		}
		source.Text = strings.TrimSpace(strings.TrimPrefix(string(file.Data), "\ufeff"))
	} else {
		ctx, cancel := context.WithTimeout(ctx, 2*time.Minute)
		defer cancel()
		engine := types.DefaultParserEngine(fileType)
		overrides := map[string]string{"anydoc_extract_images": "false"}
		reader, err := docparser.NewReader(ctx, engine, fileType, false, docparser.ReaderDeps{Remote: s.reader, Overrides: overrides})
		if err != nil {
			return source, digitalHumanInputError("文档解析器不可用，请检查文档解析服务：" + name)
		}
		result, err := reader.Read(ctx, &types.ReadRequest{FileName: name, FileType: fileType, FileContent: file.Data,
			ParserEngine: engine, ParserEngineOverrides: overrides})
		if err != nil || result == nil || result.Error != "" {
			if ctx.Err() != nil {
				return source, digitalHumanInputError("文档解析超时，请拆分文档后重试：" + name)
			}
			return source, digitalHumanInputError("无法解析文档，请检查文件是否损坏；扫描 PDF 需要可用的 OCR 解析器：" + name)
		}
		source.Text = strings.TrimSpace(result.MarkdownContent)
	}
	if source.Text == "" {
		return source, digitalHumanInputError("文档没有可读取的文字：" + name)
	}
	if !utf8.ValidString(source.Text) || utf8.RuneCountInString(source.Text) > digitalHumanSourceLimit {
		return source, digitalHumanInputError(fmt.Sprintf("文档文字超过%d字或编码无效：%s", digitalHumanSourceLimit, name))
	}
	return source, nil
}