Merge remote-tracking branch 'origin/upstream/main'

common/image/image.go

@@ -1,6 +1,8 @@
 package image
 
 import (
+	"bytes"
+	"encoding/base64"
 	"image"
 	_ "image/gif"
 	_ "image/jpeg"
@@ -8,11 +10,27 @@ import (
 	"net/http"
 	"regexp"
 	"strings"
+	"sync"
 
 	_ "golang.org/x/image/webp"
 )
 
+func IsImageUrl(url string) (bool, error) {
+	resp, err := http.Head(url)
+	if err != nil {
+		return false, err
+	}
+	if !strings.HasPrefix(resp.Header.Get("Content-Type"), "image/") {
+		return false, nil
+	}
+	return true, nil
+}
+
 func GetImageSizeFromUrl(url string) (width int, height int, err error) {
+	isImage, err := IsImageUrl(url)
+	if !isImage {
+		return
+	}
 	resp, err := http.Get(url)
 	if err != nil {
 		return
@@ -25,17 +43,51 @@ func GetImageSizeFromUrl(url string) (width int, height int, err error) {
 	return img.Width, img.Height, nil
 }
 
+func GetImageFromUrl(url string) (mimeType string, data string, err error) {
+	isImage, err := IsImageUrl(url)
+	if !isImage {
+		return
+	}
+	resp, err := http.Get(url)
+	if err != nil {
+		return
+	}
+	defer resp.Body.Close()
+	buffer := bytes.NewBuffer(nil)
+	_, err = buffer.ReadFrom(resp.Body)
+	if err != nil {
+		return
+	}
+	mimeType = resp.Header.Get("Content-Type")
+	data = base64.StdEncoding.EncodeToString(buffer.Bytes())
+	return
+}
+
 var (
 	reg = regexp.MustCompile(`data:image/([^;]+);base64,`)
 )
 
+var readerPool = sync.Pool{
+	New: func() interface{} {
+		return &bytes.Reader{}
+	},
+}
+
 func GetImageSizeFromBase64(encoded string) (width int, height int, err error) {
-	encoded = strings.TrimPrefix(encoded, "data:image/png;base64,")
-	base64 := strings.NewReader(reg.ReplaceAllString(encoded, ""))
-	img, _, err := image.DecodeConfig(base64)
+	decoded, err := base64.StdEncoding.DecodeString(reg.ReplaceAllString(encoded, ""))
 	if err != nil {
-		return
+		return 0, 0, err
 	}
+
+	reader := readerPool.Get().(*bytes.Reader)
+	defer readerPool.Put(reader)
+	reader.Reset(decoded)
+
+	img, _, err := image.DecodeConfig(reader)
+	if err != nil {
+		return 0, 0, err
+	}
+
 	return img.Width, img.Height, nil
 }
 

common/image/image_test.go

@@ -152,3 +152,20 @@ func TestGetImageSize(t *testing.T) {
 		})
 	}
 }
+
+func TestGetImageSizeFromBase64(t *testing.T) {
+	for i, c := range cases {
+		t.Run("Decode:"+strconv.Itoa(i), func(t *testing.T) {
+			resp, err := http.Get(c.url)
+			assert.NoError(t, err)
+			defer resp.Body.Close()
+			data, err := io.ReadAll(resp.Body)
+			assert.NoError(t, err)
+			encoded := base64.StdEncoding.EncodeToString(data)
+			width, height, err := img.GetImageSizeFromBase64(encoded)
+			assert.NoError(t, err)
+			assert.Equal(t, c.width, width)
+			assert.Equal(t, c.height, height)
+		})
+	}
+}

common/model-ratio.go

@@ -84,6 +84,7 @@ var ModelRatio = map[string]float64{
 	"Embedding-V1":              0.1429, // ￥0.002 / 1k tokens
 	"PaLM-2":                    1,
 	"gemini-pro":                1,      // $0.00025 / 1k characters -> $0.001 / 1k tokens
+	"gemini-pro-vision":         1,      // $0.00025 / 1k characters -> $0.001 / 1k tokens
 	"chatglm_turbo":             0.3572, // ￥0.005 / 1k tokens
 	"chatglm_pro":               0.7143, // ￥0.01 / 1k tokens
 	"chatglm_std":               0.3572, // ￥0.005 / 1k tokens
@@ -115,6 +116,9 @@ func UpdateModelRatioByJSONString(jsonStr string) error {
 }
 
 func GetModelRatio(name string) float64 {
+	if strings.HasPrefix(name, "qwen-") && strings.HasSuffix(name, "-internet") {
+		name = strings.TrimSuffix(name, "-internet")
+	}
 	ratio, ok := ModelRatio[name]
 	if !ok {
 		SysError("model ratio not found: " + name)