feat: 增加图片统一处理为二进制方法

This commit is contained in:
fuzhongyun 2025-12-16 14:06:14 +08:00
parent 7a48333efc
commit 33b6363233
2 changed files with 110 additions and 5 deletions

View File

@ -4,10 +4,12 @@ import (
"ai_scheduler/internal/data/constants"
"ai_scheduler/internal/entitys"
"ai_scheduler/internal/pkg/l_request"
"bytes"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"path/filepath"
"strings"
@ -15,12 +17,105 @@ import (
)
// HandleRecognizeFile 这里的目的是无论将什么类型的file都转为二进制格式
// 判断文件大小
// 判断文件类型
// 判断文件是否合法
// 最终输出1.将 files.FileData 填充为文件的二进制数据 2.将 files.FileType 填充为文件的类型(当前为 constants.Caller兼容写入其字符串值
// 判断文件大小统一限制为10MB判断文件类型判断文件是否合法类型在白名单映射中无法识别/非法/超限→填充unknown并兼容返回
// 若 FileData 不存在 且 FileUrl 不存在, 则直接退出
// 若 FileData 存在 FileType 存在, 则直接退出
// 若 FileData 存在 FileType 不存在, 则根据 FileData 推断文件类型并填充 FileType
// 若 FileUrl 存在, 则下载文件并填充 FileData 和 FileType
func HandleRecognizeFile(files *entitys.RecognizeFile) {
//Todo 仲云
return
if files == nil {
return
}
const maxSize = 10 * 1024 * 1024 // 10MB 上限
// 工具:根据 MIME 或扩展名映射到 FileType
mapToFileType := func(s string) constants.FileType {
if len(s) == 0 {
return constants.FileTypeUnknown
}
s = strings.ToLower(strings.TrimSpace(s))
for ft, items := range constants.FileTypeMappings {
for _, item := range items {
if !strings.HasPrefix(item, ".") { // MIME
if s == item {
return ft
}
} else { // 扩展名
if s == item {
return ft
}
}
}
}
return constants.FileTypeUnknown
}
// 分支1无数据、无URL→直接返回
if len(files.FileData) == 0 && len(files.FileUrl) == 0 {
return
}
// 分支2已有数据且已有类型→直接返回
if len(files.FileData) > 0 && len(strings.TrimSpace(files.FileType.String())) > 0 {
return
}
// 分支3仅有数据、无类型→内容检测并填充
if len(files.FileData) > 0 && len(strings.TrimSpace(files.FileType.String())) == 0 {
if len(files.FileData) > maxSize {
files.FileType = constants.Caller(constants.FileTypeUnknown)
return
}
reader := bytes.NewReader(files.FileData)
detected := detectFileType(reader, "")
if detected == constants.FileTypeUnknown {
files.FileType = constants.Caller(constants.FileTypeUnknown)
return
}
files.FileType = constants.Caller(detected)
return
}
// 分支4存在URL→下载并填充数据与类型
if len(files.FileUrl) > 0 {
fileBytes, contentType, err := downloadFile(files.FileUrl)
if err != nil || len(fileBytes) == 0 {
files.FileType = constants.Caller(constants.FileTypeUnknown)
return
}
if len(fileBytes) > maxSize {
// 超限:不写入数据,类型置 unknown
files.FileType = constants.Caller(constants.FileTypeUnknown)
return
}
// 优先使用响应头的 Content-Type 映射
detected := mapToFileType(contentType)
if detected == constants.FileTypeUnknown {
// 回退:内容检测 + URL 文件名扩展名辅助
var fname string
if u, perr := url.Parse(files.FileUrl); perr == nil {
fname = filepath.Base(u.Path)
}
reader := bytes.NewReader(fileBytes)
detected = detectFileType(reader, fname)
}
// 写入数据
files.FileData = fileBytes
if detected == constants.FileTypeUnknown {
files.FileType = constants.Caller(constants.FileTypeUnknown)
return
}
files.FileType = constants.Caller(detected)
return
}
}
// 下载文件并返回二进制数据、MIME 类型

View File

@ -10,6 +10,8 @@ const (
FileTypeWord FileType = "word"
FileTypeTxt FileType = "txt"
FileTypePDF FileType = "pdf"
FileTypePPT FileType = "ppt"
FileTypeCSV FileType = "csv"
)
var FileTypeMappings = map[FileType][]string{
@ -35,4 +37,12 @@ var FileTypeMappings = map[FileType][]string{
"text/plain",
".txt",
},
FileTypePPT: {
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
".pptx",
},
FileTypeCSV: {
"text/csv",
".csv",
},
}