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