wechat-api/hook/credit/models.go

package credit

import (
	"fmt"
	"math"
	"regexp"
	"sort"
	"strings"
	"wechat-api/ent/custom_types"
	"wechat-api/hook/dify"
)

var modelArray = []string{
	"o1",
	"gpt-4o",
	"gpt-4.1",
	"o3-mini",
	"moonshot-v1-32k",
	"deepseek-r1",
	"moonshot-v1-8k",
	"gpt-4.1-mini",
	"gpt-3.5-turbo",
	"qwen-max",
	"doubao1.5-pro-256k",
	"deepseek-v3",
	"qwq-32b-preview",
	"gpt-4o-mini",
	"qwen2.5-14b-instruct-1m",
	"gpt-4.1-nano",
	"doubao1.5-pro",
	"doubao1.5-pro-32k",
	"chatglm3",
	"qwen-turbo",
	"doubao1.5-lite-32k",
}

var priceArray = []float64{
	0.01,
	0.001667,
	0.001333,
	0.000733,
	0.000548,
	0.000365,
	0.000274,
	0.000267,
	0.00025,
	0.000219,
	0.000205,
	0.000183,
	0.000137,
	0.0001,
	0.000068,
	0.000067,
	0.000046,
	0.000046,
	0.000023,
	0.000014,
	0.000014,
}

func getModelName(modelName string) string {
	// 将字符串转换为小写
	return strings.ToLower(modelName)
}

func GetModelPrice(modelName string) (model string, price float64) {
	difyModelName := getModelName(modelName)
	for i, v := range modelArray {
		if v == difyModelName {
			return v, priceArray[i]
		}
	}

	return modelArray[0], priceArray[0]
}

func ComputePrice(price float64, tokens uint64) float64 {
	scale := float64(1000000)
	return math.Round(price*float64(tokens)*scale) / scale
}

// Subtraction 保留小数点后6位的精确减法
func Subtraction(number1, number2 float64) float64 {
	d1 := number1 * 1000000
	d2 := number2 * 1000000
	res := math.Floor(d1-d2) / 1000000
	return res
}

func ComputeModelPrice(response interface{}) (model string, price float64) {
	fmt.Printf("response=%v  \n", response)

	// 先获取所有本次响应里所有的model
	modelInputArray := make([]string, 0)

	// 如果是 gpt/chat/submit 处调用这里会是非正常的结构，所以要提前处理
	if _, ok := response.(dify.ChatResp); ok {
		return modelArray[0], priceArray[0]
	}

	// 如果是标准 VResponse 结构
	if _, ok := response.(custom_types.VResponse); ok {
		if modelName := response.(custom_types.VResponse).Model; modelName != "" {
			modelInputArray = append(modelInputArray, modelName)
		}

		// 获取 responseData 下所有model
		responseDataVal := response.(custom_types.VResponse).ResponseData
		if responseDataVal != nil {
			for _, res := range responseDataVal {
				fmt.Printf("model=%v  \n", res.Model)
				if res.Model == "" {
					continue
				} else {
					modelInputArray = append(modelInputArray, strings.ToLower(res.Model))
				}
			}
		}
	}

	fmt.Printf("modelInputArray=%v \n", modelInputArray)
	// 如果精确搜索能找到model，返回model及其价格
	if len(modelInputArray) > 0 {
		for i, providedModelName := range modelInputArray {
			for _, modelName := range modelArray {
				if modelName == providedModelName {
					return modelName, priceArray[i]
				}
			}
		}
	}

	// 如果精确搜索搜不到，则把所有model按从长倒短，挨个匹配下，看有无符合条件
	// 如果模型里带日期，这时候要在筛选一次，比如 gpt-4.1-mini-2024-07-18 其实应该匹配到 gpt-4.1-mini
	// 对模型名数组按名字长度排序，然后按长到短匹配
	sortedModelArray := make([]string, len(modelArray))
	copy(sortedModelArray, modelArray)

	// 使用長度降序排序模型陣列
	sort.Slice(sortedModelArray, func(i, j int) bool {
		return len(sortedModelArray[i]) > len(sortedModelArray[j])
	})

	// 排序后按长到短挨个匹配合适的model
	if len(sortedModelArray) > 0 {
		for _, model := range sortedModelArray {
			for _, modelName := range modelInputArray {
				re := regexp.MustCompile(`(?i)` + regexp.QuoteMeta(model))
				if re.MatchString(modelName) {
					fmt.Printf("In models.go the given model_name: %s match: %s\n", modelName, model)
					model, price := GetModelPrice(model)
					return model, price
				}
			}
		}
	}

	// 最后返回最高的价格
	return modelArray[0], priceArray[0]
}