Skip to content
Toggle navigation
P
Projects
G
Groups
S
Snippets
Help
phsl
/
new-api
This project
Loading...
Sign in
Toggle navigation
Go to a project
Project
Repository
Issues
0
Merge Requests
0
Pipelines
Wiki
Snippets
Members
Activity
Graph
Charts
Create a new issue
Jobs
Commits
Issue Boards
Files
Commits
Branches
Tags
Contributors
Graph
Compare
Charts
Commit
16997a69
authored
Jun 21, 2025
by
CaIon
Browse files
Options
Browse Files
Download
Email Patches
Plain Diff
refactor: token counter logic
parent
4f3024ad
Show whitespace changes
Inline
Side-by-side
Showing
20 changed files
with
57 additions
and
76 deletions
+57
-76
relay/audio_handler.go
+1
-1
relay/channel/claude/relay-claude.go
+3
-3
relay/channel/cloudflare/relay_cloudflare.go
+3
-3
relay/channel/cohere/relay-cohere.go
+1
-1
relay/channel/coze/relay-coze.go
+1
-1
relay/channel/dify/relay-dify.go
+1
-1
relay/channel/gemini/relay-gemini-native.go
+5
-12
relay/channel/openai/relay-openai.go
+3
-3
relay/channel/openai/relay_responses.go
+1
-1
relay/channel/palm/adaptor.go
+1
-1
relay/channel/palm/relay-palm.go
+1
-1
relay/channel/tencent/adaptor.go
+1
-1
relay/channel/xai/text.go
+1
-1
relay/embedding_handler.go
+1
-1
relay/gemini_handler.go
+4
-4
relay/relay-text.go
+3
-3
relay/rerank_handler.go
+2
-4
relay/responses_handler.go
+4
-4
service/token_counter.go
+17
-27
service/usage_helpr.go
+3
-3
No files found.
relay/audio_handler.go
View file @
16997a69
...
@@ -66,7 +66,7 @@ func AudioHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) {
...
@@ -66,7 +66,7 @@ func AudioHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) {
promptTokens
:=
0
promptTokens
:=
0
preConsumedTokens
:=
common
.
PreConsumedQuota
preConsumedTokens
:=
common
.
PreConsumedQuota
if
relayInfo
.
RelayMode
==
relayconstant
.
RelayModeAudioSpeech
{
if
relayInfo
.
RelayMode
==
relayconstant
.
RelayModeAudioSpeech
{
promptTokens
,
err
=
service
.
CountTTSToken
(
audioRequest
.
Input
,
audioRequest
.
Model
)
promptTokens
=
service
.
CountTTSToken
(
audioRequest
.
Input
,
audioRequest
.
Model
)
if
err
!=
nil
{
if
err
!=
nil
{
return
service
.
OpenAIErrorWrapper
(
err
,
"count_audio_token_failed"
,
http
.
StatusInternalServerError
)
return
service
.
OpenAIErrorWrapper
(
err
,
"count_audio_token_failed"
,
http
.
StatusInternalServerError
)
}
}
...
...
relay/channel/claude/relay-claude.go
View file @
16997a69
...
@@ -549,7 +549,7 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
...
@@ -549,7 +549,7 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
func
HandleStreamFinalResponse
(
c
*
gin
.
Context
,
info
*
relaycommon
.
RelayInfo
,
claudeInfo
*
ClaudeResponseInfo
,
requestMode
int
)
{
func
HandleStreamFinalResponse
(
c
*
gin
.
Context
,
info
*
relaycommon
.
RelayInfo
,
claudeInfo
*
ClaudeResponseInfo
,
requestMode
int
)
{
if
requestMode
==
RequestModeCompletion
{
if
requestMode
==
RequestModeCompletion
{
claudeInfo
.
Usage
,
_
=
service
.
ResponseText2Usage
(
claudeInfo
.
ResponseText
.
String
(),
info
.
UpstreamModelName
,
info
.
PromptTokens
)
claudeInfo
.
Usage
=
service
.
ResponseText2Usage
(
claudeInfo
.
ResponseText
.
String
(),
info
.
UpstreamModelName
,
info
.
PromptTokens
)
}
else
{
}
else
{
if
claudeInfo
.
Usage
.
PromptTokens
==
0
{
if
claudeInfo
.
Usage
.
PromptTokens
==
0
{
//上游出错
//上游出错
...
@@ -558,7 +558,7 @@ func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, clau
...
@@ -558,7 +558,7 @@ func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, clau
if
common
.
DebugEnabled
{
if
common
.
DebugEnabled
{
common
.
SysError
(
"claude response usage is not complete, maybe upstream error"
)
common
.
SysError
(
"claude response usage is not complete, maybe upstream error"
)
}
}
claudeInfo
.
Usage
,
_
=
service
.
ResponseText2Usage
(
claudeInfo
.
ResponseText
.
String
(),
info
.
UpstreamModelName
,
claudeInfo
.
Usage
.
PromptTokens
)
claudeInfo
.
Usage
=
service
.
ResponseText2Usage
(
claudeInfo
.
ResponseText
.
String
(),
info
.
UpstreamModelName
,
claudeInfo
.
Usage
.
PromptTokens
)
}
}
}
}
...
@@ -618,7 +618,7 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
...
@@ -618,7 +618,7 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud
}
}
}
}
if
requestMode
==
RequestModeCompletion
{
if
requestMode
==
RequestModeCompletion
{
completionTokens
,
err
:=
service
.
CountTextToken
(
claudeResponse
.
Completion
,
info
.
OriginModelName
)
completionTokens
:=
service
.
CountTextToken
(
claudeResponse
.
Completion
,
info
.
OriginModelName
)
if
err
!=
nil
{
if
err
!=
nil
{
return
service
.
OpenAIErrorWrapper
(
err
,
"count_token_text_failed"
,
http
.
StatusInternalServerError
)
return
service
.
OpenAIErrorWrapper
(
err
,
"count_token_text_failed"
,
http
.
StatusInternalServerError
)
}
}
...
...
relay/channel/cloudflare/relay_cloudflare.go
View file @
16997a69
...
@@ -71,7 +71,7 @@ func cfStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rela
...
@@ -71,7 +71,7 @@ func cfStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rela
if
err
:=
scanner
.
Err
();
err
!=
nil
{
if
err
:=
scanner
.
Err
();
err
!=
nil
{
common
.
LogError
(
c
,
"error_scanning_stream_response: "
+
err
.
Error
())
common
.
LogError
(
c
,
"error_scanning_stream_response: "
+
err
.
Error
())
}
}
usage
,
_
:=
service
.
ResponseText2Usage
(
responseText
,
info
.
UpstreamModelName
,
info
.
PromptTokens
)
usage
:=
service
.
ResponseText2Usage
(
responseText
,
info
.
UpstreamModelName
,
info
.
PromptTokens
)
if
info
.
ShouldIncludeUsage
{
if
info
.
ShouldIncludeUsage
{
response
:=
helper
.
GenerateFinalUsageResponse
(
id
,
info
.
StartTime
.
Unix
(),
info
.
UpstreamModelName
,
*
usage
)
response
:=
helper
.
GenerateFinalUsageResponse
(
id
,
info
.
StartTime
.
Unix
(),
info
.
UpstreamModelName
,
*
usage
)
err
:=
helper
.
ObjectData
(
c
,
response
)
err
:=
helper
.
ObjectData
(
c
,
response
)
...
@@ -108,7 +108,7 @@ func cfHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo)
...
@@ -108,7 +108,7 @@ func cfHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo)
for
_
,
choice
:=
range
response
.
Choices
{
for
_
,
choice
:=
range
response
.
Choices
{
responseText
+=
choice
.
Message
.
StringContent
()
responseText
+=
choice
.
Message
.
StringContent
()
}
}
usage
,
_
:=
service
.
ResponseText2Usage
(
responseText
,
info
.
UpstreamModelName
,
info
.
PromptTokens
)
usage
:=
service
.
ResponseText2Usage
(
responseText
,
info
.
UpstreamModelName
,
info
.
PromptTokens
)
response
.
Usage
=
*
usage
response
.
Usage
=
*
usage
response
.
Id
=
helper
.
GetResponseID
(
c
)
response
.
Id
=
helper
.
GetResponseID
(
c
)
jsonResponse
,
err
:=
json
.
Marshal
(
response
)
jsonResponse
,
err
:=
json
.
Marshal
(
response
)
...
@@ -150,7 +150,7 @@ func cfSTTHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayIn
...
@@ -150,7 +150,7 @@ func cfSTTHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayIn
usage
:=
&
dto
.
Usage
{}
usage
:=
&
dto
.
Usage
{}
usage
.
PromptTokens
=
info
.
PromptTokens
usage
.
PromptTokens
=
info
.
PromptTokens
usage
.
CompletionTokens
,
_
=
service
.
CountTextToken
(
cfResp
.
Result
.
Text
,
info
.
UpstreamModelName
)
usage
.
CompletionTokens
=
service
.
CountTextToken
(
cfResp
.
Result
.
Text
,
info
.
UpstreamModelName
)
usage
.
TotalTokens
=
usage
.
PromptTokens
+
usage
.
CompletionTokens
usage
.
TotalTokens
=
usage
.
PromptTokens
+
usage
.
CompletionTokens
return
nil
,
usage
return
nil
,
usage
...
...
relay/channel/cohere/relay-cohere.go
View file @
16997a69
...
@@ -162,7 +162,7 @@ func cohereStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
...
@@ -162,7 +162,7 @@ func cohereStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
}
}
})
})
if
usage
.
PromptTokens
==
0
{
if
usage
.
PromptTokens
==
0
{
usage
,
_
=
service
.
ResponseText2Usage
(
responseText
,
info
.
UpstreamModelName
,
info
.
PromptTokens
)
usage
=
service
.
ResponseText2Usage
(
responseText
,
info
.
UpstreamModelName
,
info
.
PromptTokens
)
}
}
return
nil
,
usage
return
nil
,
usage
}
}
...
...
relay/channel/coze/relay-coze.go
View file @
16997a69
...
@@ -144,7 +144,7 @@ func cozeChatStreamHandler(c *gin.Context, resp *http.Response, info *relaycommo
...
@@ -144,7 +144,7 @@ func cozeChatStreamHandler(c *gin.Context, resp *http.Response, info *relaycommo
if
usage
.
TotalTokens
==
0
{
if
usage
.
TotalTokens
==
0
{
usage
.
PromptTokens
=
info
.
PromptTokens
usage
.
PromptTokens
=
info
.
PromptTokens
usage
.
CompletionTokens
,
_
=
service
.
CountTextToken
(
"gpt-3.5-turbo"
,
responseText
)
usage
.
CompletionTokens
=
service
.
CountTextToken
(
"gpt-3.5-turbo"
,
responseText
)
usage
.
TotalTokens
=
usage
.
PromptTokens
+
usage
.
CompletionTokens
usage
.
TotalTokens
=
usage
.
PromptTokens
+
usage
.
CompletionTokens
}
}
...
...
relay/channel/dify/relay-dify.go
View file @
16997a69
...
@@ -250,7 +250,7 @@ func difyStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Re
...
@@ -250,7 +250,7 @@ func difyStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Re
}
}
if
usage
.
TotalTokens
==
0
{
if
usage
.
TotalTokens
==
0
{
usage
.
PromptTokens
=
info
.
PromptTokens
usage
.
PromptTokens
=
info
.
PromptTokens
usage
.
CompletionTokens
,
_
=
service
.
CountTextToken
(
"gpt-3.5-turbo"
,
responseText
)
usage
.
CompletionTokens
=
service
.
CountTextToken
(
"gpt-3.5-turbo"
,
responseText
)
usage
.
TotalTokens
=
usage
.
PromptTokens
+
usage
.
CompletionTokens
usage
.
TotalTokens
=
usage
.
PromptTokens
+
usage
.
CompletionTokens
}
}
usage
.
CompletionTokens
+=
nodeToken
usage
.
CompletionTokens
+=
nodeToken
...
...
relay/channel/gemini/relay-gemini-native.go
View file @
16997a69
...
@@ -9,6 +9,7 @@ import (
...
@@ -9,6 +9,7 @@ import (
relaycommon
"one-api/relay/common"
relaycommon
"one-api/relay/common"
"one-api/relay/helper"
"one-api/relay/helper"
"one-api/service"
"one-api/service"
"strings"
"github.com/gin-gonic/gin"
"github.com/gin-gonic/gin"
)
)
...
@@ -75,8 +76,7 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, resp *http.Response, info
...
@@ -75,8 +76,7 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, resp *http.Response, info
helper
.
SetEventStreamHeaders
(
c
)
helper
.
SetEventStreamHeaders
(
c
)
// 本地统计的completion tokens
responseText
:=
strings
.
Builder
{}
localCompletionTokens
:=
0
helper
.
StreamScannerHandler
(
c
,
resp
,
info
,
func
(
data
string
)
bool
{
helper
.
StreamScannerHandler
(
c
,
resp
,
info
,
func
(
data
string
)
bool
{
var
geminiResponse
GeminiChatResponse
var
geminiResponse
GeminiChatResponse
...
@@ -92,12 +92,9 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, resp *http.Response, info
...
@@ -92,12 +92,9 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, resp *http.Response, info
if
part
.
InlineData
!=
nil
&&
part
.
InlineData
.
MimeType
!=
""
{
if
part
.
InlineData
!=
nil
&&
part
.
InlineData
.
MimeType
!=
""
{
imageCount
++
imageCount
++
}
}
// 本地统计completion tokens
if
part
.
Text
!=
""
{
textTokens
,
err
:=
service
.
CountTextToken
(
part
.
Text
,
info
.
UpstreamModelName
)
responseText
.
WriteString
(
part
.
Text
)
if
err
!=
nil
{
common
.
LogError
(
c
,
"error counting text token: "
+
err
.
Error
())
}
}
localCompletionTokens
+=
textTokens
}
}
}
}
...
@@ -133,13 +130,9 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, resp *http.Response, info
...
@@ -133,13 +130,9 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, resp *http.Response, info
// 如果usage.CompletionTokens为0,则使用本地统计的completion tokens
// 如果usage.CompletionTokens为0,则使用本地统计的completion tokens
if
usage
.
CompletionTokens
==
0
{
if
usage
.
CompletionTokens
==
0
{
usage
.
CompletionTokens
=
localCompletionTokens
usage
=
service
.
ResponseText2Usage
(
responseText
.
String
(),
info
.
UpstreamModelName
,
info
.
PromptTokens
)
usage
.
TotalTokens
=
usage
.
PromptTokens
+
usage
.
CompletionTokens
}
}
// 计算最终使用量
// usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
// 移除流式响应结尾的[Done],因为Gemini API没有发送Done的行为
// 移除流式响应结尾的[Done],因为Gemini API没有发送Done的行为
//helper.Done(c)
//helper.Done(c)
...
...
relay/channel/openai/relay-openai.go
View file @
16997a69
...
@@ -8,7 +8,6 @@ import (
...
@@ -8,7 +8,6 @@ import (
"math"
"math"
"mime/multipart"
"mime/multipart"
"net/http"
"net/http"
"path/filepath"
"one-api/common"
"one-api/common"
"one-api/constant"
"one-api/constant"
"one-api/dto"
"one-api/dto"
...
@@ -16,6 +15,7 @@ import (
...
@@ -16,6 +15,7 @@ import (
"one-api/relay/helper"
"one-api/relay/helper"
"one-api/service"
"one-api/service"
"os"
"os"
"path/filepath"
"strings"
"strings"
"github.com/bytedance/gopkg/util/gopool"
"github.com/bytedance/gopkg/util/gopool"
...
@@ -181,7 +181,7 @@ func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
...
@@ -181,7 +181,7 @@ func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
}
}
if
!
containStreamUsage
{
if
!
containStreamUsage
{
usage
,
_
=
service
.
ResponseText2Usage
(
responseTextBuilder
.
String
(),
info
.
UpstreamModelName
,
info
.
PromptTokens
)
usage
=
service
.
ResponseText2Usage
(
responseTextBuilder
.
String
(),
info
.
UpstreamModelName
,
info
.
PromptTokens
)
usage
.
CompletionTokens
+=
toolCount
*
7
usage
.
CompletionTokens
+=
toolCount
*
7
}
else
{
}
else
{
if
info
.
ChannelType
==
common
.
ChannelTypeDeepSeek
{
if
info
.
ChannelType
==
common
.
ChannelTypeDeepSeek
{
...
@@ -225,7 +225,7 @@ func OpenaiHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayI
...
@@ -225,7 +225,7 @@ func OpenaiHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayI
if
simpleResponse
.
Usage
.
TotalTokens
==
0
||
(
simpleResponse
.
Usage
.
PromptTokens
==
0
&&
simpleResponse
.
Usage
.
CompletionTokens
==
0
)
{
if
simpleResponse
.
Usage
.
TotalTokens
==
0
||
(
simpleResponse
.
Usage
.
PromptTokens
==
0
&&
simpleResponse
.
Usage
.
CompletionTokens
==
0
)
{
completionTokens
:=
0
completionTokens
:=
0
for
_
,
choice
:=
range
simpleResponse
.
Choices
{
for
_
,
choice
:=
range
simpleResponse
.
Choices
{
ctkm
,
_
:=
service
.
CountTextToken
(
choice
.
Message
.
StringContent
()
+
choice
.
Message
.
ReasoningContent
+
choice
.
Message
.
Reasoning
,
info
.
UpstreamModelName
)
ctkm
:=
service
.
CountTextToken
(
choice
.
Message
.
StringContent
()
+
choice
.
Message
.
ReasoningContent
+
choice
.
Message
.
Reasoning
,
info
.
UpstreamModelName
)
completionTokens
+=
ctkm
completionTokens
+=
ctkm
}
}
simpleResponse
.
Usage
=
dto
.
Usage
{
simpleResponse
.
Usage
=
dto
.
Usage
{
...
...
relay/channel/openai/relay_responses.go
View file @
16997a69
...
@@ -110,7 +110,7 @@ func OaiResponsesStreamHandler(c *gin.Context, resp *http.Response, info *relayc
...
@@ -110,7 +110,7 @@ func OaiResponsesStreamHandler(c *gin.Context, resp *http.Response, info *relayc
tempStr
:=
responseTextBuilder
.
String
()
tempStr
:=
responseTextBuilder
.
String
()
if
len
(
tempStr
)
>
0
{
if
len
(
tempStr
)
>
0
{
// 非正常结束,使用输出文本的 token 数量
// 非正常结束,使用输出文本的 token 数量
completionTokens
,
_
:=
service
.
CountTextToken
(
tempStr
,
info
.
UpstreamModelName
)
completionTokens
:=
service
.
CountTextToken
(
tempStr
,
info
.
UpstreamModelName
)
usage
.
CompletionTokens
=
completionTokens
usage
.
CompletionTokens
=
completionTokens
}
}
}
}
...
...
relay/channel/palm/adaptor.go
View file @
16997a69
...
@@ -74,7 +74,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
...
@@ -74,7 +74,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
if
info
.
IsStream
{
if
info
.
IsStream
{
var
responseText
string
var
responseText
string
err
,
responseText
=
palmStreamHandler
(
c
,
resp
)
err
,
responseText
=
palmStreamHandler
(
c
,
resp
)
usage
,
_
=
service
.
ResponseText2Usage
(
responseText
,
info
.
UpstreamModelName
,
info
.
PromptTokens
)
usage
=
service
.
ResponseText2Usage
(
responseText
,
info
.
UpstreamModelName
,
info
.
PromptTokens
)
}
else
{
}
else
{
err
,
usage
=
palmHandler
(
c
,
resp
,
info
.
PromptTokens
,
info
.
UpstreamModelName
)
err
,
usage
=
palmHandler
(
c
,
resp
,
info
.
PromptTokens
,
info
.
UpstreamModelName
)
}
}
...
...
relay/channel/palm/relay-palm.go
View file @
16997a69
...
@@ -155,7 +155,7 @@ func palmHandler(c *gin.Context, resp *http.Response, promptTokens int, model st
...
@@ -155,7 +155,7 @@ func palmHandler(c *gin.Context, resp *http.Response, promptTokens int, model st
},
nil
},
nil
}
}
fullTextResponse
:=
responsePaLM2OpenAI
(
&
palmResponse
)
fullTextResponse
:=
responsePaLM2OpenAI
(
&
palmResponse
)
completionTokens
,
_
:=
service
.
CountTextToken
(
palmResponse
.
Candidates
[
0
]
.
Content
,
model
)
completionTokens
:=
service
.
CountTextToken
(
palmResponse
.
Candidates
[
0
]
.
Content
,
model
)
usage
:=
dto
.
Usage
{
usage
:=
dto
.
Usage
{
PromptTokens
:
promptTokens
,
PromptTokens
:
promptTokens
,
CompletionTokens
:
completionTokens
,
CompletionTokens
:
completionTokens
,
...
...
relay/channel/tencent/adaptor.go
View file @
16997a69
...
@@ -98,7 +98,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
...
@@ -98,7 +98,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
if
info
.
IsStream
{
if
info
.
IsStream
{
var
responseText
string
var
responseText
string
err
,
responseText
=
tencentStreamHandler
(
c
,
resp
)
err
,
responseText
=
tencentStreamHandler
(
c
,
resp
)
usage
,
_
=
service
.
ResponseText2Usage
(
responseText
,
info
.
UpstreamModelName
,
info
.
PromptTokens
)
usage
=
service
.
ResponseText2Usage
(
responseText
,
info
.
UpstreamModelName
,
info
.
PromptTokens
)
}
else
{
}
else
{
err
,
usage
=
tencentHandler
(
c
,
resp
)
err
,
usage
=
tencentHandler
(
c
,
resp
)
}
}
...
...
relay/channel/xai/text.go
View file @
16997a69
...
@@ -68,7 +68,7 @@ func xAIStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
...
@@ -68,7 +68,7 @@ func xAIStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
})
})
if
!
containStreamUsage
{
if
!
containStreamUsage
{
usage
,
_
=
service
.
ResponseText2Usage
(
responseTextBuilder
.
String
(),
info
.
UpstreamModelName
,
info
.
PromptTokens
)
usage
=
service
.
ResponseText2Usage
(
responseTextBuilder
.
String
(),
info
.
UpstreamModelName
,
info
.
PromptTokens
)
usage
.
CompletionTokens
+=
toolCount
*
7
usage
.
CompletionTokens
+=
toolCount
*
7
}
}
...
...
relay/embedding_handler.go
View file @
16997a69
...
@@ -15,7 +15,7 @@ import (
...
@@ -15,7 +15,7 @@ import (
)
)
func
getEmbeddingPromptToken
(
embeddingRequest
dto
.
EmbeddingRequest
)
int
{
func
getEmbeddingPromptToken
(
embeddingRequest
dto
.
EmbeddingRequest
)
int
{
token
,
_
:=
service
.
CountTokenInput
(
embeddingRequest
.
Input
,
embeddingRequest
.
Model
)
token
:=
service
.
CountTokenInput
(
embeddingRequest
.
Input
,
embeddingRequest
.
Model
)
return
token
return
token
}
}
...
...
relay/gemini_handler.go
View file @
16997a69
...
@@ -59,7 +59,7 @@ func checkGeminiInputSensitive(textRequest *gemini.GeminiChatRequest) ([]string,
...
@@ -59,7 +59,7 @@ func checkGeminiInputSensitive(textRequest *gemini.GeminiChatRequest) ([]string,
return
sensitiveWords
,
err
return
sensitiveWords
,
err
}
}
func
getGeminiInputTokens
(
req
*
gemini
.
GeminiChatRequest
,
info
*
relaycommon
.
RelayInfo
)
(
int
,
error
)
{
func
getGeminiInputTokens
(
req
*
gemini
.
GeminiChatRequest
,
info
*
relaycommon
.
RelayInfo
)
int
{
// 计算输入 token 数量
// 计算输入 token 数量
var
inputTexts
[]
string
var
inputTexts
[]
string
for
_
,
content
:=
range
req
.
Contents
{
for
_
,
content
:=
range
req
.
Contents
{
...
@@ -71,9 +71,9 @@ func getGeminiInputTokens(req *gemini.GeminiChatRequest, info *relaycommon.Relay
...
@@ -71,9 +71,9 @@ func getGeminiInputTokens(req *gemini.GeminiChatRequest, info *relaycommon.Relay
}
}
inputText
:=
strings
.
Join
(
inputTexts
,
"
\n
"
)
inputText
:=
strings
.
Join
(
inputTexts
,
"
\n
"
)
inputTokens
,
err
:=
service
.
CountTokenInput
(
inputText
,
info
.
UpstreamModelName
)
inputTokens
:=
service
.
CountTokenInput
(
inputText
,
info
.
UpstreamModelName
)
info
.
PromptTokens
=
inputTokens
info
.
PromptTokens
=
inputTokens
return
inputTokens
,
err
return
inputTokens
}
}
func
GeminiHelper
(
c
*
gin
.
Context
)
(
openaiErr
*
dto
.
OpenAIErrorWithStatusCode
)
{
func
GeminiHelper
(
c
*
gin
.
Context
)
(
openaiErr
*
dto
.
OpenAIErrorWithStatusCode
)
{
...
@@ -106,7 +106,7 @@ func GeminiHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) {
...
@@ -106,7 +106,7 @@ func GeminiHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) {
promptTokens
:=
value
.
(
int
)
promptTokens
:=
value
.
(
int
)
relayInfo
.
SetPromptTokens
(
promptTokens
)
relayInfo
.
SetPromptTokens
(
promptTokens
)
}
else
{
}
else
{
promptTokens
,
err
:=
getGeminiInputTokens
(
req
,
relayInfo
)
promptTokens
:=
getGeminiInputTokens
(
req
,
relayInfo
)
if
err
!=
nil
{
if
err
!=
nil
{
return
service
.
OpenAIErrorWrapperLocal
(
err
,
"count_input_tokens_error"
,
http
.
StatusBadRequest
)
return
service
.
OpenAIErrorWrapperLocal
(
err
,
"count_input_tokens_error"
,
http
.
StatusBadRequest
)
}
}
...
...
relay/relay-text.go
View file @
16997a69
...
@@ -251,11 +251,11 @@ func getPromptTokens(textRequest *dto.GeneralOpenAIRequest, info *relaycommon.Re
...
@@ -251,11 +251,11 @@ func getPromptTokens(textRequest *dto.GeneralOpenAIRequest, info *relaycommon.Re
case
relayconstant
.
RelayModeChatCompletions
:
case
relayconstant
.
RelayModeChatCompletions
:
promptTokens
,
err
=
service
.
CountTokenChatRequest
(
info
,
*
textRequest
)
promptTokens
,
err
=
service
.
CountTokenChatRequest
(
info
,
*
textRequest
)
case
relayconstant
.
RelayModeCompletions
:
case
relayconstant
.
RelayModeCompletions
:
promptTokens
,
err
=
service
.
CountTokenInput
(
textRequest
.
Prompt
,
textRequest
.
Model
)
promptTokens
=
service
.
CountTokenInput
(
textRequest
.
Prompt
,
textRequest
.
Model
)
case
relayconstant
.
RelayModeModerations
:
case
relayconstant
.
RelayModeModerations
:
promptTokens
,
err
=
service
.
CountTokenInput
(
textRequest
.
Input
,
textRequest
.
Model
)
promptTokens
=
service
.
CountTokenInput
(
textRequest
.
Input
,
textRequest
.
Model
)
case
relayconstant
.
RelayModeEmbeddings
:
case
relayconstant
.
RelayModeEmbeddings
:
promptTokens
,
err
=
service
.
CountTokenInput
(
textRequest
.
Input
,
textRequest
.
Model
)
promptTokens
=
service
.
CountTokenInput
(
textRequest
.
Input
,
textRequest
.
Model
)
default
:
default
:
err
=
errors
.
New
(
"unknown relay mode"
)
err
=
errors
.
New
(
"unknown relay mode"
)
promptTokens
=
0
promptTokens
=
0
...
...
relay/rerank_handler.go
View file @
16997a69
...
@@ -14,13 +14,11 @@ import (
...
@@ -14,13 +14,11 @@ import (
)
)
func
getRerankPromptToken
(
rerankRequest
dto
.
RerankRequest
)
int
{
func
getRerankPromptToken
(
rerankRequest
dto
.
RerankRequest
)
int
{
token
,
_
:=
service
.
CountTokenInput
(
rerankRequest
.
Query
,
rerankRequest
.
Model
)
token
:=
service
.
CountTokenInput
(
rerankRequest
.
Query
,
rerankRequest
.
Model
)
for
_
,
document
:=
range
rerankRequest
.
Documents
{
for
_
,
document
:=
range
rerankRequest
.
Documents
{
tkm
,
err
:=
service
.
CountTokenInput
(
document
,
rerankRequest
.
Model
)
tkm
:=
service
.
CountTokenInput
(
document
,
rerankRequest
.
Model
)
if
err
==
nil
{
token
+=
tkm
token
+=
tkm
}
}
}
return
token
return
token
}
}
...
...
relay/responses_handler.go
View file @
16997a69
...
@@ -40,10 +40,10 @@ func checkInputSensitive(textRequest *dto.OpenAIResponsesRequest, info *relaycom
...
@@ -40,10 +40,10 @@ func checkInputSensitive(textRequest *dto.OpenAIResponsesRequest, info *relaycom
return
sensitiveWords
,
err
return
sensitiveWords
,
err
}
}
func
getInputTokens
(
req
*
dto
.
OpenAIResponsesRequest
,
info
*
relaycommon
.
RelayInfo
)
(
int
,
error
)
{
func
getInputTokens
(
req
*
dto
.
OpenAIResponsesRequest
,
info
*
relaycommon
.
RelayInfo
)
int
{
inputTokens
,
err
:=
service
.
CountTokenInput
(
req
.
Input
,
req
.
Model
)
inputTokens
:=
service
.
CountTokenInput
(
req
.
Input
,
req
.
Model
)
info
.
PromptTokens
=
inputTokens
info
.
PromptTokens
=
inputTokens
return
inputTokens
,
err
return
inputTokens
}
}
func
ResponsesHelper
(
c
*
gin
.
Context
)
(
openaiErr
*
dto
.
OpenAIErrorWithStatusCode
)
{
func
ResponsesHelper
(
c
*
gin
.
Context
)
(
openaiErr
*
dto
.
OpenAIErrorWithStatusCode
)
{
...
@@ -72,7 +72,7 @@ func ResponsesHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode)
...
@@ -72,7 +72,7 @@ func ResponsesHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode)
promptTokens
:=
value
.
(
int
)
promptTokens
:=
value
.
(
int
)
relayInfo
.
SetPromptTokens
(
promptTokens
)
relayInfo
.
SetPromptTokens
(
promptTokens
)
}
else
{
}
else
{
promptTokens
,
err
:=
getInputTokens
(
req
,
relayInfo
)
promptTokens
:=
getInputTokens
(
req
,
relayInfo
)
if
err
!=
nil
{
if
err
!=
nil
{
return
service
.
OpenAIErrorWrapperLocal
(
err
,
"count_input_tokens_error"
,
http
.
StatusBadRequest
)
return
service
.
OpenAIErrorWrapperLocal
(
err
,
"count_input_tokens_error"
,
http
.
StatusBadRequest
)
}
}
...
...
service/token_counter.go
View file @
16997a69
...
@@ -171,7 +171,7 @@ func CountTokenChatRequest(info *relaycommon.RelayInfo, request dto.GeneralOpenA
...
@@ -171,7 +171,7 @@ func CountTokenChatRequest(info *relaycommon.RelayInfo, request dto.GeneralOpenA
countStr
+=
fmt
.
Sprintf
(
"%v"
,
tool
.
Function
.
Parameters
)
countStr
+=
fmt
.
Sprintf
(
"%v"
,
tool
.
Function
.
Parameters
)
}
}
}
}
toolTokens
,
err
:=
CountTokenInput
(
countStr
,
request
.
Model
)
toolTokens
:=
CountTokenInput
(
countStr
,
request
.
Model
)
if
err
!=
nil
{
if
err
!=
nil
{
return
0
,
err
return
0
,
err
}
}
...
@@ -194,7 +194,7 @@ func CountTokenClaudeRequest(request dto.ClaudeRequest, model string) (int, erro
...
@@ -194,7 +194,7 @@ func CountTokenClaudeRequest(request dto.ClaudeRequest, model string) (int, erro
// Count tokens in system message
// Count tokens in system message
if
request
.
System
!=
""
{
if
request
.
System
!=
""
{
systemTokens
,
err
:=
CountTokenInput
(
request
.
System
,
model
)
systemTokens
:=
CountTokenInput
(
request
.
System
,
model
)
if
err
!=
nil
{
if
err
!=
nil
{
return
0
,
err
return
0
,
err
}
}
...
@@ -296,10 +296,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
...
@@ -296,10 +296,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
switch
request
.
Type
{
switch
request
.
Type
{
case
dto
.
RealtimeEventTypeSessionUpdate
:
case
dto
.
RealtimeEventTypeSessionUpdate
:
if
request
.
Session
!=
nil
{
if
request
.
Session
!=
nil
{
msgTokens
,
err
:=
CountTextToken
(
request
.
Session
.
Instructions
,
model
)
msgTokens
:=
CountTextToken
(
request
.
Session
.
Instructions
,
model
)
if
err
!=
nil
{
return
0
,
0
,
err
}
textToken
+=
msgTokens
textToken
+=
msgTokens
}
}
case
dto
.
RealtimeEventResponseAudioDelta
:
case
dto
.
RealtimeEventResponseAudioDelta
:
...
@@ -311,10 +308,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
...
@@ -311,10 +308,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
audioToken
+=
atk
audioToken
+=
atk
case
dto
.
RealtimeEventResponseAudioTranscriptionDelta
,
dto
.
RealtimeEventResponseFunctionCallArgumentsDelta
:
case
dto
.
RealtimeEventResponseAudioTranscriptionDelta
,
dto
.
RealtimeEventResponseFunctionCallArgumentsDelta
:
// count text token
// count text token
tkm
,
err
:=
CountTextToken
(
request
.
Delta
,
model
)
tkm
:=
CountTextToken
(
request
.
Delta
,
model
)
if
err
!=
nil
{
return
0
,
0
,
fmt
.
Errorf
(
"error counting text token: %v"
,
err
)
}
textToken
+=
tkm
textToken
+=
tkm
case
dto
.
RealtimeEventInputAudioBufferAppend
:
case
dto
.
RealtimeEventInputAudioBufferAppend
:
// count audio token
// count audio token
...
@@ -329,10 +323,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
...
@@ -329,10 +323,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
case
"message"
:
case
"message"
:
for
_
,
content
:=
range
request
.
Item
.
Content
{
for
_
,
content
:=
range
request
.
Item
.
Content
{
if
content
.
Type
==
"input_text"
{
if
content
.
Type
==
"input_text"
{
tokens
,
err
:=
CountTextToken
(
content
.
Text
,
model
)
tokens
:=
CountTextToken
(
content
.
Text
,
model
)
if
err
!=
nil
{
return
0
,
0
,
err
}
textToken
+=
tokens
textToken
+=
tokens
}
}
}
}
...
@@ -343,10 +334,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
...
@@ -343,10 +334,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
if
!
info
.
IsFirstRequest
{
if
!
info
.
IsFirstRequest
{
if
info
.
RealtimeTools
!=
nil
&&
len
(
info
.
RealtimeTools
)
>
0
{
if
info
.
RealtimeTools
!=
nil
&&
len
(
info
.
RealtimeTools
)
>
0
{
for
_
,
tool
:=
range
info
.
RealtimeTools
{
for
_
,
tool
:=
range
info
.
RealtimeTools
{
toolTokens
,
err
:=
CountTokenInput
(
tool
,
model
)
toolTokens
:=
CountTokenInput
(
tool
,
model
)
if
err
!=
nil
{
return
0
,
0
,
err
}
textToken
+=
8
textToken
+=
8
textToken
+=
toolTokens
textToken
+=
toolTokens
}
}
...
@@ -409,7 +397,7 @@ func CountTokenMessages(info *relaycommon.RelayInfo, messages []dto.Message, mod
...
@@ -409,7 +397,7 @@ func CountTokenMessages(info *relaycommon.RelayInfo, messages []dto.Message, mod
return
tokenNum
,
nil
return
tokenNum
,
nil
}
}
func
CountTokenInput
(
input
any
,
model
string
)
(
int
,
error
)
{
func
CountTokenInput
(
input
any
,
model
string
)
int
{
switch
v
:=
input
.
(
type
)
{
switch
v
:=
input
.
(
type
)
{
case
string
:
case
string
:
return
CountTextToken
(
v
,
model
)
return
CountTextToken
(
v
,
model
)
...
@@ -432,13 +420,13 @@ func CountTokenInput(input any, model string) (int, error) {
...
@@ -432,13 +420,13 @@ func CountTokenInput(input any, model string) (int, error) {
func
CountTokenStreamChoices
(
messages
[]
dto
.
ChatCompletionsStreamResponseChoice
,
model
string
)
int
{
func
CountTokenStreamChoices
(
messages
[]
dto
.
ChatCompletionsStreamResponseChoice
,
model
string
)
int
{
tokens
:=
0
tokens
:=
0
for
_
,
message
:=
range
messages
{
for
_
,
message
:=
range
messages
{
tkm
,
_
:=
CountTokenInput
(
message
.
Delta
.
GetContentString
(),
model
)
tkm
:=
CountTokenInput
(
message
.
Delta
.
GetContentString
(),
model
)
tokens
+=
tkm
tokens
+=
tkm
if
message
.
Delta
.
ToolCalls
!=
nil
{
if
message
.
Delta
.
ToolCalls
!=
nil
{
for
_
,
tool
:=
range
message
.
Delta
.
ToolCalls
{
for
_
,
tool
:=
range
message
.
Delta
.
ToolCalls
{
tkm
,
_
:=
CountTokenInput
(
tool
.
Function
.
Name
,
model
)
tkm
:=
CountTokenInput
(
tool
.
Function
.
Name
,
model
)
tokens
+=
tkm
tokens
+=
tkm
tkm
,
_
=
CountTokenInput
(
tool
.
Function
.
Arguments
,
model
)
tkm
=
CountTokenInput
(
tool
.
Function
.
Arguments
,
model
)
tokens
+=
tkm
tokens
+=
tkm
}
}
}
}
...
@@ -446,9 +434,9 @@ func CountTokenStreamChoices(messages []dto.ChatCompletionsStreamResponseChoice,
...
@@ -446,9 +434,9 @@ func CountTokenStreamChoices(messages []dto.ChatCompletionsStreamResponseChoice,
return
tokens
return
tokens
}
}
func
CountTTSToken
(
text
string
,
model
string
)
(
int
,
error
)
{
func
CountTTSToken
(
text
string
,
model
string
)
int
{
if
strings
.
HasPrefix
(
model
,
"tts"
)
{
if
strings
.
HasPrefix
(
model
,
"tts"
)
{
return
utf8
.
RuneCountInString
(
text
)
,
nil
return
utf8
.
RuneCountInString
(
text
)
}
else
{
}
else
{
return
CountTextToken
(
text
,
model
)
return
CountTextToken
(
text
,
model
)
}
}
...
@@ -483,8 +471,10 @@ func CountAudioTokenOutput(audioBase64 string, audioFormat string) (int, error)
...
@@ -483,8 +471,10 @@ func CountAudioTokenOutput(audioBase64 string, audioFormat string) (int, error)
//}
//}
// CountTextToken 统计文本的token数量,仅当文本包含敏感词,返回错误,同时返回token数量
// CountTextToken 统计文本的token数量,仅当文本包含敏感词,返回错误,同时返回token数量
func
CountTextToken
(
text
string
,
model
string
)
(
int
,
error
)
{
func
CountTextToken
(
text
string
,
model
string
)
int
{
var
err
error
if
text
==
""
{
return
0
}
tokenEncoder
:=
getTokenEncoder
(
model
)
tokenEncoder
:=
getTokenEncoder
(
model
)
return
getTokenNum
(
tokenEncoder
,
text
)
,
err
return
getTokenNum
(
tokenEncoder
,
text
)
}
}
service/usage_helpr.go
View file @
16997a69
...
@@ -16,13 +16,13 @@ import (
...
@@ -16,13 +16,13 @@ import (
// return 0, errors.New("unknown relay mode")
// return 0, errors.New("unknown relay mode")
//}
//}
func
ResponseText2Usage
(
responseText
string
,
modeName
string
,
promptTokens
int
)
(
*
dto
.
Usage
,
error
)
{
func
ResponseText2Usage
(
responseText
string
,
modeName
string
,
promptTokens
int
)
*
dto
.
Usage
{
usage
:=
&
dto
.
Usage
{}
usage
:=
&
dto
.
Usage
{}
usage
.
PromptTokens
=
promptTokens
usage
.
PromptTokens
=
promptTokens
ctkm
,
err
:=
CountTextToken
(
responseText
,
modeName
)
ctkm
:=
CountTextToken
(
responseText
,
modeName
)
usage
.
CompletionTokens
=
ctkm
usage
.
CompletionTokens
=
ctkm
usage
.
TotalTokens
=
usage
.
PromptTokens
+
usage
.
CompletionTokens
usage
.
TotalTokens
=
usage
.
PromptTokens
+
usage
.
CompletionTokens
return
usage
,
err
return
usage
}
}
func
ValidUsage
(
usage
*
dto
.
Usage
)
bool
{
func
ValidUsage
(
usage
*
dto
.
Usage
)
bool
{
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment