Skip to content

Commit 814961a

Browse files
MartialBEMartialBE
committed
feat: refactor token counting to use strings.Builder for improved performance
1 parent 26039d7 commit 814961a

16 files changed

Lines changed: 68 additions & 44 deletions

File tree

‎common/token.go‎

Lines changed: 29 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -105,17 +105,19 @@ func CountTokenMessages(messages []types.ChatCompletionMessage, model string, pr
105105
tokensPerName = 1
106106
}
107107
tokenNum := 0
108+
var textMsg strings.Builder
109+
108110
for _, message := range messages {
109111
tokenNum += tokensPerMessage
110112
switch v := message.Content.(type) {
111113
case string:
112-
tokenNum += GetTokenNum(tokenEncoder, v)
114+
textMsg.WriteString(v + "\n")
113115
case []any:
114116
for _, it := range v {
115117
m := it.(map[string]any)
116118
switch m["type"] {
117119
case "text":
118-
tokenNum += GetTokenNum(tokenEncoder, m["text"].(string))
120+
textMsg.WriteString(m["text"].(string) + "\n")
119121
case "image_url":
120122
if preCostType == config.PreCostNotImage {
121123
continue
@@ -140,12 +142,18 @@ func CountTokenMessages(messages []types.ChatCompletionMessage, model string, pr
140142
}
141143
}
142144
}
143-
tokenNum += GetTokenNum(tokenEncoder, message.Role)
145+
textMsg.WriteString(message.Role + "\n")
146+
144147
if message.Name != nil {
145148
tokenNum += tokensPerName
146-
tokenNum += GetTokenNum(tokenEncoder, *message.Name)
149+
textMsg.WriteString(*message.Name + "\n")
147150
}
148151
}
152+
153+
if textMsg.Len() > 0 {
154+
tokenNum += GetTokenNum(tokenEncoder, textMsg.String())
155+
}
156+
149157
tokenNum += 3 // Every reply is primed with <|start|>assistant<|message|>
150158
return tokenNum
151159
}
@@ -187,6 +195,7 @@ func CountTokenInputMessages(input any, model string, preCostType int) int {
187195
return 0
188196
}
189197

198+
var textMsg strings.Builder
190199
var messages []types.ChatCompletionMessage
191200
err = json.Unmarshal(jsonStr, &messages)
192201
if err != nil {
@@ -198,13 +207,13 @@ func CountTokenInputMessages(input any, model string, preCostType int) int {
198207
tokenNum += tokensPerMessage
199208
switch v := message.Content.(type) {
200209
case string:
201-
tokenNum += GetTokenNum(tokenEncoder, v)
210+
textMsg.WriteString(v + "\n")
202211
case []any:
203212
for _, it := range v {
204213
m := it.(map[string]any)
205214
switch m["type"] {
206215
case "text":
207-
tokenNum += GetTokenNum(tokenEncoder, m["text"].(string))
216+
textMsg.WriteString(m["text"].(string) + "\n")
208217
case "image_url":
209218
if preCostType == config.PreCostNotImage {
210219
continue
@@ -229,13 +238,17 @@ func CountTokenInputMessages(input any, model string, preCostType int) int {
229238
}
230239
}
231240
}
232-
tokenNum += GetTokenNum(tokenEncoder, message.Role)
241+
textMsg.WriteString(message.Role + "\n")
233242
if message.Name != nil {
234243
tokenNum += tokensPerName
235-
tokenNum += GetTokenNum(tokenEncoder, *message.Name)
244+
textMsg.WriteString(*message.Name + "\n")
236245
}
237246
}
238247

248+
if textMsg.Len() > 0 {
249+
tokenNum += GetTokenNum(tokenEncoder, textMsg.String())
250+
}
251+
239252
tokenNum += 3 // Every reply is primed with <|start|>assistant<|message|>
240253
return tokenNum
241254
}
@@ -247,19 +260,20 @@ func CountTokenRerankMessages(messages types.RerankRequest, model string, preCos
247260

248261
tokenEncoder := GetTokenEncoder(model)
249262
tokenNum := 0
263+
var textMsg strings.Builder
250264

251-
tokenNum += GetTokenNum(tokenEncoder, messages.Query)
265+
textMsg.WriteString(messages.Query + "\n")
252266

253267
for _, document := range messages.Documents {
254268
docStr, ok := document.(string)
255269
if ok {
256-
tokenNum += GetTokenNum(tokenEncoder, docStr)
270+
textMsg.WriteString(docStr + "\n")
257271
} else {
258272
docMultimodal, ok := document.(map[string]string)
259273
if ok {
260274
text := docMultimodal["text"]
261275
if text != "" {
262-
tokenNum += GetTokenNum(tokenEncoder, text)
276+
textMsg.WriteString(text + "\n")
263277
} else {
264278
// 意思意思加点
265279
tokenNum += 10
@@ -268,6 +282,10 @@ func CountTokenRerankMessages(messages types.RerankRequest, model string, preCos
268282
}
269283
}
270284

285+
if textMsg.Len() > 0 {
286+
tokenNum += GetTokenNum(tokenEncoder, textMsg.String())
287+
}
288+
271289
return tokenNum
272290
}
273291

‎providers/claude/chat.go‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -459,9 +459,7 @@ func (h *ClaudeStreamHandler) HandlerStream(rawLine *[]byte, dataChan chan strin
459459

460460
case "content_block_delta":
461461
h.convertToOpenaiStream(&claudeResponse, dataChan)
462-
h.Usage.CompletionTokens += common.CountTokenText(claudeResponse.Delta.Text, h.Request.Model)
463-
h.Usage.TotalTokens = h.Usage.PromptTokens + h.Usage.CompletionTokens
464-
462+
h.Usage.TextBuilder.WriteString(claudeResponse.Delta.Text)
465463
case "content_block_start":
466464
h.convertToOpenaiStream(&claudeResponse, dataChan)
467465

‎providers/claude/common.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"one-api/common"
55
"one-api/types"
66
"strconv"
7+
"strings"
78
)
89

910
func StringErrorWrapper(err string, code string, statusCode int, localError bool) *ClaudeErrorWithStatusCode {
@@ -93,12 +94,13 @@ func ClaudeUsageToOpenaiUsage(cUsage *Usage, usage *types.Usage) bool {
9394
}
9495

9596
func ClaudeOutputUsage(response *ClaudeResponse) int {
96-
text := ""
97+
var textMsg strings.Builder
98+
9799
for _, c := range response.Content {
98100
if c.Type == "text" {
99-
text += c.Text
101+
textMsg.WriteString(c.Text + "\n")
100102
}
101103
}
102104

103-
return common.CountTokenText(text, response.Model)
105+
return common.CountTokenText(textMsg.String(), response.Model)
104106
}

‎providers/claude/relay.go‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"bytes"
55
"encoding/json"
66
"fmt"
7-
"one-api/common"
87
"one-api/common/requester"
98
"one-api/types"
109
"strings"
@@ -119,8 +118,7 @@ func (h *ClaudeRelayStreamHandler) HandlerStream(rawLine *[]byte, dataChan chan
119118
ClaudeUsageMerge(&claudeResponse.Usage, h.StartUsage)
120119
ClaudeUsageToOpenaiUsage(&claudeResponse.Usage, h.Usage)
121120
case "content_block_delta":
122-
h.Usage.CompletionTokens += common.CountTokenText(claudeResponse.Delta.Text, h.ModelName)
123-
h.Usage.TotalTokens = h.Usage.PromptTokens + h.Usage.CompletionTokens
121+
h.Usage.TextBuilder.WriteString(claudeResponse.Delta.Text)
124122
}
125123

126124
dataChan <- rawStr

‎providers/cloudflareAI/chat.go‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -174,9 +174,7 @@ func (h *CloudflareAIStreamHandler) convertToOpenaiStream(chatResponse *ChatResu
174174
choice.FinishReason = types.FinishReasonStop
175175
} else {
176176
choice.Delta.Content = chatResponse.Response
177-
178-
h.Usage.CompletionTokens += common.CountTokenText(chatResponse.Response, h.Request.Model)
179-
h.Usage.TotalTokens = h.Usage.PromptTokens + h.Usage.CompletionTokens
177+
h.Usage.TextBuilder.WriteString(chatResponse.Response)
180178
}
181179

182180
streamResponse.Choices = []types.ChatCompletionStreamChoice{choice}

‎providers/cohere/chat.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -215,8 +215,7 @@ func (h *CohereStreamHandler) convertToOpenaiStream(cohereResponse *ChatStreamRe
215215
choice.Delta.ToolCalls = []*types.ChatCompletionToolCalls{delta.Message.ToolCalls}
216216
}
217217

218-
h.Usage.CompletionTokens += common.CountTokenText(choice.Delta.Content, h.Request.Model)
219-
h.Usage.TotalTokens = h.Usage.PromptTokens + h.Usage.CompletionTokens
218+
h.Usage.TextBuilder.WriteString(choice.Delta.Content)
220219
}
221220

222221
chatCompletion := types.ChatCompletionStreamResponse{

‎providers/coze/chat.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -188,8 +188,7 @@ func (h *CozeStreamHandler) convertToOpenaiStream(chatResponse *CozeStreamRespon
188188
choice.FinishReason = types.FinishReasonStop
189189
} else {
190190
choice.Delta.Content = chatResponse.Message.Content
191-
h.Usage.CompletionTokens += common.CountTokenText(chatResponse.Message.Content, h.Request.Model)
192-
h.Usage.TotalTokens = h.Usage.CompletionTokens + h.Usage.PromptTokens
191+
h.Usage.TextBuilder.WriteString(chatResponse.Message.Content)
193192
}
194193

195194
streamResponse.Choices = []types.ChatCompletionStreamChoice{choice}

‎providers/openai/chat.go‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -162,9 +162,7 @@ func (h *OpenAIStreamHandler) HandlerChatStream(rawLine *[]byte, dataChan chan s
162162
if h.Usage.TotalTokens == 0 {
163163
h.Usage.TotalTokens = h.Usage.PromptTokens
164164
}
165-
countTokenText := common.CountTokenText(openaiResponse.GetResponseText(), h.ModelName)
166-
h.Usage.CompletionTokens += countTokenText
167-
h.Usage.TotalTokens += countTokenText
165+
h.Usage.TextBuilder.WriteString(openaiResponse.GetResponseText())
168166
}
169167
}
170168

‎providers/openai/completion.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ func (h *OpenAIStreamHandler) handlerCompletionStream(rawLine *[]byte, dataChan
117117
if h.Usage.TotalTokens == 0 {
118118
h.Usage.TotalTokens = h.Usage.PromptTokens
119119
}
120-
countTokenText := common.CountTokenText(openaiResponse.getResponseText(), h.ModelName)
121-
h.Usage.CompletionTokens += countTokenText
122-
h.Usage.TotalTokens += countTokenText
120+
121+
h.Usage.TextBuilder.WriteString(openaiResponse.getResponseText())
122+
123123
}

‎providers/openai/responses.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,8 +94,7 @@ func (h *OpenAIResponsesStreamHandler) HandlerChatStream(rawLine *[]byte, dataCh
9494

9595
switch openaiResponse.Type {
9696
case "response.output_text.delta":
97-
h.Usage.CompletionTokens += common.CountTokenText(openaiResponse.Delta, h.Model)
98-
h.Usage.TotalTokens = h.Usage.PromptTokens + h.Usage.CompletionTokens
97+
h.Usage.TextBuilder.WriteString(openaiResponse.Delta)
9998
case "response.completed":
10099
if openaiResponse.Response != nil {
101100
usage := openaiResponse.Response.Usage

0 commit comments

Comments
 (0)