Merge branch 'songquanpeng:main' into refactor-main

controller/relay-openai.go

@@ -111,7 +111,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)
@@ -151,5 +151,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
 }

controller/relay-text.go

@@ -302,7 +302,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 		if err != nil {
 			return errorWrapper(err, "close_request_body_failed", http.StatusInternalServerError)
 		}
-		isStream = strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
+		isStream = isStream || strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
 	}
 	if resp.StatusCode != http.StatusOK {
 		return errorWrapper(
@@ -366,7 +366,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
 			}

controller/relay-xunfei.go

@@ -63,16 +63,16 @@ type XunfeiChatResponse struct {
 			Seq    int                          `json:"seq"`
 			Text   []XunfeiChatResponseTextItem `json:"text"`
 		} `json:"choices"`
+		Usage struct {
+			//Text struct {
+			//	QuestionTokens   string `json:"question_tokens"`
+			//	PromptTokens     string `json:"prompt_tokens"`
+			//	CompletionTokens string `json:"completion_tokens"`
+			//	TotalTokens      string `json:"total_tokens"`
+			//} `json:"text"`
+			Text Usage `json:"text"`
+		} `json:"usage"`
 	} `json:"payload"`
-	Usage struct {
-		//Text struct {
-		//	QuestionTokens   string `json:"question_tokens"`
-		//	PromptTokens     string `json:"prompt_tokens"`
-		//	CompletionTokens string `json:"completion_tokens"`
-		//	TotalTokens      string `json:"total_tokens"`
-		//} `json:"text"`
-		Text Usage `json:"text"`
-	} `json:"usage"`
 }
 
 func requestOpenAI2Xunfei(request GeneralOpenAIRequest, xunfeiAppId string) *XunfeiChatRequest {
@@ -123,7 +123,7 @@ func responseXunfei2OpenAI(response *XunfeiChatResponse) *OpenAITextResponse {
 		Object:  "chat.completion",
 		Created: common.GetTimestamp(),
 		Choices: []OpenAITextResponseChoice{choice},
-		Usage:   response.Usage.Text,
+		Usage:   response.Payload.Usage.Text,
 	}
 	return &fullTextResponse
 }
@@ -222,9 +222,9 @@ func xunfeiStreamHandler(c *gin.Context, textRequest GeneralOpenAIRequest, appId
 	c.Stream(func(w io.Writer) bool {
 		select {
 		case xunfeiResponse := <-dataChan:
-			usage.PromptTokens += xunfeiResponse.Usage.Text.PromptTokens
-			usage.CompletionTokens += xunfeiResponse.Usage.Text.CompletionTokens
-			usage.TotalTokens += xunfeiResponse.Usage.Text.TotalTokens
+			usage.PromptTokens += xunfeiResponse.Payload.Usage.Text.PromptTokens
+			usage.CompletionTokens += xunfeiResponse.Payload.Usage.Text.CompletionTokens
+			usage.TotalTokens += xunfeiResponse.Payload.Usage.Text.TotalTokens
 			response := streamResponseXunfei2OpenAI(&xunfeiResponse)
 			jsonResponse, err := json.Marshal(response)
 			if err != nil {

controller/relay.go

@@ -82,8 +82,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 {