api2gpt-mid/service.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"log"
	"time"
)

// 检测key是否存在，是否超出每分钟请求次数
func checkKeyAndTimeCount(key string) (int, error) {
	var timeOut = 60 * time.Second
	var timeCount = 30
	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
	if err != nil {
		//用户不存在
		return 401, errors.New("40003")
	}
	var userInfo UserInfo
	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
	if err != nil {
		//用户状态异常
		return 401, errors.New("40004")
	}
	count, err := Redis.Incr(context.Background(), "user:count:"+key).Result()
	if err != nil {
		return 500, errors.New("系统计数器设置异常")
	}
	if count == 1 {
		_, err := Redis.Expire(context.Background(), "user:count:"+key, timeOut).Result()
		if err != nil {
			return 500, errors.New("系统计数器异常")
		}
	}
	// 如果请求次数超出限制，则中断请求
	if count > int64(timeCount) {
		//您的账户请求次数过多，超过分钟配额
		return 429, errors.New("42901")
	}
	return 200, nil
}

func queryBlance(key string) (float64, error) {
	userInfoStr, err := Redis.Get(context.Background(), "user:"+key).Result()
	var userInfo UserInfo
	err = json.Unmarshal([]byte(userInfoStr), &userInfo)
	balance, err := Redis.IncrByFloat(context.Background(), "user:"+userInfo.SID+":balance", 0).Result()
	if err != nil {
		return 0, errors.New("余额计算失败")
	}
	return balance, nil
}

// 余额查询
func checkBlance(key string, model string) (ServerInfo, 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 serverInfo, errors.New("服务器信息不存在")
	}
	err = json.Unmarshal([]byte(serverInfoStr), &serverInfo)
	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.SID+":balance", -(float64(modelInfo.ModelPrepayment) * modelInfo.ModelPrice)).Result()
	if err != nil {
		return serverInfo, errors.New("余额计算失败")
	}
	if balance < 0 {
		Redis.IncrByFloat(context.Background(), "user:"+userInfo.SID+":balance", float64(modelInfo.ModelPrepayment)*modelInfo.ModelPrice).Result()
		return serverInfo, errors.New("用户余额不足")
	}

	return serverInfo, nil
}

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

	// 余额消费日志请求
	result, err := balanceConsumption(key, model, prompt_tokens, completion_tokens, total_tokens, msg_id)
	if err != nil {
		return "", err
	}
	log.Printf("扣费KEY: %s 扣费token数: %d 扣费结果 %s", key, total_tokens, result)
	return result, nil
}