From 46096160fcabd3a625f4e36f9d95f182f12d6c66 Mon Sep 17 00:00:00 2001
From: nanako <469449812@qq.com>
Date: Sun, 9 Nov 2025 00:33:02 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=88=90=E6=9C=AC=E8=AE=A1?=
 =?UTF-8?q?=E7=AE=97=E9=80=BB=E8=BE=91=EF=BC=8C=E8=B0=83=E6=95=B4token?=
 =?UTF-8?q?=E8=AE=A1=E8=B4=B9=E6=96=B9=E5=BC=8F=EF=BC=8C=E5=B0=81=E8=A3=85?=
 =?UTF-8?q?=E6=88=90=E6=9C=AC=E8=AE=A1=E7=AE=97=E5=99=A8=E4=BB=A5=E6=8F=90?=
 =?UTF-8?q?=E9=AB=98=E4=BB=A3=E7=A0=81=E5=8F=AF=E8=AF=BB=E6=80=A7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 backend/internal/billing/calculator.go |  4 ++--
 backend/internal/router/selector.go    | 16 ++++------------
 2 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/backend/internal/billing/calculator.go b/backend/internal/billing/calculator.go
index e6a4988..cad5c2f 100644
--- a/backend/internal/billing/calculator.go
+++ b/backend/internal/billing/calculator.go
@@ -76,11 +76,11 @@ func (cc *CostCalculator) CalculateCost(
 ) float64 {
 	switch billingMethod {
 	case models.BillingMethodToken:
-		return float64(requestTokenCount)*promptTokenPrice + float64(responseTokenCount)*completionTokenPrice
+		cost := float64(requestTokenCount)*promptTokenPrice/1_000_000 + float64(responseTokenCount)*completionTokenPrice/1_000_000
+		return cost
 	case models.BillingMethodRequest:
 		return fixedPrice
 	default:
-		log.Printf("Unknown billing method: %s, defaulting to 0 cost", billingMethod)
 		return 0
 	}
 }
diff --git a/backend/internal/router/selector.go b/backend/internal/router/selector.go
index bda72f6..48ae552 100644
--- a/backend/internal/router/selector.go
+++ b/backend/internal/router/selector.go
@@ -1,6 +1,7 @@
 package router
 
 import (
+	"ai-gateway/internal/billing"
 	"ai-gateway/internal/models"
 	"errors"
 	"sort"
@@ -48,24 +49,15 @@ func SelectBackendModel(db *gorm.DB, virtualModelName string, requestTokenCount
 	})
 
 	// 选择合适的模型（考虑每个后端模型的成本阈值）
-	// 估算响应token数（假设等于请求token数）
-	estimatedResponseTokens := requestTokenCount
-
 	var selectedModel *models.BackendModel
 
 	// 按优先级遍历模型，选择第一个满足成本阈值的模型
+	costCalculator := billing.NewCostCalculator()
 	for i := range suitableModels {
 		model := &suitableModels[i]
 
-		// 计算估算成本
-		var estimatedCost float64
-		switch model.BillingMethod {
-		case models.BillingMethodToken:
-			estimatedCost = float64(requestTokenCount)*model.PromptTokenPrice +
-				float64(estimatedResponseTokens)*model.CompletionTokenPrice
-		case models.BillingMethodRequest:
-			estimatedCost = model.FixedPrice
-		}
+		// 使用封装的计算器计算估算成本
+		estimatedCost := costCalculator.CalculateModelCost(model, requestTokenCount, 0)
 
 		// 如果CostThreshold不为0，则表示设置了成本阈值
 		if model.CostThreshold != 0 {