Files
model-gateway/common/util/files.go

116 lines
3.0 KiB
Go
Raw Normal View History

package util
import (
"encoding/json"
"fmt"
"net/http"
"os"
"path/filepath"
"strings"
)
// DetectFileType 根据返回的二进制内容推断 contentType + 扩展名(尽量稳定)
func DetectFileType(data []byte) (contentType string, ext string) {
if len(data) == 0 {
return "application/octet-stream", ""
}
ct := http.DetectContentType(data)
// gateway.DetectContentType 可能带 charset 等参数text/plain; charset=utf-8
if idx := strings.Index(ct, ";"); idx > 0 {
ct = strings.TrimSpace(ct[:idx])
}
switch ct {
case "audio/mpeg":
return ct, ".mp3"
case "audio/wave", "audio/wav", "audio/x-wav":
return ct, ".wav"
case "video/mp4":
return ct, ".mp4"
case "image/png":
return ct, ".png"
case "image/jpeg":
return ct, ".jpg"
case "application/pdf":
return ct, ".pdf"
case "text/plain":
return ct, ".txt"
case "application/json":
return ct, ".json"
default:
// 兜底:尝试从 ct 截取 subtype 作为后缀(例如 application/json
if parts := strings.Split(ct, "/"); len(parts) == 2 {
sub := parts[1]
// 避免出现 "plain; charset=utf-8" 之类的后缀
if idx := strings.Index(sub, ";"); idx > 0 {
sub = strings.TrimSpace(sub[:idx])
}
return ct, "." + sub
}
return ct, ""
}
}
// SaveTmpResult 将模型输出写入临时文件,用于 OSS 上传失败后的“仅重试 OSS”。
func SaveTmpResult(taskID string, data []byte, ext string) (string, error) {
dir := filepath.Join(os.TempDir(), "model-asynch")
if err := os.MkdirAll(dir, 0o755); err != nil {
return "", err
}
if ext == "" {
ext = ".bin"
}
if ext[0] != '.' {
ext = "." + ext
}
path := filepath.Join(dir, fmt.Sprintf("%s%s", taskID, ext))
if err := os.WriteFile(path, data, 0o644); err != nil {
return "", err
}
return path, nil
}
// SaveTempFileByType
// 根据传入的数据自动判断:
// 若是 []byte 且后缀为 .mp3 → 保存二进制音频
// 若是任意结构体/map → 自动转 JSON 保存
// 返回:新临时文件路径、错误
func SaveTempFileByType(taskID string, data any, oldTmpFile string) (string, error) {
// 1. 先清理旧临时文件(统一逻辑)
if oldTmpFile != "" {
_ = os.Remove(oldTmpFile)
}
var tmpPath string
var tmpErr error
// 2. 判断是否是二进制音频([]byte + .mp3
if audioData, ok := data.([]byte); ok {
tmpPath, tmpErr = saveTmpResult(taskID, audioData, ".mp3")
} else {
// 3. 其他类型 → 序列化为 JSON 保存
mappedBytes, err := json.Marshal(data)
if err != nil {
return "", err
}
if len(mappedBytes) == 0 {
return "", nil
}
tmpPath, tmpErr = saveTmpResult(taskID, mappedBytes, ".json")
}
if tmpErr != nil || tmpPath == "" {
return "", tmpErr
}
return tmpPath, nil
}
// saveTmpResult 你原有的底层保存文件方法(保留不动)
func saveTmpResult(taskID string, data []byte, ext string) (string, error) {
// 你原来实现,比如:
filename := taskID + ext
tmpPath := filepath.Join(os.TempDir(), filename)
err := os.WriteFile(tmpPath, data, 0644)
return tmpPath, err
}