fix: prompt calculation

User will correctly get estimated prompt usage when upstream returns either zero or nothing.
This commit is contained in:
funnycups
2025-08-16 22:54:00 +08:00
committed by GitHub
parent 206ed55db4
commit e3473e3c39
2 changed files with 12 additions and 6 deletions

View File

@@ -197,22 +197,26 @@ func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo
forceFormat = true
}
if simpleResponse.Usage.TotalTokens == 0 || (simpleResponse.Usage.PromptTokens == 0 && simpleResponse.Usage.CompletionTokens == 0) {
completionTokens := 0
for _, choice := range simpleResponse.Choices {
ctkm := service.CountTextToken(choice.Message.StringContent()+choice.Message.ReasoningContent+choice.Message.Reasoning, info.UpstreamModelName)
completionTokens += ctkm
usageModified := false
if simpleResponse.Usage.PromptTokens == 0 {
completionTokens := simpleResponse.Usage.CompletionTokens
if completionTokens == 0 {
for _, choice := range simpleResponse.Choices {
ctkm := service.CountTextToken(choice.Message.StringContent()+choice.Message.ReasoningContent+choice.Message.Reasoning, info.UpstreamModelName)
completionTokens += ctkm
}
}
simpleResponse.Usage = dto.Usage{
PromptTokens: info.PromptTokens,
CompletionTokens: completionTokens,
TotalTokens: info.PromptTokens + completionTokens,
}
usageModified = true
}
switch info.RelayFormat {
case types.RelayFormatOpenAI:
if forceFormat {
if forceFormat || usageModified {
responseBody, err = common.Marshal(simpleResponse)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)