merge upstream

Signed-off-by: wozulong <>

model/ability.go

@@ -29,6 +29,13 @@ func GetGroupModels(group string) []string {
 	return models
 }
 
+func GetEnabledModels() []string {
+	var models []string
+	// Find distinct models
+	DB.Table("abilities").Where("enabled = ?", true).Distinct("model").Pluck("model", &models)
+	return models
+}
+
 func getPriority(group string, model string, retry int) (int, error) {
 	groupCol := "`group`"
 	trueVal := "1"

model/log.go

@@ -24,6 +24,7 @@ type Log struct {
 	IsStream         bool   `json:"is_stream" gorm:"default:false"`
 	ChannelId        int    `json:"channel" gorm:"index"`
 	TokenId          int    `json:"token_id" gorm:"default:0;index"`
+	Other            string `json:"other"`
 }
 
 const (
@@ -57,12 +58,13 @@ func RecordLog(userId int, logType int, content string) {
 	}
 }
 
-func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int, content string, tokenId int, userQuota int, useTimeSeconds int, isStream bool) {
+func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int, content string, tokenId int, userQuota int, useTimeSeconds int, isStream bool, other map[string]interface{}) {
 	common.LogInfo(ctx, fmt.Sprintf("record consume log: userId=%d, 用户调用前余额=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, userQuota, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content))
 	if !common.LogConsumeEnabled {
 		return
 	}
 	username, _ := CacheGetUsername(userId)
+	otherStr := common.MapToJsonStr(other)
 	log := &Log{
 		UserId:           userId,
 		Username:         username,
@@ -78,6 +80,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
 		TokenId:          tokenId,
 		UseTime:          useTimeSeconds,
 		IsStream:         isStream,
+		Other:            otherStr,
 	}
 	err := DB.Create(log).Error
 	if err != nil {

model/option.go

@@ -88,6 +88,7 @@ func InitOptionMap() {
 	common.OptionMap["ModelRatio"] = common.ModelRatio2JSONString()
 	common.OptionMap["ModelPrice"] = common.ModelPrice2JSONString()
 	common.OptionMap["GroupRatio"] = common.GroupRatio2JSONString()
+	common.OptionMap["CompletionRatio"] = common.CompletionRatio2JSONString()
 	common.OptionMap["TopUpLink"] = common.TopUpLink
 	common.OptionMap["ChatLink"] = common.ChatLink
 	common.OptionMap["ChatLink2"] = common.ChatLink2
@@ -305,6 +306,8 @@ func updateOptionMap(key string, value string) (err error) {
 		err = common.UpdateModelRatioByJSONString(value)
 	case "GroupRatio":
 		err = common.UpdateGroupRatioByJSONString(value)
+	case "CompletionRatio":
+		err = common.UpdateCompletionRatioByJSONString(value)
 	case "ModelPrice":
 		err = common.UpdateModelPriceByJSONString(value)
 	case "TopUpLink":

model/pricing.go

@@ -0,0 +1,72 @@
+package model
+
+import (
+	"one-api/common"
+	"one-api/dto"
+	"sync"
+	"time"
+)
+
+var (
+	pricingMap         []dto.ModelPricing
+	lastGetPricingTime time.Time
+	updatePricingLock  sync.Mutex
+)
+
+func GetPricing(user *User, openAIModels []dto.OpenAIModels) []dto.ModelPricing {
+	updatePricingLock.Lock()
+	defer updatePricingLock.Unlock()
+
+	if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 {
+		updatePricing(openAIModels)
+	}
+	if user != nil {
+		userPricingMap := make([]dto.ModelPricing, 0)
+		models := GetGroupModels(user.Group)
+		for _, pricing := range pricingMap {
+			if !common.StringsContains(models, pricing.ModelName) {
+				pricing.Available = false
+			}
+			userPricingMap = append(userPricingMap, pricing)
+		}
+		return userPricingMap
+	}
+	return pricingMap
+}
+
+func updatePricing(openAIModels []dto.OpenAIModels) {
+	modelRatios := common.GetModelRatios()
+	enabledModels := GetEnabledModels()
+	allModels := make(map[string]string)
+	for _, openAIModel := range openAIModels {
+		if common.StringsContains(enabledModels, openAIModel.Id) {
+			allModels[openAIModel.Id] = openAIModel.OwnedBy
+		}
+	}
+	for model, _ := range modelRatios {
+		if common.StringsContains(enabledModels, model) {
+			if _, ok := allModels[model]; !ok {
+				allModels[model] = "custom"
+			}
+		}
+	}
+	pricingMap = make([]dto.ModelPricing, 0)
+	for model, ownerBy := range allModels {
+		pricing := dto.ModelPricing{
+			Available: true,
+			ModelName: model,
+			OwnerBy:   ownerBy,
+		}
+		modelPrice, findPrice := common.GetModelPrice(model, false)
+		if findPrice {
+			pricing.ModelPrice = modelPrice
+			pricing.QuotaType = 1
+		} else {
+			pricing.ModelRatio = common.GetModelRatio(model)
+			pricing.CompletionRatio = common.GetCompletionRatio(model)
+			pricing.QuotaType = 0
+		}
+		pricingMap = append(pricingMap, pricing)
+	}
+	lastGetPricingTime = time.Now()
+}

model/usedata.go

@@ -45,6 +45,7 @@ func logQuotaDataCache(userId int, username string, modelName string, quota int,
 	if ok {
 		quotaData.Count += 1
 		quotaData.Quota += quota
+		quotaData.TokenUsed += tokenUsed
 	} else {
 		quotaData = &QuotaData{
 			UserID:    userId,