optimize: high cpu

controller/relay-openai.go

@@ -12,7 +12,7 @@ import (
 )
 
 func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*OpenAIErrorWithStatusCode, string) {
-	responseText := ""
+	var responseTextBuilder strings.Builder
 	scanner := bufio.NewScanner(resp.Body)
 	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
 		if atEOF && len(data) == 0 {
@@ -29,6 +29,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 	dataChan := make(chan string)
 	stopChan := make(chan bool)
 	go func() {
+		var streamItems []string
 		for scanner.Scan() {
 			data := scanner.Text()
 			if len(data) < 6 { // ignore blank line or wrong format
@@ -40,27 +41,33 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 			dataChan <- data
 			data = data[6:]
 			if !strings.HasPrefix(data, "[DONE]") {
-				switch relayMode {
-				case RelayModeChatCompletions:
-					var streamResponse ChatCompletionsStreamResponse
-					err := json.Unmarshal([]byte(data), &streamResponse)
-					if err != nil {
-						common.SysError("error unmarshalling stream response: " + err.Error())
-						continue // just ignore the error
-					}
-					for _, choice := range streamResponse.Choices {
-						responseText += choice.Delta.Content
-					}
-				case RelayModeCompletions:
-					var streamResponse CompletionsStreamResponse
-					err := json.Unmarshal([]byte(data), &streamResponse)
-					if err != nil {
-						common.SysError("error unmarshalling stream response: " + err.Error())
-						continue
-					}
-					for _, choice := range streamResponse.Choices {
-						responseText += choice.Text
-					}
+				streamItems = append(streamItems, data)
+			}
+		}
+		streamResp := "[" + strings.Join(streamItems, ",") + "]"
+		switch relayMode {
+		case RelayModeChatCompletions:
+			var streamResponses []ChatCompletionsStreamResponseSimple
+			err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses)
+			if err != nil {
+				common.SysError("error unmarshalling stream response: " + err.Error())
+				return // just ignore the error
+			}
+			for _, streamResponse := range streamResponses {
+				for _, choice := range streamResponse.Choices {
+					responseTextBuilder.WriteString(choice.Delta.Content)
+				}
+			}
+		case RelayModeCompletions:
+			var streamResponses []CompletionsStreamResponse
+			err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses)
+			if err != nil {
+				common.SysError("error unmarshalling stream response: " + err.Error())
+				return // just ignore the error
+			}
+			for _, streamResponse := range streamResponses {
+				for _, choice := range streamResponse.Choices {
+					responseTextBuilder.WriteString(choice.Text)
 				}
 			}
 		}
@@ -85,7 +92,7 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 	if err != nil {
 		return errorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), ""
 	}
-	return nil, responseText
+	return nil, responseTextBuilder.String()
 }
 
 func openaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model string) (*OpenAIErrorWithStatusCode, *Usage) {

controller/relay.go

@@ -189,6 +189,10 @@ type ChatCompletionsStreamResponse struct {
 	Choices []ChatCompletionsStreamResponseChoice `json:"choices"`
 }
 
+type ChatCompletionsStreamResponseSimple struct {
+	Choices []ChatCompletionsStreamResponseChoice `json:"choices"`
+}
+
 type CompletionsStreamResponse struct {
 	Choices []struct {
 		Text         string `json:"text"`