feat: impl disable billing (#5431)

labring · Mar 2, 2025 · 2d34319 · 2d34319
1 parent 050a861
commit 2d34319
Show file tree

Hide file tree

Showing 8 changed files with 54 additions and 26 deletions.
diff --git a/service/aiproxy/model/main.go b/service/aiproxy/model/main.go
@@ -42,9 +42,13 @@ func chooseDB(envName string) (*gorm.DB, error) {
 		return OpenMySQL(dsn)
 	default:
 		// Use SQLite
-		log.Info("SQL_DSN not set, using SQLite as database: ", common.SQLitePath)
+		absPath, err := filepath.Abs(common.SQLitePath)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get absolute path of SQLite database: %w", err)
+		}
+		log.Info("SQL_DSN not set, using SQLite as database: ", absPath)
 		common.UsingSQLite = true
-		return OpenSQLite(common.SQLitePath)
+		return OpenSQLite(absPath)
 	}
 }
 

diff --git a/service/aiproxy/relay/adaptor/openai/token.go b/service/aiproxy/relay/adaptor/openai/token.go
@@ -7,6 +7,7 @@ import (
 	"sync"
 	"unicode/utf8"
 
+	"github.com/labring/sealos/service/aiproxy/common/config"
 	"github.com/labring/sealos/service/aiproxy/common/image"
 	"github.com/labring/sealos/service/aiproxy/relay/model"
 	"github.com/pkoukk/tiktoken-go"
@@ -56,6 +57,9 @@ func getTokenNum(tokenEncoder *tiktoken.Tiktoken, text string) int {
 }
 
 func CountTokenMessages(messages []*model.Message, model string) int {
+	if !config.GetBillingEnabled() {
+		return 0
+	}
 	tokenEncoder := getTokenEncoder(model)
 	// Reference:
 	// https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
@@ -203,6 +207,9 @@ func countImageTokens(url string, detail string, model string) (_ int, err error
 }
 
 func CountTokenInput(input any, model string) int {
+	if !config.GetBillingEnabled() {
+		return 0
+	}
 	switch v := input.(type) {
 	case string:
 		return CountTokenText(v, model)
@@ -223,12 +230,11 @@ func CountTokenInput(input any, model string) int {
 }
 
 func CountTokenText(text string, model string) int {
+	if !config.GetBillingEnabled() {
+		return 0
+	}
 	if strings.HasPrefix(model, "tts") {
 		return utf8.RuneCountInString(text)
 	}
 	return getTokenNum(getTokenEncoder(model), text)
 }
-
-func CountToken(text string) int {
-	return CountTokenInput(text, "gpt-3.5-turbo")
-}
diff --git a/service/aiproxy/relay/controller/image.go b/service/aiproxy/relay/controller/image.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"github.com/gin-gonic/gin"
+	"github.com/labring/sealos/service/aiproxy/common/config"
 	"github.com/labring/sealos/service/aiproxy/model"
 	"github.com/labring/sealos/service/aiproxy/relay/meta"
 	relaymodel "github.com/labring/sealos/service/aiproxy/relay/model"
@@ -46,6 +47,10 @@ func getImageRequest(meta *meta.Meta, c *gin.Context) (*relaymodel.ImageRequest,
 
 func RelayImageHelper(meta *meta.Meta, c *gin.Context) *relaymodel.ErrorWithStatusCode {
 	return Handle(meta, c, func() (*PreCheckGroupBalanceReq, error) {
+		if !config.GetBillingEnabled() {
+			return &PreCheckGroupBalanceReq{}, nil
+		}
+
 		imageRequest, err := getImageRequest(meta, c)
 		if err != nil {
 			return nil, err

diff --git a/service/aiproxy/relay/controller/price.go b/service/aiproxy/relay/controller/price.go
@@ -1,21 +1,14 @@
 package controller
 
 import (
-	"github.com/labring/sealos/service/aiproxy/common/config"
 	"github.com/labring/sealos/service/aiproxy/model"
 )
 
 func GetModelPrice(modelConfig *model.ModelConfig) (float64, float64, bool) {
-	if !config.GetBillingEnabled() {
-		return 0, 0, true
-	}
 	return modelConfig.InputPrice, modelConfig.OutputPrice, true
 }
 
 func GetImageSizePrice(modelConfig *model.ModelConfig, size string) (float64, bool) {
-	if !config.GetBillingEnabled() {
-		return 0, false
-	}
 	if len(modelConfig.ImagePrices) == 0 {
 		return 0, true
 	}

diff --git a/service/aiproxy/relay/controller/rerank.go b/service/aiproxy/relay/controller/rerank.go
@@ -6,14 +6,19 @@ import (
 	"strings"
 
 	"github.com/gin-gonic/gin"
+	"github.com/labring/sealos/service/aiproxy/common/config"
 	"github.com/labring/sealos/service/aiproxy/relay/meta"
 	relaymodel "github.com/labring/sealos/service/aiproxy/relay/model"
 	"github.com/labring/sealos/service/aiproxy/relay/utils"
 )
 
 func RerankHelper(meta *meta.Meta, c *gin.Context) *relaymodel.ErrorWithStatusCode {
 	return Handle(meta, c, func() (*PreCheckGroupBalanceReq, error) {
-		price, completionPrice, ok := GetModelPrice(meta.ModelConfig)
+		if !config.GetBillingEnabled() {
+			return &PreCheckGroupBalanceReq{}, nil
+		}
+
+		inputPrice, outputPrice, ok := GetModelPrice(meta.ModelConfig)
 		if !ok {
 			return nil, fmt.Errorf("model price not found: %s", meta.OriginModel)
 		}
@@ -25,8 +30,8 @@ func RerankHelper(meta *meta.Meta, c *gin.Context) *relaymodel.ErrorWithStatusCo
 
 		return &PreCheckGroupBalanceReq{
 			InputTokens: rerankPromptTokens(rerankRequest),
-			InputPrice:  price,
-			OutputPrice: completionPrice,
+			InputPrice:  inputPrice,
+			OutputPrice: outputPrice,
 		}, nil
 	})
 }

diff --git a/service/aiproxy/relay/controller/stt.go b/service/aiproxy/relay/controller/stt.go
@@ -9,14 +9,19 @@ import (
 
 	"github.com/gin-gonic/gin"
 	"github.com/labring/sealos/service/aiproxy/common/audio"
+	"github.com/labring/sealos/service/aiproxy/common/config"
 	"github.com/labring/sealos/service/aiproxy/middleware"
 	"github.com/labring/sealos/service/aiproxy/relay/meta"
 	relaymodel "github.com/labring/sealos/service/aiproxy/relay/model"
 )
 
 func RelaySTTHelper(meta *meta.Meta, c *gin.Context) *relaymodel.ErrorWithStatusCode {
 	return Handle(meta, c, func() (*PreCheckGroupBalanceReq, error) {
-		price, completionPrice, ok := GetModelPrice(meta.ModelConfig)
+		if !config.GetBillingEnabled() {
+			return &PreCheckGroupBalanceReq{}, nil
+		}
+
+		inputPrice, outputPrice, ok := GetModelPrice(meta.ModelConfig)
 		if !ok {
 			return nil, fmt.Errorf("model price not found: %s", meta.OriginModel)
 		}
@@ -37,8 +42,8 @@ func RelaySTTHelper(meta *meta.Meta, c *gin.Context) *relaymodel.ErrorWithStatus
 
 		return &PreCheckGroupBalanceReq{
 			InputTokens: durationInt,
-			InputPrice:  price,
-			OutputPrice: completionPrice,
+			InputPrice:  inputPrice,
+			OutputPrice: outputPrice,
 		}, nil
 	})
 }
@@ -75,7 +80,7 @@ func getAudioDuration(audioFile *multipart.FileHeader) (float64, error) {
 }
 
 func getDurationFromTempFile(audioFile *multipart.FileHeader) (float64, error) {
-	tempFile, err := os.CreateTemp("", "audio.wav")
+	tempFile, err := os.CreateTemp("", "audio")
 	if err != nil {
 		return 0, fmt.Errorf("failed to create temp file: %w", err)
 	}

diff --git a/service/aiproxy/relay/controller/text.go b/service/aiproxy/relay/controller/text.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/gin-gonic/gin"
+	"github.com/labring/sealos/service/aiproxy/common/config"
 	"github.com/labring/sealos/service/aiproxy/relay/adaptor/openai"
 	"github.com/labring/sealos/service/aiproxy/relay/meta"
 	relaymodel "github.com/labring/sealos/service/aiproxy/relay/model"
@@ -12,7 +13,11 @@ import (
 
 func RelayTextHelper(meta *meta.Meta, c *gin.Context) *relaymodel.ErrorWithStatusCode {
 	return Handle(meta, c, func() (*PreCheckGroupBalanceReq, error) {
-		price, completionPrice, ok := GetModelPrice(meta.ModelConfig)
+		if !config.GetBillingEnabled() {
+			return &PreCheckGroupBalanceReq{}, nil
+		}
+
+		inputPrice, outputPrice, ok := GetModelPrice(meta.ModelConfig)
 		if !ok {
 			return nil, fmt.Errorf("model price not found: %s", meta.OriginModel)
 		}
@@ -25,8 +30,8 @@ func RelayTextHelper(meta *meta.Meta, c *gin.Context) *relaymodel.ErrorWithStatu
 		return &PreCheckGroupBalanceReq{
 			InputTokens: openai.GetPromptTokens(meta, textRequest),
 			MaxTokens:   textRequest.MaxTokens,
-			InputPrice:  price,
-			OutputPrice: completionPrice,
+			InputPrice:  inputPrice,
+			OutputPrice: outputPrice,
 		}, nil
 	})
 }
diff --git a/service/aiproxy/relay/controller/tts.go b/service/aiproxy/relay/controller/tts.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/gin-gonic/gin"
+	"github.com/labring/sealos/service/aiproxy/common/config"
 	"github.com/labring/sealos/service/aiproxy/relay/adaptor/openai"
 	"github.com/labring/sealos/service/aiproxy/relay/meta"
 	relaymodel "github.com/labring/sealos/service/aiproxy/relay/model"
@@ -12,7 +13,11 @@ import (
 
 func RelayTTSHelper(meta *meta.Meta, c *gin.Context) *relaymodel.ErrorWithStatusCode {
 	return Handle(meta, c, func() (*PreCheckGroupBalanceReq, error) {
-		price, completionPrice, ok := GetModelPrice(meta.ModelConfig)
+		if !config.GetBillingEnabled() {
+			return &PreCheckGroupBalanceReq{}, nil
+		}
+
+		inputPrice, outputPrice, ok := GetModelPrice(meta.ModelConfig)
 		if !ok {
 			return nil, fmt.Errorf("model price not found: %s", meta.OriginModel)
 		}
@@ -24,8 +29,8 @@ func RelayTTSHelper(meta *meta.Meta, c *gin.Context) *relaymodel.ErrorWithStatus
 
 		return &PreCheckGroupBalanceReq{
 			InputTokens: openai.CountTokenText(ttsRequest.Input, meta.ActualModel),
-			InputPrice:  price,
-			OutputPrice: completionPrice,
+			InputPrice:  inputPrice,
+			OutputPrice: outputPrice,
 		}, nil
 	})
 }