diff --git a/controller/relay-openai.go b/controller/relay-openai.go index 8d7b432b..298dbe95 100644 --- a/controller/relay-openai.go +++ b/controller/relay-openai.go @@ -92,7 +92,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O return nil, responseText } -func openaiHandler(c *gin.Context, resp *http.Response, consumeQuota bool) (*OpenAIErrorWithStatusCode, *Usage) { +func openaiHandler(c *gin.Context, resp *http.Response, consumeQuota bool, promptTokens int, model string) (*OpenAIErrorWithStatusCode, *Usage) { var textResponse TextResponse if consumeQuota { responseBody, err := io.ReadAll(resp.Body) @@ -132,5 +132,17 @@ func openaiHandler(c *gin.Context, resp *http.Response, consumeQuota bool) (*Ope if err != nil { return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil } + + if textResponse.Usage.TotalTokens == 0 { + completionTokens := 0 + for _, choice := range textResponse.Choices { + completionTokens += countTokenText(choice.Message.Content, model) + } + textResponse.Usage = Usage{ + PromptTokens: promptTokens, + CompletionTokens: completionTokens, + TotalTokens: promptTokens + completionTokens, + } + } return nil, &textResponse.Usage } diff --git a/controller/relay-text.go b/controller/relay-text.go index ceac4103..df4a6a99 100644 --- a/controller/relay-text.go +++ b/controller/relay-text.go @@ -362,7 +362,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode { textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model) return nil } else { - err, usage := openaiHandler(c, resp, consumeQuota) + err, usage := openaiHandler(c, resp, consumeQuota, promptTokens, textRequest.Model) if err != nil { return err } diff --git a/controller/relay.go b/controller/relay.go index 609ae2eb..dee9c94d 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -81,8 +81,9 @@ type OpenAIErrorWithStatusCode struct { } type TextResponse struct { - Usage `json:"usage"` - Error OpenAIError `json:"error"` + Choices []OpenAITextResponseChoice `json:"choices"` + Usage `json:"usage"` + Error OpenAIError `json:"error"` } type OpenAITextResponseChoice struct {