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"