sundynix-agentix/sundynix-gateway/internal/store/usage.go

package store

import (
	"context"
	"os"
	"strconv"
	"time"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	"github.com/sundynix/sundynix-shared/contract"
)

// UsageEvent 是一条持久化的任务用量明细（追加式，只增不改）——计费 / 对账 / 重算的事实源。
// 由网关消费 dispatcher 回写的 contract.UsageEvent 落库；tenant/owner 从事件显式带（background ctx 无请求租户）。
// 表名 sundynix_usage_event。对 task_id 唯一 → 幂等（防重投重复计费）。
type UsageEvent struct {
	BaseModel
	TenantID     string `gorm:"size:64;index"`
	Owner        string `gorm:"size:64;index"`      // 提交者 user.id（= 事件 UserID）
	TaskID       string `gorm:"size:64;uniqueIndex"` // 一任务一计量 → 幂等键
	Model        string `gorm:"size:64"`            // 计费模型名（空=按激活 chat 模型近似）
	PromptTok    int
	CompTok      int
	TotalTok     int
	CreditsMicro int64  `gorm:"column:credits_micro"` // 折算积分 ×10⁻⁶
	CostMicros   int64  `gorm:"column:cost_micros"`   // 折算金额（币种最小单位 ×10⁻⁶）
	Currency     string `gorm:"size:8"`
	Exceeded     bool
	TS           int64
}

func (UsageEvent) TableName() string { return "sundynix_usage_event" }
func (UsageEvent) isTenantScoped()   {}

// TokensPerCredit 读 token→积分 汇率（env TOKENS_PER_CREDIT，缺省/非法=1000）。设 1 即 token 直计。
func TokensPerCredit() float64 {
	if v := os.Getenv("TOKENS_PER_CREDIT"); v != "" {
		if n, err := strconv.ParseFloat(v, 64); err == nil && n > 0 {
			return n
		}
	}
	return 1000
}

// SaveUsageEvent 折算 credits + cost 并幂等落一条用量明细。
// 折算模型：ev.Model 优先；为空则回退当前激活 chat 模型（近似——忽略 failover 到备用模型的情形）。
// 缺 Pricing → cost=0、weight=1（计量不因缺价而丢量，可事后补价重算）。
func (p *Postgres) SaveUsageEvent(ctx context.Context, ev *contract.UsageEvent) error {
	if p.db == nil {
		return nil
	}
	model := ev.Model
	if model == "" {
		if cfg := p.ActiveConfig(ctx, contract.ConfigKindChat); cfg != nil {
			model = cfg.Model
		}
	}

	weight := 1.0
	currency := ""
	var costMicros int64
	if pr := p.pricingForModelName(ctx, model); pr != nil {
		if pr.CreditWeight > 0 {
			weight = pr.CreditWeight
		}
		currency = pr.Currency
		// cost（币种单位）= tok/1000 * per1k；×10⁶ 存微单位（整数）。
		cost := float64(ev.PromptTok)/1000*pr.InputPer1K + float64(ev.CompTok)/1000*pr.OutputPer1K
		costMicros = int64(cost * 1e6)
	}
	// credits_micro = total_tok / tokensPerCredit * weight，×10⁶ 存微积分。
	creditsMicro := int64(float64(ev.TotalTok) / TokensPerCredit() * weight * 1e6)

	row := &UsageEvent{
		TenantID: ev.TenantID, Owner: ev.UserID, TaskID: ev.TaskID, Model: model,
		PromptTok: ev.PromptTok, CompTok: ev.CompTok, TotalTok: ev.TotalTok,
		CreditsMicro: creditsMicro, CostMicros: costMicros, Currency: currency,
		Exceeded: ev.Exceeded, TS: ev.TS,
	}
	day := time.UnixMilli(ev.TS).Format("20060102")

	// 一个事务内：落明细 → 扣积分(账本+余额) → 累加 rollup。
	// 幂等锚点：usage_event 的 task_id 唯一；插入若被冲突吞掉(RowsAffected==0)→重投，跳过后续，绝不重复计费。
	return p.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		res := tx.Clauses(clause.OnConflict{
			Columns:   []clause.Column{{Name: "task_id"}},
			DoNothing: true,
		}).Create(row)
		if res.Error != nil {
			return res.Error
		}
		if res.RowsAffected == 0 {
			return nil // 重投：明细已存在，账本/余额/rollup 均不再动
		}
		if err := applyUsageCredit(tx, ev.TenantID, ev.TaskID, creditsMicro); err != nil {
			return err
		}
		return upsertRollup(tx, ev.TenantID, day, int64(ev.TotalTok), creditsMicro, costMicros, currency)
	})
}

// pricingForModelName 按模型名查计价（join model 表，pricing 以 model_id 关联）。查不到返回 nil。
func (p *Postgres) pricingForModelName(ctx context.Context, name string) *Pricing {
	if p.db == nil || name == "" {
		return nil
	}
	var pr Pricing
	err := p.db.WithContext(ctx).
		Joins("JOIN sundynix_model m ON m.id = sundynix_pricing.model_id").
		Where("m.model = ?", name).First(&pr).Error
	if err != nil {
		return nil
	}
	return &pr
}