xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-09-18 01:26:37 +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, StatusCode: resp.StatusCode,
}, nil }, nil
} }
// Reset response body // Reset response body
resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) 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. // 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. // 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, CompletionTokens: completionTokens,
TotalTokens: promptTokens + completionTokens, TotalTokens: promptTokens + completionTokens,
} }
} else { } else if textResponse.PromptTokensDetails.AudioTokens+textResponse.CompletionTokensDetails.AudioTokens > 0 {
// Convert the more expensive audio tokens to uniformly priced text tokens // Convert the more expensive audio tokens to uniformly priced text tokens.
textResponse.Usage.PromptTokens = textResponse.CompletionTokensDetails.TextTokens + // 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( int(math.Ceil(
float64(textResponse.CompletionTokensDetails.AudioTokens)* float64(textResponse.PromptTokensDetails.AudioTokens)*
ratio.GetAudioPromptRatio(modelName), ratio.GetAudioPromptRatio(modelName),
)) ))
textResponse.Usage.CompletionTokens = textResponse.CompletionTokensDetails.TextTokens + textResponse.Usage.CompletionTokens = textResponse.CompletionTokensDetails.TextTokens +
int(math.Ceil( int(math.Ceil(
float64(textResponse.CompletionTokensDetails.AudioTokens)* float64(textResponse.CompletionTokensDetails.AudioTokens)*
ratio.GetAudioPromptRatio(modelName)* ratio.GetAudioPromptRatio(modelName)*ratio.GetAudioCompletionRatio(modelName),
ratio.GetAudioCompletionRatio(modelName),
)) ))
textResponse.Usage.TotalTokens = textResponse.Usage.PromptTokens + textResponse.Usage.TotalTokens = textResponse.Usage.PromptTokens +
textResponse.Usage.CompletionTokens textResponse.Usage.CompletionTokens
} }

18
relay/model/misc.go
View File

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