From c86762b65643fa1b229aa48967a75d5309b529ad Mon Sep 17 00:00:00 2001
From: xqx333 <141810964+xqx333@users.noreply.github.com>
Date: Thu, 26 Dec 2024 02:00:04 +0800
Subject: [PATCH] Update relay-text.go
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

在上下文中存入promptTokens,避免重试过程重复计算
---
 relay/relay-text.go | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/relay/relay-text.go b/relay/relay-text.go
index 9a9e6974..86fed5f2 100644
--- a/relay/relay-text.go
+++ b/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 {