Create types/image.go

internal/types/image.go

@@ -0,0 +1,211 @@
+package types
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"monica-proxy/internal/config"
+	"monica-proxy/internal/utils"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/cespare/xxhash/v2"
+	"github.com/google/uuid"
+)
+
+const MaxFileSize = 10 * 1024 * 1024 // 10MB
+
+var imageCache sync.Map
+
+// sampleAndHash 对base64字符串进行采样并计算xxHash
+func sampleAndHash(data string) string {
+	// 如果数据长度小于1024，直接计算整个字符串的哈希
+	if len(data) <= 1024 {
+		return fmt.Sprintf("%x", xxhash.Sum64String(data))
+	}
+
+	// 采样策略：
+	// 1. 取前256字节
+	// 2. 取中间256字节
+	// 3. 取最后256字节
+	var samples []string
+	samples = append(samples, data[:256])
+	mid := len(data) / 2
+	samples = append(samples, data[mid-128:mid+128])
+	samples = append(samples, data[len(data)-256:])
+
+	// 将采样数据拼接后计算哈希
+	return fmt.Sprintf("%x", xxhash.Sum64String(strings.Join(samples, "")))
+}
+
+// UploadBase64Image 上传base64编码的图片到Monica
+func UploadBase64Image(ctx context.Context, base64Data string) (*FileInfo, error) {
+	// 1. 生成缓存key
+	cacheKey := sampleAndHash(base64Data)
+
+	// 2. 检查缓存
+	if value, exists := imageCache.Load(cacheKey); exists {
+		return value.(*FileInfo), nil
+	}
+
+	// 3. 解析base64数据
+	// 移除 "data:image/png;base64," 这样的前缀
+	parts := strings.Split(base64Data, ",")
+	if len(parts) != 2 {
+		return nil, fmt.Errorf("invalid base64 image format")
+	}
+
+	// 获取图片类型
+	mimeType := strings.TrimSuffix(strings.TrimPrefix(parts[0], "data:"), ";base64")
+	if !strings.HasPrefix(mimeType, "image/") {
+		return nil, fmt.Errorf("invalid image mime type: %s", mimeType)
+	}
+
+	// 解码base64数据
+	imageData, err := utils.Base64Decode(parts[1])
+	if err != nil {
+		return nil, fmt.Errorf("decode base64 failed: %v", err)
+	}
+
+	// 4. 验证图片格式和大小
+	fileInfo, err := validateImageBytes(imageData, mimeType)
+	if err != nil {
+		return nil, fmt.Errorf("validate image failed: %v", err)
+	}
+	log.Printf("file info: %+v", fileInfo)
+
+	// 5. 获取预签名URL
+	preSignReq := &PreSignRequest{
+		FilenameList: []string{fileInfo.FileName},
+		Module:       ImageModule,
+		Location:     ImageLocation,
+		ObjID:        uuid.New().String(),
+	}
+
+	var preSignResp PreSignResponse
+	_, err = utils.RestyDefaultClient.R().
+		SetContext(ctx).
+		SetHeader("cookie", config.MonicaConfig.MonicaCookie).
+		SetBody(preSignReq).
+		SetResult(&preSignResp).
+		Post(PreSignURL)
+
+	if err != nil {
+		return nil, fmt.Errorf("get pre-sign url failed: %v", err)
+	}
+
+	if len(preSignResp.Data.PreSignURLList) == 0 || len(preSignResp.Data.ObjectURLList) == 0 {
+		return nil, fmt.Errorf("no pre-sign url or object url returned")
+	}
+	log.Printf("preSign info: %+v", preSignResp)
+
+	// 6. 上传图片数据
+	_, err = utils.RestyDefaultClient.R().
+		SetContext(ctx).
+		SetHeader("Content-Type", fileInfo.FileType).
+		SetBody(imageData).
+		Put(preSignResp.Data.PreSignURLList[0])
+
+	if err != nil {
+		return nil, fmt.Errorf("upload file failed: %v", err)
+	}
+
+	// 7. 创建文件对象
+	fileInfo.ObjectURL = preSignResp.Data.ObjectURLList[0]
+	uploadReq := &FileUploadRequest{
+		Data: []FileInfo{*fileInfo},
+	}
+
+	var uploadResp FileUploadResponse
+	_, err = utils.RestyDefaultClient.R().
+		SetContext(ctx).
+		SetHeader("cookie", config.MonicaConfig.MonicaCookie).
+		SetBody(uploadReq).
+		SetResult(&uploadResp).
+		Post(FileUploadURL)
+
+	if err != nil {
+		return nil, fmt.Errorf("create file object failed: %v", err)
+	}
+	log.Printf("uploadResp: %+v", uploadResp)
+	if len(uploadResp.Data.Items) > 0 {
+		fileInfo.FileName = uploadResp.Data.Items[0].FileName
+		fileInfo.FileType = uploadResp.Data.Items[0].FileType
+		fileInfo.FileSize = uploadResp.Data.Items[0].FileSize
+		fileInfo.FileUID = uploadResp.Data.Items[0].FileUID
+		fileInfo.FileExt = uploadResp.Data.Items[0].FileType
+		fileInfo.FileTokens = uploadResp.Data.Items[0].FileTokens
+		fileInfo.FileChunks = uploadResp.Data.Items[0].FileChunks
+	}
+
+	fileInfo.UseFullText = true
+	fileInfo.FileURL = preSignResp.Data.CDNURLList[0]
+
+	// 8. 获取文件llm读取结果知道有返回
+	var batchResp FileBatchGetResponse
+	reqMap := make(map[string][]string)
+	reqMap["file_uids"] = []string{fileInfo.FileUID}
+	var retryCount = 1
+	for {
+		if retryCount > 5 {
+			return nil, fmt.Errorf("retry limit exceeded")
+		}
+		_, err = utils.RestyDefaultClient.R().
+			SetContext(ctx).
+			SetHeader("cookie", config.MonicaConfig.MonicaCookie).
+			SetBody(reqMap).
+			SetResult(&batchResp).
+			Post(FileGetURL)
+		if err != nil {
+			return nil, fmt.Errorf("batch get file failed: %v", err)
+		}
+		if len(batchResp.Data.Items) > 0 && batchResp.Data.Items[0].FileChunks > 0 {
+			break
+		} else {
+			retryCount++
+		}
+		time.Sleep(1 * time.Second)
+	}
+	fileInfo.FileChunks = batchResp.Data.Items[0].FileChunks
+	fileInfo.FileTokens = batchResp.Data.Items[0].FileTokens
+	fileInfo.URL = ""
+	fileInfo.ObjectURL = ""
+
+	// 9. 保存到缓存
+	imageCache.Store(cacheKey, fileInfo)
+
+	return fileInfo, nil
+}
+
+// validateImageBytes 验证图片字节数据的格式和大小
+func validateImageBytes(imageData []byte, mimeType string) (*FileInfo, error) {
+	if len(imageData) > MaxFileSize {
+		return nil, fmt.Errorf("file size exceeds limit: %d > %d", len(imageData), MaxFileSize)
+	}
+
+	contentType := http.DetectContentType(imageData)
+	if !SupportedImageTypes[contentType] {
+		return nil, fmt.Errorf("unsupported image type: %s", contentType)
+	}
+
+	// 根据MIME类型生成文件扩展名
+	ext := ".png"
+	switch mimeType {
+	case "image/jpeg":
+		ext = ".jpg"
+	case "image/gif":
+		ext = ".gif"
+	case "image/webp":
+		ext = ".webp"
+	}
+
+	fileName := fmt.Sprintf("%s%s", uuid.New().String(), ext)
+
+	return &FileInfo{
+		FileName: fileName,
+		FileSize: int64(len(imageData)),
+		FileType: contentType,
+	}, nil
+}