Merge remote-tracking branch 'origin/upstream/main'

relay/adaptor/openai/main.go

@@ -12,6 +12,7 @@ import (
 	"github.com/songquanpeng/one-api/common"
 	"github.com/songquanpeng/one-api/common/conv"
 	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/common/render"
 	"github.com/songquanpeng/one-api/relay/model"
 	"github.com/songquanpeng/one-api/relay/relaymode"
 )
@@ -25,88 +26,68 @@ const (
 func StreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*model.ErrorWithStatusCode, string, *model.Usage) {
 	responseText := ""
 	scanner := bufio.NewScanner(resp.Body)
-	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
-		if atEOF && len(data) == 0 {
-			return 0, nil, nil
-		}
-		if i := strings.Index(string(data), "\n"); i >= 0 {
-			return i + 1, data[0:i], nil
-		}
-		if atEOF {
-			return len(data), data, nil
-		}
-		return 0, nil, nil
-	})
-	dataChan := make(chan string)
-	stopChan := make(chan bool)
+	scanner.Split(bufio.ScanLines)
 	var usage *model.Usage
-	go func() {
-		for scanner.Scan() {
-			data := scanner.Text()
-			if len(data) < dataPrefixLength { // ignore blank line or wrong format
-				continue
-			}
-			if data[:dataPrefixLength] != dataPrefix && data[:dataPrefixLength] != done {
-				continue
-			}
-			if strings.HasPrefix(data[dataPrefixLength:], done) {
-				dataChan <- data
-				continue
-			}
-			switch relayMode {
-			case relaymode.ChatCompletions:
-				var streamResponse ChatCompletionsStreamResponse
-				err := json.Unmarshal([]byte(data[dataPrefixLength:]), &streamResponse)
-				if err != nil {
-					logger.SysError("error unmarshalling stream response: " + err.Error())
-					dataChan <- data // if error happened, pass the data to client
-					continue         // just ignore the error
-				}
-				if len(streamResponse.Choices) == 0 {
-					// but for empty choice, we should not pass it to client, this is for azure
-					continue // just ignore empty choice
-				}
-				dataChan <- data
-				for _, choice := range streamResponse.Choices {
-					responseText += conv.AsString(choice.Delta.Content)
-				}
-				if streamResponse.Usage != nil {
-					usage = streamResponse.Usage
-				}
-			case relaymode.Completions:
-				dataChan <- data
-				var streamResponse CompletionsStreamResponse
-				err := json.Unmarshal([]byte(data[dataPrefixLength:]), &streamResponse)
-				if err != nil {
-					logger.SysError("error unmarshalling stream response: " + err.Error())
-					continue
-				}
-				for _, choice := range streamResponse.Choices {
-					responseText += choice.Text
-				}
-			}
-		}
-		stopChan <- true
-	}()
+
 	common.SetEventStreamHeaders(c)
-	c.Stream(func(w io.Writer) bool {
-		select {
-		case data := <-dataChan:
-			if strings.HasPrefix(data, "data: [DONE]") {
-				data = data[:12]
-			}
-			// some implementations may add \r at the end of data
-			data = strings.TrimSuffix(data, "\r")
-			c.Render(-1, common.CustomEvent{Data: data})
-			return true
-		case <-stopChan:
-			return false
+
+	for scanner.Scan() {
+		data := scanner.Text()
+		if len(data) < dataPrefixLength { // ignore blank line or wrong format
+			continue
 		}
-	})
+		if data[:dataPrefixLength] != dataPrefix && data[:dataPrefixLength] != done {
+			continue
+		}
+		if strings.HasPrefix(data[dataPrefixLength:], done) {
+			render.StringData(c, data)
+			continue
+		}
+		switch relayMode {
+		case relaymode.ChatCompletions:
+			var streamResponse ChatCompletionsStreamResponse
+			err := json.Unmarshal([]byte(data[dataPrefixLength:]), &streamResponse)
+			if err != nil {
+				logger.SysError("error unmarshalling stream response: " + err.Error())
+				render.StringData(c, data) // if error happened, pass the data to client
+				continue                   // just ignore the error
+			}
+			if len(streamResponse.Choices) == 0 {
+				// but for empty choice, we should not pass it to client, this is for azure
+				continue // just ignore empty choice
+			}
+			render.StringData(c, data)
+			for _, choice := range streamResponse.Choices {
+				responseText += conv.AsString(choice.Delta.Content)
+			}
+			if streamResponse.Usage != nil {
+				usage = streamResponse.Usage
+			}
+		case relaymode.Completions:
+			render.StringData(c, data)
+			var streamResponse CompletionsStreamResponse
+			err := json.Unmarshal([]byte(data[dataPrefixLength:]), &streamResponse)
+			if err != nil {
+				logger.SysError("error unmarshalling stream response: " + err.Error())
+				continue
+			}
+			for _, choice := range streamResponse.Choices {
+				responseText += choice.Text
+			}
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		logger.SysError("error reading stream: " + err.Error())
+	}
+
+	render.Done(c)
+
 	err := resp.Body.Close()
 	if err != nil {
 		return ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), "", nil
 	}
+
 	return nil, responseText, usage
 }
 
@@ -150,7 +131,7 @@ func Handler(c *gin.Context, resp *http.Response, promptTokens int, modelName st
 		return ErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
 	}
 
-	if textResponse.Usage.TotalTokens == 0 {
+	if textResponse.Usage.TotalTokens == 0 || (textResponse.Usage.PromptTokens == 0 && textResponse.Usage.CompletionTokens == 0) {
 		completionTokens := 0
 		for _, choice := range textResponse.Choices {
 			completionTokens += CountTokenText(choice.Message.StringContent(), modelName)