api2gpt-mid/service/business.go

package service

import (
	"api2gpt-mid/api"
	"api2gpt-mid/common"
	"api2gpt-mid/model"
	"encoding/json"
	"errors"
	"log"
	"strconv"
	"time"
)

// CheckKeyAndTimeCount 检测key是否存在，是否超出每分钟请求次数
func CheckKeyAndTimeCount(key string) (int, error) {
	var timeOut = 60 * time.Second
	var timeCount = 60
	userInfoStr, err := common.RedisGet("user:" + key)
	log.Printf("用户信息 %s", userInfoStr)
	if err != nil {
		//用户不存在
		return 401, errors.New("40003")
	}
	var userInfo model.UserInfo
	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
	if err != nil {
		//用户状态异常
		return 401, errors.New("40004")
	}
	if userInfo.RPM > 0 {
		timeCount = userInfo.RPM
	}
	count, err := common.RedisIncr("user:count:" + key)
	log.Printf("用户请求次数 %d", count)
	if err != nil {
		log.Printf("系统计数器设置异常 %s", err.Error())
		return 500, errors.New("系统计数器设置异常")
	}
	if count == 1 {
		_, err := common.RedisExpire("user:count:"+key, timeOut)
		if err != nil {
			log.Printf("系统计数器异常 %s", err.Error())
			return 500, errors.New("系统计数器异常")
		}
	}
	// 如果请求次数超出限制，则中断请求
	if count > int64(timeCount) {
		//您的账户请求次数过多，超过分钟配额
		return 429, errors.New("42901")
	}
	return 200, nil
}

func QueryBlance(key string) (float64, error) {
	userInfoStr, err := common.RedisGet("user:" + key)
	var userInfo model.UserInfo
	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", 0)
	if err != nil {
		return 0, errors.New("余额计算失败")
	}
	return balance, nil
}

// 余额查询
func CheckBlance(key string, modelStr string, maxTokens int) (model.ServerInfo, error) {
	var serverInfo model.ServerInfo

	//获取模型价格
	modelPriceStr, err := common.RedisGet("model:" + modelStr)
	if err != nil {
		return serverInfo, errors.New("模型信息不存在")
	}
	var modelInfo model.ModelInfo
	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
	if err != nil {
		return serverInfo, errors.New("模型信息解析失败")
	}

	//获取用户信息
	userInfoStr, err := common.RedisGet("user:" + key)
	var userInfo model.UserInfo
	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
	//获取服务器信息
	serverInfoStr, err := common.RedisGet("server:" + strconv.Itoa(modelInfo.ServerId))
	if err != nil {
		return serverInfo, errors.New("服务器信息不存在")
	}
	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
	if err != nil {
		return serverInfo, errors.New("服务器信息解析失败")
	}

	//计算余额-先扣除指定金额
	if maxTokens == 0 {
		maxTokens = modelInfo.ModelPrepayment
	}
	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", -(float64(maxTokens) * modelInfo.ModelPrice))
	if err != nil {
		return serverInfo, errors.New("余额计算失败")
	}
	log.Printf("用户余额 %f key: %v 预扣了：%f", balance, key, (float64(maxTokens) * modelInfo.ModelPrice))
	if balance < 0 {
		_, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", float64(maxTokens)*modelInfo.ModelPrice)
		if err != nil {
			return serverInfo, errors.New("用户缓存出错")
		}
		return serverInfo, errors.New("用户余额不足")
	}

	return serverInfo, nil
}

// 余额查询 for images
func CheckBlanceForImages(key string, modelStr string, n int) (model.ServerInfo, error) {
	var serverInfo model.ServerInfo

	//获取模型价格
	modelPriceStr, err := common.RedisGet("model:" + modelStr)
	if err != nil {
		return serverInfo, errors.New("模型信息不存在")
	}
	var modelInfo model.ModelInfo
	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
	if err != nil {
		return serverInfo, errors.New("模型信息解析失败")
	}

	//获取用户信息
	userInfoStr, err := common.RedisGet("user:" + key)
	var userInfo model.UserInfo
	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
	//获取服务器信息
	serverInfoStr, err := common.RedisGet("server:" + strconv.Itoa(modelInfo.ServerId))
	if err != nil {
		return serverInfo, errors.New("服务器信息不存在")
	}
	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
	if err != nil {
		return serverInfo, errors.New("服务器信息解析失败")
	}

	//计算余额-先扣除指定金额
	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", -(float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice))
	if err != nil {
		return serverInfo, errors.New("余额计算失败")
	}
	log.Printf("用户余额 %f key: %v 预扣了：%f", balance, key, (float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice))
	if balance < 0 {
		_, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", float64(modelInfo.ModelPrepayment*n)*modelInfo.ModelPrice)
		if err != nil {
			return serverInfo, errors.New("用户缓存出错")
		}
		return serverInfo, errors.New("用户余额不足")
	}

	return serverInfo, nil
}

// 预扣返还
func CheckBlanceReturn(key string, modelStr string, maxTokens int) error {
	//获取用户信息
	userInfoStr, err := common.RedisGet("user:" + key)
	var userInfo model.UserInfo
	err = json.Unmarshal([]byte(userInfoStr), &userInfo)

	//获取模型价格
	modelPriceStr, err := common.RedisGet("model:" + modelStr)
	if err != nil {
		return errors.New("模型信息不存在")
	}
	var modelInfo model.ModelInfo
	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
	if err != nil {
		return errors.New("模型信息解析失败")
	}
	if maxTokens == 0 {
		maxTokens = modelInfo.ModelPrepayment
	}
	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", (float64(maxTokens) * modelInfo.ModelPrice))
	log.Printf("用户余额 %f key: %v 返还预扣：%f", balance, key, (float64(maxTokens) * modelInfo.ModelPrice))
	return nil
}

// 预扣返还 for images
func CheckBlanceReturnForImages(key string, modelStr string, n int) error {
	//获取用户信息
	userInfoStr, err := common.RedisGet("user:" + key)
	var userInfo model.UserInfo
	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
	//获取模型价格
	modelPriceStr, err := common.RedisGet("model:" + modelStr)
	if err != nil {
		return errors.New("模型信息不存在")
	}
	var modelInfo model.ModelInfo
	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
	if err != nil {
		return errors.New("模型信息解析失败")
	}
	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", (float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice))
	log.Printf("用户余额 %f key: %v 返还预扣：%f", balance, key, (float64(modelInfo.ModelPrepayment*n) * modelInfo.ModelPrice))
	return nil
}

// 余额消费
func Consumption(key string, modelStr string, prompt_tokens int, completion_tokens int, total_tokens int, msg_id string) (string, error) {
	//获取用户信息
	userInfoStr, err := common.RedisGet("user:" + key)
	if err != nil {
		return "", errors.New("用户信息不存在")
	}
	var userInfo model.UserInfo
	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
	if err != nil {
		return "", errors.New("用户信息解析失败")
	}
	//获取模型价格
	modelPriceStr, err := common.RedisGet("model:" + modelStr)
	if err != nil {
		return "", errors.New("模型信息不存在")
	}
	var modelInfo model.ModelInfo
	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
	if err != nil {
		return "", errors.New("模型信息解析失败")
	}
	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", float64(modelInfo.ModelPrepayment)*modelInfo.ModelPrice-(float64(prompt_tokens)*modelInfo.ModelPrice+float64(completion_tokens)*modelInfo.ModelPrice2))

	// 余额消费日志请求
	result, err := api.BalanceConsumption(key, modelStr, prompt_tokens, completion_tokens, total_tokens, msg_id)
	log.Printf("用户余额：%f 扣费KEY: %s 扣费token数: %d 扣费：%f 扣费日志发送结果 %s", balance, key, total_tokens, float64(modelInfo.ModelPrepayment)*modelInfo.ModelPrice-(float64(prompt_tokens)*modelInfo.ModelPrice+float64(completion_tokens)*modelInfo.ModelPrice2), result)
	if err != nil {
		log.Printf("%s 余额消费日志请求失败 %v", key, err)
		return "", err
	}
	return result, nil
}

// 余额消费 for images
func ConsumptionForImages(key string, modelStr string, n int, dataNum int, msg_id string) (string, error) {
	//获取用户信息
	userInfoStr, err := common.RedisGet("user:" + key)
	if err != nil {
		return "", errors.New("用户信息不存在")
	}
	var userInfo model.UserInfo
	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
	if err != nil {
		return "", errors.New("用户信息解析失败")
	}
	//获取模型价格
	modelPriceStr, err := common.RedisGet("model:" + modelStr)
	if err != nil {
		return "", errors.New("模型信息不存在")
	}
	var modelInfo model.ModelInfo
	err = json.Unmarshal([]byte(modelPriceStr), &modelInfo)
	if err != nil {
		return "", errors.New("模型信息解析失败")
	}
	balance, err := common.RedisIncrByFloat("user:"+userInfo.UID+":balance", float64(modelInfo.ModelPrepayment*n)*modelInfo.ModelPrice-(float64(1000*dataNum)*modelInfo.ModelPrice))

	// 余额消费日志请求
	result, err := api.BalanceConsumption(key, modelStr, 0, 1000*dataNum, 1000*dataNum, msg_id)
	log.Printf("用户余额：%f 扣费KEY: %s 扣费token数: %d 扣费：%f 扣费日志发送结果 %s", balance, key, 1000*dataNum, float64(1000*dataNum)*modelInfo.ModelPrice, result)
	if err != nil {
		log.Printf("%s 余额消费日志请求失败 %v", key, err)
		return "", err
	}
	return result, nil
}