fix: fix embedding

constant/sensitive.go

@@ -23,6 +23,7 @@ func SensitiveWordsToString() string {
 }
 
 func SensitiveWordsFromString(s string) {
+	SensitiveWords = []string{}
 	sw := strings.Split(s, "\n")
 	for _, w := range sw {
 		w = strings.TrimSpace(w)

controller/misc.go

@@ -33,6 +33,7 @@ func GetStatus(c *gin.Context) {
 		"success": true,
 		"message": "",
 		"data": gin.H{
+			"version":                  common.Version,
 			"start_time":               common.StartTime,
 			"email_verification":       common.EmailVerificationEnabled,
 			"github_oauth":             common.GitHubOAuthEnabled,

dto/text_response.go

@@ -1,9 +1,18 @@
 package dto
 
-type TextResponse struct {
-	Choices []*OpenAITextResponseChoice `json:"choices"`
+type TextResponseWithError struct {
+	Choices []OpenAITextResponseChoice    `json:"choices"`
+	Object  string                        `json:"object"`
+	Data    []OpenAIEmbeddingResponseItem `json:"data"`
+	Model   string                        `json:"model"`
+	Usage   `json:"usage"`
+	Error   OpenAIError `json:"error"`
+}
+
+type TextResponse struct {
+	Model   string                     `json:"model"`
+	Choices []OpenAITextResponseChoice `json:"choices"`
 	Usage   `json:"usage"`
-	Error   *OpenAIError `json:"error,omitempty"`
 }
 
 type OpenAITextResponseChoice struct {

relay/channel/ollama/adaptor.go

@@ -45,7 +45,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
 		err, responseText = openai.OpenaiStreamHandler(c, resp, info.RelayMode)
 		usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
 	} else {
-		err, usage, sensitiveResp = openai.OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName)
+		err, usage, sensitiveResp = openai.OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName, info.RelayMode)
 	}
 	return
 }

relay/channel/openai/adaptor.go

@@ -77,7 +77,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
 		err, responseText = OpenaiStreamHandler(c, resp, info.RelayMode)
 		usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
 	} else {
-		err, usage, sensitiveResp = OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName)
+		err, usage, sensitiveResp = OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName, info.RelayMode)
 	}
 	return
 }

relay/channel/openai/relay-openai.go

@@ -124,8 +124,8 @@ 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, *dto.SensitiveResponse) {
-	var textResponse dto.TextResponse
+func OpenaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model string, relayMode int) (*dto.OpenAIErrorWithStatusCode, *dto.Usage, *dto.SensitiveResponse) {
+	var responseWithError dto.TextResponseWithError
 	responseBody, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return service.OpenAIErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil, nil
@@ -134,14 +134,14 @@ func OpenaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model
 	if err != nil {
 		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil, nil
 	}
-	err = json.Unmarshal(responseBody, &textResponse)
+	err = json.Unmarshal(responseBody, &responseWithError)
 	if err != nil {
+		log.Printf("unmarshal_response_body_failed: body: %s, err: %v", string(responseBody), err)
 		return service.OpenAIErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil, nil
 	}
-	log.Printf("textResponse: %+v", textResponse)
-	if textResponse.Error != nil {
+	if responseWithError.Error.Type != "" {
 		return &dto.OpenAIErrorWithStatusCode{
-			Error:      *textResponse.Error,
+			Error:      responseWithError.Error,
 			StatusCode: resp.StatusCode,
 		}, nil, nil
 	}
@@ -150,36 +150,65 @@ func OpenaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model
 	sensitiveWords := make([]string, 0)
 	triggerSensitive := false
 
-	if textResponse.Usage.TotalTokens == 0 || checkSensitive {
-		completionTokens := 0
-		for _, choice := range textResponse.Choices {
-			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...)
+	usage := &responseWithError.Usage
+
+	//textResponse := &dto.TextResponse{
+	//	Choices: responseWithError.Choices,
+	//	Usage:   responseWithError.Usage,
+	//}
+	var doResponseBody []byte
+
+	switch relayMode {
+	case relayconstant.RelayModeEmbeddings:
+		embeddingResponse := &dto.OpenAIEmbeddingResponse{
+			Object: responseWithError.Object,
+			Data:   responseWithError.Data,
+			Model:  responseWithError.Model,
+			Usage:  *usage,
+		}
+		doResponseBody, err = json.Marshal(embeddingResponse)
+	default:
+		if responseWithError.Usage.TotalTokens == 0 || checkSensitive {
+			completionTokens := 0
+			for i, choice := range responseWithError.Choices {
+				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)
+						responseWithError.Choices[i].Message = msg
+						sensitiveWords = append(sensitiveWords, words...)
+					}
 				}
 			}
+			responseWithError.Usage = dto.Usage{
+				PromptTokens:     promptTokens,
+				CompletionTokens: completionTokens,
+				TotalTokens:      promptTokens + completionTokens,
+			}
 		}
-		textResponse.Usage = dto.Usage{
-			PromptTokens:     promptTokens,
-			CompletionTokens: completionTokens,
-			TotalTokens:      promptTokens + completionTokens,
+		textResponse := &dto.TextResponse{
+			Choices: responseWithError.Choices,
+			Model:   responseWithError.Model,
+			Usage:   *usage,
 		}
+		doResponseBody, err = json.Marshal(textResponse)
 	}
 
-	if constant.StopOnSensitiveEnabled {
-
+	if checkSensitive && triggerSensitive && constant.StopOnSensitiveEnabled {
+		sensitiveWords = common.RemoveDuplicate(sensitiveWords)
+		return service.OpenAIErrorWrapper(errors.New(fmt.Sprintf("sensitive words detected on response: %s",
+				strings.Join(sensitiveWords, ", "))), "sensitive_words_detected", http.StatusBadRequest),
+			usage, &dto.SensitiveResponse{
+				SensitiveWords: sensitiveWords,
+			}
 	} else {
-		responseBody, err = json.Marshal(textResponse)
 		// Reset response body
-		resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
+		resp.Body = io.NopCloser(bytes.NewBuffer(doResponseBody))
 		// 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.
@@ -197,12 +226,5 @@ func OpenaiHandler(c *gin.Context, resp *http.Response, promptTokens int, model
 			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
+	return nil, usage, nil
 }

relay/channel/perplexity/adaptor.go

@@ -49,7 +49,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
 		err, responseText = openai.OpenaiStreamHandler(c, resp, info.RelayMode)
 		usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
 	} else {
-		err, usage, sensitiveResp = openai.OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName)
+		err, usage, sensitiveResp = openai.OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName, info.RelayMode)
 	}
 	return
 }

relay/channel/zhipu_4v/adaptor.go

@@ -50,7 +50,7 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom
 		err, responseText = openai.OpenaiStreamHandler(c, resp, info.RelayMode)
 		usage, _ = service.ResponseText2Usage(responseText, info.UpstreamModelName, info.PromptTokens)
 	} else {
-		err, usage, sensitiveResp = openai.OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName)
+		err, usage, sensitiveResp = openai.OpenaiHandler(c, resp, info.PromptTokens, info.UpstreamModelName, info.RelayMode)
 	}
 	return
 }

relay/common/relay_utils.go

@@ -35,12 +35,12 @@ func RelayErrorHandler(resp *http.Response) (OpenAIErrorWithStatusCode *dto.Open
 	if err != nil {
 		return
 	}
-	var textResponse dto.TextResponse
+	var textResponse dto.TextResponseWithError
 	err = json.Unmarshal(responseBody, &textResponse)
 	if err != nil {
 		return
 	}
-	OpenAIErrorWithStatusCode.Error = *textResponse.Error
+	OpenAIErrorWithStatusCode.Error = textResponse.Error
 	return
 }
 

service/midjourney.go

@@ -185,7 +185,12 @@ func DoMidjourneyHttpRequest(c *gin.Context, timeout time.Duration, fullRequestU
 	req = req.WithContext(ctx)
 	req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
 	req.Header.Set("Accept", c.Request.Header.Get("Accept"))
-	req.Header.Set("mj-api-secret", strings.Split(c.Request.Header.Get("Authorization"), " ")[1])
+	auth := c.Request.Header.Get("Authorization")
+	if auth != "" {
+		auth = strings.TrimPrefix(auth, "Bearer ")
+		auth = strings.Split(auth, "-")[0]
+		req.Header.Set("mj-api-secret", auth)
+	}
 	defer cancel()
 	resp, err := GetHttpClient().Do(req)
 	if err != nil {

service/sensitive.go

@@ -40,7 +40,7 @@ func SensitiveWordReplace(text string, returnImmediately bool) (bool, []string,
 		for _, hit := range hits {
 			pos := hit.Pos
 			word := string(hit.Word)
-			text = text[:pos] + "*###*" + text[pos+len(word):]
+			text = text[:pos] + "**###**" + text[pos+len(word):]
 			words = append(words, word)
 		}
 		return true, words, text