Просмотр исходного кода

Update relay-text.go

在上下文中存入promptTokens,避免重试过程重复计算
xqx333 1 год назад
Родитель
Сommit
c86762b656
1 измененных файлов с 11 добавлено и 4 удалено
  1. 11 4
      relay/relay-text.go

+ 11 - 4
relay/relay-text.go

@@ -108,10 +108,17 @@ func TextHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) {
 		}
 	}
 
-	promptTokens, err := getPromptTokens(textRequest, relayInfo)
-	// count messages token error 计算promptTokens错误
-	if err != nil {
-		return service.OpenAIErrorWrapper(err, "count_token_messages_failed", http.StatusInternalServerError)
+	// 获取 promptTokens,如果上下文中已经存在,则直接使用
+	var promptTokens int
+	if value, exists := c.Get("prompt_tokens"); exists {
+		promptTokens = value.(int)
+	} else {
+		promptTokens, err = getPromptTokens(textRequest, relayInfo)
+		// count messages token error 计算promptTokens错误
+		if err != nil {
+			return service.OpenAIErrorWrapper(err, "count_token_messages_failed", http.StatusInternalServerError)
+		}
+		c.Set("prompt_tokens", promptTokens)
 	}
 
 	if !getModelPriceSuccess {