diff --git a/internal/controller/proxy/proxy.go b/internal/controller/proxy/proxy.go index 3c18750..8783ecd 100644 --- a/internal/controller/proxy/proxy.go +++ b/internal/controller/proxy/proxy.go @@ -26,7 +26,6 @@ import ( "github.com/lib/pq" "github.com/tidwall/gjson" "gorm.io/gorm" - "gorm.io/gorm/clause" ) type Proxy struct { @@ -150,8 +149,8 @@ func (p *Proxy) Do(llmusage *llm.TokenUsage) error { if err != nil { return err } - // 1. 记录使用记录 - if err := tx.WithContext(p.ctx).Create(&model.Usage{ + + usage := &model.Usage{ UserID: llmusage.User.ID, TokenID: llmusage.TokenID, Date: now, @@ -161,35 +160,38 @@ func (p *Proxy) Do(llmusage *llm.TokenUsage) error { CompletionTokens: llmusage.CompletionTokens, TotalTokens: llmusage.TotalTokens, Cost: fmt.Sprintf("%.8f", cost), - }).Error; err != nil { + } + // 1. 记录使用记录 + if err := tx.WithContext(p.ctx).Create(usage).Error; err != nil { 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), - } + // 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) - } + // 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) + // } // 3. 更新用户额度 if *llmusage.User.UnlimitedQuota { diff --git a/llm/openai_compatible/chat.go b/llm/openai_compatible/chat.go index 4214228..8ccea63 100644 --- a/llm/openai_compatible/chat.go +++ b/llm/openai_compatible/chat.go @@ -206,7 +206,7 @@ func (o *OpenAICompatible) StreamChat(ctx context.Context, chatReq llm.ChatReque if err := json.Unmarshal(line, &streamResp); err != nil { continue } - fmt.Printf("%#v\n", streamResp.Usage) + if streamResp.Usage != nil { o.tokenUsage.PromptTokens += streamResp.Usage.PromptTokens o.tokenUsage.CompletionTokens += streamResp.Usage.CompletionTokens