Merge branch 'songquanpeng' into sync_upstream

common/image/image.go

@@ -15,7 +15,22 @@ import (
 	_ "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
@@ -28,6 +43,26 @@ 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,`)
 )

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)