feat: 初步兼容生成内容检查

relay/channel/openai/relay-openai.go

@@ -4,8 +4,11 @@ import (
 	"bufio"
 	"bytes"
 	"encoding/json"
+	"errors"
+	"fmt"
 	"github.com/gin-gonic/gin"
 	"io"
+	"log"
 	"net/http"
 	"one-api/common"
 	"one-api/constant"
@@ -18,6 +21,7 @@ import (
 )
 
 func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*dto.OpenAIErrorWithStatusCode, string) {
+	checkSensitive := constant.ShouldCheckCompletionSensitive()
 	var responseTextBuilder strings.Builder
 	scanner := bufio.NewScanner(resp.Body)
 	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
@@ -37,11 +41,10 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*d
 	defer close(stopChan)
 	defer close(dataChan)
 	var wg sync.WaitGroup
-
 	go func() {
 		wg.Add(1)
 		defer wg.Done()
-		var streamItems []string
+		var streamItems []string // store stream items
 		for scanner.Scan() {
 			data := scanner.Text()
 			if len(data) < 6 { // ignore blank line or wrong format
@@ -50,11 +53,20 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*d
 			if data[:6] != "data: " && data[:6] != "[DONE]" {
 				continue
 			}
+			sensitive := false
+			if checkSensitive {
+				// check sensitive
+				sensitive, _, data = service.SensitiveWordReplace(data, constant.StopOnSensitiveEnabled)
+			}
 			dataChan <- data
 			data = data[6:]
 			if !strings.HasPrefix(data, "[DONE]") {
 				streamItems = append(streamItems, data)
 			}
+			if sensitive && constant.StopOnSensitiveEnabled {
+				dataChan <- "data: [DONE]"
+				break
+			}
 		}
 		streamResp := "[" + strings.Join(streamItems, ",") + "]"
 		switch relayMode {
@@ -112,50 +124,48 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*d
 	return nil, responseTextBuilder.String()
 }
 
-func OpenaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model string) (*dto.OpenAIErrorWithStatusCode, *dto.Usage) {
+func OpenaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model string) (*dto.OpenAIErrorWithStatusCode, *dto.Usage, *dto.SensitiveResponse) {
 	var textResponse dto.TextResponse
 	responseBody, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return service.OpenAIErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil
+		return service.OpenAIErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil, nil
 	}
 	err = resp.Body.Close()
 	if err != nil {
-		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
+		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil, nil
 	}
 	err = json.Unmarshal(responseBody, &textResponse)
 	if err != nil {
-		return service.OpenAIErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
+		return service.OpenAIErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil, nil
 	}
-	if textResponse.Error.Type != "" {
+	log.Printf("textResponse: %+v", textResponse)
+	if textResponse.Error != nil {
 		return &dto.OpenAIErrorWithStatusCode{
-			Error:      textResponse.Error,
+			Error:      *textResponse.Error,
 			StatusCode: resp.StatusCode,
-		}, nil
-	}
-	// Reset response body
-	resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
-	// 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.
-	// So the httpClient will be confused by the response.
-	// For example, Postman will report error, and we cannot check the response at all.
-	for k, v := range resp.Header {
-		c.Writer.Header().Set(k, v[0])
-	}
-	c.Writer.WriteHeader(resp.StatusCode)
-	_, err = io.Copy(c.Writer, resp.Body)
-	if err != nil {
-		return service.OpenAIErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError), nil
-	}
-	err = resp.Body.Close()
-	if err != nil {
-		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
+		}, nil, nil
 	}
 
-	if textResponse.Usage.TotalTokens == 0 {
+	checkSensitive := constant.ShouldCheckCompletionSensitive()
+	sensitiveWords := make([]string, 0)
+	triggerSensitive := false
+
+	if textResponse.Usage.TotalTokens == 0 || checkSensitive {
 		completionTokens := 0
 		for _, choice := range textResponse.Choices {
-			ctkm, _ := service.CountTokenText(string(choice.Message.Content), model, constant.ShouldCheckCompletionSensitive())
+			stringContent := string(choice.Message.Content)
+			ctkm, _ := service.CountTokenText(stringContent, model, false)
 			completionTokens += ctkm
+			if checkSensitive {
+				sensitive, words, stringContent := service.SensitiveWordReplace(stringContent, false)
+				if sensitive {
+					triggerSensitive = true
+					msg := choice.Message
+					msg.Content = common.StringToByteSlice(stringContent)
+					choice.Message = msg
+					sensitiveWords = append(sensitiveWords, words...)
+				}
+			}
 		}
 		textResponse.Usage = dto.Usage{
 			PromptTokens:     promptTokens,
@@ -163,5 +173,36 @@ func OpenaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model
 			TotalTokens:      promptTokens + completionTokens,
 		}
 	}
-	return nil, &textResponse.Usage
+
+	if constant.StopOnSensitiveEnabled {
+
+	} else {
+		responseBody, err = json.Marshal(textResponse)
+		// Reset response body
+		resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
+		// 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.
+		// So the httpClient will be confused by the response.
+		// For example, Postman will report error, and we cannot check the response at all.
+		for k, v := range resp.Header {
+			c.Writer.Header().Set(k, v[0])
+		}
+		c.Writer.WriteHeader(resp.StatusCode)
+		_, err = io.Copy(c.Writer, resp.Body)
+		if err != nil {
+			return service.OpenAIErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError), nil, nil
+		}
+		err = resp.Body.Close()
+		if err != nil {
+			return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil, nil
+		}
+	}
+
+	if checkSensitive && triggerSensitive {
+		sensitiveWords = common.RemoveDuplicate(sensitiveWords)
+		return service.OpenAIErrorWrapper(errors.New(fmt.Sprintf("sensitive words detected: %s", strings.Join(sensitiveWords, ", "))), "sensitive_words_detected", http.StatusBadRequest), &textResponse.Usage, &dto.SensitiveResponse{
+			SensitiveWords: sensitiveWords,
+		}
+	}
+	return nil, &textResponse.Usage, nil
 }