fix daily usage

This commit is contained in:
Sakurasan
2025-04-22 01:50:50 +08:00
parent 000162b1b1
commit 24529189d9

View File

@@ -144,6 +144,8 @@ func (p *Proxy) ProcessUsage() {
func (p *Proxy) Do(llmusage *llm.TokenUsage) error { func (p *Proxy) Do(llmusage *llm.TokenUsage) error {
err := p.db.Transaction(func(tx *gorm.DB) error { err := p.db.Transaction(func(tx *gorm.DB) error {
now := time.Now() now := time.Now()
today, _ := time.Parse("2006-01-02", now.Format("2006-01-02"))
cost := tokenizer.Cost(llmusage.Model, llmusage.PromptTokens, llmusage.CompletionTokens) cost := tokenizer.Cost(llmusage.Model, llmusage.PromptTokens, llmusage.CompletionTokens)
token, err := p.tokenDAO.GetByID(p.ctx, llmusage.TokenID) token, err := p.tokenDAO.GetByID(p.ctx, llmusage.TokenID)
if err != nil { if err != nil {
@@ -166,32 +168,32 @@ func (p *Proxy) Do(llmusage *llm.TokenUsage) error {
return fmt.Errorf("create usage error: %w", err) return fmt.Errorf("create usage error: %w", err)
} }
// 2. 更新每日统计upsert 操作) // 2. 更新每日统计
// dailyUsage := model.DailyUsage{ var dailyUsage model.DailyUsage
// UserID: llmusage.User.ID, result := tx.WithContext(p.ctx).Where("user_id = ? and date = ?", llmusage.User.ID, today).First(&dailyUsage)
// TokenID: llmusage.TokenID, if result.RowsAffected == 0 {
// Date: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()), dailyUsage.UserID = llmusage.User.ID
// Model: llmusage.Model, dailyUsage.TokenID = llmusage.TokenID
// Stream: llmusage.Stream, dailyUsage.Date = today
// PromptTokens: llmusage.PromptTokens, dailyUsage.Model = llmusage.Model
// CompletionTokens: llmusage.CompletionTokens, dailyUsage.Stream = llmusage.Stream
// TotalTokens: llmusage.TotalTokens, dailyUsage.PromptTokens = llmusage.PromptTokens
// Cost: fmt.Sprintf("%.8f", cost), dailyUsage.CompletionTokens = llmusage.CompletionTokens
// } dailyUsage.TotalTokens = llmusage.TotalTokens
dailyUsage.Cost = fmt.Sprintf("%.8f", cost)
p.dailyUsageDAO.UpsertDailyUsage(p.ctx, usage) if err := tx.WithContext(p.ctx).Create(&dailyUsage).Error; err != nil {
// 使用 OnConflict 实现 upsert return fmt.Errorf("create daily usage error: %w", err)
// if err := tx.WithContext(p.ctx).Clauses(clause.OnConflict{ }
// Columns: []clause.Column{{Name: "user_id"}, {Name: "date"}}, // 唯一键 } else {
// DoUpdates: clause.Assignments(map[string]interface{}{ if err := tx.WithContext(p.ctx).Model(&model.DailyUsage{}).Where("user_id = ? and date = ?", llmusage.User.ID, today).
// "prompt_tokens": gorm.Expr("prompt_tokens + ?", llmusage.PromptTokens), Updates(map[string]interface{}{
// "completion_tokens": gorm.Expr("completion_tokens + ?", llmusage.CompletionTokens), "prompt_tokens": gorm.Expr("prompt_tokens + ?", llmusage.PromptTokens),
// "total_tokens": gorm.Expr("total_tokens + ?", llmusage.TotalTokens), "completion_tokens": gorm.Expr("completion_tokens + ?", llmusage.CompletionTokens),
// "cost": gorm.Expr("cost + ?", fmt.Sprintf("%.8f", cost)), "total_tokens": gorm.Expr("total_tokens + ?", llmusage.TotalTokens),
// }), }).Error; err != nil {
// }).Create(&dailyUsage).Error; err != nil { return fmt.Errorf("update daily usage error: %w", err)
// return fmt.Errorf("upsert daily usage error: %w", err) }
// } }
// 3. 更新用户额度 // 3. 更新用户额度
if *llmusage.User.UnlimitedQuota { if *llmusage.User.UnlimitedQuota {