From ae00a99cf5affade2ea54827bf8d9f655b8e1b5d Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Thu, 18 Jul 2024 17:04:19 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AA=92=E4=BD=93=E8=AF=B7=E6=B1=82?=
 =?UTF-8?q?=E8=AE=A1=E8=B4=B9=E9=80=89=E9=A1=B9?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 README.md                |  2 ++
 constant/env.go          |  4 ++++
 service/token_counter.go | 10 +++++++++-
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 785cca8..ffc0e5a 100644
--- a/README.md
+++ b/README.md
@@ -74,6 +74,8 @@
 - `STREAMING_TIMEOUT`：设置流式一次回复的超时时间，默认为 30 秒
 - `DIFY_DEBUG`：设置 Dify 渠道是否输出工作流和节点信息到客户端，默认为 `true`
 - `FORCE_STREAM_OPTION`：是否覆盖客户端stream_options参数，请求上游返回流模式usage，默认为 `true`
+- `GET_MEDIA_TOKEN`：是统计图片token，默认为 `true`，关闭后将不再在本地计算图片token，可能会导致和上游计费不同，此项覆盖 `GET_MEDIA_TOKEN_NOT_STREAM` 选项作用，
+- `GET_MEDIA_TOKEN_NOT_STREAM`：是否在非流（`stream=false`）情况下统计图片token，默认为 `true`
 ## 部署
 ### 部署要求
 - 本地数据库（默认）：SQLite（Docker 部署默认使用 SQLite，必须挂载 `/data` 目录到宿主机）
diff --git a/constant/env.go b/constant/env.go
index 96483fe..a18b875 100644
--- a/constant/env.go
+++ b/constant/env.go
@@ -9,3 +9,7 @@ var DifyDebug = common.GetEnvOrDefaultBool("DIFY_DEBUG", true)
 
 // ForceStreamOption 覆盖请求参数，强制返回usage信息
 var ForceStreamOption = common.GetEnvOrDefaultBool("FORCE_STREAM_OPTION", true)
+
+var GetMediaToken = common.GetEnvOrDefaultBool("GET_MEDIA_TOKEN", true)
+
+var GetMediaTokenNotStream = common.GetEnvOrDefaultBool("GET_MEDIA_TOKEN_NOT_STREAM", true)
diff --git a/service/token_counter.go b/service/token_counter.go
index b99fc20..a1ab0dc 100644
--- a/service/token_counter.go
+++ b/service/token_counter.go
@@ -9,6 +9,7 @@ import (
 	"log"
 	"math"
 	"one-api/common"
+	"one-api/constant"
 	"one-api/dto"
 	"strings"
 	"unicode/utf8"
@@ -81,13 +82,20 @@ func getTokenNum(tokenEncoder *tiktoken.Tiktoken, text string) int {
 }
 
 func getImageToken(imageUrl *dto.MessageImageUrl, model string, stream bool) (int, error) {
-	// TODO: 非流模式下不计算图片token数量
 	if model == "glm-4v" {
 		return 1047, nil
 	}
 	if imageUrl.Detail == "low" {
 		return 85, nil
 	}
+	// TODO: 非流模式下不计算图片token数量
+	if !constant.GetMediaTokenNotStream && !stream {
+		return 1000, nil
+	}
+	// 是否统计图片token
+	if !constant.GetMediaToken {
+		return 1000, nil
+	}
 	// 同步One API的图片计费逻辑
 	if imageUrl.Detail == "auto" || imageUrl.Detail == "" {
 		imageUrl.Detail = "high"