update

main.go

@@ -134,6 +134,7 @@ type ModelInfo struct {
 	ModelPrice      float64 `json:"model_price"`
 	ModelPrice2     float64 `json:"model_price2"`
 	ModelPrepayment int     `json:"model_prepayment"`
+	ServerId        int     `json:"server_id"`
 }
 
 var (
@@ -705,39 +706,8 @@ func Cors() gin.HandlerFunc {
 	}
 }
 
-func main() {
+func test_redis() {
-
-	// 禁用控制台颜色，将日志写入文件时不需要控制台颜色。
-	gin.DisableConsoleColor()
-
-	// 记录到文件。
-	filename := time.Now().Format("20060102150405") + ".log"
-	f, _ := os.Create("logs/gin" + filename)
-	//gin.DefaultWriter = io.MultiWriter(f)
-	// 如果需要同时将日志写入文件和控制台，请使用以下代码。
-	gin.DefaultWriter = io.MultiWriter(f, os.Stdout)
-
-	log.SetOutput(gin.DefaultWriter)
-
-	r := gin.Default()
-
-	r.Use(Cors())
-
-	r.GET("/dashboard/billing/credit_grants", checkKeyMid(), balance)
-
-	r.GET("/v1/models", handleGetModels)
-	r.OPTIONS("/v1/*path", handleOptions)
-	r.POST("/v1/chat/completions", checkKeyMid(), completions)
-	r.POST("/v1/completions", checkKeyMid(), completions)
-	r.POST("/v1/embeddings", checkKeyMid(), embeddings)
-	r.POST("/v1/edits", checkKeyMid(), edit)
-	r.POST("/v1/images/generations", checkKeyMid(), images)
-
-	r.POST("/mock1", mockBalanceInquiry)
-	r.POST("/mock2", mockBalanceConsumption)
-
 	Redis = InitRedis()
-
 	//添加reids测试数据
 	//var serverInfo ServerInfo = ServerInfo{
 	//	ServerAddress: "https://gptp.any-door.cn",
@@ -797,24 +767,54 @@ func main() {
 	// modelInfoStr2, _ := json.Marshal(&modelInfo2)
 	// Redis.Set(context.Background(), "model:images-generations", modelInfoStr2, 0)
 
-	// var userInfo UserInfo = UserInfo{
+	var userInfo UserInfo = UserInfo{
-	// 	UID: "1",
+		UID: "0",
-	// 	SID: "1",
+		SID: "1",
-	// }
+	}
 	// var userInfo2 UserInfo = UserInfo{
 	// 	UID: "2",
 	// 	SID: "2",
 	// }
-	// userInfoStr, _ := json.Marshal(&userInfo)
+	userInfoStr, _ := json.Marshal(&userInfo)
 	// userInfoStr2, _ := json.Marshal(&userInfo2)
-	// Redis.Set(context.Background(), "user:8aeb3747-715c-48e8-8b80-aec815949f22", userInfoStr, 0)
+	Redis.Set(context.Background(), "user:key0", userInfoStr, 0)
 	// Redis.Set(context.Background(), "user:AK-7d8ab782-a152-4cc1-9972-568713465c96", userInfoStr2, 0)
 
-	// Redis.IncrByFloat(context.Background(), "user:1:balance", 1000).Result()
+	Redis.IncrByFloat(context.Background(), "user:0:balance", 1000).Result()
 	// Redis.IncrByFloat(context.Background(), "user:2:balance", 1000).Result()
+}
+
+func main() {
+
+	// 禁用控制台颜色，将日志写入文件时不需要控制台颜色。
+	gin.DisableConsoleColor()
+
+	// 记录到文件。
+	filename := time.Now().Format("20060102150405") + ".log"
+	f, _ := os.Create("logs/gin" + filename)
+	//gin.DefaultWriter = io.MultiWriter(f)
+	// 如果需要同时将日志写入文件和控制台，请使用以下代码。
+	gin.DefaultWriter = io.MultiWriter(f, os.Stdout)
 
-	//r.Run("127.0.0.1:8080")
+	log.SetOutput(gin.DefaultWriter)
-	//docker下使用
+
+	r := gin.Default()
+
+	//添加跨域支持
+	r.Use(Cors())
+
+	r.GET("/dashboard/billing/credit_grants", checkKeyMid(), balance)
+
+	r.GET("/v1/models", handleGetModels)
+	r.OPTIONS("/v1/*path", handleOptions)
+	r.POST("/v1/chat/completions", checkKeyMid(), completions)
+	r.POST("/v1/completions", checkKeyMid(), completions)
+	r.POST("/v1/embeddings", checkKeyMid(), embeddings)
+	r.POST("/v1/edits", checkKeyMid(), edit)
+	r.POST("/v1/images/generations", checkKeyMid(), images)
+
+	//r.POST("/mock1", mockBalanceInquiry)
+	//r.POST("/mock2", mockBalanceConsumption)
 
 	// 定义一个GET请求测试接口
 	r.GET("/ping", func(c *gin.Context) {
@@ -823,5 +823,8 @@ func main() {
 		})
 	})
 
+	//添加测试数据
+	//test_redis()
+
 	r.Run("0.0.0.0:8080")
 }

service.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"log"
+	"strconv"
 	"time"
 )
 
@@ -59,12 +60,24 @@ func queryBlance(key string) (float64, error) {
 // 余额查询
 func checkBlance(key string, model string) (ServerInfo, error) {
 	var serverInfo ServerInfo
+
+	//获取模型价格
+	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
+	if err != nil {
+		return serverInfo, errors.New("模型信息不存在")
+	}
+	var modelInfo ModelInfo
+	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
+	if err != nil {
+		return serverInfo, errors.New("模型信息解析失败")
+	}
+
 	//获取用户信息
 	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
 	var userInfo UserInfo
 	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
 	//获取服务器信息
-	serverInfoStr, err := Redis.Get(context.Background(), "server:"+userInfo.SID).Result()
+	serverInfoStr, err := Redis.Get(context.Background(), "server:"+strconv.Itoa(modelInfo.ServerId)).Result()
 	if err != nil {
 		return serverInfo, errors.New("服务器信息不存在")
 	}
@@ -72,16 +85,7 @@ func checkBlance(key string, model string) (ServerInfo, error) {
 	if err != nil {
 		return serverInfo, errors.New("服务器信息解析失败")
 	}
-	//获取模型价格
+
-	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
-	if err != nil {
-		return serverInfo, errors.New("模型信息不存在")
-	}
-	var modelInfo ModelInfo
-	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
-	if err != nil {
-		return serverInfo, errors.New("模型信息解析失败")
-	}
 	//计算余额-先扣除指定金额
 	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", -(float64(modelInfo.ModelPrepayment) * modelInfo.ModelPrice)).Result()
 	if err != nil {
@@ -99,12 +103,24 @@ func checkBlance(key string, model string) (ServerInfo, error) {
 // 余额查询 for images
 func checkBlanceForImages(key string, model string, n int) (ServerInfo, error) {
 	var serverInfo ServerInfo
+
+	//获取模型价格
+	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
+	if err != nil {
+		return serverInfo, errors.New("模型信息不存在")
+	}
+	var modelInfo ModelInfo
+	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
+	if err != nil {
+		return serverInfo, errors.New("模型信息解析失败")
+	}
+
 	//获取用户信息
 	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
 	var userInfo UserInfo
 	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
 	//获取服务器信息
-	serverInfoStr, err := Redis.Get(context.Background(), "server:"+userInfo.SID).Result()
+	serverInfoStr, err := Redis.Get(context.Background(), "server:"+strconv.Itoa(modelInfo.ServerId)).Result()
 	if err != nil {
 		return serverInfo, errors.New("服务器信息不存在")
 	}
@@ -112,16 +128,7 @@ func checkBlanceForImages(key string, model string, n int) (ServerInfo, error) {
 	if err != nil {
 		return serverInfo, errors.New("服务器信息解析失败")
 	}
-	//获取模型价格
+
-	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
-	if err != nil {
-		return serverInfo, errors.New("模型信息不存在")
-	}
-	var modelInfo ModelInfo
-	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
-	if err != nil {
-		return serverInfo, errors.New("模型信息解析失败")
-	}
 	//计算余额-先扣除指定金额
 	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.UID+":balance", -(float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice)).Result()
 	if err != nil {
@@ -138,20 +145,11 @@ func checkBlanceForImages(key string, model string, n int) (ServerInfo, error) {
 
 // 预扣返还
 func checkBlanceReturn(key string, model string) error {
-	var serverInfo ServerInfo
 	//获取用户信息
 	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
 	var userInfo UserInfo
 	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	//获取服务器信息
+
-	serverInfoStr, err := Redis.Get(context.Background(), "server:"+userInfo.SID).Result()
-	if err != nil {
-		return errors.New("服务器信息不存在")
-	}
-	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
-	if err != nil {
-		return errors.New("服务器信息解析失败")
-	}
 	//获取模型价格
 	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
 	if err != nil {
@@ -169,20 +167,10 @@ func checkBlanceReturn(key string, model string) error {
 
 // 预扣返还 for images
 func checkBlanceReturnForImages(key string, model string, n int) error {
-	var serverInfo ServerInfo
 	//获取用户信息
 	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
 	var userInfo UserInfo
 	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
-	//获取服务器信息
-	serverInfoStr, err := Redis.Get(context.Background(), "server:"+userInfo.SID).Result()
-	if err != nil {
-		return errors.New("服务器信息不存在")
-	}
-	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
-	if err != nil {
-		return errors.New("服务器信息解析失败")
-	}
 	//获取模型价格
 	modelPriceStr, err := Redis.Get(context.Background(), "model:"+model).Result()
 	if err != nil {