xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-09-17 17:16:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: enhance token usage calculations and improve logging in OpenAI handler

Browse Source
This commit is contained in:
Laisky.Cai 2025-01-14 13:37:00 +00:00
parent 2fc6caaae5
commit ca9aaaf07d
2 changed files with 21 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
relay/adaptor/openai/main.go
View File

@ -118,8 +118,10 @@ func Handler(c *gin.Context, resp *http.Response, promptTokens int, modelName st
StatusCode: resp.StatusCode,
}, nil
}
// Reset response body
resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
logger.Debugf(c.Request.Context(), "handler response: %s", string(responseBody))
// We shouldn't set the header before we parse the response body, because the parse part may fail.
// And then we will have to send an error response, but in this case, the header has already been set.
@ -148,19 +150,21 @@ func Handler(c *gin.Context, resp *http.Response, promptTokens int, modelName st
CompletionTokens: completionTokens,
TotalTokens: promptTokens + completionTokens,
}
} else {
// Convert the more expensive audio tokens to uniformly priced text tokens
textResponse.Usage.PromptTokens = textResponse.CompletionTokensDetails.TextTokens +
} else if textResponse.PromptTokensDetails.AudioTokens+textResponse.CompletionTokensDetails.AudioTokens > 0 {
// Convert the more expensive audio tokens to uniformly priced text tokens.
// Note that when there are no audio tokens in prompt and completion,
// OpenAI will return empty PromptTokensDetails and CompletionTokensDetails, which can be misleading.
textResponse.Usage.PromptTokens = textResponse.PromptTokensDetails.TextTokens +
int(math.Ceil(
float64(textResponse.CompletionTokensDetails.AudioTokens)*
float64(textResponse.PromptTokensDetails.AudioTokens)*
ratio.GetAudioPromptRatio(modelName),
))
textResponse.Usage.CompletionTokens = textResponse.CompletionTokensDetails.TextTokens +
int(math.Ceil(
float64(textResponse.CompletionTokensDetails.AudioTokens)*
ratio.GetAudioPromptRatio(modelName)*
ratio.GetAudioCompletionRatio(modelName),
ratio.GetAudioPromptRatio(modelName)*ratio.GetAudioCompletionRatio(modelName),
))
textResponse.Usage.TotalTokens = textResponse.Usage.PromptTokens +
textResponse.Usage.CompletionTokens
}

18
relay/model/misc.go
View File

@ -1,10 +1,12 @@
package model
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
PromptTokensDetails usagePromptTokensDetails `gorm:"-" json:"prompt_tokens_details,omitempty"`
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
// PromptTokensDetails may be empty for some models
PromptTokensDetails usagePromptTokensDetails `gorm:"-" json:"prompt_tokens_details,omitempty"`
// CompletionTokensDetails may be empty for some models
CompletionTokensDetails usageCompletionTokensDetails `gorm:"-" json:"completion_tokens_details,omitempty"`
ServiceTier string `gorm:"-" json:"service_tier,omitempty"`
SystemFingerprint string `gorm:"-" json:"system_fingerprint,omitempty"`
@ -25,8 +27,9 @@ type ErrorWithStatusCode struct {
type usagePromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"`
AudioTokens int `json:"audio_tokens"`
TextTokens int `json:"text_tokens"`
ImageTokens int `json:"image_tokens"`
// TextTokens could be zero for pure text chats
TextTokens int `json:"text_tokens"`
ImageTokens int `json:"image_tokens"`
}
type usageCompletionTokensDetails struct {
@ -34,5 +37,6 @@ type usageCompletionTokensDetails struct {
AudioTokens int `json:"audio_tokens"`
AcceptedPredictionTokens int `json:"accepted_prediction_tokens"`
RejectedPredictionTokens int `json:"rejected_prediction_tokens"`
TextTokens int `json:"text_tokens"`
// TextTokens could be zero for pure text chats
TextTokens int `json:"text_tokens"`
}