mirrors2/opencatd-open
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix daily usage

Browse Source
This commit is contained in:
Sakurasan
2025-04-22 01:50:50 +08:00
parent 000162b1b1
commit 24529189d9
1 changed files with 28 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
internal/controller/proxy/proxy.go
View File

@@ -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 {