322 lines
11 KiB
Go
322 lines
11 KiB
Go
package billing
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"math"
|
|
"math/big"
|
|
"strings"
|
|
"unicode"
|
|
|
|
"juhe-factory/api/internal/model"
|
|
|
|
"github.com/google/uuid"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
const (
|
|
TextBillingPerRequest = "per_request"
|
|
TextBillingPerToken = "per_token"
|
|
)
|
|
|
|
type TextPricing struct {
|
|
Mode string `json:"mode"`
|
|
PerRequest string `json:"per_request,omitempty"`
|
|
InputPerM string `json:"input_per_m,omitempty"`
|
|
OutputPerM string `json:"output_per_m,omitempty"`
|
|
}
|
|
|
|
type TextUsage struct {
|
|
Input int64
|
|
Output int64
|
|
}
|
|
|
|
func CreateTextGenerationTask(tx *gorm.DB, task *model.GenerationTask, pricing TextPricing, remark string) error {
|
|
if err := pricing.Validate(); err != nil {
|
|
return err
|
|
}
|
|
task.BillingSnapshot = pricing.MarshalSnapshot()
|
|
if pricing.Mode == TextBillingPerRequest {
|
|
return PrechargeGenerationTask(tx, task, pricing.PerRequest, 1, remark)
|
|
}
|
|
task.EstimatedPoints = "0.00"
|
|
task.PrepaidPoints = "0.00"
|
|
if task.ID == uuid.Nil {
|
|
task.ID = uuid.New()
|
|
}
|
|
return tx.Create(task).Error
|
|
}
|
|
|
|
func ChargeTextRequest(tx *gorm.DB, userID uuid.UUID, referenceID string, pricing TextPricing, remark string) (string, error) {
|
|
if pricing.Mode != TextBillingPerRequest {
|
|
return "", errors.New("文本模型不是按次计费")
|
|
}
|
|
businessID := textCallBusinessID(referenceID, "request")
|
|
idempotencyKey := "generation:hold:" + businessID
|
|
if existing, found, err := debitAmount(tx, idempotencyKey); err != nil {
|
|
return "", err
|
|
} else if found {
|
|
return existing, nil
|
|
}
|
|
if pricing.PerRequest == "0" || pricing.PerRequest == "0.00" || pricing.PerRequest == "0.0000" {
|
|
return "0.00", nil
|
|
}
|
|
_, err := DebitPoints(tx, userID, pricing.PerRequest, "generation_hold", businessID, remark, idempotencyKey)
|
|
return pricing.PerRequest, err
|
|
}
|
|
|
|
func RefundTextGenerationTask(tx *gorm.DB, task *model.GenerationTask, remark string) (bool, error) {
|
|
pricing, err := ParseTextPricingSnapshot(task.BillingSnapshot)
|
|
if err != nil || pricing.Mode == TextBillingPerRequest {
|
|
return RefundGenerationTask(tx, task, remark)
|
|
}
|
|
if task.CostRefunded {
|
|
return false, nil
|
|
}
|
|
refunded, err := RefundTextCalls(tx, task.UserID, task.ID.String(), remark)
|
|
if err == nil {
|
|
task.CostRefunded = true
|
|
}
|
|
return refunded, err
|
|
}
|
|
|
|
func LoadTextPricing(db *gorm.DB, modelID uuid.UUID) (TextPricing, error) {
|
|
var pricing TextPricing
|
|
err := db.Table("models m").
|
|
Select(`coalesce(m.text_billing_mode,'') AS mode,
|
|
coalesce(max(CASE WHEN lower(p.price_key)='default' THEN p.price::text END),'') AS per_request,
|
|
coalesce(max(CASE WHEN lower(p.price_key)='input' THEN p.price::text END),'') AS input_per_m,
|
|
coalesce(max(CASE WHEN lower(p.price_key)='output' THEN p.price::text END),'') AS output_per_m`).
|
|
Joins("LEFT JOIN model_prices p ON p.model_id=m.id").
|
|
Where("m.id=? AND m.model_type='text' AND m.enabled=true AND m.deleted_at IS NULL", modelID).
|
|
Group("m.id,m.text_billing_mode").
|
|
Take(&pricing).Error
|
|
if err != nil {
|
|
return pricing, err
|
|
}
|
|
if err := pricing.Validate(); err != nil {
|
|
return pricing, err
|
|
}
|
|
return pricing, nil
|
|
}
|
|
|
|
func (pricing TextPricing) Validate() error {
|
|
switch pricing.Mode {
|
|
case TextBillingPerRequest:
|
|
if !validNonNegativeDecimal(pricing.PerRequest) {
|
|
return errors.New("文本模型未配置按次价格")
|
|
}
|
|
case TextBillingPerToken:
|
|
if !validNonNegativeDecimal(pricing.InputPerM) || !validNonNegativeDecimal(pricing.OutputPerM) {
|
|
return errors.New("文本模型未完整配置输入、输出 Token 价格")
|
|
}
|
|
default:
|
|
return errors.New("文本模型计费模式无效")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (pricing TextPricing) MarshalSnapshot() json.RawMessage {
|
|
data, _ := json.Marshal(pricing)
|
|
return data
|
|
}
|
|
|
|
func ParseTextPricingSnapshot(value json.RawMessage) (TextPricing, error) {
|
|
var pricing TextPricing
|
|
if len(value) == 0 || string(value) == "{}" {
|
|
return pricing, errors.New("文本模型计费快照缺失")
|
|
}
|
|
if err := json.Unmarshal(value, &pricing); err != nil {
|
|
return pricing, errors.New("文本模型计费快照无效")
|
|
}
|
|
return pricing, pricing.Validate()
|
|
}
|
|
|
|
func CalculateTextPoints(pricing TextPricing, usage TextUsage) (string, error) {
|
|
if pricing.Mode != TextBillingPerToken || usage.Input < 0 || usage.Output < 0 {
|
|
return "", errors.New("Token 计费参数无效")
|
|
}
|
|
inputRate, ok := new(big.Rat).SetString(pricing.InputPerM)
|
|
if !ok || inputRate.Sign() < 0 {
|
|
return "", errors.New("输入 Token 价格无效")
|
|
}
|
|
outputRate, ok := new(big.Rat).SetString(pricing.OutputPerM)
|
|
if !ok || outputRate.Sign() < 0 {
|
|
return "", errors.New("输出 Token 价格无效")
|
|
}
|
|
amount := new(big.Rat).Mul(inputRate, new(big.Rat).SetInt64(usage.Input))
|
|
amount.Add(amount, new(big.Rat).Mul(outputRate, new(big.Rat).SetInt64(usage.Output)))
|
|
amount.Quo(amount, big.NewRat(1000000, 1))
|
|
return ceilPointCents(amount), nil
|
|
}
|
|
|
|
func EstimateTextTokens(value string) int64 {
|
|
var tokens int64
|
|
ascii := 0
|
|
flush := func() {
|
|
if ascii > 0 {
|
|
tokens += int64((ascii + 3) / 4)
|
|
ascii = 0
|
|
}
|
|
}
|
|
for _, r := range value {
|
|
if r <= unicode.MaxASCII {
|
|
ascii++
|
|
} else {
|
|
flush()
|
|
tokens++
|
|
}
|
|
}
|
|
flush()
|
|
if tokens == 0 && value != "" {
|
|
return 1
|
|
}
|
|
return tokens
|
|
}
|
|
|
|
func ReserveTextCall(tx *gorm.DB, userID uuid.UUID, referenceID, callKey string, pricing TextPricing, estimatedInputTokens int64, remark string) (string, error) {
|
|
if pricing.Mode != TextBillingPerToken || estimatedInputTokens < 0 {
|
|
return "", errors.New("Token 预授权参数无效")
|
|
}
|
|
amount, err := CalculateTextReservePoints(pricing, estimatedInputTokens)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
businessID := textCallBusinessID(referenceID, callKey) + ":reserve"
|
|
idempotencyKey := "text:reserve:" + businessID
|
|
if existing, found, err := debitAmount(tx, idempotencyKey); err != nil {
|
|
return "", err
|
|
} else if found {
|
|
return existing, nil
|
|
}
|
|
if amount == "0.00" {
|
|
return amount, nil
|
|
}
|
|
_, err = DebitPoints(tx, userID, amount, "text_token_reserve", businessID, remark+"预授权", idempotencyKey)
|
|
return amount, err
|
|
}
|
|
|
|
func CalculateTextReservePoints(pricing TextPricing, estimatedInputTokens int64) (string, error) {
|
|
if estimatedInputTokens < 0 || estimatedInputTokens > math.MaxInt64/2 {
|
|
return "", errors.New("Token 预授权参数无效")
|
|
}
|
|
return CalculateTextPoints(pricing, TextUsage{Input: estimatedInputTokens, Output: estimatedInputTokens * 2})
|
|
}
|
|
|
|
func SettleTextCall(tx *gorm.DB, userID uuid.UUID, referenceID, callKey string, pricing TextPricing, usage TextUsage, remark string) (string, error) {
|
|
if usage.Input <= 0 && usage.Output <= 0 {
|
|
return "", errors.New("上游未返回 Token 用量，无法完成计费")
|
|
}
|
|
actual, err := CalculateTextPoints(pricing, usage)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
baseID := textCallBusinessID(referenceID, callKey)
|
|
reserveID := baseID + ":reserve"
|
|
if _, found, err := debitAmount(tx, "text:reserve:"+reserveID); err != nil {
|
|
return "", err
|
|
} else if found {
|
|
if _, err := RefundDebit(tx, userID, "text:reserve:"+reserveID, "text:reserve-refund:"+reserveID,
|
|
"text_token_reserve_refund", reserveID, remark+"预授权释放"); err != nil {
|
|
return "", err
|
|
}
|
|
}
|
|
chargeKey := "generation:hold:" + baseID
|
|
if existing, found, err := debitAmount(tx, chargeKey); err != nil {
|
|
return "", err
|
|
} else if found {
|
|
return existing, nil
|
|
}
|
|
if actual == "0.00" {
|
|
return actual, nil
|
|
}
|
|
_, err = DebitPoints(tx, userID, actual, "generation_hold", baseID, remark, chargeKey)
|
|
return actual, err
|
|
}
|
|
|
|
func RefundTextCalls(tx *gorm.DB, userID uuid.UUID, referenceID, remark string) (bool, error) {
|
|
type debit struct {
|
|
BusinessType string
|
|
BusinessID string
|
|
IdempotencyKey string
|
|
}
|
|
debits := make([]debit, 0)
|
|
if err := tx.Table("point_ledger").
|
|
Select("business_type,business_id,idempotency_key").
|
|
Where("user_id=? AND change_amount<0 AND business_id LIKE ? AND business_type IN ?", userID, referenceID+":text:%", []string{"generation_hold", "text_token_reserve"}).
|
|
Order("id").Find(&debits).Error; err != nil {
|
|
return false, err
|
|
}
|
|
refunded := false
|
|
for _, debit := range debits {
|
|
refundKey := "generation:refund:" + debit.BusinessID
|
|
refundType := "generation_refund"
|
|
if debit.BusinessType == "text_token_reserve" {
|
|
refundKey = "text:reserve-refund:" + debit.BusinessID
|
|
refundType = "text_token_reserve_refund"
|
|
}
|
|
created, err := RefundDebit(tx, userID, debit.IdempotencyKey, refundKey, refundType, debit.BusinessID, remark)
|
|
if err != nil {
|
|
return refunded, err
|
|
}
|
|
refunded = refunded || created
|
|
}
|
|
return refunded, nil
|
|
}
|
|
|
|
func CommitTextReserves(tx *gorm.DB, userID uuid.UUID, referenceID, remark string) (string, error) {
|
|
type reserve struct {
|
|
BusinessID string
|
|
IdempotencyKey string
|
|
}
|
|
items := make([]reserve, 0)
|
|
if err := tx.Table("point_ledger").Select("business_id,idempotency_key").
|
|
Where("user_id=? AND business_type='text_token_reserve' AND business_id LIKE ? AND NOT EXISTS (SELECT 1 FROM point_ledger refund WHERE refund.idempotency_key='text:reserve-refund:'||point_ledger.business_id)", userID, referenceID+":text:%").
|
|
Find(&items).Error; err != nil {
|
|
return "", err
|
|
}
|
|
for _, item := range items {
|
|
baseID := strings.TrimSuffix(item.BusinessID, ":reserve")
|
|
if err := tx.Exec("UPDATE point_ledger SET business_type='generation_hold',business_id=?,idempotency_key=?,remark=? WHERE idempotency_key=?", baseID, "generation:hold:"+baseID, remark, item.IdempotencyKey).Error; err != nil {
|
|
return "", err
|
|
}
|
|
}
|
|
var amount string
|
|
if err := tx.Table("point_ledger").Select("coalesce(sum(-change_amount),0)::text").
|
|
Where("user_id=? AND business_type='generation_hold' AND business_id LIKE ?", userID, referenceID+":text:%").Scan(&amount).Error; err != nil {
|
|
return "", err
|
|
}
|
|
return amount, nil
|
|
}
|
|
|
|
func validNonNegativeDecimal(value string) bool {
|
|
number, ok := new(big.Rat).SetString(strings.TrimSpace(value))
|
|
return ok && number.Sign() >= 0
|
|
}
|
|
|
|
func ceilPointCents(value *big.Rat) string {
|
|
if value == nil || value.Sign() <= 0 {
|
|
return "0.00"
|
|
}
|
|
cents := new(big.Rat).Mul(value, big.NewRat(100, 1))
|
|
quotient := new(big.Int).Quo(cents.Num(), cents.Denom())
|
|
if new(big.Int).Mod(cents.Num(), cents.Denom()).Sign() > 0 {
|
|
quotient.Add(quotient, big.NewInt(1))
|
|
}
|
|
return fmt.Sprintf("%d.%02d", new(big.Int).Quo(quotient, big.NewInt(100)), new(big.Int).Mod(quotient, big.NewInt(100)))
|
|
}
|
|
|
|
func textCallBusinessID(referenceID, callKey string) string {
|
|
return referenceID + ":text:" + strings.TrimSpace(callKey)
|
|
}
|
|
|
|
func debitAmount(tx *gorm.DB, idempotencyKey string) (string, bool, error) {
|
|
var amount string
|
|
result := tx.Table("point_ledger").Select("(-change_amount)::text").Where("idempotency_key=? AND change_amount<0", idempotencyKey).Limit(1).Scan(&amount)
|
|
if result.Error != nil {
|
|
return "", false, result.Error
|
|
}
|
|
return amount, strings.TrimSpace(amount) != "", nil
|
|
}
|