model-gateway/common/util/billing.go

package util

import (
	"context"
	"fmt"
	"math"
	"model-gateway/service/gateway"
	"strings"

	"github.com/gogf/gf/v2/encoding/gjson"
	"github.com/gogf/gf/v2/util/gconv"
)

// ======================== 计费入口 ========================

func CalculateBilling(config map[string]any, billingData map[string]any) map[string]any {
	if len(config) == 0 {
		return nil
	}
	switch config["type"] {
	case "inference_tier": //推理模型计费
		return calculateInferenceTierBilling(config, billingData)
	case "video_resolution": //视频模型计费
		return calculateVideoResolutionBilling(config, billingData)
	}
	return nil
}

// ======================== 推理模型计费 ========================

func calculateInferenceTierBilling(config map[string]any, data map[string]any) map[string]any {
	promptTokens := gconv.Int64(data["prompt_tokens"])
	completionTokens := gconv.Int64(data["completion_tokens"])
	hasAudio := gconv.Bool(data["has_audio"])
	inputK := promptTokens / 1000

	tiers := config["pricing"].(map[string]any)["tiers"].([]any)
	var matched map[string]any
	for _, t := range tiers {
		tier := t.(map[string]any)
		if inputK >= gconv.Int64(tier["input_min"]) && inputK <= gconv.Int64(tier["input_max"]) {
			matched = tier
			break
		}
	}
	if matched == nil {
		return nil
	}

	var inputPrice float64
	if hasAudio && matched["audio_input_price"] != nil {
		inputPrice = gconv.Float64(matched["audio_input_price"])
	} else {
		inputPrice = gconv.Float64(matched["input_price"])
	}
	outputPrice := gconv.Float64(matched["output_price"])

	inputCost := float64(promptTokens) * inputPrice / 1000000
	outputCost := float64(completionTokens) * outputPrice / 1000000

	// 推理模型
	return map[string]any{
		"model_name":   data["model_name"],
		"total_tokens": promptTokens + completionTokens,
		"total_fee":    inputCost + outputCost,
		// 明细
		"prompt_tokens":     promptTokens,
		"completion_tokens": completionTokens,
		"has_audio":         hasAudio,
		"input_tier":        fmt.Sprintf("[%v, %v]", matched["input_min"], matched["input_max"]),
		"input_unit_price":  inputPrice,
		"output_unit_price": outputPrice,
		"input_cost":        inputCost,
		"output_cost":       outputCost,
	}
}

// ======================== 视频模型计费 ========================

func calculateVideoResolutionBilling(config map[string]any, data map[string]any) map[string]any {
	pricingPath := buildPricingPath(config, data)

	pricing := config["pricing"].(map[string]any)
	unitPrice := gconv.Float64(pricing[pricingPath])

	var effectiveMinToken float64
	if gconv.Bool(data["input_has_video"]) {
		effectiveMinToken = matchMinToken(config, data)
	}

	completionTokens := gconv.Float64(data["actual_tokens"])
	realChargeTokens := int64(math.Max(completionTokens, effectiveMinToken))
	totalFee := float64(realChargeTokens) * unitPrice / 1000000

	if gconv.Bool(config["enable_audio_charge"]) {
		totalFee += gconv.Float64(data["audio_word_cnt"]) * gconv.Float64(config["audio_unit_price"])
	}

	// 视频模型
	return map[string]any{
		"model_name":   data["model_name"],
		"total_tokens": realChargeTokens,
		"total_fee":    totalFee,
		// 明细
		"prompt_tokens":       0,
		"completion_tokens":   int64(completionTokens),
		"is_online":           data["is_online"],
		"video_resolution":    data["video_resolution"],
		"input_has_video":     data["input_has_video"],
		"output_duration_sec": data["output_duration_sec"],
		"aspect_ratio":        data["aspect_ratio"],
		"resolution":          data["resolution"],
		"matched_path":        pricingPath,
		"token_unit_price":    unitPrice,
		"effective_min_token": effectiveMinToken,
	}
}

// ======================== 数据提取 ========================

func ExtractRequestBilling(ctx context.Context, config map[string]any, requestPayload map[string]any) map[string]any {
	dimensions := config["dimensions"].(map[string]any)
	fields := dimensions["request"].(map[string]any)
	data := make(map[string]any)

	for key, path := range fields {
		data[key] = extractValue(requestPayload, gconv.String(path))
	}

	if defaults, ok := config["defaults"].(map[string]any); ok {
		for k, v := range defaults {
			if _, exists := data[k]; !exists || data[k] == nil {
				data[k] = v
			}
		}
	}

	// 处理 compute 字段
	if compute, ok := config["compute"].(map[string]any); ok {
		for targetField, rule := range compute {
			r := rule.(map[string]any)
			dependsOn := gconv.String(r["depends_on"])
			dependsValue := gconv.Bool(r["depends_value"])

			if gconv.Bool(data[dependsOn]) != dependsValue {
				continue
			}

			switch r["service"] {
			case "video_duration":
				urls := extractVideoUrls(requestPayload)
				if len(urls) > 0 {
					resp, err := gateway.GetVideoDuration(ctx, urls)
					if err == nil {
						data[targetField] = resp.TotalDuration
					}
				}
			}
		}
	}
	return data
}

func ExtractResponseBilling(config map[string]any, response map[string]any) map[string]any {
	dimensions := config["dimensions"].(map[string]any)
	fields := dimensions["response"].(map[string]any)
	data := make(map[string]any)

	for key, path := range fields {
		data[key] = gjson.New(response).Get(gconv.String(path)).Val()
	}
	return data
}

// ======================== 内部辅助 ========================

func buildPricingPath(config map[string]any, data map[string]any) string {
	pathConfig := config["pricing_path"].(map[string]any)
	segments := pathConfig["segments"].([]any)
	mapping := pathConfig["mapping"].(map[string]any)

	var parts []string
	for _, seg := range segments {
		segName := gconv.String(seg)
		val := gconv.String(data[segName])
		if m, ok := mapping[segName].(map[string]any); ok {
			if mapped, exists := m[val]; exists {
				val = gconv.String(mapped)
			}
		}
		parts = append(parts, val)
	}
	return strings.Join(parts, ".")
}

func matchMinToken(config map[string]any, data map[string]any) float64 {
	rule, ok := config["min_token_rule"].(map[string]any)
	if !ok {
		return 0
	}

	conditionDefs := rule["conditions"].([]any)
	rows := rule["rows"].([]any)

	for _, row := range rows {
		r := row.(map[string]any)
		conditions := r["conditions"].([]any)
		matched := true

		for i, cond := range conditions {
			condStr := gconv.String(cond)
			condDef := conditionDefs[i].(map[string]any)
			condName := gconv.String(condDef["name"])

			switch condDef["type"] {
			case "range":
				if !matchRange(condStr, gconv.Float64(data[condName])) {
					matched = false
				}
			case "enum":
				if gconv.String(data[condName]) != condStr {
					matched = false
				}
			}
		}
		if matched {
			return gconv.Float64(r["min_token"])
		}
	}
	return 0
}

func matchRange(rangeStr string, val float64) bool {
	rangeStr = strings.Trim(rangeStr, "[]()")
	parts := strings.Split(rangeStr, ",")
	if len(parts) != 2 {
		return false
	}
	return val >= gconv.Float64(strings.TrimSpace(parts[0])) && val < gconv.Float64(strings.TrimSpace(parts[1]))
}

func extractValue(source map[string]any, path string) any {
	// 数组求和：rounds.#.duration
	if strings.HasPrefix(path, "rounds.#.") && !strings.Contains(path, "==") {
		field := strings.TrimPrefix(path, "rounds.#.")
		rounds := gjson.New(source).Get("rounds").Array()
		var sum float64
		for _, r := range rounds {
			if strings.Contains(field, "#") {
				parts := strings.SplitN(field, ".#.", 2)
				arr := gjson.New(r).Get(parts[0]).Array()
				for _, item := range arr {
					sum += gconv.Float64(gjson.New(item).Get(parts[1]).Val())
				}
			} else {
				sum += gconv.Float64(gjson.New(r).Get(field).Val())
			}
		}
		return sum
	}

	// 条件判断：xxx.#.type==yyy 或 rounds.#.content.#.type==yyy
	if strings.Contains(path, "==video_url") || strings.Contains(path, "==input_audio") {
		parts := strings.Split(path, "==")
		basePath := parts[0]
		typ := parts[1]
		segments := strings.Split(basePath, ".#.")

		// 单层：content.#.type → segments = ["content", "type"]
		if len(segments) == 2 && !strings.Contains(segments[0], ".") {
			arr := gjson.New(source).Get(segments[0]).Array()
			for _, item := range arr {
				if gconv.String(gjson.New(item).Get(segments[1]).Val()) == typ {
					return true
				}
			}
			return false
		}

		// 嵌套：rounds.#.content.#.type → segments = ["rounds", "content", "type"]
		topArr := gjson.New(source).Get(segments[0]).Array()
		for _, item := range topArr {
			subArr := gjson.New(item).Get(segments[1]).Array()
			for _, sub := range subArr {
				if gconv.String(gjson.New(sub).Get("type").Val()) == typ {
					return true
				}
			}
		}
		return false
	}

	return gjson.New(source).Get(path).Val()
}

func extractVideoUrls(body map[string]any) []string {
	var urls []string
	content := gjson.New(body).Get("content").Array()
	for _, c := range content {
		item := c.(map[string]any)
		if item["type"] == "video_url" {
			if v, ok := item["video_url"].(map[string]any); ok {
				if url, ok := v["url"].(string); ok {
					urls = append(urls, url)
				}
			}
		}
	}
	return urls
}