116 lines
3.0 KiB
Go
116 lines
3.0 KiB
Go
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
|
||
}
|