🐛 fix base 64 encoded format support of gemini-pro-vision

common/image/image.go

@@ -3,6 +3,7 @@ package image
 import (
 	"bytes"
 	"encoding/base64"
+	"errors"
 	"image"
 	_ "image/gif"
 	_ "image/jpeg"
@@ -44,8 +45,26 @@ func GetImageSizeFromUrl(url string) (width int, height int, err error) {
 }
 
 func GetImageFromUrl(url string) (mimeType string, data string, err error) {
+
+	if strings.HasPrefix(url, "data:image/") {
+		dataURLPattern := regexp.MustCompile(`data:image/([^;]+);base64,(.*)`)
+
+		matches := dataURLPattern.FindStringSubmatch(url)
+		if len(matches) == 3 && matches[2] != "" {
+			mimeType = "image/" + matches[1]
+			data = matches[2]
+			return
+		}
+
+		err = errors.New("image base64 decode failed")
+		return
+	}
+
 	isImage, err := IsImageUrl(url)
 	if !isImage {
+		if err == nil {
+			err = errors.New("Invalid image link")
+		}
 		return
 	}
 	resp, err := http.Get(url)

common/image/image_test.go

@@ -169,3 +169,34 @@ func TestGetImageSizeFromBase64(t *testing.T) {
 		})
 	}
 }
+
+func TestGetImageFromUrl(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)
+
+			mimeType, base64Data, err := img.GetImageFromUrl(c.url)
+			assert.NoError(t, err)
+			assert.Equal(t, encoded, base64Data)
+			assert.Equal(t, "image/"+c.format, mimeType)
+
+			encodedBase64 := "data:image/" + c.format + ";base64," + encoded
+			mimeType, base64Data, err = img.GetImageFromUrl(encodedBase64)
+			assert.NoError(t, err)
+			assert.Equal(t, encoded, base64Data)
+			assert.Equal(t, "image/"+c.format, mimeType)
+		})
+	}
+
+	url := "https://raw.githubusercontent.com/songquanpeng/one-api/main/README.md"
+	_, _, err := img.GetImageFromUrl(url)
+	assert.Error(t, err)
+	encodedBase64 := "data:image/text;base64,"
+	_, _, err = img.GetImageFromUrl(encodedBase64)
+	assert.Error(t, err)
+}