Merge branch 'Calcium-Ion:main' into main

2025-10-21 23:03:42 +08:00 · 2024-03-23 01:10:39 +08:00
parent 492001a8b2 a70ca53449
commit a825699e9a
11 changed files with 101 additions and 49 deletions
--- a/relay/channel/ollama/adaptor.go
+++ b/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
 }
--- a/relay/channel/openai/adaptor.go
+++ b/relay/channel/openai/adaptor.go
@@ -78,7 +78,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
 }
--- a/relay/channel/openai/relay-openai.go
+++ b/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,43 +150,83 @@ 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{
+			Id:      responseWithError.Id,
+			Created: responseWithError.Created,
+			Object:  responseWithError.Object,
+			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.
 		// For example, Postman will report error, and we cannot check the response at all.
+		// Copy headers
 		for k, v := range resp.Header {
-			c.Writer.Header().Set(k, v[0])
+			// 删除任何现有的相同头部，以防止重复添加头部
+			c.Writer.Header().Del(k)
+			for _, vv := range v {
+				c.Writer.Header().Add(k, vv)
+			}
 		}
+		// reset content length
+		c.Writer.Header().Del("Content-Length")
+		c.Writer.Header().Set("Content-Length", fmt.Sprintf("%d", len(doResponseBody)))
 		c.Writer.WriteHeader(resp.StatusCode)
 		_, err = io.Copy(c.Writer, resp.Body)
 		if err != nil {
@@ -197,12 +237,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
 }
--- a/relay/channel/perplexity/adaptor.go
+++ b/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
 }
--- a/relay/channel/zhipu_4v/adaptor.go
+++ b/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
 }
--- a/relay/common/relay_utils.go
+++ b/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
 }