Browse Source

feat: thinking mode price support (#167)

zijiren 8 months ago
parent
commit
663ee9c937
4 changed files with 23 additions and 4 deletions
  1. 12 4
      core/common/consume/consume.go
  2. 5 0
      core/model/log.go
  3. 3 0
      core/relay/controller/dohelper.go
  4. 3 0
      core/relay/model/chat.go

+ 12 - 4
core/common/consume/consume.go

@@ -127,8 +127,6 @@ func CalculateAmount(
 	}
 
 	inputTokens := usage.InputTokens
-	outputTokens := usage.OutputTokens
-
 	if modelPrice.ImageInputPrice > 0 {
 		inputTokens -= usage.ImageInputTokens
 	}
@@ -139,6 +137,16 @@ func CalculateAmount(
 		inputTokens -= usage.CacheCreationTokens
 	}
 
+	outputTokens := usage.OutputTokens
+	outputPrice := modelPrice.OutputPrice
+	outputPriceUnit := modelPrice.GetOutputPriceUnit()
+	if usage.ReasoningTokens != 0 && modelPrice.ThinkingModeOutputPrice != 0 {
+		outputPrice = modelPrice.ThinkingModeOutputPrice
+		if modelPrice.ThinkingModeOutputPriceUnit != 0 {
+			outputPriceUnit = modelPrice.ThinkingModeOutputPriceUnit
+		}
+	}
+
 	inputAmount := decimal.NewFromInt(int64(inputTokens)).
 		Mul(decimal.NewFromFloat(modelPrice.InputPrice)).
 		Div(decimal.NewFromInt(modelPrice.GetInputPriceUnit()))
@@ -160,8 +168,8 @@ func CalculateAmount(
 		Div(decimal.NewFromInt(modelPrice.GetWebSearchPriceUnit()))
 
 	outputAmount := decimal.NewFromInt(int64(outputTokens)).
-		Mul(decimal.NewFromFloat(modelPrice.OutputPrice)).
-		Div(decimal.NewFromInt(modelPrice.GetOutputPriceUnit()))
+		Mul(decimal.NewFromFloat(outputPrice)).
+		Div(decimal.NewFromInt(outputPriceUnit))
 
 	return inputAmount.
 		Add(imageInputAmount).

+ 5 - 0
core/model/log.go

@@ -50,6 +50,10 @@ type Price struct {
 	OutputPrice     float64 `json:"output_price,omitempty"`
 	OutputPriceUnit int64   `json:"output_price_unit,omitempty"`
 
+	// when ThinkingModeOutputPrice and ReasoningTokens are not 0, OutputPrice and OutputPriceUnit will be overwritten
+	ThinkingModeOutputPrice     float64 `json:"thinking_mode_output_price,omitempty"`
+	ThinkingModeOutputPriceUnit int64   `json:"thinking_mode_output_price_unit,omitempty"`
+
 	CachedPrice     float64 `json:"cached_price,omitempty"`
 	CachedPriceUnit int64   `json:"cached_price_unit,omitempty"`
 
@@ -108,6 +112,7 @@ type Usage struct {
 	OutputTokens        ZeroNullInt64 `json:"output_tokens,omitempty"`
 	CachedTokens        ZeroNullInt64 `json:"cached_tokens,omitempty"`
 	CacheCreationTokens ZeroNullInt64 `json:"cache_creation_tokens,omitempty"`
+	ReasoningTokens     ZeroNullInt64 `json:"reasoning_tokens,omitempty"`
 	TotalTokens         ZeroNullInt64 `json:"total_tokens,omitempty"`
 	WebSearchCount      ZeroNullInt64 `json:"web_search_count,omitempty"`
 }

+ 3 - 0
core/relay/controller/dohelper.go

@@ -274,6 +274,9 @@ func updateUsageMetrics(usage model.Usage, log *log.Entry) {
 	if usage.CacheCreationTokens > 0 {
 		log.Data["t_cache_creation"] = usage.CacheCreationTokens
 	}
+	if usage.ReasoningTokens > 0 {
+		log.Data["t_reason"] = usage.ReasoningTokens
+	}
 	if usage.WebSearchCount > 0 {
 		log.Data["t_websearch"] = usage.WebSearchCount
 	}

+ 3 - 0
core/relay/model/chat.go

@@ -31,6 +31,9 @@ func (u *Usage) ToModelUsage() *model.Usage {
 		usage.CachedTokens = model.ZeroNullInt64(u.PromptTokensDetails.CachedTokens)
 		usage.CacheCreationTokens = model.ZeroNullInt64(u.PromptTokensDetails.CacheCreationTokens)
 	}
+	if u.CompletionTokensDetails != nil {
+		usage.ReasoningTokens = model.ZeroNullInt64(u.CompletionTokensDetails.ReasoningTokens)
+	}
 	return usage
 }