From 4825404d375622dff567deefdd69dd7495fa8c35 Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Thu, 15 May 2025 14:51:33 +0800
Subject: [PATCH] feat: enhance image decoding logic to handle base64 file
 types and improve error handling

---
 service/token_counter.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/service/token_counter.go b/service/token_counter.go
index 21b882af..d63b54ad 100644
--- a/service/token_counter.go
+++ b/service/token_counter.go
@@ -120,11 +120,12 @@ func getImageToken(info *relaycommon.RelayInfo, imageUrl *dto.MessageImageUrl, m
 	var config image.Config
 	var err error
 	var format string
+	var b64str string
 	if strings.HasPrefix(imageUrl.Url, "http") {
 		config, format, err = DecodeUrlImageData(imageUrl.Url)
 	} else {
 		common.SysLog(fmt.Sprintf("decoding image"))
-		config, format, _, err = DecodeBase64ImageData(imageUrl.Url)
+		config, format, b64str, err = DecodeBase64ImageData(imageUrl.Url)
 	}
 	if err != nil {
 		return 0, err
@@ -132,7 +133,12 @@ func getImageToken(info *relaycommon.RelayInfo, imageUrl *dto.MessageImageUrl, m
 	imageUrl.MimeType = format
 
 	if config.Width == 0 || config.Height == 0 {
-		return 0, errors.New(fmt.Sprintf("fail to decode image config: %s", imageUrl.Url))
+		// not an image
+		if format != "" && b64str != "" {
+			// file type
+			return 3 * baseTokens, nil
+		}
+		return 0, errors.New(fmt.Sprintf("fail to decode base64 config: %s", imageUrl.Url))
 	}
 
 	shortSide := config.Width