feat: Implement cache token ratio for more precise token pricing

2025-03-08 01:30:50 +08:00
parent 81137e0533
commit 4f194f4e6a
18 changed files with 258 additions and 71 deletions
--- a/relay/helper/price.go
+++ b/relay/helper/price.go
@@ -6,30 +6,33 @@ import (
 	"one-api/common"
 	relaycommon "one-api/relay/common"
 	"one-api/setting"
+	"one-api/setting/operation_setting"
 )

 type PriceData struct {
 	ModelPrice             float64
 	ModelRatio             float64
 	CompletionRatio        float64
+	CacheRatio             float64
 	GroupRatio             float64
 	UsePrice               bool
 	ShouldPreConsumedQuota int
 }

 func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, maxTokens int) (PriceData, error) {
-	modelPrice, usePrice := setting.GetModelPrice(info.OriginModelName, false)
+	modelPrice, usePrice := operation_setting.GetModelPrice(info.OriginModelName, false)
 	groupRatio := setting.GetGroupRatio(info.Group)
 	var preConsumedQuota int
 	var modelRatio float64
 	var completionRatio float64
+	var cacheRatio float64
 	if !usePrice {
 		preConsumedTokens := common.PreConsumedQuota
 		if maxTokens != 0 {
 			preConsumedTokens = promptTokens + maxTokens
 		}
 		var success bool
-		modelRatio, success = setting.GetModelRatio(info.OriginModelName)
+		modelRatio, success = operation_setting.GetModelRatio(info.OriginModelName)
 		if !success {
 			if info.UserId == 1 {
 				return PriceData{}, fmt.Errorf("模型 %s 倍率或价格未配置，请设置或开始自用模式；Model %s ratio or price not set, please set or start self-use mode", info.OriginModelName, info.OriginModelName)
@@ -37,7 +40,8 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
 				return PriceData{}, fmt.Errorf("模型 %s 倍率或价格未配置, 请联系管理员设置；Model %s ratio or price not set, please contact administrator to set", info.OriginModelName, info.OriginModelName)
 			}
 		}
-		completionRatio = setting.GetCompletionRatio(info.OriginModelName)
+		completionRatio = operation_setting.GetCompletionRatio(info.OriginModelName)
+		cacheRatio, _ = operation_setting.GetCacheRatio(info.OriginModelName)
 		ratio := modelRatio * groupRatio
 		preConsumedQuota = int(float64(preConsumedTokens) * ratio)
 	} else {
@@ -49,6 +53,7 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
 		CompletionRatio:        completionRatio,
 		GroupRatio:             groupRatio,
 		UsePrice:               usePrice,
+		CacheRatio:             cacheRatio,
 		ShouldPreConsumedQuota: preConsumedQuota,
 	}, nil
 }
--- a/relay/relay-mj.go
+++ b/relay/relay-mj.go
@@ -15,6 +15,7 @@ import (
 	relayconstant "one-api/relay/constant"
 	"one-api/service"
 	"one-api/setting"
+	"one-api/setting/operation_setting"
 	"strconv"
 	"strings"
 	"time"
@@ -157,10 +158,10 @@ func RelaySwapFace(c *gin.Context) *dto.MidjourneyResponse {
 		return service.MidjourneyErrorWrapper(constant.MjRequestError, "sour_base64_and_target_base64_is_required")
 	}
 	modelName := service.CoverActionToModelName(constant.MjActionSwapFace)
-	modelPrice, success := setting.GetModelPrice(modelName, true)
+	modelPrice, success := operation_setting.GetModelPrice(modelName, true)
 	// 如果没有配置价格，则使用默认价格
 	if !success {
-		defaultPrice, ok := setting.GetDefaultModelRatioMap()[modelName]
+		defaultPrice, ok := operation_setting.GetDefaultModelRatioMap()[modelName]
 		if !ok {
 			modelPrice = 0.1
 		} else {
@@ -463,10 +464,10 @@ func RelayMidjourneySubmit(c *gin.Context, relayMode int) *dto.MidjourneyRespons
 	fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)

 	modelName := service.CoverActionToModelName(midjRequest.Action)
-	modelPrice, success := setting.GetModelPrice(modelName, true)
+	modelPrice, success := operation_setting.GetModelPrice(modelName, true)
 	// 如果没有配置价格，则使用默认价格
 	if !success {
-		defaultPrice, ok := setting.GetDefaultModelRatioMap()[modelName]
+		defaultPrice, ok := operation_setting.GetDefaultModelRatioMap()[modelName]
 		if !ok {
 			modelPrice = 0.1
 		} else {
--- a/relay/relay-text.go
+++ b/relay/relay-text.go
@@ -110,7 +110,7 @@ func TextHelper(c *gin.Context) (openaiErr *dto.OpenAIErrorWithStatusCode) {
 	if err != nil {
 		return service.OpenAIErrorWrapperLocal(err, "model_price_error", http.StatusInternalServerError)
 	}
-	
+
 	// pre-consume quota 预消耗配额
 	preConsumedQuota, userQuota, openaiErr := preConsumeQuota(c, priceData.ShouldPreConsumedQuota, relayInfo)
 	if openaiErr != nil {
@@ -304,24 +304,26 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
 			CompletionTokens: 0,
 			TotalTokens:      relayInfo.PromptTokens,
 		}
-		extraContent += "  ，（可能是请求出错）"
+		extraContent += "（可能是请求出错）"
 	}
 	useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
 	promptTokens := usage.PromptTokens
+	cacheTokens := usage.PromptTokensDetails.CachedTokens
 	completionTokens := usage.CompletionTokens
 	modelName := relayInfo.OriginModelName

 	tokenName := ctx.GetString("token_name")
-	completionRatio := setting.GetCompletionRatio(modelName)
+	completionRatio := priceData.CompletionRatio
+	cacheRatio := priceData.CacheRatio
 	ratio := priceData.ModelRatio * priceData.GroupRatio
 	modelRatio := priceData.ModelRatio
 	groupRatio := priceData.GroupRatio
 	modelPrice := priceData.ModelPrice
-	usePrice := priceData.UsePrice

 	quota := 0
 	if !priceData.UsePrice {
-		quota = promptTokens + int(math.Round(float64(completionTokens)*completionRatio))
+		quota = (promptTokens - cacheTokens) + int(math.Round(float64(cacheTokens)*cacheRatio))
+		quota += int(math.Round(float64(completionTokens) * completionRatio))
 		quota = int(math.Round(float64(quota) * ratio))
 		if ratio != 0 && quota <= 0 {
 			quota = 1
@@ -330,8 +332,9 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
 		quota = int(modelPrice * common.QuotaPerUnit * groupRatio)
 	}
 	totalTokens := promptTokens + completionTokens
+	
 	var logContent string
-	if !usePrice {
+	if !priceData.UsePrice {
 		logContent = fmt.Sprintf("模型倍率 %.2f，补全倍率 %.2f，分组倍率 %.2f", modelRatio, completionRatio, groupRatio)
 	} else {
 		logContent = fmt.Sprintf("模型价格 %.2f，分组倍率 %.2f", modelPrice, groupRatio)
@@ -372,7 +375,7 @@ func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
 	if extraContent != "" {
 		logContent += ", " + extraContent
 	}
-	other := service.GenerateTextOtherInfo(ctx, relayInfo, modelRatio, groupRatio, completionRatio, modelPrice)
+	other := service.GenerateTextOtherInfo(ctx, relayInfo, modelRatio, groupRatio, completionRatio, cacheTokens, cacheRatio, modelPrice)
 	model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, promptTokens, completionTokens, logModel,
 		tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, relayInfo.Group, other)

--- a/relay/relay_task.go
+++ b/relay/relay_task.go
@@ -16,6 +16,7 @@ import (
 	relayconstant "one-api/relay/constant"
 	"one-api/service"
 	"one-api/setting"
+	"one-api/setting/operation_setting"
 )

 /*
@@ -37,9 +38,9 @@ func RelayTaskSubmit(c *gin.Context, relayMode int) (taskErr *dto.TaskError) {
 	}

 	modelName := service.CoverTaskActionToModelName(platform, relayInfo.Action)
-	modelPrice, success := setting.GetModelPrice(modelName, true)
+	modelPrice, success := operation_setting.GetModelPrice(modelName, true)
 	if !success {
-		defaultPrice, ok := setting.GetDefaultModelRatioMap()[modelName]
+		defaultPrice, ok := operation_setting.GetDefaultModelRatioMap()[modelName]
 		if !ok {
 			modelPrice = 0.1
 		} else {
--- a/relay/websocket.go
+++ b/relay/websocket.go
@@ -11,6 +11,7 @@ import (
 	relaycommon "one-api/relay/common"
 	"one-api/service"
 	"one-api/setting"
+	"one-api/setting/operation_setting"
 )

 func WssHelper(c *gin.Context, ws *websocket.Conn) (openaiErr *dto.OpenAIErrorWithStatusCode) {
@@ -39,7 +40,7 @@ func WssHelper(c *gin.Context, ws *websocket.Conn) (openaiErr *dto.OpenAIErrorWi
 		}
 	}
 	//relayInfo.UpstreamModelName = textRequest.Model
-	modelPrice, getModelPriceSuccess := setting.GetModelPrice(relayInfo.UpstreamModelName, false)
+	modelPrice, getModelPriceSuccess := operation_setting.GetModelPrice(relayInfo.UpstreamModelName, false)
 	groupRatio := setting.GetGroupRatio(relayInfo.Group)

 	var preConsumedQuota int
@@ -65,7 +66,7 @@ func WssHelper(c *gin.Context, ws *websocket.Conn) (openaiErr *dto.OpenAIErrorWi
 		//if realtimeEvent.Session.MaxResponseOutputTokens != 0 {
 		//	preConsumedTokens = promptTokens + int(realtimeEvent.Session.MaxResponseOutputTokens)
 		//}
-		modelRatio, _ = setting.GetModelRatio(relayInfo.UpstreamModelName)
+		modelRatio, _ = operation_setting.GetModelRatio(relayInfo.UpstreamModelName)
 		ratio = modelRatio * groupRatio
 		preConsumedQuota = int(float64(preConsumedTokens) * ratio)
 	} else {