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

fix usage

Browse Source
This commit is contained in:
Sakurasan
2025-04-22 01:06:03 +08:00
parent 6662ea5e04
commit 000162b1b1
2 changed files with 29 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

2
llm/openai_compatible/chat.go
View File

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